pi-condense 2.4.3 → 2.6.0

package/src/batch-capture.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, test } from "bun:test";
-import { serializeBatchForSummarizer } from "./batch-capture.js";
+import { captureBatch, captureUnindexedBatchesFromSession, serializeBatchForSummarizer } from "./batch-capture.js";
 import type { CapturedBatch, CapturedToolCall } from "./types.js";
 
 function toolCall(overrides: Partial<CapturedToolCall> = {}): CapturedToolCall {
@@ -49,3 +49,77 @@ describe("serializeBatchForSummarizer", () => {
     expect(result).toContain("[[3:write]] Tool:");
   });
 });
+
+describe("occurrence capture", () => {
+  test("captureBatch records the matched result's timestamp", () => {
+    const message = {
+      role: "assistant",
+      content: [{ type: "toolCall", id: "bash_23", name: "bash", input: { cmd: "ls" } }],
+      timestamp: 2100,
+    };
+    const results = [
+      { role: "toolResult", toolCallId: "bash_23", toolName: "bash", content: [{ type: "text", text: "ok" }], isError: false, timestamp: 2150 },
+    ];
+    const batch = captureBatch(message, results, 0, 9999);
+    expect(batch.toolCalls[0].resultTimestamp).toBe(2150);
+  });
+
+  test("captureBatch omits resultTimestamp when no result matched", () => {
+    const message = { role: "assistant", content: [{ type: "toolCall", id: "x", name: "bash", input: {} }], timestamp: 1 };
+    const batch = captureBatch(message, [], 0, 9999);
+    expect(batch.toolCalls[0].resultTimestamp).toBeUndefined();
+    expect("resultTimestamp" in batch.toolCalls[0]).toBe(false);
+    expect(batch.toolCalls[0].resultText).toBe("(no result)");
+  });
+
+  test("rescan pairs each assistant with the results of its OWN turn when ids repeat", () => {
+    const entry = (message: any) => ({ type: "message", message, timestamp: undefined });
+    const branch = [
+      entry({ role: "user", content: [{ type: "text", text: "go" }], timestamp: 1000 }),
+      entry({ role: "assistant", content: [{ type: "toolCall", id: "bash_23", name: "bash", input: {} }], timestamp: 1100 }),
+      entry({ role: "toolResult", toolCallId: "bash_23", toolName: "bash", content: [{ type: "text", text: "FIRST" }], isError: false, timestamp: 1150 }),
+      entry({ role: "assistant", content: [{ type: "toolCall", id: "bash_23", name: "bash", input: {} }], timestamp: 2100 }),
+      entry({ role: "toolResult", toolCallId: "bash_23", toolName: "bash", content: [{ type: "text", text: "SECOND" }], isError: false, timestamp: 2150 }),
+    ];
+    const batches = captureUnindexedBatchesFromSession(branch, { isSummarized: () => false });
+    expect(batches).toHaveLength(2);
+    expect(batches[0].toolCalls[0].resultText).toBe("FIRST");
+    expect(batches[0].toolCalls[0].resultTimestamp).toBe(1150);
+    expect(batches[1].toolCalls[0].resultText).toBe("SECOND");
+    expect(batches[1].toolCalls[0].resultTimestamp).toBe(2150);
+  });
+
+  test("rescan asks isSummarized with the occurrence key, not the bare id", () => {
+    const asked: string[] = [];
+    const entry = (message: any) => ({ type: "message", message });
+    const branch = [
+      entry({ role: "assistant", content: [{ type: "toolCall", id: "bash_23", name: "bash", input: {} }], timestamp: 1100 }),
+      entry({ role: "toolResult", toolCallId: "bash_23", toolName: "bash", content: [{ type: "text", text: "x" }], isError: false, timestamp: 1150 }),
+    ];
+    captureUnindexedBatchesFromSession(branch, { isSummarized: (id: string) => (asked.push(id), false) });
+    expect(asked).toContain("bash_23@1150");
+  });
+
+  test("rescan still skips a call whose result has not arrived", () => {
+    const entry = (message: any) => ({ type: "message", message });
+    const branch = [
+      entry({ role: "assistant", content: [{ type: "toolCall", id: "pending", name: "bash", input: {} }], timestamp: 1 }),
+    ];
+    expect(captureUnindexedBatchesFromSession(branch, { isSummarized: () => false })).toEqual([]);
+  });
+
+  test("rescan does not pair a result that falls outside its own assistant's turn window", () => {
+    // bash_23's result lands AFTER the next assistant message, i.e. in the
+    // second assistant's window, not the first's. Per-turn scanning does not
+    // fabricate a pair for the first assistant (no in-window result), and the
+    // second assistant has no bash_23 call to attach the result to either, so
+    // no batch is emitted at all.
+    const entry = (message: any) => ({ type: "message", message });
+    const branch = [
+      entry({ role: "assistant", content: [{ type: "toolCall", id: "bash_23", name: "bash", input: {} }], timestamp: 1000 }),
+      entry({ role: "assistant", content: [{ type: "toolCall", id: "other", name: "bash", input: {} }], timestamp: 1100 }),
+      entry({ role: "toolResult", toolCallId: "bash_23", toolName: "bash", content: [{ type: "text", text: "late" }], isError: false, timestamp: 1150 }),
+    ];
+    expect(captureUnindexedBatchesFromSession(branch, { isSummarized: () => false })).toEqual([]);
+  });
+});

package/src/batch-capture.ts

@@ -1,4 +1,5 @@
 import type { CapturedBatch, CapturedToolCall, BatchingMode } from "./types.js";
+import { occKey, resultTimestampOf } from "./occurrence-key.js";
 
 /** Joins the text blocks of a ToolResultMessage into a single string. */
 export function extractToolResultText(msg: any): string {
@@ -37,6 +38,7 @@ export function captureBatch(
 
       let resultText = "(no result)";
       let isError = false;
+      const resultTimestamp = match ? resultTimestampOf(match.timestamp) : undefined;
 
       if (match) {
         resultText = extractToolResultText(match);
@@ -49,6 +51,7 @@ export function captureBatch(
         args: block.input ?? block.args ?? block.arguments ?? {},
         resultText,
         isError,
+        ...(resultTimestamp !== undefined ? { resultTimestamp } : {}),
       } satisfies CapturedToolCall;
     });
 
@@ -70,14 +73,7 @@ export function captureUnindexedBatchesFromSession(
 ): CapturedBatch[] {
   // branch is SessionEntry[]. Each message entry has { type: "message", message: AgentMessage }.
   // We must unwrap the SessionEntry wrapper before accessing role/toolCallId.
-  const resultMap = new Map<string, any>();
-  for (const entry of branch) {
-    if (entry.type !== "message") continue;
-    const m = entry.message;
-    if (m.role === "toolResult" && m.toolCallId) {
-      resultMap.set(m.toolCallId, m);
-    }
-  }
+  const entries = branch.filter((entry: any) => entry.type === "message");
 
   const batches: CapturedBatch[] = [];
   // turnCounter increments for EVERY assistant message (not just prunable ones).
@@ -93,8 +89,8 @@ export function captureUnindexedBatchesFromSession(
   // a single user → final-agent-message span when batchingMode === "agent-message".
   let userTurnGroup = 0;
 
-  for (const entry of branch) {
-    if (entry.type !== "message") continue;
+  for (let i = 0; i < entries.length; i++) {
+    const entry = entries[i];
     const msg = entry.message;
 
     // Advance userTurnGroup on every user message so all subsequent assistant
@@ -109,6 +105,17 @@ export function captureUnindexedBatchesFromSession(
     // Stable turn index: count every assistant message regardless of pruning state
     const currentTurnIndex = turnCounter++;
 
+    // Per-turn result map: only the results between this assistant message and
+    // the next one. A branch-wide map is last-wins and mis-pairs repeated ids.
+    const turnResults = new Map<string, any>();
+    for (let j = i + 1; j < entries.length; j++) {
+      const m = entries[j].message;
+      if (m.role === "assistant") break;
+      if (m.role === "toolResult" && m.toolCallId && !turnResults.has(m.toolCallId)) {
+        turnResults.set(m.toolCallId, m);
+      }
+    }
+
     const content = Array.isArray(msg.content) ? msg.content : [];
     const toolCallBlocks = content.filter((c: any) => c.type === "toolCall");
 
@@ -116,13 +123,15 @@ export function captureUnindexedBatchesFromSession(
     const readyToPrune = toolCallBlocks.filter((tc: any) => {
       const id = tc.id;
       if (!id) return false;
-      if (indexer.isSummarized(id)) return false;
+      const result = turnResults.get(id);
+      if (!result) return false;
+      if (indexer.isSummarized(occKey(id, resultTimestampOf(result.timestamp)))) return false;
       if (exclude(tc.name, tc.input ?? tc.arguments)) return false;
-      return resultMap.has(id);
+      return true;
     });
 
     if (readyToPrune.length > 0) {
-      const results = readyToPrune.map((tc: any) => resultMap.get(tc.id));
+      const results = readyToPrune.map((tc: any) => turnResults.get(tc.id));
       const readyIds = new Set(readyToPrune.map((tc: any) => tc.id));
       // We pass the full message but then trim back down to only the tool calls
       // whose results already exist in the session. This lets a flush prune

package/src/chain-compressor.test.ts

@@ -3,6 +3,8 @@ import { selectEligible, compressEligible } from "./chain-compressor.js";
 import type { ChainCompressorIndexerDeps } from "./chain-compressor.js";
 import type { ChainRange, ChainCompressionEntry } from "./types.js";
 import { CUSTOM_TYPE_CHAIN } from "./types.js";
+import { detectChains } from "./chain-detector.js";
+import { inGraceRecoveryToolCallIds } from "./recovery-grace.js";
 
 function closed(startUserTimestamp: number, toolCallIds: string[] = [`tc-${startUserTimestamp}`]): ChainRange {
   return { startUserTimestamp, middleToolCallIds: toolCallIds, finalAssistantTimestamp: startUserTimestamp + 100 };
@@ -313,3 +315,115 @@ describe("selectEligible - recovery grace deferral", () => {
     expect(eligible.map((c) => c.startUserTimestamp)).toEqual([1, 2, 3]);
   });
 });
+
+describe("occurrence keys in compression", () => {
+  const chain = {
+    startUserTimestamp: 1000,
+    middleToolCallIds: ["bash_23"],
+    middleOccurrenceKeys: ["bash_23@1150"],
+    protectedToolCallIds: [],
+    finalAssistantTimestamp: 1200,
+  };
+
+  test("summary/toolRef lookups receive occurrence keys, not bare ids", async () => {
+    const asked: string[][] = [];
+    const deps = {
+      indexer: {
+        getChainEntries: () => [],
+        hasPerBatchSummaryCoveringAny: (ids: string[]) => (asked.push(ids), true),
+        getPerBatchSummariesForToolCallIds: (ids: string[]) => (asked.push(ids), ["s1"]),
+        getToolRefsForToolCallIds: (ids: string[]) => (asked.push(ids), ["t1"]),
+        registerChain: () => {},
+      },
+      blockRefs: { issue: () => "b1" },
+      appendEntry: () => {},
+      now: () => 5000,
+    };
+    await compressEligible([chain as any], 0, deps as any);
+    expect(asked).toEqual([["bash_23@1150"], ["bash_23@1150"]]);
+  });
+
+  test("persists droppedOccurrenceKeys alongside droppedToolCallIds", async () => {
+    const appended: any[] = [];
+    const deps = {
+      indexer: {
+        getChainEntries: () => [],
+        hasPerBatchSummaryCoveringAny: () => true,
+        getPerBatchSummariesForToolCallIds: () => ["s1"],
+        getToolRefsForToolCallIds: () => ["t1"],
+        registerChain: () => {},
+      },
+      blockRefs: { issue: () => "b1" },
+      appendEntry: (_t: string, data: unknown) => appended.push(data),
+      now: () => 5000,
+    };
+    const { compressedEntries } = await compressEligible([chain as any], 0, deps as any);
+    expect(compressedEntries[0].droppedToolCallIds).toEqual(["bash_23"]);
+    expect(compressedEntries[0].droppedOccurrenceKeys).toEqual(["bash_23@1150"]);
+    expect((appended[0] as any).droppedOccurrenceKeys).toEqual(["bash_23@1150"]);
+  });
+
+  test("a chain without middleOccurrenceKeys falls back to bare ids", async () => {
+    const deps = {
+      indexer: {
+        getChainEntries: () => [],
+        hasPerBatchSummaryCoveringAny: () => true,
+        getPerBatchSummariesForToolCallIds: () => ["s1"],
+        getToolRefsForToolCallIds: () => ["t1"],
+        registerChain: () => {},
+      },
+      blockRefs: { issue: () => "b1" },
+      appendEntry: () => {},
+      now: () => 5000,
+    };
+    const { compressedEntries } = await compressEligible(
+      [{ ...chain, middleOccurrenceKeys: undefined } as any],
+      0,
+      deps as any,
+    );
+    expect(compressedEntries[0].droppedOccurrenceKeys).toBeUndefined();
+  });
+});
+
+describe("selectEligible - recovery grace wiring (production boundary)", () => {
+  // Crosses the real boundary: real detectChains() + real inGraceRecoveryToolCallIds()
+  // feed real selectEligible(). Hand-injecting bare ids (as the suite above does) hides
+  // the format mismatch this test pins down.
+  it("does not compress a chain whose recovery-query result is still in grace", () => {
+    const messages = [
+      { role: "user", timestamp: 1, content: [{ type: "text", text: "u1" }] },
+      {
+        role: "assistant",
+        timestamp: 1,
+        content: [{ type: "toolCall", id: "t1", name: "context_tree_query", input: {} }],
+      },
+      {
+        role: "toolResult",
+        toolCallId: "t1",
+        toolName: "context_tree_query",
+        timestamp: 100,
+        content: [{ type: "text", text: "recovered" }],
+      },
+      { role: "assistant", timestamp: 2, content: [{ type: "text", text: "done" }] },
+      { role: "user", timestamp: 3, content: [{ type: "text", text: "u2" }] },
+      { role: "assistant", timestamp: 3, content: [{ type: "toolCall", id: "b", name: "bash", input: {} }] },
+      {
+        role: "toolResult",
+        toolCallId: "b",
+        toolName: "bash",
+        timestamp: 200,
+        content: [{ type: "text", text: "ok" }],
+      },
+      { role: "assistant", timestamp: 4, content: [{ type: "text", text: "done2" }] },
+    ];
+
+    const chains = detectChains(messages);
+    expect(chains.map((c) => c.startUserTimestamp)).toEqual([1, 3]);
+
+    const inGrace = inGraceRecoveryToolCallIds(messages, 3);
+    expect([...inGrace]).toEqual(["t1@100"]);
+
+    const eligible = selectEligible(chains, 0, new Set(), inGrace);
+    expect(eligible.map((c) => c.startUserTimestamp)).toEqual([3]);
+  });
+});

package/src/chain-compressor.ts

@@ -2,6 +2,26 @@ import { CUSTOM_TYPE_CHAIN } from "./types.js";
 import type { ChainRange, ChainCompressionEntry } from "./types.js";
 import type { ToolCallIndexer } from "./indexer.js";
 import type { BlockRefIssuer } from "./block-refs.js";
+import { bareToolCallId, parseOccKey } from "./occurrence-key.js";
+
+/**
+ * Grace ids are keyed the same way `recovery-grace.ts` keys them: occurrence
+ * (`id@timestamp`) when the recovery message carried a timestamp, bare id
+ * otherwise. A chain's middles are compared occurrence-first (exact match on
+ * `middleOccurrenceKeys`, falling back to `middleToolCallIds` for chains built
+ * before the field existed), so a graced occurrence never defers a chain
+ * holding a DIFFERENT occurrence of the same reused provider id. The only
+ * bare-to-bare fallback is for grace entries that themselves have no
+ * timestamp discriminant — there is no exact key to compare in that case.
+ */
+function chainMatchesGrace(chain: ChainRange, inGraceToolCallIds: Set<string>): boolean {
+  const keys = chain.middleOccurrenceKeys?.length ? chain.middleOccurrenceKeys : chain.middleToolCallIds;
+  if (keys.some((k) => inGraceToolCallIds.has(k))) return true;
+  for (const g of inGraceToolCallIds) {
+    if (parseOccKey(g).resultTimestamp === undefined && keys.some((k) => bareToolCallId(k) === g)) return true;
+  }
+  return false;
+}
 
 /**
  * Pure eligibility filter: given all detected chains, return the subset
@@ -31,7 +51,7 @@ export function selectEligible(
       c.middleToolCallIds.length > 0,
   );
   const toCompress = candidates.slice(0, Math.max(0, candidates.length - rollingWindow));
-  return toCompress.filter((c) => !c.middleToolCallIds.some((id) => inGraceToolCallIds.has(id)));
+  return toCompress.filter((c) => !chainMatchesGrace(c, inGraceToolCallIds));
 }
 
 /**
@@ -96,19 +116,23 @@ export async function compressEligible(
 
   const compressedEntries: ChainCompressionEntry[] = [];
   for (const chain of eligible) {
-    if (!deps.indexer.hasPerBatchSummaryCoveringAny(chain.middleToolCallIds)) {
+    // summaryBodies / toolRefs live in occurrence-key space (src/indexer.ts).
+    // Bare ids would match nothing and silently skip every chain.
+    const lookupKeys = chain.middleOccurrenceKeys?.length ? chain.middleOccurrenceKeys : chain.middleToolCallIds;
+
+    if (!deps.indexer.hasPerBatchSummaryCoveringAny(lookupKeys)) {
       skipped.push({ startUserTimestamp: chain.startUserTimestamp, reason: "no-summary" });
       continue;
     }
 
     const blockId = deps.blockRefs.issue();
-    const toolRefs = deps.indexer.getToolRefsForToolCallIds(chain.middleToolCallIds);
+    const toolRefs = deps.indexer.getToolRefsForToolCallIds(lookupKeys);
 
     // B: fuse this span's per-batch summaries into one cohesive summary.
     // Gated on >= 2 summaries (nothing to fuse otherwise). Non-fatal.
     let rangeSummaryText: string | undefined;
     if (deps.fuseRange) {
-      const summaries = deps.indexer.getPerBatchSummariesForToolCallIds(chain.middleToolCallIds);
+      const summaries = deps.indexer.getPerBatchSummariesForToolCallIds(lookupKeys);
       if (summaries.length >= 2) {
         try {
           const fused = await deps.fuseRange(summaries.join("\n\n"));
@@ -128,6 +152,7 @@ export async function compressEligible(
       compressedAt: deps.now(),
       ...(rangeSummaryText ? { rangeSummaryText } : {}),
       ...(chain.protectedToolCallIds?.length ? { protectedToolCallIds: chain.protectedToolCallIds } : {}),
+      ...(chain.middleOccurrenceKeys?.length ? { droppedOccurrenceKeys: chain.middleOccurrenceKeys } : {}),
     };
 
     deps.appendEntry(CUSTOM_TYPE_CHAIN, entry);

package/src/chain-detector.test.ts

@@ -259,6 +259,55 @@ describe("detectChains protectedToolCallIds", () => {
   });
 });
 
+describe("middleOccurrenceKeys", () => {
+  test("emits one occurrence key per middle tool result", () => {
+    const messages = [
+      { role: "user", content: [{ type: "text", text: "go" }], timestamp: 1000 },
+      { role: "assistant", content: [{ type: "toolCall", id: "bash_1", name: "bash", input: {} }], timestamp: 1100 },
+      { role: "toolResult", toolCallId: "bash_1", toolName: "bash", content: [{ type: "text", text: "x" }], isError: false, timestamp: 1150 },
+      { role: "assistant", content: [{ type: "text", text: "done" }], timestamp: 1200 },
+    ];
+    const [chain] = detectChains(messages);
+    expect(chain.middleToolCallIds).toEqual(["bash_1"]);
+    expect(chain.middleOccurrenceKeys).toEqual(["bash_1@1150"]);
+  });
+
+  test("a tool call with no result contributes no occurrence key", () => {
+    const messages = [
+      { role: "user", content: [{ type: "text", text: "go" }], timestamp: 1000 },
+      { role: "assistant", content: [{ type: "toolCall", id: "bash_1", name: "bash", input: {} }], timestamp: 1100 },
+      { role: "assistant", content: [{ type: "text", text: "done" }], timestamp: 1200 },
+    ];
+    const [chain] = detectChains(messages);
+    expect(chain.middleToolCallIds).toEqual(["bash_1"]);
+    expect(chain.middleOccurrenceKeys).toEqual([]);
+  });
+
+  test("two chains reusing one id get distinct occurrence keys", () => {
+    const turn = (base: number) => [
+      { role: "user", content: [{ type: "text", text: "go" }], timestamp: base },
+      { role: "assistant", content: [{ type: "toolCall", id: "bash_23", name: "bash", input: {} }], timestamp: base + 100 },
+      { role: "toolResult", toolCallId: "bash_23", toolName: "bash", content: [{ type: "text", text: "x" }], isError: false, timestamp: base + 150 },
+      { role: "assistant", content: [{ type: "text", text: "done" }], timestamp: base + 200 },
+    ];
+    const chains = detectChains([...turn(1000), ...turn(2000)]);
+    expect(chains.map((c) => c.middleOccurrenceKeys)).toEqual([["bash_23@1150"], ["bash_23@2150"]]);
+  });
+
+  test("an interrupted chain still reports the keys it collected", () => {
+    const messages = [
+      { role: "user", content: [{ type: "text", text: "a" }], timestamp: 1000 },
+      { role: "assistant", content: [{ type: "toolCall", id: "bash_1", name: "bash", input: {} }], timestamp: 1100 },
+      { role: "toolResult", toolCallId: "bash_1", toolName: "bash", content: [{ type: "text", text: "x" }], isError: false, timestamp: 1150 },
+      { role: "user", content: [{ type: "text", text: "b" }], timestamp: 2000 },
+      { role: "assistant", content: [{ type: "text", text: "done" }], timestamp: 2200 },
+    ];
+    const [interrupted] = detectChains(messages);
+    expect(interrupted.finalAssistantTimestamp).toBeNull();
+    expect(interrupted.middleOccurrenceKeys).toEqual(["bash_1@1150"]);
+  });
+});
+
 describe("withClosingMessage", () => {
   test("undefined closing returns the same array reference", () => {
     const msgs = [userMsg(100)];

package/src/chain-detector.ts

@@ -1,3 +1,4 @@
+import { occKey, resultTimestampOf } from "./occurrence-key.js";
 import type { ChainRange } from "./types.js";
 
 /** Prefix that identifies a synthetic chain-compression user message. */
@@ -45,6 +46,7 @@ export function detectChains(
   let state: State = "idle";
   let chainStart: { timestamp: number } | null = null;
   let middleIds = new Set<string>();
+  let middleKeys = new Set<string>();
   let protectedIds = new Set<string>();
 
   const emitInterrupted = () => {
@@ -52,6 +54,7 @@ export function detectChains(
       ranges.push({
         startUserTimestamp: chainStart.timestamp,
         middleToolCallIds: [...middleIds],
+        middleOccurrenceKeys: [...middleKeys],
         protectedToolCallIds: [...protectedIds],
         finalAssistantTimestamp: null,
       });
@@ -64,6 +67,7 @@ export function detectChains(
       emitInterrupted();
       chainStart = { timestamp: msg.timestamp };
       middleIds = new Set();
+      middleKeys = new Set();
       protectedIds = new Set();
       state = "inChain";
       continue;
@@ -82,6 +86,7 @@ export function detectChains(
     if (msg.role === "toolResult") {
       if (msg.toolCallId) {
         middleIds.add(msg.toolCallId);
+        middleKeys.add(occKey(msg.toolCallId, resultTimestampOf(msg.timestamp)));
         // toolResult fallback — results carry no args; name-only by design,
         // the assistant block always precedes its result so no protection is lost
         if (isProtected(msg.toolName, undefined)) protectedIds.add(msg.toolCallId);
@@ -93,11 +98,13 @@ export function detectChains(
       ranges.push({
         startUserTimestamp: chainStart!.timestamp,
         middleToolCallIds: [...middleIds],
+        middleOccurrenceKeys: [...middleKeys],
         protectedToolCallIds: [...protectedIds],
         finalAssistantTimestamp: msg.timestamp,
       });
       chainStart = null;
       middleIds = new Set();
+      middleKeys = new Set();
       protectedIds = new Set();
       state = "idle";
     }