pi-condense 2.0.0

package/src/chain-detector.test.ts

@@ -0,0 +1,302 @@
+import { describe, expect, test } from "bun:test";
+import { detectChains, withClosingMessage } from "./chain-detector.js";
+
+// ── Minimal message factories ──────────────────────────────────────────────
+
+function userMsg(timestamp: number, text = "do the thing"): any {
+  return { role: "user", content: [{ type: "text", text }], timestamp };
+}
+
+function syntheticChainMsg(timestamp: number, blockId = "b1"): any {
+  return {
+    role: "user",
+    content: [{ type: "text", text: `<compressed-chain id="${blockId}" tools="t1">summary</compressed-chain>` }],
+    timestamp,
+  };
+}
+
+function assistantWithTools(timestamp: number, toolCallIds: string[]): any {
+  return {
+    role: "assistant",
+    content: [
+      { type: "text", text: "working..." },
+      ...toolCallIds.map((id) => ({ type: "toolCall", id, name: "bash", arguments: {} })),
+    ],
+    timestamp,
+    usage: {},
+    stopReason: "toolUse",
+  };
+}
+
+function toolResult(timestamp: number, toolCallId: string): any {
+  return {
+    role: "toolResult",
+    toolCallId,
+    toolName: "bash",
+    content: [{ type: "text", text: "output" }],
+    isError: false,
+    timestamp,
+  };
+}
+
+function assistantText(timestamp: number, text = "done"): any {
+  return {
+    role: "assistant",
+    content: [{ type: "text", text }, { type: "thinking", thinking: "thoughts", thinkingSignature: "sig" }],
+    timestamp,
+    usage: {},
+    stopReason: "stop",
+  };
+}
+
+// ── Tests ──────────────────────────────────────────────────────────────────
+
+describe("detectChains", () => {
+  test("empty input returns empty array", () => {
+    expect(detectChains([])).toEqual([]);
+  });
+
+  test("single complete chain produces one range", () => {
+    const msgs = [
+      userMsg(100),
+      assistantWithTools(200, ["tc1", "tc2"]),
+      toolResult(300, "tc1"),
+      toolResult(310, "tc2"),
+      assistantText(400),
+    ];
+    const ranges = detectChains(msgs);
+    expect(ranges).toHaveLength(1);
+    expect(ranges[0].startUserTimestamp).toBe(100);
+    expect(ranges[0].middleToolCallIds).toContain("tc1");
+    expect(ranges[0].middleToolCallIds).toContain("tc2");
+    expect(ranges[0].finalAssistantTimestamp).toBe(400);
+  });
+
+  test("multi-chain sequence produces N ranges in order", () => {
+    const msgs = [
+      userMsg(100),
+      assistantWithTools(200, ["tc1"]),
+      toolResult(300, "tc1"),
+      assistantText(400),
+      userMsg(500),
+      assistantWithTools(600, ["tc2"]),
+      toolResult(700, "tc2"),
+      assistantText(800),
+    ];
+    const ranges = detectChains(msgs);
+    expect(ranges).toHaveLength(2);
+    expect(ranges[0].startUserTimestamp).toBe(100);
+    expect(ranges[0].finalAssistantTimestamp).toBe(400);
+    expect(ranges[1].startUserTimestamp).toBe(500);
+    expect(ranges[1].finalAssistantTimestamp).toBe(800);
+  });
+
+  test("open chain (no text-only close) is not emitted", () => {
+    const msgs = [
+      userMsg(100),
+      assistantWithTools(200, ["tc1"]),
+      toolResult(300, "tc1"),
+      // no text-only assistant close
+    ];
+    expect(detectChains(msgs)).toHaveLength(0);
+  });
+
+  test("synthetic chain message is not treated as a chain start", () => {
+    const msgs = [
+      syntheticChainMsg(50),   // should be skipped
+      assistantText(200),       // text-only but no prior real chain start
+    ];
+    expect(detectChains(msgs)).toHaveLength(0);
+  });
+
+  test("synthetic chain message in a real multi-chain sequence is a passthrough", () => {
+    // Represents a session after one chain was already compressed:
+    // synthetic summary, then the next real chain
+    const msgs = [
+      syntheticChainMsg(50),
+      userMsg(100),
+      assistantWithTools(200, ["tc1"]),
+      toolResult(300, "tc1"),
+      assistantText(400),
+    ];
+    const ranges = detectChains(msgs);
+    expect(ranges).toHaveLength(1);
+    expect(ranges[0].startUserTimestamp).toBe(100);
+  });
+
+  test("interrupted chain (user interrupts before text-only close) emits with null final", () => {
+    const msgs = [
+      userMsg(100),
+      assistantWithTools(200, ["tc1"]),
+      toolResult(300, "tc1"),
+      userMsg(400), // second user message interrupts before chain close
+      assistantText(500),
+    ];
+    const ranges = detectChains(msgs);
+    // First chain is interrupted → emitted with null finalAssistantTimestamp
+    expect(ranges).toHaveLength(2);
+    expect(ranges[0].startUserTimestamp).toBe(100);
+    expect(ranges[0].finalAssistantTimestamp).toBeNull();
+    // Second "chain" opened at ts=400, but text-only at 500 closes it (no tool calls needed)
+    // Actually, the second user at 400 starts a chain, but its assistant turn has no toolCalls
+    // so middleToolCallIds is empty and it immediately closes with the text-only assistant.
+    expect(ranges[1].startUserTimestamp).toBe(400);
+    expect(ranges[1].middleToolCallIds).toEqual([]);
+    expect(ranges[1].finalAssistantTimestamp).toBe(500);
+  });
+
+  test("agent-message flush: closing assistant absent from branch → newest chain reads as open", () => {
+    // Reproduces the runtime state at message_end: pi emits the event to extensions
+    // BEFORE appending event.message, so the just-closed assistant is missing here.
+    const branch = [
+      userMsg(100),
+      assistantWithTools(200, ["tc1"]),
+      toolResult(300, "tc1"),
+      assistantText(400),
+      userMsg(500),
+      assistantWithTools(600, ["tc2"]),
+      toolResult(700, "tc2"),
+      // closing assistant for chain 2 not yet in branch
+    ];
+    const ranges = detectChains(branch);
+    expect(ranges).toHaveLength(1); // only chain 1 closed; chain 2 open → dropped
+  });
+
+  test("middleToolCallIds contains all toolCallIds from assistant and toolResult messages", () => {
+    const msgs = [
+      userMsg(100),
+      assistantWithTools(200, ["tc1", "tc2"]),
+      toolResult(300, "tc1"),
+      toolResult(310, "tc2"),
+      assistantWithTools(400, ["tc3"]),
+      toolResult(500, "tc3"),
+      assistantText(600),
+    ];
+    const [range] = detectChains(msgs);
+    expect(range.middleToolCallIds.sort()).toEqual(["tc1", "tc2", "tc3"].sort());
+  });
+});
+
+describe("detectChains protectedToolCallIds", () => {
+  const chainMsgs = () => [
+    { role: "user", timestamp: 1 },
+    {
+      role: "assistant",
+      timestamp: 2,
+      content: [
+        { type: "toolCall", id: "tc-read", name: "read" },
+        { type: "toolCall", id: "tc-todo", name: "todowrite" },
+      ],
+    },
+    { role: "toolResult", toolCallId: "tc-read", toolName: "read", content: [{ type: "text", text: "x" }] },
+    { role: "toolResult", toolCallId: "tc-todo", toolName: "todowrite", content: [{ type: "text", text: "plan" }] },
+    { role: "assistant", timestamp: 5, content: [{ type: "text", text: "done" }] },
+  ];
+
+  test("is empty when no protectedTools are given", () => {
+    const [chain] = detectChains(chainMsgs());
+    expect(chain.protectedToolCallIds).toEqual([]);
+    expect(chain.middleToolCallIds.sort()).toEqual(["tc-read", "tc-todo"]);
+  });
+
+  test("collects protected ids by tool name from both branches", () => {
+    const [chain] = detectChains(chainMsgs(), (name) => name === "todowrite");
+    expect(chain.protectedToolCallIds).toEqual(["tc-todo"]);
+    expect(chain.middleToolCallIds.sort()).toEqual(["tc-read", "tc-todo"]);
+  });
+
+  test("does not leak protected ids across consecutive chains", () => {
+    const msgs = [
+      // chain 1: has a protected todowrite
+      { role: "user", timestamp: 1 },
+      { role: "assistant", timestamp: 2, content: [{ type: "toolCall", id: "tc-todo", name: "todowrite" }] },
+      { role: "toolResult", toolCallId: "tc-todo", toolName: "todowrite", content: [{ type: "text", text: "p" }] },
+      { role: "assistant", timestamp: 4, content: [{ type: "text", text: "done 1" }] },
+      // chain 2: only an unprotected read
+      { role: "user", timestamp: 5 },
+      { role: "assistant", timestamp: 6, content: [{ type: "toolCall", id: "tc-read", name: "read" }] },
+      { role: "toolResult", toolCallId: "tc-read", toolName: "read", content: [{ type: "text", text: "x" }] },
+      { role: "assistant", timestamp: 8, content: [{ type: "text", text: "done 2" }] },
+    ];
+    const chains = detectChains(msgs, (name) => name === "todowrite");
+    expect(chains.length).toBe(2);
+    expect(chains[0].protectedToolCallIds).toEqual(["tc-todo"]);
+    expect(chains[1].protectedToolCallIds).toEqual([]);
+    expect(chains[1].middleToolCallIds).toEqual(["tc-read"]);
+  });
+
+  test("collects protected ids by args.path via predicate", () => {
+    const msgs = [
+      { role: "user", timestamp: 1 },
+      {
+        role: "assistant",
+        timestamp: 2,
+        content: [
+          { type: "toolCall", id: "tc-skill", name: "read", input: { path: "/h/skills/x/SKILL.md" } },
+          { type: "toolCall", id: "tc-src", name: "read", input: { path: "/h/src/app.ts" } },
+        ],
+      },
+      { role: "toolResult", toolCallId: "tc-skill", toolName: "read", timestamp: 3, content: [] },
+      { role: "toolResult", toolCallId: "tc-src", toolName: "read", timestamp: 4, content: [] },
+      { role: "assistant", timestamp: 5, content: [{ type: "text", text: "done" }] },
+    ];
+    const pred = (name: string, args: unknown) =>
+      typeof (args as any)?.path === "string" && (args as any).path.includes("/skills/");
+    const [chain] = detectChains(msgs, pred);
+    expect(chain.protectedToolCallIds).toEqual(["tc-skill"]);
+  });
+
+  test("populates protectedToolCallIds on an interrupted (open→new user) chain", () => {
+    const msgs = [
+      { role: "user", timestamp: 1 },
+      { role: "assistant", timestamp: 2, content: [{ type: "toolCall", id: "tc-todo", name: "todowrite" }] },
+      { role: "toolResult", toolCallId: "tc-todo", toolName: "todowrite", content: [{ type: "text", text: "p" }] },
+      { role: "user", timestamp: 4 },
+    ];
+    const [interrupted] = detectChains(msgs, (name) => name === "todowrite");
+    expect(interrupted.finalAssistantTimestamp).toBeNull();
+    expect(interrupted.protectedToolCallIds).toEqual(["tc-todo"]);
+  });
+});
+
+describe("withClosingMessage", () => {
+  test("undefined closing returns the same array reference", () => {
+    const msgs = [userMsg(100)];
+    expect(withClosingMessage(msgs, undefined)).toBe(msgs);
+  });
+
+  test("appends closing when branch does not already end with it", () => {
+    const branch = [userMsg(100), assistantWithTools(200, ["tc1"]), toolResult(300, "tc1")];
+    const closing = assistantText(400);
+    const merged = withClosingMessage(branch, closing);
+    expect(merged).toHaveLength(4);
+    expect(merged[3]).toBe(closing);
+    expect(branch).toHaveLength(3); // original not mutated
+  });
+
+  test("does not double-append when branch already ends with the closing message (role+timestamp)", () => {
+    const closing = assistantText(400);
+    const branch = [userMsg(100), assistantWithTools(200, ["tc1"]), toolResult(300, "tc1"), closing];
+    expect(withClosingMessage(branch, closing)).toBe(branch);
+    // also dedups a distinct object with the same role+timestamp
+    const branch2 = [userMsg(100), assistantText(400)];
+    expect(withClosingMessage(branch2, assistantText(400))).toBe(branch2);
+  });
+
+  test("threading the closing message makes the newest chain close (effective window = K)", () => {
+    const branch = [
+      userMsg(100),
+      assistantWithTools(200, ["tc1"]),
+      toolResult(300, "tc1"),
+      assistantText(400),
+      userMsg(500),
+      assistantWithTools(600, ["tc2"]),
+      toolResult(700, "tc2"),
+    ];
+    const closing = assistantText(800);
+    const ranges = detectChains(withClosingMessage(branch, closing));
+    expect(ranges).toHaveLength(2);
+    expect(ranges[1].startUserTimestamp).toBe(500);
+    expect(ranges[1].finalAssistantTimestamp).toBe(800);
+  });
+});

package/src/chain-detector.ts

@@ -0,0 +1,128 @@
+import type { ChainRange } from "./types.js";
+
+/** Prefix that identifies a synthetic chain-compression user message. */
+const COMPRESSED_CHAIN_PREFIX = "<compressed-chain";
+
+function isSyntheticChainMessage(msg: any): boolean {
+  const content = msg.content;
+  if (typeof content === "string") return content.trimStart().startsWith(COMPRESSED_CHAIN_PREFIX);
+  if (!Array.isArray(content)) return false;
+  const first = content[0];
+  return first?.type === "text" && typeof first.text === "string" && first.text.trimStart().startsWith(COMPRESSED_CHAIN_PREFIX);
+}
+
+function hasToolCalls(msg: any): boolean {
+  return Array.isArray(msg.content) && msg.content.some((b: any) => b.type === "toolCall");
+}
+
+function collectToolCalls(msg: any): { id: string; name: string; args: unknown }[] {
+  if (!Array.isArray(msg.content)) return [];
+  return msg.content
+    .filter((b: any) => b.type === "toolCall" && b.id && b.name)
+    .map((b: any) => ({ id: b.id as string, name: b.name as string, args: b.input ?? b.arguments }));
+}
+
+type State = "idle" | "inChain";
+
+/**
+ * Walks an AgentMessage array and emits ChainRange records for each detectable chain.
+ *
+ * A chain is: [user message] → [assistant+toolResult turns...] → [text-only assistant].
+ * Synthetic chain messages (injected by chain-range-prune) are treated as passthroughs —
+ * not chain starts. This is defensive; the detector normally runs pre-compression.
+ *
+ * NOTE: Message identity uses `timestamp` (for user / final text-only assistant) and
+ * `toolCallId` sets (for middle tool-using turns). AgentMessage has no `.id` field.
+ *
+ * @param isProtected  Predicate over (toolName, args); matching calls are never pruned
+ *                     and their outputs are relocated verbatim into compressed chains.
+ */
+export function detectChains(
+  messages: any[],
+  isProtected: (toolName: string, args: unknown) => boolean = () => false,
+): ChainRange[] {
+  const ranges: ChainRange[] = [];
+  let state: State = "idle";
+  let chainStart: { timestamp: number } | null = null;
+  let middleIds = new Set<string>();
+  let protectedIds = new Set<string>();
+
+  const emitInterrupted = () => {
+    if (state === "inChain" && chainStart) {
+      ranges.push({
+        startUserTimestamp: chainStart.timestamp,
+        middleToolCallIds: [...middleIds],
+        protectedToolCallIds: [...protectedIds],
+        finalAssistantTimestamp: null,
+      });
+    }
+  };
+
+  for (const msg of messages) {
+    if (msg.role === "user") {
+      if (isSyntheticChainMessage(msg)) continue; // passthrough — not a chain start
+      emitInterrupted();
+      chainStart = { timestamp: msg.timestamp };
+      middleIds = new Set();
+      protectedIds = new Set();
+      state = "inChain";
+      continue;
+    }
+
+    if (state !== "inChain") continue;
+
+    if (msg.role === "assistant" && hasToolCalls(msg)) {
+      for (const { id, name, args } of collectToolCalls(msg)) {
+        middleIds.add(id);
+        if (isProtected(name, args)) protectedIds.add(id);
+      }
+      continue;
+    }
+
+    if (msg.role === "toolResult") {
+      if (msg.toolCallId) {
+        middleIds.add(msg.toolCallId);
+        // 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);
+      }
+      continue;
+    }
+
+    if (msg.role === "assistant" && !hasToolCalls(msg)) {
+      ranges.push({
+        startUserTimestamp: chainStart!.timestamp,
+        middleToolCallIds: [...middleIds],
+        protectedToolCallIds: [...protectedIds],
+        finalAssistantTimestamp: msg.timestamp,
+      });
+      chainStart = null;
+      middleIds = new Set();
+      protectedIds = new Set();
+      state = "idle";
+    }
+  }
+
+  // Open chain at end of input is intentionally dropped (in-flight).
+
+  return ranges;
+}
+
+/**
+ * Appends `closing` to a copy of `branchMessages` unless the array already ends with it.
+ *
+ * pi emits `message_end` to extensions BEFORE persisting the message to the session
+ * (agent-session.js `_processAgentEvent` runs `_emitExtensionEvent` ahead of
+ * `sessionManager.appendMessage`). At the agent-message flush boundary the just-closed
+ * final assistant is therefore still missing from `getBranch()`; without threading it
+ * in, the newest chain reads as open and the rolling window over-retains by one
+ * (effective K+1 instead of K). Identity is role+timestamp (AgentMessage has no id,
+ * matching the detector's own identity model), so a future pi that persists before
+ * emitting keeps this a no-op.
+ */
+export function withClosingMessage(branchMessages: any[], closing: any): any[] {
+  if (!closing) return branchMessages;
+  const last = branchMessages[branchMessages.length - 1];
+  if (last && last.role === closing.role && last.timestamp === closing.timestamp) return branchMessages;
+  return [...branchMessages, closing];
+}