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

@@ -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
 

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<string>();
 	private _usesDefaultTools: boolean;
 	private _allowedToolNames?: Set<string>;
 	private _excludedToolNames?: Set<string>;
@@ -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(`<sdk:${definition.name}>`, { source: "sdk" }),
 			})),
-		].filter((tool) => isAllowedTool(tool.definition.name));
+		].filter((tool) => this._isAllowedTool(tool.definition.name));
 		const definitionRegistry = new Map<string, ToolDefinitionEntry>(
 			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,

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"));
+	});
+});

packages/coding-agent/test/suite/harness.ts

@@ -119,6 +119,8 @@ export interface HarnessOptions {
 	extensionFactories?: Array<InlineExtension | CreateTestExtensionsResultInput>;
 	withConfiguredAuth?: boolean;
 	modelsJson?: Record<string, unknown>;
+	/** 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<Harne
 	const withConfiguredAuth = options.withConfiguredAuth ?? true;
 	const extensionRunnerRef: { current?: ExtensionRunner } = {};
 
-	const sessionManager = SessionManager.inMemory();
+	const sessionManager = options.sessionManager ?? SessionManager.inMemory();
 	const settingsManager = SettingsManager.inMemory(options.settings);
 
 	const authStorage = AuthStorage.inMemory();