Conversation Pruner — JavaScript source

/**
 * Conversation Pruner — compute a deterministic pruning plan for a
 * token-budgeted chat history.
 *
 * Language: JavaScript (ES2022+, ES module; runs unmodified in Node 18+
 *           and modern browsers)
 * Source:   CosmoDev polyglot showcase port of the Conversation Pruner
 *           tool (slug: conversation-pruner).
 * Port of src/lib/conversationPruner.ts (the canonical TypeScript
 *           implementation).
 * Tool page: https://dev.cosmolabs.org/tools/conversation-pruner
 * License:  display source — part of CosmoDev's polyglot tool pages.
 *
 * 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 (the opening user request), and the current (last user)
 * request. Token counts are plain integers; the only float math is the
 * summary cost: fixed framing plus 10% of the folded content, rounded up.
 */

/** Summary compression model: fixed framing tokens. */
export const SUMMARY_FIXED_TOKENS = 60;

/** Summary compression model: share of the folded content's tokens. */
export const SUMMARY_RATIO = 0.1;

/**
 * Compute the pruning plan for `messages` under `opts.budgetTokens`.
 * Each message is `{ role, content, tokens, pinned? }` with role one of
 * 'system' | 'user' | 'assistant' | 'tool'. Throws a RangeError on a
 * negative budget or any negative per-message token count. The returned
 * plan carries one `{ index, role, action, tokens }` decision per message
 * ('keep' | 'summarize' | 'drop'), the token tallies, the projected total,
 * whether it fits, and any warnings.
 */
export function planPrune(messages, opts) {
  const warnings = [];
  if (opts.budgetTokens < 0) throw new RangeError('budgetTokens must be >= 0');
  if (messages.some((m) => m.tokens < 0)) throw new RangeError('message tokens must be >= 0');

  const n = messages.length;
  const lastIndexWithRole = (role) => {
    for (let i = n - 1; i >= 0; i--) if (messages[i].role === role) return i;
    return -1;
  };
  const lastUser = lastIndexWithRole('user');

  // Untouchable: every system message, pinned messages, the first turn (the
  // opening user request that anchors the conversation), and the current
  // request (the last user message and everything after it).
  const protectedIdx = new Set();
  messages.forEach((m, i) => {
    if (m.role === 'system' || m.pinned) protectedIdx.add(i);
  });
  if (n > 0) protectedIdx.add(0);
  const firstTurn = messages.findIndex((m) => m.role !== 'system');
  if (firstTurn !== -1) protectedIdx.add(firstTurn);
  for (let i = Math.max(lastUser === -1 ? n - 1 : lastUser, 0); i < n; i++) protectedIdx.add(i);

  const protectedTokens = [...protectedIdx].reduce((sum, i) => sum + messages[i].tokens, 0);
  if (protectedTokens > opts.budgetTokens) {
    warnings.push(
      `Protected messages alone are ${protectedTokens.toLocaleString('en-US')} tokens against a ${opts.budgetTokens.toLocaleString('en-US')} budget — raise the budget (or reserve less for the reply) before pruning anything else.`,
    );
  }

  // Fill the remaining budget newest-to-oldest through the middle.
  const actions = new Array(n).fill('drop');
  for (const i of protectedIdx) actions[i] = 'keep';
  let used = protectedTokens;
  for (let i = n - 1; i >= 0; i--) {
    if (actions[i] !== 'drop') continue;
    if (used + messages[i].tokens <= opts.budgetTokens) {
      actions[i] = 'keep';
      used += messages[i].tokens;
    } else {
      break; // oldest-unfilled remain drop/summarize candidates, newest first stopped
    }
  }

  // Everything still 'drop' in the middle folds into ONE running summary when
  // the compressed form fits where the raw turns did not.
  const summarizeIdx = actions
    .map((a, i) => (a === 'drop' && !protectedIdx.has(i) ? i : -1))
    .filter((i) => i >= 0);
  const summarizeTokens = summarizeIdx.reduce((sum, i) => sum + messages[i].tokens, 0);
  const attemptedSummaryCost =
    summarizeIdx.length > 0 ? SUMMARY_FIXED_TOKENS + Math.ceil(summarizeTokens * SUMMARY_RATIO) : 0;

  // The summary only costs anything when it is actually applied — otherwise
  // those turns drop and cost zero.
  let summaryCost = 0;
  if (attemptedSummaryCost > 0 && used + attemptedSummaryCost <= opts.budgetTokens) {
    for (const i of summarizeIdx) actions[i] = 'summarize';
    summaryCost = attemptedSummaryCost;
    used += summaryCost;
  } else if (attemptedSummaryCost > 0) {
    warnings.push(
      `Even the compressed summary (${attemptedSummaryCost.toLocaleString('en-US')} tokens) does not fit the remaining budget — the oldest turns are dropped instead.`,
    );
  }

  const decisions = messages.map((m, i) => ({
    index: i,
    role: m.role,
    action: actions[i],
    tokens: m.tokens,
  }));

  let keptTokens = 0;
  let droppedTokens = 0;
  let foldedTokens = 0;
  for (const d of decisions) {
    if (d.action === 'keep') keptTokens += d.tokens;
    else if (d.action === 'drop') droppedTokens += d.tokens;
    else foldedTokens += d.tokens;
  }

  return {
    decisions,
    keptTokens,
    summarizedTokens: foldedTokens,
    droppedTokens,
    summaryCostTokens: summaryCost,
    projectedTokens: keptTokens + summaryCost,
    fitsBudget: keptTokens + summaryCost <= opts.budgetTokens,
    warnings,
  };
}

/** Human-readable one-line summary of a plan. */
export function describePrune(plan) {
  if (!plan.fitsBudget) {
    return `Does not fit: ${plan.projectedTokens.toLocaleString('en-US')} tokens projected against the budget.`;
  }
  const parts = [`${plan.keptTokens.toLocaleString('en-US')} kept`];
  if (plan.summarizedTokens > 0) {
    parts.push(
      `${plan.summarizedTokens.toLocaleString('en-US')} folded into a ${plan.summaryCostTokens.toLocaleString('en-US')}-token summary`,
    );
  }
  if (plan.droppedTokens > 0) parts.push(`${plan.droppedTokens.toLocaleString('en-US')} dropped`);
  return `${parts.join(' · ')} — fits the budget.`;
}