Fix prompt and chat code (#88)

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(

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,

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 '],

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
     ? `