Skip to content
Original file line number Diff line number Diff line change
Expand Up @@ -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<FallbackIntent | null> =>
'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();
});
});
12 changes: 8 additions & 4 deletions packages/core/src/code_assist/server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -117,7 +121,7 @@ export class CodeAssistServer implements ContentGenerator {
req,
userPromptId,
this.projectId,
this.sessionId,
this.getEffectiveSessionId(),
Comment thread
amelidev marked this conversation as resolved.
Comment thread
amelidev marked this conversation as resolved.
enabledCreditTypes,
),
req.config?.abortSignal,
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -204,7 +208,7 @@ export class CodeAssistServer implements ContentGenerator {
req,
userPromptId,
this.projectId,
this.sessionId,
this.getEffectiveSessionId(),
undefined,
),
req.config?.abortSignal,
Expand All @@ -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) {
Expand Down
4 changes: 4 additions & 0 deletions packages/core/src/config/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1861,6 +1861,10 @@
}
}

rotateSessionId(sessionId: string): void {
this._sessionId = sessionId;
}
Comment thread
amelidev marked this conversation as resolved.

resetNewSessionState(sessionId: string): void {
this.setSessionId(sessionId);
}
Expand Down Expand Up @@ -3567,9 +3571,9 @@
// Gemini API key users should have the ability to manually select the
// old preview flash model.
if (authType === AuthType.USE_GEMINI) {
setFlashModels('gemini-3-flash-preview', 'gemini-3.5-flash');

Check warning on line 3574 in packages/core/src/config/config.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3.5". Please make sure this change is appropriate to submit.
} else {
setFlashModels('gemini-3.5-flash', 'gemini-3.5-flash');

Check warning on line 3576 in packages/core/src/config/config.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3.5". Please make sure this change is appropriate to submit.

Check warning on line 3576 in packages/core/src/config/config.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3.5". Please make sure this change is appropriate to submit.
}
} else {
setFlashModels('gemini-3-flash-preview', 'gemini-2.5-flash');
Expand Down
1 change: 1 addition & 0 deletions packages/core/src/fallback/handler.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@ const createMockConfig = (overrides: Partial<Config> = {}): Config =>
setActiveModel: vi.fn(),
setModel: vi.fn(),
activateFallbackMode: vi.fn(),
rotateSessionId: vi.fn(),
getModelAvailabilityService: vi.fn(() =>
createAvailabilityServiceMock({
selectedModel: FALLBACK_MODEL,
Expand Down
6 changes: 4 additions & 2 deletions packages/core/src/fallback/handler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
*/

import type { Config } from '../config/config.js';
import { createSessionId } from '../utils/session.js';
import {
openBrowserSecurely,
shouldLaunchBrowser,
Expand Down Expand Up @@ -161,8 +162,9 @@ async function processIntent(
): Promise<boolean> {
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;

Expand Down
Loading