context-compactor

/**
 * Bounds what a tool can put into the context.
 *
 * Compaction runs once, before a reply. A tool result lands in the middle of one, and every server
 * has its own idea of a reasonable size: the official filesystem server will happily return a 6 MB
 * binary, which is a quarter of a million tokens and the instant death of any local context. No
 * window is large enough to make that safe, so the ceiling has to be here, where every tool from
 * every server passes through on its way to the model.
 *
 * This is the one thing owning the prediction loop is genuinely good for.
 */

/** Rough, and deliberately so: counting tokens per call would cost more than it saves. */
const CHARS_PER_TOKEN = 4;

/** Below this, a string is assumed to be structure — a type tag, a name, a path — and left alone. */
const STRUCTURAL_STRING_MAX = 200;

function notice(kept: number, total: number): string {
  if (kept === 0) {
    return (
      `[context-compactor: this reply has no context left for tool output, so all ${total} ` +
      "characters were dropped. The tool itself ran. Answer with what you already have rather " +
      "than calling more tools.]"
    );
  }
  return (
    `\n\n[Truncated by context-compactor: kept ${kept} of ${total} characters. The rest was ` +
    "discarded and repeating this call will not return it. Ask for a smaller piece — a section, a " +
    "range, a single record — or use a tool that summarizes rather than dumps.]"
  );
}

function capText(text: string, maxChars: number): string {
  return text.length <= maxChars ? text : text.slice(0, maxChars) + notice(maxChars, text.length);
}

/**
 * Shrinks the strings inside a value while leaving its shape exactly as it was.
 *
 * The shape is not ours to touch. An MCP tool result is `{ content: [{ type: "text", text }] }` and
 * is validated against that schema on the way back: replacing an oversized object with a truncated
 * JSON rendering of itself produces a string where a structured result belongs, and the call is
 * rejected outright — the tool having worked perfectly. Truncating the text and putting it back
 * where it came from is the only version of this that is safe on a shape we do not know.
 *
 * The budget is shared across the walk, so a result made of many medium strings is bounded just as
 * firmly as one made of a single enormous one.
 */
function capDeep(value: unknown, budget: { left: number }): unknown {
  if (typeof value === "string") {
    // Short strings are the schema, not the payload: "text" is a discriminator, "image" is an enum
    // member, and a mime type is a contract. Truncating those to a notice rebuilds the exact bug
    // this walk exists to avoid — an object of the right shape carrying the wrong values. Nothing
    // this short is worth rationing anyway.
    if (value.length <= STRUCTURAL_STRING_MAX) {
      return value;
    }
    const limit = Math.max(0, budget.left);
    const capped = capText(value, limit);
    budget.left -= Math.min(value.length, limit);
    return capped;
  }
  if (Array.isArray(value)) {
    return value.map(item => capDeep(item, budget));
  }
  if (value !== null && typeof value === "object") {
    const out: Record<string, unknown> = {};
    for (const [key, item] of Object.entries(value)) {
      out[key] = capDeep(item, budget);
    }
    return out;
  }
  return value;
}

export function capResult(result: unknown, maxTokens: number): unknown {
  return capDeep(result, { left: maxTokens * CHARS_PER_TOKEN });
}

function measureTokens(value: unknown): number {
  if (value === null || value === undefined) {
    return 0;
  }
  const text = typeof value === "string" ? value : (JSON.stringify(value) ?? "");
  return Math.ceil(text.length / CHARS_PER_TOKEN);
}

export interface ToolCapOpts {
  /** Ceiling for any single result. */
  maxPerResult: number;
  /** Ceiling for everything the tools return during one reply, together. */
  budget: number;
}

/**
 * Wraps each tool so its result passes through the ceiling. The tool keeps its name, description and
 * schema — the model sees no difference until something comes back too big.
 *
 * The budget is shared across the whole reply, and that is the point. Capping each result on its own
 * bounds nothing: a model that calls fifteen tools in one turn fills the window fifteen times four
 * thousand tokens, every one of them obediently under the per-result limit. Only the sum matters,
 * because only the sum is what the context has to hold.
 *
 * Each wrapper closes over the same counter, so the list must be built fresh for every reply.
 */
export function capTools<T extends { implementation: (...args: any[]) => any }>(
  tools: Array<T>,
  opts: ToolCapOpts,
): Array<T> {
  let spent = 0;

  return tools.map(tool => ({
    ...tool,
    implementation: async (...args: Array<unknown>) => {
      // The tool always runs, even with nothing left to spend. A call the model makes to write a
      // file must write the file; it is the result's verbosity we are rationing, not the action.
      // Capping to zero returns the tool's own shape carrying only the notice, which is both honest
      // and valid — inventing a refusal here would mean inventing a result shape we do not know.
      const result = await tool.implementation(...args);
      const remaining = Math.max(0, opts.budget - spent);
      const capped = capResult(result, Math.min(opts.maxPerResult, remaining));
      spent += measureTokens(capped);
      return capped;
    },
  }));
}