Merge PR #8866: Preserve MCP tools across PlanYolo handoff (@nick-maderight)

packages/coding-agent/src/session/agent-session.ts

@@ -1011,8 +1011,13 @@ export class AgentSession {
 			emitNotice: (level, message, source) => this.emitNotice(level, message, source),
 			setModelTemporary: (model, thinkingLevel, options) => this.setModelTemporary(model, thinkingLevel, options),
 			setActiveToolsByName: names => this.setActiveToolsByName(names),
+			setActiveToolPresentation: (toolNames, mountedToolNames) =>
+				this.setActiveToolPresentation(toolNames, mountedToolNames),
+			runToolRegistryMutation: mutation => this.runToolRegistryMutation(mutation),
 			getActiveToolNames: () => this.getActiveToolNames(),
 			getEnabledToolNames: () => this.getEnabledToolNames(),
+			getSelectedMCPToolNames: () => this.getSelectedMCPToolNames(),
+			getMountedXdevToolNames: () => this.getMountedXdevToolNames(),
 			hasBuiltInTool: name => this.hasBuiltInTool(name),
 			getPlanModeState: () => this.getPlanModeState(),
 			setPlanModeState: state => this.setPlanModeState(state),

packages/coding-agent/src/session/prewalk.ts

@@ -12,6 +12,7 @@ import prewalkContinuePrompt from "../prompts/system/prewalk-continue.md" with {
 import prewalkPlanPrompt from "../prompts/system/prewalk-plan.md" with { type: "text" };
 import { type ConfiguredThinkingLevel, prewalkWouldBeNoop } from "../thinking";
 import type { PlanProposalHandler } from "../tools/resolve";
+import { isMCPToolName } from "../tools/builtin-names";
 import { ToolError } from "../tools/tool-errors";
 import type { PlanYolo, Prewalk } from "./agent-session-types";
 import { PREWALK_PLAN_MESSAGE_TYPE } from "./messages";
@@ -65,8 +66,12 @@ export interface PrewalkCoordinatorHost {
 		options?: { ephemeral?: boolean },
 	): Promise<void>;
 	setActiveToolsByName(names: string[]): Promise<void>;
+	setActiveToolPresentation(toolNames: string[], mountedToolNames: string[]): Promise<void>;
+	runToolRegistryMutation<T>(mutation: () => Promise<T>): Promise<T>;
 	getActiveToolNames(): string[];
 	getEnabledToolNames(): string[];
+	getSelectedMCPToolNames(): string[];
+	getMountedXdevToolNames(): string[];
 	hasBuiltInTool(name: string): boolean;
 	getPlanModeState(): PlanModeState | undefined;
 	setPlanModeState(state: PlanModeState | undefined): void;
@@ -90,7 +95,7 @@ export class PrewalkCoordinator {
 	#continuePending = false;
 	#todoSeen = false;
 	#planYolo: PlanYolo | undefined;
-	#planYoloPreviousTools: string[] | undefined;
+	#planYoloPreviousNonMCPPresentation: { enabled: string[]; mounted: string[] } | undefined;
 	#planYoloArmed = false;
 
 	constructor(host: PrewalkCoordinatorHost, options: PrewalkCoordinatorOptions = {}) {
@@ -247,10 +252,14 @@ export class PrewalkCoordinator {
 	async armPlanYoloIfNeeded(): Promise<void> {
 		if (!this.#planYolo || this.#planYoloArmed) return;
 		this.#planYoloArmed = true;
-		const previousTools = this.#host.getEnabledToolNames();
+		const previousEnabledTools = this.#host.getEnabledToolNames();
+		const previousMountedTools = this.#host.getMountedXdevToolNames();
 		const augmentations = this.#host.hasBuiltInTool("write") ? ["write"] : [];
-		await this.#host.setActiveToolsByName([...new Set([...previousTools, ...augmentations])]);
-		this.#planYoloPreviousTools = previousTools;
+		await this.#host.setActiveToolsByName([...new Set([...previousEnabledTools, ...augmentations])]);
+		this.#planYoloPreviousNonMCPPresentation = {
+			enabled: previousEnabledTools.filter(name => !isMCPToolName(name)),
+			mounted: previousMountedTools.filter(name => !isMCPToolName(name)),
+		};
 		this.#host.setPlanModeState({
 			enabled: true,
 			planFilePath: this.#host.getPlanReferencePath() || "local://PLAN.md",
@@ -287,16 +296,25 @@ export class PrewalkCoordinator {
 			listPlanFiles: () => listPlanFiles({ localProtocolOptions: this.#host.localProtocolOptions() }),
 		});
 		this.#host.setPlanModeState(undefined);
-		const previousTools = this.#planYoloPreviousTools;
+		const previousPresentation = this.#planYoloPreviousNonMCPPresentation;
 		try {
-			if (previousTools) await this.#host.setActiveToolsByName(previousTools);
+			if (previousPresentation) {
+				await this.#host.runToolRegistryMutation(async () => {
+					const liveMCP = this.#host.getSelectedMCPToolNames();
+					const liveMountedMCP = this.#host.getMountedXdevToolNames().filter(isMCPToolName);
+					await this.#host.setActiveToolPresentation(
+						[...new Set([...previousPresentation.enabled, ...liveMCP])],
+						[...new Set([...previousPresentation.mounted, ...liveMountedMCP])],
+					);
+				});
+			}
 		} catch (error) {
 			this.#host.setPlanModeState(state);
 			throw error;
 		}
 		this.#host.setPlanProposalHandler(null);
 		this.#planYolo = undefined;
-		this.#planYoloPreviousTools = undefined;
+		this.#planYoloPreviousNonMCPPresentation = undefined;
 		await this.#host.setModelTemporary(planYolo.target, planYolo.thinkingLevel, { ephemeral: true });
 		this.#host.emitNotice(
 			"info",

packages/coding-agent/test/agent-session-plan-mode-convergence.test.ts

@@ -12,7 +12,7 @@
  */
 import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it } from "bun:test";
 import { type } from "@oh-my-pi/omptype";
-import { Agent, type AgentMessage, type AgentTool, type StreamFn } from "@oh-my-pi/pi-agent-core";
+import { Agent, type AgentMessage, type AgentTool, type StreamFn, type ToolApproval, type ToolLoadMode } from "@oh-my-pi/pi-agent-core";
 import { createMockModel, type MockModel, type MockResponse } from "@oh-my-pi/pi-ai/providers/mock";
 import { getBundledModel } from "@oh-my-pi/pi-catalog/models";
 import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
@@ -23,6 +23,8 @@ import { AgentRegistry } from "@oh-my-pi/pi-coding-agent/registry/agent-registry
 import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session";
 import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
 import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
+import type { CustomTool } from "@oh-my-pi/pi-coding-agent/extensibility/custom-tools/types";
+import type { XdevState } from "@oh-my-pi/pi-coding-agent/tools/xdev";
 import { TempDir } from "@oh-my-pi/pi-utils";
 import planModeReminderPrompt from "../src/prompts/system/plan-mode-tool-decision-reminder.md" with { type: "text" };
 
@@ -50,6 +52,22 @@ function makeTool(name: string): AgentTool {
 	};
 }
 
+function makeMcpTool(name: string, loadMode: ToolLoadMode, approval: ToolApproval = "read"): CustomTool {
+	return {
+		name,
+		label: name,
+		description: `Test MCP tool ${name}`,
+		parameters: type({}),
+		loadMode,
+		approval,
+		mcpServerName: name.split("__")[1] ?? "test-mcp",
+		mcpToolName: name.split("__").at(-1) ?? name,
+		async execute() {
+			return { content: [{ type: "text", text: "ok" }] };
+		},
+	};
+}
+
 /** Concatenate the text blocks of a message (string or content-array). */
 function messageText(message: AgentMessage): string {
 	if (!("content" in message)) return "";
@@ -112,6 +130,8 @@ describe("AgentSession plan-mode convergence", () => {
 			advisorResponses?: MockResponse[];
 			sideResponses?: MockResponse[];
 			planYolo?: boolean;
+			initialPlanTools?: string[];
+			xdev?: boolean;
 			rebuildGate?: { fail: boolean };
 		},
 	): Promise<PlanHarness> {
@@ -121,20 +141,38 @@ describe("AgentSession plan-mode convergence", () => {
 		const askTool = makeTool("ask");
 		const writeTool = makeTool("write");
 		const readTool = makeTool("read");
+		const toolRegistry = new Map<string, AgentTool>([
+			["ask", askTool],
+			["write", writeTool],
+			["read", readTool],
+		]);
+		const initialTools = options?.planYolo
+			? options.initialPlanTools?.includes("write")
+				? [readTool, writeTool]
+				: [readTool]
+			: [askTool, writeTool, readTool];
+		let currentAgent: Agent | undefined;
+		const xdev: XdevState | undefined = options?.xdev
+			? {
+					tools: toolRegistry,
+					mountedNames: new Set<string>(),
+					builtInNames: new Set(["ask", "write", "read"]),
+					isActive: name => currentAgent?.state.tools.some(tool => tool.name === name) ?? false,
+				}
+			: undefined;
 
 		const mock = createMockModel({ responses });
 		const agent = new Agent({
 			getApiKey: () => "test-key",
-			// All three tools active so a scripted ask/write/read call (and a
-			// forced "required" choice) can actually execute (isToolChoiceActive).
 			initialState: {
 				model,
 				systemPrompt: ["Test"],
-				tools: options?.planYolo ? [readTool] : [askTool, writeTool, readTool],
+				tools: initialTools,
 				messages: [],
 			},
 			streamFn: mock.stream,
 		});
+		currentAgent = agent;
 
 		let advisorMock: MockModel | undefined;
 		let advisorStreamFn: StreamFn | undefined;
@@ -158,16 +196,13 @@ describe("AgentSession plan-mode convergence", () => {
 				"retry.enabled": false,
 			}),
 			modelRegistry,
-			toolRegistry: new Map<string, AgentTool>([
-				["ask", askTool],
-				["write", writeTool],
-				["read", readTool],
-			]),
+			toolRegistry,
 			builtInToolNames: ["ask", "write", "read"],
 			advisorTools: [],
 			advisorStreamFn,
 			sideStreamFn,
 			planYolo: options?.planYolo ? { target: model } : undefined,
+			xdev,
 			rebuildSystemPrompt: options?.rebuildGate
 				? async () => {
 						if (options.rebuildGate?.fail) throw new Error("rebuild failed");
@@ -347,6 +382,112 @@ describe("AgentSession plan-mode convergence", () => {
 		expect(harness.session.getActiveToolNames()).toEqual(["read"]);
 	});
 
+	it("retains MCP devices discovered while PlanYolo is active", async () => {
+		const harness = await createPlanSession(
+			[{ content: ["planning A"] }, { content: ["planning B"] }, { content: ["planning C"] }],
+			{ planYolo: true, initialPlanTools: ["read", "write"], xdev: true },
+		);
+		await harness.session.prompt("make a plan");
+		await harness.session.waitForIdle();
+
+		const chromeTool = makeMcpTool("mcp__chrome_devtools_list_pages", "discoverable");
+		const contextTool = makeMcpTool("mcp__context_query_docs", "essential");
+		await harness.session.refreshMCPTools([chromeTool, contextTool]);
+		expect(harness.session.getSelectedMCPToolNames()).toEqual([
+			"mcp__context_query_docs",
+			"mcp__chrome_devtools_list_pages",
+		]);
+		expect(harness.session.getActiveToolNames()).toContain("mcp__context_query_docs");
+		expect(harness.session.getMountedXdevToolNames()).toContain("mcp__chrome_devtools_list_pages");
+
+		const planPath = resolveLocalUrlToPath("local://mcp-devices-plan.md", {
+			getArtifactsDir: () => harness.session.sessionManager.getArtifactsDir(),
+			getSessionId: () => harness.session.sessionManager.getSessionId(),
+		});
+		await Bun.write(planPath, "# MCP devices plan\n\nKeep the connected devices.\n");
+		const handler = harness.session.peekPlanProposalHandler();
+		if (!handler) throw new Error("Expected PlanYolo proposal handler");
+		await handler("mcp-devices");
+
+		expect(harness.session.getPlanModeState()).toBeUndefined();
+		expect(harness.session.getActiveToolNames()).toEqual([
+			"read",
+			"write",
+			"mcp__context_query_docs",
+		]);
+		expect(harness.session.getMountedXdevToolNames()).toEqual(["mcp__chrome_devtools_list_pages"]);
+		expect(harness.session.getSelectedMCPToolNames()).toEqual([
+			"mcp__context_query_docs",
+			"mcp__chrome_devtools_list_pages",
+		]);
+	});
+
+	it("serializes PlanYolo restoration after a pending MCP refresh", async () => {
+		const harness = await createPlanSession(
+			[{ content: ["planning A"] }, { content: ["planning B"] }, { content: ["planning C"] }],
+			{ planYolo: true, initialPlanTools: ["read", "write"], xdev: true },
+		);
+		await harness.session.prompt("make a plan");
+		await harness.session.waitForIdle();
+
+		const entered = Promise.withResolvers<void>();
+		const release = Promise.withResolvers<void>();
+		const blocker = harness.session.runToolRegistryMutation(async () => {
+			entered.resolve();
+			await release.promise;
+		});
+		await entered.promise;
+
+		const chromeTool = makeMcpTool("mcp__chrome_devtools_list_pages", "discoverable");
+		const refresh = harness.session.refreshMCPTools([chromeTool]);
+		const planPath = resolveLocalUrlToPath("local://queued-mcp-plan.md", {
+			getArtifactsDir: () => harness.session.sessionManager.getArtifactsDir(),
+			getSessionId: () => harness.session.sessionManager.getSessionId(),
+		});
+		await Bun.write(planPath, "# Queued MCP plan\n\nKeep the connected device.\n");
+		const handler = harness.session.peekPlanProposalHandler();
+		if (!handler) throw new Error("Expected PlanYolo proposal handler");
+		const approval = handler("queued-mcp");
+		release.resolve();
+		await Promise.all([blocker, refresh, approval]);
+
+		expect(harness.session.getPlanModeState()).toBeUndefined();
+		expect(harness.session.getActiveToolNames()).toContain("read");
+		expect(harness.session.getActiveToolNames()).toContain("write");
+		expect(harness.session.getMountedXdevToolNames()).toEqual(["mcp__chrome_devtools_list_pages"]);
+		expect(harness.session.getSelectedMCPToolNames()).toContain("mcp__chrome_devtools_list_pages");
+	});
+
+	it("preserves late MCP selection without leaking plan-only write", async () => {
+		const harness = await createPlanSession(
+			[{ content: ["planning A"] }, { content: ["planning B"] }, { content: ["planning C"] }],
+			{ planYolo: true, xdev: true },
+		);
+		await harness.session.prompt("make a plan");
+		await harness.session.waitForIdle();
+
+		const chromeTool = makeMcpTool("mcp__chrome_devtools_list_pages", "discoverable", "write");
+		await harness.session.refreshMCPTools([chromeTool]);
+		const registeredTool = harness.session.getToolByName("mcp__chrome_devtools_list_pages");
+		expect(registeredTool).toBeDefined();
+
+		const planPath = resolveLocalUrlToPath("local://read-only-mcp-plan.md", {
+			getArtifactsDir: () => harness.session.sessionManager.getArtifactsDir(),
+			getSessionId: () => harness.session.sessionManager.getSessionId(),
+		});
+		await Bun.write(planPath, "# Read-only MCP plan\n\nKeep the selected device.\n");
+		const handler = harness.session.peekPlanProposalHandler();
+		if (!handler) throw new Error("Expected PlanYolo proposal handler");
+		await handler("read-only-mcp");
+
+		expect(harness.session.getPlanModeState()).toBeUndefined();
+		expect(harness.session.getActiveToolNames()).toEqual(["read", "mcp__chrome_devtools_list_pages"]);
+		expect(harness.session.getActiveToolNames()).not.toContain("write");
+		expect(harness.session.getMountedXdevToolNames()).toEqual([]);
+		expect(harness.session.getSelectedMCPToolNames()).toEqual(["mcp__chrome_devtools_list_pages"]);
+		expect(harness.session.getToolByName("mcp__chrome_devtools_list_pages")).toBe(registeredTool);
+	});
+
 	it("keeps PlanYolo retryable when pre-plan tool restoration fails", async () => {
 		const rebuildGate = { fail: false };
 		const harness = await createPlanSession([{ content: ["planning"] }], { planYolo: true, rebuildGate });