From c662ec7e374563bd549dc35f47bac52dcc4bed88 Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Thu, 1 Oct 2026 13:34:54 +0200 Subject: [PATCH] fix(coding-agent): restore MCP tools loaded by tool_search on resume and reload The session restored its tool loadout before MCP servers reconnected, so deferred tools that tool_search had loaded were dropped. Restored tools that are not registered yet now stay pending and are activated when they register, until a setActiveTools() call deactivates a tool or the next prompt starts. --- packages/coding-agent/CHANGELOG.md | 1 + .../coding-agent/src/core/agent-session.ts | 45 ++++-- .../test/suite/agent-session-mcp.test.ts | 132 +++++++++++++++++- packages/coding-agent/test/suite/harness.ts | 4 +- 4 files changed, 171 insertions(+), 11 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 12fbe7252..5700e14b4 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -18,6 +18,7 @@ - Fixed `--provider` without `--model` being silently ignored and running the default model from another provider; it now fails with an error ([#10236](https://github.com/earendil-works/pi/issues/10236)). - Fixed MCP servers that ask for more scope (`insufficient_scope`) requesting sign-in over and over. The new sign-in requested only the missing scopes, so the new token lost access the previous one had; it now keeps the granted scopes. - Fixed user messages in the transcript keeping two full-width copies of every rendered line; they keep one, with identical output. +- Fixed deferred MCP tools that `tool_search` loaded being dropped on resume and `/reload` even when their server reconnected before the next prompt, because the session restored its tools before the MCP servers reconnected. ## [0.99.2] - 2026-09-30 diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index c300e123c..63aa588da 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -423,6 +423,12 @@ export class AgentSession { private _cwd: string; private _extensionRunnerRef?: { current?: ExtensionRunner }; private _initialActiveToolNames?: string[]; + /** + * Tools of the restored or reloaded loadout that are not registered yet, such as tools of MCP + * servers that are still connecting. They are activated when they are registered, and dropped when + * `setActiveToolsByName()` deactivates a tool or the next agent run starts. + */ + private _pendingToolNames = new Set(); private _usesDefaultTools: boolean; private _allowedToolNames?: Set; private _excludedToolNames?: Set; @@ -1480,10 +1486,24 @@ export class AgentSession { * Changes take effect on the next agent turn. */ setActiveToolsByName(toolNames: string[]): void { + const previous = this.getActiveToolNames(); + this._setActiveTools(toolNames); + // A loadout that deactivates a tool replaces the restored one, whose pending tools are dropped. + // One that only adds tools, like activating tool_search, keeps them. + const active = new Set(this.getActiveToolNames()); + if (previous.some((name) => !active.has(name))) this._pendingToolNames.clear(); + } + + private _setActiveTools(toolNames: string[]): void { const tools = this._applyToolLoadout(toolNames); + for (const tool of tools) this._pendingToolNames.delete(tool.name); this._rebuildSystemPrompt(tools.map((tool) => tool.name)); } + private _isAllowedTool(name: string): boolean { + return (!this._allowedToolNames || this._allowedToolNames.has(name)) && !this._excludedToolNames?.has(name); + } + private _getToolExposure(name: string): ToolExposure { return this._toolDefinitions.get(name)?.definition.exposure ?? "direct"; } @@ -1740,9 +1760,12 @@ export class AgentSession { * set, so the transcript's declarations are the whole loadout. */ private _restoreToolsFromTranscript(): void { + this._pendingToolNames.clear(); const current = getCurrentSystemMessage(this.sessionManager.buildSessionContext().messages); if (!current) return; - this.setActiveToolsByName((current.toolsAdded ?? []).map((tool) => tool.name)); + const names = (current.toolsAdded ?? []).map((tool) => tool.name); + this._pendingToolNames = new Set(names.filter((name) => this._isAllowedTool(name))); + this._setActiveTools(names); } // ========================================================================= @@ -1754,6 +1777,9 @@ export class AgentSession { // Compaction before the prompt may have scheduled a retry; the new prompt replaces it. this._failedResponse = undefined; this._recordSelection(); + // The run records the loadout in the transcript; restored tools that did not register by now + // are dropped, so a tool that never registers does not stay pending. + this._pendingToolNames.clear(); this._isAgentRunActive = true; try { await this.agent.prompt(messages); @@ -3427,9 +3453,6 @@ export class AgentSession { ); const previousActiveToolNames = this.getActiveToolNames(); const allowedToolNames = this._allowedToolNames; - const excludedToolNames = this._excludedToolNames; - const isAllowedTool = (name: string): boolean => - (!allowedToolNames || allowedToolNames.has(name)) && !excludedToolNames?.has(name); const registeredTools = this._extensionRunner.getAllRegisteredTools(); const allCustomTools = [ @@ -3438,10 +3461,10 @@ export class AgentSession { definition, sourceInfo: createSyntheticSourceInfo(``, { source: "sdk" }), })), - ].filter((tool) => isAllowedTool(tool.definition.name)); + ].filter((tool) => this._isAllowedTool(tool.definition.name)); const definitionRegistry = new Map( Array.from(this._baseToolDefinitions.entries()) - .filter(([name]) => isAllowedTool(name)) + .filter(([name]) => this._isAllowedTool(name)) .map(([name, definition]) => [ name, { @@ -3477,7 +3500,7 @@ export class AgentSession { const wrappedExtensionTools = wrapRegisteredTools(allCustomTools, runner); const wrappedBuiltInTools = wrapRegisteredTools( Array.from(this._baseToolDefinitions.values()) - .filter((definition) => isAllowedTool(definition.name)) + .filter((definition) => this._isAllowedTool(definition.name)) .map((definition) => ({ definition, sourceInfo: createSyntheticSourceInfo(`${BUILTIN_PATH_PREFIX}${definition.name}`, { @@ -3495,7 +3518,7 @@ export class AgentSession { const nextActiveToolNames = ( options?.activeToolNames ? [...options.activeToolNames] : [...previousActiveToolNames] - ).filter((name) => isAllowedTool(name)); + ).filter((name) => this._isAllowedTool(name)); if (allowedToolNames) { for (const toolName of this._toolRegistry.keys()) { @@ -3515,8 +3538,10 @@ export class AgentSession { } } } + // Pending tools that are registered now become active. + nextActiveToolNames.push(...this._pendingToolNames); - this.setActiveToolsByName([...new Set(nextActiveToolNames)]); + this._setActiveTools([...new Set(nextActiveToolNames)]); } /** Whether activating the tool declares it to the model. */ @@ -3603,6 +3628,8 @@ export class AgentSession { (name) => !previousDefaultTools.has(name), ) : []; + // Tools the new extensions register later, such as MCP tools, are pending until then. + for (const name of this.getActiveToolNames()) this._pendingToolNames.add(name); this._buildRuntime({ activeToolNames: [...this.getActiveToolNames(), ...addedDefaultTools], flagValues: previousFlagValues, diff --git a/packages/coding-agent/test/suite/agent-session-mcp.test.ts b/packages/coding-agent/test/suite/agent-session-mcp.test.ts index 0e2bb894f..30761608d 100644 --- a/packages/coding-agent/test/suite/agent-session-mcp.test.ts +++ b/packages/coding-agent/test/suite/agent-session-mcp.test.ts @@ -5,12 +5,14 @@ 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 type { SessionManager } from "../../src/core/session-manager.ts"; import { createCodemodeExtension } from "../../src/extensions/codemode/index.ts"; import type { McpExposure, McpServerEntry } from "../../src/extensions/mcp/config.ts"; import { createMcpExtension, MCP_SERVERS_SECTION } from "../../src/extensions/mcp/index.ts"; import { createMcpToolName } from "../../src/extensions/mcp/tools.ts"; import { createToolSearchExtension } from "../../src/extensions/tool-search/index.ts"; import { TOOL_SEARCH_DESCRIPTION } from "../../src/extensions/tool-search/tool.ts"; +import { createTestExtensionsResult, createTestResourceLoader } from "../utilities.ts"; import { createHarness, createTestUiContext, @@ -899,7 +901,8 @@ describe("AgentSession MCP servers registered by extensions", () => { ], }); harnesses.push(harness); - await harness.session.bindExtensions({}); + // `/reload` emits session_start only to bound extensions. + await harness.session.bindExtensions({ uiContext: createTestUiContext() }); return { harness, connected }; } @@ -987,3 +990,130 @@ describe("AgentSession MCP servers registered by extensions", () => { expect(errors).toEqual([expect.stringContaining('MCP server "orphan" is registered, but no loaded extension')]); }); }); + +describe("AgentSession MCP tools after resume and reload", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) harnesses.pop()?.cleanup(); + }); + + /** + * A deferred `docs` server that answers `initialize` after `initializeDelayMs`; `connected` counts + * its connections. `/reload` loads the extensions again. + */ + async function setup( + sessionManager?: SessionManager, + extensionFactories: ExtensionFactory[] = [], + initializeDelayMs = 0, + ) { + const connected: string[] = []; + const servers: McpServerEntry[] = [ + { name: "docs", config: { url: "http://unused.invalid", exposure: "deferred" }, source: "test" }, + ]; + const factories = [ + ...extensionFactories, + createToolSearchExtension(), + createMcpExtension({ + loadConfig: () => ({ servers, errors: [] }), + createTransport: (entry) => { + connected.push(entry.name); + const pair = createFakeServer([], { initializeDelayMs }); + void pair.server.start(); + return pair.client; + }, + }), + ]; + let extensions = await createTestExtensionsResult(factories); + const resourceLoader = { + ...createTestResourceLoader(), + getExtensions: () => extensions, + reload: async () => { + extensions = await createTestExtensionsResult(factories); + }, + }; + const harness = await createHarness({ resourceLoader, sessionManager }); + harnesses.push(harness); + // `/reload` emits session_start only to bound extensions. + await harness.session.bindExtensions({ uiContext: createTestUiContext() }); + return { harness, connected }; + } + + async function loadDocsSearch(harness: Harness) { + harness.setResponses([ + fauxAssistantMessage([fauxToolCall("tool_search", { query: "search the docs", limit: 1 })], { + stopReason: "toolUse", + }), + fauxAssistantMessage("loaded"), + ]); + await harness.session.prompt("load"); + expect(harness.session.getActiveToolNames()).toContain("mcp__docs__search"); + } + + it("declares tools tool_search loaded again on resume once their server connects", async () => { + const first = await setup(); + await loadDocsSearch(first.harness); + + // The session restores its tools before the server connects again. + const second = await setup(first.harness.sessionManager); + await vi.waitFor(() => expect(second.harness.session.getActiveToolNames()).toContain("mcp__docs__search")); + second.harness.setResponses([ + fauxAssistantMessage([fauxToolCall("mcp__docs__search", { query: "again" })], { stopReason: "toolUse" }), + fauxAssistantMessage("done"), + ]); + await second.harness.session.prompt("use it"); + + expect(getMessageText(toolResult(second.harness, "mcp__docs__search"))).toBe("again guide\nagain faq"); + const removals = second.harness.session.messages.filter( + (message) => message.role === "system" && (message.toolsRemoved ?? []).length > 0, + ); + expect(removals).toEqual([]); + }); + + it.each([ + ["drops", ["read"], false], + ["keeps", undefined, true], + ] as const)( + "%s restored tools when an extension sets the loadout before they register", + async (_, loadout, kept) => { + const first = await setup(); + await loadDocsSearch(first.harness); + + // Like plan mode restoring its tools, or an extension adding one to the current loadout. + const setLoadout: ExtensionFactory = (pi) => { + pi.on("session_start", () => pi.setActiveTools(loadout ? [...loadout] : [...pi.getActiveTools(), "read"])); + }; + const second = await setup(first.harness.sessionManager, [setLoadout]); + await vi.waitFor(() => + expect(second.harness.session.getAllTools().some((tool) => tool.name === "mcp__docs__search")).toBe(true), + ); + + expect(second.harness.session.getActiveToolNames().includes("mcp__docs__search")).toBe(kept); + }, + ); + + it("does not activate restored tools that register after the next prompt starts", async () => { + const first = await setup(); + await loadDocsSearch(first.harness); + + // The first prompt does not wait for servers without direct tools. + const second = await setup(first.harness.sessionManager, [], 200); + second.harness.setResponses([fauxAssistantMessage("done")]); + await second.harness.session.prompt("go"); + await vi.waitFor(() => + expect(second.harness.session.getAllTools().some((tool) => tool.name === "mcp__docs__search")).toBe(true), + ); + + expect(second.harness.session.getActiveToolNames()).not.toContain("mcp__docs__search"); + }); + + it("declares tools tool_search loaded again after /reload", async () => { + const { harness, connected } = await setup(); + await loadDocsSearch(harness); + + await harness.session.reload(); + + await vi.waitFor(() => expect(connected).toEqual(["docs", "docs"])); + await vi.waitFor(() => expect(harness.session.getActiveToolNames()).toContain("mcp__docs__search")); + }); +}); diff --git a/packages/coding-agent/test/suite/harness.ts b/packages/coding-agent/test/suite/harness.ts index 27bea22af..220a1558f 100644 --- a/packages/coding-agent/test/suite/harness.ts +++ b/packages/coding-agent/test/suite/harness.ts @@ -119,6 +119,8 @@ export interface HarnessOptions { extensionFactories?: Array; withConfiguredAuth?: boolean; modelsJson?: Record; + /** Session to continue, for example to test a resume. Default: a new in-memory session. */ + sessionManager?: SessionManager; } export interface Harness { @@ -156,7 +158,7 @@ export async function createHarness(options: HarnessOptions = {}): Promise