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