chore(vendor): import pi 1.3.0

Imported as a linear vendor commit so forge rebases do not replay the upstream root tree.

git-subtree-dir: third_party/pi

git-subtree-split: 71023aec1546d6d677a65669c5665c0927ce0dd0

third_party/pi/packages/coding-agent/src/extensions/codemode/execute.lazy.ts

@@ -0,0 +1,2 @@
+/** Loads the codemode sandbox runtime on the first script (see execute.ts). */
+export const loadCodemodeExecutor = () => import("./execute.ts");

third_party/pi/packages/coding-agent/src/extensions/codemode/execute.ts

@@ -0,0 +1,772 @@
+/**
+ * Runs one codemode script in the sandbox. Split from tool.ts and loaded through
+ * execute.lazy.ts so the sandbox runtime only loads when a script runs.
+ */
+
+import type { AgentTool, AgentToolCallOutcome, AgentToolResult } from "@earendil-works/pi-agent-core";
+import type {
+	AnyModel,
+	ClassifierContext,
+	ImageContent,
+	ImagesContext,
+	ModelType,
+	ModelTypeMap,
+	TextContent,
+	Usage,
+} from "@earendil-works/pi-ai";
+import {
+	type CodemodeOutputItem,
+	type CodemodeResult,
+	CodemodeSandbox,
+	type CodemodeTool,
+	loadQuickJSWasm,
+	parseCodemodeSource,
+	renderToolSample,
+	toCodemodeIdentifier,
+} from "@earendil-works/pi-codemode";
+import { getCodemodeWorkerSpecifier, getQuickJSWasmPath } from "../../config.ts";
+import type { ExtensionToolContext, ToolNamespace } from "../../core/extensions/types.ts";
+import type { SessionEntry } from "../../core/session-manager.ts";
+import { formatSize } from "../../core/tools/truncate.ts";
+import { combineUsage } from "../../core/usage-totals.ts";
+import { writeOutputFile } from "../../utils/output-files.ts";
+import { Bm25Ranker, createToolSearchDocument, DEFAULT_TOOL_SEARCH_LIMIT } from "../tool-search/tool.ts";
+import {
+	CODEMODE_DOCS_PATH,
+	CODEMODE_STORE_ENTRY_TYPE,
+	type CodemodeModelRuntime,
+	type CodemodeNestedCall,
+	type CodemodeStoreEntryData,
+	type CodemodeToolDetails,
+	type CodemodeToolInput,
+	type CodemodeToolOptions,
+	getCodemodeCallableTools,
+	toCodemodeDeclaration,
+} from "./tool.ts";
+
+const ARGS_PREVIEW_CHARS = 200;
+const ERROR_PREVIEW_CHARS = 500;
+/** `models.classify()` and `models.generateImages()` calls one script may have in flight; `Promise.all` over many items queues the rest. */
+const MAX_CONCURRENT_MODEL_CALLS = 4;
+/**
+ * Heap limit for the QuickJS VM. The worker shares pi's process, so without a limit a runaway
+ * script can grow to wasm32's 4 GiB and take the session down. Overruns throw
+ * `InternalError: out of memory` inside the script.
+ */
+const CODEMODE_MEMORY_LIMIT_BYTES = 256 * 1024 * 1024;
+// Shared across calls and sessions. Allow one replacement while a cancelled worker is still exiting.
+const MAX_CONCURRENT_CODEMODE_WORKERS = 2;
+let activeCodemodeWorkers = 0;
+const MODEL_TYPES: ReadonlySet<string> = new Set<ModelType>(["chat", "image", "classifier"]);
+
+function truncateText(text: string, maxChars: number): string {
+	return text.length > maxChars ? `${text.slice(0, maxChars - 3)}...` : text;
+}
+
+function previewArgs(args: unknown): string {
+	if (args === undefined) return "";
+	try {
+		return truncateText(JSON.stringify(args) ?? "", ARGS_PREVIEW_CHARS);
+	} catch {
+		return "";
+	}
+}
+
+function textOf(result: AgentToolResult<unknown>): string {
+	return (result.content ?? [])
+		.filter((block): block is TextContent => block.type === "text")
+		.map((block) => block.text)
+		.join("\n");
+}
+
+function toModelType(value: unknown): ModelType {
+	if (typeof value === "string" && MODEL_TYPES.has(value)) return value as ModelType;
+	throw new Error(`Unknown model type ${JSON.stringify(value)}. Use "chat", "image", or "classifier".`);
+}
+
+function toProvider(value: unknown): string | undefined {
+	if (value === undefined || value === null) return undefined;
+	if (typeof value !== "string") throw new Error("provider must be a string");
+	return value;
+}
+
+/** Catalog entry for scripts. `headers` is dropped because models.json headers can carry credentials. */
+function toModelInfo(model: AnyModel): Record<string, unknown> {
+	const info: Record<string, unknown> = { ...model };
+	delete info.headers;
+	return info;
+}
+
+/** `an image`, `a classifier`. */
+function withArticle(word: string): string {
+	return `${/^[aeiou]/.test(word) ? "an" : "a"} ${word}`;
+}
+
+/** A JSON value from a script, read field by field: values that are not objects have no fields. */
+type ScriptValue = { readonly [key: string]: unknown } | null | undefined;
+
+/** A script value in an error message: `undefined`, `a string`, `an array`, or its keys (`{ prompt }`). */
+function describeValue(value: unknown): string {
+	if (value === undefined || value === null) return String(value);
+	if (Array.isArray(value)) return value.length === 0 ? "an empty array" : "an array";
+	if (typeof value === "object") {
+		const keys = Object.keys(value);
+		if (keys.length === 0) return "{}";
+		return `{ ${keys.slice(0, 6).join(", ")}${keys.length > 6 ? ", ..." : ""} }`;
+	}
+	return typeof value === "string" ? "a string" : `a ${typeof value}`;
+}
+
+const CLASSIFIER_CONTEXT_SHAPE =
+	'{ state: { ... }, images?: [{ type: "image", data: <base64>, mimeType }], questions: { <id>: { type: "choice", instructions, criteria: { <label>: <meaning> } } | { type: "score", instructions, criteria: [<lowest level>, ..., <highest level>] } | { type: "bool", instructions, criteria: { true: <meaning>, false: <meaning> } } } }';
+
+/** Check a script's classifier context, so mistakes fail with the expected shape instead of a provider error. */
+function checkClassifierContext(context: unknown): ClassifierContext {
+	const fail = (problem: string) =>
+		new Error(
+			`models.classify() ${problem}. Expected context: ${CLASSIFIER_CONTEXT_SHAPE}. See "Classify" in ${CODEMODE_DOCS_PATH}.`,
+		);
+	if (typeof context !== "object" || context === null || Array.isArray(context)) {
+		throw fail(`expects a context object as its second argument, got ${describeValue(context)}`);
+	}
+	const { state, images, questions } = context as Record<string, unknown>;
+	if (typeof state !== "object" || state === null || Array.isArray(state)) {
+		throw fail(`context.state must be an object, got ${describeValue(state)}`);
+	}
+	if (images !== undefined) {
+		if (!Array.isArray(images)) throw fail(`context.images must be an array, got ${describeValue(images)}`);
+		images.forEach((image: ScriptValue, index) => {
+			if (image?.type !== "image" || typeof image.data !== "string" || typeof image.mimeType !== "string") {
+				throw fail(`context.images[${index}] must be an image block, got ${describeValue(image)}`);
+			}
+		});
+	}
+	if (
+		typeof questions !== "object" ||
+		questions === null ||
+		Array.isArray(questions) ||
+		Object.keys(questions).length === 0
+	) {
+		throw fail(`context.questions must map question IDs to questions, got ${describeValue(questions)}`);
+	}
+	const isStrings = (values: unknown[]) => values.length > 0 && values.every((value) => typeof value === "string");
+	for (const [id, question] of Object.entries(questions)) {
+		const at = `context.questions.${id}`;
+		if (typeof question !== "object" || question === null || Array.isArray(question)) {
+			throw fail(`${at} must be a question object, got ${describeValue(question)}`);
+		}
+		const { type, instructions, criteria } = question as Record<string, unknown>;
+		if (typeof instructions !== "string") throw fail(`${at}.instructions must be a string`);
+		if (type === "choice") {
+			if (
+				typeof criteria !== "object" ||
+				criteria === null ||
+				Array.isArray(criteria) ||
+				!isStrings(Object.values(criteria))
+			) {
+				throw fail(`${at} is a "choice" question, so criteria must map each label to its meaning`);
+			}
+		} else if (type === "score") {
+			if (!Array.isArray(criteria) || !isStrings(criteria)) {
+				throw fail(`${at} is a "score" question, so criteria must list the levels as strings, lowest first`);
+			}
+		} else if (type === "bool") {
+			const labels = criteria as ScriptValue;
+			if (typeof labels?.true !== "string" || typeof labels.false !== "string") {
+				throw fail(`${at} is a "bool" question, so criteria must be { true: string, false: string }`);
+			}
+		} else {
+			throw fail(`${at}.type must be "choice", "score", or "bool", got ${JSON.stringify(type)}`);
+		}
+	}
+	return context as unknown as ClassifierContext;
+}
+
+/** Check a script's image context, so mistakes such as `{ prompt }` fail with the expected shape. */
+function checkImagesContext(context: unknown): ImagesContext {
+	const fail = (problem: string) =>
+		new Error(
+			`models.generateImages() ${problem}. Expected context: { input: [{ type: "text", text: <prompt> }, ...optional { type: "image", data: <base64>, mimeType } references] }. See "Generate images" in ${CODEMODE_DOCS_PATH}.`,
+		);
+	if (typeof context !== "object" || context === null || Array.isArray(context)) {
+		throw fail(`expects a context object as its second argument, got ${describeValue(context)}`);
+	}
+	const { input } = context as Record<string, unknown>;
+	if (!Array.isArray(input) || input.length === 0) {
+		throw fail(`context.input must be a non-empty array of blocks, got ${describeValue(input)}`);
+	}
+	input.forEach((block: ScriptValue, index) => {
+		if (block?.type === "text" && typeof block.text === "string") return;
+		if (block?.type === "image" && typeof block.data === "string" && typeof block.mimeType === "string") return;
+		throw fail(`context.input[${index}] must be a text or image block, got ${describeValue(block)}`);
+	});
+	return context as unknown as ImagesContext;
+}
+
+/** The fields of `ClassifierResult` and `AssistantImages` that a nested call row reports. */
+interface ModelCallResult {
+	stopReason: "stop" | "error" | "aborted";
+	errorMessage?: string;
+	usage?: Usage;
+}
+
+/** Runs at most `limit` calls at once, in call order. */
+function createLimiter(limit: number): <T>(run: () => Promise<T>) => Promise<T> {
+	let active = 0;
+	const waiting: (() => void)[] = [];
+	return async (run) => {
+		if (active >= limit) await new Promise<void>((resolve) => waiting.push(resolve));
+		active++;
+		try {
+			return await run();
+		} finally {
+			active--;
+			waiting.shift()?.();
+		}
+	};
+}
+
+function isStoreEntryData(data: unknown): data is CodemodeStoreEntryData {
+	if (typeof data !== "object" || data === null) return false;
+	const { set, delete: deleted } = data as Partial<CodemodeStoreEntryData>;
+	return (
+		typeof set === "object" &&
+		set !== null &&
+		Array.isArray(deleted) &&
+		deleted.every((key: unknown) => typeof key === "string")
+	);
+}
+
+/** Values of `load()`: the `codemode-store` entries on the branch, applied from the root. */
+export function readCodemodeStore(branch: readonly SessionEntry[]): Record<string, unknown> {
+	const store = new Map<string, unknown>();
+	for (const entry of branch) {
+		if (entry.type !== "custom" || entry.customType !== CODEMODE_STORE_ENTRY_TYPE || !isStoreEntryData(entry.data)) {
+			continue;
+		}
+		for (const key of entry.data.delete) store.delete(key);
+		for (const [key, value] of Object.entries(entry.data.set)) store.set(key, value);
+	}
+	return Object.fromEntries(store);
+}
+
+/** Default token budget for script output. */
+const DEFAULT_MAX_OUTPUT_TOKENS = 10_000;
+/** Characters per token when estimating. */
+const CHARS_PER_TOKEN = 4;
+
+/** Like the script's `text()`: strings as is, other values as compact JSON. */
+function valueText(value: unknown): string {
+	if (typeof value === "string") return value;
+	return JSON.stringify(value) ?? String(value);
+}
+
+/**
+ * Lay out the script's output so the model can tell items apart: providers join adjacent text
+ * blocks with a newline or with nothing. With more than one text item (`text()` or the returned
+ * value), each starts with a `==> text N/M <==` line. `console.*` lines follow all other output in
+ * one `<console_output>` block.
+ */
+function formatOutput(output: readonly CodemodeOutputItem[]): (TextContent | ImageContent)[] {
+	const total = output.filter((item) => item.type === "text" && !item.console).length;
+	const items: (TextContent | ImageContent)[] = [];
+	const consoleLines: string[] = [];
+	let index = 0;
+	for (const item of output) {
+		if (item.type === "image") {
+			items.push(item);
+		} else if (item.console) {
+			consoleLines.push(item.text);
+		} else {
+			index++;
+			items.push({ type: "text", text: total > 1 ? `==> text ${index}/${total} <==\n${item.text}` : item.text });
+		}
+	}
+	if (consoleLines.length > 0) {
+		items.push({ type: "text", text: `<console_output>\n${consoleLines.join("\n")}\n</console_output>` });
+	}
+	return items;
+}
+
+/**
+ * Join adjacent text items into one, each part starting on its own line. The separator check looks
+ * at the last non-empty part, not the joined text: reading the end of a growing concatenation makes
+ * V8 copy it on every item, which is quadratic (#10678).
+ */
+function joinAdjacentText(items: (TextContent | ImageContent)[]): (TextContent | ImageContent)[] {
+	const joined: (TextContent | ImageContent)[] = [];
+	let parts: string[] = [];
+	let lastPart = "";
+	const flush = () => {
+		if (parts.length > 0) joined.push({ type: "text", text: parts.join("") });
+		parts = [];
+		lastPart = "";
+	};
+	for (const item of items) {
+		if (item.type !== "text") {
+			flush();
+			joined.push(item);
+			continue;
+		}
+		if (lastPart !== "" && !lastPart.endsWith("\n")) {
+			parts.push("\n");
+			lastPart = "\n";
+		}
+		parts.push(item.text);
+		if (item.text !== "") lastPart = item.text;
+	}
+	flush();
+	return joined;
+}
+
+function formatCallSummary(calls: readonly CodemodeNestedCall[]): string {
+	if (calls.length === 0) return "No tool calls were made.";
+	return `Tool calls made before the failure (they are not undone): ${calls.map((call) => `${call.name} (${call.status})`).join(", ")}`;
+}
+
+function formatError(result: Extract<CodemodeResult, { ok: false }>, calls: readonly CodemodeNestedCall[]): string {
+	const { error } = result;
+	const head =
+		error.kind === "script"
+			? (error.stack ?? `${error.name ?? "Error"}: ${error.message}`)
+			: error.kind === "timeout"
+				? `Script timed out: ${error.message}`
+				: error.kind === "aborted"
+					? `Script aborted: ${error.message}`
+					: `Script sandbox failed: ${error.message}`;
+	return `${head}\n\n${formatCallSummary(calls)}`;
+}
+
+/** Write the full text output to a temp file, like bash does for truncated output. */
+async function spillOutput(text: string): Promise<{ path: string } | { error: string }> {
+	try {
+		return { path: await writeOutputFile("pi-codemode", ".txt", text) };
+	} catch (error) {
+		return { error: error instanceof Error ? error.message : String(error) };
+	}
+}
+
+/** File extensions of the image types `image()` accepts. Must list every type the sandbox's `image()` detects. */
+const IMAGE_EXTENSIONS: Record<string, string> = {
+	"image/png": ".png",
+	"image/jpeg": ".jpg",
+	"image/gif": ".gif",
+	"image/webp": ".webp",
+};
+
+/**
+ * Save each image to a temp file and put a text item with its path before it. The model sees the
+ * image but has no other way to reach its bytes: scripts cannot write files, and `write` only takes
+ * text. Images shown more than once are saved once.
+ */
+async function saveImages(items: (TextContent | ImageContent)[]): Promise<(TextContent | ImageContent)[]> {
+	const labels = new Map<string, Promise<string>>();
+	const label = async ({ data, mimeType }: ImageContent): Promise<string> => {
+		const bytes = Buffer.from(data, "base64");
+		const kind = `${mimeType}, ${formatSize(bytes.length)}`;
+		const extension = IMAGE_EXTENSIONS[mimeType];
+		if (!extension) throw new Error(`No file extension for image type ${mimeType}`);
+		// A failed write (disk full, unwritable temp dir) must not discard the result of a script whose
+		// tool calls already ran, so it becomes part of the label.
+		try {
+			const path = await writeOutputFile("pi-codemode", extension, bytes);
+			return `[Image saved to ${path} (${kind})]`;
+		} catch (error) {
+			return `[Image (${kind}) could not be saved: ${error instanceof Error ? error.message : String(error)}]`;
+		}
+	};
+	const result = await Promise.all(
+		items.map(async (item): Promise<(TextContent | ImageContent)[]> => {
+			if (item.type !== "image") return [item];
+			let pending = labels.get(item.data);
+			if (!pending) {
+				pending = label(item);
+				labels.set(item.data, pending);
+			}
+			return [{ type: "text", text: await pending }, item];
+		}),
+	);
+	return result.flat();
+}
+
+/**
+ * Apply the token budget: when the combined text exceeds it, the text items become one
+ * item that keeps the start and end of the text, and images follow it. The full text is written to
+ * a temp file.
+ */
+async function truncateOutput(
+	items: (TextContent | ImageContent)[],
+	maxTokens: number,
+): Promise<{ items: (TextContent | ImageContent)[]; fullOutputPath?: string }> {
+	const texts = items.filter((item): item is TextContent => item.type === "text").map((item) => item.text);
+	const combined = texts.join("\n");
+	const budget = maxTokens * CHARS_PER_TOKEN;
+	if (texts.length === 0 || combined.length <= budget) return { items };
+	const headChars = Math.floor(budget / 2);
+	const tailChars = budget - headChars;
+	const removed = combined.length - headChars - tailChars;
+	const head = combined.slice(0, headChars);
+	const tail = tailChars > 0 ? combined.slice(-tailChars) : "";
+	let text = `Warning: truncated output (original token count: ${Math.ceil(combined.length / CHARS_PER_TOKEN)})\nTotal output lines: ${combined.split("\n").length}\n\n${head}…${Math.ceil(removed / CHARS_PER_TOKEN)} tokens truncated…${tail}`;
+	const spilled = await spillOutput(combined);
+	text +=
+		"path" in spilled
+			? `\n\n[Full output: ${spilled.path} (read with offset/limit)]`
+			: `\n\n[Could not save the full output: ${spilled.error}]`;
+	return {
+		items: [{ type: "text", text }, ...items.filter((item) => item.type === "image")],
+		...("path" in spilled ? { fullOutputPath: spilled.path } : {}),
+	};
+}
+
+/**
+ * The value a script receives for a nested call: a tool that declares
+ * `outputSchema` resolves to its `structuredContent`, also for error results that carry one (such
+ * as MCP results with `isError`); any other tool resolves to its text content. Other failures
+ * reject with the tool's error text.
+ */
+function toScriptValue(tool: AgentTool<any>, outcome: AgentToolCallOutcome): unknown {
+	const { result } = outcome;
+	if (tool.outputSchema && result.structuredContent !== undefined) return result.structuredContent;
+	const text = textOf(result);
+	if (outcome.isError) throw new Error(text || `Tool "${tool.name}" failed`);
+	return text;
+}
+
+/**
+ * Run one script. Without a session context (a plain Agent or a direct call) scripts cannot call
+ * tools, `store()` starts empty, and writes are dropped.
+ */
+export async function executeCodemode(
+	toolCallId: string,
+	input: CodemodeToolInput,
+	signal: AbortSignal | undefined,
+	onUpdate: ((result: AgentToolResult<CodemodeToolDetails>) => void) | undefined,
+	ctx: ExtensionToolContext | undefined,
+	options: CodemodeToolOptions = {},
+): Promise<AgentToolResult<CodemodeToolDetails>> {
+	const startedAt = performance.now();
+	const { code, options: sourceOptions } = parseCodemodeSource(input.code);
+	const calls: CodemodeNestedCall[] = [];
+	// Usage of the script's `models.*` calls. Nested tool calls report theirs through the session.
+	let modelUsage: Usage | undefined;
+	// Images returned by `models.generateImages()`, to notice a script that never shows them.
+	let generatedImages = 0;
+	const addGeneratedImages = (count: number) => {
+		generatedImages += count;
+	};
+	const addModelUsage = (usage: Usage) => {
+		modelUsage = modelUsage ? combineUsage(modelUsage, usage) : usage;
+	};
+
+	const snapshot = (): CodemodeToolDetails => ({ calls: calls.map((call) => ({ ...call })) });
+	const publish = () => onUpdate?.({ content: [], details: snapshot() });
+	// Show the script as running before its first nested call, so the row's timer starts.
+	publish();
+
+	const callable = ctx ? getCodemodeCallableTools(ctx.tools) : [];
+	// ALL_TOOLS entries carry the declaration.
+	const guidelines = options.getToolGuidelines?.();
+	const samples = new Map(
+		callable.map((tool) => [tool.name, renderToolSample(toCodemodeDeclaration(tool, guidelines?.get(tool.name)))]),
+	);
+	const sandboxTools: CodemodeTool[] = callable.map((tool) => ({
+		name: tool.name,
+		description: samples.get(tool.name),
+		execute: async (args, { signal: callSignal }) => {
+			const record: CodemodeNestedCall = {
+				id: `${toolCallId}/?`,
+				name: tool.name,
+				args: previewArgs(args),
+				status: "running",
+			};
+			calls.push(record);
+			publish();
+			const callStartedAt = performance.now();
+			// Only tools from ctx.tools are callable, so ctx is set here.
+			if (!ctx) throw new Error("Tool calls need a session");
+			const outcome = await ctx.executeTool(tool.name, args, { signal: callSignal });
+			record.id = outcome.toolCall.id;
+			record.durationMs = performance.now() - callStartedAt;
+			if (outcome.isError) {
+				record.status = callSignal.aborted ? "cancelled" : "error";
+				record.error = truncateText(textOf(outcome.result) || `Tool "${tool.name}" failed`, ERROR_PREVIEW_CHARS);
+			} else {
+				record.status = "ok";
+			}
+			publish();
+			return toScriptValue(tool, outcome);
+		},
+	}));
+
+	const sandbox = new CodemodeSandbox({
+		tools: sandboxTools,
+		globals: [
+			...createDiscoveryGlobals(callable, samples, options),
+			...(options.models && ctx
+				? createModelGlobals(ctx.modelRegistry, toolCallId, calls, publish, addModelUsage, addGeneratedImages)
+				: []),
+		],
+		timeoutMs: sourceOptions.timeoutMs ?? Number.POSITIVE_INFINITY,
+		memoryLimitBytes: CODEMODE_MEMORY_LIMIT_BYTES,
+		wasm: loadQuickJSWasm(getQuickJSWasmPath()),
+		workerUrl: getCodemodeWorkerSpecifier(),
+		acquireWorker: () => {
+			if (activeCodemodeWorkers >= MAX_CONCURRENT_CODEMODE_WORKERS) {
+				throw new Error(
+					`Codemode worker limit reached (${MAX_CONCURRENT_CODEMODE_WORKERS} workers are still running or shutting down). Retry after they exit.`,
+				);
+			}
+			activeCodemodeWorkers++;
+			return () => {
+				activeCodemodeWorkers--;
+			};
+		},
+	});
+
+	let result: CodemodeResult;
+	try {
+		const store = ctx ? readCodemodeStore(ctx.sessionManager.getBranch()) : {};
+		result = await sandbox.execute(code, { signal, store });
+	} finally {
+		await sandbox.close();
+	}
+	// Calls still marked running were cut off by the script ending, a timeout, or an abort.
+	for (const call of calls) {
+		if (call.status === "running") call.status = "cancelled";
+	}
+
+	const scriptOutput = [...result.output];
+	if (result.ok) {
+		const { set, delete: deleted } = result.storeWrites;
+		if (Object.keys(set).length > 0 || deleted.length > 0) {
+			options.appendEntry?.(CODEMODE_STORE_ENTRY_TYPE, { set, delete: deleted });
+		}
+		// pi extension: a returned value is appended like text().
+		if (result.value !== undefined) scriptOutput.push({ type: "text", text: valueText(result.value) });
+	}
+	const items = formatOutput(scriptOutput);
+	if (!result.ok) items.push({ type: "text", text: `Script error:\n${formatError(result, calls)}` });
+	if (generatedImages > 0 && !items.some((item) => item.type === "image")) {
+		items.push({
+			type: "text",
+			text: `Note: models.generateImages() returned ${generatedImages} image${generatedImages === 1 ? "" : "s"} that the script did not show. Show each image block of result.output with image(block).`,
+		});
+	}
+
+	const truncated = await truncateOutput(
+		joinAdjacentText(items),
+		sourceOptions.maxOutputTokens ?? DEFAULT_MAX_OUTPUT_TOKENS,
+	);
+	// After truncation, which joins the text items and moves images after them, so each path stays
+	// next to its image and is never cut.
+	const output = joinAdjacentText(await saveImages(truncated.items));
+	const wallTime = ((performance.now() - startedAt) / 1000).toFixed(1);
+	const header = `${result.ok ? "Script completed" : "Script failed"}\nWall time ${wallTime} seconds\nOutput:\n`;
+	const details = snapshot();
+	if (truncated.fullOutputPath) details.fullOutputPath = truncated.fullOutputPath;
+	return {
+		content: [{ type: "text", text: header }, ...output],
+		details,
+		...(modelUsage ? { usage: modelUsage } : {}),
+		...(result.ok ? {} : { isError: true }),
+	};
+}
+
+/**
+ * Whether `query` names the namespace: its name or one of its aliases, as written or as the same
+ * script identifier (`dev-radius` and `dev_radius` are both `dev_radius`).
+ */
+function isNamespaceName(namespace: ToolNamespace, query: string): boolean {
+	const queryId = toCodemodeIdentifier(query);
+	return [namespace.name, ...(namespace.aliases ?? [])].some(
+		(name) => name === query || toCodemodeIdentifier(name) === queryId,
+	);
+}
+
+/**
+ * `searchTools()`, `describeTool()`, and `describeNamespace()`: ranked search and lookup over the
+ * script's nested tools and their namespaces.
+ */
+function createDiscoveryGlobals(
+	tools: readonly AgentTool<any>[],
+	samples: ReadonlyMap<string, string>,
+	options: CodemodeToolOptions,
+): CodemodeTool[] {
+	const ranker = new Bm25Ranker();
+	const entry = (name: string) => ({ name: toCodemodeIdentifier(name), description: samples.get(name) ?? "" });
+	return [
+		{
+			name: "searchTools",
+			spread: true,
+			execute: (args) => {
+				const [query, searchOptions] = args as [unknown, { limit?: unknown; namespace?: unknown } | undefined];
+				if (typeof query !== "string") throw new Error("searchTools() expects a query string");
+				const limit = searchOptions?.limit ?? DEFAULT_TOOL_SEARCH_LIMIT;
+				if (typeof limit !== "number" || !Number.isInteger(limit) || limit <= 0) {
+					throw new Error("searchTools() limit must be a positive integer");
+				}
+				const namespace = searchOptions?.namespace;
+				if (namespace !== undefined && namespace !== null && typeof namespace !== "string") {
+					throw new Error("searchTools() namespace must be a string");
+				}
+				const documents = tools.flatMap((tool) => {
+					const toolNamespace = options.getToolNamespace?.(tool.name);
+					if (namespace && (!toolNamespace || !isNamespaceName(toolNamespace, namespace))) return [];
+					return [createToolSearchDocument(tool, toolNamespace)];
+				});
+				return ranker.rank(query, documents, limit).map((match) => entry(match.name));
+			},
+		},
+		{
+			name: "describeTool",
+			spread: true,
+			execute: (args) => {
+				const [name] = args as unknown[];
+				if (typeof name !== "string") throw new Error("describeTool() expects a tool name");
+				const tool = tools.find(
+					(candidate) => candidate.name === name || toCodemodeIdentifier(candidate.name) === name,
+				);
+				return tool ? samples.get(tool.name) : undefined;
+			},
+		},
+		{
+			name: "describeNamespace",
+			spread: true,
+			execute: (args) => {
+				const [name] = args as unknown[];
+				if (typeof name !== "string") throw new Error("describeNamespace() expects a namespace name");
+				let namespace: ToolNamespace | undefined;
+				const names: string[] = [];
+				for (const tool of tools) {
+					const toolNamespace = options.getToolNamespace?.(tool.name);
+					if (!toolNamespace || !isNamespaceName(toolNamespace, name)) continue;
+					namespace ??= toolNamespace;
+					names.push(toCodemodeIdentifier(tool.name));
+				}
+				if (!namespace) return undefined;
+				return {
+					name: namespace.name,
+					...(namespace.description ? { description: namespace.description } : {}),
+					...(namespace.instructions ? { instructions: namespace.instructions } : {}),
+					tools: names,
+				};
+			},
+		},
+	];
+}
+
+/**
+ * `models.*` for scripts: the model registry methods documented in docs/codemode.md.
+ * Classifier and image calls appear as nested call rows so the renderer shows them, and their usage
+ * goes to `addUsage`. Rows show only the model, never prompts or image data.
+ */
+function createModelGlobals(
+	models: CodemodeModelRuntime,
+	toolCallId: string,
+	calls: CodemodeNestedCall[],
+	publish: () => void,
+	addUsage: (usage: Usage) => void,
+	addGeneratedImages: (count: number) => void,
+): CodemodeTool[] {
+	const limit = createLimiter(MAX_CONCURRENT_MODEL_CALLS);
+	let callCount = 0;
+
+	/**
+	 * Resolve the script's model by provider and id only, check the context, then run the call as a
+	 * nested call row. A script-supplied baseUrl or headers must never receive the credentials.
+	 */
+	const runModelCall = async <TType extends "classifier" | "image", TContext, TResult extends ModelCallResult>(
+		name: string,
+		type: TType,
+		[model, context]: unknown[],
+		checkContext: (context: unknown) => TContext,
+		run: (resolved: ModelTypeMap[TType], context: TContext) => Promise<TResult>,
+	): Promise<TResult> => {
+		const listHint = `List the ${type} models you can use with models.getAvailableOfType("${type}").`;
+		const given = model as ScriptValue;
+		if (typeof given?.provider !== "string" || typeof given.id !== "string") {
+			// undefined arrives as null: spread arguments cross the sandbox as a JSON array.
+			const undefinedHint =
+				model === undefined || model === null
+					? " models.getModelOfType() returns undefined for an unknown provider or id."
+					: "";
+			throw new Error(
+				`${name}() expects ${withArticle(type)} model as its first argument, got ${describeValue(model)}.${undefinedHint} ${listHint}`,
+			);
+		}
+		const provider = given.provider;
+		const id = given.id;
+		const ref = `${provider}/${id}`;
+		const resolved = models.getModelOfType(type, provider, id);
+		if (!resolved) {
+			const actualType = [...MODEL_TYPES].find(
+				(other) => other !== type && models.getModelOfType(other as ModelType, provider, id) !== undefined,
+			);
+			throw new Error(
+				actualType
+					? `"${ref}" is ${withArticle(actualType)} model, not ${withArticle(type)} model. ${listHint}`
+					: `Unknown ${type} model "${ref}". ${listHint}`,
+			);
+		}
+		const checked = checkContext(context);
+
+		const record: CodemodeNestedCall = {
+			id: `${toolCallId}/${name}/${++callCount}`,
+			name,
+			args: `${resolved.provider}/${resolved.id}`,
+			status: "running",
+		};
+		calls.push(record);
+		publish();
+		const startedAt = performance.now();
+		const result = await limit(() => run(resolved, checked));
+		record.durationMs = performance.now() - startedAt;
+		record.status = result.stopReason === "stop" ? "ok" : result.stopReason === "aborted" ? "cancelled" : "error";
+		if (result.errorMessage) record.error = truncateText(result.errorMessage, ERROR_PREVIEW_CHARS);
+		if (result.usage) {
+			record.cost = result.usage.cost.total;
+			addUsage(result.usage);
+		}
+		publish();
+		return result;
+	};
+	const implementations: Record<string, CodemodeTool["execute"]> = {
+		"models.getModelsOfType": (args) => {
+			const [type, provider] = args as unknown[];
+			return models.getModelsOfType(toModelType(type), toProvider(provider)).map(toModelInfo);
+		},
+		"models.getAvailableOfType": async (args, { signal }) => {
+			const [type, provider] = args as unknown[];
+			const available = await models.getAvailableOfType(toModelType(type), toProvider(provider), { signal });
+			return available.map(toModelInfo);
+		},
+		"models.getModelOfType": (args) => {
+			const [type, provider, id] = args as unknown[];
+			if (typeof provider !== "string" || typeof id !== "string") {
+				throw new Error(
+					`models.getModelOfType(type, provider, id) expects three strings, got (${(args as unknown[]).map(describeValue).join(", ")}). The provider and the id are separate arguments, for example models.getModelOfType("classifier", "typesafe", "jev-latest").`,
+				);
+			}
+			const model = models.getModelOfType(toModelType(type), provider, id);
+			return model === undefined ? undefined : toModelInfo(model);
+		},
+		"models.classify": (args, { signal }) =>
+			runModelCall("models.classify", "classifier", args as unknown[], checkClassifierContext, (resolved, context) =>
+				models.classify(resolved, context, { signal }),
+			),
+		"models.generateImages": (args, { signal }) =>
+			runModelCall(
+				"models.generateImages",
+				"image",
+				args as unknown[],
+				checkImagesContext,
+				async (resolved, context) => {
+					const result = await models.generateImages(resolved, context, { signal });
+					addGeneratedImages(result.output.filter((block) => block.type === "image").length);
+					return result;
+				},
+			),
+	};
+	return Object.entries(implementations).map(([name, execute]) => ({ name, spread: true, execute }));
+}

third_party/pi/packages/coding-agent/src/extensions/codemode/index.ts

@@ -0,0 +1,49 @@
+/**
+ * The `codemode` tool as an extension. The CLI loads it as a built-in extension; SDK users
+ * add `createCodemodeExtension()` to their extension factories.
+ *
+ * `codemode` is registered inactive. Activate it with `--tools`, the `defaultTools` setting, or
+ * `setActiveTools()`; the MCP extension activates it when MCP tools are only reachable from scripts.
+ */
+
+import type { ExtensionAPI, ExtensionFactory } from "../../core/extensions/types.ts";
+import { SETTINGS_DEFAULTS } from "../../core/settings-defaults.ts";
+import type { CodemodeMode } from "../../core/settings-manager.ts";
+import { createCodemodeToolDefinition } from "./tool.ts";
+
+export interface CodemodeExtensionOptions {
+	/** Overrides the `codemode.mode` setting. */
+	mode?: CodemodeMode;
+	/** Overrides the `codemode.inlineBudget` setting. */
+	inlineBudget?: number;
+	/** Expose the model catalog and classifiers to scripts as `models`. Default: `true`. */
+	models?: boolean;
+}
+
+function readMode(pi: ExtensionAPI): CodemodeMode {
+	return pi.getSettings().codemode?.mode === "only" ? "only" : SETTINGS_DEFAULTS.codemode.mode;
+}
+
+function readInlineBudget(pi: ExtensionAPI): number | undefined {
+	const budget = pi.getSettings().codemode?.inlineBudget;
+	return typeof budget === "number" && Number.isFinite(budget) && budget >= 0 ? budget : undefined;
+}
+
+export function createCodemodeExtension(options: CodemodeExtensionOptions = {}): ExtensionFactory {
+	return (pi) => {
+		pi.registerTool({
+			...createCodemodeToolDefinition({
+				appendEntry: (customType, data) => pi.appendEntry(customType, data),
+				models: options.models ?? true,
+				getToolNamespace: (name) => pi.getAllTools().find((tool) => tool.name === name)?.namespace,
+				getToolGuidelines: () =>
+					new Map(pi.getAllTools().map((tool) => [tool.name, tool.promptGuidelines ?? []] as const)),
+				getMode: () => options.mode ?? readMode(pi),
+				getInlineBudget: () => options.inlineBudget ?? readInlineBudget(pi),
+			}),
+			defaultActive: false,
+		});
+	};
+}
+
+export default createCodemodeExtension();

third_party/pi/packages/coding-agent/src/extensions/codemode/renderer.ts

@@ -0,0 +1,156 @@
+/**
+ * Presentation for the codemode tool.
+ *
+ * The call shows the script; the result lists the nested tool calls with their status as they
+ * run and the cost of its model calls, followed by the script output without the "Script completed"
+ * header, and the time the script has been running (`Elapsed`) or took (`Took`), like bash rows.
+ * Nested calls are not separate tool rows because they never reach the model as tool calls.
+ */
+
+import { Container, Spacer, Text } from "@earendil-works/pi-tui";
+import type { ToolDefinition } from "../../core/extensions/types.ts";
+import { getTextOutput, replaceTabs, startToolTimer, str, toolTimerText } from "../../core/tools/render-utils.ts";
+import { keyHint } from "../../modes/interactive/components/keybinding-hints.ts";
+import { VisualLinePreview } from "../../modes/interactive/components/visual-truncate.ts";
+import { highlightCode, type Theme } from "../../modes/interactive/theme/theme.ts";
+import type { CodemodeNestedCall, CodemodeToolDetails } from "./tool.ts";
+
+const CODE_PREVIEW_LINES = 10;
+const CALL_PREVIEW_COUNT = 8;
+const OUTPUT_PREVIEW_LINES = 5;
+const COLLAPSED_ARGS_CHARS = 80;
+const SCRIPT_HEADER = /^Script (completed|failed)\nWall time [\d.]+ seconds\nOutput:\n$/;
+
+function expandHint(theme: Theme, hidden: number, noun: string): string {
+	return `${theme.fg("muted", `... (${hidden} more ${noun},`)} ${keyHint("app.tools.expand", "to expand")}${theme.fg("muted", ")")}`;
+}
+
+function formatDuration(ms: number | undefined): string {
+	if (ms === undefined) return "";
+	return ms < 1000 ? `${Math.round(ms)}ms` : `${(ms / 1000).toFixed(1)}s`;
+}
+
+/** Cents for larger amounts, two significant digits for the fractions of a cent classifier calls cost. */
+function formatCost(cost: number): string {
+	return `$${cost >= 0.01 ? cost.toFixed(2) : cost.toPrecision(2)}`;
+}
+
+function statusIcon(call: CodemodeNestedCall, theme: Theme): string {
+	switch (call.status) {
+		case "running":
+			return theme.fg("warning", "…");
+		case "ok":
+			return theme.fg("success", "✓");
+		case "error":
+			return theme.fg("error", "✗");
+		case "cancelled":
+			return theme.fg("muted", "⊘");
+	}
+}
+
+function formatCall(call: CodemodeNestedCall, theme: Theme, expanded: boolean): string {
+	const args =
+		!expanded && call.args.length > COLLAPSED_ARGS_CHARS
+			? `${call.args.slice(0, COLLAPSED_ARGS_CHARS - 3)}...`
+			: call.args;
+	const duration = formatDuration(call.durationMs);
+	let line = `${statusIcon(call, theme)} ${theme.fg("toolTitle", call.name)}`;
+	if (args) line += ` ${theme.fg("muted", args)}`;
+	if (duration) line += ` ${theme.fg("dim", duration)}`;
+	if (call.cost) line += ` ${theme.fg("dim", formatCost(call.cost))}`;
+	if (expanded && call.error) line += `\n    ${theme.fg("error", call.error.split("\n").join("\n    "))}`;
+	return line;
+}
+
+export const codemodeRenderers: Pick<
+	ToolDefinition<any, CodemodeToolDetails | undefined>,
+	"renderCall" | "renderResult"
+> = {
+	renderCall(args, theme, context) {
+		startToolTimer(context);
+		// The code includes the `// @options:` line, so options show as part of the script.
+		const code = str((args as { code?: unknown } | undefined)?.code);
+		const title = theme.fg("toolTitle", theme.bold("codemode"));
+		const component = (context.lastComponent as Container | undefined) ?? new Container();
+		component.clear();
+		if (code === null) {
+			component.addChild(new Text(`${title} ${theme.fg("error", "[invalid arg]")}`, 0, 0));
+			return component;
+		}
+		component.addChild(new Text(title, 0, 0));
+		if (code) {
+			const highlighted = highlightCode(replaceTabs(code.replace(/\r/g, "").trimEnd()), "javascript").join("\n");
+			component.addChild(
+				context.expanded
+					? new Text(highlighted, 0, 0)
+					: new VisualLinePreview({
+							text: highlighted,
+							maxVisualLines: CODE_PREVIEW_LINES,
+							keep: "start",
+							formatHint: (hidden) => expandHint(theme, hidden, "lines"),
+						}),
+			);
+		}
+		return component;
+	},
+	renderResult(result, options, theme, context) {
+		const timerText = toolTimerText(context, options.isPartial);
+		const component = (context.lastComponent as Container | undefined) ?? new Container();
+		component.clear();
+		const calls = result.details?.calls ?? [];
+		if (calls.length > 0) {
+			const shown = options.expanded ? calls : calls.slice(-CALL_PREVIEW_COUNT);
+			const lines = shown.map((call) => formatCall(call, theme, options.expanded));
+			if (shown.length < calls.length) {
+				lines.unshift(
+					`${theme.fg("muted", `... (${calls.length - shown.length} earlier calls,`)} ${keyHint("app.tools.expand", "to expand")}${theme.fg("muted", ")")}`,
+				);
+			}
+			// Collapsed rows hide earlier calls, so the total covers every call.
+			const priced = calls.filter((call) => call.cost);
+			if (priced.length > 1) {
+				const total = priced.reduce((sum, call) => sum + (call.cost ?? 0), 0);
+				lines.push(theme.fg("muted", `Model calls: ${formatCost(total)}`));
+			}
+			component.addChild(new Spacer(1));
+			component.addChild(new Text(lines.join("\n"), 0, 0));
+		}
+
+		// Drop the "Script completed\nWall time ...\nOutput:\n" header. Rejected input (invalid options)
+		// has no header.
+		const [first, ...rest] = result.content;
+		const hasHeader = first?.type === "text" && SCRIPT_HEADER.test(first.text);
+		const output = options.isPartial
+			? ""
+			: getTextOutput({ ...result, content: hasHeader ? rest : result.content }, context.showImages).trim();
+		if (output) {
+			const color = context.isError ? "error" : "toolOutput";
+			const styled = replaceTabs(output)
+				.split("\n")
+				.map((line) => theme.fg(color, line))
+				.join("\n");
+			component.addChild(new Spacer(1));
+			if (options.expanded) {
+				component.addChild(new Text(styled, 0, 0));
+			} else {
+				// Limit wrapped lines, not logical ones: script output is often one long JSON line.
+				component.addChild(
+					new VisualLinePreview({
+						text: styled,
+						maxVisualLines: OUTPUT_PREVIEW_LINES,
+						keep: "start",
+						formatHint: (hidden) => expandHint(theme, hidden, "lines"),
+					}),
+				);
+				// The collapsed preview hides the truncation notice at the end, so name the file here.
+				const fullOutputPath = result.details?.fullOutputPath;
+				if (fullOutputPath) component.addChild(new Text(theme.fg("muted", `Full output: ${fullOutputPath}`), 0, 0));
+			}
+		}
+		if (timerText) {
+			component.addChild(new Spacer(1));
+			component.addChild(new Text(theme.fg("muted", timerText), 0, 0));
+		}
+		return component;
+	},
+};

third_party/pi/packages/coding-agent/src/extensions/codemode/tool.ts

@@ -0,0 +1,411 @@
+/**
+ * The `codemode` tool: the model writes JavaScript that calls other tools. Scripts use `tools`,
+ * `ALL_TOOLS`, `text()`, `image()`, `exit()`, `store()`/`load()`, `console.*`, and `return <value>`,
+ * may start with a `// @options:` line, and reach the model catalog, classifiers, and image models
+ * through `models.*`. Results start with a "Script completed" or "Script failed" header.
+ *
+ * Scripts can call the agent loop's nested tools: active `direct` tools and every `codemode` or
+ * `deferred` tool. Nested calls run through the agent loop's tool pipeline (`ctx.executeTool`), so
+ * validation, `tool_call`/`tool_result` hooks, and permission checks apply exactly as for direct
+ * calls. Only the script's output reaches the model; nested results do not.
+ *
+ * Nested results are handed to the script as follows:
+ * - A tool that declares `outputSchema` resolves to its `structuredContent`, also for error
+ *   results that carry one.
+ * - Any other tool resolves to its text content as one string.
+ * - A failed, blocked, or invalid call rejects with an Error carrying the tool's error text.
+ *
+ * A script that fails returns a normal error result that keeps its partial output, followed by
+ * "Script error:" and the error. `store(key, value)` and `load(key)` keep JSON values across
+ * calls; successful scripts append their writes to the session as `codemode-store` custom entries,
+ * so each branch sees the values written on its own path.
+ */
+
+import { join } from "node:path";
+import type { AgentTool } from "@earendil-works/pi-agent-core";
+import type { CodemodeJsonSchema, CodemodeTool } from "@earendil-works/pi-codemode";
+import { renderToolOutputType, renderToolSample, toCodemodeIdentifier } from "@earendil-works/pi-codemode/declarations";
+import { CODEMODE_SOURCE_GRAMMAR } from "@earendil-works/pi-codemode/source";
+import { type Static, Type } from "typebox";
+import { getDocsPath } from "../../config.ts";
+import type { ToolDefinition, ToolLoadout, ToolLoadoutChanges, ToolNamespace } from "../../core/extensions/types.ts";
+import type { ModelRegistry } from "../../core/model-registry.ts";
+import { SETTINGS_DEFAULTS } from "../../core/settings-defaults.ts";
+import type { CodemodeMode } from "../../core/settings-manager.ts";
+import { wrapToolDefinition } from "../../core/tools/tool-definition-wrapper.ts";
+import { loadCodemodeExecutor } from "./execute.lazy.ts";
+import { codemodeRenderers } from "./renderer.ts";
+
+export const CODEMODE_TOOL_NAME = "codemode";
+
+/** Custom entry type holding one script's `store()` writes: {@link CodemodeStoreEntryData}. */
+export const CODEMODE_STORE_ENTRY_TYPE = "codemode-store";
+
+export interface CodemodeStoreEntryData {
+	set: Record<string, unknown>;
+	delete: string[];
+}
+
+/** The part of the model registry that scripts reach through `models`. */
+export type CodemodeModelRuntime = Pick<
+	ModelRegistry,
+	"getModelsOfType" | "getAvailableOfType" | "getModelOfType" | "classify" | "generateImages"
+>;
+
+export interface CodemodeToolOptions {
+	/** Namespace of a tool, for `searchTools()` ranking and its `namespace` filter. */
+	getToolNamespace?: (toolName: string) => ToolNamespace | undefined;
+	/** Prompt guidelines of every tool, by tool name, shown with declarations by `describeTool()` and `ALL_TOOLS`. */
+	getToolGuidelines?: () => ReadonlyMap<string, readonly string[]>;
+	/**
+	 * Expose the `models` namespace to scripts, backed by the session's model registry
+	 * (`ctx.modelRegistry`). Without it, `models` is not declared.
+	 */
+	models?: boolean;
+	/**
+	 * Persists `store()` writes as a session custom entry. Without it, writes last only for the
+	 * current script; `load()` still reads entries already on the branch.
+	 */
+	appendEntry?: (customType: string, data: CodemodeStoreEntryData) => void;
+	/** How the tool presents the loadout while active (the `codemode.mode` setting). Default: `on`. */
+	getMode?: () => CodemodeMode;
+	/** Token budget for tool declarations in the description. Default: {@link DEFAULT_CODEMODE_INLINE_BUDGET}. */
+	getInlineBudget?: () => number | undefined;
+}
+
+const TEXT_OUTPUT_SCHEMA: CodemodeJsonSchema = { type: "string" };
+
+export const codemodeSchema = Type.Object({
+	code: Type.String({
+		description: "Raw JavaScript source.",
+	}),
+});
+
+export type CodemodeToolInput = Static<typeof codemodeSchema>;
+
+export type CodemodeNestedCallStatus = "running" | "ok" | "error" | "cancelled";
+
+export interface CodemodeNestedCall {
+	/** Tool call id of the nested call, `<codemode call id>/<n>`. */
+	id: string;
+	name: string;
+	/** Compact JSON of the arguments, truncated for display. */
+	args: string;
+	status: CodemodeNestedCallStatus;
+	durationMs?: number;
+	/** Error text, truncated for display. */
+	error?: string;
+	/** Cost in USD of a `models.*` call that reported usage. */
+	cost?: number;
+}
+
+export interface CodemodeToolDetails {
+	calls: CodemodeNestedCall[];
+	/** Temp file with the full text output, when the output was truncated. */
+	fullOutputPath?: string;
+}
+
+export const codemodeToolSystemPromptContribution = {
+	snippet: "Run JavaScript that calls other tools",
+	guidelines: [
+		"Use codemode to batch independent tool calls (Promise.allSettled), chain them, or filter large output, instead of many separate calls.",
+	],
+} as const;
+
+/** The reference for scripts: globals, tool results, `store()`, the `models` API, and limits. */
+export const CODEMODE_DOCS_PATH = join(getDocsPath(), "codemode.md");
+
+const DESCRIPTION_INTRO = `Run JavaScript that calls other tools. The input is raw JavaScript (not JSON, no code fence), run as an async function body in a QuickJS sandbox: top-level \`await\` and \`return\` work. No Node, file system, network, or timers.
+- \`await tools.<name>({ ...args })\` resolves to a string, or an object if the tool's declaration says so, and rejects with an Error on failure. A tool whose declaration returns an object resolves on failure too when it returned that object; check its fields. Calls still running when the script ends are cancelled.
+- Optional first line: \`// @options: {"max_output_tokens": 10000, "timeout_ms": 60000}\``;
+
+/** One line per global. The details live in {@link CODEMODE_DOCS_PATH}. */
+function describeGlobals(models: boolean): string {
+	const lines = [
+		"Globals:",
+		"- `text(value)`, `image(dataUrlOrImageBlock)`, `console.log(...)`, and top-level `return` add output; `exit()` ends the script. With several text items, each starts with a `==> text N/M <==` line, and `console` lines follow the other output in one `<console_output>` block. `image()` also saves the image to a temp file and the result names its path.",
+		"- `store(key, value)` and `load(key)` keep JSON values across codemode calls.",
+		"- `ALL_TOOLS`, `await searchTools(query, { limit?, namespace? })`, `await describeTool(name)`, `await describeNamespace(name)`: find unlisted tools, such as MCP tools.",
+	];
+	if (models) {
+		lines.push(`- \`models\`: classifiers and image generation. Read ${CODEMODE_DOCS_PATH} first.`);
+	}
+	return lines.join("\n");
+}
+
+/** Default for {@link CodemodeDescriptionOptions.inlineBudget}, in estimated tokens. */
+export const DEFAULT_CODEMODE_INLINE_BUDGET = SETTINGS_DEFAULTS.codemode.inlineBudget;
+/** Characters per token when estimating the cost of a tool section. */
+const CHARS_PER_TOKEN = 4;
+
+/**
+ * What a script sees of a tool: its description followed by its prompt guidelines, which the system
+ * prompt only has for declared tools. Tools without an output schema resolve to their text output.
+ */
+export function toCodemodeDeclaration(
+	tool: AgentTool<any>,
+	guidelines: readonly string[] = [],
+): Omit<CodemodeTool, "execute"> {
+	const bullets = guidelines.flatMap((guideline) => (guideline.trim() ? [`- ${guideline.trim()}`] : []));
+	return {
+		name: tool.name,
+		description: bullets.length > 0 ? `${tool.description.trim()}\n\n${bullets.join("\n")}` : tool.description,
+		inputSchema: tool.parameters as CodemodeJsonSchema,
+		outputSchema: (tool.outputSchema as CodemodeJsonSchema | undefined) ?? TEXT_OUTPUT_SCHEMA,
+	};
+}
+
+/** Tools a script may call: every given tool except the codemode tool itself. */
+export function getCodemodeCallableTools(tools: readonly AgentTool<any>[]): AgentTool<any>[] {
+	return tools.filter((tool) => tool.name !== CODEMODE_TOOL_NAME);
+}
+
+export interface CodemodeDescriptionOptions {
+	/** Declare the `models` namespace; only for tools created with model access. */
+	models?: boolean;
+	/** Namespace of each tool, by tool name. Tools of one namespace are listed under one heading. */
+	namespaces?: ReadonlyMap<string, ToolNamespace>;
+	/** Tools that are callable but never listed with their declaration (`deferred` exposure). */
+	deferred?: ReadonlySet<string>;
+	/** Prompt guidelines of each tool, by tool name, listed after its description. */
+	guidelines?: ReadonlyMap<string, readonly string[]>;
+	/**
+	 * Estimated tokens (characters / 4) the tool sections may use. Tools that do not fit are left
+	 * out, like deferred tools. Unset lists every tool that is not deferred.
+	 */
+	inlineBudget?: number;
+}
+
+/** `### \`id\` (\`raw name\`)` followed by the tool's description and declaration. */
+function renderToolSection(declaration: Omit<CodemodeTool, "execute">): string {
+	const id = toCodemodeIdentifier(declaration.name);
+	const heading = id === declaration.name ? `### \`${id}\`` : `### \`${id}\` (\`${declaration.name}\`)`;
+	return `${heading}\n${renderToolSample(declaration).trim()}`;
+}
+
+interface CatalogEntry {
+	name: string;
+	section: string;
+	cost: number;
+}
+
+interface CatalogGroup {
+	namespace: ToolNamespace | undefined;
+	entries: CatalogEntry[];
+}
+
+/**
+ * Pick the tool sections that fit the budget, like OpenCode's catalog: in each round every group
+ * (tools without a namespace first, then namespaces by name) places its cheapest remaining tool; a
+ * group whose next tool does not fit drops out while the others continue. Every namespace is
+ * represented before any namespace is complete.
+ */
+function selectCatalog(groups: readonly CatalogGroup[], budget: number | undefined): Set<string> {
+	if (budget === undefined) return new Set(groups.flatMap((group) => group.entries.map((entry) => entry.name)));
+	const queues = groups.map((group) => [...group.entries].sort((a, b) => a.cost - b.cost));
+	const shown = new Set<string>();
+	let remaining = budget;
+	let active = queues.filter((queue) => queue.length > 0);
+	while (active.length > 0) {
+		active = active.filter((queue) => {
+			const next = queue[0];
+			if (next.cost > remaining) return false;
+			remaining -= next.cost;
+			shown.add(next.name);
+			queue.shift();
+			return queue.length > 0;
+		});
+	}
+	return shown;
+}
+
+/**
+ * Model-facing description: the helper list, guidance for finding tools that are not listed, the
+ * `models` API, and one section per listed tool, grouped by namespace. Deferred tools are never
+ * listed and do not affect the description at all, so it stays the same while tools are registered
+ * late (for example by MCP servers that connect). Tool sections are limited to `inlineBudget`.
+ */
+export function createCodemodeDescription(
+	tools: readonly AgentTool<any>[],
+	options: CodemodeDescriptionOptions = {},
+): string {
+	const declarations = getCodemodeCallableTools(tools)
+		.filter((tool) => !options.deferred?.has(tool.name))
+		.map((tool) => toCodemodeDeclaration(tool, options.guidelines?.get(tool.name)));
+	const groups = new Map<string, CatalogGroup>([["", { namespace: undefined, entries: [] }]]);
+	for (const declaration of declarations) {
+		const namespace = options.namespaces?.get(declaration.name);
+		const key = namespace ? `ns:${namespace.name}` : "";
+		const group = groups.get(key) ?? { namespace, entries: [] };
+		groups.set(key, group);
+		const section = renderToolSection(declaration);
+		group.entries.push({ name: declaration.name, section, cost: Math.ceil(section.length / CHARS_PER_TOKEN) });
+	}
+	const ordered = [...groups.values()].sort((a, b) =>
+		a.namespace === undefined ? -1 : b.namespace === undefined ? 1 : a.namespace.name.localeCompare(b.namespace.name),
+	);
+	const shown = selectCatalog(ordered, options.inlineBudget);
+
+	const sections = [DESCRIPTION_INTRO, describeGlobals(options.models === true)];
+	if (declarations.length === 0) return sections.join("\n\n");
+
+	const toolSections = ["Nested tools:"];
+	for (const { namespace, entries } of ordered) {
+		const visible = entries.filter((entry) => shown.has(entry.name));
+		if (namespace) {
+			// Only tools that did not fit the budget are counted as not listed here.
+			const listing =
+				visible.length === entries.length
+					? ""
+					: visible.length === 0
+						? " (tools not listed)"
+						: " (some tools not listed)";
+			const description = namespace.description?.trim();
+			toolSections.push(`## ${namespace.name}${listing}${description ? `\n${description}` : ""}`);
+		}
+		for (const entry of visible) toolSections.push(entry.section);
+	}
+	sections.push(toolSections.join("\n\n"));
+	return sections.join("\n\n");
+}
+
+/**
+ * The field names of an object schema, `{ output, exit_code, full_output_path? }`, with the fields of
+ * object-typed fields `depth` levels deep: `{ content, structuredContent?: { results } }`. Undefined
+ * for other schemas.
+ */
+function describeFields(schema: unknown, depth: number): string | undefined {
+	if (typeof schema !== "object" || schema === null) return undefined;
+	const { type, properties, required } = schema as Record<string, unknown>;
+	if (type !== "object" || typeof properties !== "object" || properties === null) return undefined;
+	const requiredNames = new Set(Array.isArray(required) ? required : []);
+	const fields = Object.entries(properties).map(([name, field]) => {
+		const nested = depth > 0 ? describeFields(field, depth - 1) : undefined;
+		return `${name}${requiredNames.has(name) ? "" : "?"}${nested ? `: ${nested}` : ""}`;
+	});
+	return `{ ${fields.join(", ")} }`;
+}
+
+/**
+ * What a script call resolves to, in one line: `a string`, the field names of an object (see
+ * {@link describeFields}), or the rendered type for anything else.
+ */
+function describeOutput(schema: CodemodeJsonSchema | undefined): string {
+	const type = renderToolOutputType(schema);
+	if (type === "string") return "a string";
+	const fields = describeFields(schema, 1);
+	return `\`${fields ?? type.replace(/\s+/g, " ")}\``;
+}
+
+/**
+ * A declared tool's description followed by how scripts call it and what the call resolves to. The
+ * arguments are the tool's declared parameters, so they are not repeated.
+ */
+function describeScriptCall(tool: AgentTool<any>): string {
+	return `${tool.description.trim()}\n\nCodemode: \`tools.${toCodemodeIdentifier(tool.name)}(args)\` resolves to ${describeOutput(toCodemodeDeclaration(tool).outputSchema)}.`;
+}
+
+/**
+ * How the codemode tool presents tools that are both declared and callable from scripts:
+ * - `on`: their descriptions say how scripts call them, and the codemode description
+ *   lists only the callable tools without `direct` exposure.
+ * - `only`: the codemode description lists every callable tool, and requests leave out the
+ *   declarations of active `direct` tools.
+ *
+ * Listed tools carry their prompt guidelines, which the system prompt only has for declared tools.
+ *
+ * Listing by exposure, not by the active set, keeps the codemode description unchanged when
+ * `tool_search` loads a tool, so loads do not redeclare codemode.
+ */
+function prepareCodemodeLoadout(loadout: ToolLoadout, options: CodemodeToolOptions): ToolLoadoutChanges {
+	const mode = options.getMode?.() ?? SETTINGS_DEFAULTS.codemode.mode;
+	const isDirect = (tool: AgentTool) => loadout.getExposure(tool.name) === "direct";
+	const callable = getCodemodeCallableTools(loadout.callable);
+	const callableNames = new Set(callable.map((tool) => tool.name));
+	const descriptions: Record<string, string> = {};
+	if (mode === "on") {
+		for (const tool of loadout.declared) {
+			if (callableNames.has(tool.name)) descriptions[tool.name] = describeScriptCall(tool);
+		}
+	}
+	const listed = mode === "only" ? callable : callable.filter((tool) => !isDirect(tool));
+	const namespaces = new Map(
+		listed.flatMap((tool) => {
+			const namespace = loadout.getNamespace(tool.name);
+			return namespace ? [[tool.name, namespace] as const] : [];
+		}),
+	);
+	const guidelines = new Map(listed.map((tool) => [tool.name, loadout.getPromptGuidelines(tool.name)] as const));
+	descriptions[CODEMODE_TOOL_NAME] = createCodemodeDescription(listed, {
+		models: options.models === true,
+		namespaces,
+		guidelines,
+		deferred: new Set(
+			listed.filter((tool) => loadout.getExposure(tool.name) === "deferred").map((tool) => tool.name),
+		),
+		inlineBudget: options.getInlineBudget?.() ?? DEFAULT_CODEMODE_INLINE_BUDGET,
+	});
+	const declaredNames = new Set(loadout.declared.map((tool) => tool.name));
+	return {
+		descriptions,
+		hiddenDeclarations:
+			mode === "only"
+				? callable.filter((tool) => isDirect(tool) && declaredNames.has(tool.name)).map((tool) => tool.name)
+				: [],
+	};
+}
+
+/** Script globals that list or look up tools and namespaces, which the script may then call by any name. */
+const DISCOVERY_GLOBALS = /\b(searchTools|describeNamespace|describeTool|ALL_TOOLS)\b/;
+
+/**
+ * The tools a script may call (`ToolDefinition.reachesTools`): every tool when it uses a discovery
+ * global, else the words of its source, so `tools.mcp__docs__search` and
+ * `tools["mcp__docs__" + name]` both name the `mcp__docs` namespace.
+ */
+function getScriptReach(code: unknown): "all" | readonly string[] {
+	if (typeof code !== "string") return [];
+	return DISCOVERY_GLOBALS.test(code) ? "all" : (code.match(/\w+/g) ?? []);
+}
+
+export function createCodemodeToolDefinition(
+	options: CodemodeToolOptions = {},
+): ToolDefinition<typeof codemodeSchema, CodemodeToolDetails | undefined> {
+	return {
+		name: CODEMODE_TOOL_NAME,
+		label: CODEMODE_TOOL_NAME,
+		// Replaced with the declarations of the callable tools when the tool is activated.
+		description: createCodemodeDescription([], { models: options.models === true }),
+		promptSnippet: codemodeToolSystemPromptContribution.snippet,
+		promptGuidelines: [...codemodeToolSystemPromptContribution.guidelines],
+		parameters: codemodeSchema,
+		// Scripts must not start other scripts.
+		exposure: "model-only",
+		prepareLoadout: (loadout) => prepareCodemodeLoadout(loadout, options),
+		reachesTools: (input) => getScriptReach(input.code),
+		// Capable models write the script as raw text instead of a JSON-escaped string.
+		constrainedSampling: { type: "grammar", variants: { openai_lark: CODEMODE_SOURCE_GRAMMAR } },
+		// The sandbox (worker, QuickJS wasm) loads on the first call, not at startup.
+		execute: async (toolCallId, params, signal, onUpdate, ctx) =>
+			(await loadCodemodeExecutor()).executeCodemode(toolCallId, params, signal, onUpdate, ctx, options),
+		...codemodeRenderers,
+	};
+}
+
+/**
+ * Create the codemode tool as an AgentTool. The description lists the given tools; the script can
+ * call whatever tools the agent loop provides at execution time.
+ */
+export function createCodemodeTool(
+	tools: readonly AgentTool<any>[] = [],
+	options: CodemodeToolOptions = {},
+): AgentTool<typeof codemodeSchema> {
+	const definition = createCodemodeToolDefinition(options);
+	const tool = wrapToolDefinition(definition);
+	Object.assign(tool, {
+		description: createCodemodeDescription(tools, { models: options.models === true }),
+		promptSnippet: definition.promptSnippet,
+		promptGuidelines: definition.promptGuidelines,
+	});
+	return tool;
+}

third_party/pi/packages/coding-agent/src/extensions/codemode/worker.ts

@@ -0,0 +1,6 @@
+/**
+ * Worker entry for the codemode sandbox in bundled builds. The Node bundle and the Bun binary
+ * build this file as a separate entrypoint because pi-codemode's own worker file is not on disk
+ * there; `getCodemodeWorkerSpecifier()` in config.ts resolves it.
+ */
+import "@earendil-works/pi-codemode/worker";

third_party/pi/packages/coding-agent/src/extensions/index.ts

@@ -1,4 +1,14 @@
 import type { InlineExtension } from "../core/extensions/types.ts";
+import codemodeExtension from "./codemode/index.ts";
 import llamaExtension from "./llama/index.ts";
+import mcpExtension from "./mcp/index.ts";
+import toolSearchExtension from "./tool-search/index.ts";
 
-export const builtInExtensions: InlineExtension[] = [{ name: "llama.cpp", factory: llamaExtension, hidden: true }];
+export const builtInExtensions: InlineExtension[] = [
+	{ name: "llama.cpp", factory: llamaExtension, builtin: true },
+	// Replaceable: an extension that registers `codemode`, `tool_search`, or `/mcp` (such as a third-party
+	// MCP extension) takes over instead of running alongside the built-in one.
+	{ name: "codemode", factory: codemodeExtension, replaceable: true, builtin: true },
+	{ name: "tool-search", factory: toolSearchExtension, replaceable: true, builtin: true },
+	{ name: "mcp", factory: mcpExtension, replaceable: true, builtin: true },
+];

third_party/pi/packages/coding-agent/src/extensions/llama/client.ts

@@ -30,6 +30,7 @@ export interface LlamaModelsResponse {
 
 export interface LlamaServerProps {
 	models_autoload?: boolean;
+	chat_template?: string;
 }
 
 export interface LlamaModelEvent {
@@ -193,11 +194,15 @@ export class LlamaClient {
 		return data;
 	}
 
-	async props(options: { signal?: AbortSignal } = {}): Promise<LlamaServerProps> {
-		const payload = await this.request("/props", { signal: options.signal });
+	async props(options: { model?: string; signal?: AbortSignal } = {}): Promise<LlamaServerProps> {
+		const query = options.model ? `?${new URLSearchParams({ model: options.model, autoload: "false" })}` : "";
+		const payload = await this.request(`/props${query}`, { signal: options.signal });
 		if (typeof payload !== "object" || payload === null) return {};
-		const { models_autoload: modelsAutoload } = payload as Record<string, unknown>;
-		return typeof modelsAutoload === "boolean" ? { models_autoload: modelsAutoload } : {};
+		const { models_autoload: modelsAutoload, chat_template: chatTemplate } = payload as Record<string, unknown>;
+		return {
+			...(typeof modelsAutoload === "boolean" ? { models_autoload: modelsAutoload } : {}),
+			...(typeof chatTemplate === "string" ? { chat_template: chatTemplate } : {}),
+		};
 	}
 
 	async load(model: string, signal?: AbortSignal): Promise<void> {

third_party/pi/packages/coding-agent/src/extensions/llama/provider.ts

@@ -1,14 +1,25 @@
-import type {
-	ApiKeyCredential,
-	AuthContext,
-	AuthResult,
-	Model,
-	Provider,
-	ProviderStreamOptions,
-	RefreshModelsContext,
+import {
+	type AnyModel,
+	type ApiKeyCredential,
+	type AuthContext,
+	type AuthResult,
+	type ClassifierModel,
+	isModelType,
+	type Model,
+	type Provider,
+	type ProviderStreamOptions,
+	type RefreshModelsContext,
 } from "@earendil-works/pi-ai";
+import { llamaCppClassifyApi } from "@earendil-works/pi-ai/api/llama-cpp-classify.lazy";
+import { typesafeSystemOneApi } from "@earendil-works/pi-ai/api/typesafe-system-one.lazy";
 import { stream, streamSimple } from "@earendil-works/pi-ai/compat";
-import { LlamaClient, type LlamaModelInfo, llamaInferenceUrl, normalizeLlamaServerUrl } from "./client.ts";
+import {
+	LlamaClient,
+	type LlamaModelInfo,
+	type LlamaServerProps,
+	llamaInferenceUrl,
+	normalizeLlamaServerUrl,
+} from "./client.ts";
 
 export const LLAMA_PROVIDER_ID = "llama.cpp";
 export const DEFAULT_LLAMA_SERVER_URL = "http://127.0.0.1:8080";
@@ -46,16 +57,91 @@ async function routerAutoloadEnabled(
 	}
 }
 
-function toPiModel(model: LlamaModelInfo, serverUrl: string): Model<"openai-completions"> {
-	const reportedContextWindow = model.meta?.n_ctx ?? model.meta?.n_ctx_train;
-	const contextWindow = reportedContextWindow && reportedContextWindow > 0 ? reportedContextWindow : 128000;
+function configuredContextWindow(model: LlamaModelInfo): number | undefined {
+	const args = model.status.args ?? [];
+	for (let index = 0; index < args.length - 1; index++) {
+		const flag = args[index];
+		if (flag !== "--ctx-size" && flag !== "-c" && flag !== "-ctx") continue;
+		const contextWindow = Number(args[index + 1]);
+		if (Number.isSafeInteger(contextWindow) && contextWindow > 0) return contextWindow;
+	}
+	return undefined;
+}
+
+function contextWindowOf(model: LlamaModelInfo, cachedContextWindow?: number): number {
+	const runtimeContextWindow = model.meta?.n_ctx;
+	if (runtimeContextWindow && runtimeContextWindow > 0) return runtimeContextWindow;
+	const configuredContext = configuredContextWindow(model);
+	if (configuredContext) return configuredContext;
+	if (cachedContextWindow && cachedContextWindow > 0) return cachedContextWindow;
+	const trainingContextWindow = model.meta?.n_ctx_train;
+	return trainingContextWindow && trainingContextWindow > 0 ? trainingContextWindow : 128000;
+}
+
+type LlamaClassifierApi = "llama-cpp-classify" | "typesafe-system-one";
+
+/**
+ * Whether llama.cpp reports a native decision model. Since llama.cpp 0.6.0, `GET /models` lists
+ * `decisions` in `architecture.output_modalities` for these models, including unloaded and sleeping ones.
+ * Older servers report `["text"]` or omit `architecture`, so their models are treated as chat models.
+ */
+function isDecisionModel(model: LlamaModelInfo): boolean {
+	return model.architecture?.output_modalities?.includes("decisions") === true;
+}
+
+/** Decision-only models cannot generate text and are not listed for chat. */
+function isChatModel(model: LlamaModelInfo): boolean {
+	return !isDecisionModel(model) || model.architecture?.output_modalities?.includes("text") === true;
+}
+
+/**
+ * A llama.cpp model used as a classifier. Decision models answer natively through llama.cpp's
+ * System One endpoint (`/v1/systemone`). Chat models fall back to `llama-cpp-classify`, which reads
+ * answers from next-token label probabilities.
+ */
+function toPiClassifierModel(
+	model: LlamaModelInfo,
+	serverUrl: string,
+	cachedContextWindow?: number,
+): ClassifierModel<LlamaClassifierApi> {
+	const decision = isDecisionModel(model);
+	return {
+		type: "classifier",
+		id: model.id,
+		name: model.id,
+		api: decision ? "typesafe-system-one" : "llama-cpp-classify",
+		provider: LLAMA_PROVIDER_ID,
+		baseUrl: decision ? llamaInferenceUrl(serverUrl) : serverUrl,
+		input: ["text"],
+		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+		contextWindow: contextWindowOf(model, cachedContextWindow),
+	};
+}
+
+function isLlamaClassifierModel(model: AnyModel): model is ClassifierModel<LlamaClassifierApi> {
+	return (
+		isModelType(model, "classifier") && (model.api === "llama-cpp-classify" || model.api === "typesafe-system-one")
+	);
+}
+
+function toPiModel(
+	model: LlamaModelInfo,
+	serverUrl: string,
+	props?: LlamaServerProps,
+	cachedContextWindow?: number,
+): Model<"openai-completions"> {
+	const contextWindow = contextWindowOf(model, cachedContextWindow);
+	const reasoning = props?.chat_template?.includes("enable_thinking") === true;
 	return {
 		id: model.id,
 		name: model.id,
 		api: "openai-completions",
 		provider: LLAMA_PROVIDER_ID,
 		baseUrl: llamaInferenceUrl(serverUrl),
-		reasoning: false,
+		reasoning,
+		...(reasoning && {
+			thinkingLevelMap: { off: "off", minimal: null, low: null, medium: "medium", high: null, xhigh: null },
+		}),
 		input: model.architecture?.input_modalities?.includes("image") ? ["text", "image"] : ["text"],
 		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
 		contextWindow,
@@ -67,6 +153,7 @@ function toPiModel(model: LlamaModelInfo, serverUrl: string): Model<"openai-comp
 			supportsUsageInStreaming: true,
 			supportsStrictMode: false,
 			maxTokensField: "max_tokens",
+			...(reasoning && { thinkingFormat: "qwen-chat-template" }),
 		},
 	};
 }
@@ -78,15 +165,18 @@ export interface LlamaProviderController {
 
 export function createLlamaProvider(): LlamaProviderController {
 	let models: readonly Model<"openai-completions">[] = [];
+	let classifiers: readonly ClassifierModel<LlamaClassifierApi>[] = [];
+	const fallbackClassifier = llamaCppClassifyApi();
+	const decisionClassifier = typesafeSystemOneApi();
 
 	const setCatalog = (
 		catalog: readonly LlamaModelInfo[],
 		serverUrl: string,
 		options: { routerAutoload?: boolean } = {},
 	): void => {
-		models = catalog
-			.filter((model) => modelIsSelectable(model, options.routerAutoload === true))
-			.map((model) => toPiModel(model, serverUrl));
+		const selectable = catalog.filter((model) => modelIsSelectable(model, options.routerAutoload === true));
+		models = selectable.filter(isChatModel).map((model) => toPiModel(model, serverUrl));
+		classifiers = selectable.map((model) => toPiClassifierModel(model, serverUrl));
 	};
 
 	const provider: Provider<"openai-completions"> = {
@@ -137,16 +227,24 @@ export function createLlamaProvider(): LlamaProviderController {
 			},
 		},
 		getModels: () => models,
+		getAllModels: () => [...models, ...classifiers],
 		refreshModels: async (context: RefreshModelsContext): Promise<void> => {
+			const cachedContextWindows = new Map<string, number>();
 			if (context.stored) {
-				const restored = context.stored.models.filter(
+				const stored = context.stored.models.filter((model) => model.provider === LLAMA_PROVIDER_ID);
+				const restored = stored.filter(
 					(model): model is Model<"openai-completions"> =>
-						model.provider === LLAMA_PROVIDER_ID && model.api === "openai-completions",
+						isModelType(model, "chat") && model.api === "openai-completions",
 				);
+				const restoredClassifiers = stored.filter(isLlamaClassifierModel);
+				for (const model of [...restored, ...restoredClassifiers]) {
+					cachedContextWindows.set(model.id, model.contextWindow);
+				}
 				if (
 					!(await context.publish({
 						update: () => {
 							models = restored;
+							classifiers = restoredClassifiers;
 						},
 					}))
 				) {
@@ -162,18 +260,39 @@ export function createLlamaProvider(): LlamaProviderController {
 			if (context.signal.aborted) return;
 			const routerAutoload = await routerAutoloadEnabled(client, catalog, context.signal);
 			if (context.signal.aborted) return;
-			const refreshed = catalog
-				.filter((model) => modelIsSelectable(model, routerAutoload))
-				.map((model) => toPiModel(model, serverUrl));
+			const selectable = catalog.filter((model) => modelIsSelectable(model, routerAutoload));
+			// Only loaded models expose their chat template without side effects. Unloaded autoload presets would
+			// need to be loaded, while querying sleeping models may wake them. Those models remain without thinking
+			// support until they are loaded and a later catalog refresh discovers it. Decision models need no
+			// template, and llama.cpp reports them in the catalog regardless of their status.
+			const refreshed = await Promise.all(
+				selectable.filter(isChatModel).map(async (model) => {
+					const cachedContextWindow = cachedContextWindows.get(model.id);
+					if (model.status.value !== "loaded") return toPiModel(model, serverUrl, undefined, cachedContextWindow);
+					const props = await client.props({ model: model.id, signal: context.signal });
+					return toPiModel(model, serverUrl, props, cachedContextWindow);
+				}),
+			);
+			const refreshedClassifiers = selectable.map((model) =>
+				toPiClassifierModel(model, serverUrl, cachedContextWindows.get(model.id)),
+			);
+			if (context.signal.aborted) return;
 			await context.publish({
-				persist: { models: refreshed, checkedAt: Date.now() },
+				persist: { models: [...refreshed, ...refreshedClassifiers], checkedAt: Date.now() },
 				update: () => {
 					models = refreshed;
+					classifiers = refreshedClassifiers;
 				},
 			});
 		},
 		stream: (model, context, options) => stream(model, context, options as ProviderStreamOptions | undefined),
 		streamSimple: (model, context, options) => streamSimple(model, context, options),
+		classify: (model, context, options) =>
+			(model.api === "typesafe-system-one" ? decisionClassifier : fallbackClassifier).classify(
+				model,
+				context,
+				options,
+			),
 	};
 
 	return { provider, setCatalog };

third_party/pi/packages/coding-agent/src/extensions/mcp/cli.lazy.ts

@@ -0,0 +1,2 @@
+/** Loads `pi mcp` and the MCP runtime only when the command runs (see cli.ts). */
+export const loadMcpCommand = () => import("./cli.ts");

third_party/pi/packages/coding-agent/src/extensions/mcp/cli.ts

@@ -0,0 +1,594 @@
+/**
+ * `pi mcp`: add, remove, and check MCP servers and sign in to them outside a session. Agents run it
+ * through bash to configure servers, verify an `mcp.json` they wrote, and start an OAuth sign-in;
+ * the user only approves access in the browser. Running sessions pick up new credentials on their
+ * next turn.
+ */
+
+import { existsSync } from "node:fs";
+import { join } from "node:path";
+import { createInterface } from "node:readline/promises";
+import chalk from "chalk";
+import { APP_NAME, CONFIG_DIR_NAME } from "../../config.ts";
+import { validateMcpServerConfig } from "../../core/mcp-servers.ts";
+import { ProjectTrustStore } from "../../core/trust-manager.ts";
+import { openBrowser } from "../../utils/open-browser.ts";
+import {
+	addMcpServerConfig,
+	getMcpToolExposure,
+	type LoadedMcpConfig,
+	loadMcpConfig,
+	type McpServerEntry,
+	removeMcpServerConfig,
+} from "./config.ts";
+import {
+	createDefaultTransport,
+	McpOAuthCredentialStore,
+	McpServerConnection,
+	McpServerLog,
+	McpSignInCancelledError,
+	signInMcpServer,
+} from "./runtime.ts";
+
+const HELP = `${chalk.bold("Usage:")}
+  ${APP_NAME} mcp add <server> [options] -- <command> [args...]
+  ${APP_NAME} mcp add <server> [options] --url <url>
+  ${APP_NAME} mcp remove <server> [-l]
+  ${APP_NAME} mcp list [--json]
+  ${APP_NAME} mcp login <server> [--timeout <seconds>]
+  ${APP_NAME} mcp logout <server>
+
+Configure and check MCP servers and sign in to OAuth servers without starting a session.
+Reads ~/${CONFIG_DIR_NAME}/agent/mcp.json and, in trusted projects, ${CONFIG_DIR_NAME}/mcp.json.
+
+Commands:
+  add <server>            Add or replace a server in mcp.json
+  remove <server>         Remove a server from mcp.json
+  list                    Show state, tools, and errors (exits 1 on failure)
+  login <server>          Sign in through the browser
+  logout <server>         Delete the stored OAuth credentials
+
+Options for add and remove:
+  -l, --local             Use ${CONFIG_DIR_NAME}/mcp.json in the current project instead of the global file
+
+Options for add:
+  --url <url>             Streamable HTTP server URL (instead of a command)
+  --env <KEY=VALUE>       Environment variable for a stdio server (repeatable)
+  --cwd <dir>             Working directory for a stdio server
+  --header <KEY=VALUE>    HTTP header (repeatable)
+  --bearer-token-env-var <NAME>
+                          Send "Authorization: Bearer \${NAME}"
+  --oauth-client-id <id>  Pre-registered OAuth client id
+  --oauth-client-secret <secret>
+                          OAuth client secret (may be \${NAME} or !command)
+  --oauth-callback-port <port>
+                          Fixed OAuth callback port
+  --oauth-client-name <name>
+                          Client name sent when registering with the OAuth server
+  --exposure <mode>       codemode (default), deferred, direct, or hidden
+  --description <text>    What the server offers, shown in the system prompt
+
+Other options:
+  --json                  Print the list as JSON
+  --timeout <seconds>     Time limit for the whole sign-in (default: 300)`;
+
+const HELP_HINT = chalk.dim(`Use "${APP_NAME} mcp --help" for usage.`);
+
+const DEFAULT_LOGIN_TIMEOUT_SECONDS = 300;
+
+export interface McpCommandOptions {
+	cwd: string;
+	agentDir: string;
+	/** Defaults to `mcp-auth.json` in the agent directory. */
+	credentials?: McpOAuthCredentialStore;
+	/** Defaults to the platform browser. */
+	openUrl?: (url: string) => void;
+	/** Defaults to console output. */
+	log?: (line: string) => void;
+	error?: (line: string) => void;
+}
+
+interface ServerReport {
+	name: string;
+	scope: string;
+	source: string;
+	/** Project `mcp.json` that overrides `enabled`, `exposure`, or `toolExposure` of this global server. */
+	override?: string;
+	enabled: boolean;
+	exposure: string;
+	transport: string;
+	state: string;
+	tools: string[];
+	/** Tools whose exposure differs from the server's, from `toolExposure`. */
+	toolExposure?: Record<string, string>;
+	/** Whether the server offers resources. They are not listed: listing every one can take long. */
+	resources?: true;
+	error?: string;
+}
+
+function errorMessage(error: unknown): string {
+	return error instanceof Error ? error.message : String(error);
+}
+
+function describeTransport(entry: McpServerEntry): string {
+	const { config } = entry;
+	return "url" in config ? config.url : [config.command, ...(config.args ?? [])].join(" ");
+}
+
+function createConnection(entry: McpServerEntry, options: McpCommandOptions, credentials: McpOAuthCredentialStore) {
+	return new McpServerConnection({
+		entry,
+		cwd: options.cwd,
+		createTransport: createDefaultTransport,
+		credentials,
+		log: new McpServerLog(join(options.agentDir, "mcp.log")),
+		onTools: () => {},
+	});
+}
+
+/** Short spellings of options. `-l`/`--local` match `pi install`. */
+const OPTION_ALIASES = new Map([["-l", "--local"]]);
+
+interface ParsedOptions {
+	positional: string[];
+	values: Map<string, string | true>;
+	/** Values of `list` options, in order. */
+	lists: Map<string, string[]>;
+}
+
+/**
+ * Parse `--name value` options; returns undefined and reports unknown ones. `--` ends the options,
+ * as does reaching `maxPositionals` positional arguments: the remaining arguments are positional,
+ * so a command's own options (`add <server> <command> --flag`) are passed through.
+ */
+function parseOptions(
+	args: string[],
+	known: Record<string, "flag" | "value" | "list">,
+	error: (line: string) => void,
+	maxPositionals = Number.POSITIVE_INFINITY,
+): ParsedOptions | undefined {
+	const positional: string[] = [];
+	const values = new Map<string, string | true>();
+	const lists = new Map<string, string[]>();
+	for (let index = 0; index < args.length; index++) {
+		const arg = OPTION_ALIASES.get(args[index]) ?? args[index];
+		if (arg === "--" || positional.length >= maxPositionals) {
+			positional.push(...args.slice(arg === "--" ? index + 1 : index));
+			break;
+		}
+		if (!arg.startsWith("--")) {
+			positional.push(arg);
+			continue;
+		}
+		const name = arg.slice(2);
+		const kind = known[name];
+		if (!kind) {
+			error(`Unknown option ${arg}.\n${HELP_HINT}`);
+			return undefined;
+		}
+		if (kind === "flag") {
+			values.set(name, true);
+			continue;
+		}
+		const value = args[++index];
+		if (value === undefined) {
+			error(`${arg} needs a value.`);
+			return undefined;
+		}
+		if (kind === "list") lists.set(name, [...(lists.get(name) ?? []), value]);
+		else values.set(name, value);
+	}
+	return { positional, values, lists };
+}
+
+/** Run `pi mcp <args>` and return the exit code. */
+export async function runMcpCommand(args: string[], options: McpCommandOptions): Promise<number> {
+	const log = options.log ?? ((line: string) => console.log(line));
+	const error = options.error ?? ((line: string) => console.error(line));
+	const [command, ...rest] = args;
+	if (command === undefined || command === "help" || args.includes("--help") || args.includes("-h")) {
+		log(HELP);
+		return 0;
+	}
+
+	const projectConfig = join(options.cwd, CONFIG_DIR_NAME, "mcp.json");
+	if (command === "add" || command === "remove") {
+		return command === "add"
+			? add(rest, projectConfig, options, log, error)
+			: remove(rest, projectConfig, options, log, error);
+	}
+	const projectTrusted = new ProjectTrustStore(options.agentDir).get(options.cwd) === true;
+	const loaded = loadMcpConfig({ agentDir: options.agentDir, cwd: options.cwd, projectTrusted });
+	const untrustedNote =
+		!projectTrusted && existsSync(projectConfig)
+			? `${projectConfig} is ignored because the project is not trusted. Start ${APP_NAME} in the project to trust it.`
+			: undefined;
+	const credentials = options.credentials ?? new McpOAuthCredentialStore();
+
+	switch (command) {
+		case "list": {
+			const parsed = parseOptions(rest, { json: "flag" }, error);
+			if (!parsed) return 1;
+			if (parsed.positional.length > 0) {
+				error(`Usage: ${APP_NAME} mcp list [--json]\n${HELP_HINT}`);
+				return 1;
+			}
+			return list(loaded, parsed.values.has("json"), untrustedNote, options, credentials, log);
+		}
+		case "login":
+		case "logout": {
+			const parsed = parseOptions(rest, command === "login" ? { timeout: "value" } : {}, error);
+			if (!parsed) return 1;
+			const [name, ...extra] = parsed.positional;
+			if (!name || extra.length > 0) {
+				error(`Usage: ${APP_NAME} mcp ${command} <server>\n${HELP_HINT}`);
+				return 1;
+			}
+			const entry = loaded.servers.find((server) => server.name === name);
+			if (!entry) {
+				error(
+					`No MCP server named "${name}".${untrustedNote ? ` ${untrustedNote}` : ""} Configured: ${loaded.servers.map((server) => server.name).join(", ") || "none"}.`,
+				);
+				return 1;
+			}
+			const connection = createConnection(entry, options, credentials);
+			const url = connection.oauthUrl;
+			if (!url) {
+				error(`MCP server "${name}" does not use OAuth. Only HTTP servers without an Authorization header do.`);
+				return 1;
+			}
+			if (command === "logout") {
+				const removed = credentials.remove(name, url);
+				log(removed ? `Signed out of MCP server "${name}".` : `No stored credentials for MCP server "${name}".`);
+				return 0;
+			}
+			const timeout = Number(parsed.values.get("timeout") ?? DEFAULT_LOGIN_TIMEOUT_SECONDS);
+			if (!Number.isFinite(timeout) || timeout <= 0) {
+				error("--timeout must be a positive number of seconds.");
+				return 1;
+			}
+			try {
+				return await login(entry, connection, url, timeout * 1000, options, credentials, log, error);
+			} finally {
+				await connection.close();
+			}
+		}
+		default:
+			error(`Unknown mcp command "${command}".\n${HELP_HINT}`);
+			return 1;
+	}
+}
+
+/** Parse `KEY=VALUE` pairs of a repeatable option into a record. */
+function parsePairs(option: string, pairs: string[] | undefined, error: (line: string) => void) {
+	if (!pairs) return {};
+	const record: Record<string, string> = {};
+	for (const pair of pairs) {
+		const separator = pair.indexOf("=");
+		if (separator <= 0) {
+			error(`--${option} expects KEY=VALUE, got "${pair}".`);
+			return undefined;
+		}
+		record[pair.slice(0, separator)] = pair.slice(separator + 1);
+	}
+	return record;
+}
+
+function add(
+	args: string[],
+	projectConfig: string,
+	options: McpCommandOptions,
+	log: (line: string) => void,
+	error: (line: string) => void,
+): number {
+	const usage = `Usage: ${APP_NAME} mcp add <server> [options] (--url <url> | -- <command> [args...])\n${HELP_HINT}`;
+	const parsed = parseOptions(
+		args,
+		{
+			local: "flag",
+			url: "value",
+			env: "list",
+			cwd: "value",
+			header: "list",
+			"bearer-token-env-var": "value",
+			"oauth-client-id": "value",
+			"oauth-client-secret": "value",
+			"oauth-callback-port": "value",
+			"oauth-client-name": "value",
+			exposure: "value",
+			description: "value",
+		},
+		error,
+		2,
+	);
+	if (!parsed) return 1;
+	const { positional, values, lists } = parsed;
+	const [name, ...command] = positional;
+	const url = values.get("url");
+	if (!name || (url === undefined) === (command.length === 0)) {
+		error(usage);
+		return 1;
+	}
+	const value = (option: string) => {
+		const found = values.get(option);
+		return typeof found === "string" ? found : undefined;
+	};
+	const exposure = value("exposure");
+	const httpOnly = [
+		"header",
+		"bearer-token-env-var",
+		"oauth-client-id",
+		"oauth-client-secret",
+		"oauth-callback-port",
+		"oauth-client-name",
+	];
+	const stdioOnly = ["env", "cwd"];
+	const misplaced = (url === undefined ? httpOnly : stdioOnly).find(
+		(option) => values.has(option) || lists.has(option),
+	);
+	if (misplaced) {
+		error(`--${misplaced} only applies to ${url === undefined ? "HTTP servers (--url)" : "stdio servers"}.`);
+		return 1;
+	}
+
+	let config: Record<string, unknown>;
+	if (typeof url === "string") {
+		const headers = parsePairs("header", lists.get("header"), error);
+		if (!headers) return 1;
+		const bearer = value("bearer-token-env-var");
+		if (bearer !== undefined) headers.Authorization = `Bearer \${${bearer}}`;
+		const port = value("oauth-callback-port");
+		const oauth = {
+			...(value("oauth-client-id") === undefined ? {} : { clientId: value("oauth-client-id") }),
+			...(value("oauth-client-secret") === undefined ? {} : { clientSecret: value("oauth-client-secret") }),
+			...(port === undefined ? {} : { callbackPort: Number(port) }),
+			...(value("oauth-client-name") === undefined ? {} : { clientName: value("oauth-client-name") }),
+		};
+		config = {
+			url,
+			...(Object.keys(headers).length > 0 ? { headers } : {}),
+			...(Object.keys(oauth).length > 0 ? { oauth } : {}),
+		};
+	} else {
+		const env = parsePairs("env", lists.get("env"), error);
+		if (!env) return 1;
+		const [executable, ...commandArgs] = command;
+		config = {
+			command: executable,
+			...(commandArgs.length > 0 ? { args: commandArgs } : {}),
+			...(Object.keys(env).length > 0 ? { env } : {}),
+			...(value("cwd") === undefined ? {} : { cwd: value("cwd") }),
+		};
+	}
+	if (exposure !== undefined) config.exposure = exposure;
+	const description = value("description");
+	if (description !== undefined) config.description = description;
+	const validated = validateMcpServerConfig(name, config);
+	if (typeof validated === "string") {
+		error(validated);
+		return 1;
+	}
+
+	const project = values.has("local");
+	const path = project ? projectConfig : join(options.agentDir, "mcp.json");
+	const scope = project ? "project" : "global";
+	let replaced: boolean;
+	try {
+		replaced = addMcpServerConfig(path, name, validated);
+	} catch (addError) {
+		error(`Could not update ${path}: ${errorMessage(addError)}`);
+		return 1;
+	}
+	log(`${replaced ? "Replaced" : "Added"} ${scope} MCP server "${name}" in ${path}.`);
+	if (project && new ProjectTrustStore(options.agentDir).get(options.cwd) !== true) {
+		log(`The project is not trusted, so ${path} is ignored until you start ${APP_NAME} in the project and trust it.`);
+	}
+	// HTTP servers without an Authorization header may use OAuth.
+	const mayNeedSignIn =
+		"url" in validated &&
+		!Object.keys(validated.headers ?? {}).some((header) => header.toLowerCase() === "authorization");
+	log(
+		`Check it with: ${APP_NAME} mcp list${mayNeedSignIn ? `. If it requires sign-in: ${APP_NAME} mcp login ${name}` : ""}`,
+	);
+	return 0;
+}
+
+function remove(
+	args: string[],
+	projectConfig: string,
+	options: McpCommandOptions,
+	log: (line: string) => void,
+	error: (line: string) => void,
+): number {
+	const parsed = parseOptions(args, { local: "flag" }, error);
+	if (!parsed) return 1;
+	const [name, ...extra] = parsed.positional;
+	if (!name || extra.length > 0) {
+		error(`Usage: ${APP_NAME} mcp remove <server> [-l]\n${HELP_HINT}`);
+		return 1;
+	}
+	const project = parsed.values.has("local");
+	const globalConfig = join(options.agentDir, "mcp.json");
+	const path = project ? projectConfig : globalConfig;
+	const scope = project ? "project" : "global";
+	let removed: boolean;
+	try {
+		removed = removeMcpServerConfig(path, name);
+	} catch (removeError) {
+		error(`Could not update ${path}: ${errorMessage(removeError)}`);
+		return 1;
+	}
+	if (removed) {
+		log(`Removed ${scope} MCP server "${name}" from ${path}.`);
+		return 0;
+	}
+	const other = loadMcpConfig({ agentDir: options.agentDir, cwd: options.cwd, projectTrusted: true }).servers.find(
+		(server) => server.name === name && server.scope !== scope,
+	);
+	error(
+		`No ${scope} MCP server named "${name}" in ${path}.${other ? ` It is defined in ${other.source}${other.scope === "project" ? "; use --local" : "; omit --local"}.` : ""}`,
+	);
+	return 1;
+}
+
+async function list(
+	loaded: LoadedMcpConfig,
+	json: boolean,
+	untrustedNote: string | undefined,
+	options: McpCommandOptions,
+	credentials: McpOAuthCredentialStore,
+	log: (line: string) => void,
+): Promise<number> {
+	const reports = await Promise.all(
+		loaded.servers.map(async (entry): Promise<ServerReport> => {
+			const report: ServerReport = {
+				name: entry.name,
+				scope: entry.scope ?? "global",
+				source: entry.source,
+				...(entry.override ? { override: entry.override } : {}),
+				enabled: entry.config.enabled !== false,
+				exposure: entry.config.exposure ?? "codemode",
+				transport: describeTransport(entry),
+				state: "disabled",
+				tools: [],
+			};
+			if (!report.enabled) return report;
+			const connection = createConnection(entry, options, credentials);
+			try {
+				await connection.getClient();
+			} catch {
+				// The connection records the state and error.
+			}
+			report.state = connection.state;
+			report.tools = connection.tools.map((tool) => tool.name);
+			const overrides = connection.tools.flatMap((tool) => {
+				const exposure = getMcpToolExposure(entry.config, tool.name);
+				return exposure === report.exposure ? [] : [[tool.name, exposure] as const];
+			});
+			if (overrides.length > 0) report.toolExposure = Object.fromEntries(overrides);
+			if (connection.hasResources) report.resources = true;
+			if (connection.state !== "connected" && connection.error) report.error = connection.error;
+			await connection.close();
+			return report;
+		}),
+	);
+	const failed = loaded.errors.length > 0 || reports.some((report) => report.enabled && report.state !== "connected");
+
+	if (json) {
+		log(
+			JSON.stringify(
+				{ servers: reports, errors: loaded.errors, ...(untrustedNote ? { note: untrustedNote } : {}) },
+				null,
+				2,
+			),
+		);
+		return failed ? 1 : 0;
+	}
+	if (reports.length === 0 && loaded.errors.length === 0) {
+		log(`No MCP servers configured. Add them to ${join(options.agentDir, "mcp.json")} or .pi/mcp.json.`);
+	}
+	for (const report of reports) {
+		const state =
+			report.state === "connected"
+				? `connected, ${report.tools.length} tool${report.tools.length === 1 ? "" : "s"}`
+				: report.state === "needs-auth"
+					? "needs sign-in"
+					: report.state;
+		log(`${report.name}: ${state} (${report.exposure}, ${report.scope})`);
+		log(`  ${report.transport}`);
+		if (report.override) log(`  project override: ${report.override}`);
+		if (report.state === "needs-auth") log(`  sign in with: ${APP_NAME} mcp login ${report.name}`);
+		if (report.tools.length > 0) {
+			const tools = report.tools.map((tool) => {
+				const exposure = report.toolExposure?.[tool];
+				return exposure ? `${tool} [${exposure}]` : tool;
+			});
+			log(`  tools: ${tools.join(", ")}`);
+		}
+		if (report.resources) log("  offers resources");
+		if (report.error) log(`  ${report.error.split("\n").join("\n  ")}`);
+	}
+	for (const configError of loaded.errors) log(`config error: ${configError}`);
+	if (untrustedNote) log(untrustedNote);
+	return failed ? 1 : 0;
+}
+
+async function login(
+	entry: McpServerEntry,
+	connection: McpServerConnection,
+	url: string,
+	timeoutMs: number,
+	options: McpCommandOptions,
+	credentials: McpOAuthCredentialStore,
+	log: (line: string) => void,
+	error: (line: string) => void,
+): Promise<number> {
+	const { name } = entry;
+	// Connecting first answers whether a sign-in is needed and records the server's challenge.
+	try {
+		await connection.getClient();
+		log(`Already signed in to MCP server "${name}" (${connection.tools.length} tools).`);
+		return 0;
+	} catch {
+		if (connection.state !== "needs-auth") {
+			error(`MCP server "${name}" failed to connect: ${connection.error ?? "unknown error"}`);
+			return 1;
+		}
+	}
+
+	const openUrl = options.openUrl ?? openBrowser;
+	const interactive = process.stdin.isTTY === true && options.openUrl === undefined;
+	try {
+		await signInMcpServer({
+			serverUrl: url,
+			store: credentials.forServer(name, url),
+			settings: connection.oauthSettings(),
+			challenge: connection.challenge,
+			prompt: {
+				showAuthorizationUrl: (authorizationUrl) => {
+					log(`Sign in to MCP server "${name}" in your browser:\n${authorizationUrl.href}`);
+					openUrl(authorizationUrl.href);
+				},
+				promptForRedirectUrl: (signal) => waitForRedirectUrl(signal, interactive),
+			},
+			signal: AbortSignal.timeout(timeoutMs),
+		});
+	} catch (signInError) {
+		error(
+			signInError instanceof McpSignInCancelledError
+				? `Sign-in to MCP server "${name}" was cancelled or not completed within ${Math.round(timeoutMs / 1000)} seconds.`
+				: `Sign-in to MCP server "${name}" failed: ${errorMessage(signInError)}`,
+		);
+		return 1;
+	}
+	connection.challenge = undefined;
+	try {
+		await connection.reconnect();
+	} catch (connectError) {
+		error(`Signed in, but ${errorMessage(connectError)}`);
+		return 1;
+	}
+	log(`Signed in to MCP server "${name}" (${connection.tools.length} tools).`);
+	return 0;
+}
+
+/**
+ * The pasted redirect URL in a terminal; otherwise only the browser callback can finish the sign-in.
+ * Resolves to undefined when `signal` aborts: the callback arrived, or the sign-in timed out.
+ */
+async function waitForRedirectUrl(signal: AbortSignal, interactive: boolean): Promise<string | undefined> {
+	if (!interactive) {
+		if (!signal.aborted) await new Promise((resolve) => signal.addEventListener("abort", resolve, { once: true }));
+		return undefined;
+	}
+	const readline = createInterface({ input: process.stdin, output: process.stderr });
+	try {
+		return await readline.question("If the browser cannot reach this machine, paste the URL it was redirected to: ", {
+			signal,
+		});
+	} catch {
+		return undefined;
+	} finally {
+		readline.close();
+	}
+}

third_party/pi/packages/coding-agent/src/extensions/mcp/config.ts

@@ -0,0 +1,242 @@
+/**
+ * MCP server configuration.
+ *
+ * Servers are read from `mcp.json` in the agent directory and, for trusted projects, from
+ * `<project>/.pi/mcp.json`. Both use the `mcpServers` shape shared by other MCP clients, so
+ * existing configurations can be copied over. Project entries replace global entries with the
+ * same name.
+ *
+ * A project entry without `command`, `url`, or `type` overrides only `enabled`, `exposure`, and
+ * `toolExposure` of the global server with the same name, for example to turn it off in one project:
+ * `{ "mcpServers": { "internal-tools": { "enabled": false } } }`. The rest of the global entry is kept,
+ * including credentials the project could not set itself.
+ *
+ * ```json
+ * {
+ *   "mcpServers": {
+ *     "filesystem": { "command": "npx", "args": ["-y", "@modelcontextprotocol/server-filesystem", "."] },
+ *     "docs": { "url": "https://example.com/mcp", "headers": { "Authorization": "Bearer ${DOCS_TOKEN}" } },
+ *     "sentry": { "url": "https://mcp.sentry.dev/mcp" }
+ *   }
+ * }
+ * ```
+ *
+ * HTTP servers without an `Authorization` header use OAuth when they answer 401 (sign in with `/mcp`).
+ * `"auth": { "provider": "<provider>" }` sends the token of a `/login` provider instead. Project files
+ * cannot use it, so a repository cannot pick where the credential goes.
+ *
+ * The top-level `autoEnableCodemode` (default true) activates the codemode tool when a server
+ * with `codemode` exposure connects. A project value overrides the global one.
+ */
+
+import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
+import { dirname, join } from "node:path";
+import { CONFIG_DIR_NAME } from "../../config.ts";
+import {
+	type McpExposure,
+	type McpServerConfig,
+	mcpNamespace,
+	validateMcpServerConfig,
+} from "../../core/mcp-servers.ts";
+
+export type {
+	McpExposure,
+	McpHttpServerConfig,
+	McpOAuthConfig,
+	McpServerConfig,
+	McpStdioServerConfig,
+} from "../../core/mcp-servers.ts";
+export { getMcpToolExposure } from "../../core/mcp-servers.ts";
+
+export interface McpServerEntry {
+	name: string;
+	config: McpServerConfig;
+	/** Config file that defined the entry, or the path of the extension that registered it. */
+	source: string;
+	/**
+	 * The global or the project `mcp.json`, or `extension` for servers registered with
+	 * `pi.registerMcpServer()`. Changes to extension servers are not saved.
+	 */
+	scope?: "global" | "project" | "extension";
+	/** Project `mcp.json` with an override of this global server's `enabled`, `exposure`, or `toolExposure`. */
+	override?: string;
+}
+
+export interface LoadedMcpConfig {
+	servers: McpServerEntry[];
+	/** Activate the codemode tool when `codemode` servers connect. Default: true. */
+	autoEnableCodemode?: boolean;
+	errors: string[];
+	/** The project `mcp.json` when the project is trusted, where `/mcp` saves project overrides. */
+	projectConfig?: string;
+}
+
+const OVERRIDE_KEYS = ["enabled", "exposure", "toolExposure"];
+
+/** Whether an entry overrides a server defined elsewhere instead of defining one. */
+function isOverride(value: Record<string, unknown>): boolean {
+	return value.command === undefined && value.url === undefined && value.type === undefined;
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+interface McpConfigState {
+	servers: Map<string, McpServerEntry>;
+	autoEnableCodemode?: boolean;
+	errors: string[];
+}
+
+function readConfigFile(path: string, scope: "global" | "project", state: McpConfigState): void {
+	const { servers, errors } = state;
+	if (!existsSync(path)) return;
+	let parsed: unknown;
+	try {
+		parsed = JSON.parse(readFileSync(path, "utf8"));
+	} catch (error) {
+		errors.push(`${path}: ${error instanceof Error ? error.message : String(error)}`);
+		return;
+	}
+	if (!isRecord(parsed) || (parsed.mcpServers !== undefined && !isRecord(parsed.mcpServers))) {
+		errors.push(`${path}: expected an object with an "mcpServers" object`);
+		return;
+	}
+	if (typeof parsed.autoEnableCodemode === "boolean") state.autoEnableCodemode = parsed.autoEnableCodemode;
+	else if (parsed.autoEnableCodemode !== undefined) errors.push(`${path}: autoEnableCodemode must be a boolean`);
+	for (const [name, value] of Object.entries(parsed.mcpServers ?? {})) {
+		if (scope === "project" && isRecord(value) && isOverride(value)) {
+			const base = servers.get(name);
+			const extra = Object.keys(value).filter((key) => !OVERRIDE_KEYS.includes(key));
+			if (!base) {
+				errors.push(`${path}: server "${name}" needs "command" or "url", or a global server to override`);
+			} else if (extra.length > 0) {
+				errors.push(`${path}: server "${name}": an override can only set ${OVERRIDE_KEYS.join(", ")}`);
+			} else {
+				const config = validateMcpServerConfig(name, { ...base.config, ...value });
+				if (typeof config === "string") errors.push(`${path}: ${config}`);
+				else servers.set(name, { ...base, config, override: path });
+			}
+			continue;
+		}
+		const config = validateMcpServerConfig(name, value);
+		if (typeof config === "string") {
+			errors.push(`${path}: ${config}`);
+			continue;
+		}
+		// Names that differ only in `-` and `_` would share a namespace.
+		const clash = [...servers.keys()].find((other) => other !== name && mcpNamespace(other) === mcpNamespace(name));
+		if (clash) {
+			errors.push(`${path}: server "${name}" conflicts with "${clash}"`);
+			continue;
+		}
+		if (scope === "project" && "url" in config && config.auth) {
+			errors.push(`${path}: server "${name}": auth is only allowed in the global mcp.json`);
+			continue;
+		}
+		servers.set(name, { name, config, source: path, scope });
+	}
+}
+
+/**
+ * Load global and (when trusted) project MCP configuration. Disabled servers are included with
+ * `enabled: false`, so they can be enabled again.
+ */
+export function loadMcpConfig(options: { agentDir: string; cwd: string; projectTrusted: boolean }): LoadedMcpConfig {
+	const state: McpConfigState = { servers: new Map(), errors: [] };
+	readConfigFile(join(options.agentDir, "mcp.json"), "global", state);
+	const projectConfig = options.projectTrusted ? join(options.cwd, CONFIG_DIR_NAME, "mcp.json") : undefined;
+	if (projectConfig) readConfigFile(projectConfig, "project", state);
+	return {
+		servers: [...state.servers.values()],
+		...(state.autoEnableCodemode === undefined ? {} : { autoEnableCodemode: state.autoEnableCodemode }),
+		errors: state.errors,
+		...(projectConfig ? { projectConfig } : {}),
+	};
+}
+
+/** Settings `/mcp` changes. `enabled: true` and `exposure: "codemode"` are the defaults and remove the key. */
+export interface McpServerConfigPatch {
+	enabled?: boolean;
+	exposure?: McpExposure;
+}
+
+/**
+ * Change one server's settings in the `mcp.json` that defines or overrides it. With `override`, a
+ * missing entry is added as an override. Overrides keep default values, since they replace the global
+ * server's. Other content is kept; the file is rewritten with its indentation.
+ */
+export function updateMcpServerConfig(
+	path: string,
+	name: string,
+	patch: McpServerConfigPatch,
+	options: { override?: boolean } = {},
+): void {
+	editMcpServers(path, (servers, parsed) => {
+		let server = servers?.[name];
+		if (server === undefined && options.override) {
+			server = {};
+			parsed.mcpServers = { ...servers, [name]: server };
+		}
+		if (!isRecord(server)) throw new Error(`${path} does not define MCP server "${name}"`);
+		const keepDefaults = isOverride(server);
+		if (patch.enabled !== undefined) {
+			if (patch.enabled && !keepDefaults) delete server.enabled;
+			else server.enabled = patch.enabled;
+		}
+		if (patch.exposure !== undefined) {
+			if (patch.exposure === "codemode" && !keepDefaults) delete server.exposure;
+			else server.exposure = patch.exposure;
+		}
+		return true;
+	});
+}
+
+/**
+ * Add a server to an `mcp.json`, creating the file when missing. An existing entry with the same
+ * name is replaced. Returns true when an entry was replaced.
+ */
+export function addMcpServerConfig(path: string, name: string, config: McpServerConfig): boolean {
+	let replaced = false;
+	editMcpServers(path, (servers, parsed) => {
+		const target = servers ?? {};
+		replaced = target[name] !== undefined;
+		target[name] = config;
+		parsed.mcpServers = target;
+		return true;
+	});
+	return replaced;
+}
+
+/** Remove a server from an `mcp.json`. Returns false when the file does not define it. */
+export function removeMcpServerConfig(path: string, name: string): boolean {
+	if (!existsSync(path)) return false;
+	let removed = false;
+	editMcpServers(path, (servers) => {
+		if (!servers || servers[name] === undefined) return false;
+		delete servers[name];
+		removed = true;
+		return true;
+	});
+	return removed;
+}
+
+/**
+ * Read an `mcp.json` (an empty config when missing), let `edit` change its `mcpServers`, and write
+ * it back with its indentation when `edit` returns true. Other content is kept.
+ */
+function editMcpServers(
+	path: string,
+	edit: (servers: Record<string, unknown> | undefined, parsed: Record<string, unknown>) => boolean,
+): void {
+	const text = existsSync(path) ? readFileSync(path, "utf8") : undefined;
+	const parsed: unknown = text === undefined ? {} : JSON.parse(text);
+	if (!isRecord(parsed) || (parsed.mcpServers !== undefined && !isRecord(parsed.mcpServers))) {
+		throw new Error(`${path}: expected an object with an "mcpServers" object`);
+	}
+	const servers = isRecord(parsed.mcpServers) ? parsed.mcpServers : undefined;
+	if (!edit(servers, parsed)) return;
+	const indent = (text && /^([ \t]+)\S/m.exec(text)?.[1]) || "  ";
+	mkdirSync(dirname(path), { recursive: true });
+	writeFileSync(path, `${JSON.stringify(parsed, null, indent)}\n`);
+}

third_party/pi/packages/coding-agent/src/extensions/mcp/log.ts

@@ -0,0 +1,69 @@
+/**
+ * Log messages MCP servers send with `notifications/message`, appended to `mcp.log` in the agent
+ * directory. Several pi processes may write to the same file, so every message is one synchronous
+ * append. The file is rotated to `mcp.log.1` once it grows past `MAX_LOG_BYTES`.
+ */
+
+import { appendFileSync, mkdirSync, renameSync, statSync } from "node:fs";
+import { dirname } from "node:path";
+
+const MAX_LOG_BYTES = 5 * 1024 * 1024;
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function formatData(data: unknown): string {
+	if (typeof data === "string") return data;
+	try {
+		return JSON.stringify(data) ?? String(data);
+	} catch {
+		return String(data);
+	}
+}
+
+/** Format one `notifications/message` from `server` as a log line; continuation lines are indented. */
+export function formatMcpLogMessage(server: string, params: unknown, now: Date = new Date()): string {
+	const message = isRecord(params) ? params : { data: params };
+	const level = typeof message.level === "string" ? message.level : "info";
+	const logger = typeof message.logger === "string" && message.logger ? ` ${message.logger}:` : "";
+	const text = formatData(message.data).replace(/\r?\n/g, "\n    ");
+	return `${now.toISOString()} [${server}] ${level}${logger} ${text}\n`;
+}
+
+/** Appends server log messages to one file. Write errors are ignored: logging must not break tools. */
+export class McpServerLog {
+	readonly path: string;
+	private size: number | undefined;
+
+	constructor(path: string) {
+		this.path = path;
+	}
+
+	write(server: string, params: unknown): void {
+		const line = formatMcpLogMessage(server, params);
+		try {
+			if (this.size === undefined) {
+				mkdirSync(dirname(this.path), { recursive: true });
+				this.size = this.currentSize();
+			}
+			if (this.size > MAX_LOG_BYTES) {
+				// Another process may have rotated it already; check before renaming.
+				if (this.currentSize() > MAX_LOG_BYTES) renameSync(this.path, `${this.path}.1`);
+				this.size = this.currentSize();
+			}
+			appendFileSync(this.path, line);
+			this.size += Buffer.byteLength(line);
+		} catch {
+			// Ignore: the log is best effort.
+		}
+	}
+
+	private currentSize(): number {
+		try {
+			return statSync(this.path).size;
+		} catch {
+			return 0;
+		}
+	}
+}

third_party/pi/packages/coding-agent/src/extensions/mcp/oauth.ts

@@ -0,0 +1,558 @@
+/**
+ * OAuth sign-in for remote MCP servers.
+ *
+ * Connections never start a browser flow on their own. They send the stored access token and, after
+ * a 401, try the stored refresh token. When that is not possible they fail with
+ * `McpOAuthAuthorizationRequiredError`, and the user signs in through `/mcp`, which runs
+ * the authorization code flow (PKCE, a Client ID Metadata Document or dynamic client registration)
+ * against a loopback callback.
+ *
+ * Credentials live in `<agent-dir>/mcp-auth.json`, keyed by server name and URL.
+ */
+
+import { createHash } from "node:crypto";
+import { mkdirSync } from "node:fs";
+import { join } from "node:path";
+import { oauthErrorHtml, oauthSuccessHtml } from "@earendil-works/pi-ai/utils/oauth-page";
+import type { AuthProvider, McpFetch } from "@earendil-works/pi-mcp";
+import {
+	type AuthorizationServerMetadata,
+	authorizeMcp,
+	McpOAuthAuthorizationRequiredError,
+	McpOAuthProvider,
+	type McpOAuthState,
+	type McpOAuthStateStore,
+	type OAuthCallback,
+	type OAuthCallbackPage,
+	OAuthCallbackServer,
+	type OAuthChallenge,
+	type OAuthClientInformationMixed,
+	type OAuthClientMetadataDocument,
+	parseWwwAuthenticate,
+	stepUpScope,
+} from "@earendil-works/pi-mcp/oauth";
+import lockfile from "proper-lockfile";
+import { APP_NAME, getAgentDir } from "../../config.ts";
+import { type AuthStorageBackend, FileAuthStorageBackend } from "../../core/auth-storage.ts";
+import { mcpNamespace } from "../../core/mcp-servers.ts";
+
+const CALLBACK_HOST = "127.0.0.1";
+const CALLBACK_PATH = "/callback";
+/** Where pi.dev serves pi's Client ID Metadata Documents: `client.json` and `<callback ID>/client.json`. */
+const CLIENT_METADATA_BASE_URL = "https://pi.dev/oauth";
+/** Redirect URI for refreshes when none is stored. Refreshing never redirects the user. */
+const FALLBACK_REDIRECT_URL = `http://${CALLBACK_HOST}${CALLBACK_PATH}`;
+/** Access tokens this close to expiry are refreshed before they are sent. */
+const REFRESH_SKEW_MS = 30_000;
+/**
+ * Bounds each request to the authorization server, so an unresponsive one cannot hold the refresh lock,
+ * delay shutdown, or keep a sign-in waiting for long.
+ */
+const OAUTH_REQUEST_TIMEOUT_MS = 15_000;
+/** A refresh lock that its holder stopped renewing (the process was killed) is taken over after this. */
+const REFRESH_LOCK_STALE_MS = 20_000;
+/** How long to wait for another process's refresh: longer than a stale lock lives. */
+const REFRESH_LOCK_WAIT_MS = 25_000;
+const REFRESH_LOCK_RETRY_MS = 100;
+
+export interface McpOAuthSettings {
+	clientId?: string;
+	/** Already resolved. */
+	clientSecret?: string;
+	callbackPort?: number;
+	/** Loopback redirect URI; see `McpOAuthConfig.callbackUrl`. */
+	callbackUrl?: string;
+	/** Scopes to request, separated by spaces. */
+	scope?: string;
+	/** `client_name` for dynamic client registration. Default: `APP_NAME`. */
+	clientName?: string;
+	/** See `McpOAuthConfig.clientRegistration`. */
+	clientRegistration?: "dcr" | "cimd";
+	/** See `McpOAuthConfig.authServerMetadataUrl`. */
+	authServerMetadataUrl?: URL;
+}
+
+/** Where the loopback callback server listens and the redirect URI it serves. */
+interface CallbackSettings {
+	/** Address to listen on. */
+	host: string;
+	/** Host name in the redirect URI. */
+	redirectHost: string;
+	port: number | undefined;
+	path: string;
+	/** The exact redirect URI, when the port is fixed. */
+	fixedRedirectUrl: string | undefined;
+}
+
+function callbackSettings(settings: McpOAuthSettings): CallbackSettings {
+	const url = new URL(settings.callbackUrl ?? `http://${CALLBACK_HOST}${CALLBACK_PATH}`);
+	const address = url.hostname.replace(/^\[|\]$/g, "");
+	const port = url.port ? Number(url.port) : settings.callbackPort;
+	let fixedRedirectUrl: string | undefined;
+	// A configured URI with a port is sent exactly as written, since servers compare it as a string.
+	if (url.port) fixedRedirectUrl = settings.callbackUrl;
+	else if (port !== undefined) {
+		url.port = String(port);
+		fixedRedirectUrl = url.href;
+	}
+	return {
+		// `localhost` is served on 127.0.0.1; browsers fall back to it when ::1 refuses.
+		host: address === "localhost" ? CALLBACK_HOST : address,
+		redirectHost: address,
+		port,
+		path: url.pathname,
+		fixedRedirectUrl,
+	};
+}
+
+/** Scopes of both lists, each once. */
+function mergeScopes(...scopes: (string | undefined)[]): string | undefined {
+	const merged = [...new Set(scopes.flatMap((scope) => scope?.split(/\s+/).filter(Boolean) ?? []))];
+	return merged.length > 0 ? merged.join(" ") : undefined;
+}
+
+type StoredStates = Record<string, McpOAuthState>;
+
+function parseStates(content: string | undefined): StoredStates {
+	if (!content?.trim()) return {};
+	const parsed: unknown = JSON.parse(content);
+	return typeof parsed === "object" && parsed !== null && !Array.isArray(parsed) ? (parsed as StoredStates) : {};
+}
+
+function serializeStates(states: StoredStates): string {
+	return `${JSON.stringify(states, null, 2)}\n`;
+}
+
+/**
+ * Keys of a server's state: by name and URL, so servers sharing a URL keep separate accounts, and the
+ * legacy key by URL alone, written by older versions.
+ */
+function storeKeys(name: string, serverUrl: string): { key: string; legacyKey: string } {
+	const legacyKey = String(new URL(serverUrl));
+	return { key: `${mcpNamespace(name)}|${legacyKey}`, legacyKey };
+}
+
+export interface McpOAuthServerStore extends McpOAuthStateStore {
+	/** Run `fn` while no other process refreshes the server's tokens. */
+	withRefreshLock<T>(fn: () => Promise<T>): Promise<T>;
+}
+
+/** Per-server OAuth state (client registration, tokens, pending PKCE verifier) in `mcp-auth.json`. */
+export class McpOAuthCredentialStore {
+	private readonly backend: AuthStorageBackend;
+	/** Directory for the refresh lock files. Without one, refreshes are only serialized in this process. */
+	private readonly lockDir: string | undefined;
+
+	constructor(backend?: AuthStorageBackend, lockDir?: string) {
+		this.backend = backend ?? new FileAuthStorageBackend(join(getAgentDir(), "mcp-auth.json"));
+		this.lockDir = backend ? lockDir : getAgentDir();
+	}
+
+	forServer(name: string, serverUrl: string): McpOAuthServerStore {
+		const { key, legacyKey } = storeKeys(name, serverUrl);
+		return {
+			// The first server to load legacy state takes it over; others with the same URL sign in again.
+			load: () =>
+				this.backend.withLock((current) => {
+					const states = parseStates(current);
+					if (states[key] || !states[legacyKey]) return { result: states[key] };
+					states[key] = states[legacyKey];
+					delete states[legacyKey];
+					return { result: states[key], next: serializeStates(states) };
+				}),
+			save: (state) =>
+				this.write((states) => {
+					states[key] = state;
+				}),
+			withRefreshLock: (fn) => this.withRefreshLock(key, fn),
+		};
+	}
+
+	/**
+	 * A lock file per server. When the process exits, proper-lockfile removes the locks it holds; when
+	 * it is killed, the lock goes stale because it is no longer renewed, and the next process takes it over.
+	 */
+	private async withRefreshLock<T>(key: string, fn: () => Promise<T>): Promise<T> {
+		if (!this.lockDir) return fn();
+		mkdirSync(this.lockDir, { recursive: true, mode: 0o700 });
+		const hash = createHash("sha256").update(key).digest("hex").slice(0, 16);
+		const release = await lockfile.lock(join(this.lockDir, `mcp-auth-refresh-${hash}`), {
+			realpath: false,
+			stale: REFRESH_LOCK_STALE_MS,
+			retries: {
+				retries: REFRESH_LOCK_WAIT_MS / REFRESH_LOCK_RETRY_MS,
+				factor: 1,
+				minTimeout: REFRESH_LOCK_RETRY_MS,
+				maxTimeout: REFRESH_LOCK_RETRY_MS,
+			},
+			// The default throws from a timer. A lost lock at worst lets two refreshes overlap.
+			onCompromised: () => {},
+		});
+		try {
+			return await fn();
+		} finally {
+			await release().catch(() => undefined);
+		}
+	}
+
+	/** The stored tokens of a server, for noticing sign-ins done by another process. Does not take over legacy state. */
+	tokens(name: string, serverUrl: string): McpOAuthState["tokens"] {
+		const { key, legacyKey } = storeKeys(name, serverUrl);
+		const states = this.backend.withLock((current) => ({ result: parseStates(current) }));
+		return (states[key] ?? states[legacyKey])?.tokens;
+	}
+
+	/** Returns whether credentials were stored for the server. Removes legacy state the server would take over. */
+	remove(name: string, serverUrl: string): boolean {
+		const { key, legacyKey } = storeKeys(name, serverUrl);
+		return this.backend.withLock((current) => {
+			const states = parseStates(current);
+			const stored = key in states ? key : legacyKey in states ? legacyKey : undefined;
+			if (!stored) return { result: false };
+			delete states[stored];
+			return { result: true, next: serializeStates(states) };
+		});
+	}
+
+	private write(update: (states: StoredStates) => void): void {
+		this.backend.withLock((current) => {
+			const states = parseStates(current);
+			update(states);
+			return { result: undefined, next: serializeStates(states) };
+		});
+	}
+}
+
+/** `fetch` with `OAUTH_REQUEST_TIMEOUT_MS` per request, on top of the request's own signal. */
+function timedFetch(fetch: McpFetch = globalThis.fetch): McpFetch {
+	return (input, init) => {
+		const timeout = AbortSignal.timeout(OAUTH_REQUEST_TIMEOUT_MS);
+		return fetch(input, { ...init, signal: init?.signal ? AbortSignal.any([init.signal, timeout]) : timeout });
+	};
+}
+
+function registeredRedirectUrls(client: OAuthClientInformationMixed | undefined): string[] {
+	return client && "redirect_uris" in client ? client.redirect_uris : [];
+}
+
+/** 12 characters identifying an MCP server URL in callback paths, computed like Codex does. */
+function callbackId(serverUrl: string): string {
+	const url = new URL(serverUrl);
+	url.hash = "";
+	return createHash("sha256").update(url.href).digest().subarray(0, 9).toString("base64url");
+}
+
+/**
+ * pi's Client ID Metadata Document, for `clientRegistration: "cimd"`, chosen like Codex chooses its own.
+ * The configuration ensures the default callback path. Without the `iss` parameter in authorization
+ * responses (RFC 9207), the redirect URI and the document are specific to the MCP server, so a response
+ * cannot be mixed up with one from another authorization server (RFC 9700 section 4.4.2.2).
+ */
+function clientMetadataDocument(
+	serverUrl: string,
+	redirectUrl: string,
+	metadata: AuthorizationServerMetadata | undefined,
+): OAuthClientMetadataDocument {
+	if (
+		!metadata?.client_id_metadata_document_supported ||
+		!metadata.token_endpoint_auth_methods_supported?.includes("none")
+	) {
+		throw new Error(
+			'The authorization server does not support Client ID Metadata Documents for public clients; remove oauth.clientRegistration "cimd"',
+		);
+	}
+	if (metadata.authorization_response_iss_parameter_supported) {
+		return { url: `${CLIENT_METADATA_BASE_URL}/client.json`, redirectUrl };
+	}
+	const id = callbackId(serverUrl);
+	const redirect = new URL(redirectUrl);
+	redirect.pathname = `${CALLBACK_PATH}/${id}`;
+	return { url: `${CLIENT_METADATA_BASE_URL}/${id}/client.json`, redirectUrl: redirect.href };
+}
+
+function createProvider(
+	serverUrl: string,
+	store: McpOAuthStateStore,
+	settings: McpOAuthSettings,
+	redirectUrl: string,
+	onRedirect: (url: URL) => void,
+): McpOAuthProvider {
+	return new McpOAuthProvider({
+		serverUrl,
+		redirectUrl,
+		clientMetadata: { client_name: settings.clientName ?? APP_NAME },
+		clientMetadataDocument:
+			settings.clientRegistration === "cimd"
+				? (metadata) => clientMetadataDocument(serverUrl, redirectUrl, metadata)
+				: undefined,
+		clientId: settings.clientId,
+		clientSecret: settings.clientSecret,
+		store,
+		onRedirect,
+	});
+}
+
+export interface McpAuthProvider extends AuthProvider {
+	/** Resolves when no refresh is running, so shutdown does not drop rotated tokens before they are saved. */
+	settled(): Promise<void>;
+}
+
+/**
+ * Auth provider for MCP connections: sends the stored access token and refreshes it when it is about
+ * to expire or after a 401. Throws `McpOAuthAuthorizationRequiredError` when the user has to sign in,
+ * including when the server asks for more scope (`insufficient_scope`). `onChallenge` receives the
+ * server's `WWW-Authenticate` challenge so sign-in can use its resource metadata URL and scope.
+ * `settings` is only called when a refresh is needed, so a secret that fails to resolve fails the
+ * refresh instead of the whole connection setup.
+ *
+ * Many servers rotate refresh tokens, so two refreshes with the same refresh token lose the grant.
+ * Requests in this process share one refresh, and other processes are kept out by the store's
+ * refresh lock, held from reading the tokens to saving new ones. Tokens that changed meanwhile
+ * (another process refreshed them, or the user signed in) are used without refreshing.
+ */
+export function createMcpAuthProvider(options: {
+	serverUrl: string;
+	store: McpOAuthServerStore;
+	settings: () => McpOAuthSettings;
+	onChallenge: (challenge: OAuthChallenge) => void;
+}): McpAuthProvider {
+	const { serverUrl, store } = options;
+	let refreshing: Promise<void> | undefined;
+
+	/** Replace `staleToken`, the access token that expired or was rejected. */
+	const refresh = (staleToken: string | undefined, fetch: McpFetch = globalThis.fetch, challenge?: OAuthChallenge) => {
+		refreshing ??= store
+			.withRefreshLock(async () => {
+				const state = await store.load();
+				if (state?.tokens?.access_token !== staleToken) return;
+				if (!state?.tokens?.refresh_token) throw new McpOAuthAuthorizationRequiredError();
+				const settings = options.settings();
+				const redirectUrl =
+					callbackSettings(settings).fixedRedirectUrl ??
+					registeredRedirectUrls(state.clientInformation)[0] ??
+					FALLBACK_REDIRECT_URL;
+				const provider = createProvider(serverUrl, store, settings, redirectUrl, () => {});
+				// Refreshes the tokens, or reports that a new sign-in is needed. Not cancellable: a refresh the
+				// server answered may have rotated the refresh token, so its answer must be saved.
+				const result = await authorizeMcp(provider, {
+					serverUrl,
+					resourceMetadataUrl: challenge?.resourceMetadataUrl,
+					authorizationServerMetadataUrl: settings.authServerMetadataUrl,
+					scope: challenge?.scope,
+					fetch: timedFetch(fetch),
+				});
+				if (result === "REDIRECT") throw new McpOAuthAuthorizationRequiredError();
+			})
+			.finally(() => {
+				refreshing = undefined;
+			});
+		return refreshing;
+	};
+
+	return {
+		token: async () => {
+			await refreshing?.catch(() => undefined);
+			const state = await store.load();
+			const token = state?.tokens?.access_token;
+			const expired = state?.tokensExpireAt !== undefined && state.tokensExpireAt - REFRESH_SKEW_MS <= Date.now();
+			if (!expired || !state?.tokens?.refresh_token) return token;
+			// Failures fall through: the request goes out with the old token and a 401 decides what happens.
+			await refresh(token).catch(() => undefined);
+			return (await store.load())?.tokens?.access_token;
+		},
+		onUnauthorized: async (context) => {
+			const challenge = parseWwwAuthenticate(context.response.headers.get("www-authenticate"));
+			options.onChallenge(challenge);
+			// A refresh keeps the granted scope, so more scope needs a new sign-in.
+			if (challenge.error === "insufficient_scope") throw new McpOAuthAuthorizationRequiredError();
+			await refresh(context.token, context.fetch, challenge);
+		},
+		settled: async () => {
+			await refreshing?.catch(() => undefined);
+		},
+	};
+}
+
+export interface McpSignInPrompt {
+	/** Show the authorization URL to the user and open it in a browser. */
+	showAuthorizationUrl(url: URL): void;
+	/**
+	 * Ask for the redirect URL from the browser address bar, for when the browser cannot reach the
+	 * loopback callback (for example over SSH). Aborted once the callback arrives or the sign-in is
+	 * cancelled. Resolves to `undefined` or an empty string when the user cancels.
+	 */
+	promptForRedirectUrl(signal: AbortSignal): Promise<string | undefined>;
+}
+
+export class McpSignInCancelledError extends Error {
+	constructor() {
+		super("Sign-in cancelled");
+		this.name = "McpSignInCancelledError";
+	}
+}
+
+type AuthorizationResponse = Pick<OAuthCallback, "code" | "iss">;
+
+function responseFromRedirectUrl(input: string, state: string, redirectUrl: URL): AuthorizationResponse {
+	let url: URL;
+	try {
+		url = new URL(input.trim());
+	} catch {
+		throw new Error("Expected the full redirect URL from the browser address bar");
+	}
+	// A server-specific redirect URI tells authorization servers apart, so it must match exactly.
+	if (url.origin !== redirectUrl.origin || url.pathname !== redirectUrl.pathname) {
+		throw new Error("The redirect URL does not match this sign-in's redirect URI");
+	}
+	const error = url.searchParams.get("error");
+	if (error) throw new Error(url.searchParams.get("error_description") ?? error);
+	if (url.searchParams.get("state") !== state) throw new Error("The redirect URL belongs to a different sign-in");
+	const code = url.searchParams.get("code");
+	if (!code) throw new Error("The redirect URL does not contain an authorization code");
+	return { code, iss: url.searchParams.get("iss") ?? undefined };
+}
+
+/**
+ * Wait for the browser callback or a pasted redirect URL, whichever comes first. Aborting `signal`
+ * aborts the prompt, which then cancels the sign-in.
+ */
+async function waitForAuthorizationResponse(
+	callback: OAuthCallbackServer,
+	state: string,
+	redirectUrl: URL,
+	prompt: McpSignInPrompt,
+	signal: AbortSignal | undefined,
+): Promise<AuthorizationResponse> {
+	const controller = new AbortController();
+	const promptSignal = signal ? AbortSignal.any([controller.signal, signal]) : controller.signal;
+	const fromBrowser = callback.waitForCallback(state, redirectUrl.pathname);
+	const fromUser = prompt.promptForRedirectUrl(promptSignal).then((input) => {
+		if (!input?.trim()) throw new McpSignInCancelledError();
+		return responseFromRedirectUrl(input, state, redirectUrl);
+	});
+	try {
+		return await Promise.race([fromBrowser, fromUser]);
+	} finally {
+		// The losing side rejects once the prompt is aborted or the callback server closes.
+		controller.abort();
+		fromBrowser.catch(() => undefined);
+		fromUser.catch(() => undefined);
+	}
+}
+
+/** Listen on `port`, or on a free port when it is taken and not `required`. */
+async function listenForCallback(
+	settings: CallbackSettings,
+	extraPaths: string[],
+	port: number | undefined,
+	required: boolean,
+): Promise<OAuthCallbackServer> {
+	const options = {
+		host: settings.host,
+		redirectHost: settings.redirectHost,
+		path: settings.path,
+		extraPaths,
+		renderPage: (page: OAuthCallbackPage) =>
+			page.ok
+				? oauthSuccessHtml("Signed in to the MCP server. You may now close this page.")
+				: oauthErrorHtml(page.message, page.details),
+	};
+	try {
+		return await OAuthCallbackServer.listen({ ...options, port: port ?? 0 });
+	} catch (error) {
+		if (required || port === undefined) throw error;
+		return OAuthCallbackServer.listen(options);
+	}
+}
+
+/**
+ * Sign in to an MCP server. Uses the stored refresh token when possible; otherwise runs the browser
+ * authorization code flow. Tokens are saved to `store`. Aborting `signal` stops the sign-in at any
+ * step with `McpSignInCancelledError`; requests to the authorization server are also time-limited.
+ */
+export async function signInMcpServer(options: {
+	serverUrl: string;
+	store: McpOAuthStateStore;
+	settings: McpOAuthSettings;
+	challenge?: OAuthChallenge;
+	prompt: McpSignInPrompt;
+	signal?: AbortSignal;
+}): Promise<void> {
+	const { serverUrl, store, settings, signal } = options;
+	if (signal?.aborted) throw new McpSignInCancelledError();
+	const stored = await store.load();
+	const stepUp = options.challenge?.error === "insufficient_scope";
+	const callbackOptions = callbackSettings(settings);
+	// Reuse the port of the registered redirect URI so the registered client stays valid.
+	const registered = registeredRedirectUrls(stored?.clientInformation)[0];
+	const preferredPort =
+		callbackOptions.port ?? (registered ? Number(new URL(registered).port) || undefined : undefined);
+	const cimd = settings.clientRegistration === "cimd";
+	const callback = await listenForCallback(
+		callbackOptions,
+		// The redirect URI of a server-specific Client ID Metadata Document.
+		cimd ? [`${CALLBACK_PATH}/${callbackId(serverUrl)}`] : [],
+		preferredPort,
+		callbackOptions.port !== undefined,
+	);
+	const redirectUrl = callbackOptions.fixedRedirectUrl ?? callback.redirectUrl;
+	try {
+		if (stored) {
+			const next: McpOAuthState = { ...stored };
+			// Every sign-in gets a fresh `state` parameter.
+			delete next.oauthState;
+			// A registered client cannot use another redirect URI, and its tokens belong to it. A Client ID
+			// Metadata Document is not stored, so with one, a stored client was registered before and is replaced.
+			const keepClient =
+				settings.clientId ||
+				(cimd ? !stored.clientInformation : registeredRedirectUrls(stored.clientInformation).includes(redirectUrl));
+			if (!keepClient) {
+				delete next.clientInformation;
+				delete next.tokens;
+				delete next.tokensExpireAt;
+			}
+			await store.save(next);
+		}
+
+		let authorizationUrl: URL | undefined;
+		const provider = createProvider(serverUrl, store, settings, redirectUrl, (url) => {
+			authorizationUrl = url;
+		});
+		const flow = {
+			serverUrl,
+			fetch: timedFetch(),
+			signal,
+			resourceMetadataUrl: options.challenge?.resourceMetadataUrl,
+			authorizationServerMetadataUrl: settings.authServerMetadataUrl,
+			// A server asking for more scope gets it on top of the configured scope and, since the challenge
+			// may list only the missing scopes, on top of the scope granted so far.
+			scope: mergeScopes(
+				settings.scope,
+				stepUp ? stepUpScope(stored?.tokens?.scope, options.challenge?.scope) : options.challenge?.scope,
+			),
+		};
+		// A refresh keeps the granted scope; a server asking for more needs the browser flow.
+		const skipRefresh = stepUp;
+		if ((await authorizeMcp(provider, { ...flow, skipRefresh })) === "AUTHORIZED") return;
+		if (!authorizationUrl) throw new Error("OAuth flow did not produce an authorization URL");
+
+		const state = await provider.state();
+		// The flow picks the redirect URI, which may be specific to the MCP server.
+		const authorizationRedirectUrl = new URL(authorizationUrl.searchParams.get("redirect_uri") ?? redirectUrl);
+		options.prompt.showAuthorizationUrl(authorizationUrl);
+		const { code, iss } = await waitForAuthorizationResponse(
+			callback,
+			state,
+			authorizationRedirectUrl,
+			options.prompt,
+			signal,
+		);
+		await authorizeMcp(provider, { ...flow, authorizationCode: code, iss });
+	} catch (error) {
+		// Aborted requests fail with the signal's reason; report them as the cancellation they are.
+		if (signal?.aborted) throw new McpSignInCancelledError();
+		throw error;
+	} finally {
+		await callback.close();
+	}
+}

third_party/pi/packages/coding-agent/src/extensions/mcp/resources.ts

@@ -0,0 +1,339 @@
+/**
+ * MCP resources, through the tools Codex and opencode use: `list_mcp_resources`,
+ * `list_mcp_resource_templates`, and `read_mcp_resource`. They take a `server` argument and cover
+ * every connected server with resources, so models trained on those tools use them unchanged.
+ *
+ * Listings are JSON, as in Codex: `{ server?, resources: [{ server, ...resource }], nextCursor? }`.
+ * With a `server`, one page is listed and `cursor` continues it; without, every page of every server.
+ * MCP App resources (`ui://` URIs and `profile=mcp-app` HTML) are left out, since they are user
+ * interfaces for hosts that render them, and so are icons. Read resources become text and images for
+ * the model; binary resources are saved to temp files. Scripts get the JSON payloads.
+ */
+
+import type { ImageContent, JsonValue, TextContent } from "@earendil-works/pi-ai";
+import type {
+	ContentBlock,
+	ListResourcesResult,
+	ListResourceTemplatesResult,
+	McpRequestOptions,
+	ReadResourceResult,
+	Resource,
+	ResourceTemplate,
+} from "@earendil-works/pi-mcp";
+import type { TSchema } from "typebox";
+import type { ToolAnnotations, ToolDefinition } from "../../core/extensions/types.ts";
+import {
+	LIST_MCP_RESOURCE_TEMPLATES_TOOL,
+	LIST_MCP_RESOURCES_TOOL,
+	READ_MCP_RESOURCE_TOOL,
+} from "../../core/mcp-servers.ts";
+import type { McpExposure } from "./config.ts";
+import { limitMcpContent, type McpToolDetails, toModelContent, toToolExposure } from "./tools.ts";
+
+export { LIST_MCP_RESOURCE_TEMPLATES_TOOL, LIST_MCP_RESOURCES_TOOL, READ_MCP_RESOURCE_TOOL };
+
+/** A connected server that offers resources. */
+export interface McpResourceServer {
+	name: string;
+	timeoutMs: number;
+	resourcesPage(cursor: string | undefined, options: McpRequestOptions): Promise<ListResourcesResult>;
+	resourceTemplatesPage(cursor: string | undefined, options: McpRequestOptions): Promise<ListResourceTemplatesResult>;
+	allResources(options: McpRequestOptions): Promise<Resource[]>;
+	allResourceTemplates(options: McpRequestOptions): Promise<ResourceTemplate[]>;
+	readResource(uri: string, options: McpRequestOptions): Promise<ReadResourceResult>;
+}
+
+/** MCP App user interfaces, which only hosts that render them can use. */
+function isMcpAppResource(item: { uri?: string; uriTemplate?: string; mimeType?: string }): boolean {
+	const uri = item.uri ?? item.uriTemplate ?? "";
+	return uri.startsWith("ui://") || /;\s*profile\s*=\s*"?mcp-app"?/i.test(item.mimeType ?? "");
+}
+
+/** A listed resource or template without `_meta` and icons, tagged with its server. */
+function listed<T extends { _meta?: unknown }>(server: string, item: T): Record<string, unknown> {
+	const { _meta, icons: _icons, ...rest } = item as T & { icons?: unknown };
+	return { server, ...rest };
+}
+
+const stringProperty = (description: string) => ({ type: "string", description });
+const SERVER_FILTER = stringProperty("MCP server name. Omit to list every server with resources.");
+const CURSOR = stringProperty("Opaque cursor from a previous call with the same server; omit for the first page.");
+const LIST_PARAMETERS = {
+	type: "object",
+	properties: { server: SERVER_FILTER, cursor: CURSOR },
+	additionalProperties: false,
+};
+const READ_PARAMETERS = {
+	type: "object",
+	properties: {
+		server: stringProperty(
+			"MCP server name exactly as configured. Must match the 'server' field returned by list_mcp_resources.",
+		),
+		uri: stringProperty("Resource URI to read. Must be one of the URIs returned by list_mcp_resources."),
+	},
+	required: ["server", "uri"],
+	additionalProperties: false,
+};
+
+const optionalString = { type: "string" };
+const LISTING_ERRORS = {
+	type: "array",
+	description: "Servers that could not be listed",
+	items: {
+		type: "object",
+		properties: { server: { type: "string" }, error: { type: "string" } },
+		required: ["server", "error"],
+	},
+};
+const LIST_OUTPUT_SCHEMA = {
+	type: "object",
+	properties: {
+		server: optionalString,
+		resources: {
+			type: "array",
+			items: {
+				type: "object",
+				properties: {
+					server: { type: "string" },
+					uri: { type: "string" },
+					name: { type: "string" },
+					title: optionalString,
+					description: optionalString,
+					mimeType: optionalString,
+					size: { type: "number" },
+				},
+				required: ["server", "uri", "name"],
+			},
+		},
+		nextCursor: optionalString,
+		errors: LISTING_ERRORS,
+	},
+	required: ["resources"],
+};
+const LIST_TEMPLATES_OUTPUT_SCHEMA = {
+	type: "object",
+	properties: {
+		server: optionalString,
+		resourceTemplates: {
+			type: "array",
+			items: {
+				type: "object",
+				properties: {
+					server: { type: "string" },
+					uriTemplate: { type: "string", description: "RFC 6570 URI template" },
+					name: { type: "string" },
+					title: optionalString,
+					description: optionalString,
+					mimeType: optionalString,
+				},
+				required: ["server", "uriTemplate", "name"],
+			},
+		},
+		nextCursor: optionalString,
+		errors: LISTING_ERRORS,
+	},
+	required: ["resourceTemplates"],
+};
+const READ_OUTPUT_SCHEMA = {
+	type: "object",
+	properties: {
+		server: { type: "string" },
+		uri: { type: "string" },
+		contents: {
+			type: "array",
+			items: {
+				anyOf: [
+					{
+						type: "object",
+						properties: { uri: { type: "string" }, mimeType: optionalString, text: { type: "string" } },
+						required: ["uri", "text"],
+					},
+					{
+						type: "object",
+						properties: {
+							uri: { type: "string" },
+							mimeType: optionalString,
+							blob: { type: "string", description: "base64" },
+						},
+						required: ["uri", "blob"],
+					},
+				],
+			},
+		},
+	},
+	required: ["server", "uri", "contents"],
+};
+
+function stringArgument(params: unknown, key: string): string | undefined {
+	const value = (params as Record<string, unknown> | undefined)?.[key];
+	if (value === undefined || value === null) return undefined;
+	if (typeof value !== "string") throw new Error(`${key} must be a string`);
+	return value.trim() || undefined;
+}
+
+function errorMessage(error: unknown): string {
+	return error instanceof Error ? error.message : String(error);
+}
+
+async function jsonResult(
+	tool: string,
+	server: string | undefined,
+	payload: Record<string, unknown>,
+): Promise<{ content: (TextContent | ImageContent)[]; details: McpToolDetails; structuredContent: JsonValue }> {
+	const { content, fullOutputPath } = await limitMcpContent([{ type: "text", text: JSON.stringify(payload) }]);
+	return {
+		content,
+		details: { server: server ?? "", tool, ...(fullOutputPath ? { fullOutputPath } : {}) },
+		structuredContent: payload as unknown as JsonValue,
+	};
+}
+
+/**
+ * The three resource tools. `servers` returns the servers whose resources they reach, at call time.
+ */
+export function createMcpResourceToolDefinitions(options: {
+	exposure: McpExposure;
+	servers: () => readonly McpResourceServer[];
+}): ToolDefinition<TSchema, McpToolDetails>[] {
+	const readOnly: ToolAnnotations = { readOnlyHint: true };
+
+	const findServer = (name: string): McpResourceServer => {
+		const servers = options.servers();
+		const server = servers.find((candidate) => candidate.name === name);
+		if (server) return server;
+		const available = servers.map((candidate) => candidate.name).join(", ");
+		throw new Error(
+			`MCP server "${name}" has no resources${available ? `. Servers with resources: ${available}` : ""}`,
+		);
+	};
+
+	/** One page of one server, or every page of every server. */
+	const list = async <T extends { _meta?: unknown }>(
+		params: unknown,
+		signal: AbortSignal | undefined,
+		key: "resources" | "resourceTemplates",
+		page: (
+			server: McpResourceServer,
+			cursor: string | undefined,
+			options: McpRequestOptions,
+		) => Promise<{
+			items: T[];
+			nextCursor?: string;
+		}>,
+		all: (server: McpResourceServer, options: McpRequestOptions) => Promise<T[]>,
+	): Promise<Record<string, unknown>> => {
+		const serverName = stringArgument(params, "server");
+		const cursor = stringArgument(params, "cursor");
+		const visible = (item: T) => !isMcpAppResource(item as { uri?: string; uriTemplate?: string; mimeType?: string });
+		if (serverName) {
+			const server = findServer(serverName);
+			const result = await page(server, cursor, { signal, timeoutMs: server.timeoutMs });
+			return {
+				server: server.name,
+				[key]: result.items.filter(visible).map((item) => listed(server.name, item)),
+				...(result.nextCursor === undefined ? {} : { nextCursor: result.nextCursor }),
+			};
+		}
+		if (cursor) throw new Error("cursor can only be used when a server is specified");
+		const servers = [...options.servers()].sort((a, b) => a.name.localeCompare(b.name));
+		const results = await Promise.allSettled(
+			servers.map((server) => all(server, { signal, timeoutMs: server.timeoutMs })),
+		);
+		const items: Record<string, unknown>[] = [];
+		const errors: { server: string; error: string }[] = [];
+		results.forEach((result, index) => {
+			const server = servers[index].name;
+			if (result.status === "fulfilled")
+				items.push(...result.value.filter(visible).map((item) => listed(server, item)));
+			else errors.push({ server, error: errorMessage(result.reason) });
+		});
+		return { [key]: items, ...(errors.length > 0 ? { errors } : {}) };
+	};
+
+	const listResources: ToolDefinition<TSchema, McpToolDetails> = {
+		name: LIST_MCP_RESOURCES_TOOL,
+		label: LIST_MCP_RESOURCES_TOOL,
+		description:
+			"Lists resources provided by MCP servers. Resources allow servers to share data that provides context to language models, such as files, database schemas, or application-specific information. Prefer resources over web search when possible.",
+		parameters: LIST_PARAMETERS as unknown as TSchema,
+		outputSchema: LIST_OUTPUT_SCHEMA as unknown as TSchema,
+		exposure: toToolExposure(options.exposure),
+		annotations: readOnly,
+		async execute(_toolCallId, params, signal) {
+			const payload = await list(
+				params,
+				signal,
+				"resources",
+				async (server, cursor, requestOptions) => {
+					const result = await server.resourcesPage(cursor, requestOptions);
+					return { items: result.resources, nextCursor: result.nextCursor };
+				},
+				(server, requestOptions) => server.allResources(requestOptions),
+			);
+			return jsonResult(LIST_MCP_RESOURCES_TOOL, stringArgument(params, "server"), payload);
+		},
+	};
+
+	const listTemplates: ToolDefinition<TSchema, McpToolDetails> = {
+		name: LIST_MCP_RESOURCE_TEMPLATES_TOOL,
+		label: LIST_MCP_RESOURCE_TEMPLATES_TOOL,
+		description:
+			"Lists resource templates provided by MCP servers. Parameterized resource templates allow servers to share data that takes parameters and provides context to language models, such as files, database schemas, or application-specific information. Prefer resource templates over web search when possible.",
+		parameters: LIST_PARAMETERS as unknown as TSchema,
+		outputSchema: LIST_TEMPLATES_OUTPUT_SCHEMA as unknown as TSchema,
+		exposure: toToolExposure(options.exposure),
+		annotations: readOnly,
+		async execute(_toolCallId, params, signal) {
+			const payload = await list(
+				params,
+				signal,
+				"resourceTemplates",
+				async (server, cursor, requestOptions) => {
+					const result = await server.resourceTemplatesPage(cursor, requestOptions);
+					return { items: result.resourceTemplates, nextCursor: result.nextCursor };
+				},
+				(server, requestOptions) => server.allResourceTemplates(requestOptions),
+			);
+			return jsonResult(LIST_MCP_RESOURCE_TEMPLATES_TOOL, stringArgument(params, "server"), payload);
+		},
+	};
+
+	const readResource: ToolDefinition<TSchema, McpToolDetails> = {
+		name: READ_MCP_RESOURCE_TOOL,
+		label: READ_MCP_RESOURCE_TOOL,
+		description: "Read a specific resource from an MCP server given the server name and resource URI.",
+		parameters: READ_PARAMETERS as unknown as TSchema,
+		outputSchema: READ_OUTPUT_SCHEMA as unknown as TSchema,
+		exposure: toToolExposure(options.exposure),
+		annotations: readOnly,
+		async execute(_toolCallId, params, signal) {
+			const serverName = stringArgument(params, "server");
+			const uri = stringArgument(params, "uri");
+			if (!serverName) throw new Error("server must be provided");
+			if (!uri) throw new Error("uri must be provided");
+			const server = findServer(serverName);
+			const result = await server.readResource(uri, { signal, timeoutMs: server.timeoutMs });
+			// Several contents (for example a directory) are labeled with their URIs.
+			const blocks: ContentBlock[] = result.contents.flatMap((contents) => [
+				...(result.contents.length > 1 ? [{ type: "text" as const, text: `${contents.uri}:` }] : []),
+				{ type: "resource" as const, resource: contents },
+			]);
+			const converted = await toModelContent(server.name, blocks);
+			const { content, fullOutputPath } = await limitMcpContent(
+				converted.length > 0 ? converted : [{ type: "text", text: `Resource ${uri} is empty.` }],
+			);
+			const contents = result.contents.map(({ _meta: _ignored, ...rest }) => rest);
+			return {
+				content,
+				details: {
+					server: server.name,
+					tool: READ_MCP_RESOURCE_TOOL,
+					...(fullOutputPath ? { fullOutputPath } : {}),
+				},
+				structuredContent: { server: server.name, uri, contents } as unknown as JsonValue,
+			};
+		},
+	};
+
+	return [listResources, listTemplates, readResource];
+}

third_party/pi/packages/coding-agent/src/extensions/mcp/runtime.lazy.ts

@@ -0,0 +1,2 @@
+/** Loads the MCP client, transports, and OAuth sign-in on first use (see runtime.ts). */
+export const loadMcpRuntime = () => import("./runtime.ts");

third_party/pi/packages/coding-agent/src/extensions/mcp/runtime.ts

@@ -0,0 +1,454 @@
+/**
+ * The part of the MCP integration that talks to servers: connections, transports, and OAuth
+ * sign-in. It pulls in the MCP client, so index.ts loads it through runtime.lazy.ts only when a
+ * server is configured.
+ */
+
+import { homedir } from "node:os";
+import { basename, join, resolve } from "node:path";
+import { setTimeout as sleep } from "node:timers/promises";
+import { pathToFileURL } from "node:url";
+import {
+	type AuthProvider,
+	type CallToolResult,
+	JSON_RPC_ERROR_CODES,
+	type ListResourcesResult,
+	type ListResourceTemplatesResult,
+	McpAuthRequiredError,
+	McpClient,
+	McpError,
+	McpHttpError,
+	type McpRequestOptions,
+	McpSessionExpiredError,
+	type Tool as McpTool,
+	type McpTransport,
+	type ReadResourceResult,
+	type Resource,
+	type ResourceTemplate,
+	StdioTransport,
+	StreamableHttpTransport,
+} from "@earendil-works/pi-mcp";
+import { McpOAuthAuthorizationRequiredError, type OAuthChallenge } from "@earendil-works/pi-mcp/oauth";
+import { VERSION } from "../../config.ts";
+import { resolveConfigValueOrThrow, resolveHeadersOrThrow } from "../../core/resolve-config-value.ts";
+import type { McpServerEntry } from "./config.ts";
+import type { McpServerLog } from "./log.ts";
+import {
+	createMcpAuthProvider,
+	type McpAuthProvider,
+	type McpOAuthCredentialStore,
+	type McpOAuthSettings,
+} from "./oauth.ts";
+import type { McpResourceServer } from "./resources.ts";
+import type { McpToolCaller } from "./tools.ts";
+
+export { McpServerLog } from "./log.ts";
+export { McpOAuthCredentialStore, McpSignInCancelledError, signInMcpServer } from "./oauth.ts";
+
+const DEFAULT_TIMEOUT_SECONDS = 60;
+const STDERR_TAIL_CHARS = 2_000;
+/** Delays between attempts to connect to an HTTP server that failed with a transient error. */
+const CONNECT_RETRY_DELAYS_MS = [250, 1_000];
+
+/**
+ * `disconnected`: the connection dropped (for example the stdio server exited); the next call
+ * reconnects.
+ */
+type ServerState = "connecting" | "connected" | "disconnected" | "needs-auth" | "failed" | "closed";
+
+export type McpTransportFactory = (
+	entry: McpServerEntry,
+	cwd: string,
+	authProvider: AuthProvider | undefined,
+) => McpTransport;
+
+function errorMessage(error: unknown): string {
+	return error instanceof Error ? error.message : String(error);
+}
+
+/** Network failures and overloaded or restarting servers, which are worth another attempt. */
+function isTransientError(error: unknown): boolean {
+	if (error instanceof McpHttpError) {
+		return error.status === 408 || error.status === 429 || (error.status >= 500 && error.status !== 501);
+	}
+	return error instanceof TypeError;
+}
+
+function signInRequiredMessage(entry: McpServerEntry): string {
+	const provider = "url" in entry.config ? entry.config.auth?.provider : undefined;
+	return `MCP server "${entry.name}" requires sign-in. Run ${provider ? `/login ${provider}` : "/mcp"} to sign in.`;
+}
+
+/** HTTP servers authenticate with OAuth unless the config supplies an `Authorization` header or `auth`. */
+function usesOAuth(entry: McpServerEntry): boolean {
+	const { config } = entry;
+	if (!("url" in config) || config.auth) return false;
+	return !Object.keys(config.headers ?? {}).some((header) => header.toLowerCase() === "authorization");
+}
+
+/** `~` and `~/…` (also `~\…` on Windows) name the home directory, like in a shell. */
+function expandHome(value: string): string {
+	if (value === "~") return homedir();
+	if (value.startsWith("~/") || (process.platform === "win32" && value.startsWith("~\\"))) {
+		return join(homedir(), value.slice(2));
+	}
+	return value;
+}
+
+export function createDefaultTransport(
+	entry: McpServerEntry,
+	cwd: string,
+	authProvider: AuthProvider | undefined,
+): McpTransport {
+	const { config, name } = entry;
+	if ("url" in config) {
+		return new StreamableHttpTransport({
+			url: config.url,
+			headers: resolveHeadersOrThrow(config.headers, `MCP server "${name}"`),
+			authProvider,
+		});
+	}
+	const env: Record<string, string> = {};
+	for (const [key, value] of Object.entries(config.env ?? {})) {
+		env[key] = resolveConfigValueOrThrow(value, `MCP server "${name}" env "${key}"`);
+	}
+	return new StdioTransport({
+		command: expandHome(config.command),
+		args: config.args?.map(expandHome),
+		cwd: resolve(cwd, expandHome(config.cwd ?? ".")),
+		env,
+		stderr: "pipe",
+	});
+}
+
+/** Servers that do not implement `resources/templates/list` have no templates. */
+async function withoutTemplates<T>(list: () => Promise<T>, empty: T): Promise<T> {
+	try {
+		return await list();
+	} catch (error) {
+		if (error instanceof McpError && error.code === JSON_RPC_ERROR_CODES.methodNotFound) return empty;
+		throw error;
+	}
+}
+
+function listTemplates(client: McpClient, options: McpRequestOptions = {}): Promise<ResourceTemplate[]> {
+	return withoutTemplates(() => client.listResourceTemplates(options), []);
+}
+
+/** One configured server. Reconnects lazily when a call finds the connection gone. */
+export class McpServerConnection implements McpToolCaller, McpResourceServer {
+	readonly entry: McpServerEntry;
+	state: ServerState = "connecting";
+	error: string | undefined;
+	tools: McpTool[] = [];
+	/**
+	 * Whether the server offers resources, which the resource tools list and read on demand. They are
+	 * not listed at connect: listing every resource can take long (a large drive, say).
+	 */
+	hasResources = false;
+	/** Server instructions from `initialize`, describing its tools as a group. */
+	instructions: string | undefined;
+	/** Last OAuth challenge from the server; sign-in uses its resource metadata URL and scope. */
+	challenge: OAuthChallenge | undefined;
+	private client: McpClient | undefined;
+	private opening: Promise<McpClient> | undefined;
+	/** Aborted by `close()`; cancels a connect in progress, including the wait between retries. */
+	private readonly shutdown = new AbortController();
+	/** Stderr of the last stdio server that failed to connect. */
+	private stderrTail: string | undefined;
+	private readonly cwd: string;
+	private readonly createTransport: McpTransportFactory;
+	private readonly authProvider: McpAuthProvider | undefined;
+	private readonly onTools: (connection: McpServerConnection) => void;
+	private readonly onChange: ((connection: McpServerConnection) => void) | undefined;
+	private readonly log: McpServerLog | undefined;
+
+	constructor(options: {
+		entry: McpServerEntry;
+		cwd: string;
+		createTransport: McpTransportFactory;
+		credentials: McpOAuthCredentialStore;
+		/** The current token of a pi provider, for servers with `auth.provider`. */
+		providerToken?: (provider: string) => Promise<string | undefined>;
+		onTools: (connection: McpServerConnection) => void;
+		/** Called when `state`, `error`, or `tools` change. */
+		onChange?: (connection: McpServerConnection) => void;
+		/** Receives the server's log messages (`notifications/message`). */
+		log?: McpServerLog;
+	}) {
+		this.entry = options.entry;
+		this.cwd = options.cwd;
+		this.createTransport = options.createTransport;
+		this.onTools = options.onTools;
+		this.onChange = options.onChange;
+		this.log = options.log;
+		const url = this.oauthUrl;
+		const provider = "url" in this.entry.config ? this.entry.config.auth?.provider : undefined;
+		this.authProvider = url
+			? createMcpAuthProvider({
+					serverUrl: url,
+					store: options.credentials.forServer(this.entry.name, url),
+					settings: () => this.oauthSettings(),
+					onChallenge: (challenge) => {
+						this.challenge = challenge;
+					},
+				})
+			: provider
+				? // Read on every request, so the provider's refreshes apply; MCP stores no copy.
+					{ token: async () => options.providerToken?.(provider), settled: async () => {} }
+				: undefined;
+	}
+
+	get name(): string {
+		return this.entry.name;
+	}
+
+	private get closed(): boolean {
+		return this.shutdown.signal.aborted;
+	}
+
+	get timeoutMs(): number {
+		return (this.entry.config.timeout ?? DEFAULT_TIMEOUT_SECONDS) * 1000;
+	}
+
+	/** Server URL when the server authenticates with OAuth. */
+	get oauthUrl(): string | undefined {
+		return usesOAuth(this.entry) && "url" in this.entry.config ? this.entry.config.url : undefined;
+	}
+
+	oauthSettings(): McpOAuthSettings {
+		const oauth = "url" in this.entry.config ? this.entry.config.oauth : undefined;
+		if (!oauth) return {};
+		return {
+			clientId: oauth.clientId,
+			clientSecret:
+				oauth.clientSecret === undefined
+					? undefined
+					: resolveConfigValueOrThrow(oauth.clientSecret, `MCP server "${this.entry.name}" oauth.clientSecret`),
+			callbackPort: oauth.callbackPort,
+			callbackUrl: oauth.callbackUrl,
+			scope: oauth.scope,
+			clientName: oauth.clientName,
+			clientRegistration: oauth.clientRegistration,
+			authServerMetadataUrl: oauth.authServerMetadataUrl ? new URL(oauth.authServerMetadataUrl) : undefined,
+		};
+	}
+
+	getClient(): Promise<McpClient> {
+		if (this.closed) return Promise.reject(new Error(`MCP server "${this.entry.name}" is shut down`));
+		if (this.client?.connectionState === "connected") return Promise.resolve(this.client);
+		this.opening ??= this.open().finally(() => {
+			this.opening = undefined;
+		});
+		return this.opening;
+	}
+
+	callTool(name: string, args: Record<string, unknown>, options: McpRequestOptions): Promise<CallToolResult> {
+		return this.withClient((client) => client.callTool(name, args, options));
+	}
+
+	readResource(uri: string, options: McpRequestOptions): Promise<ReadResourceResult> {
+		return this.withClient((client) => client.readResource(uri, options), true);
+	}
+
+	resourcesPage(cursor: string | undefined, options: McpRequestOptions): Promise<ListResourcesResult> {
+		return this.withClient((client) => client.listResourcesPage(cursor, options), true);
+	}
+
+	resourceTemplatesPage(cursor: string | undefined, options: McpRequestOptions): Promise<ListResourceTemplatesResult> {
+		return this.withClient(
+			(client) =>
+				withoutTemplates(() => client.listResourceTemplatesPage(cursor, options), { resourceTemplates: [] }),
+			true,
+		);
+	}
+
+	allResources(options: McpRequestOptions): Promise<Resource[]> {
+		return this.withClient((client) => client.listResources(options), true);
+	}
+
+	allResourceTemplates(options: McpRequestOptions): Promise<ResourceTemplate[]> {
+		return this.withClient((client) => listTemplates(client, options), true);
+	}
+
+	/**
+	 * Run a request, reconnecting when needed. `readOnly` requests are retried once after a transient
+	 * HTTP error; tool calls are not, since they may have run.
+	 */
+	private async withClient<T>(run: (client: McpClient) => Promise<T>, readOnly = false): Promise<T> {
+		for (let attempt = 1; ; attempt++) {
+			const client = await this.getClient();
+			try {
+				return await run(client);
+			} catch (error) {
+				if (readOnly && attempt === 1 && error instanceof McpHttpError && isTransientError(error)) {
+					await new Promise((resolve) => setTimeout(resolve, CONNECT_RETRY_DELAYS_MS[0]));
+					continue;
+				}
+				if (error instanceof McpSessionExpiredError && attempt === 1) {
+					// The server no longer knows the session (restart, deploy), so it did not run the request.
+					// Retry once on a new session. The old client is detached but not closed: closing would
+					// fail its other in-flight calls, which instead get the same 404 and retry the same way.
+					if (this.client === client) this.client = undefined;
+					continue;
+				}
+				if (!this.needsSignIn(error)) throw error;
+				await this.dropClient(client);
+				this.markNeedsAuth();
+				throw new Error(signInRequiredMessage(this.entry));
+			}
+		}
+	}
+
+	/** Connect again with fresh credentials, for example after signing in. */
+	async reconnect(): Promise<void> {
+		await this.opening?.catch(() => undefined);
+		if (this.client) await this.dropClient(this.client);
+		await this.getClient();
+	}
+
+	/** Disconnect after the stored credentials were removed. */
+	async signOut(): Promise<void> {
+		await this.opening?.catch(() => undefined);
+		if (this.client) await this.dropClient(this.client);
+		if (!this.closed) this.markNeedsAuth();
+	}
+
+	/** OAuth servers that still reject the request after a refresh need the user to sign in again. */
+	private needsSignIn(error: unknown): boolean {
+		return (
+			error instanceof McpOAuthAuthorizationRequiredError ||
+			(this.authProvider !== undefined && error instanceof McpAuthRequiredError)
+		);
+	}
+
+	private markNeedsAuth(): void {
+		this.state = "needs-auth";
+		this.error = undefined;
+		this.changed();
+	}
+
+	private changed(): void {
+		this.onChange?.(this);
+	}
+
+	private async dropClient(client: McpClient): Promise<void> {
+		if (this.client === client) this.client = undefined;
+		await client.close().catch(() => undefined);
+	}
+
+	private async open(): Promise<McpClient> {
+		this.state = "connecting";
+		this.changed();
+		const retries = "url" in this.entry.config ? CONNECT_RETRY_DELAYS_MS : [];
+		for (let attempt = 0; ; attempt++) {
+			this.stderrTail = undefined;
+			try {
+				return await this.connectOnce();
+			} catch (error) {
+				const delay = retries[attempt];
+				if (this.closed || delay === undefined || !isTransientError(error)) {
+					throw this.connectFailed(error);
+				}
+				await sleep(delay, undefined, { signal: this.shutdown.signal }).catch(() => undefined);
+				if (this.closed) throw this.connectFailed(error);
+			}
+		}
+	}
+
+	private async connectOnce(): Promise<McpClient> {
+		const client = new McpClient({
+			name: "pi",
+			version: VERSION,
+			requestTimeoutMs: this.timeoutMs,
+			roots: [{ uri: pathToFileURL(this.cwd).href, name: basename(this.cwd) }],
+		});
+		const log = this.log;
+		if (log) client.onNotification("notifications/message", (params) => log.write(this.entry.name, params));
+		// Closing the client fails its pending initialize or discovery requests and closes the transport.
+		// Memoized: a second `client.close()` returns before the first finished closing the transport.
+		let closing: Promise<void> | undefined;
+		const closeClient = () => {
+			closing ??= client.close().catch(() => undefined);
+			return closing;
+		};
+		this.shutdown.signal.addEventListener("abort", closeClient, { once: true });
+		let transport: McpTransport | undefined;
+		try {
+			transport = this.createTransport(this.entry, this.cwd, this.authProvider);
+			await client.connect(transport);
+			client.onNotification("notifications/tools/list_changed", () => {
+				void this.refreshTools(client);
+			});
+			const stdio = transport instanceof StdioTransport ? transport : undefined;
+			client.onClose(() => this.handleClientClose(client, stdio));
+			// Servers without the tools capability (prompts or resources only) do not answer tools/list.
+			const tools = client.serverCapabilities?.tools ? await client.listTools() : [];
+			if (this.closed) throw new Error("shut down while connecting");
+			if (client.connectionState !== "connected") throw new Error("connection closed during setup");
+			this.client = client;
+			this.tools = tools;
+			this.hasResources = client.serverCapabilities?.resources !== undefined;
+			this.instructions = client.instructions?.trim() || undefined;
+			this.state = "connected";
+			this.error = undefined;
+			this.onTools(this);
+			this.changed();
+			return client;
+		} catch (error) {
+			await closeClient();
+			if (transport instanceof StdioTransport) {
+				this.stderrTail = transport.stderr.trim().slice(-STDERR_TAIL_CHARS) || undefined;
+			}
+			throw error;
+		} finally {
+			this.shutdown.signal.removeEventListener("abort", closeClient);
+		}
+	}
+
+	private connectFailed(error: unknown): Error {
+		if (this.needsSignIn(error) && !this.closed) {
+			this.markNeedsAuth();
+			return new Error(signInRequiredMessage(this.entry));
+		}
+		this.state = this.closed ? "closed" : "failed";
+		this.error = this.stderrTail ? `${errorMessage(error)}\n${this.stderrTail}` : errorMessage(error);
+		this.changed();
+		return new Error(`MCP server "${this.entry.name}" failed to connect: ${this.error}`);
+	}
+
+	/** The transport dropped. The next call reconnects; until then the status shows why. */
+	private handleClientClose(client: McpClient, stdio: StdioTransport | undefined): void {
+		if (this.client !== client || this.closed) return;
+		this.client = undefined;
+		this.state = "disconnected";
+		const stderr = stdio?.stderr.trim().slice(-STDERR_TAIL_CHARS);
+		this.error = stderr ? `Connection closed\n${stderr}` : "Connection closed";
+		this.changed();
+	}
+
+	private async refreshTools(client: McpClient): Promise<void> {
+		try {
+			const tools = await client.listTools();
+			if (this.client !== client || this.closed) return;
+			this.tools = tools;
+			this.onTools(this);
+		} catch (error) {
+			this.error = `Failed to refresh tools: ${errorMessage(error)}`;
+		}
+		this.changed();
+	}
+
+	/** Resolves once no transport of this server is open, including one that was still connecting. */
+	async close(): Promise<void> {
+		this.shutdown.abort();
+		this.state = "closed";
+		this.changed();
+		const client = this.client;
+		this.client = undefined;
+		await Promise.all([client?.close().catch(() => undefined), this.opening?.catch(() => undefined)]);
+		// A refresh the server already answered may have rotated the refresh token; exiting before the
+		// new tokens are saved would lose the grant.
+		await this.authProvider?.settled();
+	}
+}

third_party/pi/packages/coding-agent/src/extensions/mcp/tools.ts

@@ -0,0 +1,391 @@
+/**
+ * Adapts MCP tools to pi tool definitions.
+ *
+ * Results map onto pi's model-facing content (text and images). Text over 20KB keeps its start and
+ * end with the middle cut out, like Codex does, and the full text is saved to a temp file the model
+ * can read. Binary resources other than images are saved to temp files too, and resource links name
+ * the `read_mcp_resource` tool. Codemode scripts receive the whole `CallToolResult` without `_meta`
+ * (`content` blocks as sent by the server, `structuredContent`, `isError`), never truncated: it is
+ * the tool's `structuredContent`, and every MCP tool declares a `CallToolResult` output schema. MCP
+ * errors (`isError`) are error results for the model, but scripts still resolve to the result.
+ */
+
+import { createHash } from "node:crypto";
+import type { AgentToolResult } from "@earendil-works/pi-agent-core";
+import type { ImageContent, JsonValue, TextContent } from "@earendil-works/pi-ai";
+import {
+	type CallToolResult,
+	type ContentBlock,
+	type McpRequestOptions,
+	type Tool as McpTool,
+	toLlmContent,
+} from "@earendil-works/pi-mcp";
+import { Container, Spacer, Text } from "@earendil-works/pi-tui";
+import type { TSchema } from "typebox";
+import type {
+	ToolAnnotations,
+	ToolDefinition,
+	ToolExposure,
+	ToolNamespace,
+	ToolRenderers,
+} from "../../core/extensions/types.ts";
+import { READ_MCP_RESOURCE_TOOL } from "../../core/mcp-servers.ts";
+import { formatToolCallWithArgs, getTextOutput, replaceTabs } from "../../core/tools/render-utils.ts";
+import { formatSize, truncateMiddle } from "../../core/tools/truncate.ts";
+import { keyHint } from "../../modes/interactive/components/keybinding-hints.ts";
+import { VisualLinePreview } from "../../modes/interactive/components/visual-truncate.ts";
+import { writeOutputFile } from "../../utils/output-files.ts";
+import type { McpExposure } from "./config.ts";
+
+/**
+ * Tool exposure of an MCP exposure. `codemode` and `deferred` both leave tools out of the codemode
+ * description; they differ only in which tool the MCP extension activates to reach them.
+ */
+export function toToolExposure(exposure: McpExposure): ToolExposure {
+	return exposure === "codemode" ? "deferred" : exposure;
+}
+
+/** Provider tool names are limited to 64 characters of `[A-Za-z0-9_-]`. */
+const MAX_TOOL_NAME_LENGTH = 64;
+/** Model-facing text of an MCP result beyond this is cut in the middle. */
+export const MCP_OUTPUT_MAX_BYTES = 20 * 1024;
+/** Visual (wrapped) result lines shown before the output is expanded. */
+const OUTPUT_PREVIEW_LINES = 5;
+/** Tool that reads the resources named by resource links. */
+export { READ_MCP_RESOURCE_TOOL };
+
+export interface McpToolDetails {
+	server: string;
+	tool: string;
+	/** Temp file with the full text output, when the model-facing text was truncated. */
+	fullOutputPath?: string;
+}
+
+/**
+ * Saves the full text of a truncated result, or a binary resource, and returns the file path.
+ * `extension` includes the dot, for example `.txt`.
+ */
+export type McpOutputSaver = (data: string | Uint8Array, extension: string) => Promise<string>;
+
+export function saveToTempFile(data: string | Uint8Array, extension: string): Promise<string> {
+	return writeOutputFile("pi-mcp", extension, data);
+}
+
+export interface McpToolCaller {
+	callTool(name: string, args: Record<string, unknown>, options: McpRequestOptions): Promise<CallToolResult>;
+}
+
+/**
+ * `mcp__<server>__<tool>`, sanitized and shortened with a hash suffix when too long. Like Codex,
+ * everything but `[A-Za-z0-9_]` becomes `_`, so the name is also the identifier codemode scripts
+ * call it by. `isTaken` reports names used by a different MCP tool: sanitizing can map two tools to
+ * one name (`a-b` and `a_b`), which then get the hash suffix.
+ */
+export function createMcpToolName(
+	server: string,
+	tool: string,
+	isTaken: (name: string) => boolean = () => false,
+): string {
+	const name = `mcp__${server}__${tool}`.replace(/[^A-Za-z0-9_]/g, "_");
+	if (name.length <= MAX_TOOL_NAME_LENGTH && !isTaken(name)) return name;
+	const hash = createHash("sha256").update(`${server}\0${tool}`).digest("hex").slice(0, 8);
+	return `${name.slice(0, MAX_TOOL_NAME_LENGTH - hash.length - 1)}_${hash}`;
+}
+
+function textOf(content: readonly (TextContent | ImageContent)[]): string {
+	return content
+		.filter((block): block is TextContent => block.type === "text")
+		.map((block) => block.text)
+		.join("\n");
+}
+
+/**
+ * MCP content blocks as one compact type with the fields of text and image blocks, which scripts
+ * read most (`image(block)` takes an image block as is). Rendered into every MCP tool's declaration.
+ */
+const CONTENT_BLOCK_SCHEMA = {
+	type: "object",
+	properties: {
+		type: { enum: ["text", "image", "audio", "resource_link", "resource"] },
+		text: { type: "string" },
+		data: { type: "string" },
+		mimeType: { type: "string" },
+	},
+	required: ["type"],
+	additionalProperties: true,
+};
+
+/** Keywords whose values are data, so a `$ref` key in them is not a reference. */
+const SCHEMA_DATA_KEYWORDS: ReadonlySet<string> = new Set(["const", "enum", "default", "examples"]);
+/** Keywords whose values map names to schemas, so their keys are names, not keywords. */
+const SCHEMA_MAP_KEYWORDS: ReadonlySet<string> = new Set([
+	"properties",
+	"patternProperties",
+	"dependentSchemas",
+	"dependencies",
+	"$defs",
+	"definitions",
+]);
+
+/**
+ * `schema` as a subschema at the JSON pointer `pointer` (`#/properties/structuredContent`) of
+ * another schema: its local references (`#`, `#/$defs/Item`) are rewritten to point below
+ * `pointer`, so they resolve to the same subschemas as before.
+ */
+function relocateSchema(schema: unknown, pointer: string): unknown {
+	if (Array.isArray(schema)) return schema.map((item) => relocateSchema(item, pointer));
+	if (typeof schema !== "object" || schema === null) return schema;
+	const relocated: Record<string, unknown> = {};
+	for (const [key, value] of Object.entries(schema)) {
+		if (key === "$ref" && typeof value === "string" && (value === "#" || value.startsWith("#/"))) {
+			relocated[key] = `${pointer}${value.slice(1)}`;
+		} else if (SCHEMA_DATA_KEYWORDS.has(key)) {
+			relocated[key] = value;
+		} else if (SCHEMA_MAP_KEYWORDS.has(key) && typeof value === "object" && value !== null && !Array.isArray(value)) {
+			relocated[key] = Object.fromEntries(
+				Object.entries(value).map(([name, subschema]) => [name, relocateSchema(subschema, pointer)]),
+			);
+		} else {
+			relocated[key] = relocateSchema(value, pointer);
+		}
+	}
+	return relocated;
+}
+
+/**
+ * Output schema of every MCP tool: the `CallToolResult` scripts receive (without `_meta`, which
+ * {@link convertMcpResult} drops), with the tool's own output schema as `structuredContent`.
+ */
+export function createMcpResultSchema(structuredContentSchema: Record<string, unknown> | undefined): TSchema {
+	return {
+		type: "object",
+		properties: {
+			content: { type: "array", items: CONTENT_BLOCK_SCHEMA },
+			// Servers may send structured content without declaring a schema for it, and it may be any JSON value.
+			structuredContent: structuredContentSchema
+				? relocateSchema(structuredContentSchema, "#/properties/structuredContent")
+				: {},
+			isError: { type: "boolean" },
+		},
+		required: ["content"],
+		// Fields newer protocol versions add.
+		additionalProperties: true,
+	} as unknown as TSchema;
+}
+
+/**
+ * Keep model-facing text within {@link MCP_OUTPUT_MAX_BYTES}. Longer text becomes one text block in
+ * Codex's truncation format, followed by the path of the file with the full text; images follow it.
+ */
+export async function limitMcpContent(
+	content: (TextContent | ImageContent)[],
+	saveOutput: McpOutputSaver = saveToTempFile,
+): Promise<{ content: (TextContent | ImageContent)[]; fullOutputPath?: string }> {
+	const combined = textOf(content);
+	const truncation = truncateMiddle(combined, MCP_OUTPUT_MAX_BYTES);
+	if (!truncation.truncated) return { content };
+	let fullOutputPath: string | undefined;
+	let where: string;
+	try {
+		fullOutputPath = await saveOutput(combined, ".txt");
+		where = `[Full output: ${fullOutputPath} (read it with offset/limit)]`;
+	} catch (error) {
+		where = `[Could not save the full output: ${error instanceof Error ? error.message : String(error)}]`;
+	}
+	const tokens = Math.ceil(truncation.totalBytes / 4);
+	const text = `Warning: truncated output (original token count: ${tokens})\nTotal output lines: ${truncation.totalLines}\n\n${truncation.content}\n\n${where}`;
+	return {
+		content: [{ type: "text", text }, ...content.filter((block) => block.type === "image")],
+		...(fullOutputPath ? { fullOutputPath } : {}),
+	};
+}
+
+export interface ConvertMcpResultOptions {
+	/** Saves truncated text and binary resources. Default: a temp file. */
+	saveOutput?: McpOutputSaver;
+	/** Whether the server's resources can be read with `read_mcp_resource`, which resource links then name. */
+	readableResources?: boolean;
+}
+
+/** File extension for a saved binary resource: the one its URI ends in, else `.bin`. */
+function extensionOf(uri: string): string {
+	const path = URL.canParse(uri) ? new URL(uri).pathname : uri;
+	return /\.[A-Za-z0-9]{1,8}$/.exec(path)?.[0] ?? ".bin";
+}
+
+/** Blobs of these types are shown as text. */
+function isTextMimeType(mimeType: string | undefined): boolean {
+	if (!mimeType) return false;
+	const type = mimeType.split(";", 1)[0].trim().toLowerCase();
+	return type.startsWith("text/") || type === "application/json" || type.endsWith("+json") || type.endsWith("+xml");
+}
+
+/** Model-facing content of one block of `server`'s result. */
+async function blockToContent(
+	server: string,
+	block: ContentBlock,
+	options: ConvertMcpResultOptions,
+): Promise<(TextContent | ImageContent)[]> {
+	if (block.type === "resource_link") {
+		const details = [block.mimeType, block.size === undefined ? undefined : formatSize(block.size)].filter(Boolean);
+		const read = options.readableResources ? `. Read it with ${READ_MCP_RESOURCE_TOOL} (server "${server}")` : "";
+		const description = block.description ? `: ${block.description}` : "";
+		return [
+			{
+				type: "text",
+				text: `[Resource ${block.uri} "${block.title ?? block.name}"${details.length > 0 ? ` (${details.join(", ")})` : ""}${description}${read}]`,
+			},
+		];
+	}
+	if (block.type === "resource" && "blob" in block.resource && !block.resource.mimeType?.startsWith("image/")) {
+		const { uri, mimeType, blob } = block.resource;
+		const data = Buffer.from(blob, "base64");
+		if (isTextMimeType(mimeType)) return [{ type: "text", text: data.toString("utf8") }];
+		const kind = `${mimeType ?? "unknown type"}, ${formatSize(data.length)}`;
+		try {
+			const path = await (options.saveOutput ?? saveToTempFile)(data, extensionOf(uri));
+			return [{ type: "text", text: `[Binary resource ${uri} (${kind}) saved to ${path}]` }];
+		} catch (error) {
+			const reason = error instanceof Error ? error.message : String(error);
+			return [{ type: "text", text: `[Binary resource ${uri} (${kind}) could not be saved: ${reason}]` }];
+		}
+	}
+	return toLlmContent({ content: [block] });
+}
+
+/** Model-facing content of `server`'s content blocks, before the output limit. */
+export async function toModelContent(
+	server: string,
+	blocks: readonly ContentBlock[],
+	options: ConvertMcpResultOptions = {},
+): Promise<(TextContent | ImageContent)[]> {
+	return (await Promise.all(blocks.map((block) => blockToContent(server, block, options)))).flat();
+}
+
+/** Convert an MCP result. `isError` results become error results that keep the structured result. */
+export async function convertMcpResult(
+	server: string,
+	tool: string,
+	result: CallToolResult,
+	options: ConvertMcpResultOptions = {},
+): Promise<AgentToolResult<McpToolDetails>> {
+	// Without content blocks, toLlmContent falls back to the structured content as JSON.
+	const converted: (TextContent | ImageContent)[] =
+		result.content.length > 0 ? await toModelContent(server, result.content, options) : toLlmContent(result);
+	if (result.isError && textOf(converted) === "") {
+		converted.push({ type: "text", text: `MCP tool ${server}/${tool} returned an error` });
+	}
+	const { content, fullOutputPath } = await limitMcpContent(converted, options.saveOutput);
+	const { _meta: _ignored, ...scriptResult } = result;
+	return {
+		content,
+		details: { server, tool, ...(fullOutputPath ? { fullOutputPath } : {}) },
+		structuredContent: scriptResult as unknown as JsonValue,
+		...(result.isError ? { isError: true } : {}),
+	};
+}
+
+/**
+ * Tool input schemas must be objects. MCP servers may omit `type`, and some providers reject object
+ * schemas without `properties`.
+ */
+function toParameters(schema: Record<string, unknown>): TSchema {
+	return {
+		...schema,
+		type: schema.type ?? "object",
+		...(schema.properties === undefined ? { properties: {} } : {}),
+	} as unknown as TSchema;
+}
+
+const ANNOTATION_HINTS = ["readOnlyHint", "destructiveHint", "idempotentHint", "openWorldHint"] as const;
+
+/** The boolean hints of an MCP tool's annotations, or undefined when it has none. */
+function toToolAnnotations(tool: McpTool): ToolAnnotations | undefined {
+	const annotations: ToolAnnotations = {};
+	for (const hint of ANNOTATION_HINTS) {
+		const value = tool.annotations?.[hint];
+		if (typeof value === "boolean") annotations[hint] = value;
+	}
+	return Object.keys(annotations).length > 0 ? annotations : undefined;
+}
+
+export function createMcpToolDefinition(options: {
+	server: string;
+	tool: McpTool;
+	name: string;
+	exposure: McpExposure;
+	namespace: ToolNamespace;
+	timeoutMs: number;
+	getClient: () => Promise<McpToolCaller>;
+	/** Whether `read_mcp_resource` can read the server's resources. */
+	readableResources?: () => boolean;
+}): ToolDefinition<TSchema, McpToolDetails> {
+	const { server, tool } = options;
+	const title = tool.title ?? tool.annotations?.title;
+	const annotations = toToolAnnotations(tool);
+	const label = `${server}/${tool.name}`;
+	return {
+		name: options.name,
+		label,
+		description: tool.description?.trim() || title || `MCP tool ${tool.name} from server ${server}`,
+		parameters: toParameters(tool.inputSchema),
+		outputSchema: createMcpResultSchema(tool.outputSchema),
+		exposure: toToolExposure(options.exposure),
+		namespace: options.namespace,
+		...(annotations ? { annotations } : {}),
+		...createMcpToolRenderers(label),
+		async execute(_toolCallId, params, signal, onUpdate) {
+			const client = await options.getClient();
+			const result = await client.callTool(tool.name, (params ?? {}) as Record<string, unknown>, {
+				signal,
+				timeoutMs: options.timeoutMs,
+				onProgress: (progress) => {
+					const total = progress.total === undefined ? "" : `/${progress.total}`;
+					const text = progress.message ?? `Progress ${progress.progress}${total}`;
+					onUpdate?.({ content: [{ type: "text", text }], details: { server, tool: tool.name } });
+				},
+			});
+			return convertMcpResult(server, tool.name, result, { readableResources: options.readableResources?.() });
+		},
+	};
+}
+
+/** Renderers of calls to an MCP tool, labeled `server/tool`, also used before the tool is registered. */
+export function createMcpToolRenderers(label: string): ToolRenderers {
+	return {
+		renderCall(args, theme, context) {
+			const component = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0);
+			component.setText(formatToolCallWithArgs(label, args, theme, context.expanded));
+			return component;
+		},
+		renderResult(result, options, theme, context) {
+			const component = (context.lastComponent as Container | undefined) ?? new Container();
+			component.clear();
+			const output = getTextOutput(result, context.showImages).trim();
+			if (!output) return component;
+			const color = context.isError ? "error" : "toolOutput";
+			const styled = replaceTabs(output)
+				.split("\n")
+				.map((line) => theme.fg(color, line))
+				.join("\n");
+			component.addChild(new Spacer(1));
+			if (options.expanded) {
+				component.addChild(new Text(styled, 0, 0));
+			} else {
+				// Limit wrapped lines, not logical ones: MCP results are often one long JSON line.
+				component.addChild(
+					new VisualLinePreview({
+						text: styled,
+						maxVisualLines: OUTPUT_PREVIEW_LINES,
+						keep: "start",
+						formatHint: (hidden) =>
+							`${theme.fg("muted", `... (${hidden} more lines,`)} ${keyHint("app.tools.expand", "to expand")}${theme.fg("muted", ")")}`,
+					}),
+				);
+				const fullOutputPath = (result.details as McpToolDetails | undefined)?.fullOutputPath;
+				if (fullOutputPath) component.addChild(new Text(theme.fg("muted", `Full output: ${fullOutputPath}`), 0, 0));
+			}
+			return component;
+		},
+	};
+}

third_party/pi/packages/coding-agent/src/extensions/mcp/ui.ts

@@ -0,0 +1,259 @@
+/**
+ * The `/mcp` manager view: menus that rebuild while servers connect, a read-only status screen, and
+ * the sign-in screen that accepts a pasted redirect URL.
+ */
+
+import {
+	type Component,
+	Container,
+	type Focusable,
+	Input,
+	type SelectItem,
+	SelectList,
+	Spacer,
+	Text,
+	type TUI,
+	truncateToWidth,
+	visibleWidth,
+} from "@earendil-works/pi-tui";
+import type { ExtensionCommandContext } from "../../core/extensions/types.ts";
+import type { KeybindingsManager } from "../../core/keybindings.ts";
+import { AuthUrlComponent } from "../../modes/interactive/components/auth-url.ts";
+import { DynamicBorder } from "../../modes/interactive/components/dynamic-border.ts";
+import { keyHint } from "../../modes/interactive/components/keybinding-hints.ts";
+import { getSelectListTheme, type Theme } from "../../modes/interactive/theme/theme.ts";
+
+export interface McpMenu {
+	title: string;
+	/** Shown below the title. */
+	details?: string;
+	/** Shown below the details in the error color. */
+	error?: string;
+	items: SelectItem[];
+	/** Shown when there are no items. */
+	empty?: string;
+	/** Value of the item selected when the menu opens. */
+	selected?: string;
+	/** What the confirm key does, for the key hint. */
+	confirmLabel: string;
+	/** What the cancel key does, for the key hint. */
+	cancelLabel: string;
+}
+
+export interface McpUi {
+	/**
+	 * Show a menu and resolve to the chosen item's value, or undefined when cancelled. `subscribe`
+	 * rebuilds the menu on every change, keeping the selected item.
+	 */
+	menu(build: () => McpMenu, subscribe?: (listener: () => void) => () => void): Promise<string | undefined>;
+	/** Show a message while an operation runs. With `onCancel`, the cancel key calls it. */
+	status(title: string, message: string, onCancel?: () => void): void;
+	/**
+	 * Show the authorization URL and wait for a pasted redirect URL. Resolves to undefined when
+	 * cancelled or when `signal` aborts (the browser reached the callback).
+	 */
+	redirectUrl(title: string, authorizationUrl: string, signal: AbortSignal): Promise<string | undefined>;
+}
+
+function frame(theme: Theme, title: string, body: Component[], footer?: string): Container {
+	const container = new Container();
+	container.addChild(new DynamicBorder((text: string) => theme.fg("accent", text)));
+	container.addChild(new Text(theme.fg("accent", theme.bold(title)), 1, 0));
+	for (const child of body) container.addChild(child);
+	if (footer) {
+		container.addChild(new Spacer(1));
+		container.addChild(new Text(theme.fg("dim", footer), 1, 0));
+	}
+	container.addChild(new DynamicBorder((text: string) => theme.fg("accent", text)));
+	return container;
+}
+
+const MAX_VISIBLE_ITEMS = 12;
+
+export class McpManagerView implements McpUi, Component, Focusable {
+	private readonly tui: TUI;
+	private readonly theme: Theme;
+	private readonly keybindings: KeybindingsManager;
+	private content: Container;
+	private inputHandler: ((data: string) => void) | undefined;
+	private inputTarget: Focusable | undefined;
+	private _focused = false;
+
+	constructor(tui: TUI, theme: Theme, keybindings: KeybindingsManager) {
+		this.tui = tui;
+		this.theme = theme;
+		this.keybindings = keybindings;
+		this.content = frame(theme, "MCP servers", [new Text(theme.fg("muted", "Loading…"), 1, 1)]);
+	}
+
+	get focused(): boolean {
+		return this._focused;
+	}
+
+	set focused(value: boolean) {
+		this._focused = value;
+		if (this.inputTarget) this.inputTarget.focused = value;
+	}
+
+	private setContent(content: Container, inputHandler?: (data: string) => void, inputTarget?: Focusable): void {
+		if (this.inputTarget) this.inputTarget.focused = false;
+		this.content = content;
+		this.inputHandler = inputHandler;
+		this.inputTarget = inputTarget;
+		if (this.inputTarget) this.inputTarget.focused = this._focused;
+		this.tui.requestRender();
+	}
+
+	menu(build: () => McpMenu, subscribe?: (listener: () => void) => () => void): Promise<string | undefined> {
+		return new Promise((resolve) => {
+			let unsubscribe: (() => void) | undefined;
+			let settled = false;
+			const finish = (value: string | undefined) => {
+				if (settled) return;
+				settled = true;
+				unsubscribe?.();
+				resolve(value);
+			};
+			let selected: string | undefined;
+			const render = () => {
+				const menu = build();
+				const wanted = selected ?? menu.selected;
+				const body: Component[] = [];
+				if (menu.details) body.push(new Text(this.theme.fg("muted", menu.details), 1, 0));
+				if (menu.error) body.push(new Text(this.theme.fg("error", menu.error), 1, 0));
+				body.push(new Spacer(1));
+				const footer = `${keyHint("tui.select.confirm", menu.confirmLabel)} • ${keyHint("tui.select.cancel", menu.cancelLabel)}`;
+				if (menu.items.length === 0) {
+					body.push(new Text(this.theme.fg("muted", menu.empty ?? "Nothing to show."), 1, 0));
+					this.setContent(
+						frame(this.theme, menu.title, body, keyHint("tui.select.cancel", menu.cancelLabel)),
+						(data) => {
+							if (this.keybindings.matches(data, "tui.select.cancel")) finish(undefined);
+						},
+					);
+					return;
+				}
+				const list = new SelectList(
+					menu.items,
+					Math.min(menu.items.length, MAX_VISIBLE_ITEMS),
+					getSelectListTheme(),
+				);
+				const index = menu.items.findIndex((item) => item.value === wanted);
+				if (index !== -1) list.setSelectedIndex(index);
+				selected = list.getSelectedItem()?.value;
+				list.onSelectionChange = (item) => {
+					selected = item.value;
+				};
+				list.onSelect = (item) => finish(item.value);
+				list.onCancel = () => finish(undefined);
+				body.push(list);
+				this.setContent(frame(this.theme, menu.title, body, footer), (data) => list.handleInput(data));
+			};
+			render();
+			unsubscribe = subscribe?.(() => {
+				if (!settled) render();
+			});
+		});
+	}
+
+	status(title: string, message: string, onCancel?: () => void): void {
+		const body = [new Spacer(1), new Text(this.theme.fg("muted", message), 1, 0)];
+		if (!onCancel) {
+			this.setContent(frame(this.theme, title, body));
+			return;
+		}
+		this.setContent(frame(this.theme, title, body, keyHint("tui.select.cancel", "cancel")), (data) => {
+			if (this.keybindings.matches(data, "tui.select.cancel")) onCancel();
+		});
+	}
+
+	redirectUrl(title: string, authorizationUrl: string, signal: AbortSignal): Promise<string | undefined> {
+		return new Promise((resolve) => {
+			let settled = false;
+			const finish = (value: string | undefined) => {
+				if (settled) return;
+				settled = true;
+				signal.removeEventListener("abort", onAbort);
+				resolve(value);
+			};
+			const onAbort = () => finish(undefined);
+			if (signal.aborted) {
+				finish(undefined);
+				return;
+			}
+			signal.addEventListener("abort", onAbort, { once: true });
+			const input = new Input();
+			const link = new AuthUrlComponent(this.tui, authorizationUrl);
+			const body: Component[] = [
+				new Spacer(1),
+				new Text(this.theme.fg("muted", "Approve access in your browser. If it did not open, visit:"), 1, 0),
+				link,
+				new Spacer(1),
+				new Text(
+					this.theme.fg("muted", "If the browser runs on another machine, paste the URL it was redirected to:"),
+					1,
+					0,
+				),
+				input,
+			];
+			this.setContent(
+				frame(
+					this.theme,
+					title,
+					body,
+					`${keyHint("tui.select.confirm", "submit")} • ${keyHint("tui.select.cancel", "cancel")}`,
+				),
+				(data) => {
+					if (this.keybindings.matches(data, "tui.select.confirm")) {
+						const value = input.getValue().trim();
+						if (value) finish(value);
+						return;
+					}
+					if (this.keybindings.matches(data, "tui.select.cancel")) {
+						finish(undefined);
+						return;
+					}
+					if (this.keybindings.matches(data, "app.message.copy")) {
+						void link.copy();
+						return;
+					}
+					input.handleInput(data);
+				},
+				input,
+			);
+		});
+	}
+
+	handleInput(data: string): void {
+		this.inputHandler?.(data);
+		this.tui.requestRender();
+	}
+
+	render(width: number): string[] {
+		return this.content
+			.render(width)
+			.map((line) => (visibleWidth(line) > width ? truncateToWidth(line, width, "") : line));
+	}
+
+	invalidate(): void {
+		this.content.invalidate();
+	}
+}
+
+/** Run `manage` in the manager view until it returns. */
+export async function showMcpManager(
+	ctx: ExtensionCommandContext,
+	manage: (ui: McpUi) => Promise<void>,
+): Promise<void> {
+	await ctx.ui.custom<void>((tui, theme, keybindings, done) => {
+		const view = new McpManagerView(tui, theme, keybindings);
+		void manage(view).then(
+			() => done(),
+			(error: unknown) => {
+				ctx.ui.notify(error instanceof Error ? error.message : String(error), "error");
+				done();
+			},
+		);
+		return view;
+	});
+}

third_party/pi/packages/coding-agent/src/extensions/tool-search/index.ts

@@ -0,0 +1,18 @@
+/**
+ * The `tool_search` tool as an extension. The CLI loads it as a built-in extension; SDK users add
+ * `createToolSearchExtension()` to their extension factories.
+ *
+ * `tool_search` is registered inactive. Activate it with `--tools`, the `defaultTools` setting, or
+ * `setActiveTools()`.
+ */
+
+import type { ExtensionFactory } from "../../core/extensions/types.ts";
+import { createToolSearchToolDefinition } from "./tool.ts";
+
+export function createToolSearchExtension(): ExtensionFactory {
+	return (pi) => {
+		pi.registerTool({ ...createToolSearchToolDefinition({ tools: pi }), defaultActive: false });
+	};
+}
+
+export default createToolSearchExtension();

third_party/pi/packages/coding-agent/src/extensions/tool-search/tool.ts

@@ -0,0 +1,244 @@
+/**
+ * Tool discovery: a BM25 ranker over tool metadata, shared by `searchTools()` in codemode scripts and
+ * the optional `tool_search` tool.
+ *
+ * `tool_search` searches tools that are not declared to
+ * the model (`codemode` and `deferred` exposure) and loads the matches, so they are declared for the
+ * next model call. Loading goes through the active tool set, so it is recorded in the transcript
+ * like any other tool change and survives `/tree`, resume, and fork on that branch.
+ */
+
+import { type Static, Type } from "typebox";
+import type {
+	ExtensionAPI,
+	ToolDefinition,
+	ToolExposure,
+	ToolInfo,
+	ToolNamespace,
+} from "../../core/extensions/types.ts";
+
+export const TOOL_SEARCH_TOOL_NAME = "tool_search";
+export const DEFAULT_TOOL_SEARCH_LIMIT = 8;
+
+/** A tool as the ranker sees it: its name and the text built by {@link createToolSearchDocument}. */
+export interface ToolSearchDocument {
+	name: string;
+	text: string;
+}
+
+export interface ToolSearchMatch {
+	name: string;
+	score: number;
+}
+
+/** Ranks tools for a query. BM25 today; a hybrid ranker with embeddings can replace it. */
+export interface ToolRanker {
+	rank(query: string, documents: readonly ToolSearchDocument[], limit: number): ToolSearchMatch[];
+}
+
+const STOP_WORDS: ReadonlySet<string> = new Set([
+	"a",
+	"an",
+	"and",
+	"are",
+	"as",
+	"at",
+	"be",
+	"by",
+	"for",
+	"from",
+	"in",
+	"is",
+	"it",
+	"of",
+	"on",
+	"or",
+	"that",
+	"the",
+	"this",
+	"to",
+	"with",
+]);
+
+/** Naive singular form, so `issues` matches `issue` and `searches` matches `search`. */
+function stem(term: string): string {
+	if (term.length > 4 && term.endsWith("ies")) return `${term.slice(0, -3)}y`;
+	if (term.length > 4 && /(ches|shes|sses|xes|zes)$/.test(term)) return term.slice(0, -2);
+	if (term.length > 3 && term.endsWith("s") && !term.endsWith("ss")) return term.slice(0, -1);
+	return term;
+}
+
+/** Lowercase terms, split at camelCase boundaries and non-alphanumerics, without stop words. */
+export function tokenize(text: string): string[] {
+	return text
+		.replace(/([a-z0-9])([A-Z])/g, "$1 $2")
+		.replace(/([A-Z]+)([A-Z][a-z])/g, "$1 $2")
+		.toLowerCase()
+		.split(/[^a-z0-9]+/)
+		.filter((term) => term.length > 0 && !STOP_WORDS.has(term))
+		.map(stem);
+}
+
+function isObject(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+/** Schema descriptions and property names, recursively. */
+function schemaText(schema: unknown, parts: string[]): void {
+	if (!isObject(schema)) return;
+	if (typeof schema.description === "string") parts.push(schema.description);
+	if (isObject(schema.properties)) {
+		for (const [name, property] of Object.entries(schema.properties)) {
+			parts.push(name);
+			schemaText(property, parts);
+		}
+	}
+	schemaText(schema.items, parts);
+	for (const key of ["anyOf", "oneOf", "allOf"]) {
+		const variants = schema[key];
+		if (Array.isArray(variants)) for (const variant of variants) schemaText(variant, parts);
+	}
+}
+
+/**
+ * Search text of a tool: the name, the name with `_`
+ * as spaces, the description, schema descriptions and property names, and the namespace with its
+ * description and instructions.
+ */
+export function createToolSearchDocument(
+	tool: Pick<ToolInfo, "name" | "description" | "parameters">,
+	namespace?: ToolNamespace,
+): ToolSearchDocument {
+	const parts = [tool.name, tool.name.replaceAll("_", " "), tool.description];
+	schemaText(tool.parameters, parts);
+	if (namespace) parts.push(namespace.name, namespace.description ?? "", namespace.instructions ?? "");
+	return { name: tool.name, text: parts.filter((part) => part.trim()).join(" ") };
+}
+
+/** Okapi BM25 with the usual parameters. Ties keep document order. */
+export class Bm25Ranker implements ToolRanker {
+	private readonly k1: number;
+	private readonly b: number;
+
+	constructor(options: { k1?: number; b?: number } = {}) {
+		this.k1 = options.k1 ?? 1.2;
+		this.b = options.b ?? 0.75;
+	}
+
+	rank(query: string, documents: readonly ToolSearchDocument[], limit: number): ToolSearchMatch[] {
+		const queryTerms = [...new Set(tokenize(query))];
+		if (queryTerms.length === 0 || documents.length === 0 || limit <= 0) return [];
+		const termCounts = documents.map((document) => {
+			const counts = new Map<string, number>();
+			for (const term of tokenize(document.text)) counts.set(term, (counts.get(term) ?? 0) + 1);
+			return counts;
+		});
+		const lengths = termCounts.map((counts) => [...counts.values()].reduce((sum, count) => sum + count, 0));
+		const averageLength = lengths.reduce((sum, length) => sum + length, 0) / documents.length || 1;
+		const idf = new Map(
+			queryTerms.map((term) => {
+				const frequency = termCounts.filter((counts) => counts.has(term)).length;
+				return [term, Math.log(1 + (documents.length - frequency + 0.5) / (frequency + 0.5))] as const;
+			}),
+		);
+		const matches: ToolSearchMatch[] = [];
+		documents.forEach((document, index) => {
+			let score = 0;
+			for (const term of queryTerms) {
+				const count = termCounts[index].get(term);
+				if (!count) continue;
+				const norm = this.k1 * (1 - this.b + (this.b * lengths[index]) / averageLength);
+				score += (idf.get(term) ?? 0) * ((count * (this.k1 + 1)) / (count + norm));
+			}
+			if (score > 0) matches.push({ name: document.name, score });
+		});
+		return matches.sort((a, b) => b.score - a.score).slice(0, limit);
+	}
+}
+
+export const toolSearchSchema = Type.Object({
+	query: Type.String({ description: "Search query for deferred tools." }),
+	limit: Type.Optional(
+		Type.Number({ description: `Maximum number of tools to return. Defaults to ${DEFAULT_TOOL_SEARCH_LIMIT}.` }),
+	),
+});
+
+export type ToolSearchInput = Static<typeof toolSearchSchema>;
+
+export interface ToolSearchResultTool {
+	name: string;
+	description: string;
+}
+
+export interface ToolSearchToolDetails {
+	/** Tools loaded by this call. */
+	loaded: string[];
+}
+
+export interface ToolSearchToolOptions {
+	/**
+	 * The session's tools. `tool_search` searches the tools that are not declared to the model and
+	 * activates the matches. Without it, the tool finds nothing. An `ExtensionAPI` fits.
+	 */
+	tools?: Pick<ExtensionAPI, "getAllTools" | "getActiveTools" | "setActiveTools">;
+}
+
+/** Whether `tool_search` can load a tool with this exposure. */
+function isSearchable(exposure: ToolExposure): boolean {
+	return exposure === "codemode" || exposure === "deferred";
+}
+
+/**
+ * Rank the searchable tools that are not active yet and activate the matches, so the next model
+ * call declares them. Activation is recorded in the transcript like any tool change.
+ */
+function searchAndLoad(
+	tools: NonNullable<ToolSearchToolOptions["tools"]>,
+	query: string,
+	limit: number,
+): ToolSearchResultTool[] {
+	const active = tools.getActiveTools();
+	const candidates = tools.getAllTools().filter((tool) => isSearchable(tool.exposure) && !active.includes(tool.name));
+	const documents = candidates.map((tool) => createToolSearchDocument(tool, tool.namespace));
+	const matches = new Bm25Ranker().rank(query, documents, limit);
+	if (matches.length > 0) tools.setActiveTools([...active, ...matches.map((match) => match.name)]);
+	return matches.map((match) => ({
+		name: match.name,
+		description: candidates.find((tool) => tool.name === match.name)?.description ?? "",
+	}));
+}
+
+/**
+ * The `tool_search` description. It does not list the searchable tools or their namespaces, so it
+ * stays the same while tools are registered, for example when MCP servers connect.
+ */
+export const TOOL_SEARCH_DESCRIPTION = `# Tool discovery\n\nSearches over deferred tool metadata with BM25 and exposes matching tools for the next model call.\n\nSome of the tools, such as tools of MCP servers, may not have been provided to you upfront, and you should use this tool (\`${TOOL_SEARCH_TOOL_NAME}\`) to search for the required tools. For MCP tool discovery, always use \`${TOOL_SEARCH_TOOL_NAME}\`.`;
+
+export function createToolSearchToolDefinition(
+	options: ToolSearchToolOptions = {},
+): ToolDefinition<typeof toolSearchSchema, ToolSearchToolDetails> {
+	return {
+		name: TOOL_SEARCH_TOOL_NAME,
+		label: TOOL_SEARCH_TOOL_NAME,
+		description: TOOL_SEARCH_DESCRIPTION,
+		promptSnippet: "Search for tools that are not loaded yet and load the matches",
+		parameters: toolSearchSchema,
+		// Searching is not something scripts need; it changes what the model sees.
+		exposure: "model-only",
+		// A query can match any tool that is not declared yet.
+		reachesTools: () => "all",
+		async execute(_toolCallId, { query, limit }) {
+			if (query.trim() === "") throw new Error("query must not be empty");
+			const max = limit ?? DEFAULT_TOOL_SEARCH_LIMIT;
+			if (!Number.isInteger(max) || max <= 0) throw new Error("limit must be a positive integer");
+			const tools = options.tools ? searchAndLoad(options.tools, query, max) : [];
+			const text =
+				tools.length === 0
+					? "No matching tools found."
+					: `Loaded ${tools.length} tool${tools.length === 1 ? "" : "s"}. They are available from your next call:\n${tools
+							.map((tool) => `- ${tool.name}: ${tool.description.trim().split(/\r?\n/)[0]}`)
+							.join("\n")}`;
+			return { content: [{ type: "text", text }], details: { loaded: tools.map((tool) => tool.name) } };
+		},
+	};
+}