From f25944bdc6f669cbff9281d3f604042110eab67e Mon Sep 17 00:00:00 2001 From: Tommaso Sciortino Date: Tue, 16 Dec 2025 10:24:55 -0800 Subject: [PATCH] Fix prompt and chat code (#88) --- packages/core/src/core/geminiChat.test.ts | 34 +++++++++++++++++++++++ packages/core/src/core/geminiChat.ts | 19 +++++-------- packages/core/src/core/prompts.test.ts | 11 +++++++- packages/core/src/core/prompts.ts | 4 +-- 4 files changed, 53 insertions(+), 15 deletions(-) diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 0207ad7be6..98e74631e2 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -21,6 +21,7 @@ import { DEFAULT_GEMINI_FLASH_MODEL, DEFAULT_THINKING_MODE, PREVIEW_GEMINI_MODEL, + PREVIEW_GEMINI_FLASH_MODEL, } from '../config/models.js'; import { AuthType } from './contentGenerator.js'; import { TerminalQuotaError } from '../utils/googleQuotaErrors.js'; @@ -579,6 +580,39 @@ describe('GeminiChat', () => { ); }); + it('should use maxAttempts=1 for retryWithBackoff when in Preview Model Fallback Mode (Flash)', async () => { + vi.mocked(mockConfig.isPreviewModelFallbackMode).mockReturnValue(true); + vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( + (async function* () { + yield { + candidates: [ + { + content: { parts: [{ text: 'Success' }] }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(), + ); + + const stream = await chat.sendMessageStream( + { model: PREVIEW_GEMINI_FLASH_MODEL }, + 'test', + 'prompt-id-fast-retry-flash', + new AbortController().signal, + ); + for await (const _ of stream) { + // consume stream + } + + expect(mockRetryWithBackoff).toHaveBeenCalledWith( + expect.any(Function), + expect.objectContaining({ + maxAttempts: 1, + }), + ); + }); + it('should NOT use maxAttempts=1 for other models even in Preview Model Fallback Mode', async () => { vi.mocked(mockConfig.isPreviewModelFallbackMode).mockReturnValue(true); vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 1ca9c35680..43b5ad3a8e 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -23,9 +23,9 @@ import { retryWithBackoff, isRetryableError } from '../utils/retry.js'; import type { Config } from '../config/config.js'; import { DEFAULT_THINKING_MODE, - PREVIEW_GEMINI_MODEL, resolveModel, isGemini2Model, + isPreviewModel, } from '../config/models.js'; import { hasCycleInSchema } from '../tools/tools.js'; import type { StructuredError } from './turn.js'; @@ -305,10 +305,7 @@ export class GeminiChat { let maxAttempts = INVALID_CONTENT_RETRY_OPTIONS.maxAttempts; // If we are in Preview Model Fallback Mode, we want to fail fast (1 attempt) // when probing the Preview Model. - if ( - this.config.isPreviewModelFallbackMode() && - model === PREVIEW_GEMINI_MODEL - ) { + if (this.config.isPreviewModelFallbackMode() && isPreviewModel(model)) { maxAttempts = 1; } @@ -387,7 +384,7 @@ export class GeminiChat { // Preview Model successfully used, disable fallback mode. // We only do this if we didn't bypass Preview Model (i.e. we actually used it). if ( - model === PREVIEW_GEMINI_MODEL && + isPreviewModel(model) && !this.config.isPreviewModelBypassMode() ) { this.config.setPreviewModelFallbackMode(false); @@ -492,10 +489,9 @@ export class GeminiChat { }; delete config.thinkingConfig?.thinkingLevel; } - let contentsToUse = - modelToUse === PREVIEW_GEMINI_MODEL - ? contentsForPreviewModel - : requestContents; + let contentsToUse = isPreviewModel(modelToUse) + ? contentsForPreviewModel + : requestContents; // Fire BeforeModel and BeforeToolSelection hooks if enabled const hooksEnabled = this.config.getEnableHooks(); @@ -583,8 +579,7 @@ export class GeminiChat { signal: generateContentConfig.abortSignal, maxAttempts: availabilityMaxAttempts ?? - (this.config.isPreviewModelFallbackMode() && - model === PREVIEW_GEMINI_MODEL + (this.config.isPreviewModelFallbackMode() && isPreviewModel(model) ? 1 : undefined), getAvailabilityContext, diff --git a/packages/core/src/core/prompts.test.ts b/packages/core/src/core/prompts.test.ts index b3487dac01..8ba8c89c00 100644 --- a/packages/core/src/core/prompts.test.ts +++ b/packages/core/src/core/prompts.test.ts @@ -15,9 +15,10 @@ import { CodebaseInvestigatorAgent } from '../agents/codebase-investigator.js'; import { GEMINI_DIR } from '../utils/paths.js'; import { debugLogger } from '../utils/debugLogger.js'; import { - DEFAULT_GEMINI_MODEL, PREVIEW_GEMINI_MODEL, + PREVIEW_GEMINI_FLASH_MODEL, DEFAULT_GEMINI_MODEL_AUTO, + DEFAULT_GEMINI_MODEL, } from '../config/models.js'; // Mock tool names if they are dynamically generated or complex @@ -83,6 +84,14 @@ describe('Core System Prompt (prompts.ts)', () => { expect(prompt).toMatchSnapshot(); }); + it('should use chatty system prompt for preview flash model', () => { + vi.mocked(mockConfig.getActiveModel).mockReturnValue( + PREVIEW_GEMINI_FLASH_MODEL, + ); + const prompt = getCoreSystemPrompt(mockConfig); + expect(prompt).toContain('Do not call tools in silence'); + }); + it.each([ ['empty string', ''], ['whitespace only', ' \n \t '], diff --git a/packages/core/src/core/prompts.ts b/packages/core/src/core/prompts.ts index ea4a572996..0165ad6228 100644 --- a/packages/core/src/core/prompts.ts +++ b/packages/core/src/core/prompts.ts @@ -25,7 +25,7 @@ import type { Config } from '../config/config.js'; import { GEMINI_DIR } from '../utils/paths.js'; import { debugLogger } from '../utils/debugLogger.js'; import { WriteTodosTool } from '../tools/write-todos.js'; -import { resolveModel, PREVIEW_GEMINI_MODEL } from '../config/models.js'; +import { resolveModel, isPreviewModel } from '../config/models.js'; export function resolvePathFromEnv(envVar?: string): { isSwitch: boolean; @@ -111,7 +111,7 @@ export function getCoreSystemPrompt( config.getPreviewFeatures(), ); - const isGemini3 = desiredModel === PREVIEW_GEMINI_MODEL; + const isGemini3 = isPreviewModel(desiredModel); const mandatesVariant = isGemini3 ? `