diff --git a/biome.json b/biome.json index d3e200a7b..f9a3a50df 100644 --- a/biome.json +++ b/biome.json @@ -32,6 +32,7 @@ "packages/session-backends/*/src/**/*.ts", "packages/session-backends/*/test/**/*.ts", "packages/coding-agent/examples/**/*.ts", + "packages/agent/examples/**/*.ts", "!**/node_modules/**/*", "!**/test-sessions.ts", "!**/models.generated.ts", diff --git a/package-lock.json b/package-lock.json index a282f3714..be68b57fa 100644 --- a/package-lock.json +++ b/package-lock.json @@ -731,6 +731,10 @@ "resolved": "packages/client", "link": true }, + "node_modules/@earendil-works/pi-codemode": { + "resolved": "packages/codemode", + "link": true + }, "node_modules/@earendil-works/pi-coding-agent": { "resolved": "packages/coding-agent", "link": true @@ -743,6 +747,10 @@ "resolved": "packages/evals", "link": true }, + "node_modules/@earendil-works/pi-mcp": { + "resolved": "packages/mcp", + "link": true + }, "node_modules/@earendil-works/pi-protocol": { "resolved": "packages/protocol", "link": true @@ -4538,6 +4546,12 @@ ], "license": "MIT" }, + "node_modules/quickjs-wasi": { + "version": "3.6.2", + "resolved": "https://registry.npmjs.org/quickjs-wasi/-/quickjs-wasi-3.6.2.tgz", + "integrity": "sha512-FCqGtGOrMgzUiIrMNMA2YnsOxCNwo31dzqXvclXUC6xeT35NJLKXQJsvbeCTjvoFAwZgEAPg8U6+KAPDGXn8Mg==", + "license": "MIT" + }, "node_modules/rc": { "version": "1.2.8", "resolved": "https://registry.npmjs.org/rc/-/rc-1.2.8.tgz", @@ -5650,6 +5664,21 @@ "node": ">=22.19.0" } }, + "packages/codemode": { + "name": "@earendil-works/pi-codemode", + "version": "0.87.1", + "license": "MIT", + "dependencies": { + "quickjs-wasi": "3.6.2" + }, + "devDependencies": { + "shx": "0.4.0", + "vitest": "4.1.9" + }, + "engines": { + "node": ">=22.19.0" + } + }, "packages/coding-agent": { "name": "@earendil-works/pi-coding-agent", "version": "0.87.1", @@ -5658,6 +5687,8 @@ "@earendil-works/chord": "^0.87.1", "@earendil-works/pi-agent-core": "^0.87.1", "@earendil-works/pi-ai": "^0.87.1", + "@earendil-works/pi-codemode": "^0.87.1", + "@earendil-works/pi-mcp": "^0.87.1", "@earendil-works/pi-tui": "^0.87.1", "@silvia-odwyer/photon-node": "0.3.4", "chalk": "6.0.0", @@ -5670,6 +5701,7 @@ "jiti": "2.7.0", "minimatch": "10.2.6", "proper-lockfile": "4.1.2", + "quickjs-wasi": "3.6.2", "semver": "7.8.5", "typebox": "1.3.27", "undici": "8.10.2", @@ -5804,6 +5836,22 @@ "vitest-evals": "0.15.0" } }, + "packages/mcp": { + "name": "@earendil-works/pi-mcp", + "version": "0.87.1", + "license": "MIT", + "dependencies": { + "cross-spawn": "7.0.6" + }, + "devDependencies": { + "@types/cross-spawn": "6.0.6", + "shx": "0.4.0", + "vitest": "4.1.9" + }, + "engines": { + "node": ">=22.19.0" + } + }, "packages/protocol": { "name": "@earendil-works/pi-protocol", "version": "0.87.1", diff --git a/package.json b/package.json index d9fe33331..26493c596 100644 --- a/package.json +++ b/package.json @@ -13,8 +13,8 @@ ], "scripts": { "clean": "npm run clean --workspaces", - "build": "cd packages/chord && npm run build && cd ../tui && npm run build && cd ../telemetry && npm run build && cd ../ai && npm run build && cd ../durable && npm run build && cd ../agent && npm run build && cd ../session-backends/sqlite-node && npm run build && cd ../../protocol && npm run build && cd ../client && npm run build && cd ../server && npm run build && cd ../coding-agent && npm run build", - "build:offline": "cd packages/chord && npm run build && cd ../tui && npm run build && cd ../telemetry && npm run build && cd ../ai && npm run build:offline && cd ../durable && npm run build && cd ../agent && npm run build && cd ../session-backends/sqlite-node && npm run build && cd ../../protocol && npm run build && cd ../client && npm run build && cd ../server && npm run build && cd ../coding-agent && npm run build", + "build": "cd packages/chord && npm run build && cd ../tui && npm run build && cd ../telemetry && npm run build && cd ../codemode && npm run build && cd ../mcp && npm run build && cd ../ai && npm run build && cd ../durable && npm run build && cd ../agent && npm run build && cd ../session-backends/sqlite-node && npm run build && cd ../../protocol && npm run build && cd ../client && npm run build && cd ../server && npm run build && cd ../coding-agent && npm run build", + "build:offline": "cd packages/chord && npm run build && cd ../tui && npm run build && cd ../telemetry && npm run build && cd ../codemode && npm run build && cd ../mcp && npm run build && cd ../ai && npm run build:offline && cd ../durable && npm run build && cd ../agent && npm run build && cd ../session-backends/sqlite-node && npm run build && cd ../../protocol && npm run build && cd ../client && npm run build && cd ../server && npm run build && cd ../coding-agent && npm run build", "build:native:darwin": "npm --prefix packages/tui run build:native:darwin", "build:native:linux": "npm --prefix packages/tui run build:native:linux", "build:native:win32": "npm --prefix packages/tui run build:native:win32", diff --git a/packages/agent/README.md b/packages/agent/README.md index aa43845c6..12a5d39eb 100644 --- a/packages/agent/README.md +++ b/packages/agent/README.md @@ -511,6 +511,10 @@ Thrown errors are caught by the agent and reported to the LLM as tool errors wit Return `terminate: true` from `execute()`, a blocked `beforeToolCall`, or `afterToolCall` to hint that the agent should stop after the current tool batch. This only takes effect when every finalized tool result in the batch is terminating. The hint is runtime-only; emitted `toolResult` transcript messages remain standard LLM tool results. +### MCP and Codemode + +`@earendil-works/pi-mcp` connects to MCP servers and `@earendil-works/pi-codemode` runs model-written JavaScript that calls tools. [examples/mcp-codemode](examples/mcp-codemode) wraps both as `AgentTool`s: one tool per MCP tool, and a `codemode` tool whose scripts call the agent's tools through `runToolCall()`, so `beforeToolCall` and `afterToolCall` apply to those calls too. + ## Proxy Usage For browser apps that proxy through a backend: diff --git a/packages/agent/examples/mcp-codemode/main.ts b/packages/agent/examples/mcp-codemode/main.ts new file mode 100644 index 000000000..e86988da8 --- /dev/null +++ b/packages/agent/examples/mcp-codemode/main.ts @@ -0,0 +1,50 @@ +/** + * An Agent with the tools of an MCP server plus a codemode tool that calls them from JavaScript. + * + * ANTHROPIC_API_KEY=... node packages/agent/examples/mcp-codemode/main.ts \ + * npx -y @modelcontextprotocol/server-filesystem . + * + * The arguments are the command that starts a stdio MCP server. + */ + +import { Agent } from "@earendil-works/pi-agent-core"; +import { createModels } from "@earendil-works/pi-ai"; +import { anthropicProvider } from "@earendil-works/pi-ai/providers/anthropic"; +import { McpClient, StdioTransport } from "@earendil-works/pi-mcp"; +import { createCodemodeTool, createMcpTools, createNestedToolRunner } from "./tools.ts"; + +const [command, ...args] = process.argv.slice(2); +if (!command) { + console.error("Usage: main.ts [args...]"); + process.exit(1); +} + +const models = createModels(); +models.setProvider(anthropicProvider()); +const model = models.getModel("anthropic", "claude-sonnet-4-6"); +if (!model) throw new Error("Model not found"); + +const client = new McpClient({ name: "pi-agent-example", version: "1.0.0" }); +await client.connect(new StdioTransport({ command, args })); + +try { + const agent = new Agent({ + initialState: { systemPrompt: "You are a helpful assistant.", model }, + streamFn: models.streamSimple.bind(models), + }); + const mcpTools = await createMcpTools(client, "mcp"); + // The model can call the MCP tools directly or batch them in one codemode script. + agent.state.tools = [...mcpTools, createCodemodeTool(mcpTools, createNestedToolRunner(agent))]; + + agent.subscribe((event) => { + if (event.type === "message_update" && event.assistantMessageEvent.type === "text_delta") { + process.stdout.write(event.assistantMessageEvent.delta); + } else if (event.type === "tool_execution_start") { + console.log(`\n[${event.toolName}]`); + } + }); + await agent.prompt("Use codemode to list what the MCP tools can do, then summarize it in two sentences."); + console.log(); +} finally { + await client.close(); +} diff --git a/packages/agent/examples/mcp-codemode/tools.ts b/packages/agent/examples/mcp-codemode/tools.ts new file mode 100644 index 000000000..37e358c65 --- /dev/null +++ b/packages/agent/examples/mcp-codemode/tools.ts @@ -0,0 +1,172 @@ +/** + * Adapters that give a pi-agent-core `Agent` MCP tools and a codemode tool, using + * `@earendil-works/pi-mcp` and `@earendil-works/pi-codemode`. See main.ts for how they fit together. + */ + +import { + type Agent, + type AgentTool, + type AgentToolCallOutcome, + type AgentToolResult, + runToolCall, +} from "@earendil-works/pi-agent-core"; +import type { ImageContent, JsonObject, JsonValue, TextContent } from "@earendil-works/pi-ai"; +import { + type CodemodeJsonSchema, + CodemodeSandbox, + type CodemodeTool, + renderDeclarations, +} from "@earendil-works/pi-codemode"; +import { type McpClient, toLlmContent } from "@earendil-works/pi-mcp"; +import { Type } from "typebox"; + +/** Provider tool names are limited to 64 characters of `[A-Za-z0-9_-]`. */ +function toToolName(prefix: string, name: string): string { + return `${prefix}_${name}`.replace(/[^A-Za-z0-9_-]/g, "_").slice(0, 64); +} + +function textOf(result: AgentToolResult): string { + return result.content + .filter((block): block is TextContent => block.type === "text") + .map((block) => block.text) + .join("\n"); +} + +/** + * One `AgentTool` per tool of a connected MCP server, named `_`. Tools with an MCP + * output schema declare it as `outputSchema` and return `structuredContent`, which codemode scripts + * receive instead of the text. + */ +export async function createMcpTools(client: McpClient, prefix: string): Promise { + const tools = await client.listTools(); + return tools.map( + (tool): AgentTool => ({ + name: toToolName(prefix, tool.name), + label: tool.title ?? tool.name, + description: tool.description ?? tool.title ?? `MCP tool ${tool.name}`, + // Providers require an object schema, and some reject one without `properties`. + parameters: Type.Unsafe({ + ...tool.inputSchema, + type: "object", + properties: tool.inputSchema.properties ?? {}, + }), + outputSchema: tool.outputSchema ? Type.Unsafe>(tool.outputSchema) : undefined, + execute: async (_toolCallId, params, signal, onUpdate) => { + const result = await client.callTool(tool.name, params as Record, { + signal, + onProgress: (progress) => { + const text = progress.message ?? `Progress ${progress.progress}`; + onUpdate?.({ content: [{ type: "text", text }], details: undefined }); + }, + }); + const content = toLlmContent(result); + if (result.isError && content.length === 0) { + content.push({ type: "text", text: `MCP tool ${tool.name} failed` }); + } + return { + content, + details: undefined, + structuredContent: result.structuredContent as JsonValue | undefined, + // MCP reports tool failures in the result instead of as a protocol error. + isError: result.isError === true, + }; + }, + }), + ); +} + +/** + * Run a tool call on behalf of another tool through the agent's pipeline: argument validation and + * the agent's `beforeToolCall`/`afterToolCall` hooks apply as for calls the model makes. Calling + * `tool.execute()` directly also works when the agent has no hooks. + */ +export function createNestedToolRunner( + agent: Agent, +): (name: string, args: unknown, signal: AbortSignal) => Promise { + let nextId = 1; + return (name, args, signal) => { + const assistantMessage = agent.state.messages.findLast((message) => message.role === "assistant"); + if (assistantMessage?.role !== "assistant") throw new Error("Nested tool calls need an assistant message"); + const tools = agent.state.tools; + return runToolCall( + { type: "toolCall", id: `nested-${nextId++}`, name, arguments: (args ?? {}) as JsonObject }, + { + tools, + assistantMessage, + context: { messages: agent.state.messages, tools }, + beforeToolCall: agent.beforeToolCall, + afterToolCall: agent.afterToolCall, + signal, + }, + ); + }; +} + +export const CODEMODE_TOOL_NAME = "codemode"; + +const codemodeParameters = Type.Object({ + code: Type.String({ description: "JavaScript source. Top-level await and return work." }), +}); + +/** + * A tool that runs model-written JavaScript in the codemode sandbox, where `tools.(args)` + * calls `tools`. Nested results stay out of the model context; only the script's output and + * return value come back. + */ +export function createCodemodeTool( + tools: readonly AgentTool[], + runTool: (name: string, args: unknown, signal: AbortSignal) => Promise, +): AgentTool { + const callable = tools.filter((tool) => tool.name !== CODEMODE_TOOL_NAME); + const sandboxTools: CodemodeTool[] = callable.map((tool) => ({ + name: tool.name, + description: tool.description, + // TypeBox schemas are JSON Schema. + inputSchema: tool.parameters as CodemodeJsonSchema, + // What the script receives: the structured result if the tool declares one, else its text. + outputSchema: (tool.outputSchema as CodemodeJsonSchema | undefined) ?? { type: "string" }, + execute: async (args, { signal }) => { + const outcome = await runTool(tool.name, args, signal); + if (tool.outputSchema && outcome.result.structuredContent !== undefined) { + return outcome.result.structuredContent; + } + const text = textOf(outcome.result); + if (outcome.isError) throw new Error(text || `Tool ${tool.name} failed`); + return text; + }, + })); + + return { + name: CODEMODE_TOOL_NAME, + label: "Codemode", + description: [ + "Run JavaScript that calls other tools, to chain calls, run them in parallel, or filter large results.", + "The code is the body of an async function. There is no Node, file system, network, or timers.", + "Call tools as `await tools.(args)`. Output with `text(value)` or `console.log()`, or `return` a value.", + "Only the output and the return value reach you, not the results of the tool calls.", + "", + "```ts", + renderDeclarations({ tools: sandboxTools }), + "```", + ].join("\n"), + parameters: codemodeParameters, + execute: async (_toolCallId, params, signal) => { + const sandbox = new CodemodeSandbox({ tools: sandboxTools, timeoutMs: 120_000 }); + try { + const result = await sandbox.execute(params.code, { signal }); + // Items from text(), console.*, and image(), in order. They have pi-ai's content shapes. + const content: (TextContent | ImageContent)[] = [...result.output]; + if (result.ok && result.value !== undefined) { + const value = typeof result.value === "string" ? result.value : JSON.stringify(result.value); + content.push({ type: "text", text: value }); + } + if (!result.ok) { + content.push({ type: "text", text: `Script failed: ${result.error.stack ?? result.error.message}` }); + } + return { content, details: { calls: result.calls }, isError: !result.ok }; + } finally { + await sandbox.close(); + } + }, + }; +} diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index 69c2b3180..b772985ca 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -23,6 +23,7 @@ import type { AgentMessage, AgentTool, AgentToolCall, + AgentToolCallOutcome, AgentToolResult, PrepareNextTurnContext, StreamFn, @@ -554,7 +555,7 @@ async function executeToolCallsSequential( isError: preparation.isError, }; } else { - const executed = await executePreparedToolCall(preparation, signal, emit); + const executed = await executePreparedToolCall(preparation, signal, emitToolExecutionUpdate(toolCall, emit)); finalized = await finalizeExecutedToolCall( currentContext, assistantMessage, @@ -625,7 +626,7 @@ async function executeToolCallsParallel( await emitToolExecutionEnd(finalized, emit); return finalized; } - const executed = await executePreparedToolCall(preparation, signal, emit); + const executed = await executePreparedToolCall(preparation, signal, emitToolExecutionUpdate(toolCall, emit)); const finalized = await finalizeExecutedToolCall( currentContext, assistantMessage, @@ -676,11 +677,12 @@ type ExecutedToolCallOutcome = { isError: boolean; }; -type FinalizedToolCallOutcome = { - toolCall: AgentToolCall; - result: AgentToolResult; - isError: boolean; -}; +type FinalizedToolCallOutcome = AgentToolCallOutcome; + +/** The `beforeToolCall` and `afterToolCall` hooks of {@link AgentLoopConfig}. */ +export type ToolCallHooks = Pick; + +type ToolUpdateSink = (partialResult: AgentToolResult) => Promise | void; type FinalizedToolCallEntry = FinalizedToolCallOutcome | (() => Promise); @@ -706,10 +708,11 @@ async function prepareToolCall( currentContext: AgentContext, assistantMessage: AssistantMessage, toolCall: AgentToolCall, - config: AgentLoopConfig, + config: ToolCallHooks, signal: AbortSignal | undefined, + tools: readonly AgentTool[] = currentContext.tools ?? [], ): Promise { - const tool = currentContext.tools?.find((t) => t.name === toolCall.name); + const tool = tools.find((t) => t.name === toolCall.name); if (!tool) { return { kind: "immediate", @@ -772,10 +775,52 @@ async function prepareToolCall( } } +function emitToolExecutionUpdate(toolCall: AgentToolCall, emit: AgentEventSink): ToolUpdateSink { + return (partialResult) => + emit({ + type: "tool_execution_update", + toolCallId: toolCall.id, + toolName: toolCall.name, + args: toolCall.arguments, + partialResult, + }); +} + +/** Options for {@link runToolCall}. */ +export interface RunToolCallOptions extends ToolCallHooks { + /** Tools the call resolves against. */ + tools: readonly AgentTool[]; + /** Passed to the hooks as the message that issued the call. */ + assistantMessage: AssistantMessage; + /** Passed to the hooks as the current agent context. */ + context: AgentContext; + signal?: AbortSignal; + onUpdate?: ToolUpdateSink; +} + +/** + * Run one tool call through the same steps as a model-issued call: argument preparation, schema + * validation, `beforeToolCall`, execution, and `afterToolCall`. Emits no events and adds no + * messages. Tools that call other tools use this so the hooks (for example permission checks) + * apply to those calls too. + * + * Never rejects for tool failures: unknown tools, validation errors, blocked calls, and thrown + * errors come back as `isError: true`. + */ +export async function runToolCall(toolCall: AgentToolCall, options: RunToolCallOptions): Promise { + const { assistantMessage, context, signal } = options; + const preparation = await prepareToolCall(context, assistantMessage, toolCall, options, signal, options.tools); + if (preparation.kind === "immediate") { + return { toolCall, result: preparation.result, isError: preparation.isError }; + } + const executed = await executePreparedToolCall(preparation, signal, options.onUpdate ?? (() => {})); + return finalizeExecutedToolCall(context, assistantMessage, preparation, executed, options, signal); +} + async function executePreparedToolCall( prepared: PreparedToolCall, signal: AbortSignal | undefined, - emit: AgentEventSink, + onUpdate: ToolUpdateSink, ): Promise { const updateEvents: Promise[] = []; let acceptingUpdates = true; @@ -787,22 +832,12 @@ async function executePreparedToolCall( signal, (partialResult) => { if (!acceptingUpdates) return; - updateEvents.push( - Promise.resolve( - emit({ - type: "tool_execution_update", - toolCallId: prepared.toolCall.id, - toolName: prepared.toolCall.name, - args: prepared.toolCall.arguments, - partialResult, - }), - ), - ); + updateEvents.push(Promise.resolve(onUpdate(partialResult))); }, ); acceptingUpdates = false; await Promise.all(updateEvents); - return { result, isError: false }; + return { result, isError: result.isError === true }; } catch (error) { acceptingUpdates = false; await Promise.all(updateEvents); @@ -820,7 +855,7 @@ async function finalizeExecutedToolCall( assistantMessage: AssistantMessage, prepared: PreparedToolCall, executed: ExecutedToolCallOutcome, - config: AgentLoopConfig, + config: ToolCallHooks, signal: AbortSignal | undefined, ): Promise { let result = executed.result; @@ -840,6 +875,9 @@ async function finalizeExecutedToolCall( signal, ); if (afterResult) { + // Structured content not replaced along with the content may no longer match it. + const structuredContent = + afterResult.structuredContent ?? (afterResult.content ? undefined : result.structuredContent); result = { ...result, content: afterResult.content ?? result.content, @@ -847,6 +885,8 @@ async function finalizeExecutedToolCall( usage: afterResult.usage ?? result.usage, terminate: afterResult.terminate ?? result.terminate, }; + if (structuredContent === undefined) delete result.structuredContent; + else result.structuredContent = structuredContent; isError = afterResult.isError ?? isError; } } catch (error) { diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index 8da843f56..6e17c3c8c 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -82,13 +82,17 @@ export interface BeforeToolCallResult { * - `isError`: if provided, replaces the tool result error flag * - `usage`: if provided, replaces the tool result usage * - `terminate`: if provided, replaces the early-termination hint + * - `structuredContent`: if provided, replaces the structured content. If `content` is provided + * without it, the structured content is dropped, because it may no longer match the content. + * Return it along with `content` to keep it. * - * Omitted fields keep the original executed tool result values. + * Other omitted fields keep the original executed tool result values. * There is no deep merge for `content`, `details`, or `usage`. */ export interface AfterToolCallResult { content?: (TextContent | ImageContent)[]; details?: unknown; + structuredContent?: JsonValue; isError?: boolean; /** Usage from the final tool execution itself, if available. Not used for main LLM context accounting. */ usage?: Usage; @@ -422,8 +426,18 @@ export interface AgentToolResult { content: (TextContent | ImageContent)[]; /** Arbitrary structured details for logs or UI rendering. */ details: T; + /** + * Machine-readable result matching the tool's `outputSchema`, for programmatic callers. Not sent + * to the model; `content` remains the model-facing result. + */ + structuredContent?: JsonValue; /** Usage from the final tool execution itself, if available. Not used for main LLM context accounting. */ usage?: Usage; + /** + * Report a failure without throwing. The model sees `content` as an error result, like a thrown + * error, but `details` and `structuredContent` are kept for the UI and programmatic callers. + */ + isError?: boolean; /** * Hint that the agent should stop after the current tool batch. * Early termination only happens when every finalized tool result in the batch sets this to true. @@ -431,6 +445,13 @@ export interface AgentToolResult { terminate?: boolean; } +/** Final outcome of a tool call after hooks ran. */ +export interface AgentToolCallOutcome { + toolCall: AgentToolCall; + result: AgentToolResult; + isError: boolean; +} + /** * Callback used by tools to stream partial execution updates. * @@ -448,7 +469,15 @@ export interface AgentTool Static; - /** Execute the tool call. Throw on failure instead of encoding errors in `content`. */ + /** + * JSON Schema of `structuredContent` in successful results. Tools that declare it should always + * set `structuredContent`. + */ + outputSchema?: TSchema; + /** + * Execute the tool call. Throw on failure, or return a result with `isError: true`; do not only + * describe the failure in `content`. + */ execute: ( toolCallId: string, params: Static, diff --git a/packages/agent/test/agent-loop.test.ts b/packages/agent/test/agent-loop.test.ts index 018b0f7d4..c60c04a34 100644 --- a/packages/agent/test/agent-loop.test.ts +++ b/packages/agent/test/agent-loop.test.ts @@ -8,9 +8,16 @@ import { } from "@earendil-works/pi-ai"; import { Type } from "typebox"; import { describe, expect, it } from "vitest"; -import { agentLoop, agentLoopContinue, runAgentLoop } from "../src/agent-loop.ts"; +import { agentLoop, agentLoopContinue, runAgentLoop, runToolCall } from "../src/agent-loop.ts"; import { setDefaultStreamFn } from "../src/index.ts"; -import type { AgentContext, AgentEvent, AgentLoopConfig, AgentMessage, AgentTool } from "../src/types.ts"; +import type { + AgentContext, + AgentEvent, + AgentLoopConfig, + AgentMessage, + AgentTool, + AgentToolCall, +} from "../src/types.ts"; // Mock stream for testing - mimics MockAssistantStream class MockAssistantStream extends EventStream { @@ -2081,3 +2088,106 @@ describe("agentLoopContinue with AgentMessage", () => { expect(messages[0].role).toBe("assistant"); }); }); + +describe("runToolCall", () => { + const echoSchema = Type.Object({ value: Type.String() }); + const echo: AgentTool = { + name: "echo", + label: "Echo", + description: "Echo tool", + parameters: echoSchema, + outputSchema: Type.Object({ value: Type.String() }), + async execute(_toolCallId, params, _signal, onUpdate) { + onUpdate?.({ content: [{ type: "text", text: "partial" }], details: {} }); + return { + content: [{ type: "text", text: params.value }], + details: {}, + structuredContent: { value: params.value }, + }; + }, + }; + const failing: AgentTool = { + name: "failing", + label: "Failing", + description: "Returns an error result", + parameters: Type.Object({}), + async execute() { + return { content: [{ type: "text", text: "bad" }], details: { partial: true }, isError: true }; + }, + }; + const assistantMessage = createAssistantMessage([]); + const call = (id: string, name: string, args: AgentToolCall["arguments"]): AgentToolCall => ({ + type: "toolCall", + id, + name, + arguments: args, + }); + + it("validates, runs the hooks, and reports failures as error outcomes", async () => { + const hookCalls: string[] = []; + const updates: unknown[] = []; + const options = { + tools: [echo, failing], + assistantMessage, + context: { messages: [] }, + beforeToolCall: async ({ toolCall, args }: { toolCall: { id: string }; args: unknown }) => { + hookCalls.push(`before ${toolCall.id}`); + if ((args as { value?: string }).value === "blocked") return { block: true, reason: "nope" }; + return undefined; + }, + afterToolCall: async ({ toolCall }: { toolCall: { id: string } }) => { + hookCalls.push(`after ${toolCall.id}`); + return undefined; + }, + onUpdate: (partial: unknown) => { + updates.push(partial); + }, + }; + + expect(await runToolCall(call("a", "echo", { value: "a" }), options)).toMatchObject({ + toolCall: { id: "a" }, + result: { structuredContent: { value: "a" } }, + isError: false, + }); + expect(await runToolCall(call("b", "echo", { value: { nested: true } }), options)).toMatchObject({ + isError: true, + }); + expect(await runToolCall(call("c", "echo", { value: "blocked" }), options)).toMatchObject({ + result: { content: [{ type: "text", text: "nope" }] }, + isError: true, + }); + expect(await runToolCall(call("d", "missing", {}), options)).toMatchObject({ + result: { content: [{ type: "text", text: "Tool missing not found" }] }, + isError: true, + }); + // Error results keep their details. + expect(await runToolCall(call("e", "failing", {}), options)).toMatchObject({ + result: { details: { partial: true } }, + isError: true, + }); + expect(updates).toEqual([{ content: [{ type: "text", text: "partial" }], details: {} }]); + // Validation failures and unknown tools never reach the hooks; blocked calls skip afterToolCall. + expect(hookCalls).toEqual(["before a", "after a", "before c", "before e", "after e"]); + }); + + it("lets afterToolCall replace structured content and drops it when only content is replaced", async () => { + const redacted = [{ type: "text" as const, text: "redacted" }]; + const results = [ + { content: redacted }, + { structuredContent: { value: "replaced" } }, + { content: redacted, structuredContent: { value: "both" } }, + { details: { note: "kept" } }, + ]; + const seen: unknown[] = []; + for (const afterResult of results) { + const outcome = await runToolCall(call("x", "echo", { value: "original" }), { + tools: [echo], + assistantMessage, + context: { messages: [] }, + afterToolCall: async () => afterResult, + }); + seen.push(outcome.result.structuredContent); + } + expect(seen).toEqual([undefined, { value: "replaced" }, { value: "both" }, { value: "original" }]); + }); +}); diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index bd1245528..f1cc56bf7 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -569,6 +569,28 @@ export interface AssistantMessage { timestamp: number; // Unix timestamp in milliseconds } +/** A tool call that another tool made while it ran, for example from a codemode script. */ +export interface NestedToolCallRecord { + id: string; + name: string; + /** Omitted when over the size limits; `argumentsBytes` then gives their size. */ + arguments?: JsonObject; + /** UTF-8 size of the arguments as JSON, set when `arguments` is omitted. */ + argumentsBytes?: number; + /** `unfinished`: the call was still running when the calling tool finished. */ + status: "ok" | "error" | "unfinished"; + durationMs?: number; + /** Error text, truncated. */ + error?: string; +} + +/** Bounded record of the nested calls a tool made. Results are not recorded. */ +export interface NestedToolCalls { + calls: NestedToolCallRecord[]; + /** False when calls were dropped, arguments omitted, or calls had not finished. */ + complete: boolean; +} + export type ToolResultMessage = IsJsonCompatible extends true ? { role: "toolResult"; @@ -578,6 +600,8 @@ export type ToolResultMessage = IsJsonCompatible details?: JsonRepresentation; /** Usage from the tool execution itself, if available. Not part of main LLM context accounting. */ usage?: Usage; + /** Calls this tool made to other tools. Kept for the session record; not sent to the model. */ + nestedCalls?: NestedToolCalls; isError: boolean; timestamp: number; // Unix timestamp in milliseconds } diff --git a/packages/codemode/CHANGELOG.md b/packages/codemode/CHANGELOG.md new file mode 100644 index 000000000..470fc6b81 --- /dev/null +++ b/packages/codemode/CHANGELOG.md @@ -0,0 +1,7 @@ +# Changelog + +## [Unreleased] + +### Added + +- Initial spike: `CodemodeSandbox` runs model-written JavaScript in a worker thread and exposes injected tools as `tools.(args)` async functions. diff --git a/packages/codemode/README.md b/packages/codemode/README.md new file mode 100644 index 000000000..8bfda676f --- /dev/null +++ b/packages/codemode/README.md @@ -0,0 +1,186 @@ +# @earendil-works/pi-codemode + +Runs model-written JavaScript in a QuickJS VM (compiled to WebAssembly) where the only capability is calling injected tools. Nested tool calls never enter the LLM context; only the script's output and return value do. + +Scripts use `tools`, `ALL_TOOLS`, `text`, `image`, `exit`, `store`, and `load`, and may start with a `// @options:` line. The coding agent uses it for its built-in `codemode` tool. It has no pi dependencies and can be used on its own to expose any functions (remote APIs, MCP servers, application services) to model-written scripts. + +## Usage + +```ts +import { CodemodeSandbox } from "@earendil-works/pi-codemode"; + +const sandbox = new CodemodeSandbox({ + timeoutMs: 60_000, + tools: [ + { + name: "read", + execute: async (args, { signal }) => { + const { path } = args as { path: string }; + return await readFile(path, "utf8"); + }, + }, + ], +}); + +const result = await sandbox.execute(` + const source = await tools.read({ path: "package.json" }); + text("bytes " + source.length); + return JSON.parse(source).name; +`); + +console.log(result.output); // [{ type: "text", text: "bytes 1234" }] +if (result.ok) console.log(result.value); // "@earendil-works/pi-codemode" +else console.error(result.error.kind, result.error.message); + +await sandbox.close(); +``` + +`code` is the body of an async function: `return` and top-level `await` work. Inside the script: + +- `tools.(args)` returns a promise. Arguments and results make a JSON round trip. A tool that throws rejects with an `Error` carrying the same message. Tool names are also exposed as identifiers: characters that are not valid in identifiers become `_` (`toCodemodeIdentifier`), so `my-tool` is `tools.my_tool` as well as `tools["my-tool"]`. +- `ALL_TOOLS` lists `{ name, description }` for every tool, with `name` as the identifier. +- `text(value)` appends a text item to `result.output`; values other than strings are JSON-stringified. `console.log/info/warn/error/debug` append text items too. +- `image(urlOrItem)` appends an image item. It accepts a base64 `data:` URL, `{ image_url }`, or an MCP `ImageContent` block (`{ type: "image", data, mimeType }`). Remote URLs are rejected. +- `exit()` ends the script successfully right away, keeping its output and store writes. +- `globals` passed to the sandbox are called as top-level functions, for example a host helper `image(ref)`. They behave like tools but are not recorded in `result.calls`. A name like `models.classify` puts the function on a frozen `models` object. With `spread: true`, `execute` receives all call arguments as an array instead of the first one, and `signature` replaces the declaration generated from the schemas. +- `store(key, value)` and `load(key)` read and write JSON values synchronously. See [Store](#store). +- Nothing else: no timers, `fetch`, `process`, `require`, modules, or `WebAssembly`. `eval` and `Function` work but only produce more code inside the same VM. + +`timeoutMs: Infinity` disables the deadline; the script then runs until it settles or `signal` aborts it. A script that waits on a promise nothing can settle (no tool call pending, and the VM has no timers or I/O) fails right away instead of hanging. + +`memoryLimitBytes` caps the VM's heap. Allocations beyond it fail inside the script as `InternalError: out of memory`. + +## Store + +`store`/`load` let scripts keep values across executions. The sandbox does not persist anything itself: pass the current values as `options.store`, and a successful result reports what the script changed as `result.storeWrites` (`{ set, delete }`). Failed executions report no writes. + +```ts +const result = await sandbox.execute(`store("runs", (load("runs") ?? 0) + 1)`, { store: saved }); +if (result.ok) { + for (const key of result.storeWrites.delete) delete saved[key]; + Object.assign(saved, result.storeWrites.set); +} +``` + +`load` returns a copy, so mutating it does not change the store. Storing `undefined` deletes the key. A value may be at most `MAX_STORE_VALUE_CHARS` (256 Ki) characters of JSON and all values together at most `MAX_STORE_TOTAL_CHARS` (1 Mi); larger writes throw a `RangeError` inside the script. + +## Source format + +`parseCodemodeSource()` accepts a script whose first line may be an options line: + +```js +// @options: {"max_output_tokens": 2000, "timeout_ms": 30000} +const source = await tools.read({ path: "package.json" }); +text(JSON.parse(source).name); +``` + +Supported fields are `max_output_tokens`, a token budget for the output, and `timeout_ms`, a hard deadline. The sandbox does not act on them; the caller decides. The options line is replaced by an empty line, so line numbers in stack traces still match the input. Empty input, invalid JSON, unknown fields, or an options line without code throw `CodemodeSourceError`. `CODEMODE_SOURCE_GRAMMAR` is a Lark grammar for providers that support grammar-constrained tool input. Both are also available from the lightweight `@earendil-works/pi-codemode/source` entry. + +## Bundled hosts + +By default the sandbox loads `quickjs-wasi/quickjs.wasm` from the installed package and starts the worker file that sits next to this package's module. Neither exists on disk when the host is bundled, so pass both: + +```ts +import { CodemodeSandbox, loadQuickJSWasm } from "@earendil-works/pi-codemode"; + +const sandbox = new CodemodeSandbox({ + tools, + // Compiled once per path and cached. + wasm: loadQuickJSWasm(pathToQuickJSWasm), + // A file of your build containing `import "@earendil-works/pi-codemode/worker";` + workerUrl: new URL("./codemode-worker.js", import.meta.url), +}); +``` + +## Declarations for the model + +Tools and globals can carry `description`, `inputSchema`, and `outputSchema` (JSON Schema). `renderDeclarations()` turns them into TypeScript declarations for a model-facing tool description: + +```ts +renderDeclarations({ tools: sandbox.tools, globals: sandbox.globals }); +// declare const tools: { +// /** Read a file */ +// read(args: { +// path: string; +// }): Promise; +// }; +``` + +Schemas only shape the declarations; values are not validated against them. Local references (`#/$defs/...`, `#/definitions/...`) are expanded; recursive and remote references render as `unknown`. + +## Using with pi-agent-core + +To give an `Agent` a codemode tool, expose its other tools to the sandbox and wrap `execute()` as an `AgentTool`: + +```ts +import type { AgentTool } from "@earendil-works/pi-agent-core"; +import { + type CodemodeJsonSchema, + CodemodeSandbox, + type CodemodeTool, + renderDeclarations, +} from "@earendil-works/pi-codemode"; +import { Type } from "typebox"; + +const sandboxTools: CodemodeTool[] = agentTools.map((tool) => ({ + name: tool.name, + description: tool.description, + inputSchema: tool.parameters as CodemodeJsonSchema, + outputSchema: (tool.outputSchema as CodemodeJsonSchema | undefined) ?? { type: "string" }, + execute: async (args, { signal }) => { + const result = await tool.execute("nested", args as never, signal); + if (tool.outputSchema && result.structuredContent !== undefined) return result.structuredContent; + return result.content.map((block) => (block.type === "text" ? block.text : "")).join("\n"); + }, +})); + +const codemodeTool: AgentTool = { + name: "codemode", + label: "Codemode", + description: `Run JavaScript that calls tools as \`await tools.(args)\`. Output with text() or return.\n\n${renderDeclarations({ tools: sandboxTools })}`, + parameters: Type.Object({ code: Type.String() }), + execute: async (_toolCallId, { code }, signal) => { + const sandbox = new CodemodeSandbox({ tools: sandboxTools }); + try { + const result = await sandbox.execute(code, { signal }); + const content = [...result.output]; + if (result.ok && result.value !== undefined) content.push({ type: "text", text: JSON.stringify(result.value) }); + if (!result.ok) content.push({ type: "text", text: result.error.stack ?? result.error.message }); + return { content, details: undefined, isError: !result.ok }; + } finally { + await sandbox.close(); + } + }, +}; +``` + +`result.output` items already have the shape of `@earendil-works/pi-ai`'s `TextContent` and `ImageContent`. Calling `tool.execute()` directly skips the agent's `beforeToolCall` and `afterToolCall` hooks. To apply them to nested calls too, run each call through `runToolCall()` from `@earendil-works/pi-agent-core`, as the [mcp-codemode example](https://github.com/earendil-works/pi/tree/main/packages/agent/examples/mcp-codemode) does. That example also rejects failed nested calls inside the script and combines codemode with MCP tools. + +## Results + +`execute()` never rejects for script failures. `result.error.kind` is one of: + +| kind | meaning | +| --------- | --------------------------------------------------------------------------- | +| `script` | the script threw or failed to parse; `stack` points at `codemode.js:` | +| `timeout` | the deadline expired; the worker was terminated | +| `aborted` | `options.signal` fired or `close()` was called; the worker was terminated | +| `sandbox` | the worker or VM failed, for example a wasm trap or a missing worker file | + +`result.output` holds the text and image items in the order the script produced them, also for failed executions. `result.calls` lists every tool call with `status: "ok" | "error" | "cancelled"`. A call that is still running when the script returns (not awaited) is aborted through the tool's `signal` and reported as `cancelled`. + +## How it works + +Each `execute()` starts a worker thread (about 20 ms including VM creation) that instantiates a fresh QuickJS VM from the compiled wasm module. The VM is a separate wasm instance with its own linear memory. Its only imports are a WASI shim (clock, random, and stdout/stderr writes, which the worker discards) and one host-call entry point, so the script cannot reach the host except through the functions the worker registers. + +The worker evaluates a prelude inside the VM that holds the single host bridge in a closure and builds `tools`, `console`, and globals on top of it. The script is compiled as an async function body. + +Tool calls are relayed to the host thread as messages; the host runs the tool and posts the JSON result back. The host owns the deadline and the abort signal. When either fires, it sets a shared interrupt flag that the VM polls, then calls `worker.terminate()`. The flag is needed on Bun, where `terminate()` cannot stop a thread that is spinning in wasm (`while (true) {}` or `while (true) await null`). + +The worker keeps script execution off the host thread: QuickJS runs synchronously, so a spinning script on the host thread would block its event loop. + +## Runtime notes + +- Works the same on Node and Bun, including memory limits, interrupts, and a catchable `RangeError` for deep recursion (QuickJS's stack guard is enabled; without it the wasm stack overflows and traps). +- Stack traces use QuickJS frames (`at f (codemode.js:2:31)`), prefixed with `Name: message` like V8. The wrapper prefix shares line 1 with the script, so line numbers are exact; column numbers on line 1 are shifted. +- QuickJS is an interpreter. Glue code and filtering tool results are fast enough; heavy computation is slower than in V8. diff --git a/packages/codemode/package.json b/packages/codemode/package.json new file mode 100644 index 000000000..8879900a7 --- /dev/null +++ b/packages/codemode/package.json @@ -0,0 +1,64 @@ +{ + "name": "@earendil-works/pi-codemode", + "version": "0.87.1", + "description": "Sandboxed JavaScript execution where the only capability is calling injected tools", + "type": "module", + "main": "./dist/index.js", + "types": "./dist/index.d.ts", + "exports": { + ".": { + "source": "./src/index.ts", + "types": "./dist/index.d.ts", + "import": "./dist/index.js" + }, + "./declarations": { + "source": "./src/declarations.ts", + "types": "./dist/declarations.d.ts", + "import": "./dist/declarations.js" + }, + "./source": { + "source": "./src/source.ts", + "types": "./dist/source.d.ts", + "import": "./dist/source.js" + }, + "./worker": { + "source": "./src/runtime/worker.ts", + "types": "./dist/runtime/worker.d.ts", + "import": "./dist/runtime/worker.js" + } + }, + "files": [ + "dist", + "README.md" + ], + "scripts": { + "clean": "shx rm -rf dist", + "build": "tsc -p tsconfig.build.json", + "test": "vitest --run", + "prepublishOnly": "npm run clean && npm run build" + }, + "keywords": [ + "pi", + "codemode", + "sandbox", + "quickjs", + "wasm" + ], + "author": "Earendil Works", + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/earendil-works/pi.git", + "directory": "packages/codemode" + }, + "engines": { + "node": ">=22.19.0" + }, + "dependencies": { + "quickjs-wasi": "3.6.2" + }, + "devDependencies": { + "shx": "0.4.0", + "vitest": "4.1.9" + } +} diff --git a/packages/codemode/src/declarations.ts b/packages/codemode/src/declarations.ts new file mode 100644 index 000000000..6899b63f8 --- /dev/null +++ b/packages/codemode/src/declarations.ts @@ -0,0 +1,351 @@ +import { toCodemodeIdentifier } from "./identifier.ts"; + +export { toCodemodeIdentifier }; + +import type { CodemodeJsonSchema, CodemodeTool } from "./types.ts"; + +const IDENTIFIER = /^[A-Za-z_$][A-Za-z0-9_$]*$/; +const INDENT = " "; +/** Largest rendered input type, in characters, before it becomes `unknown`. */ +export const DEFAULT_INPUT_SCHEMA_MAX_CHARS = 16_000; +/** Local `$ref` expansions per rendered schema, so shared definitions cannot blow up the output. */ +const MAX_REF_EXPANSIONS = 32; + +/** + * TypeScript types for MCP results, from the MCP `CallToolResult` schema, so `CallToolResult` + * declarations can refer to them. + */ +export const MCP_TYPESCRIPT_PREAMBLE = `type Role = "user" | "assistant"; +type MetaObject = Record; +type Annotations = { + audience?: Role[]; + priority?: number; + lastModified?: string; +}; +type Icon = { + src: string; + mimeType?: string; + sizes?: string[]; + theme?: "light" | "dark"; +}; +type TextResourceContents = { + uri: string; + mimeType?: string; + _meta?: MetaObject; + text: string; +}; +type BlobResourceContents = { + uri: string; + mimeType?: string; + _meta?: MetaObject; + blob: string; +}; +type TextContent = { + type: "text"; + text: string; + annotations?: Annotations; + _meta?: MetaObject; +}; +type ImageContent = { + type: "image"; + data: string; + mimeType: string; + annotations?: Annotations; + _meta?: MetaObject; +}; +type AudioContent = { + type: "audio"; + data: string; + mimeType: string; + annotations?: Annotations; + _meta?: MetaObject; +}; +type ResourceLink = { + icons?: Icon[]; + name: string; + title?: string; + uri: string; + description?: string; + mimeType?: string; + annotations?: Annotations; + size?: number; + _meta?: MetaObject; + type: "resource_link"; +}; +type EmbeddedResource = { + type: "resource"; + resource: TextResourceContents | BlobResourceContents; + annotations?: Annotations; + _meta?: MetaObject; +}; +type ContentBlock = + | TextContent + | ImageContent + | AudioContent + | ResourceLink + | EmbeddedResource; +type CallToolResult = { + _meta?: MetaObject; + content: ContentBlock[]; + isError?: boolean; + structuredContent?: TStructured; + [key: string]: unknown; +};`; + +export interface RenderDeclarationsOptions { + tools?: readonly CodemodeTool[]; + globals?: readonly CodemodeTool[]; +} + +/** + * Render TypeScript declarations for the script-visible API. Tools become members of + * `declare const tools`, globals become `declare function` statements, and `ns.member` globals + * members of `declare const ns`. Descriptions become doc comments; schemas become types. + */ +export function renderDeclarations(options: RenderDeclarationsOptions): string { + const sections: string[] = []; + const tools = options.tools ?? []; + if (tools.length > 0) { + const members = tools.map( + (tool) => `${docComment(tool.description, INDENT)}${INDENT}${renderToolSignature(tool)}`, + ); + sections.push(`declare const tools: {\n${members.join("\n")}\n};`); + } + const namespaces = new Map(); + for (const global of options.globals ?? []) { + const dot = global.name.indexOf("."); + if (dot === -1) { + sections.push(renderGlobal(`declare function ${global.name}`, global, "")); + continue; + } + const namespace = global.name.slice(0, dot); + const members = namespaces.get(namespace) ?? []; + if (members.length === 0) namespaces.set(namespace, members); + members.push(renderGlobal(global.name.slice(dot + 1), global, INDENT)); + } + for (const [namespace, members] of namespaces) { + sections.push(`declare const ${namespace}: {\n${members.join("\n")}\n};`); + } + return sections.join("\n\n"); +} + +/** + * One tool as a member of the `tools` object: `name(args: T): Promise;` with the + * name as the identifier scripts use. Input types longer than `inputMaxChars` render as `unknown`. + * Tools whose output schema is an MCP `CallToolResult` render as `Promise>`, + * which needs {@link MCP_TYPESCRIPT_PREAMBLE}. + */ +export function renderToolSignature( + tool: Pick, + options: { inputMaxChars?: number } = {}, +): string { + const input = + tool.inputSchema === undefined + ? "unknown" + : schemaToType(tool.inputSchema, { maxChars: options.inputMaxChars ?? DEFAULT_INPUT_SCHEMA_MAX_CHARS }); + return `${toCodemodeIdentifier(tool.name)}(args: ${input}): Promise<${outputType(tool.outputSchema)}>;`; +} + +/** + * A tool's sample: the description followed by the tool's declaration. Used for tool + * listings and `ALL_TOOLS` entries. + */ +export function renderToolSample( + tool: Pick, + options: { inputMaxChars?: number } = {}, +): string { + const declaration = `declare const tools: { ${renderToolSignature(tool, options)} };`; + return `${tool.description?.trim() ?? ""}\n\ncodemode tool declaration:\n\`\`\`ts\n${declaration}\n\`\`\``; +} + +/** + * The `structuredContent` schema of an MCP `CallToolResult` output schema (detected by a + * `content` array of objects, boolean `isError`, and object `_meta`), `true` when it declares none, or + * `undefined` when the schema is not a `CallToolResult`. + */ +export function mcpStructuredContentSchema(schema: CodemodeJsonSchema | undefined): CodemodeJsonSchema | undefined { + if (!isObject(schema) || !isObject(schema.properties)) return undefined; + const { content, isError, _meta, structuredContent } = schema.properties; + if (!isObject(content) || content.type !== "array" || !isObject(content.items) || content.items.type !== "object") { + return undefined; + } + if (!isObject(isError) || isError.type !== "boolean" || !isObject(_meta) || _meta.type !== "object") + return undefined; + return isObject(structuredContent) || typeof structuredContent === "boolean" ? structuredContent : true; +} + +function outputType(schema: CodemodeJsonSchema | undefined): string { + const structured = mcpStructuredContentSchema(schema); + if (structured !== undefined) { + const type = schemaToType(structured); + return type === "unknown" ? "CallToolResult" : `CallToolResult<${type}>`; + } + return schema === undefined ? "unknown" : schemaToType(schema); +} + +function renderGlobal(head: string, global: CodemodeTool, indent: string): string { + if (global.signature !== undefined) { + return `${docComment(global.description, indent)}${indent}${head}${global.signature};`; + } + const input = global.inputSchema === undefined ? "unknown" : schemaToType(global.inputSchema); + const output = global.outputSchema === undefined ? "unknown" : schemaToType(global.outputSchema); + return `${docComment(global.description, indent)}${indent}${head}(args: ${input}): Promise<${output}>;`; +} + +function docComment(description: string | undefined, indent: string): string { + const text = description?.trim(); + if (!text) return ""; + const lines = text.replaceAll("*/", "*\\/").split(/\r?\n/); + if (lines.length === 1) return `${indent}/** ${lines[0]} */\n`; + return `${indent}/**\n${lines.map((line) => `${indent} *${line ? ` ${line}` : ""}`).join("\n")}\n${indent} */\n`; +} + +function propertyKey(name: string): string { + return IDENTIFIER.test(name) ? name : JSON.stringify(name); +} + +function isObject(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function union(types: string[]): string { + const unique = [...new Set(types)]; + if (unique.includes("unknown")) return "unknown"; + return unique.length === 0 ? "never" : unique.join(" | "); +} + +/** + * Convert a JSON Schema to a TypeScript type expression: objects on one line (`{ a: string; b?: number; }`) with properties sorted by name, + * or one property per line with `//` comments when a property has a description; `Array` for + * arrays. Local references (`#/$defs/...`, `#/definitions/...`) resolve against `schema`; + * recursive and remote references render as `unknown`. A result longer than `maxChars` renders as + * `unknown`. + */ +export function schemaToType(schema: CodemodeJsonSchema, options: { maxChars?: number } = {}): string { + const type = toType(schema, { root: schema, resolving: new Set(), expansions: 0 }); + return options.maxChars !== undefined && type.length > options.maxChars ? "unknown" : type; +} + +interface SchemaContext { + root: CodemodeJsonSchema; + /** References being expanded on the current path, to stop at recursive types. */ + resolving: Set; + expansions: number; +} + +function resolveRef(ref: string, root: CodemodeJsonSchema): CodemodeJsonSchema | undefined { + if (ref !== "#" && !ref.startsWith("#/")) return undefined; + let current: unknown = root; + for (const segment of ref.slice(2).split("/").filter(Boolean)) { + const key = decodeURIComponent(segment).replaceAll("~1", "/").replaceAll("~0", "~"); + if (!isObject(current) || !(key in current)) return undefined; + current = current[key]; + } + return typeof current === "boolean" || isObject(current) ? current : undefined; +} + +function toType(schema: CodemodeJsonSchema, context: SchemaContext): string { + if (schema === true) return "unknown"; + if (schema === false) return "never"; + if (!isObject(schema)) return "unknown"; + if (typeof schema.$ref === "string") { + const ref = schema.$ref; + if (context.resolving.has(ref) || context.expansions >= MAX_REF_EXPANSIONS) return "unknown"; + const target = resolveRef(ref, context.root); + if (target === undefined) return "unknown"; + context.expansions++; + context.resolving.add(ref); + try { + return toType(target, context); + } finally { + context.resolving.delete(ref); + } + } + + if ("const" in schema) return JSON.stringify(schema.const) ?? "unknown"; + if (Array.isArray(schema.enum)) return union(schema.enum.map((value) => JSON.stringify(value) ?? "unknown")); + + const variants = Array.isArray(schema.anyOf) ? schema.anyOf : Array.isArray(schema.oneOf) ? schema.oneOf : undefined; + if (variants) return union(variants.map((variant) => toType(variant as CodemodeJsonSchema, context))); + if (Array.isArray(schema.allOf)) { + const parts = schema.allOf + .map((part) => toType(part as CodemodeJsonSchema, context)) + .filter((part) => part !== "unknown"); + return parts.length === 0 + ? "unknown" + : parts.map((part) => (part.includes(" | ") ? `(${part})` : part)).join(" & "); + } + + const type = schema.type; + if (Array.isArray(type)) return union(type.map((entry) => toType({ ...schema, type: entry }, context))); + switch (type) { + case "string": + return "string"; + case "number": + case "integer": + return "number"; + case "boolean": + return "boolean"; + case "null": + return "null"; + case "array": + return arrayType(schema, context); + case "object": + return objectType(schema, context); + case undefined: + if ("properties" in schema || "additionalProperties" in schema || "required" in schema) { + return objectType(schema, context); + } + if ("items" in schema || "prefixItems" in schema) return arrayType(schema, context); + return "unknown"; + default: + return "unknown"; + } +} + +function arrayType(schema: Record, context: SchemaContext): string { + if (schema.items !== undefined && !Array.isArray(schema.items)) { + return `Array<${toType(schema.items as CodemodeJsonSchema, context)}>`; + } + const tuple = Array.isArray(schema.prefixItems) + ? schema.prefixItems + : Array.isArray(schema.items) + ? schema.items + : []; + if (tuple.length > 0) return `[${tuple.map((item) => toType(item as CodemodeJsonSchema, context)).join(", ")}]`; + return "unknown[]"; +} + +function descriptionOf(property: unknown): string { + return isObject(property) && typeof property.description === "string" ? property.description.trim() : ""; +} + +function objectType(schema: Record, context: SchemaContext): string { + const properties = isObject(schema.properties) ? schema.properties : {}; + const required = new Set(Array.isArray(schema.required) ? schema.required : []); + const names = Object.keys(properties).sort(); + const members = names.map((name) => { + const optional = required.has(name) ? "" : "?"; + return `${propertyKey(name)}${optional}: ${toType(properties[name] as CodemodeJsonSchema, context)};`; + }); + const additional = schema.additionalProperties; + if (additional !== undefined && additional !== false) { + const type = additional === true ? "unknown" : toType(additional as CodemodeJsonSchema, context); + members.push(`[key: string]: ${type};`); + } else if (additional === undefined && names.length === 0) { + members.push("[key: string]: unknown;"); + } + if (members.length === 0) return "{}"; + if (!names.some((name) => descriptionOf(properties[name]))) return `{ ${members.join(" ")} }`; + + const lines = ["{"]; + names.forEach((name, index) => { + for (const line of descriptionOf(properties[name]).split(/\r?\n/)) { + if (line.trim()) lines.push(`${INDENT}// ${line.trim()}`); + } + lines.push(`${INDENT}${members[index].replaceAll("\n", `\n${INDENT}`)}`); + }); + for (const member of members.slice(names.length)) lines.push(`${INDENT}${member}`); + lines.push("}"); + return lines.join("\n"); +} diff --git a/packages/codemode/src/identifier.ts b/packages/codemode/src/identifier.ts new file mode 100644 index 000000000..75f918adf --- /dev/null +++ b/packages/codemode/src/identifier.ts @@ -0,0 +1,12 @@ +/** + * The identifier a script uses for a tool: characters that are not valid in a JavaScript + * identifier become `_`. `mcp__docs__search` stays as is, `my-tool` becomes `my_tool`. + */ +export function toCodemodeIdentifier(name: string): string { + let identifier = ""; + for (const char of name) { + const valid = identifier === "" ? /^[A-Za-z_$]$/.test(char) : /^[A-Za-z0-9_$]$/.test(char); + identifier += valid ? char : "_"; + } + return identifier === "" ? "_" : identifier; +} diff --git a/packages/codemode/src/index.ts b/packages/codemode/src/index.ts new file mode 100644 index 000000000..2ddb93625 --- /dev/null +++ b/packages/codemode/src/index.ts @@ -0,0 +1,36 @@ +export { + DEFAULT_INPUT_SCHEMA_MAX_CHARS, + MCP_TYPESCRIPT_PREAMBLE, + mcpStructuredContentSchema, + type RenderDeclarationsOptions, + renderDeclarations, + renderToolSample, + renderToolSignature, + schemaToType, +} from "./declarations.ts"; +export { toCodemodeIdentifier } from "./identifier.ts"; +export { CodemodeSandbox } from "./runtime/host.ts"; +export { MAX_STORE_TOTAL_CHARS, MAX_STORE_VALUE_CHARS } from "./runtime/prelude-source.ts"; +export { + CODEMODE_OPTIONS_PREFIX, + CODEMODE_SOURCE_GRAMMAR, + CodemodeSourceError, + type CodemodeSourceOptions, + type ParsedCodemodeSource, + parseCodemodeSource, +} from "./source.ts"; +export type { + CodemodeCall, + CodemodeCallStatus, + CodemodeError, + CodemodeErrorKind, + CodemodeExecuteOptions, + CodemodeJsonSchema, + CodemodeOutputItem, + CodemodeResult, + CodemodeSandboxOptions, + CodemodeStoreWrites, + CodemodeTool, + CodemodeToolContext, +} from "./types.ts"; +export { type CodemodeWasmModule, loadQuickJSWasm } from "./wasm.ts"; diff --git a/packages/codemode/src/runtime/host.ts b/packages/codemode/src/runtime/host.ts new file mode 100644 index 000000000..7a7308bf5 --- /dev/null +++ b/packages/codemode/src/runtime/host.ts @@ -0,0 +1,361 @@ +import { Worker } from "node:worker_threads"; +import { toCodemodeIdentifier } from "../identifier.ts"; +import type { + CodemodeCall, + CodemodeCallStatus, + CodemodeError, + CodemodeExecuteOptions, + CodemodeOutputItem, + CodemodeResult, + CodemodeSandboxOptions, + CodemodeStoreWrites, + CodemodeTool, +} from "../types.ts"; +import { type CodemodeWasmModule, loadQuickJSWasm } from "../wasm.ts"; +import { + type HostToWorkerMessage, + isWorkerToHostMessage, + type WorkerData, + type WorkerToHostMessage, +} from "./protocol.ts"; + +const DEFAULT_TIMEOUT_MS = 300_000; +const IDENTIFIER = /^[A-Za-z_$][A-Za-z0-9_$]*$/; +const RESERVED_GLOBALS: ReadonlySet = new Set([ + "tools", + "ALL_TOOLS", + "console", + "text", + "image", + "exit", + "globalThis", + "store", + "load", +]); + +function errorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} + +function serializeStore(store: Readonly> | undefined): Record { + const serialized: Record = {}; + for (const [key, value] of Object.entries(store ?? {})) { + const json = JSON.stringify(value); + if (json !== undefined) serialized[key] = json; + } + return serialized; +} + +function parseStoreWrites(json: string): CodemodeStoreWrites { + const writes: CodemodeStoreWrites = { set: {}, delete: [] }; + for (const [key, value] of JSON.parse(json) as [string, string?][]) { + if (value === undefined) writes.delete.push(key); + else writes.set[key] = JSON.parse(value); + } + return writes; +} + +function defaultWorkerUrl(): URL { + // `.ts` when running from source (tests, tsx), `.js` from the published dist. + return new URL(import.meta.url.endsWith(".ts") ? "./worker.ts" : "./worker.js", import.meta.url); +} + +interface PendingCall { + record: CodemodeCall | undefined; + startedAt: number; + controller: AbortController; +} + +interface ExecutionOptions { + code: string; + tools: ReadonlyMap; + globals: ReadonlyMap; + timeoutMs: number; + signal: AbortSignal | undefined; + memoryLimitBytes: number | undefined; + store: Record; + wasm: Promise; + workerUrl: URL; +} + +/** + * One script run in its own worker and QuickJS VM. A fresh worker per run keeps + * termination simple: a runaway script, including one that only spins the + * microtask queue, is killed with `terminate()` and cannot poison a later run. + */ +class Execution { + readonly promise: Promise; + private resolveResult!: (result: CodemodeResult) => void; + private worker: Worker | undefined; + private readonly interrupt = new SharedArrayBuffer(4); + private readonly tools: ReadonlyMap; + private readonly globals: ReadonlyMap; + private readonly signal: AbortSignal | undefined; + private readonly timer: NodeJS.Timeout | undefined; + private readonly output: CodemodeOutputItem[] = []; + private readonly calls: CodemodeCall[] = []; + private readonly pending = new Map(); + private finished = false; + + constructor(options: ExecutionOptions) { + this.promise = new Promise((resolve) => { + this.resolveResult = resolve; + }); + this.tools = options.tools; + this.globals = options.globals; + this.signal = options.signal; + + if (Number.isFinite(options.timeoutMs)) { + this.timer = setTimeout(() => { + this.finish({ kind: "timeout", message: `Execution timed out after ${options.timeoutMs} ms` }); + }, options.timeoutMs); + } + + if (options.signal) { + if (options.signal.aborted) { + this.onAbort(); + } else { + options.signal.addEventListener("abort", this.onAbort, { once: true }); + } + } + + options.wasm.then( + (wasm) => this.start(options, wasm), + (error: unknown) => { + this.finish({ kind: "sandbox", message: `Failed to load QuickJS: ${errorMessage(error)}` }); + }, + ); + } + + abort(message: string): Promise { + this.finish({ kind: "aborted", message }); + return this.promise; + } + + private start(options: ExecutionOptions, wasm: CodemodeWasmModule): void { + if (this.finished) return; + const workerData: WorkerData = { + code: options.code, + tools: [...options.tools.values()].map((tool) => ({ + name: tool.name, + jsName: toCodemodeIdentifier(tool.name), + description: tool.description ?? "", + })), + globals: [...options.globals.values()].map((global) => ({ + name: global.name, + spread: global.spread === true, + })), + wasm, + memoryLimitBytes: options.memoryLimitBytes, + store: options.store, + interrupt: this.interrupt, + }; + let worker: Worker; + try { + worker = new Worker(options.workerUrl, { workerData }); + } catch (error) { + this.finish({ kind: "sandbox", message: `Failed to start worker: ${errorMessage(error)}` }); + return; + } + this.worker = worker; + worker.on("message", (message: unknown) => this.handleMessage(message)); + worker.on("error", (error: unknown) => { + this.finish({ + kind: "sandbox", + name: error instanceof Error ? error.name : undefined, + message: errorMessage(error), + }); + }); + worker.on("exit", (code) => { + this.finish({ kind: "sandbox", message: `Worker exited with code ${code} before the script settled` }); + }); + } + + private readonly onAbort = (): void => { + const reason: unknown = this.signal?.reason; + this.finish({ kind: "aborted", message: reason instanceof Error ? reason.message : "Execution aborted" }); + }; + + private post(message: HostToWorkerMessage): void { + this.worker?.postMessage(message); + } + + private handleMessage(message: unknown): void { + if (this.finished || !isWorkerToHostMessage(message)) return; + switch (message.type) { + case "output": + this.output.push(message.item); + break; + case "call": + void this.handleCall(message); + break; + case "done": + this.handleDone(message); + break; + case "crash": + this.finish({ kind: "sandbox", message: message.message }); + break; + } + } + + private handleDone(message: Extract): void { + if (!message.ok) { + const parsed = JSON.parse(message.error) as Omit; + this.finish({ kind: "script", ...parsed }); + return; + } + this.finish(undefined, message.value === undefined ? undefined : JSON.parse(message.value), message.writes); + } + + private async handleCall(message: Extract): Promise { + const { id, name } = message; + const isTool = message.target === "tool"; + const record: CodemodeCall | undefined = isTool ? { name, status: "cancelled", durationMs: 0 } : undefined; + if (record) this.calls.push(record); + const pending: PendingCall = { record, startedAt: performance.now(), controller: new AbortController() }; + this.pending.set(id, pending); + + let status: CodemodeCallStatus; + let reply: HostToWorkerMessage; + try { + const tool = (isTool ? this.tools : this.globals).get(name); + if (!tool) throw new Error(`Unknown ${isTool ? "tool" : "global"} "${name}"`); + const args: unknown = message.args === undefined ? undefined : JSON.parse(message.args); + const value = await tool.execute(args, { signal: pending.controller.signal }); + reply = { type: "result", id, ok: true, payload: value === undefined ? undefined : JSON.stringify(value) }; + status = "ok"; + } catch (error) { + reply = { type: "result", id, ok: false, payload: errorMessage(error) }; + status = "error"; + } + + // Already cancelled by finish(): the record keeps "cancelled" and the + // worker is gone or going. + if (!this.pending.delete(id)) return; + if (record) { + record.status = status; + record.durationMs = performance.now() - pending.startedAt; + } + this.post(reply); + } + + private finish(error: CodemodeError | undefined, value?: unknown, writes?: string): void { + if (this.finished) return; + this.finished = true; + clearTimeout(this.timer); + this.signal?.removeEventListener("abort", this.onAbort); + + const now = performance.now(); + for (const pending of this.pending.values()) { + if (pending.record) pending.record.durationMs = now - pending.startedAt; + pending.controller.abort(); + } + this.pending.clear(); + + const result: CodemodeResult = error + ? { ok: false, error, output: this.output, calls: this.calls } + : { + ok: true, + value, + output: this.output, + calls: this.calls, + storeWrites: writes === undefined ? { set: {}, delete: [] } : parseStoreWrites(writes), + }; + if (!this.worker) { + this.resolveResult(result); + return; + } + Atomics.store(new Int32Array(this.interrupt), 0, 1); + this.worker + .terminate() + .catch(() => undefined) + .then(() => this.resolveResult(result)); + } +} + +/** + * Runs JavaScript in a QuickJS VM (a separate wasm instance) inside a worker + * thread. The script sees `tools.(args)` for every registered tool, `ALL_TOOLS`, + * the output helpers `text`, `image`, `exit`, and `console.*`, `store`/`load`, and the + * configured globals; nothing else (no timers, `fetch`, `process`, `require`, modules). + * + * Each `execute()` gets its own worker and VM; the sandbox only holds the tool + * table and defaults. `close()` aborts in-flight executions. + */ +export class CodemodeSandbox { + private readonly toolsByName = new Map(); + private readonly globalsByName = new Map(); + private readonly timeoutMs: number; + private readonly memoryLimitBytes: number | undefined; + private readonly wasm: CodemodeWasmModule | Promise | undefined; + private readonly workerUrl: URL; + private readonly running = new Set(); + private closed = false; + + constructor(options: CodemodeSandboxOptions = {}) { + this.timeoutMs = options.timeoutMs ?? DEFAULT_TIMEOUT_MS; + this.memoryLimitBytes = options.memoryLimitBytes; + this.wasm = options.wasm; + this.workerUrl = options.workerUrl ?? defaultWorkerUrl(); + for (const tool of options.tools ?? []) this.registerTool(tool); + const namespaces = new Set(); + for (const global of options.globals ?? []) { + const parts = global.name.split("."); + if (parts.length > 2 || !parts.every((part) => IDENTIFIER.test(part)) || RESERVED_GLOBALS.has(parts[0])) { + throw new Error(`Invalid global name "${global.name}"`); + } + if (this.globalsByName.has(global.name)) throw new Error(`Global "${global.name}" is already registered`); + if (parts.length === 2) namespaces.add(parts[0]); + this.globalsByName.set(global.name, global); + } + for (const name of namespaces) { + if (this.globalsByName.has(name)) throw new Error(`Global "${name}" conflicts with the namespace "${name}"`); + } + } + + /** Throws if a tool with the same name is already registered. */ + registerTool(tool: CodemodeTool): void { + if (this.toolsByName.has(tool.name)) throw new Error(`Tool "${tool.name}" is already registered`); + this.toolsByName.set(tool.name, tool); + } + + unregisterTool(name: string): boolean { + return this.toolsByName.delete(name); + } + + get tools(): CodemodeTool[] { + return [...this.toolsByName.values()]; + } + + get globals(): CodemodeTool[] { + return [...this.globalsByName.values()]; + } + + /** + * `code` is an async function body: `return` and top-level `await` work. + * Never rejects for script failures; those come back as `{ ok: false }`. + * The script can use `store(key, value)` and `load(key)` on `options.store`. + */ + execute(code: string, options: CodemodeExecuteOptions = {}): Promise { + if (this.closed) return Promise.reject(new Error("Sandbox is closed")); + const execution = new Execution({ + code, + tools: new Map(this.toolsByName), + globals: this.globalsByName, + timeoutMs: options.timeoutMs ?? this.timeoutMs, + signal: options.signal, + memoryLimitBytes: this.memoryLimitBytes, + store: serializeStore(options.store), + wasm: this.wasm === undefined ? loadQuickJSWasm() : Promise.resolve(this.wasm), + workerUrl: this.workerUrl, + }); + this.running.add(execution); + return execution.promise.finally(() => this.running.delete(execution)); + } + + /** Aborts in-flight executions (they resolve with `kind: "aborted"`) and rejects new ones. */ + async close(): Promise { + this.closed = true; + await Promise.all([...this.running].map((execution) => execution.abort("Sandbox closed"))); + } +} diff --git a/packages/codemode/src/runtime/prelude-source.ts b/packages/codemode/src/runtime/prelude-source.ts new file mode 100644 index 000000000..7b08b22fe --- /dev/null +++ b/packages/codemode/src/runtime/prelude-source.ts @@ -0,0 +1,324 @@ +/** + * JavaScript evaluated inside the QuickJS VM before the script runs. + * + * The VM is its own wasm instance, so nothing here guards a realm boundary. + * The prelude keeps the host bridge in a closure so the script cannot call it + * directly, and builds `tools`, `ALL_TOOLS`, the output helpers (`text`, `image`, + * `exit`, `console`), and globals on top of it. Tool arguments and results cross + * as JSON strings and are parsed on this side. + * + * Output helpers: `text(value)` appends a text item (non-strings are + * JSON-stringified), `image(urlOrItem)` appends an image from a base64 `data:` URL, an + * `{ image_url }` object, or an MCP `ImageContent` block, and `exit()` ends the script + * successfully. `console.*` appends text items like `text()`. + * + * `store(key, value)` and `load(key)` are synchronous: they work on a snapshot of + * JSON text passed in as `storeJson`, and the keys the script wrote are reported + * with a successful "done". + * + * Evaluates to a function `(bridge, toolsJson, globalsJson, storeJson) => { settle, run, stalled }`. + * `toolsJson` lists `{ name, jsName, description }`: `tools[jsName]` and `tools[name]` call the + * tool, and `ALL_TOOLS` lists `{ name: jsName, description }`. + * `stalled()` reports a script that has not finished while no host call is pending: with no timers + * or I/O in the VM, nothing can ever resume it. + * `globalsJson` lists `{ name, spread }`; `a.b` names are grouped into a frozen `a` object. + * `bridge(kind, a, b, c)` with kind "call" or "global" (id, name, argsJson), + * "output" ("text", text) or ("image", data, mimeType), or "done" (ok, valueJsonOrErrorJson, writesJson). + */ +export const MAX_STORE_VALUE_CHARS = 256 * 1024; +export const MAX_STORE_TOTAL_CHARS = 1024 * 1024; + +const IMAGE_HELPER_EXPECTS = + "image expects a non-empty image URL string, an object with image_url, or a raw MCP image block"; + +export const PRELUDE_SOURCE: string = `(function (bridge, toolsJson, globalsJson, storeJson) { + "use strict"; + const stringify = JSON.stringify; + const parse = JSON.parse; + const promiseThen = Promise.prototype.then; + const ErrorCtor = Error; + const TypeErrorCtor = TypeError; + const pending = new Map(); + let nextId = 1; + let finished = false; + // Thrown by exit() to unwind the script after it already reported success. + const EXIT = Object.freeze({}); + + function done(ok, payload, writes) { + if (finished) return; + finished = true; + bridge("done", ok, payload, writes); + } + + function serialize(value) { + return value === undefined ? undefined : stringify(value); + } + + // QuickJS stacks list frames only. Prefix "Name: message" like V8 so the + // text reads the same as a Node error, and drop this prelude's frames. + function errorText(error) { + const head = error.message ? error.name + ": " + error.message : String(error.name); + const frames = + typeof error.stack === "string" + ? error.stack.split("\\n").filter((line) => line.trim() && !line.includes("codemode-prelude.js")) + : []; + return [head, ...frames].join("\\n"); + } + + function format(value) { + if (typeof value === "string") return value; + if (value instanceof ErrorCtor) return errorText(value); + try { + const json = stringify(value); + return json === undefined ? String(value) : json; + } catch { + return String(value); + } + } + + function describeError(error) { + if (error instanceof ErrorCtor) { + return stringify({ name: error.name, message: error.message, stack: errorText(error) }); + } + return stringify({ message: format(error) }); + } + + function caller(kind, name, spread) { + return (...args) => + new Promise((resolve, reject) => { + let json; + try { + json = serialize(spread ? args : args[0]); + } catch (error) { + reject(error); + return; + } + const id = nextId++; + pending.set(id, { resolve, reject }); + bridge(kind, id, name, json); + }); + } + + const tools = Object.create(null); + const allTools = []; + for (const { name, jsName, description } of parse(toolsJson)) { + const fn = caller("call", name); + // The first tool wins when two names normalize to the same identifier. + if (!(jsName in tools)) { + tools[jsName] = fn; + allTools.push(Object.freeze({ name: jsName, description })); + } + if (!(name in tools)) tools[name] = fn; + } + Object.freeze(tools); + Object.freeze(allTools); + + const namespaces = new Map(); + for (const { name, spread } of parse(globalsJson)) { + const fn = caller("global", name, spread); + const dot = name.indexOf("."); + if (dot === -1) { + Object.defineProperty(globalThis, name, { value: fn, enumerable: true }); + continue; + } + const namespace = name.slice(0, dot); + if (!namespaces.has(namespace)) namespaces.set(namespace, Object.create(null)); + namespaces.get(namespace)[name.slice(dot + 1)] = fn; + } + for (const [namespace, members] of namespaces) { + Object.defineProperty(globalThis, namespace, { value: Object.freeze(members), enumerable: true }); + } + + // key -> JSON text. Sizes count key and JSON characters. + const stored = new Map(Object.entries(parse(storeJson))); + const writes = new Map(); + let storedChars = 0; + for (const [key, json] of stored) storedChars += key.length + json.length; + + function checkKey(name, key) { + if (typeof key !== "string") throw new TypeError(name + "() key must be a string"); + } + + function store(key, value) { + checkKey("store", key); + const previous = stored.has(key) ? key.length + stored.get(key).length : 0; + if (value === undefined) { + stored.delete(key); + storedChars -= previous; + writes.set(key, undefined); + return; + } + let json; + try { + json = stringify(value); + } catch (error) { + throw new TypeError("store(" + stringify(key) + ") value is not JSON-serializable: " + format(error)); + } + if (json === undefined) { + throw new TypeError("store(" + stringify(key) + ") value is not JSON-serializable"); + } + if (json.length > ${MAX_STORE_VALUE_CHARS}) { + throw new RangeError("store(" + stringify(key) + ") value exceeds ${MAX_STORE_VALUE_CHARS} characters of JSON"); + } + const next = storedChars - previous + key.length + json.length; + if (next > ${MAX_STORE_TOTAL_CHARS}) { + throw new RangeError("store is full: stored values would exceed ${MAX_STORE_TOTAL_CHARS} characters of JSON"); + } + stored.set(key, json); + storedChars = next; + writes.set(key, json); + } + + function load(key) { + checkKey("load", key); + const json = stored.get(key); + return json === undefined ? undefined : parse(json); + } + + function serializeWrites() { + const entries = []; + for (const [key, json] of writes) entries.push(json === undefined ? [key] : [key, json]); + return stringify(entries); + } + + Object.defineProperty(globalThis, "store", { value: store, enumerable: true }); + Object.defineProperty(globalThis, "load", { value: load, enumerable: true }); + + // Primitives become their string form, everything else JSON. + function outputText(value) { + if (value === undefined || value === null || typeof value !== "object" && typeof value !== "function") { + return String(value); + } + const json = stringify(value); + return json === undefined ? String(value) : json; + } + + function text(value) { + let rendered; + try { + rendered = outputText(value); + } catch (error) { + throw new TypeErrorCtor(error instanceof ErrorCtor ? error.message : String(error)); + } + if (!finished) bridge("output", "text", rendered); + } + + function imageUrl(value) { + if (typeof value === "string") return value; + if (typeof value !== "object" || value === null || Array.isArray(value)) { + throw new TypeErrorCtor(${JSON.stringify(IMAGE_HELPER_EXPECTS)}); + } + if (value.image_url !== undefined) { + if (typeof value.image_url !== "string") throw new TypeErrorCtor(${JSON.stringify(IMAGE_HELPER_EXPECTS)}); + return value.image_url; + } + if (typeof value.type !== "string") throw new TypeErrorCtor(${JSON.stringify(IMAGE_HELPER_EXPECTS)}); + if (value.type !== "image") { + throw new TypeErrorCtor('image only accepts MCP image blocks, got "' + value.type + '"'); + } + if (typeof value.data !== "string" || value.data === "") throw new TypeErrorCtor("image expected MCP image data"); + if (value.data.toLowerCase().startsWith("data:")) return value.data; + const mimeType = typeof value.mimeType === "string" && value.mimeType ? value.mimeType : "application/octet-stream"; + return "data:" + mimeType + ";base64," + value.data; + } + + function image(value) { + const url = imageUrl(value); + if (url === "") throw new TypeErrorCtor(${JSON.stringify(IMAGE_HELPER_EXPECTS)}); + const colon = url.indexOf(":"); + const scheme = colon === -1 ? "" : url.slice(0, colon).toLowerCase(); + if (scheme === "http" || scheme === "https") { + throw new TypeErrorCtor("remote image URLs are not supported in tool outputs. Pass a base64 data URI instead"); + } + const comma = url.indexOf(","); + const header = comma === -1 ? [] : url.slice(colon + 1, comma).split(";"); + if (scheme !== "data" || comma === -1 || header.slice(1).every((part) => part.toLowerCase() !== "base64")) { + throw new TypeErrorCtor("invalid image output. Pass a base64 data URI instead"); + } + if (!finished) bridge("output", "image", url.slice(comma + 1), header[0] || "application/octet-stream"); + } + + function exit() { + let writesJson; + try { + writesJson = serializeWrites(); + } catch (error) { + done(false, describeError(error)); + throw EXIT; + } + done(true, undefined, writesJson); + throw EXIT; + } + + const console = {}; + for (const level of ["log", "info", "warn", "error", "debug"]) { + console[level] = (...args) => { + if (!finished) bridge("output", "text", args.map(format).join(" ")); + }; + } + Object.freeze(console); + + Object.defineProperty(globalThis, "tools", { value: tools, enumerable: true }); + Object.defineProperty(globalThis, "ALL_TOOLS", { value: allTools, enumerable: true }); + Object.defineProperty(globalThis, "console", { value: console, enumerable: true }); + Object.defineProperty(globalThis, "text", { value: text, enumerable: true }); + Object.defineProperty(globalThis, "image", { value: image, enumerable: true }); + Object.defineProperty(globalThis, "exit", { value: exit, enumerable: true }); + + return { + settle(id, ok, payload) { + const entry = pending.get(id); + if (!entry) return; + pending.delete(id); + if (!ok) { + entry.reject(new ErrorCtor(payload)); + return; + } + let value; + try { + value = payload === undefined ? undefined : parse(payload); + } catch (error) { + entry.reject(error); + return; + } + entry.resolve(value); + }, + run(fn) { + let promise; + try { + promise = fn(tools, console); + } catch (error) { + done(false, describeError(error)); + return; + } + promiseThen.call( + promise, + (value) => { + let json; + try { + json = serialize(value); + } catch (error) { + done(false, describeError(error)); + return; + } + done(true, json, serializeWrites()); + }, + (error) => { + done(false, describeError(error)); + }, + ); + }, + stalled() { + if (finished || pending.size > 0) return false; + done( + false, + stringify({ + name: "Error", + message: + "The script is waiting on a promise that can never settle: no tool call is pending, and timers do not exist here.", + }), + ); + return true; + }, + }; +})`; diff --git a/packages/codemode/src/runtime/protocol.ts b/packages/codemode/src/runtime/protocol.ts new file mode 100644 index 000000000..b9e8eb63d --- /dev/null +++ b/packages/codemode/src/runtime/protocol.ts @@ -0,0 +1,51 @@ +import type { CodemodeOutputItem } from "../types.ts"; +import type { CodemodeWasmModule } from "../wasm.ts"; + +/** + * Messages between the host (main thread) and the worker. Tool arguments, + * results, and values cross as JSON strings: the worker passes them into and + * out of the QuickJS VM as strings and never builds structured values itself. + */ + +export interface WorkerData { + code: string; + /** `jsName` is the identifier the script uses; `description` is listed in `ALL_TOOLS`. */ + tools: { name: string; jsName: string; description: string }[]; + globals: { name: string; spread: boolean }[]; + /** Compiled `quickjs-wasi` module. Structured clone shares the compiled code with the worker. */ + wasm: CodemodeWasmModule; + memoryLimitBytes: number | undefined; + /** Snapshot for `load()`: key to JSON text. */ + store: Record; + /** + * One Int32 the host sets to non-zero before terminating the worker. The VM's interrupt handler + * polls it, because Bun's `worker.terminate()` cannot stop a thread that is spinning in wasm. + */ + interrupt: SharedArrayBuffer; +} + +/** JSON-encoded `{ name?, message, stack? }` of an error thrown by the script. */ +export type ScriptErrorJson = string; + +export type WorkerToHostMessage = + | { type: "call"; id: number; target: "tool" | "global"; name: string; args: string | undefined } + | { type: "output"; item: CodemodeOutputItem } + /** `writes` is a JSON array of `[key, json]` for `store()` and `[key]` for deletions. */ + | { type: "done"; ok: true; value: string | undefined; writes: string } + | { type: "done"; ok: false; error: ScriptErrorJson } + /** The VM failed outside the script's control, for example a wasm trap. */ + | { type: "crash"; message: string }; + +export type HostToWorkerMessage = + /** `payload` is the JSON result when `ok`, otherwise the error message. */ + { type: "result"; id: number; ok: boolean; payload: string | undefined }; + +export function isWorkerToHostMessage(value: unknown): value is WorkerToHostMessage { + if (typeof value !== "object" || value === null) return false; + const type = (value as { type?: unknown }).type; + return type === "call" || type === "output" || type === "done" || type === "crash"; +} + +export function isHostToWorkerMessage(value: unknown): value is HostToWorkerMessage { + return typeof value === "object" && value !== null && (value as { type?: unknown }).type === "result"; +} diff --git a/packages/codemode/src/runtime/worker.ts b/packages/codemode/src/runtime/worker.ts new file mode 100644 index 000000000..eebb674a0 --- /dev/null +++ b/packages/codemode/src/runtime/worker.ts @@ -0,0 +1,159 @@ +/** + * Worker thread entry. One worker runs one script inside a fresh QuickJS VM + * (a separate wasm instance), relays tool calls and output to the host, and + * reports the result. The host terminates the worker when the script settles, + * times out, or is aborted; the worker exists so that a spinning script never + * blocks the host thread. + * + * Importing this module starts the worker. Hosts that bundle their code (for + * example a Bun compiled executable) add a file that imports + * `@earendil-works/pi-codemode/worker` as a separate entrypoint and pass its URL + * as `workerUrl`. + */ +import { parentPort, workerData } from "node:worker_threads"; +import { JSException, type JSValueHandle, MAX_STACK_SIZE, QuickJS } from "quickjs-wasi"; +import { PRELUDE_SOURCE } from "./prelude-source.ts"; +import { isHostToWorkerMessage, type WorkerData, type WorkerToHostMessage } from "./protocol.ts"; + +function post(message: WorkerToHostMessage): void { + parentPort?.postMessage(message); +} + +function crash(error: unknown): void { + post({ type: "crash", message: error instanceof Error ? `${error.name}: ${error.message}` : String(error) }); +} + +/** + * QuickJS writes engine diagnostics to fd 1 and 2, which the default shim + * forwards to the host's stdout and stderr. That output belongs to the host + * application (for example a TUI), so it is discarded. Reporting every byte as + * written keeps libc from retrying. + */ +function discardOutput(memory: { readonly buffer: ArrayBufferLike }) { + return { + fd_write(_fd: number, iovsPtr: number, iovsLen: number, nwrittenPtr: number): number { + const view = new DataView(memory.buffer); + let written = 0; + for (let i = 0; i < iovsLen; i++) { + written += view.getUint32(iovsPtr + i * 8 + 4, true); + } + view.setUint32(nwrittenPtr, written, true); + return 0; + }, + }; +} + +function describeException(error: JSException): string { + const head = error.message ? `${error.name}: ${error.message}` : error.name; + const stack = error.stack?.trimEnd(); + return JSON.stringify({ name: error.name, message: error.message, stack: stack ? `${head}\n${stack}` : head }); +} + +async function main(data: WorkerData): Promise { + const interrupt = new Int32Array(data.interrupt); + const vm = await QuickJS.create({ + wasm: data.wasm, + memoryLimit: data.memoryLimitBytes, + // Without a guard, deep recursion overflows the wasm stack and traps instead of throwing a + // catchable RangeError. + maxStackSize: MAX_STACK_SIZE, + interruptHandler: () => Atomics.load(interrupt, 0) !== 0, + wasi: discardOutput, + }); + + // Called from the prelude with primitives only. + const bridge = vm.newFunction("bridge", (kind, a, b, c) => { + switch (kind.toString()) { + case "call": + case "global": + post({ + type: "call", + id: a.toNumber(), + target: kind.toString() === "call" ? "tool" : "global", + name: b.toString(), + args: c === undefined || c.isUndefined ? undefined : c.toString(), + }); + break; + case "output": + post({ + type: "output", + item: + a.toString() === "image" + ? { type: "image", data: b.toString(), mimeType: c.toString() } + : { type: "text", text: b.toString() }, + }); + break; + case "done": + if (a.toBoolean()) { + post({ + type: "done", + ok: true, + value: b === undefined || b.isUndefined ? undefined : b.toString(), + writes: c.toString(), + }); + } else { + post({ type: "done", ok: false, error: b.toString() }); + } + break; + } + return vm.undefined; + }); + + // The VM lives until the host terminates the worker, so these handles are never disposed. + const api = vm.withScope((scope) => + scope.escape( + vm.callFunction( + vm.evalCode(PRELUDE_SOURCE, "codemode-prelude.js"), + vm.undefined, + bridge, + vm.newString(JSON.stringify(data.tools)), + vm.newString(JSON.stringify(data.globals)), + vm.newString(JSON.stringify(data.store)), + ), + ), + ); + const settle = api.getProp("settle"); + const run = api.getProp("run"); + const stalled = api.getProp("stalled"); + /** Run queued jobs, then fail a script that waits on nothing that can ever resume it. */ + const drain = () => { + vm.executePendingJobs(); + vm.callFunction(stalled, api).dispose(); + }; + + parentPort?.on("message", (message: unknown) => { + if (!isHostToWorkerMessage(message)) return; + try { + vm.withScope(() => { + vm.callFunction( + settle, + api, + vm.newNumber(message.id), + message.ok ? vm.true : vm.false, + message.payload === undefined ? vm.undefined : vm.newString(message.payload), + ); + }); + drain(); + } catch (error) { + crash(error); + } + }); + + // The prefix shares the first line with the script so reported line numbers + // match the script as written. + let fn: JSValueHandle; + try { + fn = vm.evalCode(`(async (tools, console) => {${data.code}\n})`, "codemode.js"); + } catch (error) { + if (!(error instanceof JSException)) throw error; + post({ type: "done", ok: false, error: describeException(error) }); + return; + } + vm.callFunction(run, api, fn).dispose(); + fn.dispose(); + drain(); +} + +if (parentPort) { + main(workerData as WorkerData).catch(crash); +} diff --git a/packages/codemode/src/source.ts b/packages/codemode/src/source.ts new file mode 100644 index 000000000..86c66f20a --- /dev/null +++ b/packages/codemode/src/source.ts @@ -0,0 +1,115 @@ +/** + * Codemode source format: JavaScript, optionally preceded by one options line. + * + * ```js + * // @options: {"max_output_tokens": 2000, "timeout_ms": 30000} + * const text = await tools.read({ path: "package.json" }); + * text(JSON.parse(text).name); + * ``` + */ + +export const CODEMODE_OPTIONS_PREFIX = "// @options:"; + +const SUPPORTED_FIELDS = ["max_output_tokens", "timeout_ms"] as const; +const SUPPORTED_FIELDS_TEXT = "`max_output_tokens` and `timeout_ms`"; +/** Largest delay `setTimeout` supports, which bounds `timeout_ms`. */ +const MAX_TIMEOUT_MS = 2_147_483_647; + +/** + * Lark grammar for providers with grammar-constrained tool input. It only fixes the shape of the + * options line; the options JSON and the code are checked by {@link parseCodemodeSource}. + */ +export const CODEMODE_SOURCE_GRAMMAR = String.raw` +start: options_source | plain_source +options_source: OPTIONS_LINE NEWLINE SOURCE +plain_source: SOURCE + +OPTIONS_LINE: /[ \t]*\/\/ @options:[^\r\n]*/ +NEWLINE: /\r?\n/ +SOURCE: /[\s\S]+/ +`; + +export interface CodemodeSourceOptions { + /** Token budget for the script's output. */ + maxOutputTokens?: number; + /** Hard deadline for the whole script in milliseconds, including tool calls. */ + timeoutMs?: number; +} + +export interface ParsedCodemodeSource { + /** The script with the options line replaced by an empty line, so line numbers are unchanged. */ + code: string; + options: CodemodeSourceOptions; +} + +export class CodemodeSourceError extends Error { + constructor(message: string) { + super(message); + this.name = "CodemodeSourceError"; + } +} + +function isSafeInteger(value: unknown): value is number { + return typeof value === "number" && Number.isSafeInteger(value) && value >= 0; +} + +function parseOptions(directive: string): CodemodeSourceOptions { + if (directive === "") { + throw new CodemodeSourceError(`@options must be a JSON object with supported fields ${SUPPORTED_FIELDS_TEXT}`); + } + let value: unknown; + try { + value = JSON.parse(directive); + } catch (error) { + throw new CodemodeSourceError( + `@options must be valid JSON with supported fields ${SUPPORTED_FIELDS_TEXT}: ${error instanceof Error ? error.message : String(error)}`, + ); + } + if (typeof value !== "object" || value === null || Array.isArray(value)) { + throw new CodemodeSourceError(`@options must be a JSON object with supported fields ${SUPPORTED_FIELDS_TEXT}`); + } + const fields = value as Record; + for (const key of Object.keys(fields)) { + if (!(SUPPORTED_FIELDS as readonly string[]).includes(key)) { + throw new CodemodeSourceError(`@options only supports ${SUPPORTED_FIELDS_TEXT}; got \`${key}\``); + } + } + const options: CodemodeSourceOptions = {}; + const { max_output_tokens, timeout_ms } = fields; + if (max_output_tokens !== undefined) { + if (!isSafeInteger(max_output_tokens)) { + throw new CodemodeSourceError("@options field `max_output_tokens` must be a non-negative safe integer"); + } + options.maxOutputTokens = max_output_tokens; + } + if (timeout_ms !== undefined) { + if (!isSafeInteger(timeout_ms) || timeout_ms === 0 || timeout_ms > MAX_TIMEOUT_MS) { + throw new CodemodeSourceError( + `@options field \`timeout_ms\` must be a positive integer up to ${MAX_TIMEOUT_MS}`, + ); + } + options.timeoutMs = timeout_ms; + } + return options; +} + +/** + * Split an optional first-line `// @options: {...}` from the script. Throws + * {@link CodemodeSourceError} for empty input and invalid options. + */ +export function parseCodemodeSource(input: string): ParsedCodemodeSource { + if (input.trim() === "") { + throw new CodemodeSourceError( + 'Expected JavaScript source text (non-empty). Provide JS only, optionally with a first line `// @options: {"max_output_tokens": 1000}`.', + ); + } + const newline = input.indexOf("\n"); + const firstLine = (newline === -1 ? input : input.slice(0, newline)).replace(/\r$/, ""); + const trimmed = firstLine.trimStart(); + if (!trimmed.startsWith(CODEMODE_OPTIONS_PREFIX)) return { code: input, options: {} }; + const code = newline === -1 ? "" : input.slice(newline); + if (code.trim() === "") { + throw new CodemodeSourceError("The @options line must be followed by JavaScript source on subsequent lines"); + } + return { code, options: parseOptions(trimmed.slice(CODEMODE_OPTIONS_PREFIX.length).trim()) }; +} diff --git a/packages/codemode/src/types.ts b/packages/codemode/src/types.ts new file mode 100644 index 000000000..b3b90ca8f --- /dev/null +++ b/packages/codemode/src/types.ts @@ -0,0 +1,135 @@ +import type { CodemodeWasmModule } from "./wasm.ts"; + +export interface CodemodeToolContext { + /** + * Aborted when the script finishes (including unawaited calls), the + * execution times out, the caller aborts, or the sandbox is closed. + */ + signal: AbortSignal; +} + +/** A JSON Schema document. Only used to render declarations; values are not validated against it. */ +export type CodemodeJsonSchema = { [key: string]: unknown } | boolean; + +export interface CodemodeTool { + /** + * The script calls tools as `tools.(args)`, where `` is the name with characters that + * are not valid in identifiers replaced by `_` (see `toCodemodeIdentifier`), and also as + * `tools[""](args)`. Globals are called as `(args)` and must be identifiers, or + * `.`, which groups them into a frozen namespace object. + */ + name: string; + /** Shown as a doc comment in {@link renderDeclarations}, and listed in `ALL_TOOLS` for tools. */ + description?: string; + /** Schema of the single argument. Rendered as the parameter type; `unknown` when omitted. */ + inputSchema?: CodemodeJsonSchema; + /** Schema of the resolved value. Rendered as the promise type; `unknown` when omitted. */ + outputSchema?: CodemodeJsonSchema; + /** Globals only: `execute` receives all call arguments as an array instead of the first one. */ + spread?: boolean; + /** + * Globals only: TypeScript parameter list and return type for {@link renderDeclarations}, for + * example `(type: string, id?: string): Promise`. Replaces the rendering from the schemas. + */ + signature?: string; + /** + * `args` is whatever the script passed, after a JSON round trip. The return + * value must be JSON-serializable; a thrown error surfaces in the script as + * an `Error` with the same message. + */ + execute(args: unknown, context: CodemodeToolContext): Promise | unknown; +} + +/** + * One item of the script's output, in the order the script produced it: `text()` and `console.*` + * produce text items, `image()` image items. `data` is base64. + */ +export type CodemodeOutputItem = { type: "text"; text: string } | { type: "image"; data: string; mimeType: string }; + +export type CodemodeCallStatus = "ok" | "error" | "cancelled"; + +export interface CodemodeCall { + name: string; + status: CodemodeCallStatus; + durationMs: number; +} + +export type CodemodeErrorKind = + /** The script threw or failed to parse. `name` and `stack` come from the script's error. */ + | "script" + /** The overall deadline expired. The worker was terminated. */ + | "timeout" + /** The caller's signal fired or the sandbox was closed. The worker was terminated. */ + | "aborted" + /** The worker or VM failed outside the script's control (for example a wasm trap or a missing worker file). */ + | "sandbox"; + +export interface CodemodeError { + kind: CodemodeErrorKind; + name?: string; + message: string; + stack?: string; +} + +/** Keys the script changed with `store()`. Only successful executions report writes. */ +export interface CodemodeStoreWrites { + set: Record; + /** Keys stored as `undefined`. */ + delete: string[]; +} + +/** `output` is kept for failed executions too, up to the failure. `exit()` completes with `value: undefined`. */ +export type CodemodeResult = + | { + ok: true; + value: unknown; + output: CodemodeOutputItem[]; + calls: CodemodeCall[]; + storeWrites: CodemodeStoreWrites; + } + | { ok: false; error: CodemodeError; output: CodemodeOutputItem[]; calls: CodemodeCall[] }; + +export interface CodemodeSandboxOptions { + tools?: CodemodeTool[]; + /** + * Functions exposed as top-level identifiers instead of on `tools`, for host helpers such as + * attaching an image to the result. They behave like tools (JSON round trip, promise result) + * but are not recorded in `result.calls`. Names must be identifiers and may not shadow the + * built-in globals (`tools`, `ALL_TOOLS`, `console`, `text`, `image`, `exit`, `store`, `load`). + */ + globals?: CodemodeTool[]; + /** + * Overall deadline per execution, including time spent in tools. `Infinity` disables the + * deadline; the execution then only ends when the script settles or is aborted. + * Default: 300000. + */ + timeoutMs?: number; + /** + * Maximum memory the QuickJS VM may allocate. Allocations beyond it fail inside the script as + * `InternalError: out of memory`. Default: no limit beyond wasm32's 4 GiB address space. + */ + memoryLimitBytes?: number; + /** + * Compiled `quickjs-wasi/quickjs.wasm`, usually from {@link loadQuickJSWasm}. Default: + * `loadQuickJSWasm()`, the file in the installed `quickjs-wasi` package. Pass it when that file + * is not on disk, for example in a Bun compiled executable. + */ + wasm?: CodemodeWasmModule | Promise; + /** + * Worker entry that imports `@earendil-works/pi-codemode/worker`. Default: this package's own + * worker file. Pass it when this package is bundled, since the default is resolved relative to + * the module that creates the sandbox. + */ + workerUrl?: URL; +} + +export interface CodemodeExecuteOptions { + signal?: AbortSignal; + /** Overrides the sandbox default for this execution. */ + timeoutMs?: number; + /** + * Values the script reads with `load(key)`. Must be JSON-serializable. The script's own + * `store()` calls come back as `result.storeWrites`; persisting them is up to the caller. + */ + store?: Readonly>; +} diff --git a/packages/codemode/src/wasm.ts b/packages/codemode/src/wasm.ts new file mode 100644 index 000000000..7e2edec34 --- /dev/null +++ b/packages/codemode/src/wasm.ts @@ -0,0 +1,35 @@ +import { readFile } from "node:fs/promises"; +import { createRequire } from "node:module"; + +/** + * A compiled `WebAssembly.Module` of `quickjs-wasi/quickjs.wasm`. Typed opaquely because the Node + * type definitions do not declare the WebAssembly globals (they live in the DOM lib). + */ +export type CodemodeWasmModule = object; + +interface WebAssemblyGlobal { + compile(bytes: Uint8Array): Promise; +} + +const modules = new Map>(); + +/** + * Read and compile the QuickJS wasm once per path. `path` defaults to the file in the installed + * `quickjs-wasi` package; pass it when that file lives elsewhere, for example embedded in a Bun + * compiled executable. A failed load is retried on the next call. + */ +export function loadQuickJSWasm(path?: string): Promise { + const resolved = path ?? createRequire(import.meta.url).resolve("quickjs-wasi/quickjs.wasm"); + let module = modules.get(resolved); + if (!module) { + const { WebAssembly } = globalThis as unknown as { WebAssembly: WebAssemblyGlobal }; + module = readFile(resolved) + .then((bytes) => WebAssembly.compile(bytes)) + .catch((error: unknown) => { + modules.delete(resolved); + throw error; + }); + modules.set(resolved, module); + } + return module; +} diff --git a/packages/codemode/test/declarations.test.ts b/packages/codemode/test/declarations.test.ts new file mode 100644 index 000000000..043f69cff --- /dev/null +++ b/packages/codemode/test/declarations.test.ts @@ -0,0 +1,243 @@ +import { describe, expect, it } from "vitest"; +import { + mcpStructuredContentSchema, + renderDeclarations, + renderToolSample, + renderToolSignature, + schemaToType, +} from "../src/index.ts"; + +const execute = () => undefined; + +function mcpResultSchema(structuredContent?: unknown) { + return { + type: "object", + properties: { + content: { type: "array", items: { type: "object" } }, + ...(structuredContent === undefined ? {} : { structuredContent }), + isError: { type: "boolean" }, + _meta: { type: "object" }, + }, + required: ["content"], + }; +} + +describe("schemaToType", () => { + it("renders primitives, literals, and unions", () => { + expect(schemaToType({ type: "string" })).toBe("string"); + expect(schemaToType({ type: "integer" })).toBe("number"); + expect(schemaToType({ type: ["string", "null"] })).toBe("string | null"); + expect(schemaToType({ const: "a" })).toBe('"a"'); + expect(schemaToType({ enum: ["a", 1, null] })).toBe('"a" | 1 | null'); + expect(schemaToType({ anyOf: [{ type: "string" }, { type: "number" }] })).toBe("string | number"); + expect(schemaToType({ anyOf: [{ type: "string" }, {}] })).toBe("unknown"); + expect(schemaToType({ allOf: [{ anyOf: [{ type: "string" }, { type: "number" }] }, { const: 1 }] })).toBe( + "(string | number) & 1", + ); + expect(schemaToType({ $ref: "#/defs/x" })).toBe("unknown"); + expect(schemaToType(true)).toBe("unknown"); + expect(schemaToType(false)).toBe("never"); + }); + + it("renders objects on one line with sorted properties", () => { + expect( + schemaToType({ + type: "object", + properties: { city: { type: "string" }, "max-lines": { type: "number" } }, + required: ["city"], + additionalProperties: false, + }), + ).toBe('{ city: string; "max-lines"?: number; }'); + expect(schemaToType({ type: "object", additionalProperties: { type: "number" } })).toBe( + "{ [key: string]: number; }", + ); + expect(schemaToType({ type: "object" })).toBe("{ [key: string]: unknown; }"); + expect(schemaToType({ type: "object", properties: {}, additionalProperties: false })).toBe("{}"); + }); + + it("puts property descriptions on comment lines", () => { + expect( + schemaToType({ + type: "object", + properties: { + weather: { + type: "array", + description: "look up weather for a given list of locations", + items: { type: "object", properties: { location: { type: "string" } }, required: ["location"] }, + }, + }, + required: ["weather"], + }), + ).toBe("{\n // look up weather for a given list of locations\n weather: Array<{ location: string; }>;\n}"); + expect( + schemaToType({ + type: "object", + properties: { + outer: { + type: "object", + description: "Outer", + properties: { inner: { type: "string", description: "Inner" } }, + }, + }, + }), + ).toBe("{\n // Outer\n outer?: {\n // Inner\n inner?: string;\n };\n}"); + }); + + it("resolves local references and stops at recursive ones", () => { + const schema = { + type: "object", + properties: { + item: { $ref: "#/$defs/Item" }, + legacy: { $ref: "#/definitions/Legacy" }, + remote: { $ref: "https://example.com/schema.json" }, + }, + required: ["item"], + $defs: { + Item: { + type: "object", + properties: { id: { type: "string" }, parent: { $ref: "#/$defs/Item" } }, + required: ["id"], + }, + }, + definitions: { Legacy: { enum: ["a", "b"] } }, + }; + expect(schemaToType(schema)).toBe( + '{ item: { id: string; parent?: unknown; }; legacy?: "a" | "b"; remote?: unknown; }', + ); + }); + + it("renders arrays and tuples", () => { + expect(schemaToType({ type: "array", items: { type: "string" } })).toBe("Array"); + expect(schemaToType({ type: "array", prefixItems: [{ type: "string" }, { type: "number" }] })).toBe( + "[string, number]", + ); + expect(schemaToType({ type: "array" })).toBe("unknown[]"); + }); + + it("renders types over the budget as unknown", () => { + const properties = Object.fromEntries(Array.from({ length: 50 }, (_, i) => [`field${i}`, { type: "string" }])); + const schema = { type: "object", properties }; + expect(schemaToType(schema, { maxChars: 100 })).toBe("unknown"); + expect(schemaToType(schema)).toContain("field49?: string;"); + }); +}); + +describe("tool declarations", () => { + it("renders signatures with normalized identifiers", () => { + expect( + renderToolSignature({ + name: "hidden-dynamic-tool", + inputSchema: { + type: "object", + properties: { city: { type: "string" } }, + required: ["city"], + additionalProperties: false, + }, + outputSchema: { type: "object", properties: { ok: { type: "boolean" } }, required: ["ok"] }, + }), + ).toBe("hidden_dynamic_tool(args: { city: string; }): Promise<{ ok: boolean; }>;"); + expect(renderToolSignature({ name: "free" })).toBe("free(args: unknown): Promise;"); + }); + + it("renders MCP CallToolResult output schemas as CallToolResult", () => { + const inputSchema = { type: "object", properties: {}, additionalProperties: false }; + expect( + renderToolSignature({ + name: "mcp__sample__search", + inputSchema, + outputSchema: mcpResultSchema({ + type: "object", + properties: { results: { type: "array", items: { $ref: "#/definitions/Result~1item~0v1" } } }, + required: ["results"], + additionalProperties: false, + definitions: { + "Result/item~v1": { + type: "object", + properties: { id: { type: "string" }, score: { type: "number" } }, + required: ["id", "score"], + additionalProperties: false, + }, + }, + }), + }), + ).toBe( + "mcp__sample__search(args: {}): Promise; }>>;", + ); + expect(renderToolSignature({ name: "plain", inputSchema, outputSchema: mcpResultSchema() })).toBe( + "plain(args: {}): Promise;", + ); + expect( + mcpStructuredContentSchema({ type: "object", properties: { content: { type: "array" } } }), + ).toBeUndefined(); + }); + + it("renders the per-tool sample", () => { + expect(renderToolSample({ name: "foo", description: "bar", inputSchema: { type: "string" } })).toBe( + "bar\n\ncodemode tool declaration:\n```ts\ndeclare const tools: { foo(args: string): Promise; };\n```", + ); + }); +}); + +describe("renderDeclarations", () => { + it("renders tools and globals", () => { + const text = renderDeclarations({ + tools: [ + { + name: "read", + description: "Read a file.\nSecond line.", + inputSchema: { type: "object", properties: { path: { type: "string" } }, required: ["path"] }, + outputSchema: { type: "string" }, + execute, + }, + { name: "remote-api", execute }, + ], + globals: [{ name: "attach", description: "Attach it.", inputSchema: { type: "string" }, execute }], + }); + expect(text).toBe( + [ + "declare const tools: {", + " /**", + " * Read a file.", + " * Second line.", + " */", + " read(args: { path: string; }): Promise;", + " remote_api(args: unknown): Promise;", + "};", + "", + "/** Attach it. */", + "declare function attach(args: string): Promise;", + ].join("\n"), + ); + }); + + it("renders namespaced globals and explicit signatures", () => { + const text = renderDeclarations({ + globals: [ + { + name: "models.list", + description: "List models.", + signature: "(type: string): Promise", + execute, + }, + { name: "models.get", inputSchema: { type: "string" }, execute }, + { name: "plain", signature: "(): void", execute }, + ], + }); + expect(text).toBe( + [ + "declare function plain(): void;", + "", + "declare const models: {", + " /** List models. */", + " list(type: string): Promise;", + " get(args: string): Promise;", + "};", + ].join("\n"), + ); + }); + + it("escapes comment terminators in descriptions", () => { + const text = renderDeclarations({ tools: [{ name: "x", description: "a */ b", execute }] }); + expect(text).toContain("/** a *\\/ b */"); + }); +}); diff --git a/packages/codemode/test/sandbox.test.ts b/packages/codemode/test/sandbox.test.ts new file mode 100644 index 000000000..ed6845b0f --- /dev/null +++ b/packages/codemode/test/sandbox.test.ts @@ -0,0 +1,600 @@ +import vm from "node:vm"; +import { afterEach, describe, expect, it } from "vitest"; +import { CodemodeSandbox, type CodemodeTool } from "../src/index.ts"; +import { PRELUDE_SOURCE } from "../src/runtime/prelude-source.ts"; + +const sandboxes: CodemodeSandbox[] = []; + +function createSandbox(tools: CodemodeTool[] = [], timeoutMs = 10_000): CodemodeSandbox { + const sandbox = new CodemodeSandbox({ tools, timeoutMs }); + sandboxes.push(sandbox); + return sandbox; +} + +afterEach(async () => { + await Promise.all(sandboxes.splice(0).map((sandbox) => sandbox.close())); +}); + +const echo: CodemodeTool = { name: "echo", execute: (args) => args }; + +describe("embedded sources", () => { + it("parse as JavaScript", () => { + expect(() => new vm.Script(PRELUDE_SOURCE, { filename: "prelude.js" })).not.toThrow(); + }); +}); + +describe("script execution", () => { + it("returns the script's return value after a JSON round trip", async () => { + const sandbox = createSandbox(); + expect(await sandbox.execute("return { a: 1, b: [true, 'x'] }")).toMatchObject({ + ok: true, + value: { a: 1, b: [true, "x"] }, + output: [], + calls: [], + }); + expect(await sandbox.execute("return 'plain'")).toMatchObject({ ok: true, value: "plain" }); + expect(await sandbox.execute("")).toMatchObject({ ok: true, value: undefined }); + }); + + it("supports top-level await", async () => { + const sandbox = createSandbox(); + const result = await sandbox.execute("const x = await Promise.resolve(41); return x + 1"); + expect(result).toMatchObject({ ok: true, value: 42 }); + }); + + it("collects text(), image(), and console output in order", async () => { + const sandbox = createSandbox(); + const result = await sandbox.execute(` + console.log("hello", 1, { a: 1 }); + text({ json: true }); + text(undefined); + text(7); + image("data:image/png;base64,AAAA"); + image({ image_url: "data:image/jpeg;base64,BBBB" }); + image({ type: "image", data: "CCCC", mimeType: "image/gif" }); + console.error(new Error("bad")); + return null; + `); + expect(result.ok).toBe(true); + expect(result.output.slice(0, -1)).toEqual([ + { type: "text", text: 'hello 1 {"a":1}' }, + { type: "text", text: '{"json":true}' }, + { type: "text", text: "undefined" }, + { type: "text", text: "7" }, + { type: "image", data: "AAAA", mimeType: "image/png" }, + { type: "image", data: "BBBB", mimeType: "image/jpeg" }, + { type: "image", data: "CCCC", mimeType: "image/gif" }, + ]); + expect(result.output.at(-1)).toMatchObject({ type: "text", text: expect.stringMatching(/^Error: bad/) }); + }); + + it("rejects invalid text() and image() arguments", async () => { + const sandbox = createSandbox(); + const result = await sandbox.execute(` + const errors = []; + const circular = {}; + circular.self = circular; + for (const run of [ + () => text(circular), + () => image(""), + () => image("https://example.com/a.png"), + () => image("data:image/png,raw"), + () => image({ type: "text", text: "x" }), + () => image({ type: "image", data: "" }), + () => image(42), + ]) { + try { run(); errors.push("no error"); } catch (error) { errors.push(error.name + ": " + error.message); } + } + return errors; + `); + expect(result.ok).toBe(true); + if (!result.ok) return; + expect(result.output).toEqual([]); + const errors = result.value as string[]; + expect(errors[0]).toMatch(/^TypeError: .*circular/); + expect(errors.slice(1)).toEqual([ + "TypeError: image expects a non-empty image URL string, an object with image_url, or a raw MCP image block", + "TypeError: remote image URLs are not supported in tool outputs. Pass a base64 data URI instead", + "TypeError: invalid image output. Pass a base64 data URI instead", + 'TypeError: image only accepts MCP image blocks, got "text"', + "TypeError: image expected MCP image data", + "TypeError: image expects a non-empty image URL string, an object with image_url, or a raw MCP image block", + ]); + }); + + it("ends the script successfully on exit(), keeping output and store writes", async () => { + const sandbox = createSandbox([echo]); + const result = await sandbox.execute(` + text("before"); + store("k", 1); + await tools.echo(1); + try { exit(); } catch {} + text("after"); + return "unreachable"; + `); + expect(result).toMatchObject({ + ok: true, + value: undefined, + output: [{ type: "text", text: "before" }], + storeWrites: { set: { k: 1 }, delete: [] }, + }); + }); + + it("keeps output produced before a failure", async () => { + const sandbox = createSandbox(); + const result = await sandbox.execute('text("partial");\nthrow new Error("boom")'); + expect(result).toMatchObject({ ok: false, output: [{ type: "text", text: "partial" }] }); + }); + + it("reports syntax errors with the script's line number", async () => { + const sandbox = createSandbox(); + const result = await sandbox.execute("const a = 1;\nconst b = ;\nreturn a"); + expect(result.ok).toBe(false); + if (result.ok) return; + expect(result.error.kind).toBe("script"); + expect(result.error.name).toBe("SyntaxError"); + expect(result.error.stack).toMatch(/codemode\.js:2/); + }); + + it("reports thrown errors with the script's line number", async () => { + const sandbox = createSandbox(); + const result = await sandbox.execute("const a = 1;\nthrow new TypeError('boom ' + a)"); + expect(result.ok).toBe(false); + if (result.ok) return; + expect(result.error).toMatchObject({ kind: "script", name: "TypeError", message: "boom 1" }); + expect(result.error.stack).toMatch(/codemode\.js:2/); + }); + + it("formats stacks like V8 without prelude frames", async () => { + const sandbox = createSandbox(); + const result = await sandbox.execute("console.log(new Error('inner'));\nthrow new RangeError('outer')"); + expect(result.ok).toBe(false); + if (result.ok) return; + expect(result.error.stack).toMatch(/^RangeError: outer\n {4}at .*codemode\.js:2/); + expect(result.error.stack).not.toContain("codemode-prelude.js"); + expect(result.output[0]).toMatchObject({ + type: "text", + text: expect.stringMatching(/^Error: inner\n {4}at .*codemode\.js:1/), + }); + expect(JSON.stringify(result.output)).not.toContain("codemode-prelude.js"); + }); + + it("reports non-Error throws", async () => { + const sandbox = createSandbox(); + const result = await sandbox.execute("throw { code: 7 }"); + expect(result).toMatchObject({ ok: false, error: { kind: "script", message: '{"code":7}' } }); + }); + + it("reports a non-serializable return value as a script error", async () => { + const sandbox = createSandbox(); + const result = await sandbox.execute("return 10n"); + expect(result).toMatchObject({ ok: false, error: { kind: "script", name: "TypeError" } }); + }); +}); + +describe("tools", () => { + it("exposes tools as async functions and records calls", async () => { + const seen: unknown[] = []; + const sandbox = createSandbox([ + { + name: "add", + execute: (args) => { + seen.push(args); + const { a, b } = args as { a: number; b: number }; + return { sum: a + b }; + }, + }, + ]); + const result = await sandbox.execute(` + const first = await tools.add({ a: 1, b: 2 }); + const second = await tools.add({ a: first.sum, b: 10 }); + return second.sum; + `); + expect(result).toMatchObject({ ok: true, value: 13 }); + expect(seen).toEqual([ + { a: 1, b: 2 }, + { a: 3, b: 10 }, + ]); + expect(result.calls.map((call) => [call.name, call.status])).toEqual([ + ["add", "ok"], + ["add", "ok"], + ]); + expect(result.calls.every((call) => call.durationMs >= 0)).toBe(true); + }); + + it("runs concurrent calls and lists tool names", async () => { + const sandbox = createSandbox([ + echo, + { name: "delay", execute: (args) => new Promise((resolve) => setTimeout(() => resolve(args), 20)) }, + ]); + const result = await sandbox.execute(` + const [a, b, c] = await Promise.all([tools.delay(1), tools.delay(2), tools.echo(3)]); + return { values: [a, b, c], names: Object.keys(tools) }; + `); + expect(result).toMatchObject({ ok: true, value: { values: [1, 2, 3], names: ["echo", "delay"] } }); + }); + + it("exposes tools under normalized identifiers and lists them in ALL_TOOLS", async () => { + const sandbox = createSandbox([ + { name: "my-tool", description: "Dashes", execute: () => "dash" }, + { name: "my_tool", description: "Shadowed", execute: () => "underscore" }, + { name: "mcp__docs__search", execute: () => "mcp" }, + ]); + const result = await sandbox.execute(` + try { ALL_TOOLS.push({}); } catch {} + return { + all: ALL_TOOLS, + calls: [await tools.my_tool(), await tools["my-tool"](), await tools.mcp__docs__search()], + }; + `); + expect(result).toMatchObject({ + ok: true, + value: { + all: [ + { name: "my_tool", description: "Dashes" }, + { name: "mcp__docs__search", description: "" }, + ], + calls: ["dash", "dash", "mcp"], + }, + }); + }); + + it("passes undefined arguments and results through", async () => { + const sandbox = createSandbox([{ name: "noop", execute: (args) => args }]); + const result = await sandbox.execute("return [await tools.noop(), await tools.noop(null)]"); + expect(result).toMatchObject({ ok: true, value: [null, null] }); + }); + + it("turns tool errors into catchable Errors in the script", async () => { + const sandbox = createSandbox([ + { + name: "fail", + execute: () => { + throw new Error("tool exploded"); + }, + }, + ]); + const result = await sandbox.execute(` + try { + await tools.fail(); + return "no error"; + } catch (error) { + return { isError: error instanceof Error, message: error.message }; + } + `); + expect(result).toMatchObject({ ok: true, value: { isError: true, message: "tool exploded" } }); + expect(result.calls).toMatchObject([{ name: "fail", status: "error" }]); + }); + + it("rejects calls to unknown tools", async () => { + const sandbox = createSandbox(); + const result = await sandbox.execute("return await tools.missing()"); + expect(result).toMatchObject({ ok: false, error: { kind: "script", name: "TypeError" } }); + }); + + it("aborts unawaited calls when the script returns", async () => { + let aborted = false; + const sandbox = createSandbox([ + { + name: "slow", + execute: (_args, { signal }) => + new Promise((_resolve, reject) => { + signal.addEventListener("abort", () => { + aborted = true; + reject(new Error("aborted")); + }); + }), + }, + ]); + const result = await sandbox.execute("tools.slow(); return 'early'"); + expect(result).toMatchObject({ ok: true, value: "early" }); + expect(result.calls).toMatchObject([{ name: "slow", status: "cancelled" }]); + expect(aborted).toBe(true); + }); + + it("supports register and unregister between executions", async () => { + const sandbox = createSandbox(); + sandbox.registerTool(echo); + expect(() => sandbox.registerTool(echo)).toThrow(/already registered/); + expect(sandbox.tools.map((tool) => tool.name)).toEqual(["echo"]); + expect(await sandbox.execute("return await tools.echo('a')")).toMatchObject({ ok: true, value: "a" }); + expect(sandbox.unregisterTool("echo")).toBe(true); + expect(await sandbox.execute("return typeof tools.echo")).toMatchObject({ ok: true, value: "undefined" }); + }); +}); + +describe("store and load", () => { + it("reads the snapshot and reports writes", async () => { + const sandbox = createSandbox(); + const result = await sandbox.execute( + ` + const seen = load("counter"); + store("counter", seen + 1); + store("list", [1, { a: null }]); + store("old", undefined); + return [seen, load("counter"), load("missing"), load("old")]; + `, + { store: { counter: 41, old: "x" } }, + ); + expect(result).toMatchObject({ + ok: true, + // undefined array elements become null in the JSON round trip of the return value. + value: [41, 42, null, null], + storeWrites: { set: { counter: 42, list: [1, { a: null }] }, delete: ["old"] }, + }); + }); + + it("returns copies, so mutating a loaded value does not change the store", async () => { + const sandbox = createSandbox(); + const result = await sandbox.execute( + `const value = load("obj"); value.a = 2; const kept = { b: 1 }; store("kept", kept); kept.b = 2; + return [load("obj").a, load("kept").b];`, + { store: { obj: { a: 1 } } }, + ); + expect(result).toMatchObject({ ok: true, value: [1, 1], storeWrites: { set: { kept: { b: 1 } } } }); + }); + + it("rejects invalid keys, values, and oversized writes inside the script", async () => { + const sandbox = createSandbox(); + const result = await sandbox.execute(` + const attempt = (fn) => { try { fn(); return "ok"; } catch (error) { return error.name; } }; + return [ + attempt(() => store(1, "x")), + attempt(() => load({})), + attempt(() => store("fn", () => 1)), + attempt(() => store("big", "x".repeat(300 * 1024))), + attempt(() => { for (let i = 0; i < 8; i++) store("k" + i, "x".repeat(200 * 1024)); }), + ]; + `); + expect(result).toMatchObject({ + ok: true, + value: ["TypeError", "TypeError", "TypeError", "RangeError", "RangeError"], + }); + }); + + it("reserves the store and load names", () => { + const execute = () => undefined; + expect(() => new CodemodeSandbox({ globals: [{ name: "store", execute }] })).toThrow(/Invalid global/); + expect(() => new CodemodeSandbox({ globals: [{ name: "load", execute }] })).toThrow(/Invalid global/); + }); +}); + +describe("globals", () => { + it("exposes globals as top-level functions without recording them as calls", async () => { + const seen: unknown[] = []; + const sandbox = new CodemodeSandbox({ + tools: [echo], + globals: [{ name: "attach", execute: (args) => void seen.push(args) }], + }); + sandboxes.push(sandbox); + const result = await sandbox.execute(` + await attach({ ref: 1 }); + attach("not awaited"); + return [typeof attach, typeof globalThis.attach, await tools.echo(2)]; + `); + expect(result).toMatchObject({ ok: true, value: ["function", "function", 2] }); + expect(result.calls.map((call) => call.name)).toEqual(["echo"]); + // Messages are handled in order, so an unawaited global still runs before the script settles. + expect(seen).toEqual([{ ref: 1 }, "not awaited"]); + }); + + it("groups namespaced globals and spreads arguments on request", async () => { + const seen: unknown[] = []; + const sandbox = new CodemodeSandbox({ + globals: [ + { name: "models.list", spread: true, execute: (args) => void seen.push(args) }, + { name: "models.first", execute: (args) => args }, + ], + }); + sandboxes.push(sandbox); + const result = await sandbox.execute(` + await models.list("classifier", undefined, 3); + await models.list(); + try { models.extra = 1; } catch {} + return [Object.keys(models), await models.first("a", "ignored"), typeof models.extra]; + `); + expect(result).toMatchObject({ ok: true, value: [["list", "first"], "a", "undefined"] }); + // undefined array elements become null in the JSON round trip. + expect(seen).toEqual([["classifier", null, 3], []]); + }); + + it("rejects invalid and reserved global names", () => { + const execute = () => undefined; + for (const name of ["a.b.c", "a.", ".a", "tools.x", "store.x", "a.not-valid"]) { + expect(() => new CodemodeSandbox({ globals: [{ name, execute }] }), name).toThrow(/Invalid global/); + } + expect( + () => + new CodemodeSandbox({ + globals: [ + { name: "models", execute }, + { name: "models.list", execute }, + ], + }), + ).toThrow(/conflicts with the namespace/); + expect(() => new CodemodeSandbox({ globals: [{ name: "not-valid", execute }] })).toThrow(/Invalid global/); + expect(() => new CodemodeSandbox({ globals: [{ name: "tools", execute }] })).toThrow(/Invalid global/); + expect(() => new CodemodeSandbox({ globals: [{ name: "console", execute }] })).toThrow(/Invalid global/); + }); +}); + +describe("limits and lifetime", () => { + it("terminates a synchronous infinite loop on timeout", async () => { + const sandbox = createSandbox(); + const started = performance.now(); + const result = await sandbox.execute("while (true) {}", { timeoutMs: 200 }); + expect(result).toMatchObject({ ok: false, error: { kind: "timeout" } }); + expect(performance.now() - started).toBeLessThan(5_000); + }); + + it("runs without a deadline when timeoutMs is Infinity", async () => { + const sandbox = createSandbox([ + { name: "wait", execute: () => new Promise((r) => setTimeout(() => r("late"), 50)) }, + ]); + const result = await sandbox.execute("return await tools.wait()", { timeoutMs: Number.POSITIVE_INFINITY }); + expect(result).toMatchObject({ ok: true, value: "late" }); + }); + + it("fails a script that waits on a promise nothing can settle", async () => { + const sandbox = createSandbox([{ name: "echo", execute: (args) => args }]); + const result = await sandbox.execute("await tools.echo(1); await new Promise(() => {}); return 'never'", { + timeoutMs: Number.POSITIVE_INFINITY, + }); + expect(result).toMatchObject({ + ok: false, + error: { kind: "script", message: expect.stringContaining("can never settle") }, + }); + // Returning while a call is still pending is not a stall. + const returned = await sandbox.execute("tools.echo(2); return 'early'", { timeoutMs: Number.POSITIVE_INFINITY }); + expect(returned).toMatchObject({ ok: true, value: "early" }); + }); + + it("terminates a microtask-spinning loop on timeout", async () => { + const sandbox = createSandbox(); + const result = await sandbox.execute("while (true) await null", { timeoutMs: 200 }); + expect(result).toMatchObject({ ok: false, error: { kind: "timeout" } }); + }); + + it("aborts via signal and cancels in-flight calls", async () => { + const controller = new AbortController(); + let toolSignal: AbortSignal | undefined; + // Resolves on each call, so the test waits for the worker instead of a fixed delay. + let called: () => void = () => {}; + const nextCall = () => + new Promise((resolve) => { + called = resolve; + }); + const sandbox = createSandbox([ + { + name: "hang", + execute: (_args, { signal }) => { + toolSignal = signal; + called(); + return new Promise(() => {}); + }, + }, + ]); + const firstCall = nextCall(); + const promise = sandbox.execute("await tools.hang(); return 'never'"); + await firstCall; + controller.abort(new Error("user cancelled")); + expect(toolSignal?.aborted).toBe(false); + const result = await sandbox.execute("await tools.hang()", { signal: controller.signal }); + expect(result).toMatchObject({ ok: false, error: { kind: "aborted", message: "user cancelled" } }); + + const second = new AbortController(); + const secondCall = nextCall(); + const pending = sandbox.execute("await tools.hang(); return 'never'", { signal: second.signal }); + await secondCall; + second.abort(); + const aborted = await pending; + expect(aborted).toMatchObject({ ok: false, error: { kind: "aborted" } }); + expect(aborted.calls).toMatchObject([{ name: "hang", status: "cancelled" }]); + expect(toolSignal?.aborted).toBe(true); + + await sandbox.close(); + expect(await promise).toMatchObject({ ok: false, error: { kind: "aborted", message: "Sandbox closed" } }); + }); + + it("rejects execute after close", async () => { + const sandbox = createSandbox(); + await sandbox.close(); + await expect(sandbox.execute("return 1")).rejects.toThrow(/closed/); + }); + + it("runs executions in parallel without sharing state", async () => { + const sandbox = createSandbox(); + const results = await Promise.all([ + sandbox.execute("globalThis.shared = 'a'; await null; return globalThis.shared"), + sandbox.execute("globalThis.shared = 'b'; await null; return globalThis.shared"), + sandbox.execute("return typeof globalThis.shared"), + ]); + expect(results.map((result) => (result.ok ? result.value : result.error))).toEqual(["a", "b", "undefined"]); + }); + + it("turns deep recursion into a catchable RangeError", async () => { + const sandbox = createSandbox(); + const result = await sandbox.execute(` + let depth = 0; + function dive() { depth++; dive(); } + try { dive(); } catch (error) { return [error.name, depth > 1000]; } + `); + expect(result).toMatchObject({ ok: true, value: ["RangeError", true] }); + }); + + it("reports a missing worker file as a sandbox error", async () => { + const sandbox = new CodemodeSandbox({ workerUrl: new URL("./does-not-exist.js", import.meta.url) }); + sandboxes.push(sandbox); + expect(await sandbox.execute("return 1")).toMatchObject({ ok: false, error: { kind: "sandbox" } }); + }); + + it("reports a failing wasm module as a sandbox error", async () => { + const sandbox = new CodemodeSandbox({ wasm: Promise.reject(new Error("no wasm")) }); + sandboxes.push(sandbox); + expect(await sandbox.execute("return 1")).toMatchObject({ + ok: false, + error: { kind: "sandbox", message: "Failed to load QuickJS: no wasm" }, + }); + }); +}); + +describe("escape hatches", () => { + it("has no host globals", async () => { + const sandbox = createSandbox(); + const result = await sandbox.execute(` + return [ + typeof process, typeof require, typeof module, typeof setTimeout, typeof fetch, + typeof WebAssembly, typeof std, typeof os, typeof globalThis.constructor, + ] + `); + expect(result).toMatchObject({ + ok: true, + value: [ + "undefined", + "undefined", + "undefined", + "undefined", + "undefined", + "undefined", + "undefined", + "undefined", + "function", + ], + }); + }); + + it("keeps eval and Function inside the VM", async () => { + // Code generation is allowed: it can only produce more code in the same wasm instance. + const sandbox = createSandbox([echo]); + const result = await sandbox.execute(` + return [ + eval("typeof process"), + new Function("return typeof process")(), + tools.echo.constructor("return typeof require")(), + (async () => {}).constructor("return typeof setTimeout")() instanceof Promise, + ]; + `); + expect(result).toMatchObject({ ok: true, value: ["undefined", "undefined", "undefined", true] }); + }); + + it("rejects dynamic import", async () => { + const sandbox = createSandbox(); + const result = await sandbox.execute(` + try { await import("node:fs"); return "imported"; } catch (error) { return error.constructor.name; } + `); + expect(result.ok).toBe(true); + if (!result.ok) return; + expect(result.value).not.toBe("imported"); + }); + + it("keeps tools and console frozen", async () => { + const sandbox = createSandbox([echo]); + const result = await sandbox.execute(` + try { tools.echo = () => 'nope'; } catch {} + try { tools.extra = () => 'nope'; } catch {} + try { globalThis.tools = null; } catch {} + return [typeof tools.extra, await tools.echo('still')]; + `); + expect(result).toMatchObject({ ok: true, value: ["undefined", "still"] }); + }); +}); diff --git a/packages/codemode/test/source.test.ts b/packages/codemode/test/source.test.ts new file mode 100644 index 000000000..7db6009fe --- /dev/null +++ b/packages/codemode/test/source.test.ts @@ -0,0 +1,63 @@ +import { describe, expect, it } from "vitest"; +import { CodemodeSourceError, parseCodemodeSource } from "../src/source.ts"; + +describe("parseCodemodeSource", () => { + it("returns plain code unchanged", () => { + expect(parseCodemodeSource("text('hi')")).toEqual({ code: "text('hi')", options: {} }); + expect(parseCodemodeSource("// just a comment\nreturn 1")).toEqual({ + code: "// just a comment\nreturn 1", + options: {}, + }); + }); + + it("parses the options line and keeps line numbers", () => { + expect(parseCodemodeSource('// @options: {"timeout_ms": 10}\nconst a = 1;\ntext(a)')).toEqual({ + code: "\nconst a = 1;\ntext(a)", + options: { timeoutMs: 10 }, + }); + expect(parseCodemodeSource(' // @options:{"max_output_tokens":0,"timeout_ms":1500}\r\ntext(1)').options).toEqual( + { + maxOutputTokens: 0, + timeoutMs: 1500, + }, + ); + expect(parseCodemodeSource("// @options: {}\ntext(1)")).toEqual({ code: "\ntext(1)", options: {} }); + }); + + it("only treats the first line as an options line", () => { + const input = 'text(1)\n// @options: {"timeout_ms": 1}'; + expect(parseCodemodeSource(input)).toEqual({ code: input, options: {} }); + expect(parseCodemodeSource("// @optionsx {}\ntext(1)").options).toEqual({}); + }); + + it("rejects empty input and invalid options", () => { + const cases: [string, string | RegExp][] = [ + ["", /Expected JavaScript source text \(non-empty\)/], + [" \n", /Expected JavaScript source text \(non-empty\)/], + ["// @options:\ntext(1)", /@options must be a JSON object with supported fields/], + ["// @options: {timeout_ms: 1}\ntext(1)", /@options must be valid JSON with supported fields/], + ["// @options: [1]\ntext(1)", /@options must be a JSON object with supported fields/], + [ + '// @options: {"yield": 1}\ntext(1)', + "@options only supports `max_output_tokens` and `timeout_ms`; got `yield`", + ], + [ + '// @options: {"max_output_tokens": 1.5}\ntext(1)', + "@options field `max_output_tokens` must be a non-negative safe integer", + ], + ['// @options: {"timeout_ms": 0}\ntext(1)', /@options field `timeout_ms` must be a positive integer/], + [ + '// @options: {"timeout_ms": 1}', + "The @options line must be followed by JavaScript source on subsequent lines", + ], + [ + '// @options: {"timeout_ms": 1}\n \n', + "The @options line must be followed by JavaScript source on subsequent lines", + ], + ]; + for (const [input, message] of cases) { + expect(() => parseCodemodeSource(input), input).toThrow(CodemodeSourceError); + expect(() => parseCodemodeSource(input), input).toThrow(message); + } + }); +}); diff --git a/packages/codemode/tsconfig.build.json b/packages/codemode/tsconfig.build.json new file mode 100644 index 000000000..695dd9ad2 --- /dev/null +++ b/packages/codemode/tsconfig.build.json @@ -0,0 +1,9 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "outDir": "./dist", + "rootDir": "./src" + }, + "include": ["src/**/*.ts"], + "exclude": ["node_modules", "dist", "**/*.d.ts", "src/**/*.d.ts"] +} diff --git a/packages/codemode/vitest.config.ts b/packages/codemode/vitest.config.ts new file mode 100644 index 000000000..18733b4d4 --- /dev/null +++ b/packages/codemode/vitest.config.ts @@ -0,0 +1,11 @@ +import { defineConfig } from "vitest/config"; + +export default defineConfig({ + test: { + globals: true, + environment: "node", + reporters: process.env.GITHUB_ACTIONS ? ["dot", "github-actions"] : ["dot"], + }, + resolve: { conditions: ["source"] }, + ssr: { resolve: { conditions: ["source"] } }, +}); diff --git a/packages/coding-agent/docs/cli.md b/packages/coding-agent/docs/cli.md index a0193b2b5..b45ae962a 100644 --- a/packages/coding-agent/docs/cli.md +++ b/packages/coding-agent/docs/cli.md @@ -13,6 +13,7 @@ pi update [target] [options] pi list pi config [options] pi auth [options] +pi mcp [options] ``` @@ -137,6 +138,27 @@ Default enabled tools are `read`, `bash`, `edit`, and `write`, unless `defaultTo | `find` | Find paths using glob patterns | | `ls` | List directory contents | +Built-in extensions add two more tools. They are off by default; name them in `--tools` or `defaultTools` to enable them. + +| Built-in extension | Purpose | +|---|---| +| `codemode` | Run JavaScript that calls the other tools, for example in parallel with `Promise.allSettled`; only the script's output reaches the model | +| `tool_search` | Search tools that are not declared to the model (`codemode` and `deferred` exposure, such as MCP tools) and declare the matches for the next call | + +Codemode scripts run in a QuickJS sandbox that can only reach the other tools, through `tools.(args)`; `ALL_TOOLS` lists them. Output comes from `text(value)`, `image(dataUrlOrImageContent)`, `console.*`, and a top-level `return value`; `exit()` ends the script early. The result starts with `Script completed` or `Script failed`, the wall time, and the output; a failed script keeps its partial output, followed by `Script error:` and the error. + +A script may start with an options line such as `// @options: {"max_output_tokens": 2000, "timeout_ms": 60000}`. `max_output_tokens` (default 10000) limits the output: longer output keeps its start and end, and the full text is written to a temp file whose path is included in the result. `timeout_ms` is a hard deadline, unset by default. + +While `codemode` is active, `codemode.mode` in [settings](settings.md#tools) decides how the other tools are presented. With `on` (default) declared tools keep being declared and their descriptions show how to call them from scripts. With `only` they are hidden from the model and listed in the `codemode` description instead, so the model calls them through scripts. + +The `codemode` description lists the callable tools with their TypeScript declarations, grouped by namespace (for example one MCP server). Declarations share a budget of 3000 estimated tokens (`codemode.inlineBudget` in [settings](settings.md#tools)); every namespace is still listed with its tool count, and the description says whether the list is complete. Scripts find the rest with `await searchTools(query, { limit, namespace })`, which ranks tools with BM25, and `await describeTool(name)`, or by filtering `ALL_TOOLS`. + +`tool_search` is off by default; enable it with `--tools` or `defaultTools`. It uses the same ranking over tools that are not declared yet and declares the matches for the next model call. Loaded tools are recorded in the session like other tool changes, so they stay declared on that branch. + +Tools with an output schema resolve to structured values: `bash` to `{ output, exit_code, wall_time_seconds }`, also for non-zero exit codes, and MCP tools to their `CallToolResult`. Other tools resolve to their text output. + +`store(key, value)` and `load(key)` keep JSON values across `codemode` calls: each successful script that stores values appends a `codemode-store` custom entry to the session, so resumed sessions keep the values and each branch sees only the values written on its path. Scripts can also use `models`: `getModelsOfType`, `getAvailableOfType`, and `getModelOfType` list the model catalog, and `classify(model, context)` runs a classifier model with the session's credentials, at most four at a time per script. + ## Resources @@ -266,3 +288,20 @@ Authentication commands require `--provider ` or `--model `. Se | `--min-expiry ` | `print-bearer-token` | Require remaining token lifetime using `ms`, `s`, `m`, or `h`, such as `30m` | Credential-printing commands write secrets to stdout. + +## MCP commands + +These commands work outside a session, so agents can run them through `bash`. See [MCP Servers](mcp.md). + +| Command | Description | +|---|---| +| `pi mcp add [options] -- [args...]` | Add or replace a stdio server in `mcp.json`; `--env KEY=VALUE` (repeatable) and `--cwd ` set its environment and working directory. Arguments after the command are passed to it | +| `pi mcp add [options] --url ` | Add or replace a streamable HTTP server; `--header KEY=VALUE` (repeatable), `--bearer-token-env-var ` (sends `Authorization: Bearer ${NAME}`), `--oauth-client-id`, `--oauth-client-secret`, and `--oauth-callback-port` configure authentication | +| `pi mcp remove ` | Remove a server from `mcp.json`; stored OAuth credentials are kept | +| `pi mcp list [--json]` | Connect to every enabled server and print its state, tools, and errors; exit with `1` when a config entry is invalid or an enabled server is not connected | +| `pi mcp login [--timeout ]` | Sign in to an OAuth server: open the authorization page and wait for the browser (default 300 seconds); a terminal also accepts the pasted redirect URL | +| `pi mcp logout ` | Delete the stored OAuth credentials of a server | + +`add` and `remove` change `~/.pi/agent/mcp.json`, or `.pi/mcp.json` in the current directory with `--local` (`-l`). `add` also takes `--exposure ` (see [Exposure](mcp.md#exposure)) and does not connect; run `pi mcp list` to check the server. + +Project `.pi/mcp.json` files are only read for projects that are already trusted. diff --git a/packages/coding-agent/docs/configuration.md b/packages/coding-agent/docs/configuration.md index 40fdb0ef7..b16fd001d 100644 --- a/packages/coding-agent/docs/configuration.md +++ b/packages/coding-agent/docs/configuration.md @@ -12,6 +12,7 @@ The agent directory is shown as `` below. Set its location with the ` |---|---| | `/settings.json` | User-level [settings](settings.md), including preferences, defaults, resource paths, and Pi package declarations. | | `/keybindings.json` | Custom terminal UI and application [keybindings](keybindings.md). | +| `/mcp.json` | [MCP servers](mcp.md) available in every project. | | `/models.json` | [Compatible endpoints, models, and model overrides](models.md#configure-a-compatible-endpoint). | | `/auth.json` | Saved API keys and OAuth credentials. | | `/AGENTS.override.md`, `AGENTS.md`, `AGENTS.MD`, `CLAUDE.md`, or `CLAUDE.MD` | User instructions applied across working directories. | @@ -27,6 +28,7 @@ The agent directory is shown as `` below. Set its location with the ` | Path | Responsibility | |---|---| | `.pi/settings.json` | Project-level [settings](settings.md), resource paths, and Pi package declarations. | +| `.pi/mcp.json` | Project [MCP servers](mcp.md). | | `.pi/SYSTEM.md` | Replaces the system prompt for the project. | | `.pi/APPEND_SYSTEM.md` | Adds project-specific instructions to the system prompt. | | `.pi/extensions/` | Project extensions. | diff --git a/packages/coding-agent/docs/docs.json b/packages/coding-agent/docs/docs.json index 85f91ebc6..ad61486db 100644 --- a/packages/coding-agent/docs/docs.json +++ b/packages/coding-agent/docs/docs.json @@ -91,6 +91,10 @@ { "title": "Use Pi Packages", "path": "packages.md" + }, + { + "title": "Connect MCP Servers", + "path": "mcp.md" } ] }, diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index e611af1ff..d4a53d392 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -80,6 +80,7 @@ Automatic retries, recovery, compaction, or queued work can continue afterward. | Persist non-context session data | `pi.appendEntry()` | | Change active tools, model, or thinking level | Session control methods on `pi` | | Add a model provider | `pi.registerProvider()` | +| Add an MCP server | `pi.registerMcpServer()` | | Route each request to a model | [`pi.registerVirtualModel()`](virtual-models.md) | | Add terminal rendering | Renderer registration and `ctx.ui` | | Communicate with another extension | `pi.events` | @@ -142,14 +143,61 @@ Use sequential execution when tools share mutable in-memory state. File-mutating tools should wrap the complete read-modify-write operation with `withFileMutationQueue()`. Truncate large model-facing results and tell the model where to read the complete output. +Declare `outputSchema` and return a matching `structuredContent` when the result is data. The model still receives `content`; programmatic callers such as codemode scripts receive `structuredContent` instead of the text. Tools without `outputSchema` are passed to scripts as their text content. To report a failure that still carries data, return the result with `isError: true` instead of throwing: the model sees an error, and scripts still receive `structuredContent`. + +A tool can run other tools with `ctx.executeTool(name, args, { signal, onUpdate })`. Nested calls go through argument validation and the `tool_call` and `tool_result` handlers like model-issued calls, and emit `tool_execution_start`, `tool_execution_update`, and `tool_execution_end`; all of these events carry `parentToolCallId`, and their `toolCallId` is assigned by pi as `/`. These ids do not appear as tool calls or tool results in the transcript. Nested calls do not add transcript entries: their results only reach the calling tool, which reports them itself, for example through `onUpdate` and `details`. The session keeps a bounded record of them (name, arguments, status, duration, error; never results) as `nestedCalls` on the calling tool's result message. It is used for compaction file lists and shown in HTML exports. Arguments over 8 KiB per call or 32 KiB per tool result are omitted, at most 256 calls are kept, and `complete: false` marks a record that lost anything. `ctx.tools` lists the tools `ctx.executeTool()` can call. `tool_result` handlers that redact `content` should also replace `structuredContent`; replacing only `content` drops it. + See [`hello.ts`](../examples/extensions/hello.ts), [`todo.ts`](../examples/extensions/todo.ts), [`dynamic-tools.ts`](../examples/extensions/dynamic-tools.ts), and [`truncated-tool.ts`](../examples/extensions/truncated-tool.ts). +### Tool exposure + +`exposure` controls how the model reaches a tool. "Callable" means callable from other tools through `ctx.executeTool()` (`ctx.tools`), as the `codemode` tool's scripts do: + +- `direct` (default): declared to the model while active, and callable while active. +- `model-only`: declared to the model while active, never callable. Use it for tools that orchestrate other tools or ask the user. +- `codemode`: callable whenever registered, and listed by the `codemode` tool. Not declared to the model unless activated explicitly. +- `deferred`: like `codemode`, but codemode tools do not list it; `tool_search` can find and activate it. +- `hidden`: registered but unreachable. Re-register a tool with `exposure: "hidden"` to withdraw it, since tools cannot be unregistered. + +`namespace: { name, description }` groups related tools, as MCP servers do. Codemode tools list a namespace under one heading. + +Registering a `direct` or `model-only` tool activates it; the other exposures are not activated on registration. The active set (`pi.getActiveTools()`, `pi.setActiveTools()`) is the set of tools declared to the model. `pi.getAllTools()` reports each tool's `exposure`, `namespace`, and `annotations`. + +`annotations` are hints about what a tool does, with the meaning of MCP tool annotations: `readOnlyHint`, `destructiveHint`, `idempotentHint`, and `openWorldHint`. MCP tools carry the hints their server declares. Missing hints take the MCP defaults: a tool is not read-only, and may be destructive and reach an open world. The hints are not verified, but a permission extension can use them to decide which calls to confirm. This confirms the calls Codex asks approval for: + +```typescript +pi.on("tool_call", async (event, ctx) => { + const hints = pi.getAllTools().find((tool) => tool.name === event.toolName)?.annotations; + const needsApproval = + hints?.destructiveHint === true || + (!hints?.readOnlyHint && ((hints?.destructiveHint ?? true) || (hints?.openWorldHint ?? true))); + if (needsApproval && !(await ctx.ui.confirm("Allow tool call?", event.toolName))) { + return { block: true, reason: `${event.toolName} was not approved` }; + } +}); +``` + +A tool that orchestrates other tools can adjust what the model sees while it is active with `prepareLoadout(loadout)`. It runs whenever the active tools change and receives the declared tools, the callable tools, and every registered tool with its exposure and namespace. It returns replacement `descriptions` for declared tools (including its own) and `hiddenDeclarations`: active tools whose declarations requests leave out while they stay active and callable. `codemode` and `tool_search` use only this hook, `exposure`, and `ctx.executeTool()`, so another tool can implement the same behavior under a different name. + ### Activate tools dynamically Register every tool first, keep optional tools inactive, and use `pi.setActiveTools()` from a loader tool to select the desired active tools. Names must already be registered; unknown names are ignored. Pi records the initial prompt and tool set in the transcript's first system message, then appends tool and prompt changes before the next model request. Providers that cannot represent the transition receive a complete transcript checkpoint, which can invalidate the cached prefix. +### MCP servers + +`pi.registerMcpServer(name, config)` adds an MCP server for the current session. `config` has the shape of an `mcpServers` entry in [`mcp.json`](mcp.md): `command`, `args`, `env`, and `cwd` for stdio servers, `url`, `headers`, and `oauth` for HTTP servers, plus `exposure`, `toolExposure`, `enabled`, and `timeout`. + +```typescript +pi.registerMcpServer("jira", { url: "https://mcp.example.com/jira", exposure: "codemode" }); +pi.unregisterMcpServer("jira"); +``` + +Servers registered while the extension loads connect when the session starts, together with the `mcp.json` servers; servers registered later connect right away, and `pi.unregisterMcpServer()` closes the connection and makes the server's tools unreachable. Registrations are not saved: register again on every load, for example based on the extension's own settings. A server in `mcp.json` with the same name takes precedence, and `/mcp` shows the override. Registering the same name again replaces the extension's earlier registration; names registered by another extension, invalid names, and invalid configs throw. + +The built-in MCP support connects registered servers. When nothing does, because another extension replaced it (see [MCP](mcp.md#other-mcp-extensions)), each registration is reported as an extension error. Other MCP extensions can connect registered servers too: read them with `pi.getMcpServers()` on `session_start` and handle the `mcp_servers_change` event for later changes. + diff --git a/packages/coding-agent/docs/mcp.md b/packages/coding-agent/docs/mcp.md new file mode 100644 index 000000000..72df5d2e7 --- /dev/null +++ b/packages/coding-agent/docs/mcp.md @@ -0,0 +1,188 @@ +# MCP Servers + +Pi connects to [Model Context Protocol](https://modelcontextprotocol.io) servers over stdio or streamable HTTP and makes their tools available to the model. + +## Configure servers + +Add servers to `~/.pi/agent/mcp.json`, or to `.pi/mcp.json` in a project. The format matches other MCP clients, so existing `mcpServers` entries can be copied over: + +```json +{ + "mcpServers": { + "filesystem": { + "command": "npx", + "args": ["-y", "@modelcontextprotocol/server-filesystem", "."] + }, + "docs": { + "url": "https://example.com/mcp", + "headers": { "Authorization": "Bearer ${DOCS_TOKEN}" }, + "exposure": "direct" + } + } +} +``` + +- stdio servers take `command`, `args`, `env`, and `cwd`. Relative `cwd` resolves against the session directory. A leading `~/` in `command`, an argument, or `cwd` names the home directory. +- HTTP servers take `url`, `headers`, and `oauth` (see [Sign in with OAuth](#sign-in-with-oauth)). The legacy SSE transport is not supported. +- `env` and `headers` values can reference environment variables (`${NAME}`) or commands (`!command`), like provider API keys. +- `timeout` sets the per-request timeout in seconds (default 60). Progress notifications from the server reset it. +- `enabled: false` keeps an entry without connecting to it. + +Project entries replace global entries with the same name. A project `mcp.json` is only read after the project is trusted, because stdio servers run commands. + +`pi mcp add` and `pi mcp remove` edit the file from a shell (see [MCP commands](cli.md#mcp-commands)): + +```bash +pi mcp add filesystem -- npx -y @modelcontextprotocol/server-filesystem . +pi mcp add docs --url https://example.com/mcp --bearer-token-env-var DOCS_TOKEN --exposure direct +pi mcp add -l tools --env API_KEY='${TOOLS_KEY}' -- uvx tools-mcp +pi mcp remove docs +``` + +Rules that are easy to get wrong: + +- Server names may only contain letters, digits, `_`, and `-`. Tools are named `mcp____`. +- `type` is optional: a `command` makes a stdio server and a `url` a streamable HTTP server. When present, it must be `stdio`, `http`, or `streamable-http`. `sse` is rejected; most servers that document an SSE endpoint also serve streamable HTTP, often at `/mcp` instead of `/sse`. +- `command` is a single executable and `args` its arguments, not one shell string. +- Keep secrets out of the file: use `${NAME}` for environment variables, as in `"Authorization": "Bearer ${GITHUB_TOKEN}"`, or `!command` to run a command. A command must make up the whole value, so it has to print the header value itself: `"Authorization": "!echo Bearer $(gh auth token)"`. +- Invalid entries are skipped and reported; the other servers still connect. + +## Set up servers + +When asked to add an MCP server, the agent should: + +1. Add simple servers with `pi mcp add` (add `-l` for the project file), or edit `mcp.json` directly for settings the command does not cover. Put personal servers and servers with credentials in `~/.pi/agent/mcp.json`. Use the project `.pi/mcp.json` only for servers the project itself needs, and only in trusted projects. +2. Convert entries written for other clients: + - Claude Desktop, Claude Code, and Cursor use the same `mcpServers` shape; copy the entry. + - VS Code uses a top-level `servers` object and `inputs` prompts; move the entry under `mcpServers` and replace `${input:...}` with `${NAME}` environment variables. + - Codex uses TOML (`[mcp_servers.]` with `command`, `args`, `env`, or `url`); write the same fields as JSON. + - opencode uses `"type": "local"` with `command` as an array (split it into `command` and `args`), `"type": "remote"` for URLs, `environment` for `env`, and `{env:NAME}` for `${NAME}`. +3. Run `pi mcp list` to check the entry. It connects to every enabled server and prints the state, the tools, and errors such as the stderr of a stdio server that failed to start. It exits with 1 while anything is wrong. +4. For a server that needs a sign-in, run `pi mcp login `. It opens the authorization page in the user's browser and waits until the user approves access; tell the user to approve it. A running session uses the new credentials on its next turn. +5. Tell the user to run `/reload` (or start a new session) so the running session connects to added or changed servers. + +Pi connects when a session starts. The first prompt waits up to 10 seconds for startup connections; the tools of servers that take longer become available once they connect. HTTP connections that fail with a network error or a transient status (408, 429, 5xx) are retried twice. A server that drops its connection shows as disconnected and is reconnected on the next call. When a server announces that its tool list changed, new tools are added and withdrawn tools become unreachable until the server offers them again. + +Config errors, servers that failed to connect, and servers that need a sign-in are reported once after startup. + +Log messages servers send with MCP logging notifications are appended to `~/.pi/agent/mcp.log` as `