From 9762bf296527737ab4eeaedb26f384c1f1f5139b Mon Sep 17 00:00:00 2001
From: Adib234 <30782825+Adib234@users.noreply.github.com>
Date: Thu, 26 Mar 2026 14:45:03 -0400
Subject: [PATCH] fix(plan): after exiting plan mode switches model to a flash
 model (#23885)

---
 integration-tests/plan-mode.test.ts | 68 ++++++++++++++++++++++++++++-
 packages/core/src/config/config.ts  |  1 +
 packages/core/src/core/client.ts    |  4 ++
 3 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/integration-tests/plan-mode.test.ts b/integration-tests/plan-mode.test.ts
index 977a754f1e..d8d297c460 100644
--- a/integration-tests/plan-mode.test.ts
+++ b/integration-tests/plan-mode.test.ts
@@ -4,8 +4,10 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { writeFileSync, mkdirSync } from 'node:fs';
+import { join } from 'node:path';
 import { describe, it, expect, beforeEach, afterEach } from 'vitest';
-import { TestRig, checkModelOutputContent } from './test-helper.js';
+import { GEMINI_DIR, TestRig, checkModelOutputContent } from './test-helper.js';
 
 describe('Plan Mode', () => {
   let rig: TestRig;
@@ -227,4 +229,68 @@ describe('Plan Mode', () => {
       `Expected write_file to succeed, but it failed with error: ${planWrite?.toolRequest.error}`,
     ).toBe(true);
   });
+  it('should switch from a pro model to a flash model after exiting plan mode', async () => {
+    const plansDir = 'plans-folder';
+    const planFilename = 'my-plan.md';
+
+    await rig.setup('should-switch-to-flash', {
+      settings: {
+        model: {
+          name: 'auto-gemini-2.5',
+        },
+        experimental: { plan: true },
+        tools: {
+          core: ['exit_plan_mode', 'run_shell_command'],
+          allowed: ['exit_plan_mode', 'run_shell_command'],
+        },
+        general: {
+          defaultApprovalMode: 'plan',
+          plan: {
+            directory: plansDir,
+          },
+        },
+      },
+    });
+
+    writeFileSync(
+      join(rig.homeDir!, GEMINI_DIR, 'state.json'),
+      JSON.stringify({ terminalSetupPromptShown: true }, null, 2),
+    );
+
+    const fullPlansDir = join(rig.testDir!, plansDir);
+    mkdirSync(fullPlansDir, { recursive: true });
+    writeFileSync(join(fullPlansDir, planFilename), 'Execute echo hello');
+
+    await rig.run({
+      approvalMode: 'plan',
+      stdin: `Exit plan mode using ${planFilename} and then run a shell command \`echo hello\`.`,
+    });
+
+    const exitCallFound = await rig.waitForToolCall('exit_plan_mode');
+    expect(exitCallFound, 'Expected exit_plan_mode to be called').toBe(true);
+
+    const shellCallFound = await rig.waitForToolCall('run_shell_command');
+    expect(shellCallFound, 'Expected run_shell_command to be called').toBe(
+      true,
+    );
+
+    const apiRequests = rig.readAllApiRequest();
+    const modelNames = apiRequests.map((r) => r.attributes?.model || 'unknown');
+
+    const proRequests = apiRequests.filter((r) =>
+      r.attributes?.model?.includes('pro'),
+    );
+    const flashRequests = apiRequests.filter((r) =>
+      r.attributes?.model?.includes('flash'),
+    );
+
+    expect(
+      proRequests.length,
+      `Expected at least one Pro request. Models used: ${modelNames.join(', ')}`,
+    ).toBeGreaterThanOrEqual(1);
+    expect(
+      flashRequests.length,
+      `Expected at least one Flash request after mode switch. Models used: ${modelNames.join(', ')}`,
+    ).toBeGreaterThanOrEqual(1);
+  });
 });
diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts
index e727881a04..70ac02e22f 100644
--- a/packages/core/src/config/config.ts
+++ b/packages/core/src/config/config.ts
@@ -2413,6 +2413,7 @@ export class Config implements McpContext, AgentLoopContext {
 
     if (isPlanModeTransition || isYoloModeTransition) {
       if (this._geminiClient?.isInitialized()) {
+        this._geminiClient.clearCurrentSequenceModel();
         this._geminiClient.setTools().catch((err) => {
           debugLogger.error('Failed to update tools', err);
         });
diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts
index b37d4ad91c..8922c977f2 100644
--- a/packages/core/src/core/client.ts
+++ b/packages/core/src/core/client.ts
@@ -132,6 +132,10 @@ export class GeminiClient {
     this.updateSystemInstruction();
   };
 
+  clearCurrentSequenceModel(): void {
+    this.currentSequenceModel = null;
+  }
+
   // Hook state to deduplicate BeforeAgent calls and track response for
   // AfterAgent
   private hookStateMap = new Map<