Conversation Pruner — TypeScript source

// Pure logic for the Conversation Pruner tool (slug: conversation-pruner).
// Given a conversation with per-message token counts and a context budget,
// compute a deterministic pruning plan: which messages to keep, which to
// fold into a running summary, and which to drop outright — protecting
// system messages, pinned turns, and the current (last user) request.
export type ChatRole = 'system' | 'user' | 'assistant' | 'tool';

export interface ConversationMessage {
  role: ChatRole;
  content: string;
  /** Token count for this message (prompt-side framing included by caller). */
  tokens: number;
  /** Pinned messages are never dropped or summarized. */
  pinned?: boolean;
}

export interface PruneOptions {
  /** Total tokens available for the history (context window minus reserved reply). */
  budgetTokens: number;
}

export type PruneAction = 'keep' | 'summarize' | 'drop';

export interface PruneDecision {
  index: number;
  role: ChatRole;
  action: PruneAction;
  tokens: number;
}

export interface PrunePlan {
  decisions: PruneDecision[];
  keptTokens: number;
  summarizedTokens: number;
  droppedTokens: number;
  /** Tokens the summary placeholder itself will cost in the prompt. */
  summaryCostTokens: number;
  projectedTokens: number;
  fitsBudget: boolean;
  warnings: string[];
}

/** Summary compression model: fixed framing + 10% of the folded content. */
export const SUMMARY_FIXED_TOKENS = 60;
export const SUMMARY_RATIO = 0.1;

export function planPrune(
  messages: readonly ConversationMessage[],
  opts: PruneOptions,
): PrunePlan {
  const warnings: string[] = [];
  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: ChatRole) => {
    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<number>();
  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);

  // (reduce seeded: an empty conversation has no protected indexes at all)
  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: PruneAction[] = 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: PruneDecision[] = 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: PrunePlan): string {
  if (!plan.fitsBudget) {
    return `Does not fit: ${plan.projectedTokens.toLocaleString('en-US')} tokens projected against the budget.`;
  }
  const parts: string[] = [`${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.`;
}