diff --git a/packages/coding-agent/src/modes/interactive-mode.ts b/packages/coding-agent/src/modes/interactive-mode.ts index 9d3a2a59a..94cf50375 100644 --- a/packages/coding-agent/src/modes/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive-mode.ts @@ -2458,26 +2458,50 @@ export class InteractiveMode implements InteractiveModeContext { } const planModeState = this.session.getPlanModeState(); + const planModeTools = this.session.getEnabledToolNames(); + const planModeModelState = this.session.model + ? { model: this.session.model, thinkingLevel: this.session.configuredThinkingLevel() } + : undefined; this.session.setPlanModeState(undefined); try { if (this.#planModePreviousTools !== undefined) { await this.session.setActiveToolsByName(this.#planModePreviousTools); } - if (this.#planModePreviousModelState) { - if (!options?.deferModelRestore) { - await this.#restorePlanPreviousModel(this.#planModePreviousModelState); - } - // If #applyPlanModeModel queued a deferred switch to the plan-role model - // (because the session was streaming on entry), drop it now: we are - // leaving plan mode, so flushing it on the next agent_end would land the - // session on the plan-role model after the user has exited plan mode - // (issue #816). This runs even when deferModelRestore is set - // (compact-approval path): otherwise the stale plan switch survives and - // flushPendingModelSwitch() later clobbers the restored/execution model. - this.#clearPendingPlanModelSwitch(); + if (this.#planModePreviousModelState && !options?.deferModelRestore) { + await this.#restorePlanPreviousModel(this.#planModePreviousModelState); } + // If #applyPlanModeModel queued a deferred switch to the plan-role model + // (because the session was streaming on entry), drop it now: we are + // leaving plan mode, so flushing it on the next agent_end would land the + // session on the plan-role model after the user has exited plan mode + // (issue #816). This runs even when deferModelRestore is set + // (compact-approval path): otherwise the stale plan switch survives and + // flushPendingModelSwitch() later clobbers the restored/execution model. + if (this.#planModePreviousModelState) this.#clearPendingPlanModelSwitch(); } catch (error) { this.session.setPlanModeState(planModeState); + if ( + planModeModelState && + (!modelsAreEqual(this.session.model, planModeModelState.model) || + this.session.configuredThinkingLevel() !== planModeModelState.thinkingLevel) + ) { + try { + await this.#restorePlanPreviousModel(planModeModelState); + } catch (rollbackError) { + logger.warn("Failed to restore plan model after plan exit failure", { error: String(rollbackError) }); + } + } + const enabledTools = this.session.getEnabledToolNames(); + if ( + enabledTools.length !== planModeTools.length || + enabledTools.some((name, index) => name !== planModeTools[index]) + ) { + try { + await this.session.setActiveToolsByName(planModeTools); + } catch (rollbackError) { + logger.warn("Failed to restore plan tools after plan exit failure", { error: String(rollbackError) }); + } + } throw error; } this.session.setPlanProposalHandler?.(null); diff --git a/packages/coding-agent/test/interactive-mode-default-plan-mode.test.ts b/packages/coding-agent/test/interactive-mode-default-plan-mode.test.ts index 501f61edd..f89049cc6 100644 --- a/packages/coding-agent/test/interactive-mode-default-plan-mode.test.ts +++ b/packages/coding-agent/test/interactive-mode-default-plan-mode.test.ts @@ -11,6 +11,7 @@ import { TempDir } from "@oh-my-pi/pi-utils"; import { type } from "arktype"; import { ModelRegistry } from "../src/config/model-registry"; import { InteractiveMode } from "../src/modes/interactive-mode"; +import { XdevRegistry } from "../src/tools/xdev"; function makeTool(name: string): AgentTool { return { @@ -28,6 +29,7 @@ interface HarnessOptions { extraRegistryTools?: readonly AgentTool[]; builtInToolNames?: Iterable; rebuildGate?: { fail: boolean; calls?: number }; + xdevRegistry?: XdevRegistry; } describe("InteractiveMode plan.defaultOnStartup", () => { @@ -85,6 +87,7 @@ describe("InteractiveMode plan.defaultOnStartup", () => { toolRegistry.set(tool.name, tool); } const manager = SessionManager.create(tempDir.path(), path.join(tempDir.path(), `active-${Bun.nanoseconds()}`)); + const xdevRegistry = options.xdevRegistry; const createdSession = new AgentSession({ agent: new Agent({ initialState: { @@ -107,6 +110,7 @@ describe("InteractiveMode plan.defaultOnStartup", () => { return { systemPrompt: ["Test"] }; } : undefined, + xdevRegistry, }); session = createdSession; mode = new InteractiveMode(createdSession, "test"); @@ -195,6 +199,50 @@ describe("InteractiveMode plan.defaultOnStartup", () => { expect(session?.getActiveToolNames()).toEqual(["read"]); }); + it("restores mounted xd devices when prior-model restoration fails", async () => { + const settings = Settings.isolated({ "plan.defaultOnStartup": true, "compaction.enabled": false }); + settings.setModelRole("plan", "anthropic/claude-haiku-4-5:high"); + const writeTool = makeTool("write"); + const mountedTool = { ...makeTool("ambient_search"), loadMode: "discoverable" as const }; + const created = createHarness(settings, { + extraRegistryTools: [writeTool], + builtInToolNames: ["read", "write"], + xdevRegistry: new XdevRegistry([]), + }); + const previousModel = session?.model; + await session!.refreshRpcHostTools([mountedTool]); + await created.init({ suppressWelcomeIntro: true }); + const planModel = session?.model; + const planTools = session?.getEnabledToolNames(); + expect(planModel?.id).toBe("claude-haiku-4-5"); + expect(session?.configuredThinkingLevel()).toBe(Effort.High); + expect(session?.getMountedXdevToolNames()).toEqual([mountedTool.name]); + + const setModelTemporary = session!.setModelTemporary.bind(session); + const restoreModel = vi.spyOn(session!, "setModelTemporary").mockImplementationOnce(async (...args) => { + await setModelTemporary(...args); + throw new Error("model restore failed after switch"); + }); + await expect(created.handlePlanModeCommand()).rejects.toThrow("model restore failed after switch"); + + expect(created.planModeEnabled).toBe(true); + expect(created.planModePaused).toBe(false); + expect(session?.getPlanModeState()?.enabled).toBe(true); + expect(session?.peekPlanProposalHandler()).toBeDefined(); + expect(session?.model?.id).toBe(planModel?.id); + expect(session?.configuredThinkingLevel()).toBe(Effort.High); + expect(session?.getEnabledToolNames()).toEqual(planTools); + expect(session?.getMountedXdevToolNames()).toEqual([mountedTool.name]); + + restoreModel.mockRestore(); + await created.handlePlanModeCommand(); + expect(created.planModeEnabled).toBe(false); + expect(session?.getPlanModeState()).toBeUndefined(); + expect(session?.model?.id).toBe(previousModel?.id); + expect(session?.getEnabledToolNames()).toEqual(["read", "write", mountedTool.name]); + expect(session?.getMountedXdevToolNames()).toEqual([mountedTool.name]); + }); + it("clears old plan UI state when target-session reconciliation restore fails", async () => { const writeTool = makeTool("write"); const rebuildGate = { fail: false, calls: 0 };