pi-condense 2.0.0

package/src/stats.ts

@@ -0,0 +1,190 @@
+import type { SummarizerStats, ExternalCostUpdate, LiveReclaim } from "./types.js";
+import { CUSTOM_TYPE_STATS, EXTERNAL_COST_CHANNEL, EXTERNAL_COST_SOURCE } from "./types.js";
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+
+/**
+ * Usage shape returned by the LLM `complete()` call.
+ * Mirrors the `Usage` interface from `@earendil-works/pi-ai` but declared locally
+ * so we don't need a runtime import just for the type.
+ */
+interface Usage {
+  input: number;
+  output: number;
+  cacheRead: number;
+  cacheWrite: number;
+  totalTokens: number;
+  cost: {
+    input: number;
+    output: number;
+    cacheRead: number;
+    cacheWrite: number;
+    total: number;
+  };
+}
+
+/**
+ * Accumulates cumulative token/cost stats for summarizer LLM calls.
+ * Stats are persisted to the session via `pi.appendEntry(CUSTOM_TYPE_STATS, ...)`
+ * and reconstructed on `session_start` / `session_tree`.
+ */
+export class StatsAccumulator {
+  private stats: SummarizerStats = {
+    totalInputTokens: 0,
+    totalOutputTokens: 0,
+    totalCost: 0,
+    callCount: 0,
+    chainsCompressed: 0,
+    rangesSummarized: 0,
+  };
+  private baseline = { totalInputTokens: 0, totalOutputTokens: 0, totalCost: 0 };
+  private liveReclaim: LiveReclaim | undefined = undefined;
+
+  /** Add usage data from one summarizer LLM call. */
+  add(usage: Usage): void {
+    this.stats.totalInputTokens += usage.input ?? 0;
+    this.stats.totalOutputTokens += usage.output ?? 0;
+    this.stats.totalCost += usage.cost?.total ?? 0;
+    this.stats.callCount += 1;
+  }
+
+  /** Return session-delta spend (current stats minus baseline set at reconstructFromSession). */
+  getSessionDelta(): { totalCost: number; inputTokens: number; outputTokens: number } {
+    return {
+      totalCost: this.stats.totalCost - this.baseline.totalCost,
+      inputTokens: this.stats.totalInputTokens - this.baseline.totalInputTokens,
+      outputTokens: this.stats.totalOutputTokens - this.baseline.totalOutputTokens,
+    };
+  }
+
+  /** Store the before/after context-char measurement from the last prune. */
+  setLiveReclaim(beforeChars: number, afterChars: number): void {
+    this.liveReclaim = { beforeChars, afterChars };
+  }
+
+  /** Return the last live-reclaim measurement, or undefined if none yet. */
+  getLiveReclaim(): LiveReclaim | undefined {
+    return this.liveReclaim;
+  }
+
+  /** Return a snapshot of the current cumulative stats. */
+  getStats(): SummarizerStats {
+    return { ...this.stats };
+  }
+
+  /** Increment the chain-compression counter. */
+  addChainsCompressed(n: number): void {
+    this.stats.chainsCompressed += n;
+  }
+
+  /** Increment the fused-range-summary counter. */
+  addRangesSummarized(n: number): void {
+    this.stats.rangesSummarized += n;
+  }
+
+  /** Reset all accumulated stats to zero. Produces the same state as a fresh accumulator. */
+  reset(): void {
+    this.stats = {
+      totalInputTokens: 0,
+      totalOutputTokens: 0,
+      totalCost: 0,
+      callCount: 0,
+      chainsCompressed: 0,
+      rangesSummarized: 0,
+    };
+    this.baseline = { totalInputTokens: 0, totalOutputTokens: 0, totalCost: 0 };
+    this.liveReclaim = undefined;
+  }
+
+  /** Serialize stats for session persistence. */
+  toJSON(): SummarizerStats {
+    return { ...this.stats };
+  }
+
+  /** Restore stats from a previously persisted snapshot. */
+  fromJSON(data: SummarizerStats): void {
+    this.stats = {
+      totalInputTokens: data.totalInputTokens ?? 0,
+      totalOutputTokens: data.totalOutputTokens ?? 0,
+      totalCost: data.totalCost ?? 0,
+      callCount: data.callCount ?? 0,
+      chainsCompressed: data.chainsCompressed ?? 0,
+      rangesSummarized: data.rangesSummarized ?? 0,
+    };
+  }
+
+  /**
+   * Reconstruct stats from session history by scanning all custom entries
+   * with customType === CUSTOM_TYPE_STATS.
+   */
+  reconstructFromSession(ctx: ExtensionContext): void {
+    this.reset();
+    const branch = ctx.sessionManager.getBranch();
+    for (const entry of branch) {
+      if (
+        entry.type === "custom" &&
+        (entry as any).customType === CUSTOM_TYPE_STATS
+      ) {
+        const data = (entry as any).data as SummarizerStats;
+        if (data) {
+          this.fromJSON(data);
+        }
+      }
+    }
+    this.baseline = {
+      totalInputTokens: this.stats.totalInputTokens,
+      totalOutputTokens: this.stats.totalOutputTokens,
+      totalCost: this.stats.totalCost,
+    };
+  }
+
+  /**
+   * Persist current stats to the session.
+   * Each call appends a new entry; on reconstructFromSession we scan
+   * all entries and apply the LAST one (since each entry is a full snapshot).
+   */
+  persist(pi: ExtensionAPI): void {
+    pi.appendEntry(CUSTOM_TYPE_STATS, this.toJSON());
+  }
+}
+
+// ── Formatting helpers ──────────────────────────────────────────────────────
+
+/** Format compact counts like Pi's status line (e.g. "1.2k", "340") */
+export function formatCompactCount(n: number): string {
+  if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`;
+  if (n >= 1_000) return `${(n / 1_000).toFixed(1)}k`;
+  return String(n);
+}
+
+/** Format token counts like Pi's status line (e.g. "1.2k", "340") */
+export function formatTokens(n: number): string {
+  return formatCompactCount(n);
+}
+
+/** Format live char progress like "1.2k summary chars / 8.4k raw chars". */
+export function formatCharProgress(receivedChars: number, rawChars?: number): string {
+  const receivedLabel = `${formatCompactCount(receivedChars)} summary char${receivedChars === 1 ? "" : "s"}`;
+  if (rawChars == null) return receivedLabel;
+  return `${receivedLabel} / ${formatCompactCount(rawChars)} raw char${rawChars === 1 ? "" : "s"}`;
+}
+
+/** Format cost like "$0.003" */
+export function formatCost(n: number): string {
+  if (n < 0.001 && n > 0) return `<$0.001`;
+  return `$${n.toFixed(3)}`;
+}
+
+/**
+ * Emit the session-delta cost from `accumulator` on EXTERNAL_COST_CHANNEL.
+ * Idempotent from the aggregator's perspective: keyed by source, re-emitting overwrites.
+ */
+export function emitExternalCost(pi: ExtensionAPI, accumulator: StatsAccumulator): void {
+  const delta = accumulator.getSessionDelta();
+  const payload: ExternalCostUpdate = {
+    source: EXTERNAL_COST_SOURCE,
+    totalCost: delta.totalCost,
+    inputTokens: delta.inputTokens,
+    outputTokens: delta.outputTokens,
+  };
+  pi.events.emit(EXTERNAL_COST_CHANNEL, payload);
+}

package/src/summarizer.test.ts

@@ -0,0 +1,17 @@
+import { describe, it, expect } from "bun:test";
+import { isUsableSummary } from "./summarizer.js";
+
+describe("isUsableSummary", () => {
+  it("accepts non-empty text that stopped normally", () => {
+    expect(isUsableSummary("- did a thing", "stop")).toBe(true);
+  });
+  it("rejects empty text", () => {
+    expect(isUsableSummary("", "stop")).toBe(false);
+  });
+  it("rejects whitespace-only text", () => {
+    expect(isUsableSummary("   \n\t ", "stop")).toBe(false);
+  });
+  it("rejects truncated output even with text", () => {
+    expect(isUsableSummary("- partial", "length")).toBe(false);
+  });
+});

package/src/summarizer.ts

@@ -0,0 +1,262 @@
+import { stream } from "@earendil-works/pi-ai";
+import type { AssistantMessage } from "@earendil-works/pi-ai";
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import type {
+  CapturedBatch,
+  ContextPruneConfig,
+  SummarizerThinking,
+  SummarizeBatchOptions,
+  SummarizeBatchesOptions,
+  SummarizeResult,
+} from "./types.js";
+import { serializeBatchForSummarizer } from "./batch-capture.js";
+
+const SYSTEM_PROMPT = `You are summarizing a batch of tool calls made by an AI coding assistant.
+For each tool call provide:
+- Tool name and a one-sentence description of what it did
+- Key outcome: success/failure and the most important data returned
+- Any findings the future conversation needs to remember
+
+Keep each tool call to 1-3 bullet points. Be concise.`;
+
+const RANGE_SYSTEM_PROMPT = `You are fusing several per-step summaries of one CLOSED sub-task from an AI coding assistant's history into a SINGLE cohesive summary.
+- Merge overlapping or repeated information; do not restate each step separately.
+- Preserve concrete outcomes, decisions, file paths, identifiers, and anything later work needs to remember.
+- Keep any reference tokens like \`t12\` or \`b3\` intact.
+- Be concise: a short narrative or a few grouped bullets, not one bullet per step.`;
+
+export function summarizerThinkingOptions(config: ContextPruneConfig): Record<string, unknown> {
+  const level: SummarizerThinking = config.summarizerThinking;
+  if (level === "default") {
+    return {};
+  }
+
+  // stream()/complete() accept provider-level options. For reasoning-capable providers,
+  // pi-ai adapters translate reasoningEffort into the provider-specific field.
+  // "off" intentionally sends no effort; adapters that support explicit disable
+  // handle that the same way as an absent effort, while preserving compatibility.
+  return { reasoningEffort: level === "off" ? undefined : level };
+}
+
+/**
+ * Returns the model to use for summarization.
+ * config.summarizerModel === "default" => ctx.model
+ * "provider/model-id" => ctx.modelRegistry.find(provider, modelId), fallback to ctx.model with warning
+ */
+export function resolveModel(config: ContextPruneConfig, ctx: ExtensionContext): any {
+  if (config.summarizerModel === "default") {
+    return ctx.model;
+  }
+
+  const slashIndex = config.summarizerModel.indexOf("/");
+  if (slashIndex === -1) {
+    ctx.ui.notify(
+      `pruner: invalid summarizerModel "${config.summarizerModel}", expected "provider/model-id". Falling back to default model.`,
+      "warning"
+    );
+    return ctx.model;
+  }
+
+  const provider = config.summarizerModel.slice(0, slashIndex);
+  const modelId = config.summarizerModel.slice(slashIndex + 1);
+
+  const found = ctx.modelRegistry.find(provider, modelId);
+  if (!found) {
+    ctx.ui.notify(
+      `pruner: model "${config.summarizerModel}" not found in registry. Falling back to default model.`,
+      "warning"
+    );
+    return ctx.model;
+  }
+
+  return found;
+}
+
+function receivedTextChars(message: AssistantMessage): number {
+  return message.content.reduce((sum, content) => {
+    return content.type === "text" ? sum + content.text.length : sum;
+  }, 0);
+}
+
+/** A summary is usable only if it has non-whitespace text and was not truncated. */
+export function isUsableSummary(llmText: string, stopReason: string): boolean {
+  return llmText.trim().length > 0 && stopReason !== "length";
+}
+
+/**
+ * Shared LLM-call machinery for both per-batch and range summarization.
+ * `userMessage` already embeds the relevant system prompt as leading text
+ * (the summarizer is a single-user-message call). Returns the formatted text
+ * + usage, or null on failure. Abort errors are re-thrown so flushPending can
+ * detect options.signal.aborted and restore state without a UI error.
+ */
+async function runSummarization(
+  userMessage: string,
+  config: ContextPruneConfig,
+  ctx: ExtensionContext,
+  options: SummarizeBatchOptions
+): Promise<SummarizeResult | null> {
+  // Fast-fail if already aborted before we even start.
+  if (options.signal?.aborted) throw new Error("summarize: aborted before start");
+
+  try {
+    const model = resolveModel(config, ctx);
+
+    const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
+    if (!auth.ok) {
+      const authMessage = "error" in auth ? auth.error : "authentication failed";
+      ctx.ui.notify(`pruner: summarization failed: ${authMessage}`, "error");
+      return null;
+    }
+
+    // Pass the abort signal so the underlying fetch is cancelled immediately
+    // when the user presses Esc while the tool is running.
+    const responseStream = stream(
+      model,
+      {
+        messages: [
+          {
+            role: "user",
+            content: [{ type: "text", text: userMessage }],
+            timestamp: Date.now(),
+          },
+        ],
+      },
+      { apiKey: auth.apiKey, headers: auth.headers, signal: options.signal, ...summarizerThinkingOptions(config) }
+    );
+
+    let lastReportedChars = -1;
+    options.onTextProgress?.(0);
+    const reportTextProgress = (message: AssistantMessage) => {
+      const chars = receivedTextChars(message);
+      if (chars !== lastReportedChars) {
+        lastReportedChars = chars;
+        options.onTextProgress?.(chars);
+      }
+    };
+
+    for await (const event of responseStream) {
+      // Belt-and-suspenders: break early when signal fires mid-stream.
+      if (options.signal?.aborted) break;
+      if (event.type === "text_start" || event.type === "text_delta" || event.type === "text_end") {
+        reportTextProgress(event.partial);
+      }
+    }
+
+    // If signal fired while we were iterating, propagate the abort so
+    // flushPending can detect it and restore batches.
+    if (options.signal?.aborted) throw new Error("summarize: aborted during stream");
+
+    const response = await responseStream.result();
+    reportTextProgress(response);
+    // stopReason "aborted" means the provider cut the stream short (e.g. signal
+    // fired just before the final chunk). Treat identically to the signal check
+    // above — throw so flushPending's catch can detect options.signal.aborted.
+    if (response.stopReason === "aborted") {
+      throw new Error("summarize: stream stopped with reason aborted");
+    }
+    if (response.stopReason === "error") {
+      throw new Error(response.errorMessage ?? "Summarizer stopped with reason: error");
+    }
+
+    const llmText = response.content
+      .filter((c: any) => c.type === "text")
+      .map((c: any) => c.text)
+      .join("\n");
+
+    if (!isUsableSummary(llmText, response.stopReason)) return null;
+
+    return {
+      summaryText: llmText,
+      usage: response.usage,
+    };
+  } catch (err: any) {
+    // Propagate abort errors upward so flushPending can check signal.aborted
+    // and return { ok: false, reason: "aborted" } without showing a UI error.
+    if (options.signal?.aborted) throw err;
+    ctx.ui.notify(
+      `pruner: summarization failed: ${err.message}`,
+      "error"
+    );
+    return null;
+  }
+}
+
+/**
+ * Summarizes a captured batch. Returns formatted markdown string, or null on failure.
+ * Shows user-visible errors via ctx.ui.notify.
+ */
+export async function summarizeBatch(
+  batch: CapturedBatch,
+  config: ContextPruneConfig,
+  ctx: ExtensionContext,
+  options: SummarizeBatchOptions = {}
+): Promise<SummarizeResult | null> {
+  const serialized = serializeBatchForSummarizer(batch);
+  const userMessage =
+    SYSTEM_PROMPT + "\n\n<tool-call-batch>\n" + serialized + "\n</tool-call-batch>";
+  return runSummarization(userMessage, config, ctx, options);
+}
+
+/**
+ * Fuses a closed chain's already-computed per-batch summaries into one cohesive
+ * range summary (recursive summarization). Input is the span's per-batch summary
+ * text — small and already pruned — so this never re-sends raw tool output.
+ * Returns the fused text + usage, or null on failure. Used by chain compression
+ * to replace the concatenated per-batch body with a single coherent summary.
+ */
+export async function summarizeRange(
+  perBatchSummaryText: string,
+  config: ContextPruneConfig,
+  ctx: ExtensionContext,
+  options: SummarizeBatchOptions = {}
+): Promise<SummarizeResult | null> {
+  const userMessage =
+    RANGE_SYSTEM_PROMPT + "\n\n<sub-task-summaries>\n" + perBatchSummaryText + "\n</sub-task-summaries>";
+  return runSummarization(userMessage, config, ctx, options);
+}
+
+/**
+ * Summarizes multiple captured batches — one LLM call per batch, run in parallel.
+ *
+ * Returns an array of per-batch results. Each element is either a SummarizeResult
+ * (success) or null (that specific batch's call failed). The array length always
+ * equals batches.length so callers can zip by index.
+ *
+ * Rationale for parallel-per-batch instead of a single merged call:
+ *   • Each batch becomes its own summary message (one per turn), so they can be
+ *     rendered, browsed, and recovered independently via context_tree_query.
+ *   • Parallel calls give similar end-to-end latency to a single merged call while
+ *     keeping the summaries strictly separated.
+ */
+export async function summarizeBatches(
+  batches: CapturedBatch[],
+  config: ContextPruneConfig,
+  ctx: ExtensionContext,
+  options: SummarizeBatchesOptions = {}
+): Promise<Array<SummarizeResult | null>> {
+  if (batches.length === 0) return [];
+  // Single batch — delegate to the single-batch path (no extra overhead)
+  if (batches.length === 1) {
+    return [
+      await summarizeBatch(batches[0], config, ctx, {
+        signal: options.signal,
+        onTextProgress: (receivedChars) => {
+          options.onBatchTextProgress?.(0, 1, batches[0], receivedChars);
+        },
+      }),
+    ];
+  }
+
+  // Multiple batches — run in parallel; each produces its own SummarizeResult
+  return Promise.all(
+    batches.map((batch, index) =>
+      summarizeBatch(batch, config, ctx, {
+        signal: options.signal,
+        onTextProgress: (receivedChars) => {
+          options.onBatchTextProgress?.(index, batches.length, batch, receivedChars);
+        },
+      })
+    )
+  );
+}

package/src/summary-refs.ts

@@ -0,0 +1,61 @@
+import type { CapturedBatch } from "./types.js";
+
+export interface SummaryToolCallRef {
+  shortId: string;
+  toolCallId: string;
+}
+
+export interface SummaryMessageDetailsLike {
+  toolCallRefs?: SummaryToolCallRef[];
+  toolCallIds?: string[];
+}
+
+const SHORT_ID_PREFIX = "t";
+
+export function buildShortToolCallRefs(
+  toolCallIds: string[],
+  startIndex: number,
+): { refs: SummaryToolCallRef[]; nextIndex: number } {
+  const refs = toolCallIds.map((toolCallId, offset) => ({
+    shortId: `${SHORT_ID_PREFIX}${startIndex + offset}`,
+    toolCallId,
+  }));
+  return { refs, nextIndex: startIndex + refs.length };
+}
+
+export function normalizeSummaryToolCallRefs(details: unknown): SummaryToolCallRef[] {
+  if (!details || typeof details !== "object") return [];
+
+  const raw = details as SummaryMessageDetailsLike;
+  if (Array.isArray(raw.toolCallRefs)) {
+    return raw.toolCallRefs
+      .filter(
+        (ref): ref is SummaryToolCallRef =>
+          !!ref && typeof ref.shortId === "string" && typeof ref.toolCallId === "string",
+      )
+      .map((ref) => ({ shortId: ref.shortId, toolCallId: ref.toolCallId }));
+  }
+
+  if (Array.isArray(raw.toolCallIds)) {
+    return raw.toolCallIds.filter((id): id is string => typeof id === "string").map((id) => ({ shortId: id, toolCallId: id }));
+  }
+
+  return [];
+}
+
+export function formatSummaryToolCallRefs(refs: SummaryToolCallRef[]): string {
+  const refList = refs.map((ref) => `\`${ref.shortId}\``).join(", ");
+  return (
+    `\n\n---\n**Summarized tool refs**: ${refList}\n` +
+    `Use \`context_tree_query\` with these refs to retrieve the original full outputs.`
+  );
+}
+
+export function makeSummaryDetails(batch: CapturedBatch, refs: SummaryToolCallRef[]) {
+  return {
+    toolCallRefs: refs,
+    toolNames: batch.toolCalls.map((tc) => tc.toolName),
+    turnIndex: batch.turnIndex,
+    timestamp: batch.timestamp,
+  };
+}