one-code-extension 0.1.0

package/extensions/mcp/index.ts

@@ -0,0 +1,373 @@
+/**
+ * mcp extension — Model Context Protocol client, which pi deliberately omits.
+ *
+ * Servers are configured in Claude Code's format (`.mcp.json`, `~/.claude.json`)
+ * and their tools are registered as `mcp__<server>__<tool>`, so existing Claude
+ * Code MCP setups and permission rules work unchanged.
+ *
+ * Every MCP tool is registered deferred: a handful of servers can contribute
+ * dozens of tools, and putting all those schemas in the system prompt is exactly
+ * what `tool_search` exists to avoid.
+ *
+ * Servers are connected once per session (their tool lists are needed to
+ * register anything) and closed on `session_shutdown`. In the interactive main
+ * session the connect runs in the background so remote servers cannot delay
+ * the prompt; one-shots and subagent children await it (findings §15).
+ */
+
+import os from "node:os";
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { Type } from "typebox";
+import { DEFER_CHANNEL } from "../lib/deferred.ts";
+import { persistIfLarge } from "../lib/persisted-output.ts";
+import { ccToolRenderers } from "../lib/tui-render.ts";
+import {
+	callTool,
+	close,
+	type Connection,
+	connect,
+	type FailedConnection,
+	mcpInstructionsReminder,
+	readResource,
+	readResourceDir,
+} from "./client.ts";
+import { CONTEXT_ORDER, REMINDER_CHANNEL } from "../lib/reminders.ts";
+import { join } from "node:path";
+import { discoverPlugins } from "../lib/plugins.ts";
+import { loadServers, type McpServer } from "./config.ts";
+import {
+	describeContent,
+	describeResourceContents,
+	jsonSchemaToTypeBox,
+	type McpContentBlock,
+	type McpResourceContents,
+	namespacedToolName,
+} from "./schema.ts";
+
+/** Shared "server name not found" tool result for the resource read/list-dir tools. */
+function unknownServerError(name: string) {
+	return {
+		content: [{ type: "text" as const, text: `No connected MCP server named "${name}".` }],
+		details: {} as { server?: string; uri?: string },
+		isError: true,
+	};
+}
+
+export default function mcpExtension(pi: ExtensionAPI) {
+	const connections = new Map<string, Connection>();
+	const failures: FailedConnection[] = [];
+	const registered = new Set<string>();
+	let servers: McpServer[] = [];
+
+	/**
+	 * Where oversized results are persisted — the session dir, matching Claude
+	 * Code's `<session-dir>/tool-results/<id>.txt`. A session-less run (`--no-session`)
+	 * has no such dir, so a temp folder serves; the model gets the path either way.
+	 */
+	const resultsDir = (ctx: ExtensionContext | undefined): string => {
+		try {
+			const dir = ctx?.sessionManager.getSessionDir();
+			if (dir) return dir;
+		} catch {
+			// fall through to the temp folder
+		}
+		return join(os.tmpdir(), "one-code");
+	};
+
+	const registerToolsFor = (connection: Connection) => {
+		for (const tool of connection.tools) {
+			const name = namespacedToolName(connection.server.name, tool.name);
+			if (registered.has(name)) continue;
+			registered.add(name);
+
+			pi.registerTool({
+				name,
+				label: `${connection.server.name}: ${tool.name}`,
+				...ccToolRenderers(`${connection.server.name}: ${tool.name}`),
+				description: tool.description ?? `MCP tool "${tool.name}" from server "${connection.server.name}".`,
+				parameters: jsonSchemaToTypeBox(tool.inputSchema),
+				async execute(toolCallId, params, _signal, _onUpdate, ctx) {
+					const live = connections.get(connection.server.name);
+					if (!live) {
+						return {
+							content: [{ type: "text", text: `MCP server "${connection.server.name}" is not connected.` }],
+							details: {} as Record<string, unknown>,
+							isError: true,
+						};
+					}
+					try {
+						const result = await callTool(live, tool.name, (params ?? {}) as Record<string, unknown>);
+						const { text, images } = describeContent(result.content as McpContentBlock[] | undefined);
+						return {
+							content: [
+								{
+									type: "text",
+									text: text ? persistIfLarge(text, { dir: resultsDir(ctx), id: toolCallId }) : "(no output)",
+								},
+								...images.map((image) => ({ type: "image" as const, data: image.data, mimeType: image.mimeType })),
+							],
+							details: { server: connection.server.name, tool: tool.name } as Record<string, unknown>,
+							isError: result.isError === true,
+						};
+					} catch (error) {
+						return {
+							content: [{ type: "text", text: `MCP call failed: ${(error as Error).message}` }],
+							details: { server: connection.server.name, tool: tool.name } as Record<string, unknown>,
+							isError: true,
+						};
+					}
+				},
+			});
+
+			// Deferred: the model discovers these through tool_search.
+			pi.events.emit(DEFER_CHANNEL, {
+				name,
+				keywords: [connection.server.name, tool.name.replace(/[_-]/g, " "), "mcp"],
+			});
+		}
+	};
+
+	/** Resolves when every configured server has connected or failed. */
+	let connecting: Promise<void> | undefined;
+	let connectSettled = false;
+	let shuttingDown = false;
+
+	const connectAll = async (ctx: ExtensionContext) => {
+		servers = loadServers(ctx.cwd, os.homedir(), process.env, discoverPlugins(join(os.homedir(), ".claude")).mcpConfigs);
+		if (servers.length === 0) return;
+
+		// A server whose credentials are unset would connect with an empty token and
+		// fail with a confusing protocol error; skip it and say why instead.
+		for (const server of servers.filter((s) => s.missingEnv?.length)) {
+			failures.push({
+				server,
+				error: `not started — ${server.missingEnv?.join(", ")} not set in the environment`,
+			});
+		}
+
+		const results = await Promise.all(
+			servers
+				.filter((server) => !server.missingEnv?.length)
+				.map(async (server) => {
+				try {
+					return { connection: await connect(server) };
+				} catch (error) {
+					return { failure: { server, error: (error as Error).message } };
+				}
+			}),
+		);
+
+		// The session may have shut down while a slow server was still answering;
+		// a connection landing now would leak its transport, so close it instead.
+		if (shuttingDown) {
+			await Promise.all(results.map((result) => (result.connection ? close(result.connection) : undefined)));
+			return;
+		}
+
+		for (const result of results) {
+			if (result.connection) {
+				connections.set(result.connection.server.name, result.connection);
+				registerToolsFor(result.connection);
+				for (const warning of result.connection.warnings) {
+					failures.push({ server: result.connection.server, error: warning });
+				}
+			} else if (result.failure) {
+				failures.push(result.failure);
+			}
+		}
+
+		if (failures.length > 0 && ctx.hasUI) {
+			ctx.ui.notify(
+				failures.map((f) => `MCP server "${f.server.name}" failed: ${f.error}`).join("\n") + " (/mcp for status)",
+				"warning",
+			);
+		}
+
+		const resourceCount = [...connections.values()].reduce((n, c) => n + c.resources.length, 0);
+		if (resourceCount > 0) {
+			registerResourceTools();
+		}
+
+		// Servers' own usage instructions (initialize result) ride an every-turn
+		// reminder, as Claude Code injects them (findings §14) — without this the
+		// field is silently dropped and servers can't teach the model their tools.
+		const instructions = mcpInstructionsReminder([...connections.values()]);
+		if (instructions) {
+			pi.events.emit(REMINDER_CHANNEL, {
+				scope: "every-turn",
+				key: "mcp-instructions",
+				text: instructions,
+				placement: "first-prepend",
+				order: CONTEXT_ORDER.mcp,
+			});
+		}
+	};
+
+	let resourceToolsRegistered = false;
+	const registerResourceTools = () => {
+		if (resourceToolsRegistered) return;
+		resourceToolsRegistered = true;
+
+		pi.registerTool({
+			name: "list_mcp_resources",
+			label: "MCP Resources",
+			...ccToolRenderers("MCP Resources"),
+			description:
+				"List resources exposed by connected MCP servers. Resources are readable documents or data the server offers, addressed by uri.",
+			parameters: Type.Object({
+				server: Type.Optional(Type.String({ description: "Limit to one server by name" })),
+			}),
+			async execute(_toolCallId, params) {
+				// A bad `server` name must not read back as "this server has zero
+				// resources" — validate it like read_mcp_resource does, else a typo
+				// looks like an empty (but real) server.
+				if (params.server && !connections.has(params.server)) {
+					return {
+						content: [
+							{
+								type: "text",
+								text: `No connected MCP server named "${params.server}". Connected servers: ${[...connections.keys()].join(", ") || "none"}.`,
+							},
+						],
+						details: { count: 0 },
+						isError: true,
+					};
+				}
+				const lines: string[] = [];
+				for (const connection of connections.values()) {
+					if (params.server && connection.server.name !== params.server) continue;
+					for (const resource of connection.resources) {
+						lines.push(
+							`${connection.server.name} ${resource.uri}${resource.name ? ` — ${resource.name}` : ""}${resource.description ? `: ${resource.description}` : ""}`,
+						);
+					}
+				}
+				return {
+					content: [{ type: "text", text: lines.length ? lines.join("\n") : "No MCP resources available." }],
+					details: { count: lines.length },
+				};
+			},
+		});
+
+		pi.registerTool({
+			name: "read_mcp_resource",
+			label: "Read MCP Resource",
+			...ccToolRenderers("Read MCP Resource"),
+			description: "Read one resource from an MCP server by uri. Use list_mcp_resources to find uris.",
+			parameters: Type.Object({
+				server: Type.String({ description: "Server name that owns the resource" }),
+				uri: Type.String({ description: "Resource uri" }),
+			}),
+			async execute(toolCallId, params, _signal, _onUpdate, ctx) {
+				const connection = connections.get(params.server);
+				if (!connection) return unknownServerError(params.server);
+				try {
+					const result = await readResource(connection, params.uri);
+					const text = describeResourceContents(result.contents as McpResourceContents[] | undefined);
+					return {
+						content: [
+							{
+								type: "text",
+								text: text ? persistIfLarge(text, { dir: resultsDir(ctx), id: toolCallId }) : "(empty resource)",
+							},
+						],
+						details: { server: params.server, uri: params.uri },
+					};
+				} catch (error) {
+					return {
+						content: [{ type: "text", text: `Could not read ${params.uri}: ${(error as Error).message}` }],
+						details: { server: params.server, uri: params.uri },
+						isError: true,
+					};
+				}
+			},
+		});
+
+		pi.registerTool({
+			name: "read_mcp_resource_dir",
+			label: "List MCP Resource Directory",
+			...ccToolRenderers("List MCP Resource Directory"),
+			description:
+				'List the direct children of a directory resource on an MCP server (resources/directory/read). Not recursive: each entry carries its own uri, and subdirectories appear with mimeType "inode/directory" — call again on a subdirectory uri to descend. Only servers that support directory listing accept this; others return an error.',
+			parameters: Type.Object({
+				server: Type.String({ description: "Server name that owns the directory" }),
+				uri: Type.String({ description: "The directory resource uri to list" }),
+			}),
+			async execute(_toolCallId, params) {
+				const connection = connections.get(params.server);
+				if (!connection) return unknownServerError(params.server);
+				try {
+					const result = await readResourceDir(connection, params.uri);
+					const entries = (result.resources ?? result.entries ?? []) as Array<{ uri?: string; name?: string; mimeType?: string }>;
+					const lines = entries.map((e) => `${e.uri ?? "(no uri)"}${e.name ? ` — ${e.name}` : ""}${e.mimeType ? ` (${e.mimeType})` : ""}`);
+					return {
+						content: [{ type: "text", text: lines.length ? lines.join("\n") : "(empty directory)" }],
+						details: { server: params.server, uri: params.uri },
+					};
+				} catch (error) {
+					return {
+						content: [{ type: "text", text: `Could not list ${params.uri}: ${(error as Error).message}` }],
+						details: { server: params.server, uri: params.uri },
+						isError: true,
+					};
+				}
+			},
+		});
+
+		for (const name of ["list_mcp_resources", "read_mcp_resource", "read_mcp_resource_dir"]) {
+			pi.events.emit(DEFER_CHANNEL, { name, keywords: ["mcp", "resource", "document", "uri", "directory"] });
+		}
+	};
+
+	pi.on("session_start", async (_event, ctx) => {
+		// pi awaits session_start handlers serially before the prompt opens, and
+		// remote servers take seconds to answer — awaiting here was the entire
+		// "slow startup" (4.9s → 0.24s measured, findings §15). In the interactive
+		// main session the connect runs in the background: tools register as each
+		// server answers (tool_search handles late defers), and /mcp shows
+		// "connecting" until it settles. Print-mode one-shots and subagent RPC
+		// children still await — their single turn starts immediately and would
+		// race past tools that are not registered yet.
+		connecting = connectAll(ctx).catch((error) => {
+			// A rejection here would otherwise be an unhandled-promise crash in
+			// the background path; per-server failures are already collected in
+			// `failures`, so this only catches setup bugs — surface them loud.
+			if (ctx.hasUI) ctx.ui.notify(`MCP startup failed: ${(error as Error).message}`, "error");
+		});
+		connecting.finally(() => {
+			connectSettled = true;
+		});
+		if (!ctx.hasUI || process.env.PI_SUBAGENT_CHILD) {
+			await connecting;
+		}
+	});
+
+	pi.on("session_shutdown", async () => {
+		shuttingDown = true;
+		await Promise.all([...connections.values()].map((connection) => close(connection)));
+		connections.clear();
+	});
+
+	pi.registerCommand("mcp", {
+		description: "Show MCP server status",
+		handler: async (_args, ctx) => {
+			if (servers.length === 0) {
+				ctx.ui.notify("No MCP servers configured. Add them to .mcp.json or ~/.claude.json.", "info");
+				return;
+			}
+			const lines = servers.map((server) => {
+				const connection = connections.get(server.name);
+				if (connection) {
+					const warning = connection.warnings.length ? ` — WARNING: ${connection.warnings.join("; ")}` : "";
+					return `connected ${server.name} — ${connection.tools.length} tools, ${connection.resources.length} resources (${server.source})${warning}`;
+				}
+				const failure = failures.find((f) => f.server.name === server.name);
+				if (failure) return `failed    ${server.name} — ${failure.error} (${server.source})`;
+				return connectSettled
+					? `failed    ${server.name} — not connected (${server.source})`
+					: `…         ${server.name} — still connecting (${server.source})`;
+			});
+			ctx.ui.notify(lines.join("\n"), "info");
+		},
+	});
+}

package/extensions/mcp/schema.ts

@@ -0,0 +1,161 @@
+/**
+ * Tool naming and JSON Schema → TypeBox conversion (pure).
+ *
+ * MCP servers describe their tools with plain JSON Schema. pi expects TypeBox
+ * schemas, so rather than relying on a raw schema object happening to satisfy
+ * TypeBox's runtime checks, the common subset is converted explicitly and
+ * anything exotic degrades to a permissive object that passes arguments through.
+ */
+
+import { type TSchema, Type } from "typebox";
+
+/** Claude Code's namespacing: mcp__<server>__<tool>, sanitised for tool-name rules. */
+export function namespacedToolName(server: string, tool: string): string {
+	const clean = (part: string) => part.replace(/[^a-zA-Z0-9_]/g, "_");
+	return `mcp__${clean(server)}__${clean(tool)}`;
+}
+
+export function parseNamespacedToolName(name: string): { server: string; tool: string } | undefined {
+	if (!name.startsWith("mcp__")) return undefined;
+	const rest = name.slice("mcp__".length);
+	const separator = rest.indexOf("__");
+	if (separator === -1) return undefined;
+	return { server: rest.slice(0, separator), tool: rest.slice(separator + 2) };
+}
+
+interface JsonSchema {
+	type?: string | string[];
+	description?: string;
+	properties?: Record<string, JsonSchema>;
+	required?: string[];
+	items?: JsonSchema;
+	enum?: unknown[];
+	const?: unknown;
+	anyOf?: JsonSchema[];
+	oneOf?: JsonSchema[];
+	[key: string]: unknown;
+}
+
+function convertLeaf(schema: JsonSchema): TSchema | undefined {
+	const options = schema.description ? { description: schema.description } : {};
+	const type = Array.isArray(schema.type) ? schema.type.find((t) => t !== "null") : schema.type;
+
+	switch (type) {
+		case "string":
+			return Type.String(options);
+		case "number":
+			return Type.Number(options);
+		case "integer":
+			return Type.Integer(options);
+		case "boolean":
+			return Type.Boolean(options);
+		case "array":
+			return Type.Array(schema.items ? jsonSchemaToTypeBox(schema.items) : Type.Unknown(), options);
+		case "object":
+			return jsonSchemaToTypeBox(schema);
+		default:
+			return undefined;
+	}
+}
+
+/**
+ * Converts a JSON Schema to a TypeBox schema. Enums become literal unions;
+ * unrecognised constructs become `Unknown`, which validates anything, so an
+ * unusual server schema degrades to passing arguments through rather than
+ * rejecting every call.
+ */
+export function jsonSchemaToTypeBox(schema: JsonSchema | undefined): TSchema {
+	if (!schema || typeof schema !== "object") return Type.Unknown();
+
+	const options = schema.description ? { description: schema.description } : {};
+
+	if (Array.isArray(schema.enum) && schema.enum.length > 0) {
+		const literals = schema.enum
+			.filter((value): value is string | number | boolean => ["string", "number", "boolean"].includes(typeof value))
+			.map((value) => Type.Literal(value));
+		if (literals.length === schema.enum.length) {
+			return literals.length === 1 ? literals[0] : Type.Union(literals, options);
+		}
+		return Type.Unknown();
+	}
+
+	if (schema.properties || schema.type === "object") {
+		const required = new Set(Array.isArray(schema.required) ? schema.required : []);
+		const properties: Record<string, TSchema> = {};
+		for (const [key, value] of Object.entries(schema.properties ?? {})) {
+			const converted = jsonSchemaToTypeBox(value);
+			properties[key] = required.has(key) ? converted : Type.Optional(converted);
+		}
+		// additionalProperties stays open: servers commonly accept more than they declare.
+		return Type.Object(properties, { ...options, additionalProperties: true });
+	}
+
+	const leaf = convertLeaf(schema);
+	if (leaf) return leaf;
+
+	const variants = schema.anyOf ?? schema.oneOf;
+	if (Array.isArray(variants) && variants.length > 0) {
+		const converted = variants.map((variant) => jsonSchemaToTypeBox(variant));
+		return converted.length === 1 ? converted[0] : Type.Union(converted, options);
+	}
+
+	return Type.Unknown();
+}
+
+/** MCP tool results carry typed content blocks; map the ones pi can display. */
+export interface McpContentBlock {
+	type: string;
+	text?: string;
+	data?: string;
+	mimeType?: string;
+	resource?: { text?: string; uri?: string };
+}
+
+/**
+ * `resources/read` returns resource *contents*, not typed content blocks: each
+ * entry carries `text` or base64 `blob` plus `uri`/`mimeType`, with no `type`
+ * field. They need their own formatter — running them through describeContent
+ * yields "[undefined content]".
+ */
+export interface McpResourceContents {
+	uri?: string;
+	mimeType?: string;
+	text?: string;
+	blob?: string;
+}
+
+export function describeResourceContents(contents: McpResourceContents[] | undefined): string {
+	const parts: string[] = [];
+	for (const entry of contents ?? []) {
+		if (typeof entry.text === "string") {
+			parts.push(entry.text);
+		} else if (typeof entry.blob === "string") {
+			const size = Math.round((entry.blob.length * 3) / 4 / 1024);
+			parts.push(`[binary resource ${entry.uri ?? ""} ${entry.mimeType ?? "unknown type"}, ~${size} KB]`);
+		}
+	}
+	return parts.join("\n").trim();
+}
+
+export function describeContent(blocks: McpContentBlock[] | undefined): {
+	text: string;
+	images: Array<{ data: string; mimeType: string }>;
+} {
+	const texts: string[] = [];
+	const images: Array<{ data: string; mimeType: string }> = [];
+
+	for (const block of blocks ?? []) {
+		if (block.type === "text" && typeof block.text === "string") {
+			texts.push(block.text);
+		} else if (block.type === "image" && block.data) {
+			images.push({ data: block.data, mimeType: block.mimeType ?? "image/png" });
+		} else if (block.type === "resource" && block.resource) {
+			const { text, uri } = block.resource;
+			texts.push(text ?? `[resource ${uri ?? "(no uri)"}]`);
+		} else {
+			texts.push(`[${block.type} content]`);
+		}
+	}
+
+	return { text: texts.join("\n").trim(), images };
+}

package/extensions/memory/index.ts

@@ -0,0 +1,90 @@
+/**
+ * memory extension — filesystem side of the auto-memory system (see
+ * extensions/lib/memory.ts for the design and the prompt side).
+ *
+ * - session start: create the per-project memory directory (the system prompt
+ *   promises the model "this directory already exists") and inject the
+ *   MEMORY.md index — capped like Claude Code caps it — via the reminder queue.
+ * - write into the memory dir: stamp bookkeeping frontmatter (node_type,
+ *   originSessionId, modified) into the tool input *before* execution, so the
+ *   file lands stamped and file-tracker observes the same bytes that are on
+ *   disk. Edits are not stamped: Claude Code ties the stamp to writes, and an
+ *   edit's input carries diffs, not the file.
+ * - write/edit of MEMORY.md: check the load limits afterwards; near-limit
+ *   nudges over the reminder queue, over-limit turns the result into an error
+ *   (the write itself succeeded — the message says so), matching Claude Code.
+ */
+
+import { mkdirSync, readFileSync } from "node:fs";
+import { isAbsolute, join, resolve, sep } from "node:path";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import {
+	INDEX_NEAR_LIMIT_REMINDER,
+	INDEX_OVER_LIMIT_ERROR,
+	indexLimitStatus,
+	projectMemoryDir,
+	stampFrontmatter,
+} from "../lib/memory.ts";
+import { REMINDER_CHANNEL } from "../lib/reminders.ts";
+
+function inputPath(input: unknown, cwd: string): string | undefined {
+	const raw = (input as { path?: unknown })?.path;
+	if (typeof raw !== "string" || !raw.trim()) return undefined;
+	return isAbsolute(raw) ? raw : resolve(cwd, raw);
+}
+
+export default function memoryExtension(pi: ExtensionAPI) {
+	let dir: string | undefined;
+
+	pi.on("session_start", (_event, ctx) => {
+		// Like Claude Code: one memory dir per git repo (shared by worktrees and
+		// subdirectories); outside a repo, per cwd.
+		const candidate = projectMemoryDir(ctx.cwd);
+		try {
+			mkdirSync(candidate, { recursive: true });
+		} catch {
+			// Unwritable home (sandboxes, CI): the prompt section still renders,
+			// and writes will surface their own errors to the model.
+			return;
+		}
+		dir = candidate;
+
+		// The MEMORY.md index is injected by the claude-context extension, folded
+		// into the `# claudeMd` block exactly as Claude Code does — not as a
+		// separate reminder here. This extension owns the directory and the
+		// write-time stamping / limit checks below.
+	});
+
+	pi.on("tool_call", (event, ctx) => {
+		if (event.toolName !== "write" || !dir) return undefined;
+		const path = inputPath(event.input, ctx.cwd);
+		if (!path?.startsWith(dir + sep)) return undefined;
+		const input = event.input as { content?: unknown };
+		if (typeof input.content !== "string") return undefined;
+		input.content = stampFrontmatter(input.content, ctx.sessionManager.getSessionId(), new Date().toISOString());
+		return undefined;
+	});
+
+	pi.on("tool_result", (event, ctx) => {
+		if (event.isError || !dir) return undefined;
+		if (event.toolName !== "write" && event.toolName !== "edit") return undefined;
+		if (inputPath(event.input, ctx.cwd) !== join(dir, "MEMORY.md")) return undefined;
+
+		let index: string;
+		try {
+			index = readFileSync(join(dir, "MEMORY.md"), "utf8");
+		} catch {
+			return undefined;
+		}
+		const status = indexLimitStatus(index);
+		if (status === "near") {
+			pi.events.emit(REMINDER_CHANNEL, { text: INDEX_NEAR_LIMIT_REMINDER });
+		} else if (status === "over") {
+			return {
+				content: [...(event.content ?? []), { type: "text", text: INDEX_OVER_LIMIT_ERROR }],
+				isError: true,
+			};
+		}
+		return undefined;
+	});
+}