pi-provider-cursor-ask 0.1.0

package/src/stream/session-state.ts

@@ -0,0 +1,499 @@
+/**
+ * Conversation state that outlives a single bridge: checkpoints, blob stores,
+ * mid-pause snapshots, and the keys everything is filed under.
+ *
+ * A Cursor turn can lose its bridge while parked waiting for tool results. What
+ * survives that is stored here — an upstream checkpoint when we have one, plus
+ * enough history fingerprinting for ./recovery.ts to decide between resuming
+ * from the checkpoint and rebuilding the whole conversation.
+ *
+ * Entries are keyed by Pi session id when available, falling back to a hash of
+ * the message history, and are evicted on a TTL so an abandoned session cannot
+ * pin its blob store forever.
+ */
+import { createHash, randomUUID } from "node:crypto";
+
+import { fromBinary } from "@bufbuild/protobuf";
+
+import { ConversationStateStructureSchema } from "../proto/agent_pb.js";
+import {
+  activeBridges,
+  cleanupBridge,
+  destroyAllIdleBridges,
+  destroyIdleBridge,
+} from "./bridge-session.js";
+import { debugLog } from "./debug-log.js";
+import { textContent } from "./message-parsing.js";
+import {
+  clearStoredMidPauseMetadata as clearStoredMidPauseMetadataImpl,
+  fingerprintCompletedTurns as fingerprintCompletedTurnsImpl,
+} from "./recovery.js";
+import {
+  deleteConversationJournal,
+  evictStaleJournals,
+  readConversationJournal,
+  writeConversationJournal,
+} from "./run-journal.js";
+import {
+  CONVERSATION_TTL_MS,
+  MAX_ACTIVE_BLOB_ENTRIES,
+  MAX_CHECKPOINT_BYTES,
+  MAX_CONVERSATION_BLOB_BYTES,
+} from "./tuning.js";
+import type {
+  ChatCompletionRequest,
+  OpenAIMessage,
+  ParsedTurn,
+  StoredConversation,
+} from "./types.js";
+
+export const conversationStates = new Map<string, StoredConversation>();
+
+const sessionLocks = new Map<string, Promise<void>>();
+
+function persistJournal(convKey: string, stored: StoredConversation): void {
+  writeConversationJournal(convKey, stored);
+}
+
+/**
+ * Resolve conversation state: memory first, then durable journal.
+ * Hydrates the in-memory map on a journal hit so later turns share one object.
+ */
+export function getOrHydrateConversation(convKey: string): StoredConversation | undefined {
+  const memory = conversationStates.get(convKey);
+  if (memory) return memory;
+  const fromDisk = readConversationJournal(convKey);
+  if (!fromDisk) return undefined;
+  conversationStates.set(convKey, fromDisk);
+  return fromDisk;
+}
+
+export function cleanupAllSessionState(): void {
+  debugLog("session.cleanup_all", {
+    activeBridgeCount: activeBridges.size,
+    conversationCount: conversationStates.size,
+  });
+  for (const [bridgeKey, active] of activeBridges) {
+    cleanupBridge(active.bridge, active.heartbeatTimer, bridgeKey);
+  }
+  destroyAllIdleBridges();
+  conversationStates.clear();
+}
+
+export function evictStaleConversations(now = Date.now()): void {
+  for (const [key, stored] of conversationStates) {
+    if (!stored.sessionScoped && now - stored.lastAccessMs > CONVERSATION_TTL_MS) {
+      debugLog("conversation.evict", { key, stored, now });
+      conversationStates.delete(key);
+      deleteConversationJournal(key);
+    }
+  }
+  evictStaleJournals(now);
+}
+
+export function fingerprintCompletedTurns(turns: ParsedTurn[]): string {
+  return fingerprintCompletedTurnsImpl(turns);
+}
+
+export function clearStoredMidPauseMetadata(stored: StoredConversation): void {
+  clearStoredMidPauseMetadataImpl(stored);
+}
+
+export function clearStoredCheckpoint(stored: StoredConversation, clearBlobStore = false): void {
+  stored.checkpoint = null;
+  delete stored.checkpointSource;
+  delete stored.checkpointTurnCount;
+  delete stored.checkpointHistoryFingerprint;
+  clearStoredMidPauseMetadata(stored);
+  if (clearBlobStore) stored.blobStore.clear();
+}
+
+/**
+ * A live KV blob miss means Cursor asked for history we no longer hold.
+ * Drop the checkpoint and rotate the conversation id so the next turn rebuilds
+ * from Pi's transcript instead of replaying holes.
+ */
+export function markBlobMiss(convKey: string): void {
+  const stored = conversationStates.get(convKey);
+  if (!stored) return;
+  debugLog("conversation.blob_miss_invalidate", {
+    convKey,
+    hadCheckpoint: !!stored.checkpoint,
+  });
+  clearStoredCheckpoint(stored, false);
+  stored.conversationId = randomUUID();
+  persistJournal(convKey, stored);
+}
+
+/**
+ * Checkpoint bytes are replayed straight into `fromBinary` when a request is built. A truncated or
+ * otherwise undecodable checkpoint would throw there and fail the turn — and, because nothing
+ * clears it, every later turn in the conversation too. Decoding once here turns that permanent
+ * break into a one-time discard and a rebuild.
+ */
+function checkpointDecodes(checkpoint: Uint8Array): boolean {
+  try {
+    fromBinary(ConversationStateStructureSchema, checkpoint);
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+export function discardStaleCheckpointIfNeeded(
+  stored: StoredConversation,
+  turns: ParsedTurn[],
+  requestId: string,
+  convKey: string,
+): void {
+  // Tier 2 extends staleness validation to metadata-only mid-pause snapshots.
+  if (!stored.checkpoint) return;
+
+  const currentTurnCount = turns.length;
+  const currentHistoryFingerprint = fingerprintCompletedTurns(turns);
+  const storedCheckpointTurnCount = stored.checkpointTurnCount;
+  const storedCheckpointHistoryFingerprint = stored.checkpointHistoryFingerprint;
+
+  // A checkpoint recorded at the end of a turn stores turnCount = completedTurns.length + 1
+  // (it includes the just-finished turn in its history). When a tool-result recovery request
+  // arrives for that same turn, turns.length is still the pre-tool count. Allow the off-by-one
+  // when mid-pause metadata confirms we are in a tool continuation so the checkpoint survives
+  // long enough for planRecovery to use it.
+  const hasMidPauseForThisTurn =
+    stored.midPauseTurnCount === currentTurnCount &&
+    !!stored.midPausePendingToolCalls?.length &&
+    !!stored.midPauseHistoryFingerprint &&
+    stored.midPauseHistoryFingerprint === currentHistoryFingerprint;
+  const checkpointIsForNextTurnCount = storedCheckpointTurnCount === currentTurnCount + 1;
+  const skipTurnCountCheck = hasMidPauseForThisTurn && checkpointIsForNextTurnCount;
+
+  const reason =
+    stored.checkpoint.byteLength > MAX_CHECKPOINT_BYTES
+      ? "checkpoint_oversized"
+      : !checkpointDecodes(stored.checkpoint)
+        ? "checkpoint_undecodable"
+        : storedCheckpointTurnCount === undefined || !storedCheckpointHistoryFingerprint
+          ? "missing_checkpoint_metadata"
+          : !skipTurnCountCheck && storedCheckpointTurnCount !== currentTurnCount
+            ? "completed_turn_count_mismatch"
+            : !skipTurnCountCheck &&
+                storedCheckpointHistoryFingerprint !== currentHistoryFingerprint
+              ? "completed_history_fingerprint_mismatch"
+              : undefined;
+
+  if (!reason) return;
+
+  debugLog("chat.discard_checkpoint", {
+    requestId,
+    convKey,
+    reason,
+    checkpointBytes: stored.checkpoint.byteLength,
+    storedCheckpointTurnCount,
+    currentTurnCount,
+    hasMidPauseForThisTurn,
+    storedCheckpointHistoryFingerprint,
+    currentHistoryFingerprint,
+  });
+  clearStoredCheckpoint(stored, true);
+
+  // Compaction and other history rewrites make the Cursor conversation identity
+  // stale. Start a fresh conversation so we rebuild from Pi's (summarized)
+  // transcript instead of replaying a checkpoint with holes.
+  const historyRewritten =
+    reason === "completed_turn_count_mismatch" ||
+    reason === "completed_history_fingerprint_mismatch";
+  if (historyRewritten) {
+    stored.conversationId = randomUUID();
+    debugLog("conversation.rotated", {
+      requestId,
+      convKey,
+      reason,
+      conversationId: stored.conversationId,
+    });
+  }
+  persistJournal(convKey, stored);
+}
+
+export function trimBlobStore(
+  store: Map<string, Uint8Array>,
+  maxBytes = MAX_CONVERSATION_BLOB_BYTES,
+  maxEntries = Number.POSITIVE_INFINITY,
+): { removed: number; totalBytes: number } {
+  let totalBytes = 0;
+  for (const value of store.values()) totalBytes += value.byteLength;
+  const withinBudget = () => totalBytes <= maxBytes && store.size <= maxEntries;
+  if (withinBudget()) return { removed: 0, totalBytes };
+
+  let removed = 0;
+  // Map iteration order is insertion order — drop oldest blobs first.
+  for (const key of store.keys()) {
+    if (withinBudget()) break;
+    const value = store.get(key);
+    if (!value) continue;
+    totalBytes -= value.byteLength;
+    store.delete(key);
+    removed += 1;
+  }
+  return { removed, totalBytes };
+}
+
+export function mergeBlobStore(
+  stored: StoredConversation,
+  blobStore: Map<string, Uint8Array>,
+): void {
+  // Delete-then-set moves each still-referenced blob to the end of the insertion order, which is
+  // what `trimBlobStore` treats as newest. Plain `set` leaves an existing key in place, so the
+  // system-prompt blob — written first on every build, and pointed at by `rootPromptMessagesJson`
+  // in every checkpoint — stayed permanently oldest and would be the first thing evicted.
+  for (const [k, v] of blobStore) {
+    stored.blobStore.delete(k);
+    stored.blobStore.set(k, v);
+  }
+  const trimmed = trimBlobStore(
+    stored.blobStore,
+    MAX_CONVERSATION_BLOB_BYTES,
+    MAX_ACTIVE_BLOB_ENTRIES,
+  );
+  if (trimmed.removed > 0) {
+    debugLog("conversation.blob_store_trimmed", {
+      removed: trimmed.removed,
+      totalBytes: trimmed.totalBytes,
+      entries: stored.blobStore.size,
+      maxBytes: MAX_CONVERSATION_BLOB_BYTES,
+      maxEntries: MAX_ACTIVE_BLOB_ENTRIES,
+    });
+  }
+  stored.lastAccessMs = Date.now();
+}
+
+export function commitStoredCheckpoint(
+  stored: StoredConversation,
+  checkpointBytes: Uint8Array,
+  blobStore: Map<string, Uint8Array>,
+  completedTurns: ParsedTurn[],
+  currentTurn: ParsedTurn,
+  convKey?: string,
+): void {
+  const completedHistory = [...completedTurns, currentTurn];
+  mergeBlobStore(stored, blobStore);
+  stored.checkpoint = checkpointBytes;
+  stored.checkpointSource = "upstream";
+  stored.checkpointTurnCount = completedHistory.length;
+  stored.checkpointHistoryFingerprint = fingerprintCompletedTurns(completedHistory);
+  clearStoredMidPauseMetadata(stored);
+  stored.lastAccessMs = Date.now();
+  if (convKey) persistJournal(convKey, stored);
+}
+
+export function persistAbortedConversationState(
+  convKey: string,
+  latestCheckpoint: Uint8Array | null,
+  blobStore: Map<string, Uint8Array>,
+  completedTurns: ParsedTurn[],
+  currentTurn: ParsedTurn,
+  pendingToolCalls: Array<{ toolCallId: string; toolName: string }> = [],
+): void {
+  const stored = conversationStates.get(convKey);
+  if (!stored) return;
+
+  // An interrupted tool pause must retain its matching snapshot. Treating its
+  // checkpoint as a completed turn clears that metadata, making the next tool
+  // result reject a valid off-by-one checkpoint as stale.
+  if (pendingToolCalls.length > 0) {
+    commitStoredCheckpointMidPause(
+      stored,
+      latestCheckpoint,
+      blobStore,
+      completedTurns,
+      pendingToolCalls,
+      convKey,
+    );
+    debugLog("native.stream.abort_state_saved", {
+      convKey,
+      hasCheckpoint: !!latestCheckpoint,
+      completedTurnCount: completedTurns.length,
+      pendingToolCallIds: pendingToolCalls.map((call) => call.toolCallId),
+      currentTurn,
+    });
+    return;
+  }
+
+  // The stream did not finish this turn. `commitStoredCheckpoint` would count
+  // `currentTurn` as completed (turnCount+1) and drop mid-pause metadata, which
+  // makes the next tool-continuation / idle-retry treat a still-valid checkpoint
+  // as stale. Keep the checkpoint keyed to the completed history only.
+  if (latestCheckpoint) {
+    mergeBlobStore(stored, blobStore);
+    stored.checkpoint = latestCheckpoint;
+    stored.checkpointSource = "upstream";
+    stored.checkpointTurnCount = completedTurns.length;
+    stored.checkpointHistoryFingerprint = fingerprintCompletedTurns(completedTurns);
+    stored.lastAccessMs = Date.now();
+    persistJournal(convKey, stored);
+  } else {
+    // Blob ids referenced by the retained Pi history must outlive the cancelled
+    // bridge even when Cursor has not emitted a checkpoint yet.
+    mergeBlobStore(stored, blobStore);
+    stored.lastAccessMs = Date.now();
+    persistJournal(convKey, stored);
+  }
+
+  debugLog("native.stream.abort_state_saved", {
+    convKey,
+    hasCheckpoint: !!latestCheckpoint,
+    completedTurnCount: completedTurns.length,
+    currentTurn,
+  });
+}
+
+export function commitStoredCheckpointMidPause(
+  stored: StoredConversation,
+  checkpointBytes: Uint8Array | null,
+  blobStore: Map<string, Uint8Array>,
+  completedTurns: ParsedTurn[],
+  pendingToolCalls: Array<{ toolCallId: string; toolName: string }>,
+  convKey?: string,
+): void {
+  mergeBlobStore(stored, blobStore);
+  const completedHistoryFingerprint = fingerprintCompletedTurns(completedTurns);
+  if (checkpointBytes) {
+    stored.checkpoint = checkpointBytes;
+    stored.checkpointSource = "upstream";
+    stored.checkpointTurnCount = completedTurns.length;
+    stored.checkpointHistoryFingerprint = completedHistoryFingerprint;
+  } else {
+    // Metadata-only snapshots intentionally discard any older upstream checkpoint so later
+    // recovery code cannot accidentally treat stale bytes as authoritative for this pause.
+    stored.checkpoint = null;
+    stored.checkpointSource = "absent";
+    delete stored.checkpointTurnCount;
+    delete stored.checkpointHistoryFingerprint;
+  }
+  stored.midPausePendingToolCalls = pendingToolCalls.map((c) => ({
+    toolCallId: c.toolCallId,
+    toolName: c.toolName,
+  }));
+  stored.midPauseTurnCount = completedTurns.length;
+  stored.midPauseHistoryFingerprint = completedHistoryFingerprint;
+  stored.midPauseRecordedAtMs = Date.now();
+  stored.lastAccessMs = Date.now();
+  if (convKey) persistJournal(convKey, stored);
+}
+
+export interface HandleBridgeCloseMidPauseInput {
+  stored: StoredConversation | undefined;
+  latestCheckpoint: Uint8Array | null;
+  blobStore: Map<string, Uint8Array>;
+  completedTurns: ParsedTurn[];
+  pendingExecs: Array<{ toolCallId: string; toolName: string }>;
+  convKey?: string;
+}
+
+export function handleBridgeCloseMidPause(input: HandleBridgeCloseMidPauseInput): {
+  committed: boolean;
+} {
+  if (!input.stored) return { committed: false };
+  commitStoredCheckpointMidPause(
+    input.stored,
+    input.latestCheckpoint,
+    input.blobStore,
+    input.completedTurns,
+    input.pendingExecs,
+    input.convKey,
+  );
+  return { committed: true };
+}
+
+export function deriveRequestLockKey(body: ChatCompletionRequest): string {
+  const sessionId = derivePiSessionId(body);
+  if (sessionId) return `session:${sessionId}`;
+  return `anonymous:${deriveConversationKey(body.messages)}`;
+}
+
+export async function withSessionLock<T>(key: string, fn: () => Promise<T>): Promise<T> {
+  const previous = sessionLocks.get(key) ?? Promise.resolve();
+  let release!: () => void;
+  const current = new Promise<void>((resolve) => {
+    release = resolve;
+  });
+  const chained = previous.catch(() => {}).then(() => current);
+  sessionLocks.set(key, chained);
+  await previous.catch(() => {});
+  try {
+    return await fn();
+  } finally {
+    release();
+    if (sessionLocks.get(key) === chained) sessionLocks.delete(key);
+  }
+}
+
+export function derivePiSessionId(
+  body: Pick<ChatCompletionRequest, "pi_session_id" | "user">,
+): string | undefined {
+  const raw = body.pi_session_id ?? body.user;
+  if (typeof raw !== "string") return undefined;
+  const trimmed = raw.trim();
+  return trimmed ? trimmed : undefined;
+}
+
+export function deriveBridgeKeyFromSessionId(sessionId: string): string {
+  return createHash("sha256").update(`bridge:${sessionId}`).digest("hex").slice(0, 16);
+}
+
+export function deriveConversationKeyFromSessionId(sessionId: string): string {
+  return createHash("sha256").update(`conv:${sessionId}`).digest("hex").slice(0, 16);
+}
+
+export function deriveBridgeKey(messages: OpenAIMessage[], sessionId?: string): string {
+  if (sessionId) return deriveBridgeKeyFromSessionId(sessionId);
+  const firstSystemMsg = messages.find((m) => m.role === "system");
+  const firstUserMsg = messages.find((m) => m.role === "user");
+  const firstSystemText = firstSystemMsg ? textContent(firstSystemMsg.content) : "";
+  const firstUserText = firstUserMsg ? textContent(firstUserMsg.content) : "";
+  return createHash("sha256")
+    .update(`bridge:${firstSystemText}\0${firstUserText}`)
+    .digest("hex")
+    .slice(0, 16);
+}
+
+export function deriveConversationKey(messages: OpenAIMessage[], sessionId?: string): string {
+  if (sessionId) return deriveConversationKeyFromSessionId(sessionId);
+  const firstSystemMsg = messages.find((m) => m.role === "system");
+  const firstUserMsg = messages.find((m) => m.role === "user");
+  const firstSystemText = firstSystemMsg ? textContent(firstSystemMsg.content) : "";
+  const firstUserText = firstUserMsg ? textContent(firstUserMsg.content) : "";
+  return createHash("sha256")
+    .update(`conv:${firstSystemText}\0${firstUserText}`)
+    .digest("hex")
+    .slice(0, 16);
+}
+
+export function cleanupSessionState(sessionId?: string): void {
+  if (!sessionId) return;
+  const bridgeKey = deriveBridgeKeyFromSessionId(sessionId);
+  const convKey = deriveConversationKeyFromSessionId(sessionId);
+  const active = activeBridges.get(bridgeKey);
+  debugLog("session.cleanup", {
+    sessionId,
+    bridgeKey,
+    convKey,
+    hasActiveBridge: !!active,
+    hadConversation: conversationStates.has(convKey),
+  });
+  if (active) cleanupBridge(active.bridge, active.heartbeatTimer, bridgeKey);
+  destroyIdleBridge(bridgeKey);
+  // Drop the in-memory copy so a long-lived process cannot pin a large blob store
+  // after the user switched away. The journal stays so /resume can hydrate.
+  conversationStates.delete(convKey);
+}
+
+export function deterministicConversationId(convKey: string): string {
+  const hex = createHash("sha256").update(`cursor-conv-id:${convKey}`).digest("hex").slice(0, 32);
+  return [
+    hex.slice(0, 8),
+    hex.slice(8, 12),
+    `4${hex.slice(13, 16)}`,
+    `${(0x8 | (parseInt(hex[16], 16) & 0x3)).toString(16)}${hex.slice(17, 20)}`,
+    hex.slice(20, 32),
+  ].join("-");
+}

package/src/stream/stream-writer.ts

@@ -0,0 +1,211 @@
+/**
+ * Stream writer adapter converting internal events to Pi AssistantMessageEventStream.
+ */
+
+import type { Api, AssistantMessageEventStream, Context, Model } from "@earendil-works/pi-ai";
+import type {
+  CursorRunUsage,
+  NativeBlockKind,
+  NativeStreamWriter,
+  PendingExec,
+  StreamState,
+} from "./types.js";
+import { applyCursorUsage, createCursorAssistantMessage } from "./pi-adapter.js";
+import { parseToolCallArguments } from "./message-parsing.js";
+import { createCursorContextTracker, type CursorAssistantMessage } from "./context-usage.js";
+import { lifecycleLog, reportCursorAnomaly } from "./debug-log.js";
+
+export function createNativeStreamWriter(
+  stream: AssistantMessageEventStream,
+  model: Model<Api>,
+  context?: Context,
+  options?: { sessionId?: string; reasoning?: string },
+): NativeStreamWriter {
+  const output: CursorAssistantMessage = createCursorAssistantMessage(model);
+  const contextTracker = createCursorContextTracker(model, context, options);
+  const carriedReceipts = new Set<CursorRunUsage>();
+  let started = false;
+  let closed = false;
+  let active: { kind: NativeBlockKind; contentIndex: number; ended: boolean } | undefined;
+
+  const ensureStarted = (): void => {
+    if (started) return;
+    started = true;
+    stream.push({ type: "start", partial: output });
+  };
+
+  const endActiveBlock = (): void => {
+    if (!active || active.ended) return;
+    const block = output.content[active.contentIndex];
+    if (active.kind === "text" && block?.type === "text") {
+      stream.push({
+        type: "text_end",
+        contentIndex: active.contentIndex,
+        content: block.text,
+        partial: output,
+      });
+    } else if (active.kind === "thinking" && block?.type === "thinking") {
+      stream.push({
+        type: "thinking_end",
+        contentIndex: active.contentIndex,
+        content: block.thinking,
+        partial: output,
+      });
+    }
+    active.ended = true;
+    active = undefined;
+  };
+
+  const ensureBlock = (kind: NativeBlockKind): number => {
+    ensureStarted();
+    if (active?.kind === kind && !active.ended) return active.contentIndex;
+    endActiveBlock();
+    const contentIndex = output.content.length;
+    if (kind === "text") {
+      output.content.push({ type: "text", text: "" });
+      stream.push({ type: "text_start", contentIndex, partial: output });
+    } else {
+      output.content.push({ type: "thinking", thinking: "" });
+      stream.push({ type: "thinking_start", contentIndex, partial: output });
+    }
+    active = { kind, contentIndex, ended: false };
+    return contentIndex;
+  };
+
+  const finishUsage = (reason: string, state?: StreamState): void => {
+    const snapshot = contextTracker.finish(output);
+    const billing = applyCursorUsage(output, model, state, snapshot.tokens);
+    for (const receipt of carriedReceipts) {
+      const extra = createCursorAssistantMessage(model);
+      const info = applyCursorUsage(
+        extra,
+        model,
+        {
+          toolCallIndex: 0,
+          pendingExecs: [],
+          outputTokens: 0,
+          totalTokens: 0,
+          turnEnded: true,
+          runUsage: receipt,
+        },
+        snapshot.tokens,
+      );
+      if (info.status !== "reported" && info.status !== "partial") continue;
+      for (const key of ["input", "output", "cacheRead", "cacheWrite"] as const) {
+        output.usage[key] += extra.usage[key];
+        output.usage.cost[key] += extra.usage.cost[key];
+      }
+      output.usage.cost.total += extra.usage.cost.total;
+      billing.carriedReceipts = (billing.carriedReceipts ?? 0) + 1;
+      if (info.status === "partial") billing.status = "partial";
+    }
+    if (billing.status === "pending" && reason !== "toolUse") billing.status = "unavailable";
+    output.cursorUsage!.billing = billing;
+    if (billing.status === "partial" || billing.status === "unavailable") {
+      lifecycleLog("usage_incomplete", { modelId: model.id, reason, ...billing });
+    }
+    if (billing.status === "partial") {
+      reportCursorAnomaly(
+        "usage_incomplete",
+        "Cursor returned an incomplete billing split; displayed costs include only known buckets.",
+        { modelId: model.id, ...billing },
+        { level: "warning" },
+      );
+    }
+    if (snapshot.source === "estimate" && reason !== "toolUse") {
+      lifecycleLog("usage_context_estimated", {
+        modelId: model.id,
+        reason,
+        contextTokens: snapshot.tokens,
+      });
+    }
+  };
+
+  return {
+    output,
+    contextSnapshot(tokens: number, checkpoint?: Uint8Array) {
+      if (!closed) contextTracker.observe(tokens, output, checkpoint);
+    },
+    contextMode(mode, tokens, checkpoint) {
+      if (!closed) contextTracker.begin(mode, tokens, checkpoint);
+    },
+    carryUsage(usage) {
+      if (!closed && usage.modelId === model.id) carriedReceipts.add(usage);
+    },
+    get closed() {
+      return closed;
+    },
+    start: ensureStarted,
+    text(delta: string) {
+      if (closed || !delta) return;
+      const contentIndex = ensureBlock("text");
+      const block = output.content[contentIndex];
+      if (block?.type !== "text") return;
+      block.text += delta;
+      stream.push({ type: "text_delta", contentIndex, delta, partial: output });
+    },
+    thinking(delta: string) {
+      if (closed || !delta) return;
+      const contentIndex = ensureBlock("thinking");
+      const block = output.content[contentIndex];
+      if (block?.type !== "thinking") return;
+      block.thinking += delta;
+      stream.push({ type: "thinking_delta", contentIndex, delta, partial: output });
+    },
+    toolCall(exec: PendingExec) {
+      if (closed) return;
+      ensureStarted();
+      endActiveBlock();
+      const contentIndex = output.content.length;
+      const parsedArguments = parseToolCallArguments(exec.decodedArgs);
+      const block = {
+        type: "toolCall" as const,
+        id: exec.toolCallId,
+        name: exec.toolName,
+        arguments: {},
+      };
+      output.content.push(block);
+      stream.push({ type: "toolcall_start", contentIndex, partial: output });
+      block.arguments = parsedArguments;
+      stream.push({
+        type: "toolcall_delta",
+        contentIndex,
+        delta: exec.decodedArgs,
+        partial: output,
+      });
+      stream.push({
+        type: "toolcall_end",
+        contentIndex,
+        toolCall: {
+          type: "toolCall",
+          id: exec.toolCallId,
+          name: exec.toolName,
+          arguments: parsedArguments,
+        },
+        partial: output,
+      });
+    },
+    done(reason: "stop" | "length" | "toolUse", state?: StreamState) {
+      if (closed) return;
+      ensureStarted();
+      endActiveBlock();
+      // Intermediate tool replies publish context, not an invented per-inference bill.
+      finishUsage(reason, state);
+      output.stopReason = reason;
+      stream.push({ type: "done", reason, message: output });
+      closed = true;
+      stream.end(output);
+    },
+    error(message: string, reason: "error" | "aborted", state?: StreamState) {
+      if (closed) return;
+      ensureStarted();
+      endActiveBlock();
+      finishUsage(reason, state);
+      output.stopReason = reason;
+      output.errorMessage = message;
+      stream.push({ type: "error", reason, error: output });
+      closed = true;
+      stream.end(output);
+    },
+  };
+}