pi-claude-agent-sdk 0.7.0

package/src/mcp-server.ts

@@ -0,0 +1,86 @@
+// In-process MCP server that exposes pi tools to Claude Code.
+//
+// Pi declares tool parameters as TypeBox objects, which are already JSON
+// Schema at runtime — the same thing MCP puts on the wire. This serves them
+// verbatim instead of going through the SDK's `createSdkMcpServer`, which only
+// accepts Zod and therefore forces a JSON Schema → Zod → JSON Schema round
+// trip. That round trip is lossy below the top level: nested objects collapse
+// to open records and `anyOf`/`const` vanish, so Claude saw only the first
+// level of any tool with a nested schema — including the builtin `edit`.
+//
+// Handlers go on the underlying protocol server rather than through
+// `McpServer.registerTool`, which is the Zod-only path. Skipping registerTool
+// also skips its argument validation, which is what we want: pi validates and
+// executes tools itself, and the arguments MCP sees are discarded. A rejection
+// there would only prevent the handler from running, stranding the call.
+//
+// This rests on the Agent SDK treating what we hand it as an opaque JSON-RPC
+// endpoint: `connectSdkMcpServer` in sdk.mjs calls `instance.connect(transport)`
+// and nothing else, so none of McpServer's higher-level machinery is required.
+// The `McpServer` wrapper is kept only because the SDK's `mcpServers` option is
+// typed against that class. If this breaks after an SDK update, check whether
+// the SDK began inspecting the instance — reading registered tools, or expecting
+// tools/list_changed notifications we never send.
+
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { CallToolRequestSchema, ListToolsRequestSchema } from "@modelcontextprotocol/sdk/types.js";
+import type { McpResult } from "./extract-tool-results.js";
+
+// Claude Code stamps every tools/call with the id of the tool_use block it came
+// from. That is the only reliable way to pair a call with its result: call order
+// is not guaranteed to match the order the tool_use blocks were emitted, so
+// counting calls mispairs results as soon as the two diverge.
+//
+// This is a Claude Code extension, not part of the MCP spec — CC sets it in
+// `src/services/mcp/client.ts` (see reference-code/claude-code-rip). If CC ever
+// stops sending it, every tool call fails with the error below rather than
+// silently pairing results to the wrong call, which is the intended tradeoff.
+const TOOL_USE_ID_META = "claudecode/toolUseId";
+
+export interface McpToolDef {
+	name: string;
+	description: string;
+	inputSchema: unknown;
+	handler: (toolCallId: string) => Promise<McpResult>;
+}
+
+// MCP requires an object schema. Pi types tool parameters as any TypeBox schema,
+// so a scalar or array one typechecks but cannot go on the wire — that is a bug
+// in the tool, and reporting it at startup names the culprit. Degrading it to
+// "takes no arguments" instead would surface much later as Claude calling the
+// tool with no arguments and pi's own validation rejecting them.
+function assertObjectSchema(tool: McpToolDef): void {
+	const schema = tool.inputSchema as Record<string, unknown> | undefined;
+	if (!schema || schema.type !== "object") {
+		throw new Error(`${tool.name}: MCP tool parameters must be an object schema, got ${JSON.stringify(schema)}`);
+	}
+}
+
+export function createToolServer(name: string, tools: McpToolDef[]) {
+	const server = new McpServer({ name, version: "1.0.0" }, { capabilities: { tools: {} } });
+	const byName = new Map(tools.map((tool) => [tool.name, tool]));
+	for (const tool of tools) assertObjectSchema(tool);
+
+	server.server.setRequestHandler(ListToolsRequestSchema, () => ({
+		tools: tools.map((tool) => ({
+			name: tool.name,
+			description: tool.description,
+			inputSchema: tool.inputSchema as Record<string, unknown>,
+		})),
+	}));
+
+	server.server.setRequestHandler(CallToolRequestSchema, async (request) => {
+		const tool = byName.get(request.params.name);
+		if (!tool) throw new Error(`Unknown tool: ${request.params.name}`);
+		const toolCallId = request.params._meta?.[TOOL_USE_ID_META];
+		if (typeof toolCallId !== "string") {
+			throw new Error(`${tool.name}: tools/call is missing _meta["${TOOL_USE_ID_META}"] — cannot pair the result with its tool call`);
+		}
+		// Narrowed deliberately: McpResult also carries `toolCallId`, which is our
+		// own bookkeeping for pairing and not part of MCP's CallToolResult.
+		const { content, isError } = await tool.handler(toolCallId);
+		return { content, isError };
+	});
+
+	return { type: "sdk" as const, name, instance: server };
+}

package/src/models.ts

@@ -0,0 +1,95 @@
+// Canonical selection + display order for the model picker.
+// `resolveModel` returns the first partial match, so `opus` resolves to the first-listed opus entry.
+// Extracted from index.ts so tests can import without activating the extension.
+
+export const MODEL_IDS_IN_ORDER = ["claude-fable-5", "claude-opus-5", "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6", "claude-sonnet-5", "claude-sonnet-4-6", "claude-haiku-4-5"];
+
+// Project pi-ai's model entries down to the fields pi's registerProvider expects,
+// and keep MODEL_IDS_IN_ORDER ordering. IDs missing from pi-ai are silently dropped.
+// Context-dependent display labels are applied after plan/long-context config is known.
+export function buildModels<T extends { id: string; [key: string]: any }>(piAiModels: T[]) {
+	return MODEL_IDS_IN_ORDER
+		.map((id) => piAiModels.find((m) => m.id === id))
+		.filter((m) => m != null)
+		// Forward thinkingLevelMap so pi-ai's per-model overrides (e.g. opus-4-8
+		// mapping xhigh→xhigh and max→max) are visible to the effort lookup.
+		.map(({ id, name, reasoning, input, contextWindow, maxTokens, thinkingLevelMap }) => ({
+			id,
+			name,
+			reasoning, input, contextWindow, maxTokens,
+			thinkingLevelMap,
+			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+		}));
+}
+
+export type LongContextSettings = {
+	plan: "pro" | "max";
+	longContextExtraUsage: boolean;
+};
+
+export type ClaudeCodeRuntimeModel = {
+	cliModelId: string;
+	contextWindow: number;
+};
+
+const TWO_HUNDRED_K_CONTEXT = 200_000;
+const ONE_M_CONTEXT = 1_000_000;
+
+// Measured Claude Agent SDK subscription/OAuth behavior. Do not infer this from
+// pi-ai's advertised contextWindow: bare Opus 4.7 serves 1M, bare Opus 4.8 does
+// not, and [1m] entitlement differs by model. See diag/CONTEXT-SIZE.md.
+export function resolveClaudeCodeRuntimeModel(modelId: string, settings: LongContextSettings): ClaudeCodeRuntimeModel {
+	switch (modelId) {
+		case "claude-opus-5":
+			return { cliModelId: "claude-opus-5[1m]", contextWindow: ONE_M_CONTEXT };
+		case "claude-opus-4-8":
+			return { cliModelId: "claude-opus-4-8[1m]", contextWindow: ONE_M_CONTEXT };
+		case "claude-opus-4-7":
+			return { cliModelId: "claude-opus-4-7", contextWindow: ONE_M_CONTEXT };
+		case "claude-opus-4-6": {
+			const useOneM = settings.plan === "max" || settings.longContextExtraUsage;
+			return {
+				cliModelId: useOneM ? "claude-opus-4-6[1m]" : "claude-opus-4-6",
+				contextWindow: useOneM ? ONE_M_CONTEXT : TWO_HUNDRED_K_CONTEXT,
+			};
+		}
+		case "claude-fable-5":
+			return { cliModelId: "claude-fable-5[1m]", contextWindow: ONE_M_CONTEXT };
+		case "claude-sonnet-5":
+			return { cliModelId: "claude-sonnet-5[1m]", contextWindow: ONE_M_CONTEXT };
+		case "claude-sonnet-4-6":
+			return {
+				cliModelId: settings.longContextExtraUsage ? "claude-sonnet-4-6[1m]" : "claude-sonnet-4-6",
+				contextWindow: settings.longContextExtraUsage ? ONE_M_CONTEXT : TWO_HUNDRED_K_CONTEXT,
+			};
+		case "claude-haiku-4-5":
+			return { cliModelId: "claude-haiku-4-5", contextWindow: TWO_HUNDRED_K_CONTEXT };
+		default:
+			console.error(`claude-bridge: encountered model ${modelId} with no known context size, defaulting to 200K`);
+			return { cliModelId: modelId, contextWindow: TWO_HUNDRED_K_CONTEXT };
+	}
+}
+
+export function claudeCodeModelId(model: { id: string }, settings: LongContextSettings): string {
+	return resolveClaudeCodeRuntimeModel(model.id, settings).cliModelId;
+}
+
+export function resolveModel<T extends { id: string }>(models: T[], input: string): T | undefined {
+	const lower = input.toLowerCase();
+	return models.find((m) => m.id === lower || m.id.includes(lower));
+}
+
+// Produce the model metadata registered with pi. The registered contextWindow must
+// match the window the bridge actually requests from Claude Code, or pi's status
+// bar and auto-compaction threshold will misreport. The runtime policy is based
+// on measured SDK behavior - see diag/CONTEXT-SIZE.md
+export function applyLongContext<T extends { id: string; name: string; contextWindow?: number | null }>(
+	models: T[],
+	settings: LongContextSettings,
+): T[] {
+	return models.map((m) => {
+		const { contextWindow } = resolveClaudeCodeRuntimeModel(m.id, settings);
+		const name = contextWindow > TWO_HUNDRED_K_CONTEXT && !/\b1M\b/i.test(m.name) ? `${m.name} 1M` : m.name;
+		return contextWindow === m.contextWindow && name === m.name ? m : { ...m, contextWindow, name };
+	});
+}

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,79 @@
+// 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;
+
+	// Per-turn (reset together)
+	turnOutput: AssistantMessage | null = null;
+	turnStarted = false;
+	turnSawStreamEvent = false;
+	turnSawToolCall = false;
+
+	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;
+		// 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,24 @@
+// Skills block extraction + MCP naming constants.
+// Extracted from index.ts so tests can import without activating the extension.
+
+export const MCP_SERVER_NAME = "custom-tools";
+export const MCP_TOOL_PREFIX = `mcp__${MCP_SERVER_NAME}__`;
+
+// Extract skills block from pi's system prompt for forwarding to Claude Code.
+export function extractSkillsBlock(systemPrompt?: string): string | undefined {
+	if (!systemPrompt) return undefined;
+	const startMarker = "The following skills provide specialized instructions for specific tasks.";
+	const endMarker = "</available_skills>";
+	const start = systemPrompt.indexOf(startMarker);
+	if (start === -1) return undefined;
+	const end = systemPrompt.indexOf(endMarker, start);
+	if (end === -1) return undefined;
+	return rewriteSkillsBlock(systemPrompt.slice(start, end + endMarker.length).trim());
+}
+
+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`,
+	);
+}