diff --git a/src/cli/commands/chat.tsx b/src/cli/commands/chat.tsx index 3d122ed..36bfefe 100644 --- a/src/cli/commands/chat.tsx +++ b/src/cli/commands/chat.tsx @@ -284,6 +284,7 @@ export async function chatCommand(opts: ChatOptions): Promise { toolCount: bridge.registeredNames.length, report, host, + bridgeEnv: bridge.env, }); } catch (err) { // Per-server failure is non-fatal: one broken server shouldn't diff --git a/src/cli/ui/App.tsx b/src/cli/ui/App.tsx index a4b3d92..d8d0385 100644 --- a/src/cli/ui/App.tsx +++ b/src/cli/ui/App.tsx @@ -116,6 +116,7 @@ import { import { StatusRow } from "./layout/StatusRow.js"; import { ToastRail } from "./layout/ToastRail.js"; import { formatLoopStatus } from "./loop.js"; +import { applyMcpAppend } from "./mcp-append.js"; import { handleMcpBrowseSlash } from "./mcp-browse.js"; import { formatLongPaste } from "./paste-collapse.js"; import { resolvePreset } from "./presets.js"; @@ -3387,6 +3388,7 @@ function AppInner({ configPath={defaultConfigPath()} onClose={() => setPendingMcpBrowser(false)} postInfo={(text) => log.pushInfo(text)} + applyAppend={(target, addedTools) => applyMcpAppend(loop, target, addedTools)} /> ) : pendingPlan ? ( void; /** Pushed by the modal when a key triggers async work (`r` reconnect). */ postInfo: (text: string) => void; + /** Optional — opt-in to append-drift acceptance on `r`. Without it, append-drift refuses. */ + applyAppend?: ApplyAppend; } -export function McpBrowser({ servers, configPath, onClose, postInfo }: McpBrowserProps) { +export function McpBrowser({ + servers, + configPath, + onClose, + postInfo, + applyAppend, +}: McpBrowserProps) { const [index, setIndex] = useState(0); const max = Math.max(0, servers.length - 1); @@ -31,7 +39,7 @@ export function McpBrowser({ servers, configPath, onClose, postInfo }: McpBrowse // Hand the "starting" lifecycle line to scrollback and let the // kickoff schedule the result line via postInfo. Close the modal // so the line is visible immediately. - postInfo(kickOffMcpReconnect(target, postInfo)); + postInfo(kickOffMcpReconnect(target, postInfo, applyAppend)); onClose(); } }); diff --git a/src/cli/ui/mcp-append.ts b/src/cli/ui/mcp-append.ts new file mode 100644 index 0000000..5475090 --- /dev/null +++ b/src/cli/ui/mcp-append.ts @@ -0,0 +1,40 @@ +/** Applies an MCP append-drift mid-session: registers each new tool in the loop's registry + prefix, and updates the summary's report. */ + +import type { CacheFirstLoop } from "../../loop.js"; +import { registerSingleMcpTool } from "../../mcp/registry.js"; +import type { McpTool } from "../../mcp/types.js"; +import type { JSONSchema, ToolSpec } from "../../types.js"; +import type { McpServerSummary } from "./slash/types.js"; + +export function applyMcpAppend( + loop: CacheFirstLoop, + target: McpServerSummary, + addedTools: McpTool[], +): void { + const accepted: McpTool[] = []; + for (const mcpTool of addedTools) { + if (!mcpTool.name) continue; + const registeredName = registerSingleMcpTool(mcpTool, target.bridgeEnv); + if (!registeredName) continue; + const spec: ToolSpec = { + type: "function", + function: { + name: registeredName, + description: mcpTool.description ?? "", + parameters: mcpTool.inputSchema as unknown as JSONSchema, + }, + }; + loop.prefix.addTool(spec); + accepted.push(mcpTool); + } + if (accepted.length === 0) return; + // Refresh the summary's snapshot so `/mcp` and the browser modal show the + // new shape on their next render. + if (target.report.tools.supported) { + const merged = [...target.report.tools.items, ...accepted]; + // biome-ignore lint/suspicious/noExplicitAny: report is a typed snapshot we mutate in place; deeper refactor isn't worth it here + (target.report.tools as any).items = merged; + // biome-ignore lint/suspicious/noExplicitAny: same — toolCount mirrors items.length post-append + (target as any).toolCount = merged.length; + } +} diff --git a/src/cli/ui/mcp-reconnect-kickoff.ts b/src/cli/ui/mcp-reconnect-kickoff.ts index e1c318f..1b7db35 100644 --- a/src/cli/ui/mcp-reconnect-kickoff.ts +++ b/src/cli/ui/mcp-reconnect-kickoff.ts @@ -1,31 +1,47 @@ /** Shared async-fire-and-forget reconnect trigger — called by both `/mcp reconnect` and the McpBrowser `r` keybind. */ import { reconnectMcpServer } from "../../mcp/reconnect.js"; +import type { McpTool } from "../../mcp/types.js"; import { formatMcpLifecycleEvent } from "./mcp-lifecycle.js"; import type { McpServerSummary } from "./slash/types.js"; +/** Applies append-drift mid-session: registers each new MCP tool in the registry + prefix. */ +export type ApplyAppend = (target: McpServerSummary, addedTools: McpTool[]) => void; + /** Kicks off async reconnect; returns the start-line, schedules result via postInfo. */ export function kickOffMcpReconnect( target: McpServerSummary, postInfo: (text: string) => void, + applyAppend?: ApplyAppend, ): string { const beforeTools = target.report.tools.supported ? target.report.tools.items : []; + // Only opt into "append" when the caller wired an applyAppend handler; + // otherwise the reconnect refuses append-drift with a "restart" message. + const accept = applyAppend ? (["identity", "append"] as const) : (["identity"] as const); void (async () => { try { const result = await reconnectMcpServer({ host: target.host, spec: target.spec, beforeTools, + accept, }); if (result.ok) { + if (result.kind === "append" && applyAppend) { + applyAppend(target, result.addedTools); + } postInfo( formatMcpLifecycleEvent({ state: "connected", name: target.label, - tools: beforeTools.length, + tools: result.afterTools.length, ms: result.ms, }), ); + if (result.kind === "append") { + const names = result.addedTools.map((t) => t.name).join(", "); + postInfo(`▸ ${target.label}: added ${result.addedTools.length} tool(s) — ${names}`); + } } else { postInfo( formatMcpLifecycleEvent({ diff --git a/src/cli/ui/slash/handlers/mcp.ts b/src/cli/ui/slash/handlers/mcp.ts index 8fcfcb6..e7bf1d0 100644 --- a/src/cli/ui/slash/handlers/mcp.ts +++ b/src/cli/ui/slash/handlers/mcp.ts @@ -1,4 +1,6 @@ import { readConfig, writeConfig } from "../../../../config.js"; +import type { CacheFirstLoop } from "../../../../loop.js"; +import { applyMcpAppend } from "../../mcp-append.js"; import { kickOffMcpReconnect } from "../../mcp-reconnect-kickoff.js"; import type { SlashHandler } from "../dispatch.js"; import { appendSection } from "../helpers.js"; @@ -13,7 +15,7 @@ const mcp: SlashHandler = (args, loop, ctx) => { return toggleDisabled(sub, args[1], { servers, specs }); } if (sub === "reconnect") { - return triggerReconnect(args[1], servers, ctx.postInfo); + return triggerReconnect(args[1], servers, ctx.postInfo, loop); } // `/mcp text` (or non-TTY) falls through to the printed-card path. The // default `/mcp` opens the interactive browser modal. @@ -132,6 +134,7 @@ function triggerReconnect( rawName: string | undefined, servers: ReadonlyArray, postInfo: ((text: string) => void) | undefined, + loop: CacheFirstLoop, ): { info: string } { const name = rawName?.trim(); if (!name) { @@ -150,7 +153,14 @@ function triggerReconnect( if (!postInfo) { return { info: "/mcp reconnect requires the interactive TUI (postInfo not wired)." }; } - return { info: kickOffMcpReconnect(target, postInfo) }; + // Append-drift accepted automatically: server added new tools, we register them + // and call addTool on the prefix (cache miss only on the appended chunks per the + // benchmarks/spike-mcp-reconnect data — typically <5% loss). + return { + info: kickOffMcpReconnect(target, postInfo, (t, addedTools) => + applyMcpAppend(loop, t, addedTools), + ), + }; } export const handlers: Record = { mcp }; diff --git a/src/cli/ui/slash/types.ts b/src/cli/ui/slash/types.ts index 2528e4e..fd807a1 100644 --- a/src/cli/ui/slash/types.ts +++ b/src/cli/ui/slash/types.ts @@ -1,6 +1,6 @@ import type { EditMode } from "../../../config.js"; import type { InspectionReport } from "../../../mcp/inspect.js"; -import type { McpClientHost } from "../../../mcp/registry.js"; +import type { BridgeEnv, McpClientHost } from "../../../mcp/registry.js"; import type { JobRegistry } from "../../../tools/jobs.js"; import type { PlanStep } from "../../../tools/plan.js"; @@ -122,6 +122,8 @@ export interface McpServerSummary { report: InspectionReport; /** Mutable client handle so `/mcp reconnect` can swap the underlying socket without re-bridging tools. */ host: McpClientHost; + /** Captured at first-bridge time so append-drift reconnects can register newly-added tools with the same options. */ + bridgeEnv: BridgeEnv; } export interface SlashCommandSpec { diff --git a/src/mcp/reconnect.ts b/src/mcp/reconnect.ts index 5d4f47a..3b1d250 100644 --- a/src/mcp/reconnect.ts +++ b/src/mcp/reconnect.ts @@ -1,4 +1,4 @@ -/** `/mcp reconnect` — open a fresh client, accept identity drift only, refuse the rest cleanly. */ +/** `/mcp reconnect` — open a fresh client, accept identity (always) and append (opt-in), refuse the rest cleanly. */ import { McpClient } from "./client.js"; import { classifyToolListDrift } from "./drift.js"; @@ -16,10 +16,19 @@ export interface ReconnectArgs { spec: string; /** The current tool list, used as the drift baseline. */ beforeTools: readonly McpTool[]; + /** Drift kinds the caller is willing to accept. Default: ["identity"]. */ + accept?: ReadonlyArray<"identity" | "append">; } export type ReconnectResult = - | { ok: true; afterTools: McpTool[]; ms: number } + | { + ok: true; + kind: "identity" | "append"; + afterTools: McpTool[]; + /** Tools present in `afterTools` but not in `beforeTools` (empty for identity). */ + addedTools: McpTool[]; + ms: number; + } | { ok: false; reason: @@ -35,6 +44,7 @@ export type ReconnectResult = export async function reconnectMcpServer(args: ReconnectArgs): Promise { const t0 = Date.now(); + const accept = args.accept ?? ["identity"]; let parsed: McpSpec; try { parsed = parseMcpSpec(args.spec); @@ -57,24 +67,37 @@ export async function reconnectMcpServer(args: ReconnectArgs): Promise {}); + const refused = drift.kind as Exclude; return { ok: false, - reason: driftReason(drift.kind), + reason: driftReason(refused), message: driftMessage(drift), ms: Date.now() - t0, }; } - // Identity drift — safe to swap. + const addedTools = + acceptedKind === "append" ? listed.tools.filter((t) => drift.added.includes(t.name)) : []; + // Swap. const old = args.host.client; args.host.client = next; await old.close().catch(() => {}); - return { ok: true, afterTools: listed.tools, ms: Date.now() - t0 }; + return { + ok: true, + kind: acceptedKind, + afterTools: listed.tools, + addedTools, + ms: Date.now() - t0, + }; } catch (err) { await next.close().catch(() => {}); return { diff --git a/src/mcp/registry.ts b/src/mcp/registry.ts index 8455362..b82b932 100644 --- a/src/mcp/registry.ts +++ b/src/mcp/registry.ts @@ -49,10 +49,49 @@ export interface BridgeResult { skipped: Array<{ name: string; reason: string }>; } +/** Resolved bridge environment that `registerSingleMcpTool` needs. Stored on summaries so reconnect can append new tools later. */ +export interface BridgeEnv { + registry: ToolRegistry; + host: McpClientHost; + prefix: string; + maxResultChars: number; + tracker: LatencyTracker | null; + onProgress?: BridgeOptions["onProgress"]; +} + +/** Register one MCP tool's bridged closure into the registry. Returns the registered name (or "" if skipped). */ +export function registerSingleMcpTool( + mcpTool: import("./types.js").McpTool, + env: BridgeEnv, +): string { + if (!mcpTool.name) return ""; + const registeredName = `${env.prefix}${mcpTool.name}`; + env.registry.register({ + name: registeredName, + description: mcpTool.description ?? "", + parameters: mcpTool.inputSchema as JSONSchema, + fn: async (args: Record, ctx) => { + const t0 = env.tracker ? Date.now() : 0; + // Resolve client at call time via the host indirection so `/mcp reconnect` + // can swap a fresh client in without re-bridging tools. + const live = env.host.client; + const toolResult = await live.callTool(mcpTool.name, args, { + onProgress: env.onProgress + ? (info) => env.onProgress!({ toolName: registeredName, ...info }) + : undefined, + signal: ctx?.signal, + }); + if (env.tracker) env.tracker.record(Date.now() - t0); + return flattenMcpResult(toolResult, { maxChars: env.maxResultChars }); + }, + }); + return registeredName; +} + export async function bridgeMcpTools( client: McpClient, opts: BridgeOptions = {}, -): Promise { +): Promise { const registry = opts.registry ?? new ToolRegistry({ autoFlatten: opts.autoFlatten }); const prefix = opts.namePrefix ?? ""; const maxResultChars = opts.maxResultChars ?? DEFAULT_MAX_RESULT_CHARS; @@ -62,35 +101,29 @@ export async function bridgeMcpTools( const tracker = opts.onSlow ? new LatencyTracker(serverName, { thresholdMs: opts.slowThresholdMs, onSlow: opts.onSlow }) : null; + // Synthesize a host on the fly when the caller didn't provide one. Older + // callers (tests, single-shot non-reconnectable bridges) get the live + // `client` reference frozen in; reconnect-aware callers pass their own + // mutable host. + const host: McpClientHost = opts.host ?? { client }; + const env: BridgeEnv = { + registry, + host, + prefix, + maxResultChars, + tracker, + onProgress: opts.onProgress, + }; const listed = await client.listTools(); for (const mcpTool of listed.tools) { if (!mcpTool.name) { result.skipped.push({ name: "?", reason: "empty tool name" }); continue; } - const registeredName = `${prefix}${mcpTool.name}`; - registry.register({ - name: registeredName, - description: mcpTool.description ?? "", - parameters: mcpTool.inputSchema as JSONSchema, - fn: async (args: Record, ctx) => { - const t0 = tracker ? Date.now() : 0; - // Resolve client at call time via the host indirection (when given) so - // `/mcp reconnect` can swap a fresh client in without re-bridging tools. - const live = opts.host?.client ?? client; - const toolResult = await live.callTool(mcpTool.name, args, { - onProgress: opts.onProgress - ? (info) => opts.onProgress!({ toolName: registeredName, ...info }) - : undefined, - signal: ctx?.signal, - }); - if (tracker) tracker.record(Date.now() - t0); - return flattenMcpResult(toolResult, { maxChars: maxResultChars }); - }, - }); - result.registeredNames.push(registeredName); + const registeredName = registerSingleMcpTool(mcpTool, env); + if (registeredName) result.registeredNames.push(registeredName); } - return result; + return { ...result, env }; } export interface FlattenOptions { diff --git a/tests/mcp-append.test.ts b/tests/mcp-append.test.ts new file mode 100644 index 0000000..d20e191 --- /dev/null +++ b/tests/mcp-append.test.ts @@ -0,0 +1,117 @@ +import { describe, expect, it, vi } from "vitest"; +import { applyMcpAppend } from "../src/cli/ui/mcp-append.js"; +import type { McpServerSummary } from "../src/cli/ui/slash/types.js"; +import { CacheFirstLoop, DeepSeekClient, ImmutablePrefix } from "../src/index.js"; +import { McpClient } from "../src/mcp/client.js"; +import type { BridgeEnv, McpClientHost } from "../src/mcp/registry.js"; +import { StdioTransport } from "../src/mcp/stdio.js"; +import type { McpTool } from "../src/mcp/types.js"; +import { ToolRegistry } from "../src/tools.js"; + +function makeLoop() { + const tools = new ToolRegistry(); + const prefix = new ImmutablePrefix({ system: "s" }); + const client = new DeepSeekClient({ + apiKey: "sk-test", + fetch: vi.fn() as unknown as typeof fetch, + }); + return { loop: new CacheFirstLoop({ client, prefix, tools }), tools, prefix }; +} + +function makeFakeMcp(): { host: McpClientHost; env: BridgeEnv; registry: ToolRegistry } { + // The host's client is a real McpClient pointing at a never-spawned transport; + // applyMcpAppend doesn't actually call the tool, so this is fine. + const transport = new StdioTransport({ command: "true", args: [], shell: false }); + const host: McpClientHost = { client: new McpClient({ transport, requestTimeoutMs: 1_000 }) }; + const registry = new ToolRegistry(); + const env: BridgeEnv = { + registry, + host, + prefix: "fs_", + maxResultChars: 32_000, + tracker: null, + }; + return { host, env, registry }; +} + +function summary(env: BridgeEnv, host: McpClientHost): McpServerSummary { + return { + label: "fs", + spec: "fs=cmd", + toolCount: 0, + host, + bridgeEnv: env, + report: { + protocolVersion: "2024-11-05", + serverInfo: { name: "fs", version: "1.0" }, + capabilities: {}, + tools: { supported: true, items: [] }, + resources: { supported: false, reason: "" }, + prompts: { supported: false, reason: "" }, + elapsedMs: 50, + }, + }; +} + +const newTool: McpTool = { + name: "delete_file", + description: "Remove a file at the given path.", + inputSchema: { + type: "object", + properties: { path: { type: "string" } }, + required: ["path"], + }, +}; + +describe("applyMcpAppend", () => { + it("registers the new tool in the registry under the prefix", async () => { + const { loop, prefix } = makeLoop(); + const { env, host, registry } = makeFakeMcp(); + // Re-bind the bridgeEnv's registry to the loop's so the mutation lands there. + env.registry = loop.tools; + const target = summary(env, host); + + expect(loop.tools.has("fs_delete_file")).toBe(false); + applyMcpAppend(loop, target, [newTool]); + expect(loop.tools.has("fs_delete_file")).toBe(true); + // Sanity: the unused `registry` shows we're not mutating the wrong place. + expect(registry.has("fs_delete_file")).toBe(false); + // Prefix gained the spec, with the prefixed name. + const names = prefix.toolSpecs.map((t) => t.function.name); + expect(names).toContain("fs_delete_file"); + }); + + it("invalidates the prefix fingerprint (cache key changes for next turn)", () => { + const { loop, prefix } = makeLoop(); + const { env, host } = makeFakeMcp(); + env.registry = loop.tools; + const target = summary(env, host); + + const before = prefix.fingerprint; + applyMcpAppend(loop, target, [newTool]); + expect(prefix.fingerprint).not.toBe(before); + }); + + it("updates the summary's tool count + report.tools.items", () => { + const { loop } = makeLoop(); + const { env, host } = makeFakeMcp(); + env.registry = loop.tools; + const target = summary(env, host); + + applyMcpAppend(loop, target, [newTool]); + expect(target.toolCount).toBe(1); + if (!target.report.tools.supported) throw new Error("unreachable"); + expect(target.report.tools.items.map((t) => t.name)).toContain("delete_file"); + }); + + it("skips MCP tools without a name (defensive)", () => { + const { loop } = makeLoop(); + const { env, host } = makeFakeMcp(); + env.registry = loop.tools; + const target = summary(env, host); + + applyMcpAppend(loop, target, [{ name: "", inputSchema: { type: "object" } } as McpTool]); + expect(loop.tools.size).toBe(0); + expect(target.toolCount).toBe(0); + }); +});