From 3c1bb8c35d9b43128636cdee24af01bf717c66fb Mon Sep 17 00:00:00 2001 From: amelidev Date: Thu, 23 Jul 2026 13:36:46 -0600 Subject: [PATCH] fix(core): rotate session ID on model fallback to prevent stateful API errors (#28469) --- .../autoRoutingFallback.integration.test.ts | 82 +++++++++++++++++++ packages/core/src/code_assist/server.ts | 12 ++- packages/core/src/config/config.ts | 4 + packages/core/src/fallback/handler.test.ts | 1 + packages/core/src/fallback/handler.ts | 6 +- 5 files changed, 99 insertions(+), 6 deletions(-) diff --git a/packages/core/src/availability/autoRoutingFallback.integration.test.ts b/packages/core/src/availability/autoRoutingFallback.integration.test.ts index 9ea062e1ab..fad2eb0bbf 100644 --- a/packages/core/src/availability/autoRoutingFallback.integration.test.ts +++ b/packages/core/src/availability/autoRoutingFallback.integration.test.ts @@ -414,4 +414,86 @@ describe('Auto Routing Fallback Integration', () => { 'Pro success', ); }); + + it('should rotate session ID on fallback and retry successfully with the Flash model', async () => { + const originalSessionId = 'test-session-rotate-id'; + config = new Config({ + sessionId: originalSessionId, + targetDir: '/test', + debugMode: false, + cwd: '/test', + model: PREVIEW_GEMINI_MODEL_AUTO, + }); + + vi.spyOn(config, 'isInteractive').mockReturnValue(true); + + client = new BaseLlmClient( + fakeGenerator, + config, + AuthType.LOGIN_WITH_GOOGLE, + ); + + let attemptsPro = 0; + let attemptsFlash = 0; + + const mockGoogleApiError = { + code: 429, + message: + 'Automatically switching from gemini-2.5-pro to gemini-2.5-flash for faster responses for the remainder of this session. Possible reasons for this are...', + details: [], + }; + + vi.spyOn(fakeGenerator, 'generateContent').mockImplementation( + async (params) => { + if (params.model === PREVIEW_GEMINI_MODEL) { + attemptsPro++; + throw new RetryableQuotaError( + 'Quota exceeded for Pro', + mockGoogleApiError, + 0, + ); + } else if (params.model === PREVIEW_GEMINI_FLASH_MODEL) { + attemptsFlash++; + return { + candidates: [ + { + content: { + role: 'model', + parts: [{ text: 'Flash success after rotation' }], + }, + }, + ], + } as unknown as GenerateContentResponse; + } + throw new Error(`Unexpected model: ${params.model}`); + }, + ); + + config.setFallbackModelHandler( + async (_failed, _fallback, _error): Promise => + 'retry_always', // Approve switch to Flash + ); + + const promise = client.generateContent({ + modelConfigKey: { model: PREVIEW_GEMINI_MODEL, isChatModel: true }, + contents: [{ role: 'user', parts: [{ text: 'test query' }] }], + abortSignal: new AbortController().signal, + promptId: 'test-prompt', + role: LlmRole.UTILITY_TOOL, + }); + + await vi.runAllTimersAsync(); + const result = await promise; + + // Verify it resolved to Flash success instead of failing with Please submit a new query + expect(result.candidates?.[0]?.content?.parts?.[0]?.text).toBe( + 'Flash success after rotation', + ); + expect(attemptsPro).toBe(3); + expect(attemptsFlash).toBe(1); + + // Verify session ID has been rotated + expect(config.getSessionId()).not.toBe(originalSessionId); + expect(config.getSessionId()).toBeDefined(); + }); }); diff --git a/packages/core/src/code_assist/server.ts b/packages/core/src/code_assist/server.ts index 92fc558ebb..366a7b324f 100644 --- a/packages/core/src/code_assist/server.ts +++ b/packages/core/src/code_assist/server.ts @@ -86,6 +86,10 @@ export class CodeAssistServer implements ContentGenerator { readonly config?: Config, ) {} + getEffectiveSessionId(): string | undefined { + return this.config?.getSessionId() ?? this.sessionId; + } + async generateContentStream( req: GenerateContentParameters, userPromptId: string, @@ -117,7 +121,7 @@ export class CodeAssistServer implements ContentGenerator { req, userPromptId, this.projectId, - this.sessionId, + this.getEffectiveSessionId(), enabledCreditTypes, ), req.config?.abortSignal, @@ -153,7 +157,7 @@ export class CodeAssistServer implements ContentGenerator { translatedResponse, streamingLatency, req.config?.abortSignal, - server.sessionId, // Use sessionId as trajectoryId + server.getEffectiveSessionId(), // Use sessionId as trajectoryId ); if (response.consumedCredits) { @@ -204,7 +208,7 @@ export class CodeAssistServer implements ContentGenerator { req, userPromptId, this.projectId, - this.sessionId, + this.getEffectiveSessionId(), undefined, ), req.config?.abortSignal, @@ -224,7 +228,7 @@ export class CodeAssistServer implements ContentGenerator { translatedResponse, streamingLatency, req.config?.abortSignal, - this.sessionId, // Use sessionId as trajectoryId + this.getEffectiveSessionId(), // Use sessionId as trajectoryId ); if (response.remainingCredits) { diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index b6cf957935..56fa05e4d4 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -1861,6 +1861,10 @@ export class Config implements McpContext, AgentLoopContext { } } + rotateSessionId(sessionId: string): void { + this._sessionId = sessionId; + } + resetNewSessionState(sessionId: string): void { this.setSessionId(sessionId); } diff --git a/packages/core/src/fallback/handler.test.ts b/packages/core/src/fallback/handler.test.ts index 7931a23007..52516a9c8c 100644 --- a/packages/core/src/fallback/handler.test.ts +++ b/packages/core/src/fallback/handler.test.ts @@ -67,6 +67,7 @@ const createMockConfig = (overrides: Partial = {}): Config => setActiveModel: vi.fn(), setModel: vi.fn(), activateFallbackMode: vi.fn(), + rotateSessionId: vi.fn(), getModelAvailabilityService: vi.fn(() => createAvailabilityServiceMock({ selectedModel: FALLBACK_MODEL, diff --git a/packages/core/src/fallback/handler.ts b/packages/core/src/fallback/handler.ts index 2d26279cce..50de972652 100644 --- a/packages/core/src/fallback/handler.ts +++ b/packages/core/src/fallback/handler.ts @@ -5,6 +5,7 @@ */ import type { Config } from '../config/config.js'; +import { createSessionId } from '../utils/session.js'; import { openBrowserSecurely, shouldLaunchBrowser, @@ -161,8 +162,9 @@ async function processIntent( ): Promise { switch (intent) { case 'retry_always': - // TODO(telemetry): Implement generic fallback event logging. Existing - // logFlashFallback is specific to a single Model. + // Rotate the session ID to ensure the backend treats the retried request + // as a brand-new session, preventing stateful model-switching errors. + config.rotateSessionId(createSessionId()); config.activateFallbackMode(fallbackModel, failedModel); return true;