pi-condense 2.0.0

package/package.json

@@ -0,0 +1,56 @@
+{
+  "name": "pi-condense",
+  "version": "2.0.0",
+  "description": "Pi extension that summarizes completed tool-call batches, replaces raw outputs with short stubs in future context, and recovers any original on demand via context_tree_query.",
+  "author": "Jacek Juraszek",
+  "license": "MIT",
+  "type": "module",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/jjuraszek/pi-condense.git"
+  },
+  "homepage": "https://github.com/jjuraszek/pi-condense#readme",
+  "bugs": {
+    "url": "https://github.com/jjuraszek/pi-condense/issues"
+  },
+  "keywords": [
+    "pi-package",
+    "pi",
+    "pi-coding-agent",
+    "context",
+    "context-management",
+    "pruning",
+    "summarization",
+    "ai"
+  ],
+  "engines": {
+    "node": ">=20"
+  },
+  "files": [
+    "index.ts",
+    "src/**/*.ts",
+    "README.md",
+    "CHANGELOG.md",
+    "PRUNING.md"
+  ],
+  "scripts": {
+    "test": "bun test src/"
+  },
+  "pi": {
+    "extensions": [
+      "./index.ts"
+    ]
+  },
+  "peerDependencies": {
+    "@earendil-works/pi-coding-agent": "*",
+    "@earendil-works/pi-ai": "*",
+    "@earendil-works/pi-tui": "*",
+    "@sinclair/typebox": "*"
+  },
+  "devDependencies": {
+    "@earendil-works/pi-ai": "^0.78.0",
+    "@earendil-works/pi-coding-agent": "^0.78.0",
+    "@earendil-works/pi-tui": "^0.78.0",
+    "@sinclair/typebox": "^0.34.49"
+  }
+}

package/src/batch-capture.ts

@@ -0,0 +1,226 @@
+import type { CapturedBatch, CapturedToolCall, BatchingMode } from "./types.js";
+
+/** Joins the text blocks of a ToolResultMessage into a single string. */
+export function extractToolResultText(msg: any): string {
+  const content: any[] = Array.isArray(msg?.content) ? msg.content : [];
+  return content
+    .filter((c: any) => c.type === "text")
+    .map((c: any) => c.text)
+    .join("\n");
+}
+
+/**
+ * Converts turn_end event data into a CapturedBatch.
+ * @param message      AssistantMessage (content: Array of TextContent|ThinkingContent|ToolCall)
+ * @param toolResults  ToolResultMessage[]
+ */
+export function captureBatch(
+  message: any,
+  toolResults: any[],
+  turnIndex: number,
+  timestamp: number
+): CapturedBatch {
+  const content: any[] = Array.isArray(message?.content) ? message.content : [];
+
+  // Collect assistant prose text
+  const assistantText = content
+    .filter((block: any) => block.type === "text")
+    .map((block: any) => block.text)
+    .join("\n")
+    .trim();
+
+  // Collect tool calls, matching each to its result
+  const toolCalls: CapturedToolCall[] = content
+    .filter((block: any) => block.type === "toolCall")
+    .map((block: any) => {
+      const match = toolResults.find((result: any) => result.toolCallId === block.id);
+
+      let resultText = "(no result)";
+      let isError = false;
+
+      if (match) {
+        resultText = extractToolResultText(match);
+        isError = match.isError ?? false;
+      }
+
+      return {
+        toolCallId: block.id,
+        toolName: block.name,
+        args: block.input ?? block.args ?? block.arguments ?? {},
+        resultText,
+        isError,
+      } satisfies CapturedToolCall;
+    });
+
+  return { turnIndex, timestamp, assistantText, toolCalls };
+}
+
+/**
+ * Scans a session branch for unsummarized tool results and groups them into CapturedBatches.
+ * Useful for capturing results from the current in-progress turn when a prune is triggered.
+ *
+ * @param branch            The session message branch (from ctx.sessionManager.getBranch())
+ * @param indexer           The pruner indexer to check for already-summarized IDs
+ * @param exclude  Optional predicate; matching tool calls are skipped (user-protected tools/paths)
+ */
+export function captureUnindexedBatchesFromSession(
+  branch: any[],
+  indexer: { isSummarized(id: string): boolean },
+  exclude: (toolName: string, args: unknown) => boolean = () => false
+): 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 batches: CapturedBatch[] = [];
+  // turnCounter increments for EVERY assistant message (not just prunable ones).
+  // This makes turnIndex stable across multiple prune cycles: pruning removes
+  // ToolResultMessages from the context event but leaves AssistantMessages in the
+  // session branch, so the count of all assistant messages never decreases and
+  // always matches Pi's own event.turnIndex numbering.
+  let turnCounter = 0;
+
+  // userTurnGroup increments on every user message seen while walking the branch.
+  // All assistant tool-call batches between two consecutive user messages share the
+  // same userTurnGroup. This is used by groupBatchesByMode to merge turns within
+  // a single user → final-agent-message span when batchingMode === "agent-message".
+  let userTurnGroup = 0;
+
+  for (const entry of branch) {
+    if (entry.type !== "message") continue;
+    const msg = entry.message;
+
+    // Advance userTurnGroup on every user message so all subsequent assistant
+    // batches get a new group number.
+    if (msg.role === "user") {
+      userTurnGroup++;
+      continue;
+    }
+
+    if (msg.role !== "assistant") continue;
+
+    // Stable turn index: count every assistant message regardless of pruning state
+    const currentTurnIndex = turnCounter++;
+
+    const content = Array.isArray(msg.content) ? msg.content : [];
+    const toolCallBlocks = content.filter((c: any) => c.type === "toolCall");
+
+    // Find tool calls that have results in this branch and are not yet summarized
+    const readyToPrune = toolCallBlocks.filter((tc: any) => {
+      const id = tc.id;
+      if (!id) return false;
+      if (indexer.isSummarized(id)) return false;
+      if (exclude(tc.name, tc.input ?? tc.arguments)) return false;
+      return resultMap.has(id);
+    });
+
+    if (readyToPrune.length > 0) {
+      const results = readyToPrune.map((tc: any) => resultMap.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
+      // an intermediate completed subset in the middle of a longer tool chain
+      // without accidentally capturing later unresolved calls from the same
+      // assistant message as "(no result)" placeholders.
+      const ts = entry.timestamp ? new Date(entry.timestamp).getTime() : (msg.timestamp ?? Date.now());
+      const batch = captureBatch(msg, results, currentTurnIndex, ts);
+      batches.push({
+        ...batch,
+        toolCalls: batch.toolCalls.filter((tc) => readyIds.has(tc.toolCallId)),
+        // Tag with the current group so flushPending can merge by mode
+        userTurnGroup,
+      });
+    }
+  }
+
+  return batches;
+}
+
+/** Serializes a single CapturedBatch into readable text for the summarizer LLM. */
+export function serializeBatchForSummarizer(batch: CapturedBatch): string {
+  const parts: string[] = [];
+
+  if (batch.assistantText) {
+    parts.push(`Assistant said: ${batch.assistantText}\n`);
+  }
+
+  const toolParts = batch.toolCalls.map((tc) => {
+    const status = tc.isError ? "ERROR" : "OK";
+    const argsJson = JSON.stringify(tc.args, null, 2);
+
+    let resultText = tc.resultText;
+    const MAX_CHARS = 2000;
+    if (resultText.length > MAX_CHARS) {
+      const remaining = resultText.length - MAX_CHARS;
+      resultText = resultText.slice(0, MAX_CHARS) + ` ...[${remaining} chars truncated]`;
+    }
+
+    return `Tool: ${tc.toolName}(${argsJson})\nResult (${status}): ${resultText}`;
+  });
+
+  parts.push(toolParts.join("\n---\n"));
+
+  return parts.join("\n");
+}
+
+/**
+ * Groups CapturedBatches according to the chosen batching mode.
+ *
+ * - "turn"          : returns the input array unchanged (one summary per assistant turn).
+ * - "agent-message" : merges all consecutive batches that share the same `userTurnGroup`
+ *                     into a single CapturedBatch, producing one summary per
+ *                     user → final-agent-message span.
+ *
+ * Batches without a `userTurnGroup` (e.g. from the live `turn_end` capture path) are
+ * always passed through one-per-batch regardless of mode — grouping only applies to
+ * batches captured from the session branch scan.
+ *
+ * Merge rules:
+ *   - `assistantText` = non-empty values joined with "\n\n"
+ *   - `toolCalls`     = concatenation in original order
+ *   - `turnIndex`     = last batch's turnIndex (latest turn in the group)
+ *   - `timestamp`     = last batch's timestamp
+ *   - `userTurnGroup` = shared group value of the merged batches
+ */
+export function groupBatchesByMode(batches: CapturedBatch[], mode: BatchingMode): CapturedBatch[] {
+  if (mode !== "agent-message") return batches;
+
+  const out: CapturedBatch[] = [];
+  // current tracks the mutable merged batch being built for the current group.
+  // We spread into a plain object so we can mutate it without affecting the source.
+  let current: CapturedBatch & { userTurnGroup: number } | null = null;
+
+  for (const batch of batches) {
+    // Batches without a group key are passed through individually; they break
+    // any open merge group too since we can't confidently assign them a span.
+    if (batch.userTurnGroup === undefined) {
+      current = null;
+      out.push(batch);
+      continue;
+    }
+
+    if (current !== null && current.userTurnGroup === batch.userTurnGroup) {
+      // Same span — merge into the current accumulated batch
+      const textParts = [current.assistantText, batch.assistantText].filter(Boolean);
+      current.assistantText = textParts.join("\n\n");
+      current.toolCalls = current.toolCalls.concat(batch.toolCalls);
+      // Advance to the latest turn metadata
+      current.turnIndex = batch.turnIndex;
+      current.timestamp = batch.timestamp;
+    } else {
+      // New group — create a fresh accumulated batch (shallow copy so mutations
+      // to `current` do not bleed back into the original `batch` object)
+      current = { ...batch, userTurnGroup: batch.userTurnGroup };
+      out.push(current);
+    }
+  }
+
+  return out;
+}

package/src/block-refs.test.ts

@@ -0,0 +1,42 @@
+import { describe, expect, test } from "bun:test";
+import { BlockRefIssuer } from "./block-refs.js";
+
+describe("BlockRefIssuer", () => {
+  test("first issue() returns b1", () => {
+    const issuer = new BlockRefIssuer();
+    expect(issuer.issue()).toBe("b1");
+  });
+
+  test("subsequent calls are monotonic", () => {
+    const issuer = new BlockRefIssuer();
+    expect(issuer.issue()).toBe("b1");
+    expect(issuer.issue()).toBe("b2");
+    expect(issuer.issue()).toBe("b3");
+  });
+
+  test("rebuildFrom([]) → next is b1", () => {
+    const issuer = new BlockRefIssuer();
+    issuer.rebuildFrom([]);
+    expect(issuer.issue()).toBe("b1");
+  });
+
+  test("rebuildFrom([b1, b2]) → next is b3", () => {
+    const issuer = new BlockRefIssuer();
+    issuer.rebuildFrom(["b1", "b2"]);
+    expect(issuer.issue()).toBe("b3");
+  });
+
+  test("rebuildFrom gapped sequence → max+1, not gap-fill", () => {
+    const issuer = new BlockRefIssuer();
+    issuer.rebuildFrom(["b1", "b3"]);
+    expect(issuer.issue()).toBe("b4");
+  });
+
+  test("rebuildFrom resets counter even after prior issues", () => {
+    const issuer = new BlockRefIssuer();
+    issuer.issue(); // b1
+    issuer.issue(); // b2
+    issuer.rebuildFrom(["b5"]);
+    expect(issuer.issue()).toBe("b6");
+  });
+});

package/src/block-refs.ts

@@ -0,0 +1,16 @@
+export class BlockRefIssuer {
+  private next = 1;
+
+  issue(): string {
+    return `b${this.next++}`;
+  }
+
+  rebuildFrom(existingBlockIds: string[]): void {
+    if (existingBlockIds.length === 0) {
+      this.next = 1;
+      return;
+    }
+    const max = Math.max(...existingBlockIds.map((id) => parseInt(id.slice(1), 10)));
+    this.next = max + 1;
+  }
+}

package/src/budget.test.ts

@@ -0,0 +1,66 @@
+import { describe, it, expect } from "bun:test";
+import { shouldBudgetFlush, shouldDeltaFlush, usageFraction } from "./budget.js";
+
+const usage = (tokens: number | null, contextWindow: number) =>
+  ({ tokens, contextWindow, percent: null }) as any;
+
+describe("shouldBudgetFlush", () => {
+  it("is false when threshold is null", () => {
+    expect(shouldBudgetFlush(usage(900, 1000), null)).toBe(false);
+  });
+  it("is false for non-positive or >1 thresholds", () => {
+    expect(shouldBudgetFlush(usage(900, 1000), 0)).toBe(false);
+    expect(shouldBudgetFlush(usage(900, 1000), 1.5)).toBe(false);
+  });
+  it("is false when usage is undefined", () => {
+    expect(shouldBudgetFlush(undefined, 0.8)).toBe(false);
+  });
+  it("is false when tokens is null (post-compaction)", () => {
+    expect(shouldBudgetFlush(usage(null, 1000), 0.8)).toBe(false);
+  });
+  it("is false when contextWindow is non-positive", () => {
+    expect(shouldBudgetFlush(usage(900, 0), 0.8)).toBe(false);
+  });
+  it("is true at or over the threshold, false under", () => {
+    expect(shouldBudgetFlush(usage(800, 1000), 0.8)).toBe(true);
+    expect(shouldBudgetFlush(usage(900, 1000), 0.8)).toBe(true);
+    expect(shouldBudgetFlush(usage(799, 1000), 0.8)).toBe(false);
+  });
+
+  it("treats threshold of exactly 1.0 as valid (flush only at 100%)", () => {
+    expect(shouldBudgetFlush(usage(1000, 1000), 1)).toBe(true);
+    expect(shouldBudgetFlush(usage(999, 1000), 1)).toBe(false);
+  });
+});
+
+describe("usageFraction", () => {
+  it("returns null for undefined / null tokens / non-positive window", () => {
+    expect(usageFraction(undefined)).toBeNull();
+    expect(usageFraction(usage(null, 1000))).toBeNull();
+    expect(usageFraction(usage(900, 0))).toBeNull();
+  });
+  it("returns the 0–1 fraction", () => {
+    expect(usageFraction(usage(750, 1000))).toBe(0.75);
+  });
+});
+
+describe("shouldDeltaFlush", () => {
+  it("is false when delta is null, non-positive, or >1", () => {
+    expect(shouldDeltaFlush(usage(900, 1000), 0.5, null)).toBe(false);
+    expect(shouldDeltaFlush(usage(900, 1000), 0.5, 0)).toBe(false);
+    expect(shouldDeltaFlush(usage(900, 1000), 0.5, 1.5)).toBe(false);
+  });
+  it("is false when previousFraction is null (first turn / post-restart)", () => {
+    expect(shouldDeltaFlush(usage(900, 1000), null, 0.15)).toBe(false);
+  });
+  it("is false when usage missing or tokens null", () => {
+    expect(shouldDeltaFlush(undefined, 0.5, 0.15)).toBe(false);
+    expect(shouldDeltaFlush(usage(null, 1000), 0.5, 0.15)).toBe(false);
+  });
+  it("fires when the jump meets the delta, not below", () => {
+    expect(shouldDeltaFlush(usage(700, 1000), 0.5, 0.15)).toBe(true);  // 0.20 >= 0.15
+    expect(shouldDeltaFlush(usage(650, 1000), 0.5, 0.15)).toBe(true);  // 0.15 exactly
+    expect(shouldDeltaFlush(usage(640, 1000), 0.5, 0.15)).toBe(false); // 0.14 < 0.15
+    expect(shouldDeltaFlush(usage(600, 1000), 0.5, 0.15)).toBe(false); // 0.10 < 0.15
+  });
+});

package/src/budget.ts

@@ -0,0 +1,39 @@
+import type { ContextUsage } from "@earendil-works/pi-coding-agent";
+
+/**
+ * True iff a budget-triggered flush should fire. Computes the ratio ourselves
+ * (tokens / contextWindow, a 0–1 fraction) rather than using ContextUsage.percent
+ * (a 0–100 value, null when tokens is null). tokens is also null right after a
+ * compaction — guarded here.
+ */
+export function shouldBudgetFlush(
+  usage: ContextUsage | undefined,
+  threshold: number | null,
+): boolean {
+  if (threshold == null || threshold <= 0 || threshold > 1) return false;
+  if (!usage || usage.tokens == null || !(usage.contextWindow > 0)) return false;
+  return usage.tokens / usage.contextWindow >= threshold;
+}
+
+/** 0–1 usage fraction, or null when usage is missing / tokens null / window non-positive. */
+export function usageFraction(usage: ContextUsage | undefined): number | null {
+  if (!usage || usage.tokens == null || !(usage.contextWindow > 0)) return null;
+  return usage.tokens / usage.contextWindow;
+}
+
+/**
+ * True iff this turn's usage fraction rose by at least `delta` versus the previous turn.
+ * Mirrors shouldBudgetFlush's guards. previousFraction === null (first turn or post-restart)
+ * never fires; the absolute autoBudgetThreshold covers that gap.
+ */
+export function shouldDeltaFlush(
+  usage: ContextUsage | undefined,
+  previousFraction: number | null,
+  delta: number | null,
+): boolean {
+  if (delta == null || delta <= 0 || delta > 1) return false;
+  if (previousFraction == null) return false;
+  const current = usageFraction(usage);
+  if (current == null) return false;
+  return current - previousFraction >= delta;
+}