sscoble/pi-mono
1
0
Fork 0
mirror of https://github.com/badlogic/pi-mono.git synced 2026-11-05 23:37:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

feat(p2): expose MCP tools to code mode only, as pi does

Browse source 
MCP tools are deferred and callable only by other tools, so code mode scripts call them and neither the model nor tool_search sees them. While code mode is off or its plugin is not loaded, tool_search loads them instead, as pi falls back. The mcp_servers prompt section, as pi's, names the servers and how their tools are reached.
This commit is contained in:
Armin Ronacher 2026-11-03 18:29:00 +01:00
parent 9f7d909f8b
commit 1bc7076b25
6 changed files with 338 additions and 133 deletions
Download patch file Download diff file Expand all files Collapse all files

82
apps/p2/plugins/mcp/tools.ts
View file

@ -1,5 +1,7 @@
// MCP tools as Durable tools: one deferred tool (found through `tool_search`) per tool of the session's servers, named
// as pi names them (`mcp__<server>__<tool>`), calling the server's instance where the calling conversation works.
// MCP tools as Durable tools, exposed as pi exposes them by default: one tool per tool of the session's servers, named as
// pi names them (`mcp__<server>__<tool>`), callable only from code mode scripts (or loaded through `tool_search` while
// code mode is off), calling the server's instance where the calling conversation works; and the `mcp_servers` prompt
// section, through which the model learns of the servers.
import type { JsonValue } from "@earendil-works/chord/json";
import {
@ -7,12 +9,14 @@ import {
type DocumentReader,
defineExtension,
type Extension,
section,
type ToolRegistration,
} from "@earendil-works/pi-durable";
import {
type CallToolResult,
createMcpToolName,
limitMcpContent,
mcpNamespace,
normalizeMcpInputSchema,
toLlmContent,
} from "@earendil-works/pi-mcp";
@ -22,6 +26,20 @@ import { type Instance, MCP_EXTENSION, type McpTool } from "./contract.ts";
/** One tool of one server. */
export type ServerTool = { readonly server: string; readonly tool: McpTool };
/**
* How the model reaches MCP tools, as pi's exposures: `codemode` (pi's default) from code mode scripts, which find them
* with `searchTools()`; `tool_search` (pi's `deferred`, its fallback without code mode) by loading them with
* `tool_search`, then calling them directly.
*/
export type McpReach = "codemode" | "tool_search";
/** What the MCP extension offers: the session's servers by name, their tools, and how the model reaches them. */
export type McpOffer = {
readonly servers: readonly string[];
readonly tools: readonly ServerTool[];
readonly reach: McpReach;
};
/** A call of `server`'s `tool` for conversation `conversationId` (its documents read through `read`). */
export type McpCall = (
request: {
@ -46,18 +64,24 @@ export function serverTools(servers: readonly Instance[]): ServerTool[] {
return [...tools.values()].sort((a, b) => order(a.server, b.server) || order(a.tool.name, b.tool.name));
}
/** The servers of `instances` by name, each once (a server may have instances at several locations). */
export function serverNames(instances: readonly Instance[]): string[] {
return [...new Set(instances.map((instance) => instance.server))].sort();
}
/**
* The Durable extension offering `tools`, each calling through `call`. As in pi (and Codex), every tool whose name
* sanitizes to one another tool's also gets pi's hash suffix, so no tool's name depends on the order of the list
* and a new tool renames only those it collides with.
* The Durable extension offering `offer`'s tools, each calling through `call`, and the `mcp_servers` section. As in pi
* (and Codex), every tool whose name sanitizes to one another tool's also gets pi's hash suffix, so no tool's name
* depends on the order of the list and a new tool renames only those it collides with.
*/
export function createMcpExtension(tools: readonly ServerTool[], call: McpCall): Extension {
export function createMcpExtension({ servers, tools, reach }: McpOffer, call: McpCall): Extension {
const plain = new Map<string, number>();
for (const { server, tool } of tools) {
const name = createMcpToolName(server, tool.name);
plain.set(name, (plain.get(name) ?? 0) + 1);
}
const assigned = new Set<string>();
const listing = serversSection(servers, reach);
return defineExtension({
name: MCP_EXTENSION,
tools: tools.map((entry) => {
@ -67,19 +91,59 @@ export function createMcpExtension(tools: readonly ServerTool[], call: McpCall):
(candidate) => (plain.get(candidate) ?? 0) > 1 || assigned.has(candidate),
);
assigned.add(name);
return toolRegistration(name, entry, call);
return toolRegistration(name, entry, reach, call);
}),
sections: listing === undefined ? [] : [section(MCP_SERVERS_SECTION, () => listing)],
});
}
/** One MCP tool as a deferred Durable tool (not replay-safe: an interrupted call is not repeated). */
function toolRegistration(name: string, { server, tool }: ServerTool, call: McpCall): ToolRegistration {
/** The section listing the servers whose tools are not declared to the model, as pi's. */
export const MCP_SERVERS_SECTION = "mcp_servers";
/**
* Characters of the whole section, as in pi: when the server lines do not fit, the last servers are left out and
* counted in a closing line.
*/
const MAX_SERVERS_SECTION_CHARS = 4096;
/**
* The `mcp_servers` section, as pi renders it for servers without a description: each server's namespace and how its
* tools are reached. Undefined without servers.
*/
function serversSection(servers: readonly string[], reach: McpReach): string | undefined {
if (servers.length === 0) return undefined;
const intro =
reach === "codemode"
? "MCP servers whose tools are not declared to you. Call the tools of `codemode` servers from codemode scripts."
: "MCP servers whose tools are not declared to you. Load the tools of `tool_search` servers with `tool_search`.";
const lines = servers.map((server) => `- ${mcpNamespace(server)} (${reach})`);
const find = reach === "codemode" ? "searchTools()" : "tool_search";
const omitted = (count: number) =>
count > 0 ? [`- … ${count} more server${count === 1 ? "" : "s"}; find their tools with ${find}`] : [];
const render = (kept: number) => [intro, ...lines.slice(0, kept), ...omitted(lines.length - kept)].join("\n");
let kept = lines.length;
while (kept > 0 && render(kept).length > MAX_SERVERS_SECTION_CHARS) kept--;
return render(kept);
}
/**
* One MCP tool as a Durable tool (not replay-safe: an interrupted call is not repeated). Deferred either way, so the
* `codemode` section leaves it out; for `codemode` only other tools (scripts) may call it, so `tool_search` never
* offers it to the model either.
*/
function toolRegistration(
name: string,
{ server, tool }: ServerTool,
reach: McpReach,
call: McpCall,
): ToolRegistration {
const title = tool.title ?? tool.annotations?.title;
return {
name,
description: tool.description?.trim() || title || `MCP tool ${tool.name} from server ${server}`,
parameters: normalizeMcpInputSchema(tool.inputSchema),
deferred: true,
...(reach === "codemode" ? { callers: ["tools" as const] } : {}),
...(tool.outputSchema === undefined ? {} : { structuredOutputSchema: tool.outputSchema }),
async execute(args, api, context) {
// Durable checked `args` against `parameters`, an object schema; its type cannot follow a runtime schema.

91
apps/p2/plugins/mcp/worker.ts
View file

@ -1,8 +1,9 @@
// The MCP plugin's worker facet: `tool_search` (deferred tools, MCP's among them, are found through it) and the worker
// object holding the server's `Mcp` view. It sets the locations the session's conversations work at (the session's,
// active subagents' own), with the project `.pi/mcp.json` read at each through its environment, as its leases, again
// whenever they may have changed and whenever the view (re)activates, and offers the servers' tools as one Durable
// extension that follows the view's `servers`. No MCP process or credential is ever in a worker.
// The MCP plugin's worker facet: the worker object holding the server's `Mcp` view. It sets the locations the session's
// conversations work at (the session's, active subagents' own), with the project `.pi/mcp.json` read at each through
// its environment, as its leases, again whenever they may have changed and whenever the view (re)activates, and offers
// the servers' tools as one Durable extension that follows the view's `servers` and the code mode switch: as in pi,
// tools are for code mode scripts, and only without code mode does `tool_search` load them for the model. No MCP
// process or credential is ever in a worker.
import { setTimeout as sleep } from "node:timers/promises";
import { type ConversationId, type Extension, type Harness, ROOT_CONVERSATION_ID } from "@earendil-works/pi-durable";
@ -21,12 +22,14 @@ import {
signalOf,
until,
untilAborted,
type Watchable,
type Writable,
} from "@p2/app/v1";
import { durableContext } from "@p2/app/v1/node";
import {
CommitSignal,
type ExecutionEnvironments,
ExtensionSwitches,
Extensions,
type HarnessParts,
harnessOf,
@ -34,15 +37,16 @@ import {
workerOf,
} from "@p2/app/v1/worker";
import { Subagents } from "../subagents/contract.worker.ts";
import { type Instance, MCP_EXTENSION, Mcp, PROJECT_FILE_MAX_CHARS, type ProjectFile } from "./contract.ts";
import { createMcpExtension, type McpCall, serverTools } from "./tools.ts";
import { MCP_EXTENSION, Mcp, PROJECT_FILE_MAX_CHARS, type ProjectFile } from "./contract.ts";
import { createMcpExtension, type McpCall, type McpOffer, type McpReach, serverNames, serverTools } from "./tools.ts";
/** How long to wait before setting the locations again after a failure. */
const RETRY_MS = 2000;
export default defineFacet((host, facet) => {
const search = createToolSearchExtension();
facet.contribute(Extensions, { [search.name]: search });
/** The code mode plugin's extension: its switch is on while the `codemode` setting is, absent without the plugin. */
const CODEMODE_EXTENSION = "codemode";
export default defineFacet((host) => {
contribute(host, "p2.mcp-worker", mcpWorker);
});
@ -58,36 +62,67 @@ function mcpWorker(self: StoreSelf): void {
context.abortSignal,
);
};
const extension = self.writable<Readonly<Record<string, Extension>>>({
[MCP_EXTENSION]: createMcpExtension([], call),
});
// Code mode switched off or not loaded: the tools would be unreachable, so `tool_search` loads them, as pi falls back.
const reach = self.derive(
[self.registry(ExtensionSwitches)],
([switches]): McpReach => (switches.get(CODEMODE_EXTENSION) === true ? "codemode" : "tool_search"),
);
const extension = self.writable(extensionsFor({ servers: [], tools: [], reach: "codemode" }, call));
self.contribute(Extensions, extension);
self.onLive((ctx) => offerTools(view, extension, call, ctx));
self.onLive((ctx) => offerTools(view, reach, extension, call, ctx));
leaseLocations(self, view, harness, worker);
}
/** Follows the view's `servers` until `ctx` ends: a new tool list (not a status change) makes a new extension. */
/** The extensions offering `offer`: the servers' tools and section, and `tool_search` while it reaches the tools. */
function extensionsFor(offer: McpOffer, call: McpCall): Readonly<Record<string, Extension>> {
const extensions: Record<string, Extension> = { [MCP_EXTENSION]: createMcpExtension(offer, call) };
if (offer.reach === "tool_search" && offer.servers.length > 0) {
const search = createToolSearchExtension();
extensions[search.name] = search;
}
return extensions;
}
/**
* Follows the view's `servers` and `reach` until `ctx` ends: a new server or tool list (not a status change) or a new
* reach makes new extensions.
*/
async function offerTools(
view: Ref<typeof Mcp>,
reach: Watchable<McpReach>,
extension: Writable<Readonly<Record<string, Extension>>>,
call: McpCall,
ctx: Ctx,
): Promise<void> {
let offered = "[]";
const offer = (servers: readonly Instance[]): void => {
const tools = serverTools(servers);
const key = JSON.stringify(tools);
if (key === offered) return;
offered = key;
extension.set({ [MCP_EXTENSION]: createMcpExtension(tools, call) });
};
const watch = await expectVariant(view, "worker").servers.watch(ctx);
try {
offer(watch.value);
watch.start(async (change) => {
if (change.kind === "value") offer(change.value);
});
await untilAborted(ctx);
const reaches = await reach.watch(ctx);
try {
let servers = watch.value;
let current = reaches.value;
let offered = "";
const offer = (): void => {
const next: McpOffer = { servers: serverNames(servers), tools: serverTools(servers), reach: current };
const key = JSON.stringify(next);
if (key === offered) return;
offered = key;
extension.set(extensionsFor(next, call));
};
offer();
watch.start(async (change) => {
if (change.kind !== "value") return;
servers = change.value;
offer();
});
reaches.start(async (change) => {
if (change.kind !== "value") return;
current = change.value;
offer();
});
await untilAborted(ctx);
} finally {
await reaches.stop();
}
} finally {
await watch.stop();
}