fix(coding-agent): applied selected plan execution model after exiting plan mode

- Captured the operator's selected execution tier and passed it through plan approval options.
- Applied the stored model after exiting plan mode so #exitPlanMode's restore no longer reverts it before execution begins.
- Added a regression test that selects a different slider tier and verifies execution runs on the chosen model.
This commit is contained in:
can1357
2026-05-30 18:32:13 +02:00
parent e1a0d235ec
commit 617c74d9a7
2 changed files with 92 additions and 20 deletions
@@ -58,7 +58,7 @@ import planModeApprovedPrompt from "../prompts/system/plan-mode-approved.md" wit
import planModeCompactInstructionsPrompt from "../prompts/system/plan-mode-compact-instructions.md" with {
type: "text",
};
import type { AgentSession, AgentSessionEvent } from "../session/agent-session";
import type { AgentSession, AgentSessionEvent, ResolvedRoleModel } from "../session/agent-session";
import { HistoryStorage } from "../session/history-storage";
import type { SessionContext, SessionManager } from "../session/session-manager";
import { getRecentSessions } from "../session/session-manager";
@@ -1703,6 +1703,20 @@ export class InteractiveMode implements InteractiveModeContext {
}
}
async #applyPlanExecutionModel(entry: ResolvedRoleModel | undefined): Promise<void> {
if (!entry) return;
try {
await this.session.applyRoleModel(entry);
this.statusLine.invalidate();
this.updateEditorBorderColor();
this.showStatus(`Continuing with ${entry.role}: ${entry.model.name || entry.model.id}`);
} catch (error) {
this.showWarning(
`Could not switch to the ${entry.role} model: ${error instanceof Error ? error.message : String(error)}`,
);
}
}
async #approvePlan(
planContent: string,
options: {
@@ -1711,6 +1725,7 @@ export class InteractiveMode implements InteractiveModeContext {
title: string;
preserveContext?: boolean;
compactBeforeExecute?: boolean;
executionModel?: ResolvedRoleModel;
},
): Promise<void> {
await renameApprovedPlanFile({
@@ -1792,6 +1807,8 @@ export class InteractiveMode implements InteractiveModeContext {
return;
}
await this.#applyPlanExecutionModel(options.executionModel);
// Approved plans land in a fresh (or compacted) session whose first user-visible
// turn is the synthetic plan-approved prompt — that path bypasses the
// input-controller's title generation. Seed an auto-name from the plan title
@@ -2149,31 +2166,21 @@ export class InteractiveMode implements InteractiveModeContext {
this.showError(`Plan file not found at ${planFilePath}`);
return;
}
// Apply the operator's tier choice before dispatch so the execution turn
// — and any fresh/compacted session #approvePlan spawns — runs on it. The
// agent's model survives newSession()/compaction, so setting it here is
// sufficient for all three approve paths.
if (cycle && selectedTierIndex !== cycle.currentIndex) {
const chosen = cycle.models[selectedTierIndex];
if (chosen) {
try {
await this.session.applyRoleModel(chosen);
this.statusLine.invalidate();
this.updateEditorBorderColor();
this.showStatus(`Continuing with ${chosen.role}: ${chosen.model.name || chosen.model.id}`);
} catch (error) {
this.showWarning(
`Could not switch to the ${chosen.role} model: ${error instanceof Error ? error.message : String(error)}`,
);
}
}
}
// Capture the operator's tier choice and hand it to #approvePlan, which
// applies it AFTER #exitPlanMode. #exitPlanMode restores
// #planModePreviousModelState (the model from before plan mode), so
// applying the slider choice any earlier would be silently reverted —
// the bug that made "continue with slow" keep executing on the default
// model. Deferred application also survives newSession()/compaction.
const executionModel =
cycle && selectedTierIndex !== cycle.currentIndex ? cycle.models[selectedTierIndex] : undefined;
await this.#approvePlan(latestPlanContent, {
planFilePath,
finalPlanFilePath,
title: details.title,
preserveContext: choice !== "Approve and execute",
compactBeforeExecute: choice === "Approve and compact context",
executionModel,
});
} catch (error) {
this.showError(
@@ -9,6 +9,7 @@ import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme";
import { SILENT_ABORT_MARKER } from "@oh-my-pi/pi-coding-agent/session/messages";
import { Text } from "@oh-my-pi/pi-tui";
import { TempDir } from "@oh-my-pi/pi-utils";
import type { HookSelectorSlider } from "../src/modes/components/hook-selector";
import { ModelRegistry } from "../src/config/model-registry";
import { InteractiveMode } from "../src/modes/interactive-mode";
import { AgentSession } from "../src/session/agent-session";
@@ -141,6 +142,7 @@ describe("InteractiveMode plan review rendering", () => {
"Plan mode - next step",
["Approve and execute", "Approve and compact context", "Approve and keep context (73.2%)", "Refine plan"],
expect.any(Object),
expect.any(Object),
);
});
@@ -169,6 +171,7 @@ describe("InteractiveMode plan review rendering", () => {
"Plan mode - next step",
["Approve and execute", "Approve and compact context", "Approve and keep context", "Refine plan"],
expect.any(Object),
expect.any(Object),
);
});
@@ -234,6 +237,68 @@ describe("InteractiveMode plan review rendering", () => {
});
});
it("executes on the slider-selected tier, surviving #exitPlanMode's model restore", async () => {
// Regression: the model-tier slider's choice used to be applied BEFORE
// #approvePlan ran. #approvePlan → #exitPlanMode restores the model that
// was active before plan mode (#planModePreviousModelState), which silently
// reverted the operator's pick — sliding to "slow" still executed on the
// default model. The fix defers application until after the plan-mode exit.
authStorage.setRuntimeApiKey("anthropic", "test-key");
const slow = session.modelRegistry.find("anthropic", "claude-opus-4-5");
const def = session.modelRegistry.find("anthropic", "claude-sonnet-4-5");
if (!slow || !def) throw new Error("Expected sonnet + opus to exist in registry");
// plan === default === the session model: this is what makes plan-mode entry
// record a previous-model state for #exitPlanMode to restore. slow differs,
// so an early application would be clobbered by that restore.
session.settings.setModelRole("default", "anthropic/claude-sonnet-4-5");
session.settings.setModelRole("slow", "anthropic/claude-opus-4-5");
session.settings.setModelRole("plan", "anthropic/claude-sonnet-4-5");
const planFilePath = "local://PLAN.md";
const finalPlanFilePath = "local://APPROVED.md";
const resolvedPlanPath = resolveLocalUrlToPath(planFilePath, {
getArtifactsDir: () => session.sessionManager.getArtifactsDir(),
getSessionId: () => session.sessionManager.getSessionId(),
});
await Bun.write(resolvedPlanPath, "# Plan\n\nRun this on the slow tier.");
await mode.handlePlanModeCommand();
expect(session.getPlanModeState()?.enabled).toBe(true);
expect(session.model?.id).toBe(def.id);
// Keep-context path avoids newSession() so the assertion isolates the
// exit-plan-mode restore from session-clear effects.
vi.spyOn(session, "getContextUsage").mockReturnValue({ tokens: null, contextWindow: 200000, percent: null });
vi.spyOn(session, "prompt").mockResolvedValue(undefined as never);
let observedSegments: string[] = [];
vi.spyOn(mode, "showHookSelector").mockImplementation(
async (_title, _options, _dialogOptions, extra?: { slider?: HookSelectorSlider }) => {
const slider = extra?.slider;
expect(slider).toBeDefined();
observedSegments = slider!.segments.map(segment => segment.label);
const slowIndex = slider!.segments.findIndex(segment => segment.label === "slow");
expect(slowIndex).toBeGreaterThanOrEqual(0);
// Simulate the operator sliding the tier to "slow" before approving.
slider!.onChange?.(slowIndex);
return "Approve and keep context";
},
);
await mode.handlePlanApproval({
planFilePath,
planExists: true,
title: "PLAN",
finalPlanFilePath,
});
expect(observedSegments).toEqual(["default", "slow"]);
// The load-bearing assertion: the approved plan executes on the operator's
// selected tier, not the restored default.
expect(session.model?.id).toBe(slow.id);
});
it("re-enters plan mode on the approved titled artifact after approve-and-execute", async () => {
const planFilePath = "local://PLAN.md";
const finalPlanFilePath = "local://APPROVED.md";