From fcdd33d2fe8b337d186ec18bc04afcdf79700b38 Mon Sep 17 00:00:00 2001 From: Jeff Scott Ward Date: Sun, 26 Jul 2026 00:29:47 -0400 Subject: [PATCH] fix(mcp): map mounted tools to xd routes --- packages/coding-agent/CHANGELOG.md | 1 + .../src/prompts/system/mcp-xdev-guidance.md | 11 + packages/coding-agent/src/sdk.ts | 38 ++- .../coding-agent/src/session/session-tools.ts | 162 +++++++++--- .../agent-session-tool-rebuild-skip.test.ts | 246 +++++++++++++++++- .../test/fixtures/instructions-mcp.ts | 59 +++-- .../test/sdk-mcp-instructions.test.ts | 160 +++++++++++- 7 files changed, 600 insertions(+), 77 deletions(-) create mode 100644 packages/coding-agent/src/prompts/system/mcp-xdev-guidance.md diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index c4777f2be..e52ab04b2 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -9,6 +9,7 @@ ### Fixed +- Fixed mounted MCP tools being hard to invoke when server or plugin guidance names their original calls: sessions now include one bounded, exact original-name-to-`xd://` route map for every live mounted MCP tool—including servers without initialize instructions—and refresh it as catalogs change without disabling schema virtualization. - Fixed the Docker `natives-builder` stage failing to build releases ≥ 17.1.1: the native audio stack added bindgen (miniaudio needs libclang) and a bundled-opus CMake build (needs cmake + make), none of which were installed in the slim builder image. - Fixed `omp usage` duplicating org-less legacy accounts as "no usage data" rows whenever any sibling report carried an organization (mixed pools of pre-org-capture rows and fresh org-scoped logins): an org-less account is now covered by its own org-less report, while org-attributed sibling reports still never count as its coverage. - `omp usage` revalidates the broker credential snapshot before rendering: live usage reports were previously paired with a disk-cached account list up to an hour old, so a just-completed re-login (org-less row upserted to org-scoped) rendered as a phantom duplicate until the cache expired. diff --git a/packages/coding-agent/src/prompts/system/mcp-xdev-guidance.md b/packages/coding-agent/src/prompts/system/mcp-xdev-guidance.md new file mode 100644 index 000000000..7a36afe82 --- /dev/null +++ b/packages/coding-agent/src/prompts/system/mcp-xdev-guidance.md @@ -0,0 +1,11 @@ +## MCP Tool Routes + +{{#if tools.length}} +Execute each mounted tool by writing JSON arguments to its mounted path: +{{#each tools}} +- {{mcpToolName}} → `{{path}}` +{{/each}} +{{/if}} +{{#if hasOmittedTools}} +Additional mounted MCP tool mappings were omitted to keep this prompt bounded. Inspect `xd://` for the exact current paths. +{{/if}} diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index a08aa5683..4b80aaa6c 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -109,6 +109,7 @@ import { MCP_CONNECTION_STATUS_EVENT_CHANNEL, type McpConnectionStatusEvent } fr import { createSessionMemoryRuntimeContext, resolveMemoryBackend } from "./memory-backend"; import { MEMORY_BACKEND_TOOL_NAMES } from "./memory-backend/tool-names"; import type { MnemopiSessionState } from "./mnemopi/state"; +import mcpXdevGuidanceTemplate from "./prompts/system/mcp-xdev-guidance.md" with { type: "text" }; import lateDiagnosticTemplate from "./prompts/tools/lsp-late-diagnostic.md" with { type: "text" }; import { AgentLifecycleManager } from "./registry/agent-lifecycle"; import { type AgentRef, AgentRegistry, MAIN_AGENT_ID } from "./registry/agent-registry"; @@ -145,6 +146,7 @@ import { } from "./session/retry-fallback-chains"; import { getRestorableSessionModels } from "./session/session-context"; import { SessionManager } from "./session/session-manager"; +import { collectMountedMCPToolRoutes, projectMountedMCPXdevGuidance } from "./session/session-tools"; import { createSettingsAwareStreamFn } from "./session/settings-stream-fn"; import { SnapcompactInlineTransformer } from "./session/snapcompact-inline"; import { createSnapcompactSavingsRecorder } from "./session/snapcompact-savings-journal"; @@ -848,7 +850,6 @@ function isLegacyBuiltinToolDefinition(tool: CustomTool | ToolDefinition): boole } const TOOL_DEFINITION_MARKER = Symbol("__isToolDefinition"); - /** Matches the truncation applied to per-server instructions inside `rebuildSystemPrompt`. */ const MAX_MCP_INSTRUCTIONS_LENGTH = 4000; @@ -2632,11 +2633,11 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} : undefined; // Build combined append prompt: memory instructions + auto-learn guidance - // + MCP server instructions. For UI sessions MCP discovery is deferred, so - // `getServerInstructions()` is empty until the background connect completes; - // the rebuild that `refreshMCPTools` triggers post-discovery then picks up - // the now-connected servers' instructions, so they join the prompt for the - // rest of the session. + // + mounted MCP route guidance + optional MCP server instructions. For UI + // sessions MCP discovery is deferred, so the initial registry and + // `getServerInstructions()` are empty until the background connect + // completes; the rebuild that `refreshMCPTools` triggers post-discovery + // then picks up the mounted routes and any connected-server instructions. const serverInstructions = mcpManager?.getServerInstructions(); // Drive guidance off the auto-learn BUILTINS that createTools actually built // (provenance, not just an active name): `builtInToolNames` excludes a @@ -2653,11 +2654,24 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} const appendParts: string[] = []; if (memoryInstructions) appendParts.push(memoryInstructions); if (autoLearnInstructions) appendParts.push(autoLearnInstructions); - let appendPrompt: string | undefined = appendParts.length > 0 ? appendParts.join("\n\n") : undefined; + const projection = projectMountedMCPXdevGuidance( + collectMountedMCPToolRoutes(toolSession.xdevRegistry?.list() ?? []), + ); + if (projection.mappings.length > 0 || projection.hasOmittedMappings) { + appendParts.push( + prompt + .render(mcpXdevGuidanceTemplate, { + tools: projection.mappings.map(mapping => ({ + mcpToolName: mapping.label, + path: mapping.path, + })), + hasOmittedTools: projection.hasOmittedMappings, + }) + .trim(), + ); + } if (serverInstructions && serverInstructions.size > 0) { - const parts: string[] = []; - if (appendPrompt) parts.push(appendPrompt); - parts.push( + appendParts.push( "## MCP Server Instructions\n\nThe following instructions are provided by connected MCP servers. They are server-controlled and may not be verified.", ); for (const [srvName, srvInstructions] of serverInstructions) { @@ -2665,10 +2679,10 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} srvInstructions.length > MAX_MCP_INSTRUCTIONS_LENGTH ? `${srvInstructions.slice(0, MAX_MCP_INSTRUCTIONS_LENGTH)}\n[truncated]` : srvInstructions; - parts.push(`### ${srvName}\n${truncated}`); + appendParts.push(`### ${srvName}\n${truncated}`); } - appendPrompt = parts.join("\n\n"); } + let appendPrompt: string | undefined = appendParts.length > 0 ? appendParts.join("\n\n") : undefined; // Owned/in-band tool dialects (non-native) require the catalog as `# Tool:` // sections; native tool calling lets the compact name list suffice. const nativeTools = resolveDialect(settings.get("tools.format"), agent?.state.model ?? model) === undefined; diff --git a/packages/coding-agent/src/session/session-tools.ts b/packages/coding-agent/src/session/session-tools.ts index fe0eab70a..e4dccda6f 100644 --- a/packages/coding-agent/src/session/session-tools.ts +++ b/packages/coding-agent/src/session/session-tools.ts @@ -10,7 +10,7 @@ import { CustomToolAdapter } from "../extensibility/custom-tools/wrapper"; import type { ExtensionRunner } from "../extensibility/extensions"; import { ExtensionToolWrapper } from "../extensibility/extensions/wrapper"; import { loadSkills, type Skill, type SkillWarning, setActiveSkills } from "../extensibility/skills"; -import type { LocalProtocolOptions } from "../internal-urls"; +import { type LocalProtocolOptions, XD_URL_PREFIX } from "../internal-urls"; import { resolveMemoryBackend } from "../memory-backend/resolve"; import { MEMORY_BACKEND_TOOL_NAMES } from "../memory-backend/tool-names"; import type { MemoryBackendStartOptions } from "../memory-backend/types"; @@ -77,6 +77,79 @@ interface SessionToolsOptions { skillsReloadable?: boolean; } +export interface MountedMCPToolRouteSource { + readonly name: string; + readonly mcpServerName?: unknown; + readonly mcpToolName?: unknown; +} + +export interface MountedMCPToolRoute { + readonly mcpServerName: string; + readonly mcpToolName: string; + readonly name: string; +} + +export interface MCPXdevGuidanceMapping extends MountedMCPToolRoute { + readonly label: string; + readonly path: string; +} + +export interface MCPXdevGuidanceProjection { + readonly mappings: readonly MCPXdevGuidanceMapping[]; + readonly hasOmittedMappings: boolean; +} + +const MAX_MCP_XDEV_GUIDANCE_MAPPING_DATA_LENGTH = 4000; +const MAX_MCP_XDEV_GUIDANCE_MAPPINGS = 64; + +/** Yield exact mounted MCP ownership and route metadata. */ +export function* collectMountedMCPToolRoutes( + tools: Iterable, +): Generator { + for (const tool of tools) { + if (typeof tool.mcpServerName !== "string" || typeof tool.mcpToolName !== "string") continue; + yield { + mcpServerName: tool.mcpServerName, + mcpToolName: tool.mcpToolName, + name: tool.name, + }; + } +} + +function formatMCPXdevGuidanceLabel(label: string): string { + return (JSON.stringify(label) ?? '""') + .replaceAll("`", "\\u0060") + .replaceAll("\u2028", "\\u2028") + .replaceAll("\u2029", "\\u2029"); +} + +/** + * Project exact live MCP routes into the bounded, Markdown-safe mapping data + * rendered by the static MCP guidance prompt. + */ +export function projectMountedMCPXdevGuidance(routes: Iterable): MCPXdevGuidanceProjection { + const mappings: MCPXdevGuidanceMapping[] = []; + let remainingMappingDataLength = MAX_MCP_XDEV_GUIDANCE_MAPPING_DATA_LENGTH; + let hasOmittedMappings = false; + for (const route of routes) { + const rawMappingDataLength = route.mcpToolName.length + XD_URL_PREFIX.length + route.name.length; + if (mappings.length >= MAX_MCP_XDEV_GUIDANCE_MAPPINGS || rawMappingDataLength > remainingMappingDataLength) { + hasOmittedMappings = true; + continue; + } + const label = formatMCPXdevGuidanceLabel(route.mcpToolName); + const path = `${XD_URL_PREFIX}${route.name}`; + const mappingDataLength = label.length + path.length; + if (mappingDataLength > remainingMappingDataLength) { + hasOmittedMappings = true; + continue; + } + mappings.push({ ...route, label, path }); + remainingMappingDataLength -= mappingDataLength; + } + return { mappings, hasOmittedMappings }; +} + const XDEV_MOUNT_NOTICE_MESSAGE_TYPE = "xdev-mount-notice"; /** Owns tool registration, presentation, prompt rebuilding, skills, and permissions. */ @@ -96,6 +169,7 @@ export class SessionTools { #runtimeSelectedToolNames: ReadonlySet | undefined; #baseSystemPrompt: string[]; #lastAppliedToolSignature: string | undefined; + #mcpRefreshTail: Promise = Promise.resolve(); #promptModelKey: string | undefined; #rebuildSystemPrompt: SessionToolsOptions["rebuildSystemPrompt"]; #getLocalCalendarDate: () => string; @@ -518,6 +592,13 @@ export class SessionTools { throw error; } + if (this.#host.isDisposed()) { + this.#mountedXdevToolNames = previousMounted; + this.#xdevRegistry?.reconcile(previousMountedTools); + this.#setActiveToolNames?.(previousActiveToolNames); + return; + } + this.#notifyXdevMountDelta(previousMounted); this.#host.agent.setTools(tools); if (rebuiltSystemPrompt && rebuiltSignature) { @@ -531,15 +612,15 @@ export class SessionTools { } /** - * Record a mid-session `xd://` mount delta for the model without rewriting - * the system prompt: the prompt (and its provider cache prefix) stays - * byte-stable across MCP connects and disconnects. The delta is NOT steered - * immediately — a steered notice landing at a run's stop boundary (or while - * the session is idle) forces an unsolicited extra assistant turn — it is - * coalesced into {@link #pendingXdevMountDelta} and rides along with the - * next prompt (docs + schema stay one `read xd://` away). The full - * docs join the system prompt opportunistically on the next unrelated - * rebuild. + * Record a mid-session `xd://` mount delta for the model. Non-MCP mount + * churn remains notice-only, leaving the system prompt and provider cache + * prefix byte-stable; mounted MCP route changes additionally rebuild the + * global route guidance through the applied-tool signature. The delta is NOT + * steered immediately — a steered notice landing at a run's stop boundary + * (or while the session is idle) forces an unsolicited extra assistant turn + * — so it is coalesced into {@link #pendingXdevMountDelta} and rides along + * with the next prompt (docs + schema stay one `read xd://` away). + * Full docs join the system prompt opportunistically on a rebuild. */ #notifyXdevMountDelta(previousMounted: ReadonlySet): void { const registry = this.#xdevRegistry; @@ -813,9 +894,10 @@ export class SessionTools { * `tool.customWireName` and overrides the internal name on the model wire * (e.g. `edit` exposes itself as `apply_patch` to GPT-5 in apply_patch mode); * a stale wire name would desync prompt guidance from actual tool routing. - * 3. When MCP discovery is on, every registry tool's name+label+description+ - * customWireName, since `rebuildSystemPrompt` summarizes discoverable MCP - * tools that are not in the active set. + * 3. The bounded mounted-MCP projection: escaped original-name labels, + * actual `xd://` paths, and the omission flag in catalog order. These are + * the exact values rendered by the global transport guidance; catalog + * churn wholly behind the fallback does not change the prompt. * 4. MCP server instructions text (per server), since `rebuildSystemPrompt` * embeds these in the appended prompt under "## MCP Server Instructions". * A server upgrade can change instructions while keeping tools identical. @@ -846,8 +928,16 @@ export class SessionTools { const describeTool = (tool: AgentTool): string => `${tool.name}=${tool.label ?? ""}|${tool.description ?? ""}|${tool.customWireName ?? ""}`; const descriptionSegment = tools.map(describeTool).join("\u0002"); - let instructionsSegment = ""; + const mountedMCPProjection = projectMountedMCPXdevGuidance( + collectMountedMCPToolRoutes(this.#xdevRegistry?.list() ?? []), + ); + const mountedMCPRouteSegment = + JSON.stringify({ + mappings: mountedMCPProjection.mappings.map(mapping => [mapping.label, mapping.path] as const), + hasOmittedMappings: mountedMCPProjection.hasOmittedMappings, + }) ?? "{}"; const serverInstructions = this.#getMcpServerInstructions?.(); + let instructionsSegment = ""; if (serverInstructions && serverInstructions.size > 0) { // Sort by server name so transport flap order does not perturb the signature. const entries: string[] = []; @@ -857,22 +947,32 @@ export class SessionTools { entries.sort(); instructionsSegment = entries.join("\u0006"); } - // The xd:// device inventory is deliberately NOT part of the signature: - // a mount/unmount announces itself via `#notifyXdevMountDelta` instead of - // rewriting the system prompt, so MCP connects/disconnects keep the - // prompt (and its provider cache prefix) byte-stable. Rebuilds triggered - // by other inputs pick up the current device docs opportunistically. + // The non-MCP remainder of the xd:// inventory is deliberately NOT part + // of the signature: its mount/unmount announces itself through + // `#notifyXdevMountDelta` rather than rewriting the system prompt, keeping + // the provider cache prefix byte-stable. Mounted MCP routes are the narrow + // exception above, bounded to the exact projection rendered in the global + // route guidance so churn wholly behind its fallback does not rebuild. const date = this.#getLocalCalendarDate(); - return `${nameSegment}\u0003${descriptionSegment}\u0007${instructionsSegment}|${date}`; + return `${nameSegment}\u0003${descriptionSegment}\u0007${instructionsSegment}\u0008${mountedMCPRouteSegment}|${date}`; } /** - * Replace MCP tools in the registry and enable them immediately. Every - * connected MCP tool becomes available (mounted under `xd://` when that - * transport is active, else top-level). Lets `/mcp add/remove/reauth` take - * effect without restarting the session. + * Replace MCP tools in the registry and enable them immediately. Refreshes + * are serialized so an older asynchronous prompt rebuild cannot commit + * after a newer catalog snapshot. Every connected MCP tool becomes available + * (mounted under `xd://` when that transport is active, else top-level). */ - async refreshMCPTools(mcpTools: CustomTool[]): Promise { + refreshMCPTools(mcpTools: CustomTool[]): Promise { + const snapshot = [...mcpTools]; + const refresh = this.#mcpRefreshTail.then(() => + this.#host.isDisposed() ? undefined : this.#applyMCPToolRefresh(snapshot), + ); + this.#mcpRefreshTail = refresh.catch(() => {}); + return refresh; + } + + async #applyMCPToolRefresh(mcpTools: CustomTool[]): Promise { const existingNames = Array.from(this.#toolRegistry.keys()); const previousMcpTools = new Map( existingNames.flatMap(name => { @@ -880,6 +980,12 @@ export class SessionTools { return isMCPToolName(name) && tool ? [[name, tool] as const] : []; }), ); + const restorePreviousMcpTools = () => { + for (const name of this.#toolRegistry.keys()) { + if (isMCPToolName(name)) this.#toolRegistry.delete(name); + } + for (const [name, tool] of previousMcpTools) this.#toolRegistry.set(name, tool); + }; for (const name of existingNames) { if (isMCPToolName(name)) { this.#toolRegistry.delete(name); @@ -913,11 +1019,9 @@ export class SessionTools { const nextActive = [...new Set([...this.#getActiveNonMCPToolNames(), ...mcpTools.map(tool => tool.name)])]; try { await this.applyActiveToolsByName(nextActive); + if (this.#host.isDisposed()) restorePreviousMcpTools(); } catch (error) { - for (const name of this.#toolRegistry.keys()) { - if (isMCPToolName(name)) this.#toolRegistry.delete(name); - } - for (const [name, tool] of previousMcpTools) this.#toolRegistry.set(name, tool); + restorePreviousMcpTools(); throw error; } } diff --git a/packages/coding-agent/test/agent-session-tool-rebuild-skip.test.ts b/packages/coding-agent/test/agent-session-tool-rebuild-skip.test.ts index f9b8a5a0f..ad2261d7c 100644 --- a/packages/coding-agent/test/agent-session-tool-rebuild-skip.test.ts +++ b/packages/coding-agent/test/agent-session-tool-rebuild-skip.test.ts @@ -8,6 +8,10 @@ import type { CustomTool } from "@oh-my-pi/pi-coding-agent/extensibility/custom- import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; import { convertToLlm } from "@oh-my-pi/pi-coding-agent/session/messages"; import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; +import { + collectMountedMCPToolRoutes, + projectMountedMCPXdevGuidance, +} from "@oh-my-pi/pi-coding-agent/session/session-tools"; import { XdevRegistry } from "@oh-my-pi/pi-coding-agent/tools/xdev"; import { type } from "arktype"; @@ -96,6 +100,8 @@ describe("AgentSession refreshMCPTools rebuild skipping", () => { session: AgentSession; /** Provider-call message snapshots (LLM-converted), one per model request. */ contexts: Message[][]; + /** Mutable registry shared with the session, for lifecycle-only mount fixtures. */ + toolRegistry: Map; } { const readTool = createBasicTool("read", "Read"); const initialMcp = createMcpCustomTool("mcp__nucleus_search", "nucleus", "search", "Search nucleus"); @@ -147,7 +153,7 @@ describe("AgentSession refreshMCPTools rebuild skipping", () => { xdevRegistry: options.xdevRegistry, }); sessions.push(session); - return { session, contexts }; + return { session, contexts, toolRegistry }; } it("skips rebuild when an MCP refresh produces an identical tool set", async () => { @@ -170,10 +176,244 @@ describe("AgentSession refreshMCPTools rebuild skipping", () => { expect(rebuildCount).toBe(1); // Third refresh, again identical: still no rebuild. + await session.refreshMCPTools([initialMcp]); expect(rebuildCount).toBe(1); }); + it("serializes concurrent MCP refreshes before committing rebuilt prompts", async () => { + const firstRebuildStarted = Promise.withResolvers(); + const releaseFirstRebuild = Promise.withResolvers(); + const releaseSecondRebuild = Promise.withResolvers(); + let rebuildCount = 0; + const { session } = newSession(async toolNames => { + rebuildCount++; + if (rebuildCount === 1) { + firstRebuildStarted.resolve(); + await releaseFirstRebuild.promise; + } else { + await releaseSecondRebuild.promise; + } + return `tools:${toolNames.join(",")}`; + }); + const search = createMcpCustomTool("mcp__nucleus_search", "nucleus", "search", "Search nucleus"); + const fetch = createMcpCustomTool("mcp__nucleus_fetch", "nucleus", "fetch", "Fetch nucleus"); + + const olderRefresh = session.refreshMCPTools([search]); + await firstRebuildStarted.promise; + const newerRefresh = session.refreshMCPTools([search, fetch]); + + // Even if the newer rebuild would be allowed to resolve first, it cannot + // start until the older refresh has committed. + releaseSecondRebuild.resolve(); + expect(rebuildCount).toBe(1); + + releaseFirstRebuild.resolve(); + await Promise.all([olderRefresh, newerRefresh]); + expect(rebuildCount).toBe(2); + expect(session.systemPrompt).toEqual(["tools:read,mcp__nucleus_search,mcp__nucleus_fetch"]); + }); + + it("drops queued and in-flight MCP prompt commits when disposal begins", async () => { + const firstRebuildStarted = Promise.withResolvers(); + const releaseFirstRebuild = Promise.withResolvers(); + let rebuildCount = 0; + const { session, toolRegistry } = newSession(async toolNames => { + rebuildCount++; + firstRebuildStarted.resolve(); + await releaseFirstRebuild.promise; + return `tools:${toolNames.join(",")}`; + }); + const initialPrompt = [...session.systemPrompt]; + const initialToolNames = session.getActiveToolNames(); + const initialSearchTool = toolRegistry.get("mcp__nucleus_search"); + const search = createMcpCustomTool("mcp__nucleus_search", "nucleus", "search", "Search nucleus"); + const fetch = createMcpCustomTool("mcp__nucleus_fetch", "nucleus", "fetch", "Fetch nucleus"); + + const inFlightRefresh = session.refreshMCPTools([search]); + await firstRebuildStarted.promise; + const queuedRefresh = session.refreshMCPTools([search, fetch]); + session.beginDispose(); + releaseFirstRebuild.resolve(); + + await Promise.all([inFlightRefresh, queuedRefresh]); + expect(rebuildCount).toBe(1); + expect(session.systemPrompt).toEqual(initialPrompt); + expect(session.getActiveToolNames()).toEqual(initialToolNames); + expect(toolRegistry.get("mcp__nucleus_search")).toBe(initialSearchTool); + expect(toolRegistry.has("mcp__nucleus_fetch")).toBe(false); + }); + + it("rebuilds generated guidance when its ordered mounted MCP route projection changes", async () => { + const xdevRegistry = new XdevRegistry([]); + const serverInstructions = new Map([ + ["archive", "Archive instructions"], + ["nucleus", "Nucleus instructions"], + ]); + const renderedPrompts: string[] = []; + let rebuildCount = 0; + const { session } = newSession( + async () => { + rebuildCount++; + const projection = projectMountedMCPXdevGuidance(collectMountedMCPToolRoutes(xdevRegistry.list())); + const generatedPrompt = `mounted:${projection.mappings + .map(mapping => `${mapping.label}=${mapping.path}`) + .join(",")}`; + renderedPrompts.push(generatedPrompt); + return generatedPrompt; + }, + { xdevRegistry, getMcpServerInstructions: () => serverInstructions }, + ); + const search = createMcpCustomTool("mcp__nucleus_search", "nucleus", "search", "Search nucleus"); + const fetch = createMcpCustomTool("mcp__nucleus_fetch", "nucleus", "fetch", "Fetch nucleus"); + const uninstructed = createMcpCustomTool("mcp__silent_ping", "silent", "ping", "Ping silently"); + const searchPrompt = 'mounted:"search"=xd://mcp__nucleus_search'; + const searchAndUninstructedPrompt = 'mounted:"search"=xd://mcp__nucleus_search,"ping"=xd://mcp__silent_ping'; + const searchFetchAndUninstructedPrompt = + 'mounted:"search"=xd://mcp__nucleus_search,"fetch"=xd://mcp__nucleus_fetch,"ping"=xd://mcp__silent_ping'; + + await session.refreshMCPTools([search]); + expect(rebuildCount).toBe(1); + expect(session.systemPrompt).toEqual([searchPrompt]); + + // A new object with the same ordered route identity is still the same + // externally rendered inventory, so reconnecting it must preserve the + // cached prompt rather than rebuilding. + const equivalentSearch = createMcpCustomTool("mcp__nucleus_search", "nucleus", "search", "Search nucleus"); + await session.refreshMCPTools([equivalentSearch]); + expect(rebuildCount).toBe(1); + expect(session.systemPrompt).toEqual([searchPrompt]); + + // Global route guidance is independent of optional server instructions. + // Adding a route for a server absent from the instructions map must rebuild. + await session.refreshMCPTools([equivalentSearch, uninstructed]); + expect(rebuildCount).toBe(2); + expect(session.systemPrompt).toEqual([searchAndUninstructedPrompt]); + + await session.refreshMCPTools([equivalentSearch, fetch, uninstructed]); + expect(rebuildCount).toBe(3); + expect(session.systemPrompt).toEqual([searchFetchAndUninstructedPrompt]); + + const fetchSearchAndUninstructedPrompt = + 'mounted:"fetch"=xd://mcp__nucleus_fetch,"search"=xd://mcp__nucleus_search,"ping"=xd://mcp__silent_ping'; + await session.refreshMCPTools([fetch, equivalentSearch, uninstructed]); + expect(rebuildCount).toBe(4); + expect(session.systemPrompt).toEqual([fetchSearchAndUninstructedPrompt]); + + const replacementSearch = createMcpCustomTool("mcp__nucleus_search", "nucleus", "search", "Search nucleus"); + await session.refreshMCPTools([replacementSearch, uninstructed]); + expect(rebuildCount).toBe(5); + expect(session.systemPrompt).toEqual([searchAndUninstructedPrompt]); + const stableLabel = replacementSearch.label; + const reownedSearch = { + ...createMcpCustomTool("mcp__nucleus_search", "archive", "search", "Search nucleus"), + label: stableLabel, + }; + // Ownership alone is not rendered in the global route projection. + await session.refreshMCPTools([reownedSearch, uninstructed]); + expect(rebuildCount).toBe(5); + expect(session.systemPrompt).toEqual([searchAndUninstructedPrompt]); + + const renamedOriginalSearch = { + ...createMcpCustomTool("mcp__nucleus_search", "archive", "lookup", "Search nucleus"), + label: stableLabel, + }; + const renamedOriginalAndUninstructedPrompt = + 'mounted:"lookup"=xd://mcp__nucleus_search,"ping"=xd://mcp__silent_ping'; + await session.refreshMCPTools([renamedOriginalSearch, uninstructed]); + expect(rebuildCount).toBe(6); + expect(session.systemPrompt).toEqual([renamedOriginalAndUninstructedPrompt]); + + const remountedSearch = { + ...createMcpCustomTool("mcp__archive_lookup", "archive", "lookup", "Search nucleus"), + label: stableLabel, + }; + const remountedAndUninstructedPrompt = 'mounted:"lookup"=xd://mcp__archive_lookup,"ping"=xd://mcp__silent_ping'; + await session.refreshMCPTools([remountedSearch, uninstructed]); + expect(rebuildCount).toBe(7); + expect(session.systemPrompt).toEqual([remountedAndUninstructedPrompt]); + + const equivalentRemountedSearch = { + ...createMcpCustomTool("mcp__archive_lookup", "archive", "lookup", "Search nucleus"), + label: stableLabel, + }; + await session.refreshMCPTools([equivalentRemountedSearch, uninstructed]); + expect(rebuildCount).toBe(7); + expect(session.systemPrompt).toEqual([remountedAndUninstructedPrompt]); + + // Removing the uninstructed server's rendered route also changes guidance. + const remountedPrompt = 'mounted:"lookup"=xd://mcp__archive_lookup'; + await session.refreshMCPTools([equivalentRemountedSearch]); + expect(rebuildCount).toBe(8); + expect(session.systemPrompt).toEqual([remountedPrompt]); + expect(renderedPrompts).toEqual([ + searchPrompt, + searchAndUninstructedPrompt, + searchFetchAndUninstructedPrompt, + fetchSearchAndUninstructedPrompt, + searchAndUninstructedPrompt, + renamedOriginalAndUninstructedPrompt, + remountedAndUninstructedPrompt, + remountedPrompt, + ]); + }); + + it("skips rebuild when only an omitted mounted MCP mapping changes", async () => { + const xdevRegistry = new XdevRegistry([]); + let rebuildCount = 0; + const { session } = newSession( + async () => { + rebuildCount++; + return "bounded mounted MCP guidance"; + }, + { xdevRegistry }, + ); + const tools = Array.from({ length: 65 }, (_, index) => + createMcpCustomTool(`mcp__archive_tool_${index}`, "archive", `tool_${index}`, "Archive tool"), + ); + + await session.refreshMCPTools(tools); + expect(rebuildCount).toBe(1); + + const changedOmittedTool = { + ...createMcpCustomTool("mcp__archive_tool_64", "archive", "renamed_tail", "Archive tool"), + label: tools[64]!.label, + }; + await session.refreshMCPTools([...tools.slice(0, 64), changedOmittedTool]); + expect(rebuildCount).toBe(1); + }); + + it("skips rebuild when only non-MCP xd mounts change", async () => { + const xdevRegistry = new XdevRegistry([]); + let rebuildCount = 0; + const { session, toolRegistry } = newSession( + async toolNames => { + rebuildCount++; + return `tools:${toolNames.join(",")}`; + }, + { xdevRegistry }, + ); + const search = createMcpCustomTool("mcp__nucleus_search", "nucleus", "search", "Search nucleus"); + const catalog = { + ...createBasicTool("catalog_lookup", "Catalog Lookup"), + loadMode: "discoverable" as const, + }; + toolRegistry.set(catalog.name, catalog); + + await session.refreshMCPTools([search]); + expect(rebuildCount).toBe(1); + + // Ordinary xd:// inventory changes travel through mount notices and do + // not affect the global MCP-route guidance or its rebuild signature. + await session.setActiveToolPresentation(["read", search.name, catalog.name], [search.name, catalog.name]); + expect(session.getMountedXdevToolNames()).toContain(catalog.name); + expect(rebuildCount).toBe(1); + + await session.setActiveToolPresentation(["read", search.name], [search.name]); + expect(session.getMountedXdevToolNames()).not.toContain(catalog.name); + expect(rebuildCount).toBe(1); + }); + it("rebuilds when an MCP tool's description changes", async () => { let rebuildCount = 0; const { session } = newSession(async toolNames => { @@ -575,7 +815,7 @@ describe("AgentSession refreshMCPTools rebuild skipping", () => { await session.refreshMCPTools([search, fetch]); releaseFirstCall.resolve(); await firstPrompt; - expect(rebuildCount).toBe(1); + expect(rebuildCount).toBe(2); expect(contexts).toHaveLength(1); expect(mountNoticesIn(contexts[0])).toHaveLength(0); @@ -591,7 +831,7 @@ describe("AgentSession refreshMCPTools rebuild skipping", () => { // A later unmount is likewise held for the following user prompt. await session.refreshMCPTools([search]); - expect(rebuildCount).toBe(1); + expect(rebuildCount).toBe(3); expect(contexts).toHaveLength(2); await session.prompt("third"); const allNotices = mountNoticesIn(contexts[2]); diff --git a/packages/coding-agent/test/fixtures/instructions-mcp.ts b/packages/coding-agent/test/fixtures/instructions-mcp.ts index 93ae657a2..dc91a58ea 100755 --- a/packages/coding-agent/test/fixtures/instructions-mcp.ts +++ b/packages/coding-agent/test/fixtures/instructions-mcp.ts @@ -1,20 +1,20 @@ #!/usr/bin/env bun /** - * Test fixture: a minimal, well-behaved stdio MCP server that reports - * server-provided `instructions` on `initialize` and exposes a single tool. + * Test fixture: a minimal, well-behaved stdio MCP server that exposes + * deterministic tools. By default it reports server-provided `instructions` + * on `initialize`; its Context Mode fixture mode omits that field entirely. * - * Used by `sdk-mcp-instructions.test.ts` to prove that a deferred interactive - * (`hasUI`) session, whose MCP discovery runs in the background, still folds - * each connected server's instructions into the system prompt once the - * connection completes — see issue: instructions were previously dropped - * permanently for deferred UI sessions. + * Used by `sdk-mcp-instructions.test.ts` to prove that deferred interactive + * (`hasUI`) discovery rebuilds global mounted-route guidance independently of + * optional server instructions, while still folding instructions into the + * prompt when a connected server provides them. * * Speaks newline-delimited JSON-RPC 2.0 (the wire format of `StdioTransport`): * one JSON object per line on stdin, one JSON response per line on stdout. * Only requests (objects with an `id`) get a response; notifications are * dropped. Server-to-client requests are never sent — the client side only * needs `initialize` + `tools/list` answered to register the tool and capture - * the instructions. + * any optional instructions. * * Exported `SERVER_INSTRUCTIONS` is imported by the test for the assertion; * the server only starts when run as the entry module (`import.meta.main`), so @@ -26,9 +26,14 @@ import * as readline from "node:readline"; export const SERVER_INSTRUCTIONS = "INSTR_FIXTURE_SENTINEL_3f9a2c: when this server is connected, always greet in Latin."; -/** Single tool advertised by the fixture so `tools/list` is non-empty. */ -export const TOOL_NAME = "do_thing"; +/** Default advertised tool; bounded and Context Mode fixture modes replace it. */ +export const TOOL_NAME = "do`thing"; export const TOOL_RESULT = "MCP_DEFERRED_SMOKE_OK_5c92"; +export const BOUNDED_GUIDANCE_MODE = "--bounded-guidance"; +export const CONTEXT_MODE_NO_INSTRUCTIONS_MODE = "--context-mode-no-instructions"; +const CONTEXT_MODE_TOOL_NAME = "ctx_execute"; +/** One more tool than the 64-row prompt budget, forcing the static fallback. */ +export const BOUNDED_GUIDANCE_TOOL_COUNT = 65; type JsonRpcRequest = { jsonrpc: "2.0"; @@ -38,6 +43,7 @@ type JsonRpcRequest = { }; function buildResult(method: string): Record { + const contextModeWithoutInstructions = process.argv.includes(CONTEXT_MODE_NO_INSTRUCTIONS_MODE); switch (method) { case "initialize": return { @@ -46,18 +52,29 @@ function buildResult(method: string): Record { // Declare only the tools capability so the client never probes // resources/list or prompts/list — keeps the fixture minimal. capabilities: { tools: {} }, - instructions: SERVER_INSTRUCTIONS, - }; - case "tools/list": - return { - tools: [ - { - name: TOOL_NAME, - description: "Fixture tool returning a deterministic sentinel.", - inputSchema: { type: "object", properties: {}, additionalProperties: false }, - }, - ], + ...(contextModeWithoutInstructions ? {} : { instructions: SERVER_INSTRUCTIONS }), }; + case "tools/list": { + const tools = process.argv.includes(BOUNDED_GUIDANCE_MODE) + ? Array.from({ length: BOUNDED_GUIDANCE_TOOL_COUNT }, (_, index) => { + const suffix = String.fromCharCode(97 + Math.floor(index / 26), 97 + (index % 26)); + return { + name: `row_${suffix}`, + description: `Bounded guidance fixture tool ${suffix}.`, + inputSchema: { type: "object", properties: {}, additionalProperties: false }, + }; + }) + : [ + { + name: contextModeWithoutInstructions ? CONTEXT_MODE_TOOL_NAME : TOOL_NAME, + description: contextModeWithoutInstructions + ? "Execute code through the Context Mode fixture." + : "Fixture tool returning a deterministic sentinel.", + inputSchema: { type: "object", properties: {}, additionalProperties: false }, + }, + ]; + return { tools }; + } case "tools/call": return { content: [{ type: "text", text: TOOL_RESULT }], isError: false }; default: diff --git a/packages/coding-agent/test/sdk-mcp-instructions.test.ts b/packages/coding-agent/test/sdk-mcp-instructions.test.ts index 557d8ea68..065b1e9b1 100644 --- a/packages/coding-agent/test/sdk-mcp-instructions.test.ts +++ b/packages/coding-agent/test/sdk-mcp-instructions.test.ts @@ -9,39 +9,58 @@ import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { createAgentSession } from "@oh-my-pi/pi-coding-agent/sdk"; import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; import { removeSyncWithRetries, Snowflake } from "@oh-my-pi/pi-utils"; -import { SERVER_INSTRUCTIONS, TOOL_NAME, TOOL_RESULT } from "./fixtures/instructions-mcp"; +import { getAgentDir, setAgentDir } from "@oh-my-pi/pi-utils/dirs"; +import { + BOUNDED_GUIDANCE_MODE, + CONTEXT_MODE_NO_INSTRUCTIONS_MODE, + SERVER_INSTRUCTIONS, + TOOL_RESULT, +} from "./fixtures/instructions-mcp"; // Contract: a deferred interactive (`hasUI`) session runs MCP discovery off the -// first-paint path, so an MCP server's `instructions` are not available when the -// prompt is first built. Once the background connection completes — and the -// resulting `refreshMCPTools` rebuilds the system prompt — that server's -// instructions MUST join the prompt for the rest of the session. Regression -// guard: a prior version gated instruction inclusion on `!deferMCPDiscoveryForUI`, -// which dropped server instructions permanently for every UI session. +// first-paint path. Once the background connection completes, the resulting +// `refreshMCPTools` rebuild must add one global bounded route section for every +// mounted MCP tool, whether or not its server returned optional `instructions`. +// Any supplied server instructions join their separately framed section for the +// rest of the session. Regression guards cover both previously dropped deferred +// instructions and the installed Context Mode server's absent instructions. const FIXTURE_PATH = path.join(import.meta.dir, "fixtures", "instructions-mcp.ts"); -const MCP_TOOL_NAME = `mcp__instr_${TOOL_NAME}`; +const MCP_TOOL_NAME = "mcp__instr_do_thing"; +const MCP_MAPPING_FALLBACK = + "Additional mounted MCP tool mappings were omitted to keep this prompt bounded. Inspect `xd://` for the exact current paths."; +const MCP_EXECUTION_GUIDANCE = "Execute each mounted tool by writing JSON arguments to its mounted path:"; +const MCP_ROUTE_SECTION = "## MCP Tool Routes"; +const CONTEXT_MODE_ROUTE = '- "ctx_execute" → `xd://mcp__context_mode_ctx_execute`'; +const CONTEXT_MODE_MCP_TOOL_NAME = "mcp__context_mode_ctx_execute"; describe("createAgentSession MCP server instructions (deferred UI)", () => { let registryDir: string; let tempDir: string; let authStorage: AuthStorage; let modelRegistry: ModelRegistry; - // Discovery resolves user-level MCP config from `os.homedir()`; redirect it - // to an empty dir so the test connects ONLY to the fixture server and never - // spawns the developer's real MCP servers. + // Discovery resolves user-level MCP config through the process-global agent + // directory. Redirect both that path and os.homedir() so the test connects + // only to the fixture and never spawns the developer's real MCP servers. + let originalAgentDir: string; let isolatedHome: string; + let isolatedAgentDir: string; beforeAll(async () => { registryDir = path.join(os.tmpdir(), `pi-sdk-mcp-instr-registry-${Snowflake.next()}`); fs.mkdirSync(registryDir, { recursive: true }); isolatedHome = path.join(os.tmpdir(), `pi-sdk-mcp-instr-home-${Snowflake.next()}`); fs.mkdirSync(isolatedHome, { recursive: true }); + isolatedAgentDir = path.join(isolatedHome, ".omp", "agent"); + fs.mkdirSync(isolatedAgentDir, { recursive: true }); + originalAgentDir = getAgentDir(); + setAgentDir(isolatedAgentDir); authStorage = await AuthStorage.create(path.join(registryDir, "auth.db")); modelRegistry = new ModelRegistry(authStorage); }); afterAll(() => { authStorage.close(); + setAgentDir(originalAgentDir); for (const dir of [registryDir, isolatedHome]) { if (dir && fs.existsSync(dir)) { removeSyncWithRetries(dir); @@ -108,8 +127,122 @@ describe("createAgentSession MCP server instructions (deferred UI)", () => { } expect(prompt).toContain(SERVER_INSTRUCTIONS); - // The instructions are framed under the MCP section, not pasted raw. + // The instructions are framed under the MCP section, and guidance keeps + // the escaped original tool name while routing through the exact + // normalized name actually mounted in the live xd:// registry. expect(prompt).toContain("MCP Server Instructions"); + expect(prompt).toContain('- "do\\u0060thing" → `xd://mcp__instr_do_thing`'); + expect(prompt).toContain(MCP_EXECUTION_GUIDANCE); + } finally { + await session.dispose(); + } + }, 20_000); + + it("renders a mounted Context Mode route when initialize omits instructions", async () => { + fs.writeFileSync( + path.join(tempDir, ".mcp.json"), + JSON.stringify({ + mcpServers: { + "context-mode": { + type: "stdio", + command: process.execPath, + args: [FIXTURE_PATH, CONTEXT_MODE_NO_INSTRUCTIONS_MODE], + }, + }, + }), + ); + const { session } = await createAgentSession({ + cwd: tempDir, + agentDir: tempDir, + modelRegistry, + sessionManager: SessionManager.inMemory(), + settings: Settings.isolated({}), + model: getBundledModel("openai", "gpt-4o-mini"), + disableExtensionDiscovery: true, + skills: [], + contextFiles: [], + promptTemplates: [], + slashCommands: [], + enableLsp: false, + skipPythonPreflight: true, + enableMCP: true, + hasUI: true, + }); + try { + // Context Mode advertises mounted MCP tools but currently supplies no + // `connection.instructions`. Deferred discovery must still rebuild the + // prompt with the globally rendered route guidance. The SDK exposes no + // completion signal for this real child-process handshake, and fake + // timers cannot drive it, so poll only until the route becomes visible. + let prompt = session.systemPrompt.join("\n"); + expect(prompt).not.toContain(CONTEXT_MODE_ROUTE); + const deadline = Date.now() + 12_000; + while (!prompt.includes(CONTEXT_MODE_ROUTE) && Date.now() < deadline) { + await Bun.sleep(50); + prompt = session.systemPrompt.join("\n"); + } + + expect(prompt).toContain(CONTEXT_MODE_ROUTE); + expect(session.getXdevToolEntries().map(entry => entry.name)).toContain(CONTEXT_MODE_MCP_TOOL_NAME); + expect(session.getActiveToolNames()).not.toContain(CONTEXT_MODE_MCP_TOOL_NAME); + expect(prompt.split(MCP_EXECUTION_GUIDANCE)).toHaveLength(2); + expect(prompt.split(MCP_ROUTE_SECTION)).toHaveLength(2); + expect(prompt).not.toContain(SERVER_INSTRUCTIONS); + expect(prompt).not.toContain("## MCP Server Instructions"); + expect(prompt).not.toContain("### context-mode"); + } finally { + await session.dispose(); + } + }, 20_000); + + it("bounds mounted route guidance deterministically and points to the live xd:// inventory", async () => { + fs.writeFileSync( + path.join(tempDir, ".mcp.json"), + JSON.stringify({ + mcpServers: { + instr: { + type: "stdio", + command: process.execPath, + args: [FIXTURE_PATH, BOUNDED_GUIDANCE_MODE], + }, + }, + }), + ); + const { session } = await createAgentSession({ + cwd: tempDir, + agentDir: tempDir, + modelRegistry, + sessionManager: SessionManager.inMemory(), + settings: Settings.isolated({}), + model: getBundledModel("openai", "gpt-4o-mini"), + disableExtensionDiscovery: true, + skills: [], + contextFiles: [], + promptTemplates: [], + slashCommands: [], + enableLsp: false, + skipPythonPreflight: true, + enableMCP: true, + hasUI: true, + }); + try { + // Deferred discovery is a real child-process handshake with no + // completion signal exposed to this integration harness; fake timers + // cannot advance it, so retain the established polling bounds above. + const deadline = Date.now() + 12_000; + let prompt = session.systemPrompt.join("\n"); + while (!prompt.includes(SERVER_INSTRUCTIONS) && Date.now() < deadline) { + await Bun.sleep(50); + prompt = session.systemPrompt.join("\n"); + } + + expect(prompt).toContain(SERVER_INSTRUCTIONS); + const renderedMappings = prompt.split("\n").filter(line => line.startsWith('- "row_')); + expect(renderedMappings).toHaveLength(64); + expect(renderedMappings[0]).toBe('- "row_aa" → `xd://mcp__instr_row_aa`'); + expect(renderedMappings[63]).toBe('- "row_cl" → `xd://mcp__instr_row_cl`'); + expect(prompt).not.toContain('- "row_cm" → `xd://mcp__instr_row_cm`'); + expect(prompt).toContain(MCP_MAPPING_FALLBACK); } finally { await session.dispose(); } @@ -193,6 +326,9 @@ describe("createAgentSession MCP server instructions (deferred UI)", () => { expect(activeNames).toContain(MCP_TOOL_NAME); expect(session.getXdevToolEntries().map(entry => entry.name)).not.toContain(MCP_TOOL_NAME); + expect(prompt).toContain("## MCP Server Instructions"); + expect(prompt).toContain(SERVER_INSTRUCTIONS); + expect(prompt).not.toContain(`xd://${MCP_TOOL_NAME}`); } finally { await session.dispose(); }