diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index cb490b418..5cc68b613 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/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), diff --git a/packages/coding-agent/src/session/prewalk.ts b/packages/coding-agent/src/session/prewalk.ts index bd361ba58..cc2ed0063 100644 --- a/packages/coding-agent/src/session/prewalk.ts +++ b/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; setActiveToolsByName(names: string[]): Promise; + setActiveToolPresentation(toolNames: string[], mountedToolNames: string[]): Promise; + runToolRegistryMutation(mutation: () => Promise): Promise; 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 { 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", diff --git a/packages/coding-agent/test/agent-session-plan-mode-convergence.test.ts b/packages/coding-agent/test/agent-session-plan-mode-convergence.test.ts index c7e0bc972..eceda4000 100644 --- a/packages/coding-agent/test/agent-session-plan-mode-convergence.test.ts +++ b/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 { @@ -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([ + ["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(), + 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([ - ["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(); + const release = Promise.withResolvers(); + 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 });