Merge pull request #1407 from corbitsdev/cl-10535-excerpt-once

refactor(session): stop excerpting the same tail turn twice in turnCost

Author: TheGreatAxios

src/session/compactor-excerpt-once.test.ts

@@ -0,0 +1,97 @@
+import { describe, expect, test } from "bun:test";
+import {
+  createPruningCompactor,
+  excerptTailTurnCallCount,
+} from "./compactor.js";
+import type {
+  ConversationTurn,
+  ReactorState,
+  StrategyContext,
+} from "@intx/types/runtime";
+
+const mockStrategyCtx: StrategyContext = {
+  state: {} as ReactorState,
+  trigger: "test",
+};
+
+function textTurn(
+  role: ConversationTurn["role"],
+  text: string,
+): ConversationTurn {
+  return {
+    role,
+    content: [{ type: "text", text }],
+    timestamp: Date.now(),
+  };
+}
+
+function pairTurns(id: string, resultText: string): ConversationTurn[] {
+  return [
+    {
+      role: "assistant",
+      content: [
+        {
+          type: "tool_call",
+          id,
+          name: "read_file",
+          arguments: { path: `src/${id}.ts` },
+        },
+      ],
+      timestamp: Date.now(),
+    },
+    {
+      role: "user",
+      content: [
+        {
+          type: "tool_result",
+          callId: id,
+          content: [{ type: "text", text: resultText }],
+        },
+      ],
+      timestamp: Date.now(),
+    },
+  ];
+}
+
+function liveResultText(turns: ConversationTurn[]): string {
+  return turns
+    .flatMap((t) =>
+      t.content.flatMap((b) => {
+        if (b.type === "text") return [b.text];
+        if (b.type === "tool_result")
+          return b.content.map((c) => (c.type === "text" ? c.text : ""));
+        return [];
+      }),
+    )
+    .join("\n");
+}
+
+describe("excerpt once per tail turn", () => {
+  test("compactor excerpts each tail turn exactly once with identical values", async () => {
+    const dump = "z".repeat(12_000);
+    const turns: ConversationTurn[] = [
+      textTurn("user", "do the work"),
+      ...pairTurns("a", `a-result:${dump}`),
+      ...pairTurns("b", `b-result:${dump}`),
+      ...pairTurns("c", `c-result:${dump}`),
+      textTurn("user", "newest ask"),
+      textTurn("assistant", "newest reply"),
+    ];
+
+    const before = excerptTailTurnCallCount();
+    const result = await createPruningCompactor({
+      compactionShape: { tailBudgetTokens: 1_000_000 },
+      summarize: async () => {
+        throw new Error("must not invent a summary: nothing excluded");
+      },
+    }).apply(turns, mockStrategyCtx);
+
+    const calls = excerptTailTurnCallCount() - before;
+
+    expect(calls).toBe(turns.length);
+
+    const live = liveResultText(result.output);
+    expect(live.match(/\[tail-shortened \d+→/g)?.length ?? 0).toBe(3);
+    expect(live).not.toContain(dump);
+  });
+});

src/session/compactor.ts

@@ -647,6 +647,11 @@ function excerptTailText(
   };
 }
 
+let excerptTailTurnCallCountValue = 0;
+export function excerptTailTurnCallCount(): number {
+  return excerptTailTurnCallCountValue;
+}
+
 // Excerpted live copy of a tail turn: thinking/redacted_thinking drop except
 // on the last assistant with a tool_call (Anthropic continuation needs the
 // unmodified thinking+signature on that tool_use). Whole blocks drop so
@@ -658,6 +663,7 @@ function excerptTailTurn(
   shape: CompactionShape,
   keepThinking = false,
 ): { turn: ConversationTurn; shortenedOutputs: number; changed: boolean } {
+  excerptTailTurnCallCountValue += 1;
   let shortenedOutputs = 0;
   let changed = false;
   const content = turn.content.flatMap((block): ConversationTurn["content"] => {
@@ -716,12 +722,7 @@ function lastAssistantToolCallIndex(
   return undefined;
 }
 
-function excerptedTurnTokens(
-  turn: ConversationTurn,
-  shape: CompactionShape,
-  keepThinking = false,
-): number {
-  const { turn: live } = excerptTailTurn(turn, shape, keepThinking);
+function liveTurnTokens(live: ConversationTurn): number {
   let tokens = 0;
   for (const block of live.content) {
     tokens += estimateContentBlockTokens(block);
@@ -769,18 +770,24 @@ function selectTail(
   const budgetTokens = shape.tailBudgetTokens;
   const keepThinkingAt = lastAssistantToolCallIndex(turns);
 
+  type ExcerptResult = ReturnType<typeof excerptTailTurn>;
+  const excerptCache = new Map<number, ExcerptResult>();
+  const excerptFor = (idx: number, turn: ConversationTurn): ExcerptResult => {
+    let result = excerptCache.get(idx);
+    if (result === undefined) {
+      result = excerptTailTurn(turn, shape, idx === keepThinkingAt);
+      excerptCache.set(idx, result);
+    }
+    return result;
+  };
+
   const turnCost = (idx: number): { tokens: number; shortened: number } => {
     const turn = turns[idx];
     if (turn === undefined) return { tokens: 0, shortened: 0 };
-    const keepThinking = idx === keepThinkingAt;
-    const {
-      turn: live,
-      shortenedOutputs,
-      changed,
-    } = excerptTailTurn(turn, shape, keepThinking);
+    const { turn: live, shortenedOutputs, changed } = excerptFor(idx, turn);
     if (changed) excerpted.set(idx, live);
     return {
-      tokens: excerptedTurnTokens(turn, shape, keepThinking),
+      tokens: liveTurnTokens(live),
       shortened: shortenedOutputs,
     };
   };
@@ -836,7 +843,7 @@ function selectTail(
     for (const idx of closure) {
       const t = turns[idx];
       if (t === undefined || isFoldableHandoffTurn(t)) continue;
-      closureTokens += excerptedTurnTokens(t, shape, idx === keepThinkingAt);
+      closureTokens += liveTurnTokens(excerptFor(idx, t).turn);
     }
     if (picked.size > 0 && usedTokens + closureTokens > budgetTokens) break;
     for (const idx of closure) pick(idx);