mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-06 23:44:36 +00:00
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
parent
b485fa3128
commit
8562bcf66a
143 changed files with 17620 additions and 225 deletions
|
|
@ -0,0 +1,2 @@
|
|||
/** Loads the codemode sandbox runtime on the first script (see execute.ts). */
|
||||
export const loadCodemodeExecutor = () => import("./execute.ts");
|
||||
428
packages/coding-agent/src/extensions/codemode/execute.ts
Normal file
428
packages/coding-agent/src/extensions/codemode/execute.ts
Normal 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],
|
||||
}));
|
||||
}
|
||||
46
packages/coding-agent/src/extensions/codemode/index.ts
Normal file
46
packages/coding-agent/src/extensions/codemode/index.ts
Normal 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();
|
||||
114
packages/coding-agent/src/extensions/codemode/renderer.ts
Normal file
114
packages/coding-agent/src/extensions/codemode/renderer.ts
Normal 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;
|
||||
},
|
||||
};
|
||||
451
packages/coding-agent/src/extensions/codemode/tool.ts
Normal file
451
packages/coding-agent/src/extensions/codemode/tool.ts
Normal 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;
|
||||
}
|
||||
6
packages/coding-agent/src/extensions/codemode/worker.ts
Normal file
6
packages/coding-agent/src/extensions/codemode/worker.ts
Normal 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";
|
||||
|
|
@ -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 },
|
||||
];
|
||||
|
|
|
|||
2
packages/coding-agent/src/extensions/mcp/cli.lazy.ts
Normal file
2
packages/coding-agent/src/extensions/mcp/cli.lazy.ts
Normal 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");
|
||||
596
packages/coding-agent/src/extensions/mcp/cli.ts
Normal file
596
packages/coding-agent/src/extensions/mcp/cli.ts
Normal 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);
|
||||
}
|
||||
}
|
||||
182
packages/coding-agent/src/extensions/mcp/config.ts
Normal file
182
packages/coding-agent/src/extensions/mcp/config.ts
Normal 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`);
|
||||
}
|
||||
985
packages/coding-agent/src/extensions/mcp/index.ts
Normal file
985
packages/coding-agent/src/extensions/mcp/index.ts
Normal 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();
|
||||
69
packages/coding-agent/src/extensions/mcp/log.ts
Normal file
69
packages/coding-agent/src/extensions/mcp/log.ts
Normal 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;
|
||||
}
|
||||
}
|
||||
}
|
||||
355
packages/coding-agent/src/extensions/mcp/oauth.ts
Normal file
355
packages/coding-agent/src/extensions/mcp/oauth.ts
Normal 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();
|
||||
}
|
||||
}
|
||||
342
packages/coding-agent/src/extensions/mcp/resources.ts
Normal file
342
packages/coding-agent/src/extensions/mcp/resources.ts
Normal 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];
|
||||
}
|
||||
2
packages/coding-agent/src/extensions/mcp/runtime.lazy.ts
Normal file
2
packages/coding-agent/src/extensions/mcp/runtime.lazy.ts
Normal 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");
|
||||
456
packages/coding-agent/src/extensions/mcp/runtime.ts
Normal file
456
packages/coding-agent/src/extensions/mcp/runtime.ts
Normal 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);
|
||||
}
|
||||
}
|
||||
284
packages/coding-agent/src/extensions/mcp/tools.ts
Normal file
284
packages/coding-agent/src/extensions/mcp/tools.ts
Normal 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?.() });
|
||||
},
|
||||
};
|
||||
}
|
||||
246
packages/coding-agent/src/extensions/mcp/ui.ts
Normal file
246
packages/coding-agent/src/extensions/mcp/ui.ts
Normal 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;
|
||||
});
|
||||
}
|
||||
18
packages/coding-agent/src/extensions/tool-search/index.ts
Normal file
18
packages/coding-agent/src/extensions/tool-search/index.ts
Normal 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();
|
||||
269
packages/coding-agent/src/extensions/tool-search/tool.ts
Normal file
269
packages/coding-agent/src/extensions/tool-search/tool.ts
Normal 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) } };
|
||||
},
|
||||
};
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue