Conversation Pruner — Ruby source

# Conversation Pruner — compute a deterministic pruning plan for a
# token-budgeted chat history.
#
# Language: Ruby (3.x, zero dependencies)
# Port of src/lib/conversationPruner.ts (the canonical TypeScript
#           implementation) for the CosmoDev polyglot showcase
#           (slug: conversation-pruner).
#
# Given per-message token counts and a context budget, decide which messages
# to keep verbatim, which to fold into one running summary, and which to drop
# outright — protecting system messages, pinned turns, the first turn, and
# the current (last user) request.

module ConversationPruner
  # Summary compression model: fixed framing tokens.
  SUMMARY_FIXED_TOKENS = 60

  # Summary compression model: share of the folded content.
  SUMMARY_RATIO = 0.1

  Message = Struct.new(:role, :content, :tokens, :pinned, keyword_init: true) do
    def initialize(role:, content:, tokens:, pinned: false)
      super
    end
  end

  Decision = Struct.new(:index, :role, :action, :tokens, keyword_init: true)

  Plan = Struct.new(
    :decisions, :kept_tokens, :summarized_tokens, :dropped_tokens,
    :summary_cost_tokens, :projected_tokens, :fits_budget, :warnings,
    keyword_init: true
  )

  module_function

  # Group an integer with thousands separators (toLocaleString stand-in).
  def fmt(v)
    v.to_s.reverse.gsub(/(\d{3})(?=\d)/, '\1,').reverse
  end

  # Compute the pruning plan for +messages+ under +budget_tokens+.
  # Raises ArgumentError on a negative budget or any negative per-message
  # token count.
  def plan_prune(messages, budget_tokens:)
    raise ArgumentError, 'budgetTokens must be >= 0' if budget_tokens.negative?
    raise ArgumentError, 'message tokens must be >= 0' if messages.any? { |m| m.tokens.negative? }

    warnings = []
    n = messages.length
    last_user = messages.rindex { |m| m.role == 'user' } || -1

    # Untouchable: every system message, pinned messages, the first turn
    # (the opening user request), and the current request (the last user
    # message and everything after it).
    protected_idx = {}
    messages.each_with_index do |m, i|
      protected_idx[i] = true if m.role == 'system' || m.pinned
    end
    protected_idx[0] = true if n.positive?
    first_turn = messages.index { |m| m.role != 'system' }
    protected_idx[first_turn] = true unless first_turn.nil?
    from = last_user == -1 ? n - 1 : last_user
    ([from, 0].max...n).each { |i| protected_idx[i] = true }

    protected_tokens = protected_idx.keys.sum { |i| messages[i].tokens }
    if protected_tokens > budget_tokens
      warnings << "Protected messages alone are #{fmt(protected_tokens)} tokens against a " \
                  "#{fmt(budget_tokens)} budget — raise the budget (or reserve less for the " \
                  'reply) before pruning anything else.'
    end

    # Fill the remaining budget newest-to-oldest through the middle.
    actions = Array.new(n, 'drop')
    protected_idx.each_key { |i| actions[i] = 'keep' }
    used = protected_tokens
    (n - 1).downto(0) do |i|
      next unless actions[i] == 'drop'

      if used + messages[i].tokens <= budget_tokens
        actions[i] = 'keep'
        used += messages[i].tokens
      else
        break # oldest-unfilled remain drop/summarize candidates
      end
    end

    # Everything still 'drop' in the middle folds into ONE running summary
    # when the compressed form fits where the raw turns did not.
    summarize_idx = (0...n).select { |i| actions[i] == 'drop' && !protected_idx.key?(i) }
    summarize_tokens = summarize_idx.sum { |i| messages[i].tokens }
    attempted = summarize_idx.empty? ? 0 : SUMMARY_FIXED_TOKENS + (summarize_tokens * SUMMARY_RATIO).ceil

    # The summary only costs anything when it is actually applied.
    summary_cost = 0
    if attempted.positive? && used + attempted <= budget_tokens
      summarize_idx.each { |i| actions[i] = 'summarize' }
      summary_cost = attempted
      used += summary_cost
    elsif attempted.positive?
      warnings << "Even the compressed summary (#{fmt(attempted)} tokens) does not fit the " \
                  'remaining budget — the oldest turns are dropped instead.'
    end

    decisions = messages.each_with_index.map do |m, i|
      Decision.new(index: i, role: m.role, action: actions[i], tokens: m.tokens)
    end

    kept = dropped = folded = 0
    decisions.each do |d|
      case d.action
      when 'keep' then kept += d.tokens
      when 'drop' then dropped += d.tokens
      else folded += d.tokens
      end
    end

    Plan.new(
      decisions: decisions, kept_tokens: kept, summarized_tokens: folded,
      dropped_tokens: dropped, summary_cost_tokens: summary_cost,
      projected_tokens: kept + summary_cost,
      fits_budget: kept + summary_cost <= budget_tokens,
      warnings: warnings
    )
  end

  # Human-readable one-line summary of a plan.
  def describe_prune(plan)
    unless plan.fits_budget
      return "Does not fit: #{fmt(plan.projected_tokens)} tokens projected against the budget."
    end

    parts = ["#{fmt(plan.kept_tokens)} kept"]
    if plan.summarized_tokens.positive?
      parts << "#{fmt(plan.summarized_tokens)} folded into a #{fmt(plan.summary_cost_tokens)}-token summary"
    end
    parts << "#{fmt(plan.dropped_tokens)} dropped" if plan.dropped_tokens.positive?
    "#{parts.join(' · ')} — fits the budget."
  end
end