@viniciosrab/pi-claude-bridge 0.9.0

package/src/prompt-stream.ts

@@ -0,0 +1,102 @@
+// Long-lived streaming-input prompt for query().
+//
+// The SDK accepts `prompt: AsyncIterable<SDKUserMessage>` and pumps it to the
+// CLI's stdin. Keeping that iterable parked for the life of the query lets us
+// write a steer to stdin *while* a tool is running, which is what makes CC
+// drain it at the next tool boundary (`priority: "next"`) instead of treating
+// it as a follow-up turn.
+//
+// The ack is the load-bearing part. `push()` resolves on the line *after*
+// `yield`, and the SDK's pump is `for await (m of stream) { await
+// transport.write(m) }` — so resuming past the yield proves the write to stdin
+// completed. Callers await that before releasing the MCP tool result, which
+// travels back over the same stdin FIFO; winning that race is what guarantees
+// CC sees the steer before the tool result.
+//
+// The drain and the FIFO ordering are CC CLI internals, not SDK contract, so
+// this can break under a CC upgrade without any type error.
+
+import type { SDKUserMessage } from "@anthropic-ai/claude-agent-sdk";
+
+export interface PromptStream {
+	stream: AsyncGenerator<SDKUserMessage>;
+	/** Enqueue a message; resolves once the SDK has written it to stdin.
+	 *  Rejects (never hangs) if the stream is already ended or failed. */
+	push: (msg: SDKUserMessage) => Promise<void>;
+	/** Close the input: the generator returns, the SDK closes the CLI's stdin. */
+	end: () => void;
+	/** Abandon the input, rejecting every queued and in-flight ack. */
+	fail: (error: Error) => void;
+}
+
+export function makePromptStream(): PromptStream {
+	type Item = { msg: SDKUserMessage; resolve: () => void; reject: (e: Error) => void };
+	const queue: Item[] = [];
+	// The item currently parked at `yield`. Tracked separately so fail() can
+	// settle it — a dying CLI may abandon the pump without ever resuming us,
+	// and an unsettled ack would wedge tool-result delivery forever.
+	let inflight: Item | null = null;
+	let wake: (() => void) | null = null;
+	let done = false;
+	let failure: Error | null = null;
+
+	const kick = () => { wake?.(); wake = null; };
+
+	async function* gen(): AsyncGenerator<SDKUserMessage> {
+		try {
+			while (true) {
+				while (queue.length === 0 && !done && !failure) {
+					await new Promise<void>((resolve) => { wake = resolve; });
+				}
+				if (failure) throw failure;
+				const item = queue.shift();
+				if (!item) return; // ended and drained
+				inflight = item;
+				try {
+					yield item.msg;
+					item.resolve();
+				} finally {
+					// Reached either normally (no-op, already resolved) or when the
+					// pump abandons iteration — a `for await` break/throw calls
+					// gen.return(), which resumes the yield as a return.
+					item.reject(new Error("prompt stream closed"));
+					inflight = null;
+				}
+			}
+		} finally {
+			// No consumer left to drain the queue, so nothing would ever settle a
+			// later push. Closing here keeps the reject-never-hang contract a
+			// property of this module rather than of every call site.
+			done = true;
+		}
+	}
+
+	return {
+		stream: gen(),
+		push: (msg) => failure || done
+			? Promise.reject(failure ?? new Error("prompt stream closed"))
+			: new Promise<void>((resolve, reject) => { queue.push({ msg, resolve, reject }); kick(); }),
+		end: () => { done = true; kick(); },
+		fail: (error) => {
+			// First failure wins: the query's `finally` fails the stream a second
+			// time with a generic "query ended", which would otherwise mask the
+			// real cause its `catch` recorded.
+			if (failure) return;
+			failure = error;
+			queue.splice(0).forEach((item) => item.reject(error));
+			inflight?.reject(error);
+			kick();
+		},
+	};
+}
+
+/** `uuid` is deliberately omitted: we need no dedup, and supplying one makes
+ *  CC's stdin loop do a session lookup on the message. */
+export function userMessage(content: SDKUserMessage["message"]["content"], priority?: SDKUserMessage["priority"]): SDKUserMessage {
+	return {
+		type: "user",
+		message: { role: "user", content } as SDKUserMessage["message"],
+		parent_tool_use_id: null,
+		...(priority ? { priority } : {}),
+	};
+}

package/src/query-state.ts

@@ -0,0 +1,112 @@
+// Query state: QueryContext class.
+//
+// All per-query and per-turn mutable state lives here. Reentrant queries
+// (subagents) each get their own QueryContext instance, managed by index.ts.
+// Adding a new field = one property on the class.
+//
+// Extracted from index.ts so tests can import without activating the extension.
+
+import type { AssistantMessage, AssistantMessageEventStream, Model } from "@earendil-works/pi-ai";
+import type { McpResult } from "./extract-tool-results.js";
+import type { PromptStream } from "./prompt-stream.js";
+
+export interface PendingToolCall {
+	toolName: string;
+	resolve: (result: McpResult) => void;
+}
+
+export class QueryContext {
+	// Query-scoped (fully isolated per query)
+	activeQuery: unknown | null = null;
+	currentPiStream: AssistantMessageEventStream | null = null;
+	latestCursor = 0;
+	pendingToolCalls = new Map<string, PendingToolCall>();
+	pendingResults = new Map<string, McpResult>();
+	/** tool_use ids emitted this turn. Sole purpose is routing a delivered result
+	 *  to the owning query when several queries are in flight — pairing a result
+	 *  to its call is done by id from Claude's tools/call _meta, not from here. */
+	turnToolCallIds: string[] = [];
+	/** Streaming-input handle for the active query — how steers reach CC mid-turn. */
+	promptStream: PromptStream | null = null;
+	/** Last rate-limit rejection seen on this query. Claude Code sends it just before the
+	 *  failure it caused, which is the only thing tying the two together. */
+	rateLimitRejection: { rateLimitType?: string; resetsAt?: number } | null = null;
+	/** Highest 5% utilization bucket we notified for, so repeat rate_limit_event spam is suppressed. */
+	lastRateLimitWarnStep: number | null = null;
+	lastRateLimitWarnThreshold: number | undefined;
+	/** pi session this query serves, from SimpleStreamOptions.sessionId at fresh-query
+	 *  setup. A bridge process serves several pi sessions at once (subagents run their
+	 *  own AgentSessions), and history rewrites must only discard the rewriting
+	 *  session's parked queries — this is the match key. Null when the host did not
+	 *  supply an id.
+	 */
+	piSessionId: string | null = null;
+	/** pi rewrote the history this query was built from (session_compact,
+	 *  session_tree in its own pi session). Set by markRebuildForSession, consumed
+	 *  by the tool-result delivery that discards the query. Not session-wide state:
+	 *  it dies with the context it belongs to, so it cannot leak into later turns
+	 *  the way a module flag does.
+	 */
+	historyStale = false;
+	/** A steer never reached CC. A first query has no session mirror yet, so
+	 *  completion must carry this into the mirror it creates. */
+	missedSteer = false;
+
+	// Per-turn (reset together)
+	turnOutput: AssistantMessage | null = null;
+	turnStarted = false;
+	turnSawStreamEvent = false;
+	turnSawToolCall = false;
+	/** API message id from the last message_start, and whether its message_stop has
+	 *  arrived. An `assistant` message under a different id while the stream is still
+	 *  open is Claude Code's non-streaming fallback for a stalled stream. */
+	turnStreamMessageId: string | undefined;
+	turnStreamOpen = false;
+	/** turnBlocks length at that message_start: where an abandoned attempt's blocks begin. */
+	turnStreamBlockStart = 0;
+
+	get turnBlocks(): Array<any> {
+		if (!this.turnOutput) throw new Error("turnBlocks accessed before resetTurnState");
+		return this.turnOutput.content;
+	}
+
+	/** Answer every parked MCP handler with `reason` and forget the turn's queued
+	 *  results. Called when the query it belongs to is going away (abort, error,
+	 *  normal end). Handlers must be *resolved*, not rejected: an error reply is
+	 *  still a reply, and a handler left awaiting a subprocess that is gone keeps
+	 *  CC's tools/call open forever, which wedges pi's turn behind it. */
+	releasePendingToolCalls(reason: string): void {
+		for (const pending of this.pendingToolCalls.values()) pending.resolve({ content: [{ type: "text", text: reason }] });
+		this.pendingToolCalls.clear();
+		this.pendingResults.clear();
+	}
+
+	resetTurnState(model: Model<any>): void {
+		this.turnOutput = {
+			role: "assistant", content: [],
+			api: model.api, provider: model.provider, model: model.id,
+			usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0,
+				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } },
+			stopReason: "stop", timestamp: Date.now(),
+		};
+		this.turnStarted = false;
+		this.turnSawStreamEvent = false;
+		this.turnSawToolCall = false;
+		this.turnStreamMessageId = undefined;
+		this.turnStreamOpen = false;
+		this.turnStreamBlockStart = 0;
+		// turnToolCallIds is NOT reset — it persists across tool-result delivery
+		// callbacks within the same assistant message so results can be routed to
+		// this query while its handlers are still pending.
+	}
+}
+
+let _ctx = new QueryContext();
+
+export function ctx(): QueryContext { return _ctx; }
+
+// Test-only: replace the module-level context so test files start clean.
+// Not called from production.
+export function resetCtx(): void {
+	_ctx = new QueryContext();
+}

package/src/session-verify.ts

@@ -0,0 +1,38 @@
+// Pure session-file integrity check. Returns an array of warning strings;
+// callers decide how to surface them (debug log, piUI, diagDump, etc.).
+// Extracted from index.ts so tests can import without activating the extension.
+
+import { statSync, readFileSync } from "fs";
+
+export function verifyWrittenSession(jsonlPath: string, expectedSessionId: string, expectedRecordCount: number): string[] {
+	const warnings = [];
+	let st;
+	try {
+		st = statSync(jsonlPath);
+	} catch (e) {
+		warnings.push(`file missing after save — path=${jsonlPath} err=${e.message}`);
+		return warnings;
+	}
+	let content;
+	try {
+		content = readFileSync(jsonlPath, "utf8");
+	} catch (e) {
+		warnings.push(`file unreadable — path=${jsonlPath} size=${st.size} err=${e.message}`);
+		return warnings;
+	}
+	const lines = content.split("\n").filter((l) => l.trim().length > 0);
+	if (lines.length !== expectedRecordCount) {
+		warnings.push(`record count mismatch — expected=${expectedRecordCount} actual=${lines.length} path=${jsonlPath} bytes=${content.length}`);
+		return warnings;
+	}
+	try {
+		const firstRec = JSON.parse(lines[0]);
+		const lastRec = JSON.parse(lines[lines.length - 1]);
+		if (firstRec.sessionId !== expectedSessionId || lastRec.sessionId !== expectedSessionId) {
+			warnings.push(`sessionId drift — expected=${expectedSessionId} first=${firstRec.sessionId} last=${lastRec.sessionId}`);
+		}
+	} catch (e) {
+		warnings.push(`malformed JSONL — path=${jsonlPath} err=${e.message}`);
+	}
+	return warnings;
+}

package/src/skills.ts

@@ -0,0 +1,20 @@
+import { formatSkillsForPrompt, type Skill } from "@earendil-works/pi-coding-agent";
+
+export const MCP_SERVER_NAME = "custom-tools";
+export const MCP_TOOL_PREFIX = `mcp__${MCP_SERVER_NAME}__`;
+
+export type SkillReadTool = "mcp" | "native" | "none";
+
+export function renderSkillsBlock(skills: Skill[], readTool: SkillReadTool): string | undefined {
+	if (readTool === "none" || skills.length === 0) return undefined;
+	const block = formatSkillsForPrompt(skills).trim();
+	if (!block) return undefined;
+	return readTool === "mcp" ? rewriteSkillsBlock(block) : block;
+}
+
+export function rewriteSkillsBlock(skillsBlock: string): string {
+	return skillsBlock.replace(
+		"Use the read tool to load a skill's file",
+		`Use the read tool (mcp__${MCP_SERVER_NAME}__read) to load a skill's file`,
+	);
+}

package/src/transcript.ts

@@ -0,0 +1,74 @@
+/**
+ * Translate pi's transcript-shaped provider input into the prompt/tools fields used by
+ * the bridge's downstream consumers.
+ *
+ * System messages carry the base prompt and tool set plus later section patches and tool
+ * deltas. pi-ai replays that state, but preserves section replay order. Prompt capture
+ * keys come from pi's canonical section builder, so a deleted and re-added section must
+ * be ranked back into canonical order before the bridge performs its exact-key lookup.
+ */
+import {
+	contentText,
+	getCurrentSystemMessage,
+	getCurrentTools,
+	type Context,
+	type SystemMessage,
+} from "@earendil-works/pi-ai";
+
+/** pi's canonical built-in section order; custom sections follow their replay position.
+ *  Known limitation, inherited from the pre-swap replay: if pi reorders existing custom
+ *  sections while patching an unrelated one, the replay diverges from getSystemPrompt()
+ *  and the exact-key capture lookup throws on that legitimate turn. Untriggered today:
+ *  pi builds custom sections in one patch per render. */
+const SECTION_RANK = new Map<string, number>([
+	["preamble", 0], ["tools", 1], ["rules", 2], ["docs", 3], ["addendum", 4],
+	["project_context", 5], ["skills", 6], ["cwd", 7],
+]);
+
+/**
+ * The map's entries, stably sorted by canonical rank. A name absent from the rank table
+ * inherits its replayed predecessor's rank, so an already-canonical replay remains
+ * unchanged and custom sections retain their position.
+ */
+function stableRanked(sections: Map<string, string>, ranks: Map<string, number>): Map<string, string> {
+	let predecessorRank = -1;
+	const ranked = [...sections].map(([name, value]) => {
+		const rank = ranks.get(name) ?? predecessorRank;
+		predecessorRank = rank;
+		return { name, value, rank };
+	});
+	ranked.sort((a, b) => a.rank - b.rank);
+	return new Map(ranked.map(({ name, value }) => [name, value]));
+}
+
+/** Render replayed system state in the same section order as pi's prompt builder. */
+function canonicalSystemPrompt(message: SystemMessage | undefined): string | undefined {
+	if (!message) return undefined;
+	const sections = new Map<string, string>(
+		Object.entries(message.sections ?? {}).filter((entry): entry is [string, string] => entry[1] !== null),
+	);
+	const parts = [contentText(message.content), ...stableRanked(sections, SECTION_RANK).values()]
+		.filter((part) => part.length > 0);
+	return parts.length > 0 ? parts.join("\n\n") : undefined;
+}
+
+/**
+ * Restore the prompt and tools fields expected by the bridge and remove prompt-state
+ * messages from conversation history. Contexts without system messages are returned
+ * unchanged because systemless one-off calls already use the bridge-compatible shape.
+ */
+export function toBridgeContext(context: Context): Context {
+	if (!context.messages.some((message) => message.role === "system")) return context;
+	const tools = getCurrentTools(context.messages);
+	return {
+		...context,
+		systemPrompt: canonicalSystemPrompt(getCurrentSystemMessage(context.messages)),
+		tools: tools.length > 0 ? tools : undefined,
+		messages: nonSystemMessages(context.messages),
+	};
+}
+
+/** `messages` with every prompt-state system message removed from conversation history. */
+export function nonSystemMessages<T extends { role: string }>(messages: readonly T[]): T[] {
+	return messages.filter((message) => message.role !== "system");
+}