smashimo/oh-my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(coding-agent): preserve plan exit rollback state

Browse Source
This commit is contained in:
Victor Araújo
2026-07-18 16:33:45 -03:00
parent 39c95e5e29
commit 8a8ff498b3
2 changed files with 84 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/coding-agent/src/modes/interactive-mode.ts
+36 -12
View File
@@ -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);
packages/coding-agent/test/interactive-mode-default-plan-mode.test.ts
+48
View File
@@ -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<string>;
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 };