feat(coding-agent): codemode and MCP

This adds codemode and MCP support to pi. Codemode (packages/codemode) runs
model-written JavaScript in a QuickJS wasm VM inside a worker. From there the
script calls pi's tools as async functions, uses a per-session store, and can
read the model catalog and run classifiers. The codemode tool in coding-agent
is compatible with Codex's exec tool, and codemode.mode works like Codex's
tool modes (on or only). Large results are written to a JSON temp file. MCP
(packages/mcp) is a standalone client for stdio and streamable HTTP servers,
with OAuth.

The core only gets general mechanisms; codemode, tool_search and MCP are
built-in extensions that use them. Other extensions can replace these
built-ins. Tools declare an exposure: direct, model-only, codemode, deferred
or hidden. prepareLoadout() lets an orchestrating tool change which tool
declarations the model sees. ctx.executeTool() runs nested calls through
validation and the tool_call/tool_result hooks. Those calls emit events with
parentToolCallId and are recorded as bounded nestedCalls on the parent
result, which compaction and HTML export use. Tools can also return
structuredContent with an outputSchema, and can report errors with isError
instead of throwing. MCP servers come from mcp.json (global, or per project
after the project is trusted) or from pi.registerMcpServer(). They are
managed with /mcp and pi mcp list|login|logout, and their runtimes load on
first use.

Closes #10040
This commit is contained in:
Armin Ronacher 2026-09-29 09:12:50 +02:00
parent b485fa3128
commit 8562bcf66a
143 changed files with 17620 additions and 225 deletions

View file

@ -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",

48
package-lock.json generated
View file

@ -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",

View file

@ -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",

View file

@ -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:

View file

@ -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 <mcp server command> [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();
}

View file

@ -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<unknown>): 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 `<prefix>_<tool>`. 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<AgentTool[]> {
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<Record<string, unknown>>(tool.outputSchema) : undefined,
execute: async (_toolCallId, params, signal, onUpdate) => {
const result = await client.callTool(tool.name, params as Record<string, unknown>, {
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<AgentToolCallOutcome> {
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.<name>(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<AgentToolCallOutcome>,
): AgentTool<typeof codemodeParameters> {
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.<name>(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();
}
},
};
}

View file

@ -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<any>;
isError: boolean;
};
type FinalizedToolCallOutcome = AgentToolCallOutcome;
/** The `beforeToolCall` and `afterToolCall` hooks of {@link AgentLoopConfig}. */
export type ToolCallHooks = Pick<AgentLoopConfig, "beforeToolCall" | "afterToolCall">;
type ToolUpdateSink = (partialResult: AgentToolResult<any>) => Promise<void> | void;
type FinalizedToolCallEntry = FinalizedToolCallOutcome | (() => Promise<FinalizedToolCallOutcome>);
@ -706,10 +708,11 @@ async function prepareToolCall(
currentContext: AgentContext,
assistantMessage: AssistantMessage,
toolCall: AgentToolCall,
config: AgentLoopConfig,
config: ToolCallHooks,
signal: AbortSignal | undefined,
tools: readonly AgentTool<any>[] = currentContext.tools ?? [],
): Promise<PreparedToolCall | ImmediateToolCallOutcome> {
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<any>[];
/** 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<AgentToolCallOutcome> {
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<ExecutedToolCallOutcome> {
const updateEvents: Promise<void>[] = [];
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<FinalizedToolCallOutcome> {
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) {

View file

@ -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<T = JsonValue | undefined> {
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<T = JsonValue | undefined> {
terminate?: boolean;
}
/** Final outcome of a tool call after hooks ran. */
export interface AgentToolCallOutcome {
toolCall: AgentToolCall;
result: AgentToolResult<any>;
isError: boolean;
}
/**
* Callback used by tools to stream partial execution updates.
*
@ -448,7 +469,15 @@ export interface AgentTool<TParameters extends TSchema = TSchema, TDetails = any
* Must return an object that matches `TParameters`.
*/
prepareArguments?: (args: unknown) => Static<TParameters>;
/** 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<TParameters>,

View file

@ -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<AssistantMessageEvent, AssistantMessage> {
@ -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<typeof echoSchema> = {
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" }]);
});
});

View file

@ -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<TDetails = JsonValue> = IsJsonCompatible<TDetails> extends true
? {
role: "toolResult";
@ -578,6 +600,8 @@ export type ToolResultMessage<TDetails = JsonValue> = IsJsonCompatible<TDetails>
details?: JsonRepresentation<TDetails>;
/** 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
}

View file

@ -0,0 +1,7 @@
# Changelog
## [Unreleased]
### Added
- Initial spike: `CodemodeSandbox` runs model-written JavaScript in a worker thread and exposes injected tools as `tools.<name>(args)` async functions.

186
packages/codemode/README.md Normal file
View file

@ -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.<name>(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<string>;
// };
```
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.<name>(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:<line>` |
| `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.

View file

@ -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"
}
}

View file

@ -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<T>`
* declarations can refer to them.
*/
export const MCP_TYPESCRIPT_PREAMBLE = `type Role = "user" | "assistant";
type MetaObject = Record<string, unknown>;
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<TStructured = { [key: string]: unknown }> = {
_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<string, string[]>();
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<R>;` 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<CallToolResult<T>>`,
* which needs {@link MCP_TYPESCRIPT_PREAMBLE}.
*/
export function renderToolSignature(
tool: Pick<CodemodeTool, "name" | "inputSchema" | "outputSchema">,
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<CodemodeTool, "name" | "description" | "inputSchema" | "outputSchema">,
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<string, unknown> {
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<T>` 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<string>;
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<string, unknown>, 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<string, unknown>, 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");
}

View file

@ -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;
}

View file

@ -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";

View file

@ -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<string> = 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<Record<string, unknown>> | undefined): Record<string, string> {
const serialized: Record<string, string> = {};
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<string, CodemodeTool>;
globals: ReadonlyMap<string, CodemodeTool>;
timeoutMs: number;
signal: AbortSignal | undefined;
memoryLimitBytes: number | undefined;
store: Record<string, string>;
wasm: Promise<CodemodeWasmModule>;
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<CodemodeResult>;
private resolveResult!: (result: CodemodeResult) => void;
private worker: Worker | undefined;
private readonly interrupt = new SharedArrayBuffer(4);
private readonly tools: ReadonlyMap<string, CodemodeTool>;
private readonly globals: ReadonlyMap<string, CodemodeTool>;
private readonly signal: AbortSignal | undefined;
private readonly timer: NodeJS.Timeout | undefined;
private readonly output: CodemodeOutputItem[] = [];
private readonly calls: CodemodeCall[] = [];
private readonly pending = new Map<number, PendingCall>();
private finished = false;
constructor(options: ExecutionOptions) {
this.promise = new Promise<CodemodeResult>((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<CodemodeResult> {
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<WorkerToHostMessage, { type: "done" }>): void {
if (!message.ok) {
const parsed = JSON.parse(message.error) as Omit<CodemodeError, "kind">;
this.finish({ kind: "script", ...parsed });
return;
}
this.finish(undefined, message.value === undefined ? undefined : JSON.parse(message.value), message.writes);
}
private async handleCall(message: Extract<WorkerToHostMessage, { type: "call" }>): Promise<void> {
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.<name>(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<string, CodemodeTool>();
private readonly globalsByName = new Map<string, CodemodeTool>();
private readonly timeoutMs: number;
private readonly memoryLimitBytes: number | undefined;
private readonly wasm: CodemodeWasmModule | Promise<CodemodeWasmModule> | undefined;
private readonly workerUrl: URL;
private readonly running = new Set<Execution>();
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<string>();
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<CodemodeResult> {
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<void> {
this.closed = true;
await Promise.all([...this.running].map((execution) => execution.abort("Sandbox closed")));
}
}

View file

@ -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;
},
};
})`;

View file

@ -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<string, string>;
/**
* 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";
}

View file

@ -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<void> {
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);
}

View file

@ -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<string, unknown>;
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()) };
}

View file

@ -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.<id>(args)`, where `<id>` is the name with characters that
* are not valid in identifiers replaced by `_` (see `toCodemodeIdentifier`), and also as
* `tools["<name>"](args)`. Globals are called as `<name>(args)` and must be identifiers, or
* `<namespace>.<member>`, 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<Model[]>`. 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> | 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<string, unknown>;
/** 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<CodemodeWasmModule>;
/**
* 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<Record<string, unknown>>;
}

View file

@ -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<CodemodeWasmModule>;
}
const modules = new Map<string, Promise<CodemodeWasmModule>>();
/**
* 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<CodemodeWasmModule> {
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;
}

View file

@ -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<string>");
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<unknown>;");
});
it("renders MCP CallToolResult output schemas as CallToolResult<T>", () => {
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<CallToolResult<{ results: Array<{ id: string; score: number; }>; }>>;",
);
expect(renderToolSignature({ name: "plain", inputSchema, outputSchema: mcpResultSchema() })).toBe(
"plain(args: {}): Promise<CallToolResult>;",
);
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<unknown>; };\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<string>;",
" remote_api(args: unknown): Promise<unknown>;",
"};",
"",
"/** Attach it. */",
"declare function attach(args: string): Promise<unknown>;",
].join("\n"),
);
});
it("renders namespaced globals and explicit signatures", () => {
const text = renderDeclarations({
globals: [
{
name: "models.list",
description: "List models.",
signature: "(type: string): Promise<string[]>",
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<string[]>;",
" get(args: string): Promise<unknown>;",
"};",
].join("\n"),
);
});
it("escapes comment terminators in descriptions", () => {
const text = renderDeclarations({ tools: [{ name: "x", description: "a */ b", execute }] });
expect(text).toContain("/** a *\\/ b */");
});
});

View file

@ -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<void>((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"] });
});
});

View file

@ -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);
}
});
});

View file

@ -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"]
}

View file

@ -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"] } },
});

View file

@ -13,6 +13,7 @@ pi update [target] [options]
pi list
pi config [options]
pi auth <check|print-api-key|print-bearer-token> [options]
pi mcp <list|login|logout> [options]
```
<a id="modes"></a>
@ -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.<name>(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.
<a id="resource-options"></a>
## Resources
@ -266,3 +288,20 @@ Authentication commands require `--provider <provider>` or `--model <model>`. Se
| `--min-expiry <duration>` | `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 <server> [options] -- <command> [args...]` | Add or replace a stdio server in `mcp.json`; `--env KEY=VALUE` (repeatable) and `--cwd <dir>` set its environment and working directory. Arguments after the command are passed to it |
| `pi mcp add <server> [options] --url <url>` | Add or replace a streamable HTTP server; `--header KEY=VALUE` (repeatable), `--bearer-token-env-var <NAME>` (sends `Authorization: Bearer ${NAME}`), `--oauth-client-id`, `--oauth-client-secret`, and `--oauth-callback-port` configure authentication |
| `pi mcp remove <server>` | 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 <server> [--timeout <seconds>]` | 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 <server>` | 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 <mode>` (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.

View file

@ -12,6 +12,7 @@ The agent directory is shown as `<agent-dir>` below. Set its location with the `
|---|---|
| `<agent-dir>/settings.json` | User-level [settings](settings.md), including preferences, defaults, resource paths, and Pi package declarations. |
| `<agent-dir>/keybindings.json` | Custom terminal UI and application [keybindings](keybindings.md). |
| `<agent-dir>/mcp.json` | [MCP servers](mcp.md) available in every project. |
| `<agent-dir>/models.json` | [Compatible endpoints, models, and model overrides](models.md#configure-a-compatible-endpoint). |
| `<agent-dir>/auth.json` | Saved API keys and OAuth credentials. |
| `<agent-dir>/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 `<agent-dir>` 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. |

View file

@ -91,6 +91,10 @@
{
"title": "Use Pi Packages",
"path": "packages.md"
},
{
"title": "Connect MCP Servers",
"path": "mcp.md"
}
]
},

View file

@ -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 `<parent id>/<n>`. 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.
<a id="extensioncontext"></a>
<a id="extensioncommandcontext"></a>
<a id="use-extension-context"></a>

View file

@ -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__<server>__<tool>`.
- `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.<name>]` 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 <server>`. 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 `<time> [<server>] <level> <logger>: <message>`. The file is moved to `mcp.log.1` when it grows past 5 MB.
## Manage servers
`/mcp` opens the server manager. It lists every configured server with its state, tool count, exposure, and whether it comes from the global or the project `mcp.json`; servers that need attention come first. Select a server to:
- sign in, for OAuth servers that need it (see [Sign in with OAuth](#sign-in-with-oauth))
- see its tools, its command or URL, and the full connection error, including the tail of a stdio server's stderr
- reconnect
- sign out, which deletes the stored OAuth credentials
- change its exposure (see [Exposure](#exposure))
- disable or enable it
Exposure changes and enabling or disabling are saved to the `mcp.json` that defines the server; other content of the file is kept. Disabled servers stay listed so they can be enabled again.
Outside the interactive TUI, `/mcp` prints the server status. `/mcp login <server>`, `/mcp logout <server>`, and `/mcp reconnect <server>` run those actions directly.
From a shell, `pi mcp add`, `pi mcp remove`, `pi mcp list`, `pi mcp login <server>`, and `pi mcp logout <server>` manage servers without a session (see [MCP commands](cli.md#mcp-commands)).
Stopping a stdio server closes its stdin, then sends SIGTERM and finally SIGKILL to its whole process group, so servers started through wrappers such as `npx` or `uvx` do not linger.
## Sign in with OAuth
Remote servers that use OAuth, such as Sentry, need no credentials in `mcp.json`:
```json
{
"mcpServers": {
"sentry": { "url": "https://mcp.sentry.dev/mcp" }
}
}
```
When such a server rejects the connection, `/mcp` shows it as needing sign-in. Select it and choose "Sign in" (or run `/mcp login sentry`, or `pi mcp login sentry` in a shell) to open the authorization page in your browser. After you approve access, the browser redirects to a temporary server on `127.0.0.1` and pi connects. If the browser runs on another machine, for example over SSH, paste the URL it was redirected to into the sign-in screen instead.
Pi registers itself with the authorization server (dynamic client registration), stores tokens in `~/.pi/agent/mcp-auth.json`, and refreshes access tokens automatically when they expire or the server rejects them. If the server later asks for more scope than was granted, it shows as needing sign-in again, and signing in requests the new scope. "Sign out" in `/mcp` (or `/mcp logout sentry`) deletes the stored credentials.
OAuth applies to HTTP servers without an `Authorization` header. For authorization servers that do not support dynamic client registration, configure a pre-registered client:
```json
{
"mcpServers": {
"example": {
"url": "https://mcp.example.com/mcp",
"oauth": { "clientId": "my-client", "clientSecret": "${EXAMPLE_SECRET}", "callbackPort": 8765 }
}
}
}
```
The redirect URI must match the one registered for the client. `callbackPort` fixes it to `http://127.0.0.1:<port>/callback`. For another redirect URI, set `callbackUrl`, for example `"callbackUrl": "http://localhost:8080/oauth/callback"`. It must be an `http` URI on `localhost`, `127.0.0.1`, or `[::1]`, and is sent exactly as written. Without a port in `callbackUrl`, pi listens on `callbackPort`, or on a free port, and adds it to the URI; authorization servers accept any port for loopback redirects (RFC 8252). `clientSecret` is optional and can reference environment variables or commands.
`scope` sets the scopes to request, separated by spaces, for servers that do not advertise the ones they need. Without it, pi requests the scopes the server advertises. When a server later asks for more scope, pi requests those on top of `scope`.
## Exposure
Each server's tools are registered as `mcp__<server>__<tool>`. The `exposure` setting controls how the model reaches them:
- `codemode` (default): the tools are callable from [`codemode`](cli.md#tools) scripts and listed in the `codemode` tool's description, but are not declared to the model. Large MCP tool lists stay out of the model's tool declarations, and scripts can call several MCP tools, in parallel if needed, while returning only the part of the result the model needs. Pi activates the `codemode` tool when such a server connects. Large servers do not fill the description: declarations share a token budget, and scripts find the remaining tools with `searchTools()` (see [`codemode`](cli.md#tools)).
- `codemode-deferred`: like `codemode`, but the tools are not listed in the `codemode` tool's description either; it only names the server and its tool count. Scripts call them by name and find them with `searchTools()` or in `ALL_TOOLS`. Use it for large servers that codemode scripts use rarely.
- `deferred`: the tools are not declared to the model until the [`tool_search`](cli.md#tools) tool loads them. The model searches, and the matches are declared from its next call on and called directly, without codemode. Pi activates the `tool_search` tool when such a server connects. Use it for large servers without codemode.
- `direct`: the tools are declared to the model like built-in tools, and are also callable from codemode.
- `hidden`: the tools are registered but cannot be called.
`toolExposure` sets the exposure of single tools and overrides `exposure` for them. Keys are tool names as the server offers them, or patterns where `*` matches any characters. An exact name wins over patterns; among patterns, the first match in the object wins. With `hidden` as the server's exposure, only the listed tools are reachable:
```json
{
"mcpServers": {
"github": {
"url": "https://api.githubcopilot.com/mcp/",
"exposure": "deferred",
"toolExposure": {
"search_code": "direct",
"get_*": "codemode",
"delete_*": "hidden"
}
}
}
}
```
`pi mcp list` marks tools whose exposure differs from the server's, and the Tools view in `/mcp` shows it too.
Tools that are not declared (`codemode`, `codemode-deferred`, and `deferred` exposure) are reachable through either tool: codemode scripts can call all of them, and `tool_search` can load any of them. For example, with `codemode` active, scripts can call the tools of a `deferred` server, and with `tool_search` active, the model can load the tools of a `codemode` server.
Tools called from codemode scripts do not depend on the active tool set, so they stay callable after `/tree`, resume, and fork. Tools loaded by `tool_search` are recorded in the transcript like any other tool change and stay declared on that branch. To keep pi from activating the `codemode` tool, set `"autoEnableCodemode": false` at the top level of `mcp.json`, next to `mcpServers`. A project `mcp.json` value overrides the global one. Pi warns once when neither `codemode` nor `tool_search` is active, since the tools then cannot be called.
Text results over 20KB reach the model with the middle cut out, in the format Codex uses: the start and end of the text around a `…N chars truncated…` marker. The full text is saved to a temp file whose path the result names. Codemode scripts always receive the whole result, so a script can filter a large result down to what the model needs.
Codemode scripts receive an MCP tool's whole `CallToolResult` (`content` blocks as sent by the server, `structuredContent`, and `isError`), and the `codemode` description declares it as `CallToolResult<T>`. A result with `isError` resolves in scripts and is reported to the model as an error for direct calls. `image(result.content[0])` forwards an image block to the model. The server's `instructions` describe its tools in the `codemode` description.
## Resources
When a connected server offers [resources](https://modelcontextprotocol.io/specification/2025-11-25/server/resources), pi adds the resource tools Codex and opencode use:
- `list_mcp_resources` lists resources as JSON: `{ server?, resources: [{ server, uri, name, ... }], nextCursor? }`. With `server`, it lists one page of that server, and `cursor` continues with the next one. Without, it lists every resource of every server.
- `list_mcp_resource_templates` lists URI templates for resources the servers do not list, in the same way.
- `read_mcp_resource` reads a resource given `server` and `uri`. Text resources reach the model as text and images as images; other binary resources are saved to temp files, and the model sees the file path. Scripts receive `{ server, uri, contents }`.
The tools reach every enabled server with resources whose exposure is not `hidden`, and take the widest exposure among them: `direct` if one of the servers is direct, else `codemode`, else `codemode-deferred`, else `deferred`. Resource links in tool results name `read_mcp_resource` and the server.
Resources for MCP Apps (`ui://` URIs or `text/html;profile=mcp-app`) are left out of the listings, since pi does not render them, and so are resource icons.
Reading and listing resources is retried once after a transient HTTP error (408, 429, 5xx). Tool calls are not retried, since the server may have run them.
## Permissions
Every MCP call goes through pi's tool pipeline, so `tool_call` and `tool_result` extension handlers, including permission gates, apply to MCP tools. Calls made from codemode scripts carry the `codemode` call's id as `parentToolCallId`. `pi.getAllTools()` reports the tool annotations servers declare (`readOnlyHint`, `destructiveHint`, `idempotentHint`, `openWorldHint`), so a permission extension can confirm only calls that change something (see [Extensions](extensions.md#tool-exposure)). The resource tools are marked read-only.
## Servers from extensions
Extensions can add servers for the current session with `pi.registerMcpServer(name, config)`, using the same config shape as `mcp.json` (see [Extensions](extensions.md#mcp-servers)). They connect like configured servers and appear in `/mcp` with the extension as their source. Enabling, disabling, and exposure changes for them apply to the current session only. A server in `mcp.json` with the same name takes precedence; `/mcp` lists the overridden registration. `pi mcp` shell commands do not load extensions and only see `mcp.json` servers.
## Other MCP extensions
An installed extension that registers the `/mcp` command, such as `pi-mcp-adapter`, replaces the built-in MCP support: pi then neither reads `mcp.json` in sessions nor connects servers, and `/mcp` belongs to that extension. Remove the extension to use the built-in support. Likewise, an extension that registers a tool named `codemode` or `tool_search` replaces the built-in tool of that name. `pi mcp` shell commands always use the built-in support.
## SDK
SDK sessions do not load the built-in extensions. Add the MCP extension, the codemode extension for `codemode` and `codemode-deferred` servers, and the tool search extension for `deferred` servers to the resource loader. See [SDK](sdk.md#codemode-mcp).

View file

@ -109,7 +109,11 @@ Use `DefaultResourceLoader` when you want standard discovery with selected overr
<a id="inlineextension"></a>
Inline extension factories can be supplied through `DefaultResourceLoader`. Give one an `InlineExtension` name only when it needs a stable name in diagnostics and startup output.
Inline extension factories can be supplied through `DefaultResourceLoader`. Give one an `InlineExtension` name only when it needs a stable name in diagnostics and startup output. A named inline extension with `replaceable: true` is left out when another extension registers a tool, command, or flag with a name it registers during loading, instead of both loading with a conflict. The CLI's built-in codemode, tool search, and MCP extensions are replaceable.
<a id="codemode-mcp"></a>
The CLI loads `codemode`, `tool_search`, and MCP as built-in extensions. SDK sessions do not; add `createCodemodeExtension()`, `createToolSearchExtension()`, and `createMcpExtension()` to the `extensionFactories` of `DefaultResourceLoader`. `codemode` and `tool_search` are registered inactive: enable them through the `defaultTools` setting, or let the MCP extension activate them: `codemode` for servers with `codemode` or `codemode-deferred` exposure, `tool_search` for servers with `deferred` exposure. The MCP extension connects its servers on `session_start`, so call `session.bindExtensions()`. See [Codemode and MCP](../examples/sdk/14-codemode-mcp.ts).
See the focused examples for [models](../examples/sdk/02-custom-model.ts), [tools](../examples/sdk/05-tools.ts), [extensions](../examples/sdk/06-extensions.ts), and [full control](../examples/sdk/12-full-control.ts).
@ -130,6 +134,7 @@ See the focused examples for [models](../examples/sdk/02-custom-model.ts), [tool
| [Sessions](../examples/sdk/11-sessions.ts) | Control session persistence and restoration |
| [Full control](../examples/sdk/12-full-control.ts) | Replace default discovery and state services |
| [Session runtime](../examples/sdk/13-session-runtime.ts) | Replace the active session safely |
| [Codemode and MCP](../examples/sdk/14-codemode-mcp.ts) | Add the `codemode`, `tool_search`, and MCP extensions |
<a id="exports"></a>

View file

@ -37,6 +37,7 @@ Project trust does not limit what tool calls can access or affect. After Pi star
Pi requires a project-trust decision when it finds any of these resources from the current working directory:
- `.pi/settings.json`
- `.pi/mcp.json`
- `.pi/extensions`, `.pi/skills`, `.pi/prompts`, or `.pi/themes`
- `.pi/SYSTEM.md` or `.pi/APPEND_SYSTEM.md`
- project `.agents/skills` in the current directory or an ancestor directory
@ -46,6 +47,7 @@ A bare `.pi` directory does not require project trust.
Granting project trust allows Pi to load:
- project settings
- project MCP servers from `.pi/mcp.json`
- extensions, skills, prompt templates, themes, and system-prompt files under `.pi`
- missing packages configured through project settings
- project-local and project-package extensions

View file

@ -38,8 +38,10 @@ See [Choose a Model](models.md) for model selection and thinking controls.
| Setting | Type | Default | Description |
|---|---|---|---|
| `defaultTools` | `string[]` | `read`, `bash`, `edit`, `write` | Built-in tools enabled at startup. An empty array disables all built-in tools but not extension or SDK tools. |
| `codemode.mode` | `"on"` \| `"only"` | `"on"` | How the `codemode` tool presents tools while it is active. `on`: declared tools get their `codemode` declaration appended to their description, and `codemode` lists only tools that are not declared (MCP `codemode` exposure). `only`: `codemode` lists every tool scripts can call, and active built-in and extension tools are hidden from the model, so it reaches them through `codemode`. |
| `codemode.inlineBudget` | number | `3000` | Estimated tokens (characters / 4) the `codemode` tool's description may spend on tool declarations. Tools that do not fit are left out and found with `searchTools()`. `0` lists only namespaces. |
Available built-in tools are `read`, `bash`, `powershell`, `edit`, `write`, `grep`, `find`, and `ls`. CLI tool options override this setting for one invocation. See [Command Line](cli.md#tools).
Available built-in tools are `read`, `bash`, `powershell`, `edit`, `write`, `grep`, `find`, and `ls`. `defaultTools` can also name `codemode` and `tool_search`, which built-in extensions register inactive. CLI tool options override this setting for one invocation. See [Command Line](cli.md#tools).
## Sessions and context

View file

@ -0,0 +1,60 @@
/**
* Codemode, Tool Search, and MCP
*
* The CLI loads `codemode`, `tool_search`, and MCP as built-in extensions. SDK sessions add
* them to the resource loader's extension factories.
*
* `codemode` and `tool_search` are registered inactive. Enable them through the `defaultTools` setting
* (`tools` would also restrict the session to the named tools, which hides MCP tools), or let the
* MCP extension activate them: `codemode` for servers with `codemode` or `codemode-deferred` exposure,
* `tool_search` for servers with `deferred` exposure.
*/
import {
createAgentSession,
createCodemodeExtension,
createMcpExtension,
createToolSearchExtension,
DefaultResourceLoader,
getAgentDir,
SessionManager,
SettingsManager,
} from "@earendil-works/pi-coding-agent";
const cwd = process.cwd();
const resourceLoader = new DefaultResourceLoader({
cwd,
agentDir: getAgentDir(),
extensionFactories: [
createCodemodeExtension({ mode: "on" }),
createToolSearchExtension(),
// Reads mcp.json from the agent directory and the trusted project, like the CLI.
createMcpExtension(),
],
});
await resourceLoader.reload();
const settingsManager = SettingsManager.create(cwd);
settingsManager.applyOverrides({ defaultTools: ["read", "bash", "edit", "write", "codemode", "tool_search"] });
const { session } = await createAgentSession({
resourceLoader,
settingsManager,
sessionManager: SessionManager.inMemory(),
});
try {
// Emits session_start, which connects the MCP servers in the background.
await session.bindExtensions({});
console.log("Active tools:", session.getActiveToolNames().join(", "));
session.subscribe((event) => {
if (event.type === "message_update" && event.assistantMessageEvent.type === "text_delta") {
process.stdout.write(event.assistantMessageEvent.delta);
}
});
await session.prompt("Use codemode to count the TypeScript files in src/ and list the three largest.");
console.log();
} finally {
session.dispose();
}

View file

@ -21,6 +21,7 @@ The runtime example shows how to build a recreate function that closes over proc
| `11-sessions.ts` | In-memory, persistent, continue, list sessions |
| `12-full-control.ts` | Replace everything, no discovery |
| `13-session-runtime.ts` | Manage runtime-backed session replacement |
| `14-codemode-mcp.ts` | Add the `codemode`, `tool_search`, and MCP extensions |
## Running

View file

@ -496,6 +496,17 @@
"node": ">=22.0.0"
}
},
"node_modules/@earendil-works/pi-codemode": {
"version": "0.87.1",
"resolved": "https://registry.npmjs.org/@earendil-works/pi-codemode/-/pi-codemode-0.87.1.tgz",
"license": "MIT",
"dependencies": {
"quickjs-wasi": "3.6.2"
},
"engines": {
"node": ">=22.19.0"
}
},
"node_modules/@earendil-works/pi-coding-agent": {
"version": "0.87.1",
"resolved": "https://registry.npmjs.org/@earendil-works/pi-coding-agent/-/pi-coding-agent-0.87.1.tgz",
@ -504,6 +515,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",
@ -516,6 +529,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",
@ -528,6 +542,17 @@
"node": ">=22.19.0"
}
},
"node_modules/@earendil-works/pi-mcp": {
"version": "0.87.1",
"resolved": "https://registry.npmjs.org/@earendil-works/pi-mcp/-/pi-mcp-0.87.1.tgz",
"license": "MIT",
"dependencies": {
"cross-spawn": "7.0.6"
},
"engines": {
"node": ">=22.19.0"
}
},
"node_modules/@earendil-works/pi-telemetry": {
"version": "0.87.1",
"resolved": "https://registry.npmjs.org/@earendil-works/pi-telemetry/-/pi-telemetry-0.87.1.tgz",
@ -1761,6 +1786,12 @@
"node": ">= 20"
}
},
"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/retry": {
"version": "0.13.1",
"resolved": "https://registry.npmjs.org/retry/-/retry-0.13.1.tgz",

View file

@ -12,6 +12,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",
@ -24,6 +26,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",
@ -518,6 +521,28 @@
"node": ">=22.0.0"
}
},
"node_modules/@earendil-works/pi-codemode": {
"version": "0.87.1",
"resolved": "https://registry.npmjs.org/@earendil-works/pi-codemode/-/pi-codemode-0.87.1.tgz",
"license": "MIT",
"dependencies": {
"quickjs-wasi": "3.6.2"
},
"engines": {
"node": ">=22.19.0"
}
},
"node_modules/@earendil-works/pi-mcp": {
"version": "0.87.1",
"resolved": "https://registry.npmjs.org/@earendil-works/pi-mcp/-/pi-mcp-0.87.1.tgz",
"license": "MIT",
"dependencies": {
"cross-spawn": "7.0.6"
},
"engines": {
"node": ">=22.19.0"
}
},
"node_modules/@earendil-works/pi-telemetry": {
"version": "0.87.1",
"resolved": "https://registry.npmjs.org/@earendil-works/pi-telemetry/-/pi-telemetry-0.87.1.tgz",
@ -1751,6 +1776,12 @@
"node": ">= 20"
}
},
"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/retry": {
"version": "0.13.1",
"resolved": "https://registry.npmjs.org/retry/-/retry-0.13.1.tgz",

View file

@ -41,7 +41,7 @@
"clean": "shx rm -rf dist",
"build": "npm run build:unbundled && node ../../scripts/build-coding-agent-bundle.mjs",
"build:unbundled": "tsc -p tsconfig.build.json && shx chmod +x dist/cli.js dist/rpc-entry.js && npm run copy-assets",
"build:binary": "npm --prefix ../tui run build && npm --prefix ../telemetry run build && npm --prefix ../ai run build && npm --prefix ../agent run build && npm --prefix ../protocol run build && npm --prefix ../client run build && npm run build && bun build --compile --no-compile-autoload-bunfig ./src/bun/cli.ts ./src/utils/image-resize-worker.ts --outfile dist/pi && npm run copy-binary-assets",
"build:binary": "npm --prefix ../tui run build && npm --prefix ../telemetry run build && npm --prefix ../ai run build && npm --prefix ../agent run build && npm --prefix ../protocol run build && npm --prefix ../client run build && npm run build && bun build --compile --no-compile-autoload-bunfig ./dist/bun/cli.js ./src/utils/image-resize-worker.ts ./src/extensions/codemode/worker.ts --outfile dist/pi && npm run copy-binary-assets",
"copy-assets": "shx mkdir -p dist/modes/interactive/theme && shx cp src/modes/interactive/theme/*.json dist/modes/interactive/theme/ && shx mkdir -p dist/modes/interactive/assets && shx cp src/modes/interactive/assets/*.png dist/modes/interactive/assets/ && shx mkdir -p dist/core/export-html/vendor && shx cp src/core/export-html/template.html src/core/export-html/template.css src/core/export-html/template.js dist/core/export-html/ && shx cp src/core/export-html/vendor/*.js dist/core/export-html/vendor/",
"copy-binary-assets": "shx cp package.json dist/ && shx cp README.md dist/ && shx cp CHANGELOG.md dist/ && shx mkdir -p dist/theme && shx cp src/modes/interactive/theme/*.json dist/theme/ && shx mkdir -p dist/assets && shx cp src/modes/interactive/assets/*.png dist/assets/ && shx mkdir -p dist/export-html/vendor && shx cp src/core/export-html/template.html dist/export-html/ && shx cp src/core/export-html/vendor/*.js dist/export-html/vendor/ && shx cp -r docs dist/ && shx cp -r examples dist/ && shx cp ../../node_modules/@silvia-odwyer/photon-node/photon_rs_bg.wasm dist/",
"test": "vitest --run",
@ -52,6 +52,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",
@ -64,6 +66,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",

View file

@ -0,0 +1,5 @@
// Imported by runtime-setup.ts. Bun's default loader for .wasm evaluates to the file's path.
declare module "quickjs-wasi/quickjs.wasm" {
const path: string;
export default path;
}

View file

@ -1,9 +1,12 @@
import { bedrockProviderModule } from "@earendil-works/pi-ai/bedrock-provider";
import { registerBunOAuthFlows } from "@earendil-works/pi-ai/bun-oauth";
import { setBedrockProviderModule } from "@earendil-works/pi-ai/compat";
import { APP_NAME } from "../config.ts";
// Bun loads .wasm imports as files: embedded in compiled executables, evaluating to a readable path.
import quickjsWasmPath from "quickjs-wasi/quickjs.wasm";
import { APP_NAME, setEmbeddedQuickJSWasmPath } from "../config.ts";
process.title = APP_NAME;
process.emitWarning = (() => {}) as typeof process.emitWarning;
registerBunOAuthFlows();
setBedrockProviderModule(bedrockProviderModule);
setEmbeddedQuickJSWasmPath(quickjsWasmPath);

View file

@ -282,7 +282,8 @@ ${chalk.bold("Commands:")}
${APP_NAME} list List installed extensions from settings
${APP_NAME} config [-l] Open TUI to enable/disable package resources (Tab switches scope)
${APP_NAME} auth <command> Print credentials or check provider readiness
${APP_NAME} <command> --help Show help for install/remove/uninstall/update/list/config/auth
${APP_NAME} mcp <command> Check MCP servers, sign in to or out of OAuth servers
${APP_NAME} <command> --help Show help for install/remove/uninstall/update/list/config/auth/mcp
${chalk.bold("Options:")}
--provider <name> Provider name (default: google)

View file

@ -1,4 +1,5 @@
import { accessSync, constants, existsSync, readFileSync, realpathSync } from "fs";
import { createRequire } from "module";
import { homedir } from "os";
import { basename, dirname, join, resolve, sep, win32 } from "path";
import { fileURLToPath } from "url";
@ -476,6 +477,33 @@ export function getBundledInteractiveAssetPath(name: string): string {
return join(getInteractiveAssetsDir(), name);
}
let embeddedQuickJSWasmPath: string | undefined;
/** Called by the Bun entry with the path of the QuickJS wasm file embedded in the compiled executable. */
export function setEmbeddedQuickJSWasmPath(path: string): void {
embeddedQuickJSWasmPath = path;
}
/** Get path to `quickjs-wasi/quickjs.wasm`, the VM that runs codemode scripts. */
export function getQuickJSWasmPath(): string {
return embeddedQuickJSWasmPath ?? createRequire(import.meta.url).resolve("quickjs-wasi/quickjs.wasm");
}
/**
* Get the URL of the codemode worker entry (`src/extensions/codemode/worker.ts`), or undefined to use
* the worker file that ships next to pi-codemode's own module.
* - For Bun binary: the build passes the worker as an extra entrypoint. Bun embeds it at its path
* relative to the common directory of all entrypoints (the package root, since the main entry is
* dist/bun/cli.js) with a .js extension, and all bundled code sees the executable as import.meta.url.
* - For the Node bundle: the build emits codemode-worker.js next to the chunk that contains this module.
* - For Node.js (dist/) and tsx (src/): pi-codemode's own worker.
*/
export function getCodemodeWorkerUrl(): URL | undefined {
if (isBunBinary) return new URL("./src/extensions/codemode/worker.js", import.meta.url);
if (isBundledNode) return new URL("./codemode-worker.js", import.meta.url);
return undefined;
}
// =============================================================================
// App Config (from package.json piConfig)
// =============================================================================

View file

@ -15,15 +15,21 @@
import { readFileSync } from "node:fs";
import { basename, dirname } from "node:path";
import type {
Agent,
AgentContext,
AgentEvent,
AgentMessage,
AgentState,
AgentTool,
PrepareNextTurnContext,
ThinkingLevel,
import {
type AfterToolCallContext,
type AfterToolCallResult,
type Agent,
type AgentContext,
type AgentEvent,
type AgentMessage,
type AgentState,
type AgentTool,
type AgentToolCallOutcome,
type BeforeToolCallContext,
type BeforeToolCallResult,
type PrepareNextTurnContext,
runToolCall,
type ThinkingLevel,
} from "@earendil-works/pi-agent-core";
import { contentText, getCurrentSystemMessage, retryDelayMs } from "@earendil-works/pi-ai";
import type {
@ -78,6 +84,7 @@ import {
type AgentActivityOutcome,
type BoundaryContextPreview,
type ContextUsage,
type ExecuteToolOptions,
type ExtensionCommandContextActions,
type ExtensionErrorListener,
type ExtensionMode,
@ -98,7 +105,9 @@ import {
type ToolExecutionEndEvent,
type ToolExecutionStartEvent,
type ToolExecutionUpdateEvent,
type ToolExposure,
type ToolInfo,
type ToolLoadout,
type TreePreparation,
type TurnStartEvent,
wrapRegisteredTools,
@ -107,6 +116,7 @@ import { emitSessionShutdownEvent } from "./extensions/runner.ts";
import { type BashExecutionMessage, type CustomMessage, convertToLlm } from "./messages.ts";
import { ModelRegistry } from "./model-registry.ts";
import type { ModelRuntime } from "./model-runtime.ts";
import { NestedToolCallRunner } from "./nested-tool-calls.ts";
import { expandPromptTemplate, type PromptTemplate } from "./prompt-templates.ts";
import type { ResourceExtensionPaths, ResourceLoader } from "./resource-loader.ts";
import { exportSessionToJsonl } from "./session-export.ts";
@ -169,9 +179,16 @@ export function parseSkillBlock(text: string): ParsedSkillBlock | null {
};
}
/** Tool execution events of calls a tool made through `ctx.executeTool()` carry `parentToolCallId`. */
type WithParentToolCallId<E> = E extends {
type: "tool_execution_start" | "tool_execution_update" | "tool_execution_end";
}
? E & { parentToolCallId?: string }
: E;
/** Session-specific events that extend the core AgentEvent */
export type AgentSessionEvent =
| Exclude<AgentEvent, { type: "agent_end" }>
| WithParentToolCallId<Exclude<AgentEvent, { type: "agent_end" }>>
| {
type: "agent_end";
messages: AgentMessage[];
@ -418,6 +435,10 @@ export class AgentSession {
// Tool registry for extension getTools/setTools
private _toolRegistry: Map<string, AgentTool> = new Map();
/** Created on the first `ctx.executeTool()` call. */
private _nestedToolCalls: NestedToolCallRunner | undefined;
/** Declared tools whose declarations requests leave out, from `prepareLoadout` hooks. */
private _hiddenDeclarations: ReadonlySet<string> = new Set();
private _toolDefinitions: Map<string, ToolDefinitionEntry> = new Map();
private _toolPromptSnippets: Map<string, string> = new Map();
private _toolPromptGuidelines: Map<string, string[]> = new Map();
@ -453,6 +474,7 @@ export class AgentSession {
this._installAgentNextTurnRefresh();
this._installAgentRequestProjection();
this._installAgentBoundaryHooks();
this._installHiddenDeclarationsProjection();
this._installAgentForcedPromptProjection();
this._buildRuntime({
@ -587,61 +609,117 @@ export class AgentSession {
* happens here instead of in wrappers.
*/
private _installAgentToolHooks(): void {
this.agent.beforeToolCall = async ({ toolCall, args }) => {
const runner = this._extensionRunner;
if (!runner.hasHandlers("tool_call")) {
return undefined;
}
this.agent.beforeToolCall = (context) => this._beforeToolCall(context);
this.agent.afterToolCall = (context) => this._afterToolCall(context);
}
try {
return await runner.emitToolCall({
type: "tool_call",
/** `tool_call` handlers. `parentToolCallId` is set for calls another tool made. */
private async _beforeToolCall(
{ toolCall, args }: BeforeToolCallContext,
parentToolCallId?: string,
): Promise<BeforeToolCallResult | undefined> {
const runner = this._extensionRunner;
if (!runner.hasHandlers("tool_call")) {
return undefined;
}
try {
return await runner.emitToolCall({
type: "tool_call",
toolName: toolCall.name,
toolCallId: toolCall.id,
...(parentToolCallId ? { parentToolCallId } : {}),
input: args as Record<string, unknown>,
});
} catch (err) {
if (err instanceof Error) {
throw err;
}
throw new Error(`Extension failed, blocking execution: ${String(err)}`);
}
}
/** `tool_result` handlers and image normalization. `parentToolCallId` is set for calls another tool made. */
private async _afterToolCall(
{ toolCall, args, result, isError }: AfterToolCallContext,
parentToolCallId?: string,
): Promise<AfterToolCallResult | undefined> {
const runner = this._extensionRunner;
const hookResult = runner.hasHandlers("tool_result")
? await runner.emitToolResult({
type: "tool_result",
toolName: toolCall.name,
toolCallId: toolCall.id,
...(parentToolCallId ? { parentToolCallId } : {}),
input: args as Record<string, unknown>,
});
} catch (err) {
if (err instanceof Error) {
throw err;
content: result.content,
details: result.details,
...(result.structuredContent === undefined ? {} : { structuredContent: result.structuredContent }),
isError,
usage: result.usage,
})
: undefined;
const content = hookResult?.content ?? result.content ?? [];
// Runs after the extension hook so images injected or replaced by extensions are normalized too.
const resizeOptions = this._limitsModel()?.inputLimits?.images?.resize;
const normalizedContent = await normalizeToolResultImages(content, {
autoResizeImages: this.settingsManager.getImageAutoResize(),
...(resizeOptions ? { resizeOptions } : {}),
});
if (!hookResult && normalizedContent === content) {
return undefined;
}
// The hook result already dropped structured content that replaced content no longer matches.
return {
content: normalizedContent,
details: hookResult?.details,
structuredContent: hookResult ? hookResult.structuredContent : result.structuredContent,
isError: hookResult?.isError ?? isError,
usage: hookResult?.usage,
};
}
/**
* Run a call that the tool call `parentToolCallId` made through `ctx.executeTool()`. It goes
* through the agent's tool pipeline with the session's hooks, against the callable tools.
*/
private async _executeNestedToolCall(
parentToolCallId: string,
name: string,
args: unknown,
options: ExecuteToolOptions,
): Promise<AgentToolCallOutcome> {
this._nestedToolCalls ??= new NestedToolCallRunner({
getTools: () => this._getCallableTools(),
isSequential: () => this.agent.toolExecution === "sequential",
runToolCall: (toolCall, parentId, signal, onUpdate) => {
const assistantMessage = this._findLastAssistantMessage();
if (!assistantMessage) {
return Promise.resolve({
toolCall,
result: { content: [{ type: "text", text: "No assistant message issued this call" }], details: {} },
isError: true,
});
}
throw new Error(`Extension failed, blocking execution: ${String(err)}`);
}
};
this.agent.afterToolCall = async ({ toolCall, args, result, isError }) => {
const runner = this._extensionRunner;
const hookResult = runner.hasHandlers("tool_result")
? await runner.emitToolResult({
type: "tool_result",
toolName: toolCall.name,
toolCallId: toolCall.id,
input: args as Record<string, unknown>,
content: result.content,
details: result.details,
isError,
usage: result.usage,
})
: undefined;
const content = hookResult?.content ?? result.content ?? [];
// Runs after the extension hook so images injected or replaced by extensions are normalized too.
const resizeOptions = this._limitsModel()?.inputLimits?.images?.resize;
const normalizedContent = await normalizeToolResultImages(content, {
autoResizeImages: this.settingsManager.getImageAutoResize(),
...(resizeOptions ? { resizeOptions } : {}),
});
if (!hookResult && normalizedContent === content) {
return undefined;
}
return {
content: normalizedContent,
details: hookResult?.details,
isError: hookResult?.isError ?? isError,
usage: hookResult?.usage,
};
};
return runToolCall(toolCall, {
tools: this._getCallableTools(),
assistantMessage,
context: { messages: this.agent.state.messages, tools: this.agent.state.tools },
beforeToolCall: (context) => this._beforeToolCall(context, parentId),
afterToolCall: (context) => this._afterToolCall(context, parentId),
signal,
onUpdate,
});
},
emit: async (event) => {
await this._extensionRunner.emit(event);
this._emit(event);
},
});
return this._nestedToolCalls.execute(parentToolCallId, name, args, options);
}
/** Whether `projection`, the current session projection, exceeds the compaction threshold of `model`. */
@ -981,6 +1059,15 @@ export class AgentSession {
/** Internal handler for agent events - shared by subscribe and reconnect */
private _handleAgentEvent = async (event: AgentEvent): Promise<void> => {
// Record the calls a tool made through ctx.executeTool() on its result message.
if (this._nestedToolCalls) {
if (event.type === "message_start" && event.message.role === "toolResult") {
const nestedCalls = this._nestedToolCalls.takeRecord(event.message.toolCallId);
if (nestedCalls) event.message.nestedCalls = nestedCalls;
} else if (event.type === "agent_end") {
this._nestedToolCalls.clear();
}
}
// When a user message starts, check if it's from either queue and remove it BEFORE emitting
// This ensures the UI sees the updated queue state
if (event.type === "message_start" && event.message.role === "user") {
@ -1343,13 +1430,18 @@ export class AgentSession {
}
/**
* Get the names of currently active tools.
* Returns the names of tools currently set on the agent.
* Get the names of currently active tools, which are the tools declared to the model.
* Tools with `codemode` or `deferred` exposure are callable from other tools without being active.
*/
getActiveToolNames(): string[] {
return this.agent.state.tools.map((t) => t.name);
}
/** Get the names of the tools that tools can call through `ctx.executeTool()`. */
getCallableToolNames(): string[] {
return this._getCallableTools().map((t) => t.name);
}
/**
* Get all configured tools with name, description, parameter schema, prompt guidelines, and source metadata.
*/
@ -1359,6 +1451,9 @@ export class AgentSession {
description: definition.description,
parameters: definition.parameters,
promptGuidelines: definition.promptGuidelines,
exposure: this._getToolExposure(definition.name),
...(definition.namespace ? { namespace: definition.namespace } : {}),
...(definition.annotations ? { annotations: { ...definition.annotations } } : {}),
sourceInfo,
}));
}
@ -1369,22 +1464,80 @@ export class AgentSession {
/**
* Set active tools by name.
* Only tools in the registry can be enabled. Unknown tool names are ignored.
* Only tools in the registry can be enabled. Unknown and hidden tool names are ignored.
* Also rebuilds the system prompt to reflect the new tool set.
* Changes take effect on the next agent turn.
*/
setActiveToolsByName(toolNames: string[]): void {
const tools: AgentTool[] = [];
const validToolNames: string[] = [];
for (const name of toolNames) {
const tools = this._applyToolLoadout(toolNames);
this._rebuildSystemPrompt(tools.map((tool) => tool.name));
}
private _getToolExposure(name: string): ToolExposure {
return this._toolDefinitions.get(name)?.definition.exposure ?? "direct";
}
/**
* Tools callable through `ctx.executeTool()`: the active `direct` tools and every registered
* `codemode` or `deferred` tool.
*/
private _getCallableTools(active: ReadonlySet<string> = new Set(this.getActiveToolNames())): AgentTool[] {
return [...this._toolRegistry.values()].filter((tool) => {
const exposure = this._getToolExposure(tool.name);
return exposure === "codemode" || exposure === "deferred" || (exposure === "direct" && active.has(tool.name));
});
}
/**
* Set the agent's tools for the given active tool names and return them. The active tools are
* the registered, non-hidden ones; they are declared to the model. Active tools with a
* `prepareLoadout` hook can change the declared descriptions and hide declarations from
* requests (see {@link _installHiddenDeclarationsProjection}).
*/
private _applyToolLoadout(toolNames: string[]): AgentTool[] {
const tools = [...new Set(toolNames)].flatMap((name) => {
const tool = this._toolRegistry.get(name);
if (tool) {
tools.push(tool);
validToolNames.push(name);
return tool && this._getToolExposure(name) !== "hidden" ? [tool] : [];
});
const hooks = tools.flatMap((tool) => {
const entry = this._toolDefinitions.get(tool.name);
return entry?.definition.prepareLoadout ? [entry] : [];
});
const hidden = new Set<string>();
let declared = tools;
if (hooks.length > 0) {
const loadout: ToolLoadout = {
declared: tools,
callable: this._getCallableTools(new Set(tools.map((tool) => tool.name))),
registered: [...this._toolRegistry.values()],
getExposure: (name) => this._getToolExposure(name),
getNamespace: (name) => this._toolDefinitions.get(name)?.definition.namespace,
};
const descriptions = new Map<string, string>();
for (const { definition, sourceInfo } of hooks) {
try {
const changes = definition.prepareLoadout?.(loadout);
for (const [name, description] of Object.entries(changes?.descriptions ?? {})) {
descriptions.set(name, description);
}
for (const name of changes?.hiddenDeclarations ?? []) hidden.add(name);
} catch (error) {
this._extensionRunner.emitError({
extensionPath: sourceInfo.path,
event: "prepare_loadout",
error: error instanceof Error ? error.message : String(error),
stack: error instanceof Error ? error.stack : undefined,
});
}
}
declared = tools.map((tool) => {
const description = descriptions.get(tool.name);
return description === undefined ? tool : { ...tool, description };
});
}
this.agent.state.tools = tools;
this._rebuildSystemPrompt(validToolNames);
this._hiddenDeclarations = hidden;
this.agent.state.tools = declared;
return declared;
}
/** Whether compaction or branch summarization is currently running */
@ -1505,11 +1658,7 @@ export class AgentSession {
options: NormalizedBuildSystemPromptOptions,
messages: AgentMessage[] = this.agent.state.messages,
): SystemMessage | undefined {
options.selectedTools = [...new Set(options.selectedTools)].filter((name) => this._toolRegistry.has(name));
this.agent.state.tools = options.selectedTools.flatMap((name) => {
const tool = this._toolRegistry.get(name);
return tool ? [tool] : [];
});
options.selectedTools = this._applyToolLoadout(options.selectedTools).map((tool) => tool.name);
const sections = diffSystemPromptSections(
getCurrentSystemMessage(messages)?.sections ?? {},
buildSystemPromptSections(options),
@ -1527,6 +1676,31 @@ export class AgentSession {
* collapse into one head holding the forced text and the current tools. Runs after the
* `context` extension handlers.
*/
/**
* Remove the declarations that `prepareLoadout` hooks hide from every request. The whole
* transcript is filtered with the current set, so the projected declarations stay consistent
* across requests and only change when the loadout does.
*/
private _installHiddenDeclarationsProjection(): void {
const previousTransformContext = this.agent.transformContext;
this.agent.transformContext = async (messages, signal) => {
const transformed = previousTransformContext ? await previousTransformContext(messages, signal) : messages;
const hidden = this._hiddenDeclarations;
if (hidden.size === 0) return transformed;
return transformed.map((message) => {
if (message.role !== "system" || (!message.toolsAdded && !message.toolsRemoved)) return message;
const { toolsAdded, toolsRemoved, ...rest } = message;
const added = toolsAdded?.filter((tool) => !hidden.has(tool.name)) ?? [];
const removed = toolsRemoved?.filter((tool) => !hidden.has(tool.name)) ?? [];
return {
...rest,
...(added.length > 0 ? { toolsAdded: added } : {}),
...(removed.length > 0 ? { toolsRemoved: removed } : {}),
};
});
};
}
private _installAgentForcedPromptProjection(): void {
const previousTransformContext = this.agent.transformContext;
this.agent.transformContext = async (messages, signal) => {
@ -1544,18 +1718,15 @@ export class AgentSession {
};
}
/** Restore the active tool loadout declared by the session transcript, if it declares one. */
/**
* Restore the active tool loadout declared by the session transcript, if it declares one.
* Tools reachable only from other tools are never declared, but they do not depend on the active
* set, so the transcript's declarations are the whole loadout.
*/
private _restoreToolsFromTranscript(): void {
const current = getCurrentSystemMessage(this.sessionManager.buildSessionContext().messages);
if (!current) return;
const toolNames = (current.toolsAdded ?? [])
.map((tool) => tool.name)
.filter((name) => this._toolRegistry.has(name));
this.agent.state.tools = toolNames.flatMap((name) => {
const registered = this._toolRegistry.get(name);
return registered ? [registered] : [];
});
this._rebuildSystemPrompt(toolNames);
this.setActiveToolsByName((current.toolsAdded ?? []).map((tool) => tool.name));
}
// =========================================================================
@ -3017,6 +3188,7 @@ export class AgentSession {
this._applyExtensionBindings(this._extensionRunner);
await this._extensionRunner.emit(this._sessionStartEvent);
this._extensionRunner.reportUnhandledMcpServers();
await this.extendResourcesFromExtensions(this._sessionStartEvent.reason === "reload" ? "reload" : "startup");
}
@ -3160,6 +3332,7 @@ export class AgentSession {
},
getActiveTools: () => this.getActiveToolNames(),
getAllTools: () => this.getAllTools(),
getSettings: () => this.settingsManager.getSettings(),
setActiveTools: (toolNames) => this.setActiveToolsByName(toolNames),
refreshTools: () => this._refreshToolRegistry(),
getCommands,
@ -3202,6 +3375,8 @@ export class AgentSession {
},
getSystemPrompt: () => this.systemPrompt,
getSystemPromptOptions: () => this._baseSystemPromptOptions,
executeTool: (callerId, name, args, options) => this._executeNestedToolCall(callerId, name, args, options),
getCallableTools: () => this._getCallableTools(),
},
{
registerProvider: (name, config) => {
@ -3229,7 +3404,11 @@ export class AgentSession {
}
private _refreshToolRegistry(options?: { activeToolNames?: string[]; includeAllExtensionTools?: boolean }): void {
const previousRegistryNames = new Set(this._toolRegistry.keys());
// Tools that were already activated on registration. A tool whose exposure changes to
// `direct` or `model-only` (for example from `hidden`) is activated like a new tool.
const previousActivatedOnRegistration = new Set(
[...this._toolRegistry.keys()].filter((name) => this._isActivatedOnRegistration(name)),
);
const previousActiveToolNames = this.getActiveToolNames();
const allowedToolNames = this._allowedToolNames;
const excludedToolNames = this._excludedToolNames;
@ -3302,17 +3481,18 @@ export class AgentSession {
if (allowedToolNames) {
for (const toolName of this._toolRegistry.keys()) {
if (allowedToolNames.has(toolName)) {
// Naming a tool activates it even when it is not active by default.
if (allowedToolNames.has(toolName) && this._isDeclarable(toolName)) {
nextActiveToolNames.push(toolName);
}
}
} else if (options?.includeAllExtensionTools) {
for (const tool of wrappedExtensionTools) {
nextActiveToolNames.push(tool.name);
if (this._isActivatedOnRegistration(tool.name)) nextActiveToolNames.push(tool.name);
}
} else if (!options?.activeToolNames) {
for (const toolName of this._toolRegistry.keys()) {
if (!previousRegistryNames.has(toolName)) {
if (!previousActivatedOnRegistration.has(toolName) && this._isActivatedOnRegistration(toolName)) {
nextActiveToolNames.push(toolName);
}
}
@ -3321,6 +3501,17 @@ export class AgentSession {
this.setActiveToolsByName([...new Set(nextActiveToolNames)]);
}
/** Whether activating the tool declares it to the model. */
private _isDeclarable(name: string): boolean {
const exposure = this._getToolExposure(name);
return exposure === "direct" || exposure === "model-only";
}
/** Whether registering the tool activates it, which declares it to the model. */
private _isActivatedOnRegistration(name: string): boolean {
return this._isDeclarable(name) && this._toolDefinitions.get(name)?.definition.defaultActive !== false;
}
private _buildRuntime(options: {
activeToolNames?: string[];
flagValues?: Map<string, boolean | string>;
@ -3398,6 +3589,7 @@ export class AgentSession {
if (hasBindings) {
await options?.beforeSessionStart?.();
await this._extensionRunner.emit({ type: "session_start", reason: "reload" });
this._extensionRunner.reportUnhandledMcpServers();
await this.extendResourcesFromExtensions("reload");
}
}

View file

@ -24,9 +24,15 @@ export function createFileOps(): FileOperations {
}
/**
* Extract file operations from tool calls in an assistant message.
* Extract file operations from tool calls in an assistant message, or from the nested calls
* recorded on a tool result.
*/
export function extractFileOpsFromMessage(message: AgentMessage, fileOps: FileOperations): void {
if (message.role === "toolResult") {
// Calls made from codemode scripts are recorded on the script's result.
for (const call of message.nestedCalls?.calls ?? []) addFileOp(call.name, call.arguments, fileOps);
return;
}
if (message.role !== "assistant") return;
if (!("content" in message) || !Array.isArray(message.content)) return;
@ -34,24 +40,23 @@ export function extractFileOpsFromMessage(message: AgentMessage, fileOps: FileOp
if (typeof block !== "object" || block === null) continue;
if (!("type" in block) || block.type !== "toolCall") continue;
if (!("arguments" in block) || !("name" in block)) continue;
addFileOp(block.name, block.arguments as Record<string, unknown> | undefined, fileOps);
}
}
const args = block.arguments as Record<string, unknown> | undefined;
if (!args) continue;
const path = typeof args.path === "string" ? args.path : undefined;
if (!path) continue;
switch (block.name) {
case "read":
fileOps.read.add(path);
break;
case "write":
fileOps.written.add(path);
break;
case "edit":
fileOps.edited.add(path);
break;
}
function addFileOp(toolName: string, args: Record<string, unknown> | undefined, fileOps: FileOperations): void {
const path = typeof args?.path === "string" ? args.path : undefined;
if (!path) return;
switch (toolName) {
case "read":
fileOps.read.add(path);
break;
case "write":
fileOps.written.add(path);
break;
case "edit":
fileOps.edited.add(path);
break;
}
}

View file

@ -930,6 +930,21 @@
'</div>';
};
// Calls this tool made to other tools (for example from a codemode script), recorded without results.
const renderNestedCalls = () => {
const nested = result?.nestedCalls;
if (!nested || !Array.isArray(nested.calls) || nested.calls.length === 0) return '';
const icons = { ok: '✓', error: '✗', unfinished: '…' };
const lines = nested.calls.map(c => {
const args = c.arguments ? JSON.stringify(c.arguments) : `[arguments omitted, ${c.argumentsBytes} bytes]`;
const duration = c.durationMs !== undefined ? ` ${c.durationMs}ms` : '';
const error = c.error ? `\n ${c.error.split('\n').join('\n ')}` : '';
return `${icons[c.status] || '?'} ${c.name} ${args}${duration}${error}`;
});
const title = `Nested calls: ${nested.calls.length}${nested.complete ? '' : ' (incomplete record)'}`;
return formatExpandableOutput([title, ...lines].join('\n'), 1);
};
const toolDomId = `tool-call-${escapeHtml(call.id)}`;
let html = `<div class="tool-execution ${statusClass}" id="${toolDomId}">`;
const args = call.arguments || {};
@ -1063,6 +1078,7 @@
}
}
html += renderNestedCalls();
html += '</div>';
return html;
}

View file

@ -70,6 +70,7 @@ export type {
EntryRenderOptions,
ExecOptions,
ExecResult,
ExecuteToolOptions,
Extension,
ExtensionActions,
// API
@ -88,6 +89,7 @@ export type {
// Runtime
ExtensionRuntime,
ExtensionShortcut,
ExtensionToolContext,
ExtensionUIContext,
ExtensionUIDialogOptions,
ExtensionVirtualModel,
@ -111,6 +113,7 @@ export type {
LsToolResultEvent,
MarkdownTransformContext,
MarkdownTransformer,
McpServersChangeEvent,
// Events - Message
MessageEndEvent,
MessageEndEventResult,
@ -167,6 +170,7 @@ export type {
SetThinkingLevelHandler,
TerminalInputHandler,
ThinkingLevelSelectEvent,
ToolAnnotations,
// Events - Tool
ToolCallEvent,
ToolCallEventResult,
@ -178,7 +182,11 @@ export type {
ToolExecutionMode,
ToolExecutionStartEvent,
ToolExecutionUpdateEvent,
ToolExposure,
ToolInfo,
ToolLoadout,
ToolLoadoutChanges,
ToolNamespace,
ToolRenderResultOptions,
ToolResultEvent,
ToolResultEventResult,

View file

@ -15,6 +15,7 @@ import { resolvePath } from "../../utils/paths.ts";
import { createEventBus, type EventBus } from "../event-bus.ts";
import type { ExecOptions } from "../exec.ts";
import { execCommand } from "../exec.ts";
import { type McpServerConfig, McpServerRegistry, validateMcpServerConfig } from "../mcp-servers.ts";
import { readPiManifest } from "../pi-manifest.ts";
import { createSyntheticSourceInfo } from "../source-info.ts";
import { time } from "../timings.ts";
@ -173,6 +174,7 @@ export function createExtensionRuntime(): ExtensionRuntime {
setLabel: notInitialized,
getActiveTools: notInitialized,
getAllTools: notInitialized,
getSettings: notInitialized,
setActiveTools: notInitialized,
// registerTool() is valid during extension load; refresh is only needed post-bind.
refreshTools: () => {},
@ -183,6 +185,7 @@ export function createExtensionRuntime(): ExtensionRuntime {
flagValues: new Map(),
pendingProviderRegistrations: [],
pendingNativeProviderRegistrations: [],
mcpServers: new McpServerRegistry(),
pendingVirtualModelRegistrations: [],
createContext: notInitialized,
assertActive,
@ -405,6 +408,11 @@ function createExtensionAPI(
return runtime.getAllTools();
},
getSettings() {
assertActive();
return runtime.getSettings();
},
setActiveTools(toolNames: string[]): void {
assertActive();
runtime.setActiveTools(toolNames);
@ -445,6 +453,30 @@ function createExtensionAPI(
applyRuntimeChange(() => runtime.unregisterProvider(name, extension.path));
},
registerMcpServer(name: string, config: McpServerConfig) {
assertActive();
const validated = validateMcpServerConfig(name, config);
if (typeof validated === "string") {
throw new Error(`Invalid MCP server registered by extension "${extension.path}": ${validated}`);
}
const owner = runtime.mcpServers.get(name)?.extensionPath;
if (owner !== undefined && owner !== extension.path) {
throw new Error(`MCP server "${name}" is already registered by extension "${owner}"`);
}
const server = { name, config: structuredClone(validated), extensionPath: extension.path };
applyRuntimeChange(() => runtime.mcpServers.register(server));
},
unregisterMcpServer(name: string) {
assertActive();
applyRuntimeChange(() => runtime.mcpServers.unregister(name, extension.path));
},
getMcpServers() {
assertActive();
return runtime.mcpServers.list();
},
registerVirtualModel<TState>(model: ExtensionVirtualModel<TState>) {
assertActive();
// Routing runs after the runner binds, so the context is created per request. The state

View file

@ -2,7 +2,7 @@
* Extension runner - executes extensions and manages their lifecycle.
*/
import type { AgentMessage } from "@earendil-works/pi-agent-core";
import type { AgentMessage, AgentTool } from "@earendil-works/pi-agent-core";
import {
getCurrentSystemMessage,
type ImageContent,
@ -41,6 +41,7 @@ import type {
ContextUsage,
ContextWithSystemEvent,
EntryRenderer,
ExecuteToolOptions,
Extension,
ExtensionActions,
ExtensionCommandContext,
@ -53,6 +54,7 @@ import type {
ExtensionMode,
ExtensionRuntime,
ExtensionShortcut,
ExtensionToolContext,
ExtensionUIContext,
InputEvent,
InputEventResult,
@ -373,6 +375,10 @@ export class ExtensionRunner {
private getSystemPromptFn: () => string = () => "";
private getSystemPromptOptionsFn: () => BuildSystemPromptOptions = () =>
normalizeBuildSystemPromptOptions({ cwd: this.cwd });
private executeToolFn: ExtensionContextActions["executeTool"];
private getCallableToolsFn: () => readonly AgentTool[] = () => [];
/** Registered MCP servers already reported as unhandled. */
private readonly reportedMcpServers = new Set<string>();
private newSessionHandler: NewSessionHandler = async () => ({ cancelled: false });
private forkHandler: ForkHandler = async () => ({ cancelled: false });
private navigateTreeHandler: NavigateTreeHandler = async () => ({ cancelled: false });
@ -420,6 +426,7 @@ export class ExtensionRunner {
this.runtime.setLabel = actions.setLabel;
this.runtime.getActiveTools = actions.getActiveTools;
this.runtime.getAllTools = actions.getAllTools;
this.runtime.getSettings = actions.getSettings;
this.runtime.setActiveTools = actions.setActiveTools;
this.runtime.refreshTools = actions.refreshTools;
this.runtime.getCommands = actions.getCommands;
@ -442,6 +449,15 @@ export class ExtensionRunner {
this.getSystemPromptFn = contextActions.getSystemPrompt;
this.getSystemPromptOptionsFn =
contextActions.getSystemPromptOptions ?? (() => normalizeBuildSystemPromptOptions({ cwd: this.cwd }));
this.executeToolFn = contextActions.executeTool;
this.getCallableToolsFn = contextActions.getCallableTools ?? (() => []);
// Servers registered from now on reach the extension that connects them right away. Servers
// registered during loading are read on session_start.
this.runtime.mcpServers.setChangeListener(() => {
void this.emit({ type: "mcp_servers_change", servers: this.runtime.mcpServers.list() });
this.reportUnhandledMcpServers();
});
// Flush provider registrations queued during extension loading
for (const { name, config, extensionPath } of this.runtime.pendingProviderRegistrations) {
@ -728,6 +744,23 @@ export class ExtensionRunner {
}
}
/**
* Report registered MCP servers when no extension handles `mcp_servers_change`, which means
* nothing connects them (for example when another MCP extension replaced the built-in one).
*/
reportUnhandledMcpServers(): void {
if (this.hasHandlers("mcp_servers_change")) return;
for (const server of this.runtime.mcpServers.list()) {
if (this.reportedMcpServers.has(server.name)) continue;
this.reportedMcpServers.add(server.name);
this.emitError({
extensionPath: server.extensionPath,
event: "register_mcp_server",
error: `MCP server "${server.name}" is registered, but no loaded extension connects MCP servers; another extension may have replaced the built-in MCP support`,
});
}
}
hasHandlers(eventType: string): boolean {
for (const ext of this.extensions) {
const handlers = ext.handlers.get(eventType);
@ -912,6 +945,39 @@ export class ExtensionRunner {
};
}
/**
* Create the context for executing the tool call `toolCallId`: the extension context plus
* `tools` and `executeTool()`. `signal` is the default signal of nested calls.
*/
createToolContext(toolCallId: string, signal: AbortSignal | undefined): ExtensionToolContext {
const runner = this;
// createContext() returns a fresh object, so adding properties does not affect other contexts.
return Object.defineProperties(this.createContext() as ExtensionToolContext, {
tools: {
get() {
runner.assertActive();
return runner.getCallableToolsFn();
},
},
executeTool: {
value: async (name: string, args: unknown, options: ExecuteToolOptions = {}) => {
runner.assertActive();
if (!runner.executeToolFn) {
return {
toolCall: { type: "toolCall", id: `${toolCallId}/0`, name, arguments: {} },
result: {
content: [{ type: "text", text: "Nested tool calls are not available in this context" }],
details: {},
},
isError: true,
};
}
return runner.executeToolFn(toolCallId, name, args, { ...options, signal: options.signal ?? signal });
},
},
});
}
createCommandContext(): ExtensionCommandContext {
// Use property descriptors instead of object spread so the guarded getters from
// createContext() stay lazy. A spread would eagerly read them once and freeze the
@ -1118,12 +1184,18 @@ export class ExtensionRunner {
if (handlerResult.content !== undefined) {
currentEvent.content = handlerResult.content;
// Structured content that is not replaced along with the content may no longer match it.
if (handlerResult.structuredContent === undefined) delete currentEvent.structuredContent;
modified = true;
}
if (handlerResult.details !== undefined) {
currentEvent.details = handlerResult.details;
modified = true;
}
if (handlerResult.structuredContent !== undefined) {
currentEvent.structuredContent = handlerResult.structuredContent;
modified = true;
}
if (handlerResult.isError !== undefined) {
currentEvent.isError = handlerResult.isError;
modified = true;
@ -1152,6 +1224,7 @@ export class ExtensionRunner {
return {
content: currentEvent.content,
details: currentEvent.details,
structuredContent: currentEvent.structuredContent,
isError: currentEvent.isError,
usage: currentEvent.usage,
};

View file

@ -10,6 +10,8 @@
import type {
AgentMessage,
AgentTool,
AgentToolCallOutcome,
AgentToolResult,
AgentToolUpdateCallback,
ThinkingLevel,
@ -24,6 +26,7 @@ import type {
ConstrainedSamplingConfig,
ImageApi,
ImageContent,
JsonValue,
Message,
Model,
OAuthCredentials,
@ -60,6 +63,7 @@ import type { EventBus } from "../event-bus.ts";
import type { ExecOptions, ExecResult } from "../exec.ts";
import type { ReadonlyFooterDataProvider } from "../footer-data-provider.ts";
import type { KeybindingsManager } from "../keybindings.ts";
import type { McpServerConfig, McpServerRegistry, RegisteredMcpServer } from "../mcp-servers.ts";
import type { CustomMessage } from "../messages.ts";
import type { ModelRegistry } from "../model-registry.ts";
import type { ScopedModel } from "../model-resolver.ts";
@ -73,6 +77,7 @@ import type {
SessionEntry,
SessionManager,
} from "../session-manager.ts";
import type { Settings } from "../settings-manager.ts";
import type { SlashCommandInfo } from "../slash-commands.ts";
import type { SourceInfo } from "../source-info.ts";
import type { BuildSystemPromptOptions, NormalizedBuildSystemPromptOptions } from "../system-prompt.ts";
@ -359,6 +364,36 @@ export interface ExtensionContext {
getSystemPrompt(): string;
}
/** Options for {@link ExtensionToolContext.executeTool}. */
export interface ExecuteToolOptions {
/** Defaults to the calling tool's signal. */
signal?: AbortSignal;
/** Receives partial results of the nested tool, in addition to `tool_execution_update` events. */
onUpdate?: AgentToolUpdateCallback;
}
/**
* Context passed to tool `execute()` in a session: the extension context plus `executeTool()`
* for running other tools through the same validation, hooks, and permission checks as
* model-issued calls.
*
* A tool wrapped with `wrapToolDefinition()` without a context factory, such as a built-in tool
* created with `createBashTool()` and run in a plain `Agent` or called directly, gets no context.
*/
export interface ExtensionToolContext extends ExtensionContext {
/** Tools {@link executeTool} can call. */
readonly tools: readonly AgentTool[];
/**
* Run another tool. The call gets the id `<calling id>/<n>`, and the `tool_call`, `tool_result`,
* and `tool_execution_*` events carry `parentToolCallId`. It does not appear in the transcript;
* a bounded record of it is kept as `nestedCalls` on the calling tool's result message.
*
* Never rejects for tool failures: unknown tools, validation errors, blocked calls, and thrown
* errors come back as `isError: true`.
*/
executeTool(name: string, args: unknown, options?: ExecuteToolOptions): Promise<AgentToolCallOutcome>;
}
/**
* Extended context for command handlers.
* Includes session control methods only safe in user-initiated commands.
@ -456,6 +491,69 @@ export interface ToolRenderContext<TState = any, TArgs = any> {
isError: boolean;
}
/**
* How the model reaches a tool. "Callable" means callable from other tools through
* `ctx.executeTool()`, as the `codemode` tool does.
*
* - `direct`: declared to the model while active, and callable while active.
* - `model-only`: declared to the model while active, never callable. Use it for orchestrating or
* interactive tools.
* - `codemode`: callable whenever registered. Not declared to the model unless explicitly
* activated. Codemode tools list it in their description.
* - `deferred`: like `codemode`, but codemode tools do not list it; tool search can find it.
* - `hidden`: registered but unreachable. Activating it has no effect.
*
* `direct` and `model-only` tools are activated when they are registered; the others are not.
* The active tool set (`getActiveTools`/`setActiveTools`) is the set declared to the model.
*/
export type ToolExposure = "direct" | "model-only" | "codemode" | "deferred" | "hidden";
/**
* Hints about what a tool does, with the meaning of MCP tool annotations. They come from the tool's
* author and are not verified; permission extensions can use them to decide which calls to confirm.
*/
export interface ToolAnnotations {
/** The tool does not modify its environment. */
readOnlyHint?: boolean;
/** The tool may delete or overwrite data, rather than only add to it. Meaningful when not read-only. */
destructiveHint?: boolean;
/** Repeating a call with the same arguments has no further effect. Meaningful when not read-only. */
idempotentHint?: boolean;
/** The tool reaches an open world of external entities, such as the web, rather than a closed domain. */
openWorldHint?: boolean;
}
/** A group of related tools, such as the tools of one MCP server. Codemode tools list them together. */
export interface ToolNamespace {
/** For example `mcp__docs`. */
name: string;
/** Shown once above the group's tools. */
description?: string;
}
/** The tools of a session as {@link ToolDefinition.prepareLoadout} sees them. */
export interface ToolLoadout {
/** Tools declared to the model (the active tools), in order, with their original descriptions. */
readonly declared: readonly AgentTool[];
/** Tools callable through `ctx.executeTool()`. */
readonly callable: readonly AgentTool[];
/** Every registered tool. */
readonly registered: readonly AgentTool[];
getExposure(name: string): ToolExposure;
getNamespace(name: string): ToolNamespace | undefined;
}
/** Changes {@link ToolDefinition.prepareLoadout} makes to what the model sees. */
export interface ToolLoadoutChanges {
/** Model-facing descriptions of declared tools, by tool name. */
descriptions?: Readonly<Record<string, string>>;
/**
* Declared tools whose declarations requests leave out. They stay active and callable, and the
* transcript still declares them, so the active set survives `/tree` and resume.
*/
hiddenDeclarations?: readonly string[];
}
/**
* Tool definition for registerTool().
*/
@ -480,6 +578,37 @@ export interface ToolDefinition<TParams extends TSchema = TSchema, TDetails = un
/** Optional compatibility shim to prepare raw tool call arguments before schema validation. Must return an object conforming to TParams. */
prepareArguments?: (args: unknown) => Static<TParams>;
/**
* JSON Schema of `structuredContent` in successful results. Tools that declare it should always
* set `structuredContent`; codemode scripts then receive it instead of the text content.
*/
outputSchema?: TSchema;
/**
* How the model reaches the tool. Default: `"direct"`. See {@link ToolExposure}.
*/
exposure?: ToolExposure;
/** Group the tool belongs to, for example its MCP server. */
namespace?: ToolNamespace;
/** Hints about what the tool does, for example from an MCP server. */
annotations?: ToolAnnotations;
/**
* Whether registering the tool activates it. Default: `true` for `direct` and `model-only` tools;
* other exposures are never activated on registration. A tool with `defaultActive: false` is
* activated by naming it in `--tools` or the `defaultTools` setting, or with `setActiveTools()`.
*/
defaultActive?: boolean;
/**
* Adjust how the loadout is presented to the model while this tool is active. Called whenever
* the active tools change. Tools that orchestrate other tools use it, for example to list the
* callable tools in their own description.
*/
prepareLoadout?: (loadout: ToolLoadout) => ToolLoadoutChanges | undefined;
/**
* Per-tool execution mode override.
* - "sequential": this tool must execute one at a time with other tool calls.
@ -495,7 +624,7 @@ export interface ToolDefinition<TParams extends TSchema = TSchema, TDetails = un
params: Static<TParams>,
signal: AbortSignal | undefined,
onUpdate: AgentToolUpdateCallback<TDetails> | undefined,
ctx: ExtensionContext,
ctx: ExtensionToolContext,
): Promise<AgentToolResult<TDetails>>;
/** Custom rendering for tool call display */
@ -567,6 +696,18 @@ export interface ResourcesDiscoverResult {
themePaths?: string[];
}
/**
* Fired when an extension registers or unregisters an MCP server after the extensions are bound
* (see {@link ExtensionAPI.registerMcpServer}). Servers registered while extensions load are read
* with `pi.getMcpServers()` on `session_start`. Handling this event marks an extension as the one
* that connects registered servers.
*/
export interface McpServersChangeEvent {
type: "mcp_servers_change";
/** Every registered server after the change. */
servers: RegisteredMcpServer[];
}
// ============================================================================
// Session Events
// ============================================================================
@ -900,6 +1041,8 @@ export interface ToolExecutionStartEvent {
toolCallId: string;
toolName: string;
args: any;
/** Set when another tool (for example a codemode script) made this call. */
parentToolCallId?: string;
}
/** Fired during tool execution with partial/streaming output */
@ -909,6 +1052,8 @@ export interface ToolExecutionUpdateEvent {
toolName: string;
args: any;
partialResult: any;
/** Set when another tool (for example a codemode script) made this call. */
parentToolCallId?: string;
}
/** Fired when a tool finishes executing */
@ -918,6 +1063,8 @@ export interface ToolExecutionEndEvent {
toolName: string;
result: any;
isError: boolean;
/** Set when another tool (for example a codemode script) made this call. */
parentToolCallId?: string;
}
// ============================================================================
@ -988,7 +1135,14 @@ export type InputEventResult =
interface ToolCallEventBase {
type: "tool_call";
/**
* The call's id. For calls another tool made (with `parentToolCallId` set), pi assigns
* `<parent id>/<n>`; such ids never appear as tool calls or tool results in the transcript, only
* in the parent result's `nestedCalls` record.
*/
toolCallId: string;
/** Set when another tool (for example a codemode script) issued this call. */
parentToolCallId?: string;
}
export interface BashToolCallEvent extends ToolCallEventBase {
@ -1055,9 +1209,17 @@ export type ToolCallEvent =
interface ToolResultEventBase {
type: "tool_result";
/** The call's id; `<parent id>/<n>` for nested calls, see `ToolCallEvent`. */
toolCallId: string;
/** Set when another tool (for example a codemode script) issued this call. */
parentToolCallId?: string;
input: Record<string, unknown>;
content: (TextContent | ImageContent)[];
/**
* Machine-readable result for tools that declare an `outputSchema`. Handlers that redact
* `content` should also replace this; replacing `content` alone drops it.
*/
structuredContent?: JsonValue;
isError: boolean;
/** Usage from the tool execution itself, if available. */
usage?: Usage;
@ -1186,6 +1348,7 @@ export function isToolCallEventType(toolName: string, event: ToolCallEvent): boo
export type ExtensionEvent =
| ProjectTrustEvent
| ResourcesDiscoverEvent
| McpServersChangeEvent
| SessionEvent
| ContextEvent
| ContextWithSystemEvent
@ -1255,9 +1418,15 @@ export type UserBashEventResult =
result: BashResult;
};
/**
* Changes a `tool_result` handler makes. Omitted fields stay as they are, except that replacing
* `content` without returning `structuredContent` drops the structured content, because it may no
* longer match. Return it along with `content` to keep it.
*/
export interface ToolResultEventResult {
content?: (TextContent | ImageContent)[];
details?: unknown;
structuredContent?: JsonValue;
isError?: boolean;
usage?: Usage;
}
@ -1390,6 +1559,7 @@ export interface ExtensionAPI {
on(event: "session_compact", handler: ExtensionHandler<SessionCompactEvent>): () => void;
on(event: "session_compact_failed", handler: ExtensionHandler<SessionCompactFailedEvent>): () => void;
on(event: "session_shutdown", handler: ExtensionHandler<SessionShutdownEvent>): () => void;
on(event: "mcp_servers_change", handler: ExtensionHandler<McpServersChangeEvent>): () => void;
on(
event: "session_before_tree",
handler: ExtensionHandler<SessionBeforeTreeEvent, SessionBeforeTreeResult>,
@ -1532,13 +1702,19 @@ export interface ExtensionAPI {
/** Execute a shell command. */
exec(command: string, args: string[], options?: ExecOptions): Promise<ExecResult>;
/** Get the list of currently active tool names. */
/** Get the names of the active tools, which are the tools declared to the model. */
getActiveTools(): string[];
/** Get all configured tools with parameter schema, prompt guidelines, and source metadata. */
/** Get all configured tools with parameter schema, prompt guidelines, exposure, and source metadata. */
getAllTools(): ToolInfo[];
/** Set the active tools by name. */
/** Get a copy of the effective settings (global and project settings merged, with overrides). */
getSettings(): Settings;
/**
* Set the active tools by name. Unknown and `hidden` tools are ignored. Tools with `codemode` or
* `deferred` exposure stay callable from codemode scripts whether active or not.
*/
setActiveTools(toolNames: string[]): void;
/** Get available slash commands in the current session. */
@ -1637,6 +1813,32 @@ export interface ExtensionAPI {
*/
unregisterProvider(name: string): void;
// =========================================================================
// MCP Servers
// =========================================================================
/**
* Register an MCP server for this session, with the same config as an `mcpServers` entry in
* `mcp.json`. The server connects next to the configured servers: on `session_start` when
* registered during extension load, right away when registered later. Registering a name again
* replaces the extension's earlier registration.
*
* The registration is not saved; register again on every load. A server of the same name in
* `mcp.json` takes precedence. Throws for invalid configs and for names another extension
* registered. When no loaded extension handles MCP servers (for example because another MCP
* extension replaced the built-in one), the registration is reported as an extension error.
*
* @example
* pi.registerMcpServer("jira", { url: "https://mcp.example.com/jira" });
*/
registerMcpServer(name: string, config: McpServerConfig): void;
/** Remove an MCP server this extension registered and close its connection. */
unregisterMcpServer(name: string): void;
/** Every MCP server registered by extensions. For extensions that connect MCP servers. */
getMcpServers(): RegisteredMcpServer[];
/**
* Register a virtual model: a selectable catalog entry that routes each request to a physical
* model. The selection (`ctx.model`, `model_change` entries) names the virtual model; assistant
@ -1790,6 +1992,14 @@ export type InlineExtension =
factory: ExtensionFactory;
/** Omit this extension from the startup Extensions list. */
hidden?: boolean;
/**
* Leave this extension out when another extension registers a tool, command, or flag with a
* name it registers during loading, instead of reporting a conflict. The CLI's built-in MCP,
* codemode, and tool search extensions use it, so for example an MCP extension that registers
* `/mcp` replaces the built-in MCP support. The factory still runs, so it should only register
* tools, commands, flags, and event handlers.
*/
replaceable?: boolean;
};
// ============================================================================
@ -1838,11 +2048,16 @@ export type GetActiveToolsHandler = () => string[];
/** Tool info with name, description, parameter schema, prompt guidelines, and source metadata. */
export type ToolInfo = Pick<ToolDefinition, "name" | "description" | "parameters" | "promptGuidelines"> & {
exposure: ToolExposure;
namespace?: ToolNamespace;
annotations?: ToolAnnotations;
sourceInfo: SourceInfo;
};
export type GetAllToolsHandler = () => ToolInfo[];
export type GetSettingsHandler = () => Settings;
export type GetCommandsHandler = () => SlashCommandInfo[];
export type SetActiveToolsHandler = (toolNames: string[]) => void;
@ -1886,6 +2101,8 @@ export interface ExtensionRuntimeState {
registerProvider: (name: string, config: ProviderConfig, extensionPath?: string) => void;
registerNativeProvider: (provider: Provider, extensionPath?: string) => void;
unregisterProvider: (name: string, extensionPath?: string) => void;
/** Servers registered with `pi.registerMcpServer()`. */
mcpServers: McpServerRegistry;
registerVirtualModel: (definition: VirtualModelDefinition, extensionPath?: string) => void;
unregisterVirtualModel: (provider: string, id: string) => void;
}
@ -1903,6 +2120,7 @@ export interface ExtensionActions {
setLabel: SetLabelHandler;
getActiveTools: GetActiveToolsHandler;
getAllTools: GetAllToolsHandler;
getSettings: GetSettingsHandler;
setActiveTools: SetActiveToolsHandler;
refreshTools: RefreshToolsHandler;
getCommands: GetCommandsHandler;
@ -1928,6 +2146,15 @@ export interface ExtensionContextActions {
compact: (options?: CompactOptions) => void;
getSystemPrompt: () => string;
getSystemPromptOptions?: () => BuildSystemPromptOptions;
/** Backs `ExtensionToolContext.executeTool()`. Without it, nested calls fail. */
executeTool?: (
callerId: string,
name: string,
args: unknown,
options: ExecuteToolOptions,
) => Promise<AgentToolCallOutcome>;
/** Backs `ExtensionToolContext.tools`. */
getCallableTools?: () => readonly AgentTool[];
}
/**
@ -1967,6 +2194,8 @@ export interface Extension {
path: string;
resolvedPath: string;
hidden?: boolean;
/** See {@link InlineExtension}. */
replaceable?: boolean;
sourceInfo: SourceInfo;
handlers: Map<string, HandlerFn[]>;
tools: Map<string, RegisteredTool>;

View file

@ -12,15 +12,17 @@ import type { RegisteredTool } from "./types.ts";
/**
* Wrap a RegisteredTool into an AgentTool.
* Uses the runner's createContext() for consistent context across tools and event handlers.
* Uses the runner's createToolContext() for consistent context across tools and event handlers.
*/
export function wrapRegisteredTool(registeredTool: RegisteredTool, runner: ExtensionRunner): AgentTool {
return wrapToolDefinition(registeredTool.definition, () => runner.createContext());
return wrapToolDefinition(registeredTool.definition, (toolCallId, signal) =>
runner.createToolContext(toolCallId, signal),
);
}
/**
* Wrap all registered tools into AgentTools.
* Uses the runner's createContext() for consistent context across tools and event handlers.
* Uses the runner's createToolContext() for consistent context across tools and event handlers.
*/
export function wrapRegisteredTools(registeredTools: RegisteredTool[], runner: ExtensionRunner): AgentTool[] {
return registeredTools.map((tool) => wrapRegisteredTool(tool, runner));

View file

@ -0,0 +1,237 @@
/**
* MCP server configuration and the servers extensions register with `pi.registerMcpServer()`.
*
* The core only validates and stores registrations. The MCP extension (built in, or another
* extension that handles `mcp_servers_change`) connects them next to the servers from `mcp.json`.
*/
/**
* - `codemode`: tools are callable from codemode scripts and listed in its description, but not
* declared to the model.
* - `codemode-deferred`: like `codemode`, but not listed in the codemode description. Scripts find
* them with `searchTools()`.
* - `deferred`: not declared to the model until the `tool_search` tool loads them; the model then
* calls them directly. Does not need codemode.
* - `direct`: tools are declared to the model like any other tool (and callable from codemode).
* - `hidden`: tools are registered but unreachable.
*/
export type McpExposure = "codemode" | "codemode-deferred" | "deferred" | "direct" | "hidden";
const MCP_EXPOSURES: readonly string[] = [
"codemode",
"codemode-deferred",
"deferred",
"direct",
"hidden",
] satisfies McpExposure[];
interface McpServerConfigBase {
/** Default: `codemode`. */
exposure?: McpExposure;
/**
* Exposure of single tools, overriding `exposure`. Keys are tool names as the server offers them,
* or patterns where `*` matches any characters. An exact name wins over patterns; among patterns
* the first match in the object wins. `hidden` removes tools, so `"exposure": "hidden"` with
* overrides for a few tools exposes only those.
*/
toolExposure?: Record<string, McpExposure>;
/** Set to false to keep the entry without connecting. Default: true. */
enabled?: boolean;
/** Per-request timeout in seconds. Progress notifications from the server reset it. Default: 60. */
timeout?: number;
}
export interface McpStdioServerConfig extends McpServerConfigBase {
type?: "stdio";
command: string;
args?: string[];
/** Values may reference environment variables (`${NAME}`) or commands (`!cmd`). */
env?: Record<string, string>;
/** Relative paths resolve against the session working directory. */
cwd?: string;
}
/** OAuth client settings for servers that do not support dynamic client registration. */
export interface McpOAuthConfig {
/** Pre-registered client id. Without it, pi registers a client with the authorization server. */
clientId?: string;
/** May reference environment variables (`${NAME}`) or commands (`!cmd`). */
clientSecret?: string;
/**
* Port of the loopback callback server, for clients registered with a fixed redirect URI. Without
* `callbackUrl`, the redirect URI is `http://127.0.0.1:<port>/callback`.
*/
callbackPort?: number;
/**
* Redirect URI registered for `clientId`, for example `http://localhost:8080/oauth/callback`. It must
* be an `http` URI on `localhost`, `127.0.0.1`, or `[::1]`. Without a port, the callback server
* listens on `callbackPort` or a free port, which is added to the URI (RFC 8252).
*/
callbackUrl?: string;
/** Scopes to request, separated by spaces. Default: the scopes the server advertises. */
scope?: string;
}
const LOOPBACK_HOSTS = ["localhost", "127.0.0.1", "[::1]"];
/** Whether a redirect URI can be served by pi's loopback callback server. */
export function isLoopbackRedirectUri(value: string): boolean {
if (!URL.canParse(value)) return false;
const url = new URL(value);
return url.protocol === "http:" && LOOPBACK_HOSTS.includes(url.hostname) && url.search === "" && url.hash === "";
}
export interface McpHttpServerConfig extends McpServerConfigBase {
type?: "http";
url: string;
/** Values may reference environment variables (`${NAME}`) or commands (`!cmd`). */
headers?: Record<string, string>;
oauth?: McpOAuthConfig;
}
export type McpServerConfig = McpStdioServerConfig | McpHttpServerConfig;
const SERVER_NAME = /^[A-Za-z0-9_-]+$/;
function isRecord(value: unknown): value is Record<string, unknown> {
return typeof value === "object" && value !== null && !Array.isArray(value);
}
function isStringRecord(value: unknown): value is Record<string, string> {
return isRecord(value) && Object.values(value).every((entry) => typeof entry === "string");
}
function validateOAuth(value: unknown): string | undefined {
if (value === undefined) return undefined;
if (!isRecord(value)) return "oauth must be an object";
if (value.clientId !== undefined && typeof value.clientId !== "string") return "oauth.clientId must be a string";
if (value.clientSecret !== undefined && typeof value.clientSecret !== "string") {
return "oauth.clientSecret must be a string";
}
const port = value.callbackPort;
if (port !== undefined && (typeof port !== "number" || !Number.isInteger(port) || port < 1 || port > 65535)) {
return "oauth.callbackPort must be a port number";
}
if (value.callbackUrl !== undefined) {
if (typeof value.callbackUrl !== "string" || !isLoopbackRedirectUri(value.callbackUrl)) {
return "oauth.callbackUrl must be an http URI on localhost, 127.0.0.1, or [::1] without query or fragment";
}
const urlPort = new URL(value.callbackUrl).port;
if (urlPort && port !== undefined && Number(urlPort) !== port) {
return "oauth.callbackUrl and oauth.callbackPort name different ports";
}
}
if (value.scope !== undefined && typeof value.scope !== "string") return "oauth.scope must be a string";
return undefined;
}
function isExposure(value: unknown): value is McpExposure {
return typeof value === "string" && MCP_EXPOSURES.includes(value);
}
function toolPatternRegExp(pattern: string): RegExp {
const source = pattern
.split("*")
.map((part) => part.replace(/[.+?^${}()|[\]\\]/g, "\\$&"))
.join(".*");
return new RegExp(`^${source}$`);
}
/** Exposure of one tool of a server: its `toolExposure` entry, else the server's `exposure`. */
export function getMcpToolExposure(config: McpServerConfig, toolName: string): McpExposure {
const overrides = config.toolExposure ?? {};
const exact = overrides[toolName];
if (exact !== undefined) return exact;
for (const [pattern, exposure] of Object.entries(overrides)) {
if (pattern.includes("*") && toolPatternRegExp(pattern).test(toolName)) return exposure;
}
return config.exposure ?? "codemode";
}
/** Validate one server entry of the `mcpServers` shape. Returns the config or an error message. */
export function validateMcpServerConfig(name: string, value: unknown): McpServerConfig | string {
if (!SERVER_NAME.test(name)) return `invalid server name "${name}" (use letters, digits, "_" and "-")`;
if (!isRecord(value)) return `server "${name}" must be an object`;
const { type, exposure, enabled, timeout, toolExposure } = value;
const exposures = MCP_EXPOSURES.map((value) => `"${value}"`).join(", ");
if (exposure !== undefined && !isExposure(exposure)) {
return `server "${name}": exposure must be one of ${exposures}`;
}
if (toolExposure !== undefined) {
if (!isRecord(toolExposure)) return `server "${name}": toolExposure must map tool names to exposures`;
for (const [tool, value] of Object.entries(toolExposure)) {
if (!isExposure(value)) return `server "${name}": toolExposure "${tool}" must be one of ${exposures}`;
}
}
if (enabled !== undefined && typeof enabled !== "boolean") return `server "${name}": enabled must be a boolean`;
if (timeout !== undefined && (typeof timeout !== "number" || !(timeout > 0))) {
return `server "${name}": timeout must be a positive number of seconds`;
}
if (type === "sse") return `server "${name}": legacy SSE transport is not supported; use the streamable HTTP URL`;
if (typeof value.url === "string" && (type === undefined || type === "http" || type === "streamable-http")) {
if (!URL.canParse(value.url) || !/^https?:$/.test(new URL(value.url).protocol)) {
return `server "${name}": url must be an http or https URL`;
}
if (value.headers !== undefined && !isStringRecord(value.headers)) {
return `server "${name}": headers must map names to strings`;
}
const oauthError = validateOAuth(value.oauth);
if (oauthError) return `server "${name}": ${oauthError}`;
return value as unknown as McpHttpServerConfig;
}
if (typeof value.command === "string" && (type === undefined || type === "stdio")) {
if (
value.args !== undefined &&
!(Array.isArray(value.args) && value.args.every((arg) => typeof arg === "string"))
) {
return `server "${name}": args must be an array of strings`;
}
if (value.env !== undefined && !isStringRecord(value.env))
return `server "${name}": env must map names to strings`;
if (value.cwd !== undefined && typeof value.cwd !== "string") return `server "${name}": cwd must be a string`;
return value as unknown as McpStdioServerConfig;
}
return `server "${name}" needs either "command" (stdio) or "url" (streamable HTTP)`;
}
/** A server an extension registered with `pi.registerMcpServer()`. */
export interface RegisteredMcpServer {
name: string;
config: McpServerConfig;
/** Path of the extension that registered the server. */
extensionPath: string;
}
/** Servers registered by the extensions of one runtime. */
export class McpServerRegistry {
private readonly servers = new Map<string, RegisteredMcpServer>();
private changeListener: (() => void) | undefined;
/** Register or replace a server. The caller checks ownership. */
register(server: RegisteredMcpServer): void {
this.servers.set(server.name, server);
this.changeListener?.();
}
/** Remove a server registered by `extensionPath`. Servers of other extensions are left alone. */
unregister(name: string, extensionPath: string): void {
if (this.servers.get(name)?.extensionPath !== extensionPath) return;
this.servers.delete(name);
this.changeListener?.();
}
get(name: string): RegisteredMcpServer | undefined {
return this.servers.get(name);
}
/** Copies of the registered servers, in registration order. */
list(): RegisteredMcpServer[] {
return [...this.servers.values()].map((server) => ({ ...server, config: structuredClone(server.config) }));
}
/** Called after every change. The runner sets it when it binds, to emit `mcp_servers_change`. */
setChangeListener(listener: (() => void) | undefined): void {
this.changeListener = listener;
}
}

View file

@ -2,6 +2,7 @@ import type {
Api,
AssistantMessage,
AssistantMessageEventStream,
AuthOperationOptions,
AuthResult,
ClassifierApi,
ClassifierContext,
@ -141,6 +142,28 @@ export class ModelRegistry {
return this.runtime.complete(model, context, options);
}
/** Every known model of a type (chat, image, classifier), optionally for one provider. */
getModelsOfType<TType extends ModelType>(type: TType, provider?: string): readonly ModelTypeMap[TType][] {
return this.runtime.getModelsOfType(type, provider);
}
/** Models of a type whose provider has working credentials. */
getAvailableOfType<TType extends ModelType>(
type: TType,
provider?: string,
options?: AuthOperationOptions,
): Promise<readonly ModelTypeMap[TType][]> {
return this.runtime.getAvailableOfType(type, provider, options);
}
getModelOfType<TType extends ModelType>(
type: TType,
provider: string,
modelId: string,
): ModelTypeMap[TType] | undefined {
return this.runtime.getModelOfType(type, provider, modelId);
}
/** Classify structured state with request-time authentication. Never rejects. */
classify(
model: ClassifierModel<ClassifierApi>,

View file

@ -0,0 +1,238 @@
/**
* Tool calls that a tool makes while it runs (`ctx.executeTool()`), for example from codemode
* scripts. The agent loop does not know about them: the session runs each one through the agent's
* tool pipeline (`runToolCall`) with its own hooks, emits `tool_execution_*` events with
* `parentToolCallId`, and records the calls on the model-issued call's tool result message.
*
* Nothing here runs until a tool calls `ctx.executeTool()`.
*/
import type {
AgentTool,
AgentToolCall,
AgentToolCallOutcome,
AgentToolResult,
AgentToolUpdateCallback,
} from "@earendil-works/pi-agent-core";
import type { JsonObject, NestedToolCallRecord, NestedToolCalls, TextContent } from "@earendil-works/pi-ai";
/**
* Limits of the nested-call record on a tool result: arguments
* over the per-call or total size are omitted, calls beyond the count are dropped, and the record
* is marked incomplete when any of that happens.
*/
export const NESTED_CALL_LIMITS = {
maxCalls: 256,
maxArgumentBytesPerCall: 8 * 1024,
maxArgumentBytesTotal: 32 * 1024,
maxErrorChars: 500,
} as const;
const encoder = new TextEncoder();
/**
* Collects the nested calls of one model-issued tool call, including calls made by nested tools.
* The snapshot becomes `nestedCalls` on the tool result message.
*/
export class NestedCallRecorder {
private readonly calls: NestedToolCallRecord[] = [];
private readonly startedAt = new Map<NestedToolCallRecord, number>();
private complete = true;
private argumentBytes = 0;
/** Record a call as it starts. Returns undefined when the call is dropped. */
start(toolCall: AgentToolCall): NestedToolCallRecord | undefined {
if (this.calls.length >= NESTED_CALL_LIMITS.maxCalls) {
this.complete = false;
return undefined;
}
const record: NestedToolCallRecord = { id: toolCall.id, name: toolCall.name, status: "unfinished" };
const json = JSON.stringify(toolCall.arguments ?? {});
const bytes = encoder.encode(json).length;
if (
bytes > NESTED_CALL_LIMITS.maxArgumentBytesPerCall ||
this.argumentBytes + bytes > NESTED_CALL_LIMITS.maxArgumentBytesTotal
) {
record.argumentsBytes = bytes;
this.complete = false;
} else {
record.arguments = JSON.parse(json) as JsonObject;
this.argumentBytes += bytes;
}
this.calls.push(record);
this.startedAt.set(record, performance.now());
return record;
}
finish(record: NestedToolCallRecord | undefined, isError: boolean, errorText: string): void {
if (!record) return;
record.status = isError ? "error" : "ok";
record.durationMs = Math.round(performance.now() - (this.startedAt.get(record) ?? performance.now()));
this.startedAt.delete(record);
if (isError && errorText) record.error = errorText.slice(0, NESTED_CALL_LIMITS.maxErrorChars);
}
/** Copy of the record so far, or undefined when no nested call was made. */
snapshot(): NestedToolCalls | undefined {
if (this.calls.length === 0 && this.complete) return undefined;
const calls = this.calls.map((call) => ({ ...call }));
return { calls, complete: this.complete && calls.every((call) => call.status !== "unfinished") };
}
}
export interface NestedToolCallOptions {
/** Defaults to the calling tool's signal. */
signal?: AbortSignal;
/** Receives partial results of the nested tool, in addition to `tool_execution_update` events. */
onUpdate?: AgentToolUpdateCallback;
}
/** `tool_execution_*` events of nested calls. */
export type NestedToolExecutionEvent =
| { type: "tool_execution_start"; toolCallId: string; toolName: string; args: unknown; parentToolCallId: string }
| {
type: "tool_execution_update";
toolCallId: string;
toolName: string;
args: unknown;
partialResult: AgentToolResult<unknown>;
parentToolCallId: string;
}
| {
type: "tool_execution_end";
toolCallId: string;
toolName: string;
result: AgentToolResult<unknown>;
isError: boolean;
parentToolCallId: string;
};
export interface NestedToolCallHost {
/** Tools nested calls resolve against. */
getTools(): readonly AgentTool[];
/** Whether every nested call runs exclusively, as when the agent executes tool calls sequentially. */
isSequential(): boolean;
/** Run the call through the tool pipeline, with hooks that report `parentToolCallId`. */
runToolCall(
toolCall: AgentToolCall,
parentToolCallId: string,
signal: AbortSignal | undefined,
onUpdate: (partialResult: AgentToolResult<unknown>) => Promise<void>,
): Promise<AgentToolCallOutcome>;
emit(event: NestedToolExecutionEvent): Promise<void>;
}
/** Calls below one model-issued call share its recorder. */
interface CallScope {
recorder: NestedCallRecorder;
nextId: number;
/** Set inside a call that holds the exclusive queue, so its own nested calls do not wait on it. */
holdsQueue: boolean;
}
function textOf(result: AgentToolResult<unknown>): string {
return (result.content ?? [])
.filter((block): block is TextContent => block.type === "text")
.map((block) => block.text)
.join("\n");
}
export class NestedToolCallRunner {
private readonly host: NestedToolCallHost;
/** Scopes by the id of the calling tool call. */
private readonly scopes = new Map<string, CallScope>();
/** Serializes nested calls that must not run concurrently. */
private queueTail: Promise<void> = Promise.resolve();
constructor(host: NestedToolCallHost) {
this.host = host;
}
/**
* Run `name` on behalf of the call `callerId`. The nested call gets the id `<callerId>/<n>`.
* Never rejects for tool failures: they come back as `isError: true`.
*/
async execute(
callerId: string,
name: string,
args: unknown,
options: NestedToolCallOptions = {},
): Promise<AgentToolCallOutcome> {
let scope = this.scopes.get(callerId);
if (!scope) {
scope = { recorder: new NestedCallRecorder(), nextId: 1, holdsQueue: false };
this.scopes.set(callerId, scope);
}
const toolCall: AgentToolCall = {
type: "toolCall",
id: `${callerId}/${scope.nextId++}`,
name,
arguments: (args ?? {}) as AgentToolCall["arguments"],
};
const record = scope.recorder.start(toolCall);
await this.host.emit({
type: "tool_execution_start",
toolCallId: toolCall.id,
toolName: name,
args: toolCall.arguments,
parentToolCallId: callerId,
});
const exclusive =
!scope.holdsQueue &&
(this.host.isSequential() ||
this.host.getTools().find((tool) => tool.name === name)?.executionMode === "sequential");
let release: (() => void) | undefined;
if (exclusive) {
const previous = this.queueTail;
this.queueTail = new Promise((resolve) => {
release = resolve;
});
await previous;
}
this.scopes.set(toolCall.id, {
recorder: scope.recorder,
nextId: 1,
holdsQueue: scope.holdsQueue || exclusive,
});
let outcome: AgentToolCallOutcome;
try {
outcome = await this.host.runToolCall(toolCall, callerId, options.signal, async (partialResult) => {
options.onUpdate?.(partialResult);
await this.host.emit({
type: "tool_execution_update",
toolCallId: toolCall.id,
toolName: name,
args: toolCall.arguments,
partialResult,
parentToolCallId: callerId,
});
});
} finally {
this.scopes.delete(toolCall.id);
release?.();
}
scope.recorder.finish(record, outcome.isError, textOf(outcome.result));
await this.host.emit({
type: "tool_execution_end",
toolCallId: toolCall.id,
toolName: name,
result: outcome.result,
isError: outcome.isError,
parentToolCallId: callerId,
});
return outcome;
}
/** Remove and return the record of the nested calls a model-issued call made. */
takeRecord(toolCallId: string): NestedToolCalls | undefined {
const scope = this.scopes.get(toolCallId);
this.scopes.delete(toolCallId);
return scope?.recorder.snapshot();
}
clear(): void {
this.scopes.clear();
}
}

View file

@ -93,6 +93,21 @@ function mergeExtensionWarnings(
];
}
/**
* Leave out replaceable extensions (see `InlineExtension`) that share a tool, command, or flag name
* with another extension. For example, a third-party MCP extension that registers `/mcp` replaces
* the built-in MCP extension instead of both connecting the same servers.
*/
function omitReplacedExtensions(extensions: Extension[]): Extension[] {
const names = (extension: Extension) => [
...[...extension.tools.keys()].map((name) => `tool:${name}`),
...[...extension.commands.keys()].map((name) => `command:${name}`),
...[...extension.flags.keys()].map((name) => `flag:${name}`),
];
const taken = new Set(extensions.filter((extension) => !extension.replaceable).flatMap(names));
return extensions.filter((extension) => !extension.replaceable || !names(extension).some((name) => taken.has(name)));
}
export interface ResourceLoader {
getExtensions(): LoadExtensionsResult;
getSkills(): { skills: Skill[]; diagnostics: ResourceDiagnostic[] };
@ -630,6 +645,7 @@ export class DefaultResourceLoader implements ResourceLoader {
const inlineExtensions = await this.loadExtensionFactories(extensionsResult.runtime);
extensionsResult.extensions.push(...inlineExtensions.extensions);
extensionsResult.errors.push(...inlineExtensions.errors);
extensionsResult.extensions = omitReplacedExtensions(extensionsResult.extensions);
return extensionsResult;
}
@ -644,7 +660,10 @@ export class DefaultResourceLoader implements ResourceLoader {
if (!preTrustExtensions) {
const extensionsResult = await loadExtensionsCached(extensionPaths, this.cwd, this.eventBus);
const inlineExtensions = await this.loadExtensionFactories(extensionsResult.runtime);
extensionsResult.extensions.push(...inlineExtensions.extensions);
extensionsResult.extensions = omitReplacedExtensions([
...extensionsResult.extensions,
...inlineExtensions.extensions,
]);
extensionsResult.errors.push(...inlineExtensions.errors);
this.addExtensionConflictDiagnostics(extensionsResult);
return extensionsResult;
@ -682,7 +701,7 @@ export class DefaultResourceLoader implements ResourceLoader {
orderedExtensions.push(...inlineExtensions);
const extensionsResult: LoadExtensionsResult = {
extensions: orderedExtensions,
extensions: omitReplacedExtensions(orderedExtensions),
errors: [...preTrustExtensions.errors, ...remainingExtensions.errors],
warnings: [...(preTrustExtensions.warnings ?? []), ...(remainingExtensions.warnings ?? [])],
runtime: preTrustExtensions.runtime,
@ -1044,6 +1063,7 @@ export class DefaultResourceLoader implements ResourceLoader {
try {
const extension = await loadExtensionFromFactory(factory, this.cwd, this.eventBus, runtime, extensionPath);
extension.hidden = isNamed && input.hidden;
extension.replaceable = isNamed && input.replaceable === true;
extensions.push(extension);
} catch (error) {
const message = error instanceof Error ? error.message : "failed to load extension";

View file

@ -86,6 +86,22 @@ export interface WarningSettings {
anthropicExtraUsage?: boolean; // default: true
}
/**
* How the codemode tool presents tools while it is active.
* - `on`: declared tools that scripts can call get their codemode declaration appended to their
* description; the codemode description lists only the tools without `direct` exposure.
* - `only`: the codemode description lists every tool scripts can call, and active `direct` tools are
* not declared to the model.
*/
export type CodemodeMode = "on" | "only";
export interface CodemodeSettings {
/** Default: `on`. */
mode?: CodemodeMode;
/** Estimated tokens (characters / 4) the codemode description may spend on tool declarations. Default: 3000. */
inlineBudget?: number;
}
export type DefaultProjectTrust = "ask" | "always" | "never";
export type TransportSetting = Transport;
@ -151,6 +167,7 @@ export interface Settings {
showHardwareCursor?: boolean; // Show terminal cursor while still positioning it for IME
markdown?: MarkdownSettings;
warnings?: WarningSettings;
codemode?: CodemodeSettings;
sessionDir?: string; // Custom session storage directory (same format as --session-dir CLI flag)
httpProxy?: string; // Proxy URL applied as HTTP_PROXY and HTTPS_PROXY for Pi-managed HTTP clients
httpIdleTimeoutMs?: number; // HTTP header/body idle timeout in milliseconds; 0 disables it
@ -499,6 +516,11 @@ export class SettingsManager {
return settings as Settings;
}
/** A copy of the effective settings: global and project settings merged, with overrides. */
getSettings(): Settings {
return structuredClone(this.settings);
}
getGlobalSettings(): Settings {
return structuredClone(this.globalSettings);
}

View file

@ -155,7 +155,7 @@ export function buildSystemPromptSections(input: BuildSystemPromptOptions): Syst
- Additional docs: ${getDocsPath()}
- Examples: ${getExamplesPath()} (extensions, custom tools, SDK)
- When reading pi docs or examples, resolve docs/... under Additional docs and examples/... under Examples, not the current working directory
- When asked about: extensions (docs/extensions.md, examples/extensions/), themes (docs/themes.md), skills (docs/skills.md), prompt templates (docs/prompt-templates.md), TUI components (docs/tui.md), keybindings (docs/keybindings.md), SDK integrations (docs/sdk.md), custom providers (docs/custom-provider.md), adding models (docs/models.md), pi packages (docs/packages.md), environment variables (docs/environment-variables.md)
- When asked about: extensions (docs/extensions.md, examples/extensions/), themes (docs/themes.md), skills (docs/skills.md), prompt templates (docs/prompt-templates.md), TUI components (docs/tui.md), keybindings (docs/keybindings.md), SDK integrations (docs/sdk.md), custom providers (docs/custom-provider.md), adding models (docs/models.md), pi packages (docs/packages.md), environment variables (docs/environment-variables.md), MCP servers (docs/mcp.md)
- When working on pi topics, read the docs and examples, and follow .md cross-references before implementing
- Always read pi .md files completely and follow links to related docs (e.g., tui.md for TUI API details)`;
}

View file

@ -47,6 +47,17 @@ export const bashToolSystemPromptContribution = {
export type BashToolInput = Static<typeof bashSchema>;
/**
* Result for programmatic callers such as codemode scripts. A non-zero exit code is an error result for the model, but scripts still resolve to this value.
*/
const bashOutputSchema = Type.Object({
output: Type.String({ description: "Combined stdout and stderr, truncated like the model-facing output" }),
exit_code: Type.Number(),
wall_time_seconds: Type.Number(),
});
export type BashToolOutput = Static<typeof bashOutputSchema>;
export interface BashToolDetails {
truncation?: TruncationResult;
fullOutputPath?: string;
@ -240,6 +251,7 @@ export function createShellToolDefinition(
promptSnippet: config.promptSnippet,
promptGuidelines: exposeSessionEnvironment && config.promptGuidelines ? [...config.promptGuidelines] : undefined,
parameters: bashSchema,
outputSchema: bashOutputSchema,
constrainedSampling: { type: "json_schema", strict: "prefer" },
async execute(
_toolCallId,
@ -339,6 +351,7 @@ export function createShellToolDefinition(
};
const appendStatus = (text: string, status: string) => `${text ? `${text}\n\n` : ""}${status}`;
const startedAt = performance.now();
try {
let exitCode: number | null;
@ -368,10 +381,20 @@ export function createShellToolDefinition(
if (exitCode === null) {
throw new Error(appendStatus(outputText, "Command terminated without an exit code"));
}
const structuredContent: BashToolOutput = {
output: outputText,
exit_code: exitCode,
wall_time_seconds: Math.round((performance.now() - startedAt) / 100) / 10,
};
if (exitCode !== 0) {
throw new Error(appendStatus(outputText, `Command exited with code ${exitCode}`));
return {
content: [{ type: "text", text: appendStatus(outputText, `Command exited with code ${exitCode}`) }],
details,
structuredContent,
isError: true,
};
}
return { content: [{ type: "text", text: outputText }], details };
return { content: [{ type: "text", text: outputText }], details, structuredContent };
} finally {
clearUpdateTimer();
}

View file

@ -1,28 +1,38 @@
import type { AgentTool } from "@earendil-works/pi-agent-core";
import type { ExtensionContext, ToolDefinition } from "../extensions/types.ts";
import type { ExtensionToolContext, ToolDefinition } from "../extensions/types.ts";
/** Creates the context for one tool call. */
export type ToolContextFactory = (toolCallId: string, signal: AbortSignal | undefined) => ExtensionToolContext;
/** Wrap a ToolDefinition into an AgentTool for the core runtime. */
export function wrapToolDefinition<TDetails = unknown>(
definition: ToolDefinition<any, TDetails>,
ctxFactory?: () => ExtensionContext,
ctxFactory?: ToolContextFactory,
): AgentTool<any, TDetails> {
return {
name: definition.name,
label: definition.label,
description: definition.description,
parameters: definition.parameters,
outputSchema: definition.outputSchema,
constrainedSampling: definition.constrainedSampling,
prepareArguments: definition.prepareArguments,
executionMode: definition.executionMode,
execute: (toolCallId, params, signal, onUpdate, ctx?: ExtensionContext) =>
definition.execute(toolCallId, params, signal, onUpdate, ctx ?? (ctxFactory?.() as ExtensionContext)),
execute: (toolCallId, params, signal, onUpdate, ctx?: ExtensionToolContext) =>
definition.execute(
toolCallId,
params,
signal,
onUpdate,
ctx ?? (ctxFactory?.(toolCallId, signal) as ExtensionToolContext),
),
};
}
/** Wrap multiple ToolDefinitions into AgentTools for the core runtime. */
export function wrapToolDefinitions(
definitions: ToolDefinition<any, any>[],
ctxFactory?: () => ExtensionContext,
ctxFactory?: ToolContextFactory,
): AgentTool<any>[] {
return definitions.map((definition) => wrapToolDefinition(definition, ctxFactory));
}
@ -39,6 +49,7 @@ export function createToolDefinitionFromAgentTool(tool: AgentTool<any>): ToolDef
label: tool.label,
description: tool.description,
parameters: tool.parameters as any,
outputSchema: tool.outputSchema,
constrainedSampling: tool.constrainedSampling,
prepareArguments: tool.prepareArguments,
executionMode: tool.executionMode,

View file

@ -274,3 +274,41 @@ export function truncateLine(
}
return { text: `${line.slice(0, maxChars)}... [truncated]`, wasTruncated: true };
}
export interface MiddleTruncationResult {
/** The start and end of the content with a `…N chars truncated…` marker between them. */
content: string;
truncated: boolean;
/** Characters left out. */
removedChars: number;
totalBytes: number;
totalLines: number;
}
/**
* Keep the start and the end of `content`, half of `maxBytes` each, and replace the middle with a
* `…N chars truncated…` marker, like Codex does for tool output. Cuts only at character boundaries.
*/
export function truncateMiddle(content: string, maxBytes: number): MiddleTruncationResult {
const buf = Buffer.from(content, "utf-8");
const totalLines = splitLinesForCounting(content).length;
if (buf.length <= maxBytes) {
return { content, truncated: false, removedChars: 0, totalBytes: buf.length, totalLines };
}
// Continuation bytes (10xxxxxx) are not character starts.
const isBoundary = (index: number) => index >= buf.length || (buf[index] & 0xc0) !== 0x80;
let headEnd = Math.floor(maxBytes / 2);
while (headEnd > 0 && !isBoundary(headEnd)) headEnd--;
let tailStart = buf.length - (maxBytes - Math.floor(maxBytes / 2));
while (tailStart < buf.length && !isBoundary(tailStart)) tailStart++;
const head = buf.subarray(0, headEnd).toString("utf-8");
const tail = buf.subarray(tailStart).toString("utf-8");
const removedChars = Array.from(buf.subarray(headEnd, tailStart).toString("utf-8")).length;
return {
content: `${head}…${removedChars} chars truncated…${tail}`,
truncated: true,
removedChars,
totalBytes: buf.length,
totalLines,
};
}

View file

@ -29,6 +29,7 @@ type TrustFile = Record<string, boolean | null | undefined>;
const TRUST_REQUIRING_PROJECT_CONFIG_RESOURCES = [
"settings.json",
"mcp.json",
"extensions",
"skills",
"prompts",

View file

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

View file

@ -0,0 +1,428 @@
/**
* Runs one codemode script in the sandbox. Split from tool.ts and loaded through
* execute.lazy.ts so the sandbox runtime only loads when a script runs.
*/
import { randomBytes } from "node:crypto";
import { writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { AgentTool, AgentToolCallOutcome, AgentToolResult } from "@earendil-works/pi-agent-core";
import type { AnyModel, ClassifierContext, ImageContent, ModelType, TextContent } from "@earendil-works/pi-ai";
import {
type CodemodeResult,
CodemodeSandbox,
type CodemodeTool,
loadQuickJSWasm,
parseCodemodeSource,
renderToolSample,
toCodemodeIdentifier,
} from "@earendil-works/pi-codemode";
import { getCodemodeWorkerUrl, getQuickJSWasmPath } from "../../config.ts";
import type { ExtensionToolContext } from "../../core/extensions/types.ts";
import type { SessionEntry } from "../../core/session-manager.ts";
import { Bm25Ranker, createToolSearchDocument, DEFAULT_TOOL_SEARCH_LIMIT } from "../tool-search/tool.ts";
import {
CODEMODE_STORE_ENTRY_TYPE,
type CodemodeModelRuntime,
type CodemodeNestedCall,
type CodemodeStoreEntryData,
type CodemodeToolDetails,
type CodemodeToolInput,
type CodemodeToolOptions,
getCodemodeCallableTools,
MODEL_GLOBAL_DECLARATIONS,
toCodemodeDeclaration,
} from "./tool.ts";
const ARGS_PREVIEW_CHARS = 200;
const ERROR_PREVIEW_CHARS = 500;
/** Classifier calls one script may have in flight; `Promise.all` over many items queues the rest. */
const MAX_CONCURRENT_MODEL_CALLS = 4;
/**
* Heap limit for the QuickJS VM. The worker shares pi's process, so without a limit a runaway
* script can grow to wasm32's 4 GiB and take the session down. Overruns throw
* `InternalError: out of memory` inside the script.
*/
const CODEMODE_MEMORY_LIMIT_BYTES = 256 * 1024 * 1024;
const MODEL_TYPES: ReadonlySet<string> = new Set<ModelType>(["chat", "image", "classifier"]);
function truncateText(text: string, maxChars: number): string {
return text.length > maxChars ? `${text.slice(0, maxChars - 3)}...` : text;
}
function previewArgs(args: unknown): string {
if (args === undefined) return "";
try {
return truncateText(JSON.stringify(args) ?? "", ARGS_PREVIEW_CHARS);
} catch {
return "";
}
}
function textOf(result: AgentToolResult<unknown>): string {
return (result.content ?? [])
.filter((block): block is TextContent => block.type === "text")
.map((block) => block.text)
.join("\n");
}
function toModelType(value: unknown): ModelType {
if (typeof value === "string" && MODEL_TYPES.has(value)) return value as ModelType;
throw new Error(`Unknown model type ${JSON.stringify(value)}. Use "chat", "image", or "classifier".`);
}
function toProvider(value: unknown): string | undefined {
if (value === undefined || value === null) return undefined;
if (typeof value !== "string") throw new Error("provider must be a string");
return value;
}
/** Catalog entry for scripts. `headers` is dropped because models.json headers can carry credentials. */
function toModelInfo(model: AnyModel): Record<string, unknown> {
const info: Record<string, unknown> = { ...model };
delete info.headers;
return info;
}
/** Runs at most `limit` calls at once, in call order. */
function createLimiter(limit: number): <T>(run: () => Promise<T>) => Promise<T> {
let active = 0;
const waiting: (() => void)[] = [];
return async (run) => {
if (active >= limit) await new Promise<void>((resolve) => waiting.push(resolve));
active++;
try {
return await run();
} finally {
active--;
waiting.shift()?.();
}
};
}
function isStoreEntryData(data: unknown): data is CodemodeStoreEntryData {
if (typeof data !== "object" || data === null) return false;
const { set, delete: deleted } = data as Partial<CodemodeStoreEntryData>;
return (
typeof set === "object" &&
set !== null &&
Array.isArray(deleted) &&
deleted.every((key: unknown) => typeof key === "string")
);
}
/** Values of `load()`: the `codemode-store` entries on the branch, applied from the root. */
export function readCodemodeStore(branch: readonly SessionEntry[]): Record<string, unknown> {
const store = new Map<string, unknown>();
for (const entry of branch) {
if (entry.type !== "custom" || entry.customType !== CODEMODE_STORE_ENTRY_TYPE || !isStoreEntryData(entry.data)) {
continue;
}
for (const key of entry.data.delete) store.delete(key);
for (const [key, value] of Object.entries(entry.data.set)) store.set(key, value);
}
return Object.fromEntries(store);
}
/** Default token budget for script output. */
const DEFAULT_MAX_OUTPUT_TOKENS = 10_000;
/** Characters per token when estimating. */
const CHARS_PER_TOKEN = 4;
/** Like the script's `text()`: strings as is, other values as compact JSON. */
function valueText(value: unknown): string {
if (typeof value === "string") return value;
return JSON.stringify(value) ?? String(value);
}
function formatCallSummary(calls: readonly CodemodeNestedCall[]): string {
if (calls.length === 0) return "No tool calls were made.";
return `Tool calls made before the failure (they are not undone): ${calls.map((call) => `${call.name} (${call.status})`).join(", ")}`;
}
function formatError(result: Extract<CodemodeResult, { ok: false }>, calls: readonly CodemodeNestedCall[]): string {
const { error } = result;
const head =
error.kind === "script"
? (error.stack ?? `${error.name ?? "Error"}: ${error.message}`)
: error.kind === "timeout"
? `Script timed out: ${error.message}`
: error.kind === "aborted"
? `Script aborted: ${error.message}`
: `Script sandbox failed: ${error.message}`;
return `${head}\n\n${formatCallSummary(calls)}`;
}
/** Write the full text output to a temp file, like bash does for truncated output. */
async function spillOutput(text: string): Promise<{ path: string } | { error: string }> {
const path = join(tmpdir(), `pi-codemode-${randomBytes(8).toString("hex")}.txt`);
try {
await writeFile(path, text);
return { path };
} catch (error) {
return { error: error instanceof Error ? error.message : String(error) };
}
}
/**
* Apply the token budget: when the combined text exceeds it, the text items become one
* item that keeps the start and end of the text, and images follow it. The full text is written to
* a temp file.
*/
async function truncateOutput(
items: (TextContent | ImageContent)[],
maxTokens: number,
): Promise<{ items: (TextContent | ImageContent)[]; fullOutputPath?: string }> {
const texts = items.filter((item): item is TextContent => item.type === "text").map((item) => item.text);
const combined = texts.join("\n");
const budget = maxTokens * CHARS_PER_TOKEN;
if (texts.length === 0 || combined.length <= budget) return { items };
const headChars = Math.floor(budget / 2);
const tailChars = budget - headChars;
const removed = combined.length - headChars - tailChars;
const head = combined.slice(0, headChars);
const tail = tailChars > 0 ? combined.slice(-tailChars) : "";
let text = `Warning: truncated output (original token count: ${Math.ceil(combined.length / CHARS_PER_TOKEN)})\nTotal output lines: ${combined.split("\n").length}\n\n${head}…${Math.ceil(removed / CHARS_PER_TOKEN)} tokens truncated…${tail}`;
const spilled = await spillOutput(combined);
text +=
"path" in spilled
? `\n\n[Full output: ${spilled.path} (read with offset/limit)]`
: `\n\n[Could not save the full output: ${spilled.error}]`;
return {
items: [{ type: "text", text }, ...items.filter((item) => item.type === "image")],
...("path" in spilled ? { fullOutputPath: spilled.path } : {}),
};
}
/**
* The value a script receives for a nested call: a tool that declares
* `outputSchema` resolves to its `structuredContent`, also for error results that carry one (such
* as MCP results with `isError`); any other tool resolves to its text content. Other failures
* reject with the tool's error text.
*/
function toScriptValue(tool: AgentTool<any>, outcome: AgentToolCallOutcome): unknown {
const { result } = outcome;
if (tool.outputSchema && result.structuredContent !== undefined) return result.structuredContent;
const text = textOf(result);
if (outcome.isError) throw new Error(text || `Tool "${tool.name}" failed`);
return text;
}
/**
* Run one script. Without a session context (a plain Agent or a direct call) scripts cannot call
* tools, `store()` starts empty, and writes are dropped.
*/
export async function executeCodemode(
toolCallId: string,
input: CodemodeToolInput,
signal: AbortSignal | undefined,
onUpdate: ((result: AgentToolResult<CodemodeToolDetails>) => void) | undefined,
ctx: ExtensionToolContext | undefined,
options: CodemodeToolOptions = {},
): Promise<AgentToolResult<CodemodeToolDetails>> {
const startedAt = performance.now();
const { code, options: sourceOptions } = parseCodemodeSource(input.code);
const calls: CodemodeNestedCall[] = [];
const snapshot = (): CodemodeToolDetails => ({ calls: calls.map((call) => ({ ...call })) });
const publish = () => onUpdate?.({ content: [], details: snapshot() });
const callable = ctx ? getCodemodeCallableTools(ctx.tools) : [];
// ALL_TOOLS entries carry the declaration.
const samples = new Map(callable.map((tool) => [tool.name, renderToolSample(toCodemodeDeclaration(tool))]));
const sandboxTools: CodemodeTool[] = callable.map((tool) => ({
name: tool.name,
description: samples.get(tool.name),
execute: async (args, { signal: callSignal }) => {
const record: CodemodeNestedCall = {
id: `${toolCallId}/?`,
name: tool.name,
args: previewArgs(args),
status: "running",
};
calls.push(record);
publish();
const callStartedAt = performance.now();
// Only tools from ctx.tools are callable, so ctx is set here.
if (!ctx) throw new Error("Tool calls need a session");
const outcome = await ctx.executeTool(tool.name, args, { signal: callSignal });
record.id = outcome.toolCall.id;
record.durationMs = performance.now() - callStartedAt;
if (outcome.isError) {
record.status = callSignal.aborted ? "cancelled" : "error";
record.error = truncateText(textOf(outcome.result) || `Tool "${tool.name}" failed`, ERROR_PREVIEW_CHARS);
} else {
record.status = "ok";
}
publish();
return toScriptValue(tool, outcome);
},
}));
const sandbox = new CodemodeSandbox({
tools: sandboxTools,
globals: [
...createDiscoveryGlobals(callable, samples, options),
...(options.models && ctx ? createModelGlobals(ctx.modelRegistry, toolCallId, calls, publish) : []),
],
timeoutMs: sourceOptions.timeoutMs ?? Number.POSITIVE_INFINITY,
memoryLimitBytes: CODEMODE_MEMORY_LIMIT_BYTES,
wasm: loadQuickJSWasm(getQuickJSWasmPath()),
workerUrl: getCodemodeWorkerUrl(),
});
let result: CodemodeResult;
try {
const store = ctx ? readCodemodeStore(ctx.sessionManager.getBranch()) : {};
result = await sandbox.execute(code, { signal, store });
} finally {
await sandbox.close();
}
// Calls still marked running were cut off by the script ending, a timeout, or an abort.
for (const call of calls) {
if (call.status === "running") call.status = "cancelled";
}
const items: (TextContent | ImageContent)[] = result.output.map((item) =>
item.type === "text" ? { type: "text", text: item.text } : item,
);
if (result.ok) {
const { set, delete: deleted } = result.storeWrites;
if (Object.keys(set).length > 0 || deleted.length > 0) {
options.appendEntry?.(CODEMODE_STORE_ENTRY_TYPE, { set, delete: deleted });
}
// pi extension: a returned value is appended like text().
if (result.value !== undefined) items.push({ type: "text", text: valueText(result.value) });
} else {
items.push({ type: "text", text: `Script error:\n${formatError(result, calls)}` });
}
const truncated = await truncateOutput(items, sourceOptions.maxOutputTokens ?? DEFAULT_MAX_OUTPUT_TOKENS);
const wallTime = ((performance.now() - startedAt) / 1000).toFixed(1);
const header = `${result.ok ? "Script completed" : "Script failed"}\nWall time ${wallTime} seconds\nOutput:\n`;
const details = snapshot();
if (truncated.fullOutputPath) details.fullOutputPath = truncated.fullOutputPath;
return {
content: [{ type: "text", text: header }, ...truncated.items],
details,
...(result.ok ? {} : { isError: true }),
};
}
/** `searchTools()` and `describeTool()`: ranked search and lookup over the script's nested tools. */
function createDiscoveryGlobals(
tools: readonly AgentTool<any>[],
samples: ReadonlyMap<string, string>,
options: CodemodeToolOptions,
): CodemodeTool[] {
const ranker = new Bm25Ranker();
const entry = (name: string) => ({ name: toCodemodeIdentifier(name), description: samples.get(name) ?? "" });
return [
{
name: "searchTools",
spread: true,
execute: (args) => {
const [query, searchOptions] = args as [unknown, { limit?: unknown; namespace?: unknown } | undefined];
if (typeof query !== "string") throw new Error("searchTools() expects a query string");
const limit = searchOptions?.limit ?? DEFAULT_TOOL_SEARCH_LIMIT;
if (typeof limit !== "number" || !Number.isInteger(limit) || limit <= 0) {
throw new Error("searchTools() limit must be a positive integer");
}
const namespace = searchOptions?.namespace;
if (namespace !== undefined && namespace !== null && typeof namespace !== "string") {
throw new Error("searchTools() namespace must be a string");
}
const documents = tools.flatMap((tool) => {
const toolNamespace = options.getToolNamespace?.(tool.name);
if (namespace && toolNamespace?.name !== namespace) return [];
return [createToolSearchDocument(tool, toolNamespace)];
});
return ranker.rank(query, documents, limit).map((match) => entry(match.name));
},
},
{
name: "describeTool",
spread: true,
execute: (args) => {
const [name] = args as unknown[];
if (typeof name !== "string") throw new Error("describeTool() expects a tool name");
const tool = tools.find(
(candidate) => candidate.name === name || toCodemodeIdentifier(candidate.name) === name,
);
return tool ? samples.get(tool.name) : undefined;
},
},
];
}
/**
* `models.*` for scripts: the model registry methods declared in {@link MODEL_GLOBAL_DECLARATIONS}.
* Classifier calls appear as nested call rows so the renderer shows them.
*/
function createModelGlobals(
models: CodemodeModelRuntime,
toolCallId: string,
calls: CodemodeNestedCall[],
publish: () => void,
): CodemodeTool[] {
const limit = createLimiter(MAX_CONCURRENT_MODEL_CALLS);
let classifyCount = 0;
const implementations: Record<string, CodemodeTool["execute"]> = {
"models.getModelsOfType": (args) => {
const [type, provider] = args as unknown[];
return models.getModelsOfType(toModelType(type), toProvider(provider)).map(toModelInfo);
},
"models.getAvailableOfType": async (args, { signal }) => {
const [type, provider] = args as unknown[];
const available = await models.getAvailableOfType(toModelType(type), toProvider(provider), { signal });
return available.map(toModelInfo);
},
"models.getModelOfType": (args) => {
const [type, provider, id] = args as unknown[];
if (typeof provider !== "string" || typeof id !== "string") {
throw new Error("models.getModelOfType() expects a type, a provider, and an id");
}
const model = models.getModelOfType(toModelType(type), provider, id);
return model === undefined ? undefined : toModelInfo(model);
},
"models.classify": async (args, { signal }) => {
const [model, context] = args as unknown[];
const ref = model as { provider?: unknown; id?: unknown } | null;
if (
typeof ref !== "object" ||
ref === null ||
typeof ref.provider !== "string" ||
typeof ref.id !== "string"
) {
throw new Error(
"models.classify() expects a model from models.getModelOfType() or models.getAvailableOfType()",
);
}
// Only provider and id count. A script-supplied baseUrl or headers must never receive the credentials.
const resolved = models.getModelOfType("classifier", ref.provider, ref.id);
if (!resolved) throw new Error(`Unknown classifier model "${ref.provider}/${ref.id}"`);
const record: CodemodeNestedCall = {
id: `${toolCallId}/models.classify/${++classifyCount}`,
name: "models.classify",
args: `${resolved.provider}/${resolved.id}`,
status: "running",
};
calls.push(record);
publish();
const startedAt = performance.now();
const result = await limit(() => models.classify(resolved, context as ClassifierContext, { signal }));
record.durationMs = performance.now() - startedAt;
record.status = result.stopReason === "stop" ? "ok" : result.stopReason === "aborted" ? "cancelled" : "error";
if (result.errorMessage) record.error = truncateText(result.errorMessage, ERROR_PREVIEW_CHARS);
publish();
return result;
},
};
return MODEL_GLOBAL_DECLARATIONS.map((declaration) => ({
name: declaration.name,
spread: true,
execute: implementations[declaration.name],
}));
}

View file

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

View file

@ -0,0 +1,114 @@
/**
* Presentation for the codemode tool.
*
* The call shows the script; the result lists the nested tool calls with their status as they
* run, followed by the script output without the "Script completed" header. Nested calls are not
* separate tool rows because they never reach the model as tool calls.
*/
import { Text } from "@earendil-works/pi-tui";
import type { ToolDefinition } from "../../core/extensions/types.ts";
import { getTextOutput, replaceTabs, str } from "../../core/tools/render-utils.ts";
import { keyHint } from "../../modes/interactive/components/keybinding-hints.ts";
import { highlightCode, type Theme } from "../../modes/interactive/theme/theme.ts";
import type { CodemodeNestedCall, CodemodeToolDetails } from "./tool.ts";
const CODE_PREVIEW_LINES = 10;
const CALL_PREVIEW_COUNT = 8;
const OUTPUT_PREVIEW_LINES = 5;
const COLLAPSED_ARGS_CHARS = 80;
const SCRIPT_HEADER = /^Script (completed|failed)\nWall time [\d.]+ seconds\nOutput:\n$/;
function expandHint(theme: Theme, hidden: number, noun: string): string {
return `${theme.fg("muted", `... (${hidden} more ${noun},`)} ${keyHint("app.tools.expand", "to expand")}${theme.fg("muted", ")")}`;
}
function formatDuration(ms: number | undefined): string {
if (ms === undefined) return "";
return ms < 1000 ? `${Math.round(ms)}ms` : `${(ms / 1000).toFixed(1)}s`;
}
function statusIcon(call: CodemodeNestedCall, theme: Theme): string {
switch (call.status) {
case "running":
return theme.fg("warning", "…");
case "ok":
return theme.fg("success", "✓");
case "error":
return theme.fg("error", "✗");
case "cancelled":
return theme.fg("muted", "⊘");
}
}
function formatCall(call: CodemodeNestedCall, theme: Theme, expanded: boolean): string {
const args =
!expanded && call.args.length > COLLAPSED_ARGS_CHARS
? `${call.args.slice(0, COLLAPSED_ARGS_CHARS - 3)}...`
: call.args;
const duration = formatDuration(call.durationMs);
let line = `${statusIcon(call, theme)} ${theme.fg("toolTitle", call.name)}`;
if (args) line += ` ${theme.fg("muted", args)}`;
if (duration) line += ` ${theme.fg("dim", duration)}`;
if (expanded && call.error) line += `\n ${theme.fg("error", call.error.split("\n").join("\n "))}`;
return line;
}
export const codemodeRenderers: Pick<
ToolDefinition<any, CodemodeToolDetails | undefined>,
"renderCall" | "renderResult"
> = {
renderCall(args, theme, context) {
// The code includes the `// @options:` line, so options show as part of the script.
const code = str((args as { code?: unknown } | undefined)?.code);
let text = theme.fg("toolTitle", theme.bold("codemode"));
if (code === null) {
text += ` ${theme.fg("error", "[invalid arg]")}`;
} else if (code) {
const lines = highlightCode(replaceTabs(code.replace(/\r/g, "").trimEnd()), "javascript");
const shown = context.expanded ? lines : lines.slice(0, CODE_PREVIEW_LINES);
text += `\n${shown.join("\n")}`;
if (shown.length < lines.length) text += `\n${expandHint(theme, lines.length - shown.length, "lines")}`;
}
const component = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0);
component.setText(text);
return component;
},
renderResult(result, options, theme, context) {
const sections: string[] = [];
const calls = result.details?.calls ?? [];
if (calls.length > 0) {
const shown = options.expanded ? calls : calls.slice(-CALL_PREVIEW_COUNT);
const lines = shown.map((call) => formatCall(call, theme, options.expanded));
if (shown.length < calls.length) {
lines.unshift(
`${theme.fg("muted", `... (${calls.length - shown.length} earlier calls,`)} ${keyHint("app.tools.expand", "to expand")}${theme.fg("muted", ")")}`,
);
}
sections.push(lines.join("\n"));
}
// Drop the "Script completed\nWall time ...\nOutput:\n" header. Rejected input (invalid options)
// has no header.
const [first, ...rest] = result.content;
const hasHeader = first?.type === "text" && SCRIPT_HEADER.test(first.text);
const output = options.isPartial
? ""
: getTextOutput({ ...result, content: hasHeader ? rest : result.content }, context.showImages).trim();
if (output) {
const color = context.isError ? "error" : "toolOutput";
const lines = replaceTabs(output).split("\n");
const shown = options.expanded ? lines : lines.slice(0, OUTPUT_PREVIEW_LINES);
let text = shown.map((line) => theme.fg(color, line)).join("\n");
if (shown.length < lines.length) text += `\n${expandHint(theme, lines.length - shown.length, "lines")}`;
// The collapsed preview hides the truncation notice at the end, so name the file here.
const fullOutputPath = result.details?.fullOutputPath;
if (fullOutputPath && !options.expanded) text += `\n${theme.fg("muted", `Full output: ${fullOutputPath}`)}`;
sections.push(text);
}
const component = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0);
component.setText(sections.length > 0 ? `\n${sections.join("\n\n")}` : "");
return component;
},
};

View file

@ -0,0 +1,451 @@
/**
* The `codemode` tool: the model writes JavaScript that calls other tools. Scripts use `tools`,
* `ALL_TOOLS`, `text()`, `image()`, `exit()`, `store()`/`load()`, `console.*`, and `return <value>`,
* may start with a `// @options:` line, and reach the model catalog and classifiers through
* `models.*`. Results start with a "Script completed" or "Script failed" header.
*
* Scripts can call the agent loop's nested tools: active `direct` tools and every `codemode` or
* `deferred` tool. Nested calls run through the agent loop's tool pipeline (`ctx.executeTool`), so
* validation, `tool_call`/`tool_result` hooks, and permission checks apply exactly as for direct
* calls. Only the script's output reaches the model; nested results do not.
*
* Nested results are handed to the script as follows:
* - A tool that declares `outputSchema` resolves to its `structuredContent`, also for error
* results that carry one (MCP tools resolve to their `CallToolResult`, including `isError`).
* - Any other tool resolves to its text content as one string.
* - A failed, blocked, or invalid call rejects with an Error carrying the tool's error text.
*
* A script that fails returns a normal error result that keeps its partial output, followed by
* "Script error:" and the error. `store(key, value)` and `load(key)` keep JSON values across
* calls; successful scripts append their writes to the session as `codemode-store` custom entries,
* so each branch sees the values written on its own path.
*/
import type { AgentTool } from "@earendil-works/pi-agent-core";
import type { CodemodeJsonSchema, CodemodeTool } from "@earendil-works/pi-codemode";
import {
MCP_TYPESCRIPT_PREAMBLE,
mcpStructuredContentSchema,
renderDeclarations,
renderToolSample,
toCodemodeIdentifier,
} from "@earendil-works/pi-codemode/declarations";
import { CODEMODE_SOURCE_GRAMMAR } from "@earendil-works/pi-codemode/source";
import { type Static, Type } from "typebox";
import type {
ToolDefinition,
ToolInfo,
ToolLoadout,
ToolLoadoutChanges,
ToolNamespace,
} from "../../core/extensions/types.ts";
import type { ModelRegistry } from "../../core/model-registry.ts";
import type { CodemodeMode } from "../../core/settings-manager.ts";
import { wrapToolDefinition } from "../../core/tools/tool-definition-wrapper.ts";
import { loadCodemodeExecutor } from "./execute.lazy.ts";
import { codemodeRenderers } from "./renderer.ts";
export const CODEMODE_TOOL_NAME = "codemode";
/** Custom entry type holding one script's `store()` writes: {@link CodemodeStoreEntryData}. */
export const CODEMODE_STORE_ENTRY_TYPE = "codemode-store";
export interface CodemodeStoreEntryData {
set: Record<string, unknown>;
delete: string[];
}
/** The part of the model registry that scripts reach through `models`. */
export type CodemodeModelRuntime = Pick<
ModelRegistry,
"getModelsOfType" | "getAvailableOfType" | "getModelOfType" | "classify"
>;
export interface CodemodeToolOptions {
/** Namespace of a tool, for `searchTools()` ranking and its `namespace` filter. */
getToolNamespace?: (toolName: string) => ToolNamespace | undefined;
/**
* Expose the `models` namespace to scripts, backed by the session's model registry
* (`ctx.modelRegistry`). Without it, `models` is not declared.
*/
models?: boolean;
/**
* Persists `store()` writes as a session custom entry. Without it, writes last only for the
* current script; `load()` still reads entries already on the branch.
*/
appendEntry?: (customType: string, data: CodemodeStoreEntryData) => void;
/** How the tool presents the loadout while active (the `codemode.mode` setting). Default: `on`. */
getMode?: () => CodemodeMode;
/** Token budget for tool declarations in the description. Default: {@link DEFAULT_CODEMODE_INLINE_BUDGET}. */
getInlineBudget?: () => number | undefined;
}
const TEXT_OUTPUT_SCHEMA: CodemodeJsonSchema = { type: "string" };
export const codemodeSchema = Type.Object({
code: Type.String({
description:
'Raw JavaScript source. Top-level await and return work. May start with a `// @options: {"max_output_tokens": 1000}` line.',
}),
});
export type CodemodeToolInput = Static<typeof codemodeSchema>;
/**
* Whether a registered tool is this package's `codemode` tool rather than another extension's tool
* with the same name. Compares the parameter schema, which the definition passes through by reference.
*/
export function isCodemodeTool(tool: Pick<ToolInfo, "name" | "parameters">): boolean {
return tool.name === CODEMODE_TOOL_NAME && tool.parameters === codemodeSchema;
}
export type CodemodeNestedCallStatus = "running" | "ok" | "error" | "cancelled";
export interface CodemodeNestedCall {
/** Tool call id of the nested call, `<codemode call id>/<n>`. */
id: string;
name: string;
/** Compact JSON of the arguments, truncated for display. */
args: string;
status: CodemodeNestedCallStatus;
durationMs?: number;
/** Error text, truncated for display. */
error?: string;
}
export interface CodemodeToolDetails {
calls: CodemodeNestedCall[];
/** Temp file with the full text output, when the output was truncated. */
fullOutputPath?: string;
}
export const codemodeToolSystemPromptContribution = {
snippet: "Run JavaScript that calls other tools (chains, loops, Promise.all, filtering large results)",
guidelines: [
"Use codemode to batch or chain several tool calls, or to filter large tool output down to what you need, instead of issuing many individual tool calls. Batch independent calls in one codemode call using await Promise.allSettled([...]).",
],
} as const;
const DESCRIPTION_INTRO = `Run JavaScript code to orchestrate/compose tool calls
- Evaluates the provided JavaScript code in a fresh QuickJS sandbox as the body of an async function: top-level \`await\` and \`return\` work.
- All nested tools are available on the global \`tools\` object, for example \`await tools.read(...)\`. Tool names are exposed as normalized JavaScript identifiers, for example \`await tools.mcp__ologs__get_profile(...)\`.
- Nested tool methods take an object as their input argument.
- Nested tools return either an object or a string, based on the description.
- A nested tool call that fails, is blocked, or gets invalid arguments rejects with an Error carrying the tool's error text.
- Runs raw JavaScript -- no Node, no file system, no network access, no timers.
- Accepts raw JavaScript source text, not JSON, quoted strings, or markdown code fences.
- You may optionally start the tool input with a first line like \`// @options: {"max_output_tokens": 1000, "timeout_ms": 60000}\`.
- \`max_output_tokens\` sets the token budget for the script's output. Defaults to 10000 tokens.
- \`timeout_ms\` sets a hard deadline for the whole script. By default there is none.
- When the JS code is fully evaluated, calls that are still running are cancelled and unawaited promises are silently discarded.
- Tool calls are real and have side effects. If the script fails partway, earlier calls are not undone.
- Scripts have a 256 MB memory limit; exceeding it throws \`InternalError: out of memory\`. Filter or aggregate large data instead of accumulating it.
- Global helpers:
- \`exit()\`: Immediately ends the current script successfully (like an early return from the top level).
- \`text(value: string | number | boolean | undefined | null)\`: Appends a text item. Non-string values are stringified with \`JSON.stringify(...)\` when possible.
- \`image(imageUrlOrItem: string | { image_url: string } | ImageContent)\`: Appends an image item. \`image_url\` should be a base64-encoded \`data:\` URL. To forward an MCP tool image, pass an individual \`ImageContent\` block from \`result.content\`, for example \`image(result.content[0])\`.
- \`store(key: string, value: any)\`: stores a serializable value under a string key for later \`codemode\` calls in the same session. Storing \`undefined\` deletes the key. Writes are kept only if the script succeeds.
- \`load(key: string)\`: returns the stored value for a string key, or \`undefined\` if it is missing.
- \`ALL_TOOLS\`: metadata for the enabled nested tools as \`{ name, description }\` entries.
- \`searchTools(query: string, options?: { limit?: number; namespace?: string })\`: resolves to the nested tools that best match the query (BM25, default limit 8), as \`{ name, description }\` entries like \`ALL_TOOLS\`.
- \`describeTool(name: string)\`: resolves to the description and declaration of a nested tool, or \`undefined\`.
- \`console.log(...)\` and the other \`console\` methods append a text item like \`text()\`.
- \`return value\` at the top level appends the value like \`text()\`.`;
const MODEL_TYPES = `type ModelType = "chat" | "image" | "classifier";
/** A model catalog entry. \`provider\` and \`id\` identify it; the other fields depend on the type. */
interface ModelInfo {
type?: ModelType;
provider: string;
id: string;
name: string;
api: string;
input: ("text" | "image")[];
contextWindow?: number;
[key: string]: unknown;
}
type ClassifierQuestion =
| { type: "choice"; instructions: string; criteria: Record<string, string> }
| { type: "score"; instructions: string; criteria: string[] }
| { type: "bool"; instructions: string; criteria: { true: string; false: string } };
type ClassifierAnswer =
| { type: "choice"; choice: string; probabilities: Record<string, number>; confidence: number }
| { type: "score"; score: number; confidence: number }
| { type: "bool"; probability: number };
interface ClassifierContext {
state: Record<string, unknown>;
questions: Record<string, ClassifierQuestion>;
}
interface ClassifierResult {
api: string;
provider: string;
model: string;
answers: Record<string, ClassifierAnswer>;
stopReason: "stop" | "error" | "aborted";
errorMessage?: string;
timestamp: number;
}`;
/** Declarations of the `models` globals; codemode-execute.ts implements them. */
export const MODEL_GLOBAL_DECLARATIONS: readonly Omit<CodemodeTool, "execute">[] = [
{
name: "models.getModelsOfType",
description: "Every known model of a type, optionally for one provider.",
signature: "(type: ModelType, provider?: string): Promise<ModelInfo[]>",
},
{
name: "models.getAvailableOfType",
description: "Models of a type whose provider has working credentials.",
signature: "(type: ModelType, provider?: string): Promise<ModelInfo[]>",
},
{
name: "models.getModelOfType",
description: "One catalog entry, or undefined.",
signature: "(type: ModelType, provider: string, id: string): Promise<ModelInfo | undefined>",
},
{
name: "models.classify",
description:
"Run a classifier model on one state. Only `provider` and `id` of `model` are used. Provider errors do not throw: check `stopReason` and `errorMessage`.",
signature: "(model: ModelInfo, context: ClassifierContext): Promise<ClassifierResult>",
},
];
const DEFERRED_TOOLS_GUIDANCE = `Some deferred nested tools may be omitted from this description. They are still available on the global \`tools\` object and listed in \`ALL_TOOLS\`.
To find one, call \`await searchTools(query)\`, or filter \`ALL_TOOLS\` by \`name\` and \`description\`.`;
/** Default for {@link CodemodeDescriptionOptions.inlineBudget}, in estimated tokens. */
export const DEFAULT_CODEMODE_INLINE_BUDGET = 3000;
/** Characters per token when estimating the cost of a tool section. */
const CHARS_PER_TOKEN = 4;
/** What a script sees of a tool. Tools without an output schema resolve to their text output. */
export function toCodemodeDeclaration(tool: AgentTool<any>): Omit<CodemodeTool, "execute"> {
return {
name: tool.name,
description: tool.description,
inputSchema: tool.parameters as CodemodeJsonSchema,
outputSchema: (tool.outputSchema as CodemodeJsonSchema | undefined) ?? TEXT_OUTPUT_SCHEMA,
};
}
/** Tools a script may call: every given tool except the codemode tool itself. */
export function getCodemodeCallableTools(tools: readonly AgentTool<any>[]): AgentTool<any>[] {
return tools.filter((tool) => tool.name !== CODEMODE_TOOL_NAME);
}
export interface CodemodeDescriptionOptions {
/** Declare the `models` namespace; only for tools created with model access. */
models?: boolean;
/** Namespace of each tool, by tool name. Tools of one namespace are listed under one heading. */
namespaces?: ReadonlyMap<string, ToolNamespace>;
/** Tools that are callable but never listed with their declaration (`deferred` exposure). */
deferred?: ReadonlySet<string>;
/**
* Estimated tokens (characters / 4) the tool sections may use. Tools that do not fit are left
* out, like deferred tools. Unset lists every tool that is not deferred.
*/
inlineBudget?: number;
}
/** `### \`id\` (\`raw name\`)` followed by the tool's description and declaration. */
function renderToolSection(declaration: Omit<CodemodeTool, "execute">): string {
const id = toCodemodeIdentifier(declaration.name);
const heading = id === declaration.name ? `### \`${id}\`` : `### \`${id}\` (\`${declaration.name}\`)`;
return `${heading}\n${renderToolSample(declaration).trim()}`;
}
interface CatalogEntry {
name: string;
section: string;
cost: number;
deferred: boolean;
}
interface CatalogGroup {
namespace: ToolNamespace | undefined;
entries: CatalogEntry[];
}
/**
* Pick the tool sections that fit the budget, like OpenCode's catalog: in each round every group
* (tools without a namespace first, then namespaces by name) places its cheapest remaining tool; a
* group whose next tool does not fit drops out while the others continue. Every namespace is
* represented before any namespace is complete.
*/
function selectCatalog(groups: readonly CatalogGroup[], budget: number | undefined): Set<string> {
const listable = groups.map((group) => group.entries.filter((entry) => !entry.deferred));
if (budget === undefined) return new Set(listable.flat().map((entry) => entry.name));
const queues = listable.map((entries) => [...entries].sort((a, b) => a.cost - b.cost));
const shown = new Set<string>();
let remaining = budget;
let active = queues.filter((queue) => queue.length > 0);
while (active.length > 0) {
active = active.filter((queue) => {
const next = queue[0];
if (next.cost > remaining) return false;
remaining -= next.cost;
shown.add(next.name);
queue.shift();
return queue.length > 0;
});
}
return shown;
}
/**
* Model-facing description: the helper list, guidance for omitted tools, the shared MCP types when
* MCP tools are callable, the `models` API, and one section per tool, grouped by namespace. Tool
* sections are limited to `inlineBudget`; every namespace is listed with its tool count either
* way, and the listing states whether it is complete.
*/
export function createCodemodeDescription(
tools: readonly AgentTool<any>[],
options: CodemodeDescriptionOptions = {},
): string {
const declarations = getCodemodeCallableTools(tools).map(toCodemodeDeclaration);
const groups = new Map<string, CatalogGroup>([["", { namespace: undefined, entries: [] }]]);
for (const declaration of declarations) {
const namespace = options.namespaces?.get(declaration.name);
const key = namespace ? `ns:${namespace.name}` : "";
const group = groups.get(key) ?? { namespace, entries: [] };
groups.set(key, group);
const section = renderToolSection(declaration);
group.entries.push({
name: declaration.name,
section,
cost: Math.ceil(section.length / CHARS_PER_TOKEN),
deferred: options.deferred?.has(declaration.name) === true,
});
}
const ordered = [...groups.values()].sort((a, b) =>
a.namespace === undefined ? -1 : b.namespace === undefined ? 1 : a.namespace.name.localeCompare(b.namespace.name),
);
const shown = selectCatalog(ordered, options.inlineBudget);
const complete = shown.size === declarations.length;
const sections = [DESCRIPTION_INTRO];
if (!complete) sections.push(DEFERRED_TOOLS_GUIDANCE);
if (declarations.some((declaration) => mcpStructuredContentSchema(declaration.outputSchema) !== undefined)) {
sections.push(`Shared MCP Types:\n\`\`\`ts\n${MCP_TYPESCRIPT_PREAMBLE}\n\`\`\``);
}
if (options.models) {
const noop = () => undefined;
const models = renderDeclarations({
globals: MODEL_GLOBAL_DECLARATIONS.map((global) => ({ ...global, execute: noop })),
});
sections.push(`Model API:\n\`\`\`ts\n${MODEL_TYPES}\n\n${models}\n\`\`\``);
}
if (declarations.length === 0) return sections.join("\n\n");
const toolSections = [
complete
? `Nested tools: COMPLETE list (${declarations.length} tool${declarations.length === 1 ? "" : "s"}).`
: `Nested tools: PARTIAL - ${shown.size} of ${declarations.length} shown.`,
];
for (const { namespace, entries } of ordered) {
const visible = entries.filter((entry) => shown.has(entry.name));
if (namespace) {
const count = `${entries.length} tool${entries.length === 1 ? "" : "s"}`;
const suffix =
visible.length === entries.length
? ""
: visible.length === 0
? ", none shown"
: `, ${visible.length} shown`;
const description = namespace.description?.trim();
toolSections.push(`## ${namespace.name} (${count}${suffix})${description ? `\n${description}` : ""}`);
}
for (const entry of visible) toolSections.push(entry.section);
}
sections.push(toolSections.join("\n\n"));
return sections.join("\n\n");
}
/**
* How the codemode tool presents tools that are both declared and callable from scripts:
* - `on`: their descriptions get the codemode declaration appended, and the codemode description
* lists only the callable tools without `direct` exposure.
* - `only`: the codemode description lists every callable tool, and requests leave out the
* declarations of active `direct` tools.
*
* Listing by exposure, not by the active set, keeps the codemode description unchanged when
* `tool_search` loads a tool, so loads do not redeclare codemode.
*/
function prepareCodemodeLoadout(loadout: ToolLoadout, options: CodemodeToolOptions): ToolLoadoutChanges {
const mode = options.getMode?.() ?? "on";
const isDirect = (tool: AgentTool) => loadout.getExposure(tool.name) === "direct";
const callable = getCodemodeCallableTools(loadout.callable);
const callableNames = new Set(callable.map((tool) => tool.name));
const descriptions: Record<string, string> = {};
if (mode === "on") {
for (const tool of loadout.declared) {
if (callableNames.has(tool.name)) descriptions[tool.name] = renderToolSample(toCodemodeDeclaration(tool));
}
}
const listed = mode === "only" ? callable : callable.filter((tool) => !isDirect(tool));
const namespaces = new Map(
listed.flatMap((tool) => {
const namespace = loadout.getNamespace(tool.name);
return namespace ? [[tool.name, namespace] as const] : [];
}),
);
descriptions[CODEMODE_TOOL_NAME] = createCodemodeDescription(listed, {
models: options.models === true,
namespaces,
deferred: new Set(
listed.filter((tool) => loadout.getExposure(tool.name) === "deferred").map((tool) => tool.name),
),
inlineBudget: options.getInlineBudget?.() ?? DEFAULT_CODEMODE_INLINE_BUDGET,
});
const declaredNames = new Set(loadout.declared.map((tool) => tool.name));
return {
descriptions,
hiddenDeclarations:
mode === "only"
? callable.filter((tool) => isDirect(tool) && declaredNames.has(tool.name)).map((tool) => tool.name)
: [],
};
}
export function createCodemodeToolDefinition(
options: CodemodeToolOptions = {},
): ToolDefinition<typeof codemodeSchema, CodemodeToolDetails | undefined> {
return {
name: CODEMODE_TOOL_NAME,
label: CODEMODE_TOOL_NAME,
// Replaced with the declarations of the callable tools when the tool is activated.
description: createCodemodeDescription([], { models: options.models === true }),
promptSnippet: codemodeToolSystemPromptContribution.snippet,
promptGuidelines: [...codemodeToolSystemPromptContribution.guidelines],
parameters: codemodeSchema,
// Scripts must not start other scripts.
exposure: "model-only",
prepareLoadout: (loadout) => prepareCodemodeLoadout(loadout, options),
// Capable models write the script as raw text instead of a JSON-escaped string.
constrainedSampling: { type: "grammar", variants: { openai_lark: CODEMODE_SOURCE_GRAMMAR } },
// The sandbox (worker, QuickJS wasm) loads on the first call, not at startup.
execute: async (toolCallId, params, signal, onUpdate, ctx) =>
(await loadCodemodeExecutor()).executeCodemode(toolCallId, params, signal, onUpdate, ctx, options),
...codemodeRenderers,
};
}
/**
* Create the codemode tool as an AgentTool. The description lists the given tools; the script can
* call whatever tools the agent loop provides at execution time.
*/
export function createCodemodeTool(
tools: readonly AgentTool<any>[] = [],
options: CodemodeToolOptions = {},
): AgentTool<typeof codemodeSchema> {
const definition = createCodemodeToolDefinition(options);
const tool = wrapToolDefinition(definition);
Object.assign(tool, {
description: createCodemodeDescription(tools, { models: options.models === true }),
promptSnippet: definition.promptSnippet,
promptGuidelines: definition.promptGuidelines,
});
return tool;
}

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -0,0 +1,985 @@
/**
* Built-in MCP integration.
*
* Connects the servers from `mcp.json` and the servers extensions register with
* `pi.registerMcpServer()` when a session starts, and servers registered later right away. A server
* in `mcp.json` takes precedence over a registered server of the same name. Tools are registered as
* `mcp__<server>__<tool>`. By default (`"exposure": "codemode"`) the tools are only callable from
* codemode scripts, which keeps large MCP tool lists out of the model's tool declarations; the
* codemode tool is activated for that unless `autoEnableCodemode` is false. `"codemode-deferred"`
* leaves them out of the codemode description as well. `"deferred"` declares them to the model once
* the `tool_search` tool loads them, and activates `tool_search` instead of codemode.
* `"exposure": "direct"` declares them to the model right away, and `"hidden"` makes them
* unreachable. `toolExposure` overrides the exposure of single tools. Servers with resources are
* reached through Codex's `list_mcp_resources`, `list_mcp_resource_templates`, and
* `read_mcp_resource` tools (resources.ts).
*
* Every call runs through pi's tool pipeline, so `tool_call`/`tool_result` hooks and permission
* extensions apply to MCP tools the same way they do to built-in tools.
*
* Problems found at startup (config errors, failed connections, servers that need a sign-in) are
* reported once. `/mcp` opens a manager to sign in, reconnect, enable or disable servers, and change
* their exposure; the last two are saved to the `mcp.json` that defines the server, or apply to the
* current session for registered servers.
*/
import { join, resolve } from "node:path";
import type { SelectItem } from "@earendil-works/pi-tui";
import type { TSchema } from "typebox";
import { getAgentDir } from "../../config.ts";
import type {
ExtensionAPI,
ExtensionCommandContext,
ExtensionContext,
ExtensionFactory,
ToolDefinition,
} from "../../core/extensions/types.ts";
import { openBrowser } from "../../utils/open-browser.ts";
import { CODEMODE_TOOL_NAME, isCodemodeTool } from "../codemode/tool.ts";
import { isToolSearchTool, TOOL_SEARCH_TOOL_NAME } from "../tool-search/tool.ts";
import {
getMcpToolExposure,
type LoadedMcpConfig,
loadMcpConfig,
type McpExposure,
type McpServerConfigPatch,
type McpServerEntry,
updateMcpServerConfig,
} from "./config.ts";
import type { McpOAuthCredentialStore, McpSignInPrompt } from "./oauth.ts";
import { createMcpResourceToolDefinitions } from "./resources.ts";
import { loadMcpRuntime } from "./runtime.lazy.ts";
import type * as McpRuntime from "./runtime.ts";
import type { McpServerConnection, McpServerLog, McpTransportFactory } from "./runtime.ts";
import { createMcpToolDefinition, createMcpToolName, type McpToolDetails } from "./tools.ts";
import { type McpMenu, type McpUi, showMcpManager } from "./ui.ts";
export type { McpTransportFactory } from "./runtime.ts";
export interface McpExtensionOptions {
/** Defaults to reading `mcp.json` from the agent directory and the trusted project. */
loadConfig?: (ctx: ExtensionContext) => LoadedMcpConfig;
/** Defaults to stdio and streamable HTTP transports built from the server config. */
createTransport?: McpTransportFactory;
/** Defaults to `mcp-auth.json` in the agent directory. */
credentials?: McpOAuthCredentialStore;
/** File server log messages are appended to. Defaults to `mcp.log` in the agent directory. */
logPath?: string;
/** Opens the OAuth authorization URL. Defaults to the platform browser. */
openUrl?: (url: string) => void;
/** Saves `/mcp` changes to the server's config file. Defaults to editing its `mcp.json`. */
updateConfig?: (entry: McpServerEntry, patch: McpServerConfigPatch) => void;
/**
* How long the first prompt waits for servers that are still connecting at startup, in
* milliseconds. Their tools become available when they connect. Default: 10000.
*/
startupWaitMs?: number;
}
const DEFAULT_STARTUP_WAIT_MS = 10_000;
/** A configured server. Disabled servers have no connection. */
interface McpServer {
entry: McpServerEntry;
connection?: McpServerConnection;
/** For servers extensions registered: the config as registered, to detect re-registrations. */
registeredConfig?: string;
/** Result of the last `/mcp` action that failed, shown in the manager. */
message?: string;
}
const EXPOSURE_DESCRIPTIONS: Record<Exclude<McpExposure, "hidden">, string> = {
codemode: "called from codemode scripts, listed in the codemode description",
"codemode-deferred": "called from codemode scripts, not listed; scripts find them with searchTools()",
deferred: "not declared until tool_search loads them, then called directly; no codemode needed",
direct: "declared to the model like built-in tools",
};
function errorMessage(error: unknown): string {
return error instanceof Error ? error.message : String(error);
}
function firstLine(text: string): string {
return text.split("\n", 1)[0] ?? "";
}
function isEnabled(server: McpServer): boolean {
return server.entry.config.enabled !== false;
}
function exposureOf(entry: McpServerEntry): McpExposure {
return entry.config.exposure ?? "codemode";
}
/** Short state for lists and the startup report. `withError` appends the first line of a failure. */
function describeState(server: McpServer, withError = true): string {
if (!isEnabled(server)) return "disabled";
const connection = server.connection;
if (!connection) return "starting";
switch (connection.state) {
case "needs-auth":
return "needs sign-in";
case "failed":
return withError ? `failed: ${firstLine(connection.error ?? "unknown error")}` : "failed";
case "connected": {
const { tools, resources } = connection;
const count = resources.length;
const resourceCount = count > 0 ? ` · ${count} resource${count === 1 ? "" : "s"}` : "";
return `connected · ${tools.length} tool${tools.length === 1 ? "" : "s"}${resourceCount}`;
}
case "connecting":
return "connecting…";
default:
return connection.state;
}
}
/** Servers that need the user first. */
function attentionRank(server: McpServer): number {
if (!isEnabled(server)) return 5;
switch (server.connection?.state) {
case "needs-auth":
return 0;
case "failed":
return 1;
case "disconnected":
return 2;
case "connected":
return 4;
default:
return 3;
}
}
function describeTransport(entry: McpServerEntry): string {
const { config } = entry;
if ("url" in config) return config.url;
return [config.command, ...(config.args ?? [])].join(" ");
}
const MCP_USAGE = "Usage: /mcp, /mcp login [server], /mcp logout [server], /mcp reconnect [server]";
export function createMcpExtension(options: McpExtensionOptions = {}): ExtensionFactory {
return (pi: ExtensionAPI) => {
let servers: McpServer[] = [];
/** Servers from `mcp.json`, which take precedence over registered servers of the same name. */
let configuredEntries: McpServerEntry[] = [];
let configErrors: string[] = [];
/** Registered servers that `mcp.json` overrides, shown in `/mcp`. */
let overridden: string[] = [];
/** Between session_start and session_shutdown. Registrations before that are read on session_start. */
let sessionActive = false;
let autoEnableCodemode = true;
/** Whether the "codemode tools unreachable" warning was shown since the session started. */
let warnedUnreachable = false;
let pending: Promise<unknown> | undefined;
/** Whether a prompt already waited for the startup connections since the session started. */
let waitedForStartup = false;
const startupWaitMs = options.startupWaitMs ?? DEFAULT_STARTUP_WAIT_MS;
/** Bumped on every session start and shutdown so a runtime load that resolves late is dropped. */
let generation = 0;
/** Working directory of the session, for stdio servers. */
let sessionCwd = process.cwd();
let credentials = options.credentials;
let serverLog: McpServerLog | undefined;
const openUrl = options.openUrl ?? openBrowser;
const updateConfig =
options.updateConfig ?? ((entry, patch) => updateMcpServerConfig(entry.source, entry.name, patch));
const listeners = new Set<() => void>();
const emitChange = () => {
for (const listener of listeners) listener();
};
const subscribe = (listener: () => void) => {
listeners.add(listener);
return () => {
listeners.delete(listener);
};
};
const connections = () => servers.flatMap((server) => (server.connection ? [server.connection] : []));
const findServer = (name: string) => servers.find((server) => server.entry.name === name);
/** Servers extensions registered, except names `mcp.json` defines, which take precedence. */
const registeredServers = (): { servers: McpServer[]; overridden: string[] } => {
const registered: McpServer[] = [];
const overriddenNames: string[] = [];
for (const { name, config, extensionPath } of pi.getMcpServers()) {
const configured = configuredEntries.find((entry) => entry.name === name);
if (configured) {
overriddenNames.push(`"${name}" registered by ${extensionPath} is overridden by ${configured.source}`);
continue;
}
registered.push({
entry: { name, config, source: extensionPath, scope: "extension" },
registeredConfig: JSON.stringify(config),
});
}
return { servers: registered, overridden: overriddenNames };
};
const getCredentials = (runtime: typeof McpRuntime): McpOAuthCredentialStore => {
credentials ??= new runtime.McpOAuthCredentialStore();
return credentials;
};
const getServerLog = (runtime: typeof McpRuntime): McpServerLog => {
serverLog ??= new runtime.McpServerLog(options.logPath ?? join(getAgentDir(), "mcp.log"));
return serverLog;
};
/** pi tool name to the `<server>\0<tool>` it was assigned to, so names stay unique and stable. */
const toolOwners = new Map<string, string>();
/** Tool names currently offered by each server. */
const serverTools = new Map<string, Set<string>>();
/** Last definition registered under each tool name, to re-register withdrawn tools as hidden. */
const definitions = new Map<string, ToolDefinition<TSchema, McpToolDetails>>();
const registerTools = (connection: McpServerConnection) => {
const server = connection.entry.name;
const entry = findServer(server)?.entry ?? connection.entry;
const namespaceName = `mcp__${server}`;
const namespace = {
name: namespaceName,
description: connection.instructions ?? `Tools in the ${namespaceName} namespace.`,
};
const previous = serverTools.get(server) ?? new Set<string>();
const current = new Set<string>();
const assignName = (tool: string, owner: string) => {
const name = createMcpToolName(server, tool, (candidate) => {
const existing = toolOwners.get(candidate);
return (existing !== undefined && existing !== owner) || current.has(candidate);
});
toolOwners.set(name, owner);
current.add(name);
return name;
};
for (const tool of connection.tools) {
const definition = createMcpToolDefinition({
server,
tool,
name: assignName(tool.name, `${server}\0${tool.name}`),
exposure: getMcpToolExposure(entry.config, tool.name),
namespace,
timeoutMs: connection.timeoutMs,
getClient: async () => connection,
readableResources: () => resourceServers().includes(connection),
});
definitions.set(definition.name, definition);
pi.registerTool(definition);
}
serverTools.set(server, current);
// Tools cannot be unregistered, so tools the server dropped are re-registered as hidden. When
// the server offers them again they are registered with their configured exposure above.
for (const name of previous) {
const definition = definitions.get(name);
if (!current.has(name) && definition) pi.registerTool({ ...definition, exposure: "hidden" });
}
syncResourceTools();
};
/** Make a disabled server's tools unreachable. */
const hideTools = (server: string) => {
for (const name of serverTools.get(server) ?? []) {
const definition = definitions.get(name);
if (definition) pi.registerTool({ ...definition, exposure: "hidden" });
}
serverTools.set(server, new Set());
syncResourceTools();
};
/** Enabled servers with resources whose exposure is not `hidden`, which the resource tools reach. */
const serversWithResources = (): McpServer[] =>
servers.filter(
(server) => server.connection?.hasResources && isEnabled(server) && exposureOf(server.entry) !== "hidden",
);
const resourceServers = (): McpServerConnection[] =>
serversWithResources().flatMap((server) => (server.connection ? [server.connection] : []));
/** Exposure the resource tools were last registered with; undefined until a server has resources. */
let resourceToolsExposure: McpExposure | undefined;
/**
* Register the resource tools with the widest exposure of the servers they reach: `direct` when
* one of them is direct, and so on. They are hidden when no server has resources.
*/
const syncResourceTools = () => {
const exposures = new Set(serversWithResources().map((server) => exposureOf(server.entry)));
const exposure = (["direct", "codemode", "codemode-deferred", "deferred"] as const).find((candidate) =>
exposures.has(candidate),
);
const next = exposure ?? "hidden";
if (next === resourceToolsExposure || (resourceToolsExposure === undefined && next === "hidden")) return;
const wasDirect = resourceToolsExposure === "direct";
resourceToolsExposure = next;
const resourceDefinitions = createMcpResourceToolDefinitions({ exposure: next, servers: resourceServers });
for (const definition of resourceDefinitions) pi.registerTool(definition);
if (wasDirect) {
const names = new Set(resourceDefinitions.map((definition) => definition.name));
pi.setActiveTools(pi.getActiveTools().filter((name) => !names.has(name)));
}
};
/**
* Tools that are not declared to the model are reached through the codemode tool (scripts call
* them) or the tool_search tool (it declares them). Either reaches every such tool. Activate the
* one the tools' exposure asks for: codemode for `codemode` and `codemode-deferred` unless
* `autoEnableCodemode` is false, tool_search for `deferred`.
*/
const ensureDiscoveryActive = (ctx: ExtensionContext) => {
const exposures = new Set<McpExposure>();
for (const { connection, entry } of servers) {
if (connection?.state !== "connected") continue;
// Resource tools share the server's exposure.
if (connection.hasResources) exposures.add(exposureOf(entry));
for (const tool of connection.tools) exposures.add(getMcpToolExposure(entry.config, tool.name));
}
const needsCodemode = exposures.has("codemode") || exposures.has("codemode-deferred");
const needsToolSearch = exposures.has("deferred");
if (!needsCodemode && !needsToolSearch) return;
// Other extensions' tools of the same names cannot reach MCP tools, so never activate them.
const tools = pi.getAllTools();
const hasCodemode = tools.some(isCodemodeTool);
const hasToolSearch = tools.some(isToolSearchTool);
const active = pi.getActiveTools();
const activate: string[] = [];
if (needsCodemode && hasCodemode && autoEnableCodemode && !active.includes(CODEMODE_TOOL_NAME)) {
activate.push(CODEMODE_TOOL_NAME);
}
if (needsToolSearch && hasToolSearch && !active.includes(TOOL_SEARCH_TOOL_NAME)) {
activate.push(TOOL_SEARCH_TOOL_NAME);
}
if (activate.length > 0) pi.setActiveTools([...active, ...activate]);
const reachable = [...active, ...activate];
if (hasCodemode && reachable.includes(CODEMODE_TOOL_NAME)) return;
if (hasToolSearch && reachable.includes(TOOL_SEARCH_TOOL_NAME)) return;
if (warnedUnreachable) return;
warnedUnreachable = true;
const reason = needsCodemode && hasCodemode && !autoEnableCodemode ? " (autoEnableCodemode is false)" : "";
ctx.ui.notify(
`MCP tools are only reachable from the codemode or tool_search tool, but neither is active${reason}; they cannot be called.`,
"warning",
);
};
/**
* Stored tokens of servers waiting for a sign-in, as they were when the sign-in was needed.
* `pi mcp login` in another process (for example run by the agent) changes them.
*/
const tokensAtSignIn = new Map<McpServerConnection, string>();
const storedTokens = (connection: McpServerConnection): string => {
const url = connection.oauthUrl;
return url && credentials ? JSON.stringify(credentials.tokens(url) ?? null) : "null";
};
const onConnectionChange = (connection: McpServerConnection) => {
if (connection.state !== "needs-auth") tokensAtSignIn.delete(connection);
else if (!tokensAtSignIn.has(connection)) tokensAtSignIn.set(connection, storedTokens(connection));
emitChange();
};
/** Reconnect servers that need a sign-in when their credentials were stored since. */
const reconnectSignedIn = async (ctx: ExtensionContext) => {
const signedIn = [...tokensAtSignIn].filter(([connection, tokens]) => storedTokens(connection) !== tokens);
if (signedIn.length === 0) return;
for (const [connection] of signedIn) tokensAtSignIn.delete(connection);
await Promise.allSettled(signedIn.map(([connection]) => connection.reconnect()));
ensureDiscoveryActive(ctx);
};
/** Create the server's connection, loading the MCP runtime on first use. */
const createConnection = async (server: McpServer): Promise<McpServerConnection> => {
const runtime = await loadMcpRuntime();
const connection = new runtime.McpServerConnection({
entry: server.entry,
cwd: sessionCwd,
createTransport: options.createTransport ?? runtime.createDefaultTransport,
credentials: getCredentials(runtime),
log: getServerLog(runtime),
onTools: registerTools,
onChange: onConnectionChange,
});
server.connection = connection;
emitChange();
return connection;
};
/**
* One message for everything that needs the user after startup, or only for `only`, servers
* that connected later.
*/
const reportProblems = (ctx: ExtensionContext, only?: readonly McpServer[]) => {
const lines = only ? [] : configErrors.map((error) => `config: ${error}`);
for (const server of only ?? servers) {
const state = server.connection?.state;
if (state === "needs-auth" || state === "failed")
lines.push(`${server.entry.name}: ${describeState(server)}`);
}
if (lines.length === 0) return;
ctx.ui.notify(
`MCP servers need attention:\n${lines.map((line) => ` ${line}`).join("\n")}\nRun /mcp to fix.`,
"warning",
);
};
/**
* Save a config change; returns an error message when the file could not be updated. Changes to
* registered servers only apply to the current session.
*/
const saveConfig = (server: McpServer, patch: McpServerConfigPatch): string | undefined => {
if (server.entry.scope !== "extension") {
try {
updateConfig(server.entry, patch);
} catch (error) {
return `Could not update ${server.entry.source}: ${errorMessage(error)}`;
}
}
server.entry = { ...server.entry, config: { ...server.entry.config, ...patch } };
return undefined;
};
const signIn = async (server: McpServer, prompt: McpSignInPrompt): Promise<string | undefined> => {
const connection = server.connection;
const url = connection?.oauthUrl;
if (!connection || !url) return `MCP server "${server.entry.name}" does not use OAuth.`;
const runtime = await loadMcpRuntime();
try {
await runtime.signInMcpServer({
serverUrl: url,
store: getCredentials(runtime).forServer(url),
settings: connection.oauthSettings(),
challenge: connection.challenge,
prompt,
});
} catch (error) {
if (error instanceof runtime.McpSignInCancelledError) return "Sign-in cancelled.";
return `Sign-in failed: ${errorMessage(error)}`;
}
// The challenge that asked for this sign-in (for example for more scope) is answered.
connection.challenge = undefined;
try {
await connection.reconnect();
} catch (error) {
return `Signed in, but ${errorMessage(error)}`;
}
return undefined;
};
const signOut = async (server: McpServer): Promise<boolean> => {
const connection = server.connection;
const url = connection?.oauthUrl;
if (!connection || !url) return false;
const removed = getCredentials(await loadMcpRuntime()).remove(url);
await connection.signOut();
return removed;
};
const reconnect = async (server: McpServer): Promise<string | undefined> => {
const connection = server.connection;
if (!connection) return `MCP server "${server.entry.name}" is disabled.`;
try {
await connection.reconnect();
return undefined;
} catch (error) {
return errorMessage(error);
}
};
/** Returns an error message when the config could not be saved; connection errors show in the state. */
const setEnabled = async (server: McpServer, enabled: boolean): Promise<string | undefined> => {
const failed = saveConfig(server, { enabled });
if (failed) return failed;
if (!enabled) {
const connection = server.connection;
server.connection = undefined;
hideTools(server.entry.name);
emitChange();
await connection?.close();
return undefined;
}
const connection = await createConnection(server);
await connection.getClient().catch(() => undefined);
return undefined;
};
const setExposure = (server: McpServer, exposure: McpExposure): string | undefined => {
const failed = saveConfig(server, { exposure });
if (failed) return failed;
if (server.connection?.state === "connected") registerTools(server.connection);
syncResourceTools();
// Tools no longer exposed directly leave the declared set; direct tools are activated on registration.
const indirect = new Set(
pi
.getAllTools()
.filter((tool) => tool.exposure !== "direct")
.map((tool) => tool.name),
);
const tools = serverTools.get(server.entry.name) ?? new Set<string>();
pi.setActiveTools(pi.getActiveTools().filter((name) => !tools.has(name) || !indirect.has(name)));
emitChange();
return undefined;
};
// ---------------------------------------------------------------------------------------
// Manager (`/mcp` in the TUI)
// ---------------------------------------------------------------------------------------
const notices = () => [
...configErrors.map((error) => `config: ${error}`),
...overridden.map((line) => `overridden: ${line}`),
];
const serversMenu = (): McpMenu => ({
title: "MCP servers",
error: notices().join("\n") || undefined,
items: [...servers]
.sort((a, b) => attentionRank(a) - attentionRank(b) || a.entry.name.localeCompare(b.entry.name))
.map((server) => ({
value: server.entry.name,
label: server.entry.name,
description: `${describeState(server)} · ${exposureOf(server.entry)} · ${server.entry.scope ?? server.entry.source}`,
})),
empty: `No MCP servers configured. Add them to ${resolve(getAgentDir(), "mcp.json")} or .pi/mcp.json.`,
confirmLabel: "manage",
cancelLabel: "close",
});
const serverMenu = (name: string): McpMenu => {
const server = findServer(name);
if (!server) {
return {
title: name,
items: [],
empty: "This server is no longer configured.",
confirmLabel: "",
cancelLabel: "back",
};
}
const { entry, connection } = server;
const saved =
entry.scope === "extension"
? "for this session"
: entry.scope
? `saved to the ${entry.scope} mcp.json`
: "saved to mcp.json";
const items: SelectItem[] = [];
if (!isEnabled(server)) {
items.push({ value: "enable", label: "Enable", description: saved });
} else {
const state = connection?.state;
if (state === "needs-auth")
items.push({ value: "signin", label: "Sign in", description: "opens the browser" });
if (state === "connected" && connection) {
items.push({ value: "tools", label: "Tools", description: `${connection.tools.length} offered` });
}
if (state === "failed" || state === "disconnected" || state === "connected" || state === "needs-auth") {
items.push({ value: "reconnect", label: "Reconnect" });
}
if (state === "connected" && connection?.oauthUrl) {
items.push({ value: "signout", label: "Sign out", description: "deletes the stored credentials" });
}
items.push({ value: "exposure", label: "Exposure", description: exposureOf(entry) });
items.push({ value: "disable", label: "Disable", description: saved });
}
const details = [
describeTransport(entry),
`${entry.scope ?? "config"}: ${entry.source}`,
`State: ${describeState(server, false)}`,
];
const error = [server.message, connection?.state === "connected" ? undefined : connection?.error]
.filter((line): line is string => line !== undefined)
.join("\n");
return {
title: `MCP server ${name}`,
details: details.join("\n"),
error: error || undefined,
items,
selected: items[0]?.value,
confirmLabel: "select",
cancelLabel: "back",
};
};
const showTools = async (ui: McpUi, server: McpServer) => {
const exposure = exposureOf(server.entry);
const overridden = Object.keys(server.entry.config.toolExposure ?? {}).length > 0;
await ui.menu(() => ({
title: `Tools of ${server.entry.name}`,
details: `Exposure ${exposure}: ${exposure === "hidden" ? "unreachable" : EXPOSURE_DESCRIPTIONS[exposure]}${overridden ? "\nSome tools override it with toolExposure." : ""}`,
items: (server.connection?.tools ?? []).map((tool) => {
const toolExposure = getMcpToolExposure(server.entry.config, tool.name);
const description = firstLine(tool.description ?? "");
return {
value: tool.name,
label: tool.name,
description: toolExposure === exposure ? description : `[${toolExposure}] ${description}`,
};
}),
empty: "The server offers no tools.",
confirmLabel: "back",
cancelLabel: "back",
}));
};
const chooseExposure = async (ui: McpUi, server: McpServer): Promise<string | undefined> => {
const current = exposureOf(server.entry);
const choice = await ui.menu(() => ({
title: `Exposure of ${server.entry.name}`,
details:
server.entry.scope === "extension"
? `Applies to this session; the server is registered by ${server.entry.source}.`
: `Saved to ${server.entry.source}.`,
items: (Object.keys(EXPOSURE_DESCRIPTIONS) as (keyof typeof EXPOSURE_DESCRIPTIONS)[]).map((exposure) => ({
value: exposure,
label: `${exposure === current ? "✓ " : " "}${exposure}`,
description: EXPOSURE_DESCRIPTIONS[exposure],
})),
selected: current,
confirmLabel: "save",
cancelLabel: "back",
}));
if (!choice || choice === current) return undefined;
return setExposure(server, choice as McpExposure);
};
const runAction = async (ui: McpUi, ctx: ExtensionContext, server: McpServer, action: string) => {
const { name } = server.entry;
let message: string | undefined;
switch (action) {
case "signin": {
const title = `Sign in to ${name}`;
let authorizationUrl = "";
ui.status(title, "Contacting the authorization server…");
message = await signIn(server, {
showAuthorizationUrl: (url) => {
authorizationUrl = url.href;
openUrl(url.href);
},
promptForRedirectUrl: async (signal) => {
const value = await ui.redirectUrl(title, authorizationUrl, signal);
ui.status(title, "Connecting…");
return value;
},
});
break;
}
case "reconnect":
// A failure shows as the connection's state and error.
ui.status(`MCP server ${name}`, "Reconnecting…");
await reconnect(server);
break;
case "signout":
await signOut(server);
break;
case "tools":
await showTools(ui, server);
break;
case "exposure":
message = await chooseExposure(ui, server);
break;
case "enable":
case "disable":
ui.status(`MCP server ${name}`, action === "enable" ? "Connecting…" : "Disconnecting…");
message = await setEnabled(server, action === "enable");
break;
}
server.message = message;
ensureDiscoveryActive(ctx);
emitChange();
};
const manage = async (ui: McpUi, ctx: ExtensionContext) => {
for (;;) {
const name = await ui.menu(serversMenu, subscribe);
if (!name) return;
for (;;) {
const action = await ui.menu(() => serverMenu(name), subscribe);
const server = findServer(name);
if (!action || !server) break;
await runAction(ui, ctx, server, action);
}
}
};
// ---------------------------------------------------------------------------------------
// Subcommands and plain status (no TUI)
// ---------------------------------------------------------------------------------------
const formatStatus = (): string => {
if (servers.length === 0 && configErrors.length === 0 && overridden.length === 0) {
return `No MCP servers configured. Add them to ${resolve(getAgentDir(), "mcp.json")} or .pi/mcp.json.`;
}
const lines = servers.map((server) => {
const { name } = server.entry;
const exposure = exposureOf(server.entry);
const connection = server.connection;
if (connection?.state === "needs-auth")
return `${name}: needs sign-in, run /mcp login ${name} (${exposure})`;
const tools = connection?.state === "connected" ? `, ${connection.tools.length} tools` : "";
const state = !isEnabled(server)
? "disabled"
: connection?.state === "disconnected"
? "disconnected, reconnects on next call"
: (connection?.state ?? "starting");
const error =
connection?.error && connection.state !== "connected"
? `\n ${connection.error.split("\n").join("\n ")}`
: "";
return `${name}: ${state}${tools} (${exposure})${error}`;
});
for (const error of configErrors) lines.push(`config error: ${error}`);
for (const line of overridden) lines.push(`overridden: ${line}`);
return lines.join("\n");
};
/** Resolve the server for a subcommand, asking when the name is omitted and ambiguous. */
const pickServer = async (
name: string | undefined,
ctx: ExtensionCommandContext,
options: { eligible: (server: McpServer) => boolean; preferred: (server: McpServer) => boolean; none: string },
): Promise<McpServer | undefined> => {
if (name) {
const server = findServer(name);
if (!server) ctx.ui.notify(`No MCP server named "${name}".`, "error");
else if (!options.eligible(server)) ctx.ui.notify(options.none, "error");
return server && options.eligible(server) ? server : undefined;
}
const candidates = servers.filter(options.eligible);
if (candidates.length === 0) {
ctx.ui.notify(options.none, "info");
return undefined;
}
const preferred = candidates.filter(options.preferred);
if (candidates.length === 1) return candidates[0];
if (preferred.length === 1) return preferred[0];
const choice = await ctx.ui.select(
"MCP server",
candidates.map((server) => server.entry.name),
);
return candidates.find((server) => server.entry.name === choice);
};
const usesOAuth = (server: McpServer) => server.connection?.oauthUrl !== undefined;
const oauthPick = {
eligible: usesOAuth,
preferred: (server: McpServer) => server.connection?.state === "needs-auth",
none: "No enabled MCP server uses OAuth. Only HTTP servers without an Authorization header do.",
};
const loginCommand = async (server: McpServer, ctx: ExtensionCommandContext) => {
const { name } = server.entry;
if (!ctx.hasUI) {
ctx.ui.notify(`Signing in to MCP server "${name}" requires interactive mode.`, "error");
return;
}
const failure = await signIn(server, {
showAuthorizationUrl: (url) => {
ctx.ui.notify(`Sign in to MCP server "${name}" in your browser:\n${url.href}`, "info");
openUrl(url.href);
},
promptForRedirectUrl: (signal) =>
ctx.ui.input(
`Waiting for sign-in to "${name}". If the browser cannot reach this machine, paste the URL it was redirected to.`,
"http://127.0.0.1:.../callback?code=...",
{ signal },
),
});
if (failure) {
ctx.ui.notify(failure, failure === "Sign-in cancelled." ? "info" : "error");
return;
}
ensureDiscoveryActive(ctx);
ctx.ui.notify(`Signed in to MCP server "${name}" (${server.connection?.tools.length ?? 0} tools).`, "info");
};
pi.on("session_start", (_event, ctx) => {
const loaded = (options.loadConfig ?? defaultLoadConfig)(ctx);
configErrors = loaded.errors;
autoEnableCodemode = loaded.autoEnableCodemode ?? true;
warnedUnreachable = false;
waitedForStartup = false;
sessionCwd = ctx.cwd;
const current = ++generation;
sessionActive = true;
configuredEntries = loaded.servers;
const registered = registeredServers();
overridden = registered.overridden;
servers = [...loaded.servers.map((entry) => ({ entry })), ...registered.servers];
emitChange();
const enabled = servers.filter(isEnabled);
if (enabled.length === 0) {
reportProblems(ctx);
return;
}
// The MCP client loads only now, so sessions without servers never pay for it. Waiting one
// event loop turn lets the first render happen before loading and connecting.
pending = new Promise((resolve) => setImmediate(resolve))
.then(() => loadMcpRuntime())
.then(async () => {
if (current !== generation) return;
const started = await Promise.all(enabled.map((server) => createConnection(server)));
if (current !== generation) return;
await Promise.allSettled(started.map((connection) => connection.getClient()));
if (current !== generation) return;
ensureDiscoveryActive(ctx);
reportProblems(ctx);
})
.catch((error: unknown) => {
// The session may have been disposed meanwhile, which makes ctx stale.
try {
ctx.ui.notify(`MCP failed to load: ${errorMessage(error)}`, "error");
} catch {}
});
});
// The first prompt waits for startup connections so their tools are available to it, but not
// indefinitely: a slow or hanging server must not hold up the prompt.
pi.on("before_agent_start", async (_event, ctx) => {
const startup = pending;
if (!startup || waitedForStartup) return;
waitedForStartup = true;
let timer: NodeJS.Timeout | undefined;
const finished = await Promise.race([
startup.then(() => true),
new Promise<boolean>((resolve) => {
timer = setTimeout(() => resolve(false), startupWaitMs);
}),
]);
clearTimeout(timer);
if (!finished) {
ctx.ui.notify("MCP servers are still connecting; their tools become available once connected.", "info");
}
});
// Pick up sign-ins done outside the session, such as `pi mcp login` run by the agent.
pi.on("turn_start", async (_event, ctx) => {
if (tokensAtSignIn.size > 0) await reconnectSignedIn(ctx);
});
// Servers registered or unregistered during the session connect or disconnect right away.
pi.on("mcp_servers_change", async (_event, ctx) => {
if (!sessionActive) return;
const current = generation;
const registered = registeredServers();
overridden = registered.overridden;
const next = new Map(registered.servers.map((server) => [server.entry.name, server]));
// Unregistered servers and re-registered ones with a new config are dropped; the latter come back below.
const removed = servers.filter(
(server) =>
server.entry.scope === "extension" &&
next.get(server.entry.name)?.registeredConfig !== server.registeredConfig,
);
servers = servers.filter((server) => !removed.includes(server));
for (const server of removed) hideTools(server.entry.name);
const added = registered.servers.filter((server) => !findServer(server.entry.name));
servers.push(...added);
emitChange();
await Promise.all(removed.map((server) => server.connection?.close()));
const connecting = added.filter(isEnabled);
if (current !== generation || connecting.length === 0) return;
try {
const started = await Promise.all(connecting.map((server) => createConnection(server)));
if (current !== generation) {
await Promise.all(started.map((connection) => connection.close()));
return;
}
await Promise.allSettled(started.map((connection) => connection.getClient()));
} catch (error) {
ctx.ui.notify(`MCP failed to load: ${errorMessage(error)}`, "error");
return;
}
if (current !== generation) return;
ensureDiscoveryActive(ctx);
reportProblems(ctx, connecting);
});
pi.on("session_shutdown", async () => {
sessionActive = false;
generation++;
const closing = connections();
servers = [];
emitChange();
await Promise.all(closing.map((connection) => connection.close()));
});
pi.registerCommand("mcp", {
description: "Manage MCP servers: sign in, reconnect, enable or disable, and change exposure",
getArgumentCompletions: (prefix) => {
const [action, server, ...rest] = prefix.trimStart().split(/\s+/);
if (rest.length > 0) return null;
if (server === undefined) {
return ["login", "logout", "reconnect"]
.filter((item) => item.startsWith(action ?? ""))
.map((item) => ({ value: `${item} `, label: item }));
}
if (action !== "login" && action !== "logout" && action !== "reconnect") return null;
const items = servers
.filter((candidate) =>
action === "reconnect" ? candidate.connection !== undefined : usesOAuth(candidate),
)
.filter((candidate) => candidate.entry.name.startsWith(server))
.map((candidate) => ({
value: `${action} ${candidate.entry.name}`,
label: candidate.entry.name,
description: describeState(candidate),
}));
return items.length > 0 ? items : null;
},
handler: async (args, ctx) => {
await pending;
const [action, name, ...extra] = args.trim().split(/\s+/).filter(Boolean);
if (action === undefined) {
if (ctx.mode === "tui") await showMcpManager(ctx, (ui) => manage(ui, ctx));
else ctx.ui.notify(formatStatus(), "info");
return;
}
if (extra.length > 0) {
ctx.ui.notify(MCP_USAGE, "warning");
return;
}
switch (action) {
case "login": {
const server = await pickServer(name, ctx, oauthPick);
if (server) await loginCommand(server, ctx);
return;
}
case "logout": {
const server = await pickServer(name, ctx, oauthPick);
if (!server) return;
const removed = await signOut(server);
ctx.ui.notify(
removed
? `Signed out of MCP server "${server.entry.name}".`
: `No stored credentials for MCP server "${server.entry.name}".`,
"info",
);
return;
}
case "reconnect": {
const server = await pickServer(name, ctx, {
eligible: (candidate) => candidate.connection !== undefined,
preferred: (candidate) =>
candidate.connection?.state === "failed" || candidate.connection?.state === "disconnected",
none: "No enabled MCP server to reconnect.",
});
if (!server) return;
const failure = await reconnect(server);
if (failure) ctx.ui.notify(failure, "error");
else {
ensureDiscoveryActive(ctx);
ctx.ui.notify(
`Reconnected to MCP server "${server.entry.name}" (${describeState(server)}).`,
"info",
);
}
return;
}
default:
ctx.ui.notify(MCP_USAGE, "warning");
}
},
});
};
}
function defaultLoadConfig(ctx: ExtensionContext): LoadedMcpConfig {
return loadMcpConfig({ agentDir: getAgentDir(), cwd: ctx.cwd, projectTrusted: ctx.isProjectTrusted() });
}
export default createMcpExtension();

View file

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

View file

@ -0,0 +1,355 @@
/**
* OAuth sign-in for remote MCP servers.
*
* Connections never start a browser flow on their own. They send the stored access token and, after
* a 401, try the stored refresh token. When that is not possible they fail with
* `McpOAuthAuthorizationRequiredError`, and the user signs in through `/mcp`, which runs
* the authorization code flow (PKCE, dynamic client registration) against a loopback callback.
*
* Credentials live in `<agent-dir>/mcp-auth.json`, keyed by server URL.
*/
import { join } from "node:path";
import type { AuthProvider } from "@earendil-works/pi-mcp";
import {
adaptOAuthProvider,
authorizeMcp,
McpOAuthAuthorizationRequiredError,
McpOAuthProvider,
type McpOAuthState,
type McpOAuthStateStore,
OAuthCallbackServer,
type OAuthChallenge,
type OAuthClientInformationMixed,
parseWwwAuthenticate,
} from "@earendil-works/pi-mcp/oauth";
import { APP_NAME, getAgentDir } from "../../config.ts";
import { type AuthStorageBackend, FileAuthStorageBackend } from "../../core/auth-storage.ts";
const CALLBACK_HOST = "127.0.0.1";
const CALLBACK_PATH = "/callback";
/** Redirect URI for refreshes when none is stored. Refreshing never redirects the user. */
const FALLBACK_REDIRECT_URL = `http://${CALLBACK_HOST}${CALLBACK_PATH}`;
/** Access tokens this close to expiry are refreshed before they are sent. */
const REFRESH_SKEW_MS = 30_000;
export interface McpOAuthSettings {
clientId?: string;
/** Already resolved. */
clientSecret?: string;
callbackPort?: number;
/** Loopback redirect URI; see `McpOAuthConfig.callbackUrl`. */
callbackUrl?: string;
/** Scopes to request, separated by spaces. */
scope?: string;
}
/** Where the loopback callback server listens and the redirect URI it serves. */
interface CallbackSettings {
/** Address to listen on. */
host: string;
/** Host name in the redirect URI. */
redirectHost: string;
port: number | undefined;
path: string;
/** The exact redirect URI, when the port is fixed. */
fixedRedirectUrl: string | undefined;
}
function callbackSettings(settings: McpOAuthSettings): CallbackSettings {
const url = new URL(settings.callbackUrl ?? `http://${CALLBACK_HOST}${CALLBACK_PATH}`);
const address = url.hostname.replace(/^\[|\]$/g, "");
const port = url.port ? Number(url.port) : settings.callbackPort;
let fixedRedirectUrl: string | undefined;
// A configured URI with a port is sent exactly as written, since servers compare it as a string.
if (url.port) fixedRedirectUrl = settings.callbackUrl;
else if (port !== undefined) {
url.port = String(port);
fixedRedirectUrl = url.href;
}
return {
// `localhost` is served on 127.0.0.1; browsers fall back to it when ::1 refuses.
host: address === "localhost" ? CALLBACK_HOST : address,
redirectHost: address,
port,
path: url.pathname,
fixedRedirectUrl,
};
}
/** Scopes of both lists, each once. */
function mergeScopes(...scopes: (string | undefined)[]): string | undefined {
const merged = [...new Set(scopes.flatMap((scope) => scope?.split(/\s+/).filter(Boolean) ?? []))];
return merged.length > 0 ? merged.join(" ") : undefined;
}
type StoredStates = Record<string, McpOAuthState>;
function parseStates(content: string | undefined): StoredStates {
if (!content?.trim()) return {};
const parsed: unknown = JSON.parse(content);
return typeof parsed === "object" && parsed !== null && !Array.isArray(parsed) ? (parsed as StoredStates) : {};
}
/** Per-server OAuth state (client registration, tokens, pending PKCE verifier) in `mcp-auth.json`. */
export class McpOAuthCredentialStore {
private readonly backend: AuthStorageBackend;
constructor(backend: AuthStorageBackend = new FileAuthStorageBackend(join(getAgentDir(), "mcp-auth.json"))) {
this.backend = backend;
}
forServer(serverUrl: string): McpOAuthStateStore {
const key = String(new URL(serverUrl));
return {
load: () => this.read()[key],
save: (state) =>
this.write((states) => {
states[key] = state;
}),
};
}
/** The stored tokens of a server, for noticing sign-ins done by another process. */
tokens(serverUrl: string): McpOAuthState["tokens"] {
return this.read()[String(new URL(serverUrl))]?.tokens;
}
/** Returns whether credentials were stored for the server. */
remove(serverUrl: string): boolean {
const key = String(new URL(serverUrl));
if (!(key in this.read())) return false;
this.write((states) => {
delete states[key];
});
return true;
}
private read(): StoredStates {
return this.backend.withLock((current) => ({ result: parseStates(current) }));
}
private write(update: (states: StoredStates) => void): void {
this.backend.withLock((current) => {
const states = parseStates(current);
update(states);
return { result: undefined, next: `${JSON.stringify(states, null, 2)}\n` };
});
}
}
function registeredRedirectUrls(client: OAuthClientInformationMixed | undefined): string[] {
return client && "redirect_uris" in client ? client.redirect_uris : [];
}
function createProvider(
serverUrl: string,
store: McpOAuthStateStore,
settings: McpOAuthSettings,
redirectUrl: string,
onRedirect: (url: URL) => void,
): McpOAuthProvider {
return new McpOAuthProvider({
serverUrl,
redirectUrl,
clientMetadata: { client_name: APP_NAME },
clientId: settings.clientId,
clientSecret: settings.clientSecret,
store,
onRedirect,
});
}
/**
* Auth provider for MCP connections: sends the stored access token and refreshes it when it is about
* to expire or after a 401. Throws `McpOAuthAuthorizationRequiredError` when the user has to sign in,
* including when the server asks for more scope (`insufficient_scope`). `onChallenge` receives the
* server's `WWW-Authenticate` challenge so sign-in can use its resource metadata URL and scope.
* `settings` is only called when a refresh is needed, so a secret that fails to resolve fails the
* refresh instead of the whole connection setup.
*
* Concurrent requests share one refresh, and a 401 for a token that was already replaced just
* retries: many servers rotate refresh tokens, so refreshing twice with the same one would fail
* and discard the new grant.
*/
export function createMcpAuthProvider(options: {
serverUrl: string;
store: McpOAuthStateStore;
settings: () => McpOAuthSettings;
onChallenge: (challenge: OAuthChallenge) => void;
}): AuthProvider {
const { serverUrl, store } = options;
let refreshing: Promise<void> | undefined;
const refresh = (context: Parameters<NonNullable<AuthProvider["onUnauthorized"]>>[0] | undefined) => {
refreshing ??= (async () => {
const state = await store.load();
if (!state?.tokens?.refresh_token) throw new McpOAuthAuthorizationRequiredError();
const settings = options.settings();
const redirectUrl =
callbackSettings(settings).fixedRedirectUrl ??
registeredRedirectUrls(state.clientInformation)[0] ??
FALLBACK_REDIRECT_URL;
const provider = createProvider(serverUrl, store, settings, redirectUrl, () => {});
// Refreshes the tokens, or reports that a new sign-in is needed.
const result = context
? await adaptOAuthProvider(provider).onUnauthorized?.(context)
: await authorizeMcp(provider, { serverUrl });
if (result === "REDIRECT") throw new McpOAuthAuthorizationRequiredError();
})().finally(() => {
refreshing = undefined;
});
return refreshing;
};
return {
token: async () => {
await refreshing?.catch(() => undefined);
const state = await store.load();
const expired = state?.tokensExpireAt !== undefined && state.tokensExpireAt - REFRESH_SKEW_MS <= Date.now();
if (!expired || !state?.tokens?.refresh_token) return state?.tokens?.access_token;
// Failures fall through: the request goes out with the old token and a 401 decides what happens.
await refresh(undefined).catch(() => undefined);
return (await store.load())?.tokens?.access_token;
},
onUnauthorized: async (context) => {
const challenge = parseWwwAuthenticate(context.response.headers.get("www-authenticate"));
options.onChallenge(challenge);
// A refresh keeps the granted scope, so more scope needs a new sign-in.
if (challenge.error === "insufficient_scope") throw new McpOAuthAuthorizationRequiredError();
if (!refreshing && context.token !== undefined) {
const current = (await store.load())?.tokens?.access_token;
if (current !== undefined && current !== context.token) return;
}
await refresh(context);
},
};
}
export interface McpSignInPrompt {
/** Show the authorization URL to the user and open it in a browser. */
showAuthorizationUrl(url: URL): void;
/**
* Ask for the redirect URL from the browser address bar, for when the browser cannot reach the
* loopback callback (for example over SSH). Aborted once the callback arrives. Resolves to
* `undefined` or an empty string when the user cancels.
*/
promptForRedirectUrl(signal: AbortSignal): Promise<string | undefined>;
}
export class McpSignInCancelledError extends Error {
constructor() {
super("Sign-in cancelled");
this.name = "McpSignInCancelledError";
}
}
function codeFromRedirectUrl(input: string, state: string): string {
let url: URL;
try {
url = new URL(input.trim());
} catch {
throw new Error("Expected the full redirect URL from the browser address bar");
}
const error = url.searchParams.get("error");
if (error) throw new Error(url.searchParams.get("error_description") ?? error);
if (url.searchParams.get("state") !== state) throw new Error("The redirect URL belongs to a different sign-in");
const code = url.searchParams.get("code");
if (!code) throw new Error("The redirect URL does not contain an authorization code");
return code;
}
/** Wait for the browser callback or a pasted redirect URL, whichever comes first. */
async function waitForAuthorizationCode(
callback: OAuthCallbackServer,
state: string,
prompt: McpSignInPrompt,
): Promise<string> {
const controller = new AbortController();
const fromBrowser = callback.waitForCallback(state).then((result) => result.code);
const fromUser = prompt.promptForRedirectUrl(controller.signal).then((input) => {
if (!input?.trim()) throw new McpSignInCancelledError();
return codeFromRedirectUrl(input, state);
});
try {
return await Promise.race([fromBrowser, fromUser]);
} finally {
// The losing side rejects once the prompt is aborted or the callback server closes.
controller.abort();
fromBrowser.catch(() => undefined);
fromUser.catch(() => undefined);
}
}
/** Listen on `port`, or on a free port when it is taken and not `required`. */
async function listenForCallback(
settings: CallbackSettings,
port: number | undefined,
required: boolean,
): Promise<OAuthCallbackServer> {
const options = { host: settings.host, redirectHost: settings.redirectHost, path: settings.path };
try {
return await OAuthCallbackServer.listen({ ...options, port: port ?? 0 });
} catch (error) {
if (required || port === undefined) throw error;
return OAuthCallbackServer.listen(options);
}
}
/**
* Sign in to an MCP server. Uses the stored refresh token when possible; otherwise runs the browser
* authorization code flow. Tokens are saved to `store`.
*/
export async function signInMcpServer(options: {
serverUrl: string;
store: McpOAuthStateStore;
settings: McpOAuthSettings;
challenge?: OAuthChallenge;
prompt: McpSignInPrompt;
}): Promise<void> {
const { serverUrl, store, settings } = options;
const stored = await store.load();
const callbackOptions = callbackSettings(settings);
// Reuse the port of the registered redirect URI so the registered client stays valid.
const registered = registeredRedirectUrls(stored?.clientInformation)[0];
const preferredPort =
callbackOptions.port ?? (registered ? Number(new URL(registered).port) || undefined : undefined);
const callback = await listenForCallback(callbackOptions, preferredPort, callbackOptions.port !== undefined);
const redirectUrl = callbackOptions.fixedRedirectUrl ?? callback.redirectUrl;
try {
if (stored) {
const next: McpOAuthState = { ...stored };
// Every sign-in gets a fresh `state` parameter.
delete next.oauthState;
// A registered client cannot use another redirect URI, and its tokens belong to it.
if (!settings.clientId && !registeredRedirectUrls(stored.clientInformation).includes(redirectUrl)) {
delete next.clientInformation;
delete next.tokens;
delete next.tokensExpireAt;
}
await store.save(next);
}
let authorizationUrl: URL | undefined;
const provider = createProvider(serverUrl, store, settings, redirectUrl, (url) => {
authorizationUrl = url;
});
const flow = {
serverUrl,
resourceMetadataUrl: options.challenge?.resourceMetadataUrl,
// A server asking for more scope gets it on top of the configured scope.
scope: mergeScopes(settings.scope, options.challenge?.scope),
};
// A refresh keeps the granted scope; a server asking for more needs the browser flow.
const skipRefresh = options.challenge?.error === "insufficient_scope";
if ((await authorizeMcp(provider, { ...flow, skipRefresh })) === "AUTHORIZED") return;
if (!authorizationUrl) throw new Error("OAuth flow did not produce an authorization URL");
const state = await provider.state();
options.prompt.showAuthorizationUrl(authorizationUrl);
const code = await waitForAuthorizationCode(callback, state, options.prompt);
await authorizeMcp(provider, { ...flow, authorizationCode: code });
} finally {
await callback.close();
}
}

View file

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

View file

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

View file

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

View file

@ -0,0 +1,284 @@
/**
* Adapts MCP tools to pi tool definitions.
*
* Results map onto pi's model-facing content (text and images). Text over 20KB keeps its start and
* end with the middle cut out, like Codex does, and the full text is saved to a temp file the model
* can read. Binary resources other than images are saved to temp files too, and resource links name
* the `read_mcp_resource` tool. Codemode scripts receive the whole `CallToolResult` without `_meta`
* (`content` blocks as sent by the server, `structuredContent`, `isError`), never truncated: it is
* the tool's `structuredContent`, and every MCP tool declares a `CallToolResult` output schema. MCP
* errors (`isError`) are error results for the model, but scripts still resolve to the result.
*/
import { createHash, randomBytes } from "node:crypto";
import { writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { AgentToolResult } from "@earendil-works/pi-agent-core";
import type { ImageContent, JsonValue, TextContent } from "@earendil-works/pi-ai";
import {
type CallToolResult,
type ContentBlock,
type McpRequestOptions,
type Tool as McpTool,
toLlmContent,
} from "@earendil-works/pi-mcp";
import type { TSchema } from "typebox";
import type { ToolAnnotations, ToolDefinition, ToolExposure, ToolNamespace } from "../../core/extensions/types.ts";
import { formatSize, truncateMiddle } from "../../core/tools/truncate.ts";
import type { McpExposure } from "./config.ts";
/**
* Tool exposure of an MCP exposure. Both deferred exposures leave tools out of the codemode
* description; they differ only in which tool the MCP extension activates to reach them.
*/
export function toToolExposure(exposure: McpExposure): ToolExposure {
return exposure === "codemode-deferred" ? "deferred" : exposure;
}
/** Provider tool names are limited to 64 characters of `[A-Za-z0-9_-]`. */
const MAX_TOOL_NAME_LENGTH = 64;
/** Model-facing text of an MCP result beyond this is cut in the middle. */
export const MCP_OUTPUT_MAX_BYTES = 20 * 1024;
/** Tool that reads the resources named by resource links. */
export const READ_MCP_RESOURCE_TOOL = "read_mcp_resource";
export interface McpToolDetails {
server: string;
tool: string;
/** Temp file with the full text output, when the model-facing text was truncated. */
fullOutputPath?: string;
}
/**
* Saves the full text of a truncated result, or a binary resource, and returns the file path.
* `extension` includes the dot, for example `.txt`.
*/
export type McpOutputSaver = (data: string | Uint8Array, extension: string) => Promise<string>;
export async function saveToTempFile(data: string | Uint8Array, extension: string): Promise<string> {
const path = join(tmpdir(), `pi-mcp-${randomBytes(8).toString("hex")}${extension}`);
// Results can carry private data, so only the user may read the file.
await writeFile(path, data, { mode: 0o600 });
return path;
}
export interface McpToolCaller {
callTool(name: string, args: Record<string, unknown>, options: McpRequestOptions): Promise<CallToolResult>;
}
/**
* `mcp__<server>__<tool>`, sanitized and shortened with a hash suffix when too long. `isTaken`
* reports names already used by a different MCP tool: sanitizing can map two tools to one name
* (`a.b` and `a_b`), and the second then gets the hash suffix too.
*/
export function createMcpToolName(
server: string,
tool: string,
isTaken: (name: string) => boolean = () => false,
): string {
const name = `mcp__${server}__${tool}`.replace(/[^A-Za-z0-9_-]/g, "_");
if (name.length <= MAX_TOOL_NAME_LENGTH && !isTaken(name)) return name;
const hash = createHash("sha256").update(`${server}\0${tool}`).digest("hex").slice(0, 8);
return `${name.slice(0, MAX_TOOL_NAME_LENGTH - hash.length - 1)}_${hash}`;
}
function textOf(content: readonly (TextContent | ImageContent)[]): string {
return content
.filter((block): block is TextContent => block.type === "text")
.map((block) => block.text)
.join("\n");
}
/**
* Output schema of every MCP tool: the `CallToolResult` scripts receive, with the tool's own output
* schema as `structuredContent`. Codemode detects this shape to render `CallToolResult<T>`
* declarations.
*/
export function createMcpResultSchema(structuredContentSchema: Record<string, unknown> | undefined): TSchema {
return {
type: "object",
properties: {
content: { type: "array", items: { type: "object" } },
...(structuredContentSchema ? { structuredContent: structuredContentSchema } : {}),
isError: { type: "boolean" },
_meta: { type: "object" },
},
required: ["content"],
} as unknown as TSchema;
}
/**
* Keep model-facing text within {@link MCP_OUTPUT_MAX_BYTES}. Longer text becomes one text block in
* Codex's truncation format, followed by the path of the file with the full text; images follow it.
*/
export async function limitMcpContent(
content: (TextContent | ImageContent)[],
saveOutput: McpOutputSaver = saveToTempFile,
): Promise<{ content: (TextContent | ImageContent)[]; fullOutputPath?: string }> {
const combined = textOf(content);
const truncation = truncateMiddle(combined, MCP_OUTPUT_MAX_BYTES);
if (!truncation.truncated) return { content };
let fullOutputPath: string | undefined;
let where: string;
try {
fullOutputPath = await saveOutput(combined, ".txt");
where = `[Full output: ${fullOutputPath} (read it with offset/limit)]`;
} catch (error) {
where = `[Could not save the full output: ${error instanceof Error ? error.message : String(error)}]`;
}
const tokens = Math.ceil(truncation.totalBytes / 4);
const text = `Warning: truncated output (original token count: ${tokens})\nTotal output lines: ${truncation.totalLines}\n\n${truncation.content}\n\n${where}`;
return {
content: [{ type: "text", text }, ...content.filter((block) => block.type === "image")],
...(fullOutputPath ? { fullOutputPath } : {}),
};
}
export interface ConvertMcpResultOptions {
/** Saves truncated text and binary resources. Default: a temp file. */
saveOutput?: McpOutputSaver;
/** Whether the server's resources can be read with `read_mcp_resource`, which resource links then name. */
readableResources?: boolean;
}
/** File extension for a saved binary resource: the one its URI ends in, else `.bin`. */
function extensionOf(uri: string): string {
const path = URL.canParse(uri) ? new URL(uri).pathname : uri;
return /\.[A-Za-z0-9]{1,8}$/.exec(path)?.[0] ?? ".bin";
}
/** Blobs of these types are shown as text. */
function isTextMimeType(mimeType: string | undefined): boolean {
if (!mimeType) return false;
const type = mimeType.split(";", 1)[0].trim().toLowerCase();
return type.startsWith("text/") || type === "application/json" || type.endsWith("+json") || type.endsWith("+xml");
}
/** Model-facing content of one block of `server`'s result. */
async function blockToContent(
server: string,
block: ContentBlock,
options: ConvertMcpResultOptions,
): Promise<(TextContent | ImageContent)[]> {
if (block.type === "resource_link") {
const details = [block.mimeType, block.size === undefined ? undefined : formatSize(block.size)].filter(Boolean);
const read = options.readableResources ? `. Read it with ${READ_MCP_RESOURCE_TOOL} (server "${server}")` : "";
const description = block.description ? `: ${block.description}` : "";
return [
{
type: "text",
text: `[Resource ${block.uri} "${block.title ?? block.name}"${details.length > 0 ? ` (${details.join(", ")})` : ""}${description}${read}]`,
},
];
}
if (block.type === "resource" && "blob" in block.resource && !block.resource.mimeType?.startsWith("image/")) {
const { uri, mimeType, blob } = block.resource;
const data = Buffer.from(blob, "base64");
if (isTextMimeType(mimeType)) return [{ type: "text", text: data.toString("utf8") }];
const kind = `${mimeType ?? "unknown type"}, ${formatSize(data.length)}`;
try {
const path = await (options.saveOutput ?? saveToTempFile)(data, extensionOf(uri));
return [{ type: "text", text: `[Binary resource ${uri} (${kind}) saved to ${path}]` }];
} catch (error) {
const reason = error instanceof Error ? error.message : String(error);
return [{ type: "text", text: `[Binary resource ${uri} (${kind}) could not be saved: ${reason}]` }];
}
}
return toLlmContent({ content: [block] });
}
/** Model-facing content of `server`'s content blocks, before the output limit. */
export async function toModelContent(
server: string,
blocks: readonly ContentBlock[],
options: ConvertMcpResultOptions = {},
): Promise<(TextContent | ImageContent)[]> {
return (await Promise.all(blocks.map((block) => blockToContent(server, block, options)))).flat();
}
/** Convert an MCP result. `isError` results become error results that keep the structured result. */
export async function convertMcpResult(
server: string,
tool: string,
result: CallToolResult,
options: ConvertMcpResultOptions = {},
): Promise<AgentToolResult<McpToolDetails>> {
// Without content blocks, toLlmContent falls back to the structured content as JSON.
const converted: (TextContent | ImageContent)[] =
result.content.length > 0 ? await toModelContent(server, result.content, options) : toLlmContent(result);
if (result.isError && textOf(converted) === "") {
converted.push({ type: "text", text: `MCP tool ${server}/${tool} returned an error` });
}
const { content, fullOutputPath } = await limitMcpContent(converted, options.saveOutput);
const { _meta: _ignored, ...scriptResult } = result;
return {
content,
details: { server, tool, ...(fullOutputPath ? { fullOutputPath } : {}) },
structuredContent: scriptResult as unknown as JsonValue,
...(result.isError ? { isError: true } : {}),
};
}
/**
* Tool input schemas must be objects. MCP servers may omit `type`, and some providers reject object
* schemas without `properties`.
*/
function toParameters(schema: Record<string, unknown>): TSchema {
return {
...schema,
type: schema.type ?? "object",
...(schema.properties === undefined ? { properties: {} } : {}),
} as unknown as TSchema;
}
const ANNOTATION_HINTS = ["readOnlyHint", "destructiveHint", "idempotentHint", "openWorldHint"] as const;
/** The boolean hints of an MCP tool's annotations, or undefined when it has none. */
function toToolAnnotations(tool: McpTool): ToolAnnotations | undefined {
const annotations: ToolAnnotations = {};
for (const hint of ANNOTATION_HINTS) {
const value = tool.annotations?.[hint];
if (typeof value === "boolean") annotations[hint] = value;
}
return Object.keys(annotations).length > 0 ? annotations : undefined;
}
export function createMcpToolDefinition(options: {
server: string;
tool: McpTool;
name: string;
exposure: McpExposure;
namespace: ToolNamespace;
timeoutMs: number;
getClient: () => Promise<McpToolCaller>;
/** Whether `read_mcp_resource` can read the server's resources. */
readableResources?: () => boolean;
}): ToolDefinition<TSchema, McpToolDetails> {
const { server, tool } = options;
const title = tool.title ?? tool.annotations?.title;
const annotations = toToolAnnotations(tool);
return {
name: options.name,
label: `${server}/${tool.name}`,
description: tool.description?.trim() || title || `MCP tool ${tool.name} from server ${server}`,
parameters: toParameters(tool.inputSchema),
outputSchema: createMcpResultSchema(tool.outputSchema),
exposure: toToolExposure(options.exposure),
namespace: options.namespace,
...(annotations ? { annotations } : {}),
async execute(_toolCallId, params, signal, onUpdate) {
const client = await options.getClient();
const result = await client.callTool(tool.name, (params ?? {}) as Record<string, unknown>, {
signal,
timeoutMs: options.timeoutMs,
onProgress: (progress) => {
const total = progress.total === undefined ? "" : `/${progress.total}`;
const text = progress.message ?? `Progress ${progress.progress}${total}`;
onUpdate?.({ content: [{ type: "text", text }], details: { server, tool: tool.name } });
},
});
return convertMcpResult(server, tool.name, result, { readableResources: options.readableResources?.() });
},
};
}

View file

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

View file

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

View file

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

View file

@ -90,6 +90,7 @@ export type {
EntryRenderOptions,
ExecOptions,
ExecResult,
ExecuteToolOptions,
Extension,
ExtensionActions,
ExtensionAPI,
@ -104,6 +105,7 @@ export type {
ExtensionHandler,
ExtensionRuntime,
ExtensionShortcut,
ExtensionToolContext,
ExtensionUIContext,
ExtensionUIDialogOptions,
ExtensionVirtualModel,
@ -119,6 +121,7 @@ export type {
LsToolCallEvent,
MarkdownTransformContext,
MarkdownTransformer,
McpServersChangeEvent,
MessageEndEvent,
MessageEndEventResult,
MessageRenderer,
@ -163,6 +166,7 @@ export type {
SourceInfo,
TerminalInputHandler,
ThinkingLevelSelectEvent,
ToolAnnotations,
ToolCallEvent,
ToolCallEventResult,
ToolDefinition,
@ -170,7 +174,11 @@ export type {
ToolExecutionMode,
ToolExecutionStartEvent,
ToolExecutionUpdateEvent,
ToolExposure,
ToolInfo,
ToolLoadout,
ToolLoadoutChanges,
ToolNamespace,
ToolRenderResultOptions,
ToolResultEvent,
ToolResultEventResult,
@ -205,6 +213,7 @@ export {
} from "./core/extensions/index.ts";
// Footer data provider (git branch + extension statuses - data not otherwise available to extensions)
export type { ReadonlyFooterDataProvider } from "./core/footer-data-provider.ts";
export type { RegisteredMcpServer } from "./core/mcp-servers.ts";
export { convertToLlm } from "./core/messages.ts";
export { ModelRegistry } from "./core/model-registry.ts";
export {
@ -391,6 +400,12 @@ export {
type VirtualModelDefinition,
type VirtualModelStateData,
} from "./core/virtual-models.ts";
// Built-in extensions. The CLI loads them; SDK sessions add them to their extension factories.
export { type CodemodeExtensionOptions, createCodemodeExtension } from "./extensions/codemode/index.ts";
export type { CodemodeStoreEntryData, CodemodeToolDetails } from "./extensions/codemode/tool.ts";
export type { LoadedMcpConfig, McpExposure, McpServerConfig, McpServerEntry } from "./extensions/mcp/config.ts";
export { createMcpExtension, type McpExtensionOptions, type McpTransportFactory } from "./extensions/mcp/index.ts";
export { createToolSearchExtension } from "./extensions/tool-search/index.ts";
// Main entry point
export { type MainOptions, main } from "./main.ts";
// Run modes for programmatic SDK usage

View file

@ -62,6 +62,7 @@ import { SettingsManager } from "./core/settings-manager.ts";
import { printTimings, resetTimings, time } from "./core/timings.ts";
import { hasTrustRequiringProjectResources, ProjectTrustStore } from "./core/trust-manager.ts";
import { builtInExtensions } from "./extensions/index.ts";
import { loadMcpCommand } from "./extensions/mcp/cli.lazy.ts";
import { runMigrations, showDeprecationWarnings } from "./migrations.ts";
import { InteractiveMode, runPrintMode, runRpcMode } from "./modes/index.ts";
import { initTheme, setThemeJsonValidator, stopThemeWatcher } from "./modes/interactive/theme/theme.ts";
@ -604,6 +605,12 @@ export async function main(args: string[], options?: MainOptions) {
return;
}
if (args[0] === "mcp") {
const { runMcpCommand } = await loadMcpCommand();
process.exitCode = await runMcpCommand(args.slice(1), { cwd, agentDir });
return;
}
const parsed = parseArgs(args);
if (parsed.diagnostics.length > 0) {
for (const d of parsed.diagnostics) {

View file

@ -3495,6 +3495,8 @@ export class InteractiveMode {
break;
case "tool_execution_start": {
// Nested calls (from codemode scripts) are shown inside their parent's row.
if (event.parentToolCallId) break;
let component = this.pendingTools.get(event.toolCallId);
if (!component) {
component = new ToolExecutionComponent(

View file

@ -8,6 +8,7 @@ import go from "highlight.js/lib/languages/go.js";
import groovy from "highlight.js/lib/languages/groovy.js";
import java from "highlight.js/lib/languages/java.js";
import javascript from "highlight.js/lib/languages/javascript.js";
import json from "highlight.js/lib/languages/json.js";
import kotlin from "highlight.js/lib/languages/kotlin.js";
import lua from "highlight.js/lib/languages/lua.js";
import nix from "highlight.js/lib/languages/nix.js";
@ -26,6 +27,7 @@ const eagerLanguages = {
java,
go,
javascript,
json,
cpp,
typescript,
php,

View file

@ -0,0 +1,62 @@
import { stripVTControlCharacters } from "node:util";
import type { AgentToolResult } from "@earendil-works/pi-agent-core";
import type { Text } from "@earendil-works/pi-tui";
import { beforeAll, describe, expect, it } from "vitest";
import type { ToolRenderContext } from "../src/core/extensions/types.ts";
import { codemodeRenderers } from "../src/extensions/codemode/renderer.ts";
import type { CodemodeToolDetails } from "../src/extensions/codemode/tool.ts";
import { initTheme, theme } from "../src/modes/interactive/theme/theme.ts";
function render(result: AgentToolResult<CodemodeToolDetails | undefined>, isError = false): string {
const context = {
args: { code: "" },
toolCallId: "call",
invalidate: () => {},
lastComponent: undefined,
state: {},
cwd: "/",
executionStarted: true,
argsComplete: true,
isPartial: false,
expanded: true,
showImages: false,
isError,
} satisfies ToolRenderContext;
const component = codemodeRenderers.renderResult?.(
result,
{ expanded: true, isPartial: false },
theme,
context,
) as Text;
return stripVTControlCharacters(component.render(200).join("\n"))
.split("\n")
.map((line) => line.trimEnd())
.join("\n")
.trim();
}
describe("codemode renderer", () => {
beforeAll(() => initTheme("dark"));
it("hides the script header and shows the output", () => {
const text = render({
content: [
{ type: "text", text: "Script completed\nWall time 0.1 seconds\nOutput:\n" },
{ type: "text", text: "hello" },
],
details: { calls: [{ id: "call/1", name: "read", args: '{"path":"a"}', status: "ok", durationMs: 5 }] },
});
expect(text).toBe('✓ read {"path":"a"} 5ms\n\nhello');
});
it("shows results without a header, such as rejected options", () => {
const text = render(
{
content: [{ type: "text", text: "The @options line must be followed by JavaScript source" }],
details: undefined,
},
true,
);
expect(text).toBe("The @options line must be followed by JavaScript source");
});
});

View file

@ -0,0 +1,27 @@
import type { ToolResultMessage } from "@earendil-works/pi-ai/compat";
import { describe, expect, it } from "vitest";
import { computeFileLists, createFileOps, extractFileOpsFromMessage } from "../src/core/compaction/utils.ts";
describe("compaction file operations", () => {
it("include files touched by nested calls recorded on tool results", () => {
const result: ToolResultMessage = {
role: "toolResult",
toolCallId: "codemode-1",
toolName: "codemode",
content: [],
isError: false,
timestamp: 0,
nestedCalls: {
calls: [
{ id: "codemode-1/1", name: "read", arguments: { path: "a.ts" }, status: "ok" },
{ id: "codemode-1/2", name: "edit", arguments: { path: "b.ts", edits: [] }, status: "ok" },
{ id: "codemode-1/3", name: "write", argumentsBytes: 40000, status: "ok" },
],
complete: false,
},
};
const fileOps = createFileOps();
extractFileOpsFromMessage(result, fileOps);
expect(computeFileLists(fileOps)).toEqual({ readFiles: ["a.ts"], modifiedFiles: ["b.ts"] });
});
});

View file

@ -2,7 +2,7 @@ import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { afterEach, describe, expect, it } from "vitest";
import type { ExtensionContext } from "../src/core/extensions/types.ts";
import type { ExtensionToolContext } from "../src/core/extensions/types.ts";
import { createEditToolDefinition } from "../src/core/tools/edit.ts";
const tempDirs: string[] = [];
@ -83,7 +83,7 @@ describe("edit tool prepareArguments", () => {
newText: "after",
});
const result = await definition.execute("tool-1", prepared, undefined, undefined, {} as ExtensionContext);
const result = await definition.execute("tool-1", prepared, undefined, undefined, {} as ExtensionToolContext);
expect(result.content).toEqual([{ type: "text", text: "Successfully replaced 1 block(s) in legacy.txt." }]);
expect(await readFile(filePath, "utf8")).toBe("after\n");
});

View file

@ -88,6 +88,7 @@ describe("ExtensionRunner", () => {
setLabel: () => {},
getActiveTools: () => [],
getAllTools: () => [],
getSettings: () => ({}),
setActiveTools: () => {},
refreshTools: () => {},
getCommands: () => [],

View file

@ -10,7 +10,7 @@ vi.mock("../src/utils/image-resize.js", () => ({
}));
import { processFileArguments } from "../src/cli/file-processor.ts";
import type { ExtensionContext } from "../src/core/extensions/types.ts";
import type { ExtensionToolContext } from "../src/core/extensions/types.ts";
import { createReadTool, createReadToolDefinition } from "../src/core/tools/read.ts";
import { resizeImage } from "../src/utils/image-resize.ts";
@ -70,7 +70,7 @@ describe("image resize callers", () => {
contextWindow: 1000,
maxTokens: 100,
};
const ctx = { cwd: testDir, model } as unknown as ExtensionContext;
const ctx = { cwd: testDir, model } as unknown as ExtensionToolContext;
await createReadToolDefinition(testDir).execute(
"test-read-model-profile",

View file

@ -0,0 +1,185 @@
import { existsSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join, resolve } from "node:path";
import { afterEach, describe, expect, it } from "vitest";
import { runMcpCommand } from "../src/extensions/mcp/cli.ts";
const FIXTURE = resolve(import.meta.dirname, "../../mcp/test/fixtures/stdio-server.mjs");
describe("pi mcp", () => {
const dirs: string[] = [];
afterEach(() => {
while (dirs.length > 0) rmSync(dirs.pop() ?? "", { recursive: true, force: true });
});
async function run(args: string[], servers: Record<string, unknown> | undefined, dir?: string) {
const agentDir = dir ?? mkdtempSync(join(tmpdir(), "pi-mcp-command-"));
if (!dir) dirs.push(agentDir);
if (servers) writeFileSync(join(agentDir, "mcp.json"), JSON.stringify({ mcpServers: servers }));
const output: string[] = [];
const exitCode = await runMcpCommand(args, {
cwd: agentDir,
agentDir,
log: (line) => output.push(line),
error: (line) => output.push(line),
});
return { exitCode, output: output.join("\n"), agentDir };
}
const readConfig = (path: string) => JSON.parse(readFileSync(path, "utf8")) as Record<string, unknown>;
const servers = {
fixture: { command: process.execPath, args: [FIXTURE] },
broken: { command: "pi-test-missing-mcp-server" },
parked: { command: process.execPath, args: [FIXTURE], enabled: false },
bad: { args: ["no command"] },
};
it("lists servers with their state, tools, and errors, and fails while anything is wrong", async () => {
const { exitCode, output } = await run(["list"], servers);
expect(exitCode).toBe(1);
expect(output).toContain("fixture: connected, 1 tool (codemode, global)\n");
expect(output).toContain(" tools: echo");
expect(output).toContain(
"broken: failed (codemode, global)\n pi-test-missing-mcp-server\n spawn pi-test-missing-mcp-server ENOENT",
);
expect(output).toContain("parked: disabled (codemode, global)");
expect(output).toContain("config error: ");
expect(output).toContain('server "bad" needs either "command"');
const ok = await run(["list"], { fixture: servers.fixture });
expect(ok.exitCode).toBe(0);
});
it("prints JSON for scripts", async () => {
const { exitCode, output } = await run(["list", "--json"], { fixture: servers.fixture, parked: servers.parked });
expect(exitCode).toBe(0);
const parsed = JSON.parse(output) as { servers: { name: string; state: string; tools: string[] }[] };
expect(parsed.servers.map(({ name, state, tools }) => ({ name, state, tools }))).toEqual([
{ name: "fixture", state: "connected", tools: ["echo"] },
{ name: "parked", state: "disabled", tools: [] },
]);
});
it("rejects unknown servers and servers without OAuth for login and logout", async () => {
expect(await run(["login", "nope"], servers)).toMatchObject({
exitCode: 1,
output: 'No MCP server named "nope". Configured: fixture, broken, parked.',
});
expect(await run(["logout", "fixture"], servers)).toMatchObject({
exitCode: 1,
output: 'MCP server "fixture" does not use OAuth. Only HTTP servers without an Authorization header do.',
});
expect((await run(["frobnicate"], servers)).exitCode).toBe(1);
});
it("adds stdio servers and passes options after the command through", async () => {
const added = await run(
["add", "--env", "A=1", "--env", "B=x=y", "files", "--", "npx", "-y", "server", "--root", "."],
undefined,
);
expect(added.exitCode).toBe(0);
expect(added.output).toContain('Added global MCP server "files"');
expect(readConfig(join(added.agentDir, "mcp.json"))).toEqual({
mcpServers: {
files: { command: "npx", args: ["-y", "server", "--root", "."], env: { A: "1", B: "x=y" } },
},
});
// Without `--`, options after the command belong to the command too.
const replaced = await run(["add", "files", "node", "server.js", "--port", "1"], undefined, added.agentDir);
expect(replaced.output).toContain('Replaced global MCP server "files"');
expect(readConfig(join(added.agentDir, "mcp.json"))).toEqual({
mcpServers: { files: { command: "node", args: ["server.js", "--port", "1"] } },
});
});
it("adds HTTP servers and keeps other content of the file", async () => {
const { exitCode, output, agentDir } = await run(
[
"add",
"docs",
"--url",
"https://example.com/mcp",
"--bearer-token-env-var",
"DOCS_TOKEN",
"--header",
"X-Team=core",
"--exposure",
"direct",
],
{ fixture: servers.fixture },
);
expect(exitCode).toBe(0);
expect(output).not.toContain("mcp login");
expect(readConfig(join(agentDir, "mcp.json"))).toEqual({
mcpServers: {
fixture: servers.fixture,
docs: {
url: "https://example.com/mcp",
// biome-ignore lint/suspicious/noTemplateCurlyInString: literal config value reference
headers: { "X-Team": "core", Authorization: "Bearer ${DOCS_TOKEN}" },
exposure: "direct",
},
},
});
const oauth = await run(
["add", "sentry", "--url", "https://mcp.sentry.dev/mcp", "--oauth-client-id", "pi"],
undefined,
agentDir,
);
expect(oauth.output).toContain("If it requires sign-in: pi mcp login sentry");
expect(readConfig(join(agentDir, "mcp.json")).mcpServers).toMatchObject({
sentry: { url: "https://mcp.sentry.dev/mcp", oauth: { clientId: "pi" } },
});
});
it("rejects invalid add invocations without writing", async () => {
const cases = [
["add", "x"],
["add", "x", "--url", "https://example.com", "--", "cmd"],
["add", "bad name", "--", "cmd"],
["add", "x", "--url", "ftp://example.com"],
["add", "x", "--env", "A=1", "--url", "https://example.com"],
["add", "x", "--header", "A=1", "--", "cmd"],
["add", "x", "--env", "NOVALUE", "--", "cmd"],
["add", "x", "--exposure", "loud", "--", "cmd"],
];
for (const args of cases) {
const result = await run(args, undefined);
expect(result.exitCode, args.join(" ")).toBe(1);
expect(existsSync(join(result.agentDir, "mcp.json"))).toBe(false);
}
});
it("adds and removes project servers", async () => {
const added = await run(["add", "-l", "local", "--", "node", "server.js"], undefined);
expect(added.output).toContain("The project is not trusted");
const projectConfig = join(added.agentDir, ".pi", "mcp.json");
expect(readConfig(projectConfig)).toEqual({ mcpServers: { local: { command: "node", args: ["server.js"] } } });
const wrongScope = await run(["remove", "local"], undefined, added.agentDir);
expect(wrongScope.exitCode).toBe(1);
expect(wrongScope.output).toContain(`It is defined in ${projectConfig}; use --local.`);
const removed = await run(["remove", "local", "--local"], undefined, added.agentDir);
expect(removed.exitCode).toBe(0);
expect(removed.output).toContain('Removed project MCP server "local"');
expect(readConfig(projectConfig)).toEqual({ mcpServers: {} });
});
it("removes global servers", async () => {
const { exitCode, agentDir } = await run(["remove", "broken"], servers);
expect(exitCode).toBe(0);
expect(Object.keys(readConfig(join(agentDir, "mcp.json")).mcpServers as object)).toEqual([
"fixture",
"parked",
"bad",
]);
const missing = await run(["remove", "broken"], undefined, agentDir);
expect(missing).toMatchObject({ exitCode: 1 });
expect(missing.output).toContain('No global MCP server named "broken"');
});
});

View file

@ -0,0 +1,560 @@
import { mkdirSync, mkdtempSync, readFileSync, realpathSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import {
type JsonRpcMessage,
LATEST_PROTOCOL_VERSION,
McpAuthRequiredError,
McpHttpError,
McpSessionExpiredError,
} from "@earendil-works/pi-mcp";
import { createInMemoryTransportPair, type InMemoryTransport } from "@earendil-works/pi-mcp/testing";
import { afterEach, describe, expect, it } from "vitest";
import { InMemoryAuthStorageBackend } from "../src/core/auth-storage.ts";
import { truncateMiddle } from "../src/core/tools/truncate.ts";
import { getMcpToolExposure, loadMcpConfig, type McpServerEntry } from "../src/extensions/mcp/config.ts";
import {
createDefaultTransport,
McpOAuthCredentialStore,
McpServerConnection,
McpServerLog,
} from "../src/extensions/mcp/runtime.ts";
import { convertMcpResult, createMcpToolName } from "../src/extensions/mcp/tools.ts";
// Config values are resolved at connect time, so the literal reference must survive loading.
// biome-ignore lint/suspicious/noTemplateCurlyInString: literal config value reference
const TOKEN_HEADER = "Bearer ${TOKEN}";
describe("MCP config", () => {
const dirs: string[] = [];
afterEach(() => {
for (const dir of dirs.splice(0)) rmSync(dir, { recursive: true, force: true });
});
function setup(global: unknown, project: unknown) {
const root = mkdtempSync(join(tmpdir(), "pi-mcp-config-"));
dirs.push(root);
const agentDir = join(root, "agent");
const cwd = join(root, "project");
mkdirSync(agentDir, { recursive: true });
mkdirSync(join(cwd, ".pi"), { recursive: true });
writeFileSync(join(agentDir, "mcp.json"), JSON.stringify(global));
writeFileSync(join(cwd, ".pi", "mcp.json"), JSON.stringify(project));
return { agentDir, cwd };
}
it("merges global and trusted project servers and validates entries", () => {
const paths = setup(
{
mcpServers: {
shared: { command: "global-cmd" },
remote: { url: "https://example.com/mcp", headers: { Authorization: TOKEN_HEADER } },
off: { command: "x", enabled: false },
bad: { args: ["no command"] },
legacy: { type: "sse", url: "https://example.com/sse" },
badUrl: { url: "example.com/mcp" },
"bad name": { command: "x" },
},
},
{ mcpServers: { shared: { command: "project-cmd", exposure: "direct" } } },
);
const trusted = loadMcpConfig({ ...paths, projectTrusted: true });
// Disabled servers are kept so /mcp can enable them again.
expect(trusted.servers.map((server) => [server.name, server.scope, server.config])).toEqual([
["shared", "project", { command: "project-cmd", exposure: "direct" }],
["remote", "global", { url: "https://example.com/mcp", headers: { Authorization: TOKEN_HEADER } }],
["off", "global", { command: "x", enabled: false }],
]);
expect(trusted.errors).toHaveLength(4);
expect(trusted.errors[0]).toContain('server "bad" needs either "command"');
expect(trusted.errors[1]).toContain("legacy SSE transport is not supported");
expect(trusted.errors[2]).toContain('server "badUrl": url must be an http or https URL');
expect(trusted.errors[3]).toContain('invalid server name "bad name"');
// Untrusted projects cannot add or override servers, since stdio servers run commands.
const untrusted = loadMcpConfig({ ...paths, projectTrusted: false });
expect(untrusted.servers.find((server) => server.name === "shared")?.config).toEqual({ command: "global-cmd" });
});
it("validates exposure and reads autoEnableCodemode with project precedence", () => {
const paths = setup(
{
autoEnableCodemode: false,
mcpServers: {
later: { command: "x", exposure: "deferred" },
scripts: { command: "x", exposure: "codemode-deferred" },
off: { command: "x", exposure: "hidden" },
wrong: { command: "x", exposure: "model-only" },
},
},
{ autoEnableCodemode: "yes", mcpServers: {} },
);
const untrusted = loadMcpConfig({ ...paths, projectTrusted: false });
expect(untrusted.autoEnableCodemode).toBe(false);
expect(untrusted.servers.map((server) => [server.name, server.config.exposure])).toEqual([
["later", "deferred"],
["scripts", "codemode-deferred"],
["off", "hidden"],
]);
expect(untrusted.errors).toEqual([expect.stringContaining('server "wrong": exposure must be one of')]);
const trusted = loadMcpConfig({ ...paths, projectTrusted: true });
expect(trusted.autoEnableCodemode).toBe(false);
expect(trusted.errors).toContainEqual(expect.stringContaining("autoEnableCodemode must be a boolean"));
});
it("validates the OAuth callback URL and scope", () => {
const paths = setup(
{
mcpServers: {
ok: {
url: "https://a.example/mcp",
oauth: { callbackUrl: "http://localhost:8080/callback", scope: "a b" },
},
ipv6: { url: "https://a.example/mcp", oauth: { callbackUrl: "http://[::1]/cb", callbackPort: 9000 } },
same: { url: "https://a.example/mcp", oauth: { callbackUrl: "http://127.0.0.1:2/cb", callbackPort: 2 } },
remote: { url: "https://a.example/mcp", oauth: { callbackUrl: "https://example.com/callback" } },
both: { url: "https://a.example/mcp", oauth: { callbackUrl: "http://127.0.0.1:1/cb", callbackPort: 2 } },
scope: { url: "https://a.example/mcp", oauth: { scope: ["a"] } },
},
},
{},
);
const { servers, errors } = loadMcpConfig({ ...paths, projectTrusted: false });
expect(servers.map((server) => server.name)).toEqual(["ok", "ipv6", "same"]);
expect(errors).toEqual([
expect.stringContaining('server "remote": oauth.callbackUrl must be an http URI on localhost'),
expect.stringContaining('server "both": oauth.callbackUrl and oauth.callbackPort name different ports'),
expect.stringContaining('server "scope": oauth.scope must be a string'),
]);
});
it("resolves per-tool exposure from exact names, then patterns in order", () => {
const paths = setup(
{
mcpServers: {
gh: {
command: "x",
exposure: "deferred",
toolExposure: { "get_*": "codemode", get_me: "direct", "*delete*": "hidden", "get_file.*": "direct" },
},
bad: { command: "x", toolExposure: { a: "visible" } },
},
},
{},
);
const { servers, errors } = loadMcpConfig({ ...paths, projectTrusted: false });
expect(errors).toEqual([expect.stringContaining('server "bad": toolExposure "a" must be one of')]);
const config = servers[0].config;
expect(getMcpToolExposure(config, "get_me")).toBe("direct");
expect(getMcpToolExposure(config, "get_issue")).toBe("codemode");
expect(getMcpToolExposure(config, "get_delete_hint")).toBe("codemode");
expect(getMcpToolExposure(config, "delete_repo")).toBe("hidden");
expect(getMcpToolExposure(config, "list_issues")).toBe("deferred");
// Only `*` is special.
expect(getMcpToolExposure({ command: "x", toolExposure: { "get_file.*": "direct" } }, "get_file_x")).toBe(
"codemode",
);
});
});
describe("MCP tools", () => {
it("creates provider-safe tool names", () => {
expect(createMcpToolName("docs", "search")).toBe("mcp__docs__search");
expect(createMcpToolName("my-server", "get.item/v2")).toBe("mcp__my-server__get_item_v2");
const long = createMcpToolName("server", "x".repeat(100));
expect(long).toHaveLength(64);
expect(long).toMatch(/^mcp__server__x+_[0-9a-f]{8}$/);
expect(createMcpToolName("server", `${"x".repeat(100)}y`)).not.toBe(long);
// Names that sanitize to one already taken by another tool get a hash suffix.
const taken = createMcpToolName("s", "a_b");
const second = createMcpToolName("s", "a.b", (name) => name === taken);
expect(second).toMatch(/^mcp__s__a_b_[0-9a-f]{8}$/);
});
it("converts results, passing the CallToolResult to scripts and flagging errors", async () => {
const blocks = [
{ type: "resource_link" as const, uri: "file:///a", name: "a" },
{ type: "resource" as const, resource: { uri: "file:///b", text: "b text" } },
{ type: "audio" as const, data: "", mimeType: "audio/wav" },
];
expect(
await convertMcpResult("docs", "t", {
content: blocks,
structuredContent: { ok: true },
_meta: { trace: "x" },
}),
).toEqual({
content: [
{ type: "text", text: '[Resource file:///a "a"]' },
{ type: "text", text: "b text" },
{ type: "text", text: "[audio audio/wav omitted]" },
],
details: { server: "docs", tool: "t" },
// Scripts get the server's blocks as sent, without `_meta`.
structuredContent: { content: blocks, structuredContent: { ok: true } },
});
expect((await convertMcpResult("docs", "t", { content: [], structuredContent: { n: 1 } })).content).toEqual([
{ type: "text", text: '{\n "n": 1\n}' },
]);
expect(await convertMcpResult("docs", "t", { content: [{ type: "text", text: "nope" }], isError: true })).toEqual(
{
content: [{ type: "text", text: "nope" }],
details: { server: "docs", tool: "t" },
structuredContent: { content: [{ type: "text", text: "nope" }], isError: true },
isError: true,
},
);
expect((await convertMcpResult("docs", "t", { content: [], isError: true })).content).toEqual([
{ type: "text", text: "MCP tool docs/t returned an error" },
]);
});
it("points resource links to read_mcp_resource and saves binary resources", async () => {
const saved: [string | Uint8Array, string][] = [];
const saveOutput = async (data: string | Uint8Array, extension: string) => {
saved.push([data, extension]);
return `/tmp/saved${extension}`;
};
const converted = await convertMcpResult(
"docs",
"t",
{
content: [
{
type: "resource_link",
uri: "docs://guide",
name: "guide",
title: "The Guide",
mimeType: "text/markdown",
size: 2048,
description: "How to use it",
},
{
type: "resource",
resource: { uri: "file:///r/report.pdf", mimeType: "application/pdf", blob: "JVBERg==" },
},
{ type: "resource", resource: { uri: "docs://logo", mimeType: "image/png", blob: "AAAA" } },
],
},
{ saveOutput, readableResources: true },
);
expect(converted.content).toEqual([
{
type: "text",
text: '[Resource docs://guide "The Guide" (text/markdown, 2.0KB): How to use it. Read it with read_mcp_resource (server "docs")]',
},
{ type: "text", text: "[Binary resource file:///r/report.pdf (application/pdf, 4B) saved to /tmp/saved.pdf]" },
{ type: "image", data: "AAAA", mimeType: "image/png" },
]);
expect(saved).toEqual([[Buffer.from("%PDF"), ".pdf"]]);
});
it("cuts the middle of model-facing text over 20KB and keeps the full result for scripts", async () => {
const saved: (string | Uint8Array)[] = [];
const saveOutput = async (data: string | Uint8Array) => {
saved.push(data);
return "/tmp/full.txt";
};
const lines = Array.from({ length: 3000 }, (_, index) => `line ${index + 1}`);
const full = lines.join("\n");
const image = { type: "image" as const, data: "AAAA", mimeType: "image/png" };
const result = { content: [{ type: "text" as const, text: full }, image] };
const converted = await convertMcpResult("docs", "snapshot", result, { saveOutput });
expect(converted.content).toHaveLength(2);
const text = (converted.content[0] as { text: string }).text;
// Codex's format: a header, the start and end of the text, then the file with the full text.
expect(text).toMatch(
new RegExp(
`^Warning: truncated output \\(original token count: ${Math.ceil(full.length / 4)}\\)\nTotal output lines: 3000\n\nline 1\nline 2\n`,
),
);
expect(text).toMatch(/…\d+ chars truncated…/);
expect(text.endsWith("line 3000\n\n[Full output: /tmp/full.txt (read it with offset/limit)]")).toBe(true);
expect(Buffer.byteLength(text)).toBeLessThan(21 * 1024);
expect(converted.content[1]).toEqual(image);
expect(converted.details).toEqual({ server: "docs", tool: "snapshot", fullOutputPath: "/tmp/full.txt" });
expect(saved).toEqual([full]);
expect(converted.structuredContent).toEqual(result);
// Text within the limit is not saved.
await convertMcpResult("docs", "small", { content: [{ type: "text", text: "ok" }] }, { saveOutput });
expect(saved).toHaveLength(1);
});
it("cuts multi-byte text only at character boundaries", () => {
const text = `${"é".repeat(20_000)}end`;
const result = truncateMiddle(text, 1001);
expect(result.truncated).toBe(true);
expect(result.content).not.toContain("\uFFFD");
expect(result.content.endsWith("end")).toBe(true);
const [head, tail] = result.content.split(/…\d+ chars truncated…/);
expect(Buffer.byteLength(head)).toBeLessThanOrEqual(500);
expect(Buffer.byteLength(tail)).toBeLessThanOrEqual(501);
expect(Array.from(head).length + Array.from(tail).length + result.removedChars).toBe(Array.from(text).length);
});
});
describe("MCP connections", () => {
const servers: InMemoryTransport[] = [];
/** In-memory server that answers initialize, tools/list, and tools/call with "ok". */
function createTransport(options: { expireFirstCall?: boolean; methods?: string[]; noTools?: boolean } = {}) {
const pair = createInMemoryTransportPair();
servers.push(pair.server);
pair.server.onMessage((message) => {
if (!("id" in message) || !("method" in message)) return;
options.methods?.push(message.method);
const response: JsonRpcMessage =
message.method === "initialize"
? {
jsonrpc: "2.0",
id: message.id,
result: {
protocolVersion: LATEST_PROTOCOL_VERSION,
capabilities: options.noTools ? { prompts: {} } : { tools: {} },
serverInfo: { name: "fake", version: "1.0.0" },
},
}
: message.method === "tools/list"
? options.noTools
? { jsonrpc: "2.0", id: message.id, error: { code: -32601, message: "Method not found" } }
: { jsonrpc: "2.0", id: message.id, result: { tools: [] } }
: message.method === "resources/read"
? { jsonrpc: "2.0", id: message.id, result: { contents: [{ uri: "docs://a", text: "ok" }] } }
: { jsonrpc: "2.0", id: message.id, result: { content: [{ type: "text", text: "ok" }] } };
queueMicrotask(() => void pair.server.send(response));
});
void pair.server.start();
if (options.expireFirstCall) {
const send = pair.client.send.bind(pair.client);
// Simulates the HTTP transport's 404 for a session the server no longer knows.
pair.client.send = async (message: JsonRpcMessage) => {
if ("method" in message && message.method === "tools/call") throw new McpSessionExpiredError("gone");
return send(message);
};
}
return pair.client;
}
function connect(
entry: McpServerEntry,
transports: (() => ReturnType<typeof createTransport>)[],
log?: McpServerLog,
) {
let opened = 0;
const connection = new McpServerConnection({
entry,
cwd: process.cwd(),
createTransport: () => transports[opened++](),
credentials: new McpOAuthCredentialStore(new InMemoryAuthStorageBackend()),
log,
onTools: () => {},
});
return { connection, opened: () => opened };
}
it("starts a new session and retries once when the session expired", async () => {
const { connection, opened } = connect({ name: "fake", config: { command: "unused" }, source: "test" }, [
() => createTransport({ expireFirstCall: true }),
() => createTransport(),
]);
const results = await Promise.all([connection.callTool("echo", {}, {}), connection.callTool("echo", {}, {})]);
expect(results).toEqual([
{ content: [{ type: "text", text: "ok" }] },
{ content: [{ type: "text", text: "ok" }] },
]);
expect(opened()).toBe(2);
await connection.close();
});
it.skipIf(process.platform === "win32")(
"expands ~ in the command, arguments, and cwd of stdio servers",
async () => {
const home = mkdtempSync(join(tmpdir(), "pi-mcp-home-"));
const previousHome = process.env.HOME;
process.env.HOME = home;
mkdirSync(join(home, "work"));
// Answers every tool call with its working directory.
writeFileSync(
join(home, "server.mjs"),
`import { createInterface } from "node:readline";
for await (const line of createInterface({ input: process.stdin })) {
const message = JSON.parse(line);
if (!("id" in message)) continue;
const result = message.method === "initialize"
? { protocolVersion: "2025-06-18", capabilities: { tools: {} }, serverInfo: { name: "cwd", version: "1" } }
: message.method === "tools/list"
? { tools: [{ name: "cwd", inputSchema: { type: "object" } }] }
: { content: [{ type: "text", text: process.cwd() }] };
process.stdout.write(JSON.stringify({ jsonrpc: "2.0", id: message.id, result }) + "\\n");
}`,
);
const connection = new McpServerConnection({
entry: {
name: "home",
config: { command: process.execPath, args: ["~/server.mjs"], cwd: "~/work" },
source: "test",
},
cwd: tmpdir(),
createTransport: createDefaultTransport,
credentials: new McpOAuthCredentialStore(new InMemoryAuthStorageBackend()),
onTools: () => {},
});
try {
const result = await connection.callTool("cwd", {}, {});
expect(realpathSync((result.content[0] as { text: string }).text)).toBe(realpathSync(join(home, "work")));
} finally {
await connection.close();
process.env.HOME = previousHome;
rmSync(home, { recursive: true, force: true });
}
},
);
it("connects to servers without the tools capability without listing tools", async () => {
const methods: string[] = [];
const { connection } = connect({ name: "fake", config: { command: "unused" }, source: "test" }, [
() => createTransport({ noTools: true, methods }),
]);
await connection.getClient();
expect(connection.state).toBe("connected");
expect(connection.tools).toEqual([]);
expect(methods).toEqual(["initialize"]);
await connection.close();
});
it("marks a dropped connection and reconnects on the next call", async () => {
const { connection, opened } = connect({ name: "fake", config: { command: "unused" }, source: "test" }, [
() => createTransport(),
() => createTransport(),
]);
await connection.getClient();
await servers.at(-1)?.close();
await new Promise((resolve) => setTimeout(resolve, 0));
expect(connection.state).toBe("disconnected");
expect(connection.error).toBe("Connection closed");
expect(await connection.callTool("echo", {}, {})).toEqual({ content: [{ type: "text", text: "ok" }] });
expect(connection.state).toBe("connected");
expect(opened()).toBe(2);
await connection.close();
});
it("retries HTTP connections that fail with a transient error", async () => {
const { connection, opened } = connect(
{ name: "fake", config: { url: "http://unused.invalid", headers: { Authorization: "x" } }, source: "test" },
[
() => {
const transport = createTransport();
transport.send = async () => {
throw new McpHttpError(503, "MCP HTTP request failed with status 503");
};
return transport;
},
() => createTransport(),
],
);
await connection.getClient();
expect(connection.state).toBe("connected");
expect(opened()).toBe(2);
await connection.close();
const failing = connect(
{ name: "fake", config: { url: "http://unused.invalid", headers: { Authorization: "x" } }, source: "test" },
[
() => {
const transport = createTransport();
transport.send = async () => {
throw new McpHttpError(400, "MCP HTTP request failed with status 400: bad");
};
return transport;
},
],
);
await expect(failing.connection.getClient()).rejects.toThrow("status 400: bad");
expect(failing.connection.state).toBe("failed");
expect(failing.opened()).toBe(1);
});
it("retries resource reads, but not tool calls, after a transient HTTP error", async () => {
const transport = createTransport();
const send = transport.send.bind(transport);
const failed = new Set<string>();
// The first read and the first call fail with 502.
transport.send = async (message) => {
const method = "method" in message ? message.method : "";
if ((method === "resources/read" || method === "tools/call") && !failed.has(method)) {
failed.add(method);
throw new McpHttpError(502, "MCP HTTP request failed with status 502");
}
return send(message);
};
const { connection } = connect({ name: "fake", config: { command: "unused" }, source: "test" }, [
() => transport,
]);
expect(await connection.readResource("docs://a", {})).toEqual({ contents: [{ uri: "docs://a", text: "ok" }] });
await expect(connection.callTool("echo", {}, {})).rejects.toThrow("status 502");
await connection.close();
});
it("asks OAuth servers that keep rejecting requests for a new sign-in", async () => {
const { connection } = connect({ name: "fake", config: { url: "http://unused.invalid" }, source: "test" }, [
() => {
const transport = createTransport();
transport.send = async () => {
throw new McpAuthRequiredError(new Response(null, { status: 401 }));
};
return transport;
},
]);
await expect(connection.getClient()).rejects.toThrow('MCP server "fake" requires sign-in. Run /mcp to sign in.');
expect(connection.state).toBe("needs-auth");
await connection.close();
});
it("appends server log messages to the log file", async () => {
const dir = mkdtempSync(join(tmpdir(), "pi-mcp-log-"));
try {
const path = join(dir, "mcp.log");
const { connection } = connect(
{ name: "fake", config: { command: "unused" }, source: "test" },
[() => createTransport()],
new McpServerLog(path),
);
await connection.getClient();
const server = servers.at(-1);
await server?.send({
jsonrpc: "2.0",
method: "notifications/message",
params: { level: "warning", logger: "db", data: "slow\nquery" },
});
await server?.send({
jsonrpc: "2.0",
method: "notifications/message",
params: { level: "error", data: { code: 7 } },
});
await new Promise((resolve) => setTimeout(resolve, 0));
const lines = readFileSync(path, "utf8").replace(/^\S+ /gm, "");
expect(lines).toBe('[fake] warning db: slow\n query\n[fake] error {"code":7}\n');
await connection.close();
} finally {
rmSync(dir, { recursive: true, force: true });
}
});
it("resolves the OAuth client secret lazily", async () => {
const { connection } = connect(
{
name: "fake",
config: { url: "http://unused.invalid", oauth: { clientSecret: "!exit 1" } },
source: "test",
},
[() => createTransport()],
);
expect(await connection.callTool("echo", {}, {})).toEqual({ content: [{ type: "text", text: "ok" }] });
expect(() => connection.oauthSettings()).toThrow("oauth.clientSecret");
await connection.close();
});
});

View file

@ -0,0 +1,185 @@
import type { AgentTool, AgentToolCall } from "@earendil-works/pi-agent-core";
import { Type } from "typebox";
import { describe, expect, it } from "vitest";
import {
NESTED_CALL_LIMITS,
NestedCallRecorder,
type NestedToolCallHost,
NestedToolCallRunner,
type NestedToolExecutionEvent,
} from "../src/core/nested-tool-calls.ts";
function createRunner(tools: AgentTool[], options: { sequential?: boolean } = {}) {
const events: NestedToolExecutionEvent[] = [];
const host: NestedToolCallHost = {
getTools: () => tools,
isSequential: () => options.sequential ?? false,
runToolCall: async (toolCall, _parentId, signal, onUpdate) => {
const tool = tools.find((candidate) => candidate.name === toolCall.name);
if (!tool) {
return {
toolCall,
result: { content: [{ type: "text", text: `Tool ${toolCall.name} not found` }], details: {} },
isError: true,
};
}
const result = await tool.execute(
toolCall.id,
toolCall.arguments,
signal,
(partial) => void onUpdate(partial),
);
return { toolCall, result, isError: result.isError === true };
},
emit: async (event) => {
events.push(event);
},
};
return { runner: new NestedToolCallRunner(host), events };
}
describe("NestedToolCallRunner", () => {
it("assigns ids below the caller, emits events with the parent id, and records the calls", async () => {
const echo: AgentTool = {
name: "echo",
label: "Echo",
description: "Echo",
parameters: Type.Object({}),
async execute(_id, _params, _signal, onUpdate) {
onUpdate?.({ content: [{ type: "text", text: "partial" }], details: {} });
return { content: [{ type: "text", text: "ok" }], details: {} };
},
};
const { runner, events } = createRunner([echo]);
const updates: unknown[] = [];
const first = await runner.execute("call", "echo", { a: 1 }, { onUpdate: (partial) => updates.push(partial) });
const missing = await runner.execute("call", "missing", {});
expect(first.toolCall.id).toBe("call/1");
expect(missing).toMatchObject({ toolCall: { id: "call/2" }, isError: true });
expect(updates).toHaveLength(1);
expect(events.map((event) => [event.type, event.toolCallId, event.parentToolCallId])).toEqual([
["tool_execution_start", "call/1", "call"],
["tool_execution_update", "call/1", "call"],
["tool_execution_end", "call/1", "call"],
["tool_execution_start", "call/2", "call"],
["tool_execution_end", "call/2", "call"],
]);
expect(runner.takeRecord("call")).toEqual({
calls: [
{ id: "call/1", name: "echo", arguments: { a: 1 }, status: "ok", durationMs: expect.any(Number) },
{
id: "call/2",
name: "missing",
arguments: {},
status: "error",
durationMs: expect.any(Number),
error: "Tool missing not found",
},
],
complete: true,
});
// The record is taken once.
expect(runner.takeRecord("call")).toBeUndefined();
expect(runner.takeRecord("other")).toBeUndefined();
});
it("records calls of nested tools on the model-issued call", async () => {
const leaf: AgentTool = {
name: "leaf",
label: "Leaf",
description: "Leaf",
parameters: Type.Object({}),
async execute() {
return { content: [], details: {} };
},
};
const tools: AgentTool[] = [leaf];
const { runner } = createRunner(tools);
tools.push({
name: "middle",
label: "Middle",
description: "Calls leaf",
parameters: Type.Object({}),
async execute(toolCallId) {
await runner.execute(toolCallId, "leaf", {});
return { content: [], details: {} };
},
});
await runner.execute("call", "middle", {});
expect(runner.takeRecord("call")?.calls.map((call) => call.id)).toEqual(["call/1", "call/1/1"]);
});
it("serializes concurrent calls to sequential tools", async () => {
let active = 0;
let maxActive = { sequential: 0, parallel: 0 };
const makeTool = (name: "sequential" | "parallel"): AgentTool => ({
name,
label: name,
description: name,
parameters: Type.Object({}),
executionMode: name === "sequential" ? "sequential" : undefined,
async execute() {
active++;
maxActive = { ...maxActive, [name]: Math.max(maxActive[name], active) };
await new Promise((resolve) => setTimeout(resolve, 5));
active--;
return { content: [], details: {} };
},
});
const { runner } = createRunner([makeTool("sequential"), makeTool("parallel")]);
await Promise.all([1, 2, 3].map(() => runner.execute("call", "sequential", {})));
await Promise.all([1, 2, 3].map(() => runner.execute("call", "parallel", {})));
expect(maxActive).toEqual({ sequential: 1, parallel: 3 });
});
});
describe("NestedCallRecorder", () => {
const call = (id: string, args: AgentToolCall["arguments"]): AgentToolCall => ({
type: "toolCall",
id,
name: "t",
arguments: args,
});
it("omits oversized arguments and drops calls beyond the limit", () => {
const recorder = new NestedCallRecorder();
expect(recorder.snapshot()).toBeUndefined();
const small = recorder.start(call("a", { x: 1 }));
recorder.finish(small, false, "");
expect(recorder.snapshot()).toEqual({
calls: [{ id: "a", name: "t", arguments: { x: 1 }, status: "ok", durationMs: expect.any(Number) }],
complete: true,
});
const big = recorder.start(call("b", { text: "x".repeat(NESTED_CALL_LIMITS.maxArgumentBytesPerCall) }));
recorder.finish(big, true, "e".repeat(1000));
const snapshot = recorder.snapshot();
expect(snapshot?.complete).toBe(false);
expect(snapshot?.calls[1]).toMatchObject({ id: "b", status: "error" });
expect(snapshot?.calls[1].arguments).toBeUndefined();
expect(snapshot?.calls[1].argumentsBytes).toBeGreaterThan(NESTED_CALL_LIMITS.maxArgumentBytesPerCall);
expect(snapshot?.calls[1].error).toHaveLength(NESTED_CALL_LIMITS.maxErrorChars);
for (let i = 0; i < NESTED_CALL_LIMITS.maxCalls; i++)
recorder.finish(recorder.start(call(`c${i}`, {})), false, "");
expect(recorder.snapshot()?.calls).toHaveLength(NESTED_CALL_LIMITS.maxCalls);
});
it("caps the total argument size and marks unfinished calls incomplete", () => {
const recorder = new NestedCallRecorder();
const chunk = { text: "x".repeat(7000) };
const records = Array.from({ length: 6 }, (_, i) => recorder.start(call(`c${i}`, chunk)));
const snapshot = recorder.snapshot();
// 32 KiB fits four 7000-byte argument objects.
expect(snapshot?.calls.filter((entry) => entry.arguments !== undefined)).toHaveLength(4);
expect(snapshot?.calls.every((entry) => entry.status === "unfinished")).toBe(true);
expect(snapshot?.complete).toBe(false);
expect(records).toHaveLength(6);
});
});

View file

@ -1037,6 +1037,56 @@ export default function(pi: ExtensionAPI) {
expect(runner.getCommand("deploy:2")?.description).toBe("global command");
expect(runner.getToolDefinition("duplicate-tool")?.description).toBe("explicit tool");
});
it("should leave out replaceable extensions whose names another extension registers", async () => {
// A third-party MCP extension registering /mcp replaces the built-in one instead of both running.
const globalExtDir = join(agentDir, "extensions");
mkdirSync(globalExtDir, { recursive: true });
writeFileSync(
join(globalExtDir, "other-mcp.ts"),
`
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
export default function(pi: ExtensionAPI) {
pi.registerCommand("mcp", { description: "other mcp", handler: async () => {} });
}`,
);
const loader = new DefaultResourceLoader({
cwd,
agentDir,
extensionFactories: [
{
name: "mcp",
replaceable: true,
factory: (pi) => pi.registerCommand("mcp", { description: "built-in mcp", handler: async () => {} }),
},
{
name: "llama",
replaceable: true,
factory: (pi) =>
pi.registerCommand("llama", { description: "built-in llama", handler: async () => {} }),
},
],
});
await loader.reload();
const extensionsResult = loader.getExtensions();
expect(extensionsResult.extensions.map((extension) => extension.path)).toEqual([
join(globalExtDir, "other-mcp.ts"),
"<inline:llama>",
]);
expect(extensionsResult.errors).toEqual([]);
const runner = new ExtensionRunner(
extensionsResult.extensions,
extensionsResult.runtime,
cwd,
SessionManager.inMemory(),
await createModelRegistry(AuthStorage.create(join(tempDir, "auth-replaceable.json"))),
);
expect(runner.getCommand("mcp")?.description).toBe("other mcp");
expect(runner.getCommand("llama")?.description).toBe("built-in llama");
});
});
describe("loadProjectContextFiles - nested worktree dedup", () => {

View file

@ -0,0 +1,646 @@
import { readFileSync, rmSync } from "node:fs";
import type { AgentTool } from "@earendil-works/pi-agent-core";
import {
type ClassifierModel,
type ClassifierResult,
fauxAssistantMessage,
fauxToolCall,
getCurrentTools,
type TranscriptContext,
} from "@earendil-works/pi-ai";
import type { ToolResultMessage } from "@earendil-works/pi-ai/compat";
import { Type } from "typebox";
import { afterEach, describe, expect, it } from "vitest";
import type { ExtensionAPI } from "../../src/core/extensions/types.ts";
import type { CustomEntry } from "../../src/core/session-manager.ts";
import { createToolDefinitionFromAgentTool } from "../../src/core/tools/tool-definition-wrapper.ts";
import { readCodemodeStore } from "../../src/extensions/codemode/execute.ts";
import { createCodemodeExtension } from "../../src/extensions/codemode/index.ts";
import {
CODEMODE_STORE_ENTRY_TYPE,
type CodemodeToolDetails,
createCodemodeTool,
} from "../../src/extensions/codemode/tool.ts";
import { createHarness, getToolResult, type Harness, type HarnessOptions } from "./harness.ts";
const TINY_PNG_BASE64 =
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==";
const echoSchema = Type.Object({ text: Type.String({ description: "Text to echo" }) });
const echoTool: AgentTool<typeof echoSchema> = {
name: "echo",
label: "Echo",
description: "Echo text back.\n\nSecond paragraph.",
parameters: echoSchema,
execute: async (_id, params) => ({
content: [{ type: "text", text: `echo: ${params.text}` }],
details: {},
}),
};
const statsTool: AgentTool = {
name: "stats",
label: "Stats",
description: "Return structured stats",
parameters: Type.Object({}),
outputSchema: Type.Object({ files: Type.Number(), names: Type.Array(Type.String()) }),
execute: async () => ({
content: [{ type: "text", text: "2 files" }],
details: {},
structuredContent: { files: 2, names: ["a", "b"] },
}),
};
const screenshotTool: AgentTool = {
name: "screenshot",
label: "Screenshot",
description: "Return a screenshot",
parameters: Type.Object({}),
execute: async () => ({
content: [
{ type: "text", text: "captured" },
{ type: "image", data: TINY_PNG_BASE64, mimeType: "image/png" },
],
details: {},
}),
};
const codemodeResult = (harness: Harness) => getToolResult(harness, "codemode");
/** The output after the script header, which is checked on the way. */
function resultText(message: ToolResultMessage): string {
const [header, ...items] = message.content;
expect(header).toEqual({
type: "text",
text: expect.stringMatching(/^Script (completed|failed)\nWall time \d+\.\d seconds\nOutput:\n$/),
});
return items.map((block) => (block.type === "text" ? block.text : `<${block.type}>`)).join("\n");
}
describe("AgentSession codemode tool", () => {
const harnesses: Harness[] = [];
afterEach(() => {
while (harnesses.length > 0) harnesses.pop()?.cleanup();
});
async function setup(extensionFactories?: HarnessOptions["extensionFactories"]) {
// Registered by an extension so they run with the session's tool context, next to the built-in codemode.
const registerTools = (pi: ExtensionAPI) => {
for (const tool of [echoTool as AgentTool, statsTool, screenshotTool]) {
pi.registerTool(createToolDefinitionFromAgentTool(tool));
}
};
const harness = await createHarness({
initialActiveToolNames: ["codemode"],
extensionFactories: [createCodemodeExtension(), registerTools, ...(extensionFactories ?? [])],
});
harnesses.push(harness);
return harness;
}
it("presents callable tools per codemode.mode", async () => {
const harness = await setup();
const description = (name: string) =>
harness.session.agent.state.tools.find((tool) => tool.name === name)?.description ?? "";
const requestTools: string[][] = [];
const record = (context: TranscriptContext) => {
requestTools.push(getCurrentTools(context.messages).map((tool) => tool.name));
return fauxAssistantMessage("ok");
};
// on: declared tools carry their codemode declaration and are not listed again in codemode.
harness.session.setActiveToolsByName(["echo", "codemode"]);
expect(description("echo")).toContain("codemode tool declaration:");
expect(description("codemode")).not.toContain("### `echo`");
harness.setResponses([record]);
await harness.session.prompt("on");
expect(requestTools[0]).toEqual(expect.arrayContaining(["echo", "codemode"]));
// only: codemode lists echo, which stays active but is left out of requests.
harness.settingsManager.applyOverrides({ codemode: { mode: "only" } });
harness.session.setActiveToolsByName(["echo", "codemode"]);
expect(description("echo")).not.toContain("codemode tool declaration:");
expect(description("codemode")).toContain("### `echo`");
expect(description("codemode")).not.toContain("### `stats`");
harness.setResponses([record]);
await harness.session.prompt("only");
expect(requestTools[1]).toContain("codemode");
expect(requestTools[1]).not.toContain("echo");
// Without codemode, tools keep their plain descriptions.
harness.session.setActiveToolsByName(["echo"]);
expect(description("echo")).toBe("Echo text back.\n\nSecond paragraph.");
});
it("runs nested calls in parallel and returns only the script result", async () => {
const harness = await setup();
harness.setResponses([
fauxAssistantMessage(
[
fauxToolCall("codemode", {
code: `
const [a, b, stats] = await Promise.all([
tools.echo({ text: "one" }),
tools.echo({ text: "two" }),
tools.stats({}),
]);
console.log("files", stats.files);
text(ALL_TOOLS.map((tool) => tool.name).join(","));
return { a, b, names: stats.names };
`,
}),
],
{ stopReason: "toolUse" },
),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("go");
const result = codemodeResult(harness);
expect(result.isError).toBe(false);
expect(resultText(result)).toBe(
'files 2\necho,stats,screenshot\n{"a":"echo: one","b":"echo: two","names":["a","b"]}',
);
const details = result.details as unknown as CodemodeToolDetails;
expect(details.calls.map((call) => [call.name, call.status])).toEqual([
["echo", "ok"],
["echo", "ok"],
["stats", "ok"],
]);
expect(details.calls.every((call) => call.id.startsWith(`${result.toolCallId}/`))).toBe(true);
// Nested calls never become transcript tool results; their events carry the parent id.
const toolResults = harness.session.messages.filter((message) => message.role === "toolResult");
expect(toolResults).toHaveLength(1);
});
it("routes nested calls through extension hooks", async () => {
const harness = await setup([
(pi) => {
pi.on("tool_call", (event) => {
if (event.toolName === "echo" && (event.input as { text: string }).text === "forbidden") {
return { block: true, reason: "echo of forbidden text is blocked" };
}
return undefined;
});
pi.on("tool_result", (event) => {
if (event.toolName === "stats") {
return { content: [{ type: "text", text: "redacted" }] };
}
return undefined;
});
},
]);
harness.setResponses([
fauxAssistantMessage(
[
fauxToolCall("codemode", {
code: `
let blocked;
try {
await tools.echo({ text: "forbidden" });
} catch (error) {
blocked = error.message;
}
const stats = await tools.stats({});
return { blocked, stats };
`,
}),
],
{ stopReason: "toolUse" },
),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("go");
const result = codemodeResult(harness);
// Replacing content without replacing structured content drops the structured result.
expect(JSON.parse(resultText(result))).toEqual({
blocked: "echo of forbidden text is blocked",
stats: "redacted",
});
const details = result.details as unknown as CodemodeToolDetails;
expect(details.calls.map((call) => call.status)).toEqual(["error", "ok"]);
});
it("keeps structured content that tool_result handlers replace along with the content", async () => {
const harness = await setup([
(pi) => {
pi.on("tool_result", (event) =>
event.toolName === "stats"
? { content: [{ type: "text", text: "0 files" }], structuredContent: { files: 0, names: [] } }
: undefined,
);
// A later handler that only touches details keeps what the first one set.
pi.on("tool_result", (event) => (event.toolName === "stats" ? { details: { audited: true } } : undefined));
},
]);
harness.setResponses([
fauxAssistantMessage([fauxToolCall("codemode", { code: "return await tools.stats({});" })], {
stopReason: "toolUse",
}),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("go");
expect(JSON.parse(resultText(codemodeResult(harness)))).toEqual({ files: 0, names: [] });
});
it("attaches only the images the script passes to image(), in output order", async () => {
const harness = await setup();
harness.setResponses([
fauxAssistantMessage(
[
fauxToolCall("codemode", {
code: `
// Tools without an outputSchema resolve to their text; images are not passed on.
const shot = await tools.screenshot({});
text(shot);
image("data:image/png;base64,${TINY_PNG_BASE64}");
text("after");
`,
}),
],
{ stopReason: "toolUse" },
),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("go");
const result = codemodeResult(harness);
expect(resultText(result)).toBe("captured\n<image>\nafter");
expect(result.content[2]).toEqual({ type: "image", data: TINY_PNG_BASE64, mimeType: "image/png" });
});
it("reports script failures as results that keep partial output and the calls that ran", async () => {
const harness = await setup();
harness.setResponses([
fauxAssistantMessage(
[
fauxToolCall("codemode", {
code: `text("partial");\nawait tools.echo({ text: "x" });\nthrow new Error("boom");`,
}),
],
{ stopReason: "toolUse" },
),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("go");
const result = codemodeResult(harness);
expect(result.isError).toBe(true);
expect((result.content[0] as { text: string }).text).toMatch(/^Script failed\n/);
const text = resultText(result);
expect(text).toMatch(/^partial\nScript error:\nError: boom\n/);
expect(text).toContain("codemode.js:3");
expect(text).toContain("Tool calls made before the failure (they are not undone): echo (ok)");
expect((result.details as unknown as CodemodeToolDetails).calls.map((call) => call.name)).toEqual(["echo"]);
});
});
describe("codemode options and store", () => {
const harnesses: Harness[] = [];
afterEach(() => {
while (harnesses.length > 0) harnesses.pop()?.cleanup();
});
// No tools override: the session builds its own codemode tool, including the store writer.
async function setup() {
const harness = await createHarness({
initialActiveToolNames: ["codemode"],
extensionFactories: [createCodemodeExtension()],
});
harnesses.push(harness);
return harness;
}
async function run(harness: Harness, code: string): Promise<ToolResultMessage> {
harness.setResponses([
fauxAssistantMessage([fauxToolCall("codemode", { code })], { stopReason: "toolUse" }),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("go");
return codemodeResult(harness);
}
function storeEntries(harness: Harness): unknown[] {
return harness.sessionManager
.getBranch()
.filter(
(entry): entry is CustomEntry => entry.type === "custom" && entry.customType === CODEMODE_STORE_ENTRY_TYPE,
)
.map((entry) => entry.data);
}
const increment = `const next = (load("count") ?? 0) + 1;\nstore("count", next);\nreturn next;`;
it("applies the timeout_ms option and rejects invalid options", async () => {
const harness = await setup();
const timedOut = await run(harness, '// @options: {"timeout_ms": 200}\nwhile (true) {}');
expect(timedOut.isError).toBe(true);
expect(resultText(timedOut)).toContain("Script error:\nScript timed out");
const invalid = await run(harness, '// @options: {"yield": 1}\ntext(1)');
expect(invalid.isError).toBe(true);
expect(invalid.content).toEqual([
{
type: "text",
text: "@options only supports `max_output_tokens` and `timeout_ms`; got `yield`",
},
]);
});
it("limits script memory so runaway allocations fail inside the script", async () => {
const harness = await setup();
const result = await run(
harness,
'// @options: {"timeout_ms": 30000}\nlet a = [];\ntry { while (true) a.push("x".repeat(1 << 20) + a.length); } catch (error) { const n = a.length; a = null; return { n, error: String(error) }; }',
);
expect(result.isError).toBe(false);
const { n, error } = JSON.parse(resultText(result)) as { n: number; error: string };
expect(error).toContain("out of memory");
// Each entry holds at least 1 MiB, so the limit stops the script well before wasm32's 4 GiB.
expect(n).toBeLessThan(512);
});
it("truncates output to the token budget and spills the full text", async () => {
const harness = await setup();
const result = await run(
harness,
'// @options: {"max_output_tokens": 10}\nfor (let i = 0; i < 100; i++) text("row " + i);\nimage("data:image/png;base64,AAAA");',
);
const details = result.details as unknown as CodemodeToolDetails;
const path = details.fullOutputPath;
if (!path) throw new Error("No spill file");
try {
const text = resultText(result);
expect(text).toMatch(/^Warning: truncated output/);
expect(text).toContain("row 0\n");
expect(text).toContain("tokens truncated");
expect(text).toContain("row 99\n");
expect(text).not.toContain("row 50\n");
expect(text).toContain(`[Full output: ${path} (read with offset/limit)]`);
// Images follow the truncated text.
expect(result.content.at(-1)).toEqual({ type: "image", data: "AAAA", mimeType: "image/png" });
expect(readFileSync(path, "utf8")).toBe(Array.from({ length: 100 }, (_, i) => `row ${i}`).join("\n"));
} finally {
rmSync(path, { force: true });
}
const small = await run(harness, "return { ok: true };");
expect((small.details as unknown as CodemodeToolDetails).fullOutputPath).toBeUndefined();
});
it("resolves bash calls to structured results, also for non-zero exit codes", async () => {
const harness = await createHarness({
initialActiveToolNames: ["codemode", "bash"],
extensionFactories: [createCodemodeExtension()],
});
harnesses.push(harness);
const result = await run(
harness,
'const r = await tools.bash({ command: "echo out; exit 3" });\ntext(JSON.stringify([r.output, r.exit_code, typeof r.wall_time_seconds]));',
);
expect(result.isError).toBe(false);
expect(resultText(result)).toBe('["out\\n",3,"number"]');
});
it("persists store() writes as custom entries for later calls", async () => {
const harness = await setup();
expect(resultText(await run(harness, increment))).toBe("1");
expect(resultText(await run(harness, increment))).toBe("2");
expect(storeEntries(harness)).toEqual([
{ set: { count: 1 }, delete: [] },
{ set: { count: 2 }, delete: [] },
]);
const appended = harness
.eventsOfType("entry_appended")
.filter((event) => event.entry.type === "custom" && event.entry.customType === CODEMODE_STORE_ENTRY_TYPE);
expect(appended).toHaveLength(2);
expect(resultText(await run(harness, 'store("count", undefined);\nreturn load("count") === undefined;'))).toBe(
"true",
);
expect(storeEntries(harness).at(-1)).toEqual({ set: {}, delete: ["count"] });
});
it("appends nothing for failed scripts or scripts without writes", async () => {
const harness = await setup();
expect((await run(harness, 'store("count", 5);\nthrow new Error("boom");')).isError).toBe(true);
expect((await run(harness, 'return load("count") ?? "missing";')).isError).toBe(false);
expect(storeEntries(harness)).toEqual([]);
});
it("runs without a session, starting from an empty store", async () => {
const tool = createCodemodeTool();
const result = await tool.execute("direct", { code: increment });
expect(result.content[1]).toEqual({ type: "text", text: "1" });
});
it("loads the values written on the current branch", async () => {
const harness = await setup();
await run(harness, increment);
const firstPrompt = harness.sessionManager.getBranch().find((entry) => entry.type === "message");
if (!firstPrompt) throw new Error("No first prompt entry");
expect(resultText(await run(harness, increment))).toBe("2");
// Branch from the first prompt: the store entries written after it are on another path.
harness.sessionManager.branch(firstPrompt.id);
expect(resultText(await run(harness, increment))).toBe("1");
});
it("folds store entries from the root, ignoring malformed data", () => {
const entry = (data: unknown, customType = CODEMODE_STORE_ENTRY_TYPE): CustomEntry => ({
type: "custom",
customType,
data,
id: Math.random().toString(36).slice(2),
parentId: null,
timestamp: new Date(0).toISOString(),
});
expect(
readCodemodeStore([
entry({ set: { a: 1, b: { c: 2 } }, delete: [] }),
entry({ set: { a: 3 }, delete: ["b"] }),
entry({ set: { z: 1 } }),
entry({ set: { other: 1 }, delete: [] }, "other-extension"),
]),
).toEqual({ a: 3 });
});
});
describe("codemode models", () => {
const harnesses: Harness[] = [];
afterEach(() => {
while (harnesses.length > 0) harnesses.pop()?.cleanup();
});
const scorerModel: ClassifierModel<"test-classifier"> = {
type: "classifier",
id: "judge",
name: "Judge",
api: "test-classifier",
provider: "scorer",
baseUrl: "https://classifier.test/v1",
input: ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 1000,
headers: { "X-Secret": "hunter2" },
};
interface ClassifyObservation {
baseUrl: string;
apiKey: string | undefined;
text: unknown;
}
async function setup() {
const harness = await createHarness({
initialActiveToolNames: ["codemode"],
extensionFactories: [createCodemodeExtension()],
});
harnesses.push(harness);
const observed: ClassifyObservation[] = [];
let active = 0;
let maxActive = 0;
harness.session.modelRuntime.registerProvider("scorer", {
apiKey: "secret-key",
models: [scorerModel],
classifiers: {
"test-classifier": {
classify: async (model, context, options): Promise<ClassifierResult> => {
active++;
maxActive = Math.max(maxActive, active);
await new Promise((resolve) => setTimeout(resolve, 10));
active--;
const text = context.state.text;
observed.push({ baseUrl: model.baseUrl, apiKey: options?.apiKey, text });
if (text === "explode") {
return {
api: model.api,
provider: model.provider,
model: model.id,
answers: {},
stopReason: "error",
errorMessage: "classifier exploded",
timestamp: 0,
};
}
return {
api: model.api,
provider: model.provider,
model: model.id,
answers: { approved: { type: "bool", probability: text === "good" ? 0.9 : 0.1 } },
stopReason: "stop",
timestamp: 0,
};
},
},
},
});
harness.session.setActiveToolsByName(["codemode"]);
return { harness, observed, maxActive: () => maxActive };
}
async function run(harness: Harness, code: string): Promise<ToolResultMessage> {
harness.setResponses([
fauxAssistantMessage([fauxToolCall("codemode", { code })], { stopReason: "toolUse" }),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("go");
return codemodeResult(harness);
}
const questions = `{ approved: { type: "bool", instructions: "Approval?", criteria: { true: "yes", false: "no" } } }`;
it("declares models only for the session's own codemode tool", async () => {
const { harness } = await setup();
const codemode = harness.session.agent.state.tools.find((tool) => tool.name === "codemode");
expect(codemode?.description).toContain("declare const models: {");
expect(codemode?.description).toContain(
"classify(model: ModelInfo, context: ClassifierContext): Promise<ClassifierResult>;",
);
expect(codemode?.description).toContain("interface ClassifierResult {");
const overridden = await createHarness({ tools: [createCodemodeTool() as AgentTool] });
harnesses.push(overridden);
const plain = overridden.session.agent.state.tools.find((tool) => tool.name === "codemode");
expect(plain?.description).not.toContain("declare const models");
});
it("lists models and classifies with catalog auth, ignoring script-supplied fields", async () => {
const { harness, observed, maxActive } = await setup();
const result = await run(
harness,
`
const [model] = await models.getAvailableOfType("classifier", "scorer");
const listed = await models.getModelsOfType("classifier");
const same = await models.getModelOfType("classifier", "scorer", "judge");
const texts = ["good", "bad", "good", "bad", "good", "bad"];
const results = await Promise.all(
texts.map((text) => models.classify({ ...model, baseUrl: "https://evil.test" }, { state: { text }, questions: ${questions} })),
);
return {
id: model.id,
headers: "headers" in model,
listed: listed.some((entry) => entry.provider === "scorer" && entry.id === "judge"),
same: same.id,
missing: (await models.getModelOfType("classifier", "scorer", "nope")) === undefined,
probabilities: results.map((r) => r.answers.approved.probability),
};
`,
);
expect(result.isError).toBe(false);
expect(JSON.parse(resultText(result))).toEqual({
id: "judge",
headers: false,
listed: true,
same: "judge",
missing: true,
probabilities: [0.9, 0.1, 0.9, 0.1, 0.9, 0.1],
});
expect(observed).toHaveLength(6);
expect(
observed.every((entry) => entry.baseUrl === "https://classifier.test/v1" && entry.apiKey === "secret-key"),
).toBe(true);
// Six classifications with at most four in flight.
expect(maxActive()).toBe(4);
const details = result.details as unknown as CodemodeToolDetails;
expect(details.calls.map((call) => [call.name, call.args, call.status])).toEqual(
Array.from({ length: 6 }, () => ["models.classify", "scorer/judge", "ok"]),
);
});
it("reports provider errors as results and invalid arguments as exceptions", async () => {
const { harness } = await setup();
const result = await run(
harness,
`
const model = await models.getModelOfType("classifier", "scorer", "judge");
const failed = await models.classify(model, { state: { text: "explode" }, questions: ${questions} });
const attempt = async (fn) => { try { await fn(); return "ok"; } catch (error) { return error.message; } };
return {
failed: [failed.stopReason, failed.errorMessage],
badType: await attempt(() => models.getModelsOfType("video")),
unknown: await attempt(() => models.classify({ provider: "scorer", id: "nope" }, {})),
noModel: await attempt(() => models.classify("judge", {})),
};
`,
);
expect(result.isError).toBe(false);
const value = JSON.parse(resultText(result));
expect(value.failed).toEqual(["error", "classifier exploded"]);
expect(value.badType).toContain('Unknown model type "video"');
expect(value.unknown).toBe('Unknown classifier model "scorer/nope"');
expect(value.noModel).toContain("expects a model");
const details = result.details as unknown as CodemodeToolDetails;
expect(details.calls.map((call) => [call.name, call.status, call.error])).toEqual([
["models.classify", "error", "classifier exploded"],
]);
});
});

View file

@ -0,0 +1,185 @@
import { mkdtempSync, rmSync, writeFileSync } from "node:fs";
import { type AddressInfo, createServer } from "node:net";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai";
import type { ToolResultMessage } from "@earendil-works/pi-ai/compat";
import { afterEach, describe, expect, it } from "vitest";
import { InMemoryAuthStorageBackend } from "../../src/core/auth-storage.ts";
import { runMcpCommand } from "../../src/extensions/mcp/cli.ts";
import type { McpOAuthConfig, McpServerEntry } from "../../src/extensions/mcp/config.ts";
import { createMcpExtension } from "../../src/extensions/mcp/index.ts";
import { McpOAuthCredentialStore } from "../../src/extensions/mcp/oauth.ts";
import { createHarness, createTestUiContext, getMessageText, type Harness } from "./harness.ts";
import { startOAuthMcpServer } from "./mcp-oauth-server.ts";
describe("AgentSession MCP OAuth", () => {
const cleanups: (() => Promise<void> | void)[] = [];
afterEach(async () => {
while (cleanups.length > 0) await cleanups.pop()?.();
});
async function setup(browser: "follow" | "paste", oauth?: McpOAuthConfig) {
const server = await startOAuthMcpServer();
cleanups.push(server.close);
const backend = new InMemoryAuthStorageBackend();
const entry: McpServerEntry = {
name: "issues",
config: { url: server.url, exposure: "direct", ...(oauth ? { oauth } : {}) },
source: "test",
};
const notifications: string[] = [];
const opened: URL[] = [];
let redirectLocation: Promise<string> | undefined;
const harness: Harness = await createHarness({
initialActiveToolNames: [],
extensionFactories: [
createMcpExtension({
loadConfig: () => ({ servers: [entry], errors: [] }),
credentials: new McpOAuthCredentialStore(backend),
openUrl: (url) => {
opened.push(new URL(url));
if (browser === "follow") {
// The browser follows the authorization redirect to the loopback callback.
void fetch(url);
} else {
// The browser cannot reach the callback; the user pastes the redirect URL.
redirectLocation = fetch(url, { redirect: "manual" }).then(
(response) => response.headers.get("location") ?? "",
);
}
},
}),
],
});
cleanups.push(() => harness.cleanup());
await harness.session.bindExtensions({
uiContext: createTestUiContext({
notify: (message) => notifications.push(message),
// The paste prompt waits until sign-in completes unless the user pastes the redirect URL.
input: (_title, _placeholder, opts) =>
browser === "paste"
? Promise.resolve(redirectLocation)
: new Promise((resolve) =>
opts?.signal?.addEventListener("abort", () => resolve(undefined), { once: true }),
),
}),
});
return { harness, server, notifications, backend, opened };
}
async function callWhoami(harness: Harness): Promise<ToolResultMessage> {
harness.setResponses([
fauxAssistantMessage([fauxToolCall("mcp__issues__whoami", {})], { stopReason: "toolUse" }),
fauxAssistantMessage("done"),
]);
const before = harness.session.messages.length;
await harness.session.prompt("who am i");
const result = harness.session.messages
.slice(before)
.find((message): message is ToolResultMessage => message.role === "toolResult");
if (!result) throw new Error("no tool result");
return result;
}
it("signs in through the browser, refreshes expired tokens, and signs out", async () => {
const { harness, server, notifications, backend } = await setup("follow");
await harness.session.prompt("/mcp");
// Startup problems are reported once, pointing to /mcp.
expect(notifications).toContain("MCP servers need attention:\n issues: needs sign-in\nRun /mcp to fix.");
expect(notifications.at(-1)).toBe("issues: needs sign-in, run /mcp login issues (direct)");
await harness.session.prompt("/mcp login issues");
expect(notifications.at(-1)).toBe('Signed in to MCP server "issues" (1 tools).');
expect(server.log).toEqual(["401 none", "register", "token code"]);
expect(backend.withLock((current) => ({ result: current }))).toContain('"access_token": "access-1"');
expect(getMessageText(await callWhoami(harness))).toBe("token access-1");
// An expired access token is refreshed without user interaction.
server.expireAccessTokens();
expect(getMessageText(await callWhoami(harness))).toBe("token access-2");
expect(server.log.slice(-3)).toEqual(["401 access-1", "token refresh", "call access-2"]);
// A token past its expiry is refreshed before the request, without a 401 round trip.
backend.withLock((current) => {
const states = JSON.parse(current ?? "{}") as Record<string, { tokensExpireAt?: number }>;
for (const state of Object.values(states)) state.tokensExpireAt = Date.now() - 1_000;
return { result: undefined, next: JSON.stringify(states) };
});
expect(getMessageText(await callWhoami(harness))).toBe("token access-3");
expect(server.log.slice(-2)).toEqual(["token refresh", "call access-3"]);
await harness.session.prompt("/mcp logout issues");
expect(notifications.at(-1)).toBe('Signed out of MCP server "issues".');
const result = await callWhoami(harness);
expect(result.isError).toBe(true);
expect(getMessageText(result)).toBe('MCP server "issues" requires sign-in. Run /mcp to sign in.');
});
it("accepts a pasted redirect URL when the browser cannot reach the callback", async () => {
const { harness, server, notifications } = await setup("paste");
await harness.session.prompt("/mcp login");
expect(notifications.at(-1)).toBe('Signed in to MCP server "issues" (1 tools).');
expect(getMessageText(await callWhoami(harness))).toBe(`token access-1`);
expect(server.log).toContain("token code");
});
async function freePort(): Promise<number> {
return new Promise<number>((resolve) => {
const probe = createServer().listen(0, "127.0.0.1", () => {
const address = probe.address() as AddressInfo;
probe.close(() => resolve(address.port));
});
});
}
it("uses the configured callback URL and scope", async () => {
const callbackUrl = `http://localhost:${await freePort()}/callback`;
const { harness, notifications, opened } = await setup("follow", { callbackUrl, scope: "issues:read" });
await harness.session.prompt("/mcp login issues");
expect(notifications.at(-1)).toBe('Signed in to MCP server "issues" (1 tools).');
expect(opened[0].searchParams.get("redirect_uri")).toBe(callbackUrl);
expect(opened[0].searchParams.get("scope")).toBe("issues:read");
expect(getMessageText(await callWhoami(harness))).toBe("token access-1");
});
it("adds the listening port to a callback URL without one", async () => {
const { harness, notifications, opened } = await setup("follow", { callbackUrl: "http://127.0.0.1/oauth/done" });
await harness.session.prompt("/mcp login issues");
expect(notifications.at(-1)).toBe('Signed in to MCP server "issues" (1 tools).');
expect(opened[0].searchParams.get("redirect_uri")).toMatch(/^http:\/\/127\.0\.0\.1:\d+\/oauth\/done$/);
const port = await freePort();
const fixed = await setup("follow", { callbackUrl: "http://127.0.0.1/oauth/done", callbackPort: port });
await fixed.harness.session.prompt("/mcp login issues");
expect(fixed.opened[0].searchParams.get("redirect_uri")).toBe(`http://127.0.0.1:${port}/oauth/done`);
});
it("uses credentials from pi mcp login on the next turn", async () => {
const { harness, server, backend } = await setup("follow");
const agentDir = mkdtempSync(join(tmpdir(), "pi-mcp-login-"));
cleanups.push(() => rmSync(agentDir, { recursive: true, force: true }));
writeFileSync(join(agentDir, "mcp.json"), JSON.stringify({ mcpServers: { issues: { url: server.url } } }));
// The agent runs `pi mcp login issues` through bash; the user approves in the browser.
const output: string[] = [];
const exitCode = await runMcpCommand(["login", "issues"], {
cwd: agentDir,
agentDir,
credentials: new McpOAuthCredentialStore(backend),
openUrl: (url) => void fetch(url),
log: (line) => output.push(line),
error: (line) => output.push(line),
});
expect(exitCode).toBe(0);
expect(output.at(-1)).toBe('Signed in to MCP server "issues" (1 tools).');
// The session still waits for a sign-in, and reconnects when the next turn starts.
expect(getMessageText(await callWhoami(harness))).toBe("token access-1");
});
});

View file

@ -0,0 +1,733 @@
import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai";
import type { SystemMessage, ToolResultMessage } from "@earendil-works/pi-ai/compat";
import { type JsonRpcRequest, LATEST_PROTOCOL_VERSION } from "@earendil-works/pi-mcp";
import { createInMemoryTransportPair } from "@earendil-works/pi-mcp/testing";
import { Type } from "typebox";
import { afterEach, describe, expect, it, vi } from "vitest";
import type { ExtensionAPI, ExtensionFactory } from "../../src/core/extensions/types.ts";
import { createCodemodeExtension } from "../../src/extensions/codemode/index.ts";
import type { McpExposure, McpServerEntry } from "../../src/extensions/mcp/config.ts";
import { createMcpExtension } from "../../src/extensions/mcp/index.ts";
import { createMcpToolName } from "../../src/extensions/mcp/tools.ts";
import { createToolSearchExtension } from "../../src/extensions/tool-search/index.ts";
import {
createHarness,
createTestUiContext,
getAssistantTexts,
getMessageText,
type Harness,
getToolResult as toolResult,
} from "./harness.ts";
const TINY_PNG_BASE64 =
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==";
const SERVER_TOOLS = [
{
name: "search",
description: "Search the docs.",
inputSchema: { type: "object", properties: { query: { type: "string" } }, required: ["query"] },
outputSchema: {
type: "object",
properties: { hits: { type: "array", items: { type: "string" } } },
required: ["hits"],
},
},
{
name: "fail",
description: "Always fails.",
inputSchema: { type: "object", properties: {} },
annotations: { title: "Fail", destructiveHint: true, readOnlyHint: false, idempotentHint: "yes" },
},
{ name: "shot", description: "Returns an image.", inputSchema: { type: "object", properties: {} } },
];
/** Minimal MCP server over an in-memory transport. Records the tool calls it receives. */
function createFakeServer(calls: string[], listTools: () => unknown[] = () => SERVER_TOOLS, resources = false) {
const pair = createInMemoryTransportPair();
const respond = (request: JsonRpcRequest): unknown => {
switch (request.method) {
case "initialize":
return {
protocolVersion: LATEST_PROTOCOL_VERSION,
capabilities: { tools: {}, ...(resources ? { resources: {} } : {}) },
serverInfo: { name: "docs", version: "1.0.0" },
};
case "tools/list":
return { tools: listTools() };
case "resources/list":
return {
resources: [
{ uri: "docs://intro", name: "intro", mimeType: "text/markdown", _meta: { x: 1 } },
// MCP App user interfaces are left out.
{ uri: "ui://docs/viewer", name: "viewer", mimeType: "text/html;profile=mcp-app" },
],
};
case "resources/templates/list":
return {
resourceTemplates: [
{ uriTemplate: "docs://pages/{slug}", name: "page", icons: [{ src: "data:image/png;base64,AAAA" }] },
],
};
case "resources/read": {
const { uri } = request.params as { uri: string };
calls.push(`read:${uri}`);
return { contents: [{ uri, mimeType: "text/markdown", text: `# ${uri}` }] };
}
case "tools/call": {
const params = request.params as { name: string; arguments?: { query?: string } };
calls.push(`${params.name}:${JSON.stringify(params.arguments ?? {})}`);
if (params.name === "search") {
const hits = [`${params.arguments?.query} guide`, `${params.arguments?.query} faq`];
return { content: [{ type: "text", text: hits.join("\n") }], structuredContent: { hits } };
}
if (params.name === "shot") {
return { content: [{ type: "image", data: TINY_PNG_BASE64, mimeType: "image/png" }] };
}
return { content: [{ type: "text", text: "server exploded" }], isError: true };
}
default:
return {};
}
};
pair.server.onMessage((message) => {
if (!("id" in message) || !("method" in message)) return;
const request = message as JsonRpcRequest;
queueMicrotask(() => {
void pair.server.send({ jsonrpc: "2.0", id: request.id, result: respond(request) });
});
});
return pair;
}
describe("AgentSession MCP integration", () => {
const harnesses: Harness[] = [];
afterEach(() => {
while (harnesses.length > 0) harnesses.pop()?.cleanup();
});
async function setup(
exposure: McpExposure,
listTools?: () => unknown[],
options: {
autoEnableCodemode?: boolean;
builtInTools?: string[];
extensionFactories?: ExtensionFactory[];
toolExposure?: Record<string, McpExposure>;
resources?: boolean;
withoutToolSearch?: boolean;
} = {},
) {
const {
builtInTools,
extensionFactories = [],
toolExposure,
resources,
withoutToolSearch,
...configOptions
} = options;
const calls: string[] = [];
const notifications: string[] = [];
const servers: ReturnType<typeof createFakeServer>["server"][] = [];
const entry: McpServerEntry = {
name: "docs",
config: { url: "http://unused.invalid", exposure, ...(toolExposure ? { toolExposure } : {}) },
source: "test",
};
// `builtInTools` are the built-in tools active at the start. The MCP extension activates codemode
// or tool_search.
const harness = await createHarness({
initialActiveToolNames: builtInTools ?? [],
extensionFactories: [
...extensionFactories,
createCodemodeExtension(),
...(withoutToolSearch ? [] : [createToolSearchExtension()]),
createMcpExtension({
loadConfig: () => ({ servers: [entry], errors: [], ...configOptions }),
createTransport: () => {
const pair = createFakeServer(calls, listTools, resources);
servers.push(pair.server);
void pair.server.start();
return pair.client;
},
}),
],
});
harnesses.push(harness);
await harness.session.bindExtensions({
uiContext: createTestUiContext({ notify: (message) => notifications.push(message) }),
});
return { harness, calls, servers, notifications };
}
function declaredToolNames(harness: Harness): string[] {
return harness.session.messages
.filter((message): message is SystemMessage => message.role === "system")
.flatMap((message) => (message.toolsAdded ?? []).map((tool) => tool.name));
}
function nestedToolNames(harness: Harness): string[] {
return harness.session.getCallableToolNames();
}
it("exposes codemode-only MCP tools through codemode and hides them from the model", async () => {
const { harness, calls } = await setup("codemode");
const searchName = createMcpToolName("docs", "search");
// MCP tools resolve to their CallToolResult, errors included.
harness.setResponses([
fauxAssistantMessage(
[
fauxToolCall("codemode", {
code: `
const [a, b] = await Promise.allSettled([
tools.${searchName}({ query: "mcp" }),
tools.${searchName}({ query: "pi" }),
]);
const failure = await tools.mcp__docs__fail({});
const shot = await tools.mcp__docs__shot({});
image(shot.content[0]);
text(JSON.stringify({
hits: [...a.value.structuredContent.hits, ...b.value.structuredContent.hits],
failed: failure.isError,
failure: failure.content[0].text,
found: ALL_TOOLS.filter((tool) => tool.name.includes("search")).map((tool) => tool.name),
}));
`,
}),
],
{ stopReason: "toolUse" },
),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("search the docs");
// Exec was activated for the codemode-exposed server; MCP tools are never declared.
expect(harness.session.getActiveToolNames()).toEqual(["codemode"]);
expect(declaredToolNames(harness)).toEqual(["codemode"]);
expect(nestedToolNames(harness)).toEqual([searchName, "mcp__docs__fail", "mcp__docs__shot"]);
const codemode = harness.session.agent.state.tools.find((tool) => tool.name === "codemode");
expect(codemode?.description).toContain("Shared MCP Types:");
expect(codemode?.description).toContain("### `mcp__docs__search`");
const result = toolResult(harness, "codemode");
expect(result.isError).toBe(false);
// Output items keep the order the script produced them in.
expect(result.content[1]).toEqual({ type: "image", data: TINY_PNG_BASE64, mimeType: "image/png" });
expect(JSON.parse((result.content[2] as { text: string }).text)).toEqual({
hits: ["mcp guide", "mcp faq", "pi guide", "pi faq"],
failed: true,
failure: "server exploded",
found: [searchName],
});
expect(result.content).toHaveLength(3);
expect(calls).toEqual(['search:{"query":"mcp"}', 'search:{"query":"pi"}', "fail:{}", "shot:{}"]);
});
it("keeps codemode-only MCP tools callable across tree navigation", async () => {
const { harness } = await setup("codemode");
const searchName = createMcpToolName("docs", "search");
harness.setResponses([fauxAssistantMessage("one"), fauxAssistantMessage("two")]);
await harness.session.prompt("first");
await harness.session.prompt("second");
expect(harness.session.getActiveToolNames()).toEqual(["codemode"]);
expect(nestedToolNames(harness)).toContain(searchName);
const firstAssistant = harness.sessionManager
.getBranch()
.find((entry) => entry.type === "message" && entry.message.role === "assistant");
if (!firstAssistant) throw new Error("No assistant entry");
await harness.session.navigateTree(firstAssistant.id);
expect(harness.session.getActiveToolNames()).toEqual(["codemode"]);
expect(nestedToolNames(harness)).toContain(searchName);
const codemode = harness.session.agent.state.tools.find((tool) => tool.name === "codemode");
expect(codemode?.description).toContain(`${searchName}(args: {`);
});
it("rejects direct model calls to codemode-only MCP tools", async () => {
const { harness, calls } = await setup("codemode");
harness.setResponses([
fauxAssistantMessage([fauxToolCall(createMcpToolName("docs", "search"), { query: "x" })], {
stopReason: "toolUse",
}),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("go");
const result = toolResult(harness, "mcp__docs__search");
expect(result.isError).toBe(true);
expect(getMessageText(result)).toBe("Tool mcp__docs__search not found");
expect(calls).toEqual([]);
});
it("declares directly exposed MCP tools to the model", async () => {
const { harness } = await setup("direct");
harness.setResponses([
fauxAssistantMessage([fauxToolCall("mcp__docs__search", { query: "direct" })], { stopReason: "toolUse" }),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("go");
expect(declaredToolNames(harness)).toEqual(["mcp__docs__search", "mcp__docs__fail", "mcp__docs__shot"]);
expect(harness.session.getActiveToolNames()).not.toContain("codemode");
// Boolean annotation hints are passed on for permission extensions.
const annotations = new Map(harness.session.getAllTools().map((tool) => [tool.name, tool.annotations]));
expect(annotations.get("mcp__docs__fail")).toEqual({ destructiveHint: true, readOnlyHint: false });
expect(annotations.get("mcp__docs__search")).toBeUndefined();
const result = toolResult(harness, "mcp__docs__search");
expect(getMessageText(result)).toBe("direct guide\ndirect faq");
});
it.each(["direct", "codemode"] as const)(
"withdraws and restores %s MCP tools the server changes",
async (exposure) => {
let tools = SERVER_TOOLS;
const { harness, servers } = await setup(exposure, () => tools);
harness.setResponses([fauxAssistantMessage("ready")]);
await harness.session.prompt("start");
// Direct tools are declared to the model, codemode tools are only callable from codemode.
const reachable = () =>
exposure === "direct" ? harness.session.getActiveToolNames() : nestedToolNames(harness);
const listChanged = async () => {
await servers[0].send({ jsonrpc: "2.0", method: "notifications/tools/list_changed" });
await new Promise((resolve) => setTimeout(resolve, 10));
};
expect(reachable()).toContain("mcp__docs__fail");
tools = SERVER_TOOLS.filter((tool) => tool.name !== "fail");
await listChanged();
expect(reachable()).not.toContain("mcp__docs__fail");
expect(reachable()).toContain("mcp__docs__search");
const codemode = harness.session.agent.state.tools.find((tool) => tool.name === "codemode");
expect(codemode?.description ?? "").not.toContain("mcp__docs__fail");
tools = SERVER_TOOLS;
await listChanged();
expect(reachable()).toContain("mcp__docs__fail");
},
);
it("lists and reads resources with Codex's resource tools", async () => {
const { harness, calls } = await setup("direct", undefined, { resources: true });
harness.setResponses([
fauxAssistantMessage(
[
fauxToolCall("list_mcp_resources", {}),
fauxToolCall("list_mcp_resource_templates", { server: "docs" }),
fauxToolCall("read_mcp_resource", { server: "docs", uri: "docs://pages/setup" }),
],
{ stopReason: "toolUse" },
),
fauxAssistantMessage([fauxToolCall("read_mcp_resource", { server: "nope", uri: "docs://x" })], {
stopReason: "toolUse",
}),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("read");
// The resource tools take the exposure of the servers they reach.
expect(harness.session.getActiveToolNames()).toEqual(
expect.arrayContaining(["list_mcp_resources", "list_mcp_resource_templates", "read_mcp_resource"]),
);
expect(JSON.parse(getMessageText(toolResult(harness, "list_mcp_resources")))).toEqual({
resources: [{ server: "docs", uri: "docs://intro", name: "intro", mimeType: "text/markdown" }],
});
expect(JSON.parse(getMessageText(toolResult(harness, "list_mcp_resource_templates")))).toEqual({
server: "docs",
resourceTemplates: [{ server: "docs", uriTemplate: "docs://pages/{slug}", name: "page" }],
});
const results = harness.session.messages.filter(
(message): message is ToolResultMessage =>
message.role === "toolResult" && message.toolName === "read_mcp_resource",
);
expect(getMessageText(results[0])).toBe("# docs://pages/setup");
expect(results[1].isError).toBe(true);
expect(getMessageText(results[1])).toBe('MCP server "nope" has no resources. Servers with resources: docs');
expect(calls).toEqual(["read:docs://pages/setup"]);
const annotations = harness.session.getAllTools().find((tool) => tool.name === "read_mcp_resource");
expect(annotations?.annotations).toEqual({ readOnlyHint: true });
});
it("makes the resource tools callable from codemode for codemode servers", async () => {
const { harness, calls } = await setup("codemode", undefined, { resources: true });
harness.setResponses([
fauxAssistantMessage(
[
fauxToolCall("codemode", {
code: `const listed = await tools.list_mcp_resources({ server: "docs" });
const read = await tools.read_mcp_resource({ server: "docs", uri: listed.resources[0].uri });
return { uris: listed.resources.map((r) => r.uri), text: read.contents[0].text };`,
}),
],
{ stopReason: "toolUse" },
),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("read");
expect(harness.session.getActiveToolNames()).toEqual(["codemode"]);
expect(JSON.parse(getMessageText(toolResult(harness, "codemode")).split("\n").at(-1) ?? "")).toEqual({
uris: ["docs://intro"],
text: "# docs://intro",
});
expect(calls).toEqual(["read:docs://intro"]);
});
it("applies per-tool exposure overrides", async () => {
const { harness } = await setup("hidden", undefined, { toolExposure: { search: "direct", "s*": "codemode" } });
harness.setResponses([fauxAssistantMessage("ready")]);
await harness.session.prompt("start");
// `search` is declared, `shot` is only callable from codemode, `fail` keeps the server's `hidden`.
expect(declaredToolNames(harness)).toEqual(["mcp__docs__search", "codemode"]);
expect(nestedToolNames(harness)).toEqual(["mcp__docs__search", "mcp__docs__shot"]);
const codemode = harness.session.agent.state.tools.find((tool) => tool.name === "codemode");
expect(codemode?.description).toContain("### `mcp__docs__shot`");
expect(codemode?.description).not.toContain("mcp__docs__fail");
});
it("keeps codemode-deferred MCP tools callable from codemode but out of its description", async () => {
const { harness, calls } = await setup("codemode-deferred");
const searchName = createMcpToolName("docs", "search");
harness.setResponses([
fauxAssistantMessage(
[
fauxToolCall("codemode", {
code: `return (await tools.${searchName}({ query: "d" })).structuredContent;`,
}),
],
{
stopReason: "toolUse",
},
),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("go");
expect(harness.session.getActiveToolNames()).toEqual(["codemode"]);
const codemode = harness.session.agent.state.tools.find((tool) => tool.name === "codemode");
expect(codemode?.description).not.toContain(searchName);
const result = toolResult(harness, "codemode");
expect(result.isError).toBe(false);
expect(JSON.parse((result.content[1] as { text: string }).text)).toEqual({ hits: ["d guide", "d faq"] });
expect(calls).toEqual(['search:{"query":"d"}']);
});
it("does not activate codemode when autoEnableCodemode is false", async () => {
const { harness, notifications } = await setup("codemode", undefined, { autoEnableCodemode: false });
harness.setResponses([fauxAssistantMessage("ready")]);
await harness.session.prompt("start");
expect(harness.session.getActiveToolNames()).toEqual([]);
expect(nestedToolNames(harness)).toContain("mcp__docs__search");
expect(notifications).toEqual([
"MCP tools are only reachable from the codemode or tool_search tool, but neither is active (autoEnableCodemode is false); they cannot be called.",
]);
});
it("treats codemode MCP tools as reachable through an active tool_search", async () => {
const { harness, notifications } = await setup("codemode", undefined, {
autoEnableCodemode: false,
builtInTools: ["tool_search"],
});
harness.setResponses([
fauxAssistantMessage([fauxToolCall("tool_search", { query: "search the docs", limit: 1 })], {
stopReason: "toolUse",
}),
fauxAssistantMessage([fauxToolCall("mcp__docs__search", { query: "loaded" })], { stopReason: "toolUse" }),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("go");
expect(harness.session.getActiveToolNames()).toEqual(["tool_search", "mcp__docs__search"]);
expect(getMessageText(toolResult(harness, "mcp__docs__search"))).toBe("loaded guide\nloaded faq");
expect(notifications).toEqual([]);
});
it("reaches deferred MCP tools through codemode when tool_search is not available", async () => {
const { harness, notifications } = await setup("deferred", undefined, {
builtInTools: ["codemode"],
withoutToolSearch: true,
});
harness.setResponses([fauxAssistantMessage("ready")]);
await harness.session.prompt("start");
expect(harness.session.getActiveToolNames()).toEqual(["codemode"]);
expect(nestedToolNames(harness)).toContain("mcp__docs__search");
expect(notifications).toEqual([]);
});
it("warns when deferred MCP tools have neither tool_search nor codemode", async () => {
const { harness, notifications } = await setup("deferred", undefined, { withoutToolSearch: true });
harness.setResponses([fauxAssistantMessage("ready")]);
await harness.session.prompt("start");
expect(harness.session.getActiveToolNames()).toEqual([]);
expect(notifications).toEqual([
"MCP tools are only reachable from the codemode or tool_search tool, but neither is active; they cannot be called.",
]);
});
it("does not activate another extension's tool named codemode", async () => {
// Registered first, so it wins over the codemode extension's codemode.
const otherExec: ExtensionFactory = (pi) => {
pi.registerTool({
name: "codemode",
label: "codemode",
description: "Another extension's codemode tool.",
parameters: Type.Object({}),
defaultActive: false,
execute: async () => ({ content: [], details: undefined }),
});
};
const { harness } = await setup("codemode", undefined, { extensionFactories: [otherExec] });
harness.setResponses([fauxAssistantMessage("ready")]);
await harness.session.prompt("start");
expect(harness.session.getActiveToolNames()).toEqual([]);
});
it("does not hold the first prompt for servers that are still connecting", async () => {
const entry: McpServerEntry = { name: "slow", config: { url: "http://unused.invalid" }, source: "test" };
const harness = await createHarness({
initialActiveToolNames: [],
extensionFactories: [
createCodemodeExtension(),
createMcpExtension({
loadConfig: () => ({ servers: [entry], errors: [] }),
// The server never answers `initialize`.
createTransport: () => {
const pair = createInMemoryTransportPair();
void pair.server.start();
return pair.client;
},
startupWaitMs: 20,
}),
],
});
harnesses.push(harness);
await harness.session.bindExtensions({});
harness.setResponses([fauxAssistantMessage("ready")]);
await harness.session.prompt("start");
expect(getAssistantTexts(harness)).toEqual(["ready"]);
expect(harness.session.getActiveToolNames()).toEqual([]);
});
it("does not let codemode call itself or inactive direct tools", async () => {
const { harness } = await setup("direct");
harness.setResponses([fauxAssistantMessage("ready")]);
await harness.session.prompt("start");
harness.session.setActiveToolsByName(["codemode", "mcp__docs__search"]);
expect(nestedToolNames(harness)).toEqual(["mcp__docs__search"]);
});
it("finds tools from scripts with searchTools() and describeTool()", async () => {
const { harness, calls } = await setup("codemode-deferred");
harness.setResponses([
fauxAssistantMessage(
[
fauxToolCall("codemode", {
code: `
const [match] = await searchTools("search the docs", { limit: 1 });
const none = await searchTools("docs", { namespace: "mcp__other" });
const declaration = await describeTool(match.name);
const result = await tools[match.name]({ query: "found" });
text(JSON.stringify({
name: match.name,
sameAsAllTools: ALL_TOOLS.find((tool) => tool.name === match.name).description === match.description,
none: none.length,
declared: declaration.includes("codemode tool declaration:"),
missing: (await describeTool("nope")) === undefined,
hits: result.structuredContent.hits,
}));
`,
}),
],
{ stopReason: "toolUse" },
),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("go");
const result = toolResult(harness, "codemode");
expect(result.isError).toBe(false);
expect(JSON.parse((result.content[1] as { text: string }).text)).toEqual({
name: "mcp__docs__search",
sameAsAllTools: true,
none: 0,
declared: true,
missing: true,
hits: ["found guide", "found faq"],
});
expect(calls).toEqual(['search:{"query":"found"}']);
});
it("activates tool_search for deferred MCP tools and keeps loaded tools declared on the branch", async () => {
// No built-in tools are active; the MCP extension activates tool_search, not codemode.
const { harness, calls } = await setup("deferred");
const searchName = createMcpToolName("docs", "search");
harness.setResponses([
fauxAssistantMessage([fauxToolCall("tool_search", { query: "search the docs", limit: 1 })], {
stopReason: "toolUse",
}),
fauxAssistantMessage([fauxToolCall(searchName, { query: "loaded" })], { stopReason: "toolUse" }),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("find a docs tool");
expect(harness.session.getActiveToolNames()).toEqual(["tool_search", searchName]);
const toolSearch = harness.session.agent.state.tools.find((tool) => tool.name === "tool_search");
expect(toolSearch?.description).toContain("mcp__docs");
const search = toolResult(harness, "tool_search");
expect(getMessageText(search)).toBe(
`Loaded 1 tool. They are available from your next call:\n- ${searchName}: Search the docs.`,
);
// Only the loaded tool is added; earlier declarations are not repeated.
const loadMessages = harness.session.messages.filter(
(message): message is SystemMessage =>
message.role === "system" && (message.toolsAdded ?? []).some((tool) => tool.name === searchName),
);
expect(loadMessages).toHaveLength(1);
expect(loadMessages[0].toolsAdded?.map((tool) => tool.name)).toEqual([searchName]);
expect(getMessageText(toolResult(harness, searchName))).toBe("loaded guide\nloaded faq");
expect(calls).toEqual(['search:{"query":"loaded"}']);
// Loads are recorded in the transcript: navigating back before the load drops the tool,
// navigating to a later entry restores it.
const branch = harness.sessionManager.getBranch();
const firstUser = branch.find((entry) => entry.type === "message" && entry.message.role === "user");
const last = branch.at(-1);
if (!firstUser || !last) throw new Error("Missing entries");
await harness.session.navigateTree(firstUser.id);
expect(harness.session.getActiveToolNames()).not.toContain(searchName);
await harness.session.navigateTree(last.id);
expect(harness.session.getActiveToolNames()).toContain(searchName);
});
it("finds nothing to load when every matching tool is already declared", async () => {
const { harness } = await setup("direct", undefined, { builtInTools: ["tool_search"] });
harness.setResponses([
fauxAssistantMessage([fauxToolCall("tool_search", { query: "docs" })], { stopReason: "toolUse" }),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("go");
expect(getMessageText(toolResult(harness, "tool_search"))).toBe("No matching tools found.");
});
});
describe("AgentSession MCP servers registered by extensions", () => {
const harnesses: Harness[] = [];
afterEach(() => {
while (harnesses.length > 0) harnesses.pop()?.cleanup();
});
/** `configured` are the mcp.json servers; `plugins` register servers through the extension API. */
async function setup(plugins: ExtensionFactory | ExtensionFactory[], configured: McpServerEntry[] = []) {
const connected: McpServerEntry[] = [];
const harness = await createHarness({
initialActiveToolNames: [],
extensionFactories: [
...(Array.isArray(plugins) ? plugins : [plugins]),
createCodemodeExtension(),
createMcpExtension({
loadConfig: () => ({ servers: configured, errors: [] }),
createTransport: (entry) => {
connected.push(entry);
const pair = createFakeServer([]);
void pair.server.start();
return pair.client;
},
}),
],
});
harnesses.push(harness);
await harness.session.bindExtensions({});
return { harness, connected };
}
it("connects servers registered while extensions load", async () => {
const { harness, connected } = await setup((pi) => {
pi.registerMcpServer("plugin", { url: "http://plugin.invalid", exposure: "direct" });
});
harness.setResponses([fauxAssistantMessage("ready")]);
await harness.session.prompt("start");
expect(connected.map((entry) => [entry.name, entry.scope])).toEqual([["plugin", "extension"]]);
expect(harness.session.getActiveToolNames()).toContain("mcp__plugin__search");
});
it("connects and disconnects servers registered during the session", async () => {
let api: ExtensionAPI | undefined;
const { harness, connected } = await setup((pi) => {
api = pi;
});
if (!api) throw new Error("No extension API");
const pi = api;
pi.registerMcpServer("late", { url: "http://late.invalid" });
await vi.waitFor(() => expect(harness.session.getCallableToolNames()).toContain("mcp__late__search"));
expect(connected.map((entry) => entry.name)).toEqual(["late"]);
// Codemode-exposed tools need the codemode tool, which is activated for them.
expect(harness.session.getActiveToolNames()).toContain("codemode");
pi.unregisterMcpServer("late");
await vi.waitFor(() => expect(harness.session.getCallableToolNames()).not.toContain("mcp__late__search"));
});
it("prefers the mcp.json server over a registered server of the same name", async () => {
const configured: McpServerEntry = { name: "docs", config: { url: "http://config.invalid" }, source: "mcp.json" };
const { harness, connected } = await setup(
(pi) => {
pi.registerMcpServer("docs", { url: "http://plugin.invalid" });
},
[configured],
);
harness.setResponses([fauxAssistantMessage("ready")]);
await harness.session.prompt("start");
expect(connected).toEqual([configured]);
});
it("rejects names another extension registered", async () => {
let error: unknown;
await setup([
(pi) => {
pi.registerMcpServer("taken", { url: "http://x.invalid" });
// Registering again replaces the extension's own registration.
pi.registerMcpServer("taken", { url: "http://y.invalid" });
},
(pi) => {
try {
pi.registerMcpServer("taken", { url: "http://z.invalid" });
} catch (caught) {
error = caught;
}
},
]);
expect(String(error)).toMatch(/MCP server "taken" is already registered by extension/);
});
it("reports registered servers when no extension connects them", async () => {
const harness = await createHarness({
extensionFactories: [(pi) => pi.registerMcpServer("orphan", { url: "http://orphan.invalid" })],
});
harnesses.push(harness);
const errors: string[] = [];
await harness.session.bindExtensions({ onError: (error) => errors.push(error.error) });
expect(errors).toEqual([expect.stringContaining('MCP server "orphan" is registered, but no loaded extension')]);
});
});

View file

@ -0,0 +1,145 @@
import { fauxAssistantMessage, fauxToolCall, getCurrentTools, type TranscriptContext } from "@earendil-works/pi-ai";
import type { ToolResultMessage } from "@earendil-works/pi-ai/compat";
import { Type } from "typebox";
import { afterEach, describe, expect, it } from "vitest";
import type { ExtensionAPI } from "../../src/core/extensions/types.ts";
import { createCodemodeExtension } from "../../src/extensions/codemode/index.ts";
import { createToolSearchExtension } from "../../src/extensions/tool-search/index.ts";
import { createHarness, type Harness } from "./harness.ts";
/**
* A tool that calls other tools, built only on the extension API: its own name, exposure, loadout
* hook, and ctx.executeTool(). Codemode and tool search use the same mechanisms.
*/
function orchestratorExtension(pi: ExtensionAPI): void {
pi.registerTool({
name: "echo",
label: "echo",
description: "Echo text.",
parameters: Type.Object({ text: Type.String() }),
execute: async (_id, { text }) => ({ content: [{ type: "text", text: `echo: ${text}` }], details: {} }),
});
pi.registerTool({
name: "helper",
label: "helper",
description: "Only reachable from other tools.",
parameters: Type.Object({}),
exposure: "codemode",
execute: async () => ({ content: [{ type: "text", text: "helped" }], details: {} }),
});
pi.registerTool({
name: "run_tools",
label: "run_tools",
description: "Runs tools.",
parameters: Type.Object({}),
exposure: "model-only",
prepareLoadout: (loadout) => ({
descriptions: {
run_tools: `Runs tools: ${loadout.callable.map((tool) => tool.name).join(", ")}`,
echo: "Echo text (also callable from run_tools).",
},
hiddenDeclarations: ["echo"],
}),
execute: async (_id, _params, _signal, _onUpdate, ctx) => {
const helper = await ctx.executeTool("helper", {});
const echo = await ctx.executeTool("echo", { text: "hi" });
const self = await ctx.executeTool("run_tools", {});
const text = [helper, echo, self]
.map((outcome) => (outcome.result.content[0] as { text: string }).text)
.join(" | ");
return { content: [{ type: "text", text }], details: {} };
},
});
}
describe("AgentSession tool orchestration", () => {
const harnesses: Harness[] = [];
afterEach(() => {
while (harnesses.length > 0) harnesses.pop()?.cleanup();
});
it("supports tools that call other tools under any name through the extension API", async () => {
const toolCalls: string[] = [];
const harness = await createHarness({
initialActiveToolNames: [],
extensionFactories: [
orchestratorExtension,
(pi) => {
pi.on("tool_call", (event) => {
toolCalls.push(`${event.toolName}:${event.parentToolCallId ?? "top"}`);
});
},
],
});
harnesses.push(harness);
await harness.session.bindExtensions({});
expect(harness.session.getActiveToolNames()).toEqual(["echo", "run_tools"]);
expect(harness.session.getCallableToolNames()).toEqual(["echo", "helper"]);
const find = (name: string) => harness.session.agent.state.tools.find((tool) => tool.name === name);
expect(find("run_tools")?.description).toBe("Runs tools: echo, helper");
expect(find("echo")?.description).toBe("Echo text (also callable from run_tools).");
const requestTools: string[][] = [];
harness.setResponses([
(context: TranscriptContext) => {
requestTools.push(getCurrentTools(context.messages).map((tool) => tool.name));
return fauxAssistantMessage([fauxToolCall("run_tools", {})], { stopReason: "toolUse" });
},
fauxAssistantMessage("done"),
]);
await harness.session.prompt("go");
// echo stays active, but its declaration is left out of requests.
expect(requestTools[0]).toEqual(["run_tools"]);
const result = harness.session.messages.find(
(message): message is ToolResultMessage => message.role === "toolResult",
);
if (!result) throw new Error("No tool result");
expect(result.content).toEqual([{ type: "text", text: "helped | echo: hi | Tool run_tools not found" }]);
const parent = result.toolCallId;
expect(toolCalls).toEqual(["run_tools:top", `helper:${parent}`, `echo:${parent}`]);
expect(result.nestedCalls?.calls.map((call) => [call.id, call.name, call.status])).toEqual([
[`${parent}/1`, "helper", "ok"],
[`${parent}/2`, "echo", "ok"],
[`${parent}/3`, "run_tools", "error"],
]);
// The record is persisted with the session.
const persisted = harness.sessionManager
.getBranch()
.find((entry) => entry.type === "message" && entry.message.role === "toolResult");
expect(persisted?.type === "message" && persisted.message).toMatchObject({ nestedCalls: result.nestedCalls });
});
it("registers codemode and tool_search inactive until they are named", async () => {
const extensionFactories = [createCodemodeExtension(), createToolSearchExtension()];
const plain = await createHarness({ extensionFactories });
harnesses.push(plain);
expect(plain.session.getAllTools().map((tool) => tool.name)).toEqual(
expect.arrayContaining(["codemode", "tool_search"]),
);
expect(plain.session.getActiveToolNames()).toEqual(["read", "bash", "edit", "write"]);
// --tools and the defaultTools setting name them explicitly.
const allowed = await createHarness({ allowedToolNames: ["read", "codemode"], extensionFactories });
harnesses.push(allowed);
expect(allowed.session.getActiveToolNames()).toEqual(["read", "codemode"]);
const initial = await createHarness({ initialActiveToolNames: ["tool_search"], extensionFactories });
harnesses.push(initial);
expect(initial.session.getActiveToolNames()).toEqual(["tool_search"]);
});
it("leaves results without nested calls unchanged", async () => {
const harness = await createHarness({ initialActiveToolNames: [], extensionFactories: [orchestratorExtension] });
harnesses.push(harness);
await harness.session.bindExtensions({});
harness.setResponses([
fauxAssistantMessage([fauxToolCall("echo", { text: "x" })], { stopReason: "toolUse" }),
fauxAssistantMessage("done"),
]);
await harness.session.prompt("go");
const result = harness.session.messages.find((message) => message.role === "toolResult");
expect(result && "nestedCalls" in result).toBe(false);
});
});

Some files were not shown because too many files have changed in this diff Show more