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

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