@office-ai/aioncli-core 0.2.3 → 0.18.4

package/dist/src/core/geminiChat.test.js

@@ -4,420 +4,1665 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
-import { GeminiChat } from './geminiChat.js';
+import { ApiError, ThinkingLevel } from '@google/genai';
+import { GeminiChat, InvalidStreamError, StreamEventType, SYNTHETIC_THOUGHT_SIGNATURE, } from './geminiChat.js';
 import { setSimulate429 } from '../utils/testUtils.js';
-// Mocks
-const mockModelsModule = {
-    generateContent: vi.fn(),
-    generateContentStream: vi.fn(),
-    countTokens: vi.fn(),
-    embedContent: vi.fn(),
-    batchEmbedContents: vi.fn(),
-};
+import { DEFAULT_GEMINI_FLASH_MODEL, DEFAULT_GEMINI_MODEL, DEFAULT_THINKING_MODE, PREVIEW_GEMINI_MODEL, } from '../config/models.js';
+import { AuthType } from './contentGenerator.js';
+import { TerminalQuotaError } from '../utils/googleQuotaErrors.js';
+import { retryWithBackoff } from '../utils/retry.js';
+import { uiTelemetryService } from '../telemetry/uiTelemetry.js';
+// Mock fs module to prevent actual file system operations during tests
+const mockFileSystem = new Map();
+vi.mock('node:fs', () => {
+    const fsModule = {
+        mkdirSync: vi.fn(),
+        writeFileSync: vi.fn((path, data) => {
+            mockFileSystem.set(path, data);
+        }),
+        readFileSync: vi.fn((path) => {
+            if (mockFileSystem.has(path)) {
+                return mockFileSystem.get(path);
+            }
+            throw Object.assign(new Error('ENOENT: no such file or directory'), {
+                code: 'ENOENT',
+            });
+        }),
+        existsSync: vi.fn((path) => mockFileSystem.has(path)),
+    };
+    return {
+        default: fsModule,
+        ...fsModule,
+    };
+});
+const { mockHandleFallback } = vi.hoisted(() => ({
+    mockHandleFallback: vi.fn(),
+}));
+// Add mock for the retry utility
+const { mockRetryWithBackoff } = vi.hoisted(() => ({
+    mockRetryWithBackoff: vi.fn(),
+}));
+vi.mock('../utils/retry.js', () => ({
+    retryWithBackoff: mockRetryWithBackoff,
+}));
+vi.mock('../fallback/handler.js', () => ({
+    handleFallback: mockHandleFallback,
+}));
+const { mockLogContentRetry, mockLogContentRetryFailure } = vi.hoisted(() => ({
+    mockLogContentRetry: vi.fn(),
+    mockLogContentRetryFailure: vi.fn(),
+}));
+vi.mock('../telemetry/loggers.js', () => ({
+    logContentRetry: mockLogContentRetry,
+    logContentRetryFailure: mockLogContentRetryFailure,
+}));
+vi.mock('../telemetry/uiTelemetry.js', () => ({
+    uiTelemetryService: {
+        setLastPromptTokenCount: vi.fn(),
+    },
+}));
 describe('GeminiChat', () => {
+    let mockContentGenerator;
     let chat;
     let mockConfig;
-    const config = {};
     beforeEach(() => {
         vi.clearAllMocks();
+        vi.mocked(uiTelemetryService.setLastPromptTokenCount).mockClear();
+        mockContentGenerator = {
+            generateContent: vi.fn(),
+            generateContentStream: vi.fn(),
+            countTokens: vi.fn(),
+            embedContent: vi.fn(),
+            batchEmbedContents: vi.fn(),
+        };
+        mockHandleFallback.mockClear();
+        // Default mock implementation for tests that don't care about retry logic
+        mockRetryWithBackoff.mockImplementation(async (apiCall) => apiCall());
         mockConfig = {
             getSessionId: () => 'test-session-id',
             getTelemetryLogPromptsEnabled: () => true,
             getUsageStatisticsEnabled: () => true,
             getDebugMode: () => false,
-            getContentGeneratorConfig: () => ({
-                authType: 'oauth-personal',
+            getPreviewFeatures: () => false,
+            getContentGeneratorConfig: vi.fn().mockReturnValue({
+                authType: 'oauth-personal', // Ensure this is set for fallback tests
                 model: 'test-model',
             }),
             getModel: vi.fn().mockReturnValue('gemini-pro'),
             setModel: vi.fn(),
+            isInFallbackMode: vi.fn().mockReturnValue(false),
             getQuotaErrorOccurred: vi.fn().mockReturnValue(false),
             setQuotaErrorOccurred: vi.fn(),
             flashFallbackHandler: undefined,
+            getProjectRoot: vi.fn().mockReturnValue('/test/project/root'),
+            storage: {
+                getProjectTempDir: vi.fn().mockReturnValue('/test/temp'),
+            },
+            getToolRegistry: vi.fn().mockReturnValue({
+                getTool: vi.fn(),
+            }),
+            getContentGenerator: vi.fn().mockReturnValue(mockContentGenerator),
+            getRetryFetchErrors: vi.fn().mockReturnValue(false),
+            modelConfigService: {
+                getResolvedConfig: vi.fn().mockImplementation((modelConfigKey) => {
+                    const thinkingConfig = modelConfigKey.model.startsWith('gemini-3')
+                        ? {
+                            thinkingLevel: ThinkingLevel.HIGH,
+                        }
+                        : {
+                            thinkingBudget: DEFAULT_THINKING_MODE,
+                        };
+                    return {
+                        model: modelConfigKey.model,
+                        generateContentConfig: {
+                            temperature: 0,
+                            thinkingConfig,
+                        },
+                    };
+                }),
+            },
+            isPreviewModelBypassMode: vi.fn().mockReturnValue(false),
+            setPreviewModelBypassMode: vi.fn(),
+            isPreviewModelFallbackMode: vi.fn().mockReturnValue(false),
+            setPreviewModelFallbackMode: vi.fn(),
+            isInteractive: vi.fn().mockReturnValue(false),
         };
         // Disable 429 simulation for tests
         setSimulate429(false);
         // Reset history for each test by creating a new instance
-        chat = new GeminiChat(mockConfig, mockModelsModule, config, []);
+        chat = new GeminiChat(mockConfig);
     });
     afterEach(() => {
         vi.restoreAllMocks();
         vi.resetAllMocks();
     });
-    describe('sendMessage', () => {
-        it('should call generateContent with the correct parameters', async () => {
-            const response = {
-                candidates: [
-                    {
-                        content: {
-                            parts: [{ text: 'response' }],
-                            role: 'model',
-                        },
-                        finishReason: 'STOP',
-                        index: 0,
-                        safetyRatings: [],
-                    },
-                ],
-                text: () => 'response',
-            };
-            vi.mocked(mockModelsModule.generateContent).mockResolvedValue(response);
-            await chat.sendMessage({ message: 'hello' }, 'prompt-id-1');
-            expect(mockModelsModule.generateContent).toHaveBeenCalledWith({
-                model: 'gemini-pro',
-                contents: [{ role: 'user', parts: [{ text: 'hello' }] }],
-                config: {},
-            }, 'prompt-id-1');
+    describe('constructor', () => {
+        it('should initialize lastPromptTokenCount based on history size', () => {
+            const history = [
+                { role: 'user', parts: [{ text: 'Hello' }] },
+                { role: 'model', parts: [{ text: 'Hi there' }] },
+            ];
+            const chatWithHistory = new GeminiChat(mockConfig, '', [], history);
+            const estimatedTokens = Math.ceil(JSON.stringify(history).length / 4);
+            expect(chatWithHistory.getLastPromptTokenCount()).toBe(estimatedTokens);
+        });
+        it('should initialize lastPromptTokenCount for empty history', () => {
+            const chatEmpty = new GeminiChat(mockConfig);
+            expect(chatEmpty.getLastPromptTokenCount()).toBe(Math.ceil(JSON.stringify([]).length / 4));
         });
     });
     describe('sendMessageStream', () => {
-        it('should call generateContentStream with the correct parameters', async () => {
-            const response = (async function* () {
+        it('should succeed if a tool call is followed by an empty part', async () => {
+            // 1. Mock a stream that contains a tool call, then an invalid (empty) part.
+            const streamWithToolCall = (async function* () {
                 yield {
                     candidates: [
                         {
                             content: {
-                                parts: [{ text: 'response' }],
                                 role: 'model',
+                                parts: [{ functionCall: { name: 'test_tool', args: {} } }],
+                            },
+                        },
+                    ],
+                };
+                // This second chunk is invalid according to isValidResponse
+                yield {
+                    candidates: [
+                        {
+                            content: {
+                                role: 'model',
+                                parts: [{ text: '' }],
                             },
-                            finishReason: 'STOP',
-                            index: 0,
-                            safetyRatings: [],
                         },
                     ],
-                    text: () => 'response',
                 };
             })();
-            vi.mocked(mockModelsModule.generateContentStream).mockResolvedValue(response);
-            await chat.sendMessageStream({ message: 'hello' }, 'prompt-id-1');
-            expect(mockModelsModule.generateContentStream).toHaveBeenCalledWith({
-                model: 'gemini-pro',
-                contents: [{ role: 'user', parts: [{ text: 'hello' }] }],
-                config: {},
-            }, 'prompt-id-1');
-        });
-    });
-    describe('recordHistory', () => {
-        const userInput = {
-            role: 'user',
-            parts: [{ text: 'User input' }],
-        };
-        it('should add user input and a single model output to history', () => {
-            const modelOutput = [
-                { role: 'model', parts: [{ text: 'Model output' }] },
-            ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(userInput, modelOutput);
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithToolCall);
+            // 2. Action & Assert: The stream processing should complete without throwing an error
+            // because the presence of a tool call makes the empty final chunk acceptable.
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-tool-call-empty-end', new AbortController().signal);
+            await expect((async () => {
+                for await (const _ of stream) {
+                    /* consume stream */
+                }
+            })()).resolves.not.toThrow();
+            // 3. Verify history was recorded correctly
             const history = chat.getHistory();
-            expect(history).toEqual([userInput, modelOutput[0]]);
+            expect(history.length).toBe(2); // user turn + model turn
+            const modelTurn = history[1];
+            expect(modelTurn?.parts?.length).toBe(1); // The empty part is discarded
+            expect(modelTurn?.parts[0].functionCall).toBeDefined();
         });
-        it('should consolidate adjacent model outputs', () => {
-            const modelOutputParts = [
-                { role: 'model', parts: [{ text: 'Model part 1' }] },
-                { role: 'model', parts: [{ text: 'Model part 2' }] },
-            ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(userInput, modelOutputParts);
-            const history = chat.getHistory();
-            expect(history.length).toBe(2);
-            expect(history[0]).toEqual(userInput);
-            expect(history[1].role).toBe('model');
-            expect(history[1].parts).toEqual([{ text: 'Model part 1Model part 2' }]);
-        });
-        it('should handle a mix of user and model roles in outputContents (though unusual)', () => {
-            const mixedOutput = [
-                { role: 'model', parts: [{ text: 'Model 1' }] },
-                { role: 'user', parts: [{ text: 'Unexpected User' }] }, // This should be pushed as is
-                { role: 'model', parts: [{ text: 'Model 2' }] },
-            ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(userInput, mixedOutput);
-            const history = chat.getHistory();
-            expect(history.length).toBe(4); // user, model1, user_unexpected, model2
-            expect(history[0]).toEqual(userInput);
-            expect(history[1]).toEqual(mixedOutput[0]);
-            expect(history[2]).toEqual(mixedOutput[1]);
-            expect(history[3]).toEqual(mixedOutput[2]);
-        });
-        it('should consolidate multiple adjacent model outputs correctly', () => {
-            const modelOutputParts = [
-                { role: 'model', parts: [{ text: 'M1' }] },
-                { role: 'model', parts: [{ text: 'M2' }] },
-                { role: 'model', parts: [{ text: 'M3' }] },
-            ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(userInput, modelOutputParts);
-            const history = chat.getHistory();
-            expect(history.length).toBe(2);
-            expect(history[1].parts).toEqual([{ text: 'M1M2M3' }]);
+        it('should fail if the stream ends with an empty part and has no finishReason', async () => {
+            // 1. Mock a stream that ends with an invalid part and has no finish reason.
+            const streamWithNoFinish = (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: {
+                                role: 'model',
+                                parts: [{ text: 'Initial content...' }],
+                            },
+                        },
+                    ],
+                };
+                // This second chunk is invalid and has no finishReason, so it should fail.
+                yield {
+                    candidates: [
+                        {
+                            content: {
+                                role: 'model',
+                                parts: [{ text: '' }],
+                            },
+                        },
+                    ],
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithNoFinish);
+            // 2. Action & Assert: The stream should fail because there's no finish reason.
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-no-finish-empty-end', new AbortController().signal);
+            await expect((async () => {
+                for await (const _ of stream) {
+                    /* consume stream */
+                }
+            })()).rejects.toThrow(InvalidStreamError);
         });
-        it('should not consolidate if roles are different between model outputs', () => {
-            const modelOutputParts = [
-                { role: 'model', parts: [{ text: 'M1' }] },
-                { role: 'user', parts: [{ text: 'Interjecting User' }] },
-                { role: 'model', parts: [{ text: 'M2' }] },
-            ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(userInput, modelOutputParts);
+        it('should succeed if the stream ends with an invalid part but has a finishReason and contained a valid part', async () => {
+            // 1. Mock a stream that sends a valid chunk, then an invalid one, but has a finish reason.
+            const streamWithInvalidEnd = (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: {
+                                role: 'model',
+                                parts: [{ text: 'Initial valid content...' }],
+                            },
+                        },
+                    ],
+                };
+                // This second chunk is invalid, but the response has a finishReason.
+                yield {
+                    candidates: [
+                        {
+                            content: {
+                                role: 'model',
+                                parts: [{ text: '' }], // Invalid part
+                            },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithInvalidEnd);
+            // 2. Action & Assert: The stream should complete without throwing an error.
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-valid-then-invalid-end', new AbortController().signal);
+            await expect((async () => {
+                for await (const _ of stream) {
+                    /* consume stream */
+                }
+            })()).resolves.not.toThrow();
+            // 3. Verify history was recorded correctly with only the valid part.
             const history = chat.getHistory();
-            expect(history.length).toBe(4); // user, M1, Interjecting User, M2
-            expect(history[1].parts).toEqual([{ text: 'M1' }]);
-            expect(history[3].parts).toEqual([{ text: 'M2' }]);
-        });
-        it('should merge with last history entry if it is also a model output', () => {
-            // @ts-expect-error Accessing private property for test setup
-            chat.history = [
-                userInput,
-                { role: 'model', parts: [{ text: 'Initial Model Output' }] },
-            ]; // Prime the history
-            const newModelOutput = [
-                { role: 'model', parts: [{ text: 'New Model Part 1' }] },
-                { role: 'model', parts: [{ text: 'New Model Part 2' }] },
-            ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(userInput, newModelOutput); // userInput here is for the *next* turn, but history is already primed
-            // Reset and set up a more realistic scenario for merging with existing history
-            chat = new GeminiChat(mockConfig, mockModelsModule, config, []);
-            const firstUserInput = {
-                role: 'user',
-                parts: [{ text: 'First user input' }],
-            };
-            const firstModelOutput = [
-                { role: 'model', parts: [{ text: 'First model response' }] },
-            ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(firstUserInput, firstModelOutput);
-            const secondUserInput = {
-                role: 'user',
-                parts: [{ text: 'Second user input' }],
-            };
-            const secondModelOutput = [
-                { role: 'model', parts: [{ text: 'Second model response part 1' }] },
-                { role: 'model', parts: [{ text: 'Second model response part 2' }] },
-            ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(secondUserInput, secondModelOutput);
-            const finalHistory = chat.getHistory();
-            expect(finalHistory.length).toBe(4); // user1, model1, user2, model2(consolidated)
-            expect(finalHistory[0]).toEqual(firstUserInput);
-            expect(finalHistory[1]).toEqual(firstModelOutput[0]);
-            expect(finalHistory[2]).toEqual(secondUserInput);
-            expect(finalHistory[3].role).toBe('model');
-            expect(finalHistory[3].parts).toEqual([
-                { text: 'Second model response part 1Second model response part 2' },
-            ]);
+            expect(history.length).toBe(2); // user turn + model turn
+            const modelTurn = history[1];
+            expect(modelTurn?.parts?.length).toBe(1);
+            expect(modelTurn?.parts[0].text).toBe('Initial valid content...');
         });
-        it('should correctly merge consolidated new output with existing model history', () => {
-            // Setup: history ends with a model turn
-            const initialUser = {
-                role: 'user',
-                parts: [{ text: 'Initial user query' }],
-            };
-            const initialModel = {
-                role: 'model',
-                parts: [{ text: 'Initial model answer.' }],
-            };
-            chat = new GeminiChat(mockConfig, mockModelsModule, config, [
-                initialUser,
-                initialModel,
-            ]);
-            // New interaction
-            const currentUserInput = {
-                role: 'user',
-                parts: [{ text: 'Follow-up question' }],
-            };
-            const newModelParts = [
-                { role: 'model', parts: [{ text: 'Part A of new answer.' }] },
-                { role: 'model', parts: [{ text: 'Part B of new answer.' }] },
-            ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(currentUserInput, newModelParts);
+        it('should consolidate subsequent text chunks after receiving an empty text chunk', async () => {
+            // 1. Mock the API to return a stream where one chunk is just an empty text part.
+            const multiChunkStream = (async function* () {
+                yield {
+                    candidates: [
+                        { content: { role: 'model', parts: [{ text: 'Hello' }] } },
+                    ],
+                };
+                // FIX: The original test used { text: '' }, which is invalid.
+                // A chunk can be empty but still valid. This chunk is now removed
+                // as the important part is consolidating what comes after.
+                yield {
+                    candidates: [
+                        {
+                            content: { role: 'model', parts: [{ text: ' World!' }] },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(multiChunkStream);
+            // 2. Action: Send a message and consume the stream.
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-empty-chunk-consolidation', new AbortController().signal);
+            for await (const _ of stream) {
+                // Consume the stream
+            }
+            // 3. Assert: Check that the final history was correctly consolidated.
             const history = chat.getHistory();
-            // Expected: initialUser, initialModel, currentUserInput, consolidatedNewModelParts
-            expect(history.length).toBe(4);
-            expect(history[0]).toEqual(initialUser);
-            expect(history[1]).toEqual(initialModel);
-            expect(history[2]).toEqual(currentUserInput);
-            expect(history[3].role).toBe('model');
-            expect(history[3].parts).toEqual([
-                { text: 'Part A of new answer.Part B of new answer.' },
-            ]);
+            expect(history.length).toBe(2);
+            const modelTurn = history[1];
+            expect(modelTurn?.parts?.length).toBe(1);
+            expect(modelTurn?.parts[0].text).toBe('Hello World!');
         });
-        it('should handle empty modelOutput array', () => {
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(userInput, []);
+        it('should consolidate adjacent text parts that arrive in separate stream chunks', async () => {
+            // 1. Mock the API to return a stream of multiple, adjacent text chunks.
+            const multiChunkStream = (async function* () {
+                yield {
+                    candidates: [
+                        { content: { role: 'model', parts: [{ text: 'This is the ' }] } },
+                    ],
+                };
+                yield {
+                    candidates: [
+                        { content: { role: 'model', parts: [{ text: 'first part.' }] } },
+                    ],
+                };
+                // This function call should break the consolidation.
+                yield {
+                    candidates: [
+                        {
+                            content: {
+                                role: 'model',
+                                parts: [{ functionCall: { name: 'do_stuff', args: {} } }],
+                            },
+                        },
+                    ],
+                };
+                yield {
+                    candidates: [
+                        {
+                            content: {
+                                role: 'model',
+                                parts: [{ text: 'This is the second part.' }],
+                            },
+                        },
+                    ],
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(multiChunkStream);
+            // 2. Action: Send a message and consume the stream.
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-multi-chunk', new AbortController().signal);
+            for await (const _ of stream) {
+                // Consume the stream to trigger history recording.
+            }
+            // 3. Assert: Check that the final history was correctly consolidated.
             const history = chat.getHistory();
-            // If modelOutput is empty, it might push a default empty model part depending on isFunctionResponse
-            // Assuming isFunctionResponse(userInput) is false for this simple text input
+            // The history should contain the user's turn and ONE consolidated model turn.
             expect(history.length).toBe(2);
-            expect(history[0]).toEqual(userInput);
-            expect(history[1].role).toBe('model');
-            expect(history[1].parts).toEqual([]);
-        });
-        it('should handle aggregating modelOutput', () => {
-            const modelOutputUndefinedParts = [
-                { role: 'model', parts: [{ text: 'First model part' }] },
-                { role: 'model', parts: [{ text: 'Second model part' }] },
-                { role: 'model', parts: undefined }, // Test undefined parts
-                { role: 'model', parts: [{ text: 'Third model part' }] },
-                { role: 'model', parts: [] }, // Test empty parts array
-            ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(userInput, modelOutputUndefinedParts);
-            const history = chat.getHistory();
-            expect(history.length).toBe(5);
-            expect(history[0]).toEqual(userInput);
-            expect(history[1].role).toBe('model');
-            expect(history[1].parts).toEqual([
-                { text: 'First model partSecond model part' },
-            ]);
-            expect(history[2].role).toBe('model');
-            expect(history[2].parts).toBeUndefined();
-            expect(history[3].role).toBe('model');
-            expect(history[3].parts).toEqual([{ text: 'Third model part' }]);
-            expect(history[4].role).toBe('model');
-            expect(history[4].parts).toEqual([]);
-        });
-        it('should handle modelOutput with parts being undefined or empty (if they pass initial every check)', () => {
-            const modelOutputUndefinedParts = [
-                { role: 'model', parts: [{ text: 'Text part' }] },
-                { role: 'model', parts: undefined }, // Test undefined parts
-                { role: 'model', parts: [] }, // Test empty parts array
-            ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(userInput, modelOutputUndefinedParts);
-            const history = chat.getHistory();
-            expect(history.length).toBe(4); // userInput, model1 (text), model2 (undefined parts), model3 (empty parts)
-            expect(history[0]).toEqual(userInput);
-            expect(history[1].role).toBe('model');
-            expect(history[1].parts).toEqual([{ text: 'Text part' }]);
-            expect(history[2].role).toBe('model');
-            expect(history[2].parts).toBeUndefined();
-            expect(history[3].role).toBe('model');
-            expect(history[3].parts).toEqual([]);
-        });
-        it('should correctly handle automaticFunctionCallingHistory', () => {
-            const afcHistory = [
-                { role: 'user', parts: [{ text: 'AFC User' }] },
-                { role: 'model', parts: [{ text: 'AFC Model' }] },
-            ];
-            const modelOutput = [
-                { role: 'model', parts: [{ text: 'Regular Model Output' }] },
-            ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(userInput, modelOutput, afcHistory);
-            const history = chat.getHistory();
-            expect(history.length).toBe(3);
-            expect(history[0]).toEqual(afcHistory[0]);
-            expect(history[1]).toEqual(afcHistory[1]);
-            expect(history[2]).toEqual(modelOutput[0]);
-        });
-        it('should add userInput if AFC history is present but empty', () => {
-            const modelOutput = [
-                { role: 'model', parts: [{ text: 'Model Output' }] },
-            ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(userInput, modelOutput, []); // Empty AFC history
+            const modelTurn = history[1];
+            expect(modelTurn.role).toBe('model');
+            // The model turn should have 3 distinct parts: the merged text, the function call, and the final text.
+            expect(modelTurn?.parts?.length).toBe(3);
+            expect(modelTurn?.parts[0].text).toBe('This is the first part.');
+            expect(modelTurn.parts[1].functionCall).toBeDefined();
+            expect(modelTurn.parts[2].text).toBe('This is the second part.');
+        });
+        it('should preserve text parts that stream in the same chunk as a thought', async () => {
+            // 1. Mock the API to return a single chunk containing both a thought and visible text.
+            const mixedContentStream = (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: {
+                                role: 'model',
+                                parts: [
+                                    { thought: 'This is a thought.' },
+                                    { text: 'This is the visible text that should not be lost.' },
+                                ],
+                            },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(mixedContentStream);
+            // 2. Action: Send a message and fully consume the stream to trigger history recording.
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-mixed-chunk', new AbortController().signal);
+            for await (const _ of stream) {
+                // This loop consumes the stream.
+            }
+            // 3. Assert: Check the final state of the history.
             const history = chat.getHistory();
+            // The history should contain two turns: the user's message and the model's response.
             expect(history.length).toBe(2);
-            expect(history[0]).toEqual(userInput);
-            expect(history[1]).toEqual(modelOutput[0]);
+            const modelTurn = history[1];
+            expect(modelTurn.role).toBe('model');
+            // CRUCIAL ASSERTION:
+            // The buggy code would fail here, resulting in parts.length being 0.
+            // The corrected code will pass, preserving the single visible text part.
+            expect(modelTurn?.parts?.length).toBe(1);
+            expect(modelTurn?.parts[0].text).toBe('This is the visible text that should not be lost.');
         });
-        it('should skip "thought" content from modelOutput', () => {
-            const modelOutputWithThought = [
-                { role: 'model', parts: [{ thought: true }, { text: 'Visible text' }] },
-                { role: 'model', parts: [{ text: 'Another visible text' }] },
-            ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(userInput, modelOutputWithThought);
-            const history = chat.getHistory();
-            expect(history.length).toBe(2); // User input + consolidated model output
-            expect(history[0]).toEqual(userInput);
-            expect(history[1].role).toBe('model');
-            // The 'thought' part is skipped, 'Another visible text' becomes the first part.
-            expect(history[1].parts).toEqual([{ text: 'Another visible text' }]);
-        });
-        it('should skip "thought" content even if it is the only content', () => {
-            const modelOutputOnlyThought = [
-                { role: 'model', parts: [{ thought: true }] },
-            ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(userInput, modelOutputOnlyThought);
-            const history = chat.getHistory();
-            expect(history.length).toBe(1); // User input + default empty model part
-            expect(history[0]).toEqual(userInput);
+        it('should use maxAttempts=1 for retryWithBackoff when in Preview Model Fallback Mode', async () => {
+            vi.mocked(mockConfig.isPreviewModelFallbackMode).mockReturnValue(true);
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue((async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { parts: [{ text: 'Success' }] },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })());
+            const stream = await chat.sendMessageStream({ model: PREVIEW_GEMINI_MODEL }, 'test', 'prompt-id-fast-retry', 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((async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { parts: [{ text: 'Success' }] },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })());
+            const stream = await chat.sendMessageStream({ model: DEFAULT_GEMINI_FLASH_MODEL }, 'test', 'prompt-id-normal-retry', new AbortController().signal);
+            for await (const _ of stream) {
+                // consume stream
+            }
+            expect(mockRetryWithBackoff).toHaveBeenCalledWith(expect.any(Function), expect.objectContaining({
+                maxAttempts: undefined, // Should use default
+            }));
+        });
+        it('should pass DEFAULT_GEMINI_MODEL to handleFallback when Preview Model is bypassed (downgraded)', async () => {
+            // ARRANGE
+            vi.mocked(mockConfig.isPreviewModelBypassMode).mockReturnValue(true);
+            // Mock retryWithBackoff to simulate catching the error and calling onPersistent429
+            vi.mocked(retryWithBackoff).mockImplementation(async (apiCall, options) => {
+                const onPersistent429 = options?.onPersistent429;
+                try {
+                    await apiCall();
+                }
+                catch (error) {
+                    if (onPersistent429) {
+                        await onPersistent429(AuthType.LOGIN_WITH_GOOGLE, error);
+                    }
+                    throw error;
+                }
+            });
+            // We need the API call to fail so retryWithBackoff calls the callback.
+            vi.mocked(mockContentGenerator.generateContentStream).mockRejectedValue(new TerminalQuotaError('Simulated Quota Error', {
+                code: 429,
+                message: 'Simulated Quota Error',
+                details: [],
+            }));
+            // ACT
+            const consumeStream = async () => {
+                const stream = await chat.sendMessageStream({ model: PREVIEW_GEMINI_MODEL }, 'test', 'prompt-id-bypass', new AbortController().signal);
+                // Consume the stream to trigger execution
+                for await (const _ of stream) {
+                    // do nothing
+                }
+            };
+            await expect(consumeStream()).rejects.toThrow('Simulated Quota Error');
+            expect(retryWithBackoff).toHaveBeenCalled();
+            // ASSERT
+            // handleFallback is called via onPersistent429Callback
+            // We verify it was called with DEFAULT_GEMINI_MODEL
+            expect(mockHandleFallback).toHaveBeenCalledWith(expect.anything(), DEFAULT_GEMINI_MODEL, // This is the key assertion
+            expect.anything(), expect.anything());
         });
-        it('should correctly consolidate text parts when a thought part is in between', () => {
-            const modelOutputMixed = [
-                { role: 'model', parts: [{ text: 'Part 1.' }] },
+        it('should throw an error when a tool call is followed by an empty stream response', async () => {
+            // 1. Setup: A history where the model has just made a function call.
+            const initialHistory = [
+                {
+                    role: 'user',
+                    parts: [{ text: 'Find a good Italian restaurant for me.' }],
+                },
                 {
                     role: 'model',
-                    parts: [{ thought: true }, { text: 'Should be skipped' }],
+                    parts: [
+                        {
+                            functionCall: {
+                                name: 'find_restaurant',
+                                args: { cuisine: 'Italian' },
+                            },
+                        },
+                    ],
                 },
-                { role: 'model', parts: [{ text: 'Part 2.' }] },
-            ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(userInput, modelOutputMixed);
-            const history = chat.getHistory();
-            expect(history.length).toBe(2);
-            expect(history[0]).toEqual(userInput);
-            expect(history[1].role).toBe('model');
-            expect(history[1].parts).toEqual([{ text: 'Part 1.Part 2.' }]);
-        });
-        it('should handle multiple thought parts correctly', () => {
-            const modelOutputMultipleThoughts = [
-                { role: 'model', parts: [{ thought: true }] },
-                { role: 'model', parts: [{ text: 'Visible 1' }] },
-                { role: 'model', parts: [{ thought: true }] },
-                { role: 'model', parts: [{ text: 'Visible 2' }] },
             ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(userInput, modelOutputMultipleThoughts);
-            const history = chat.getHistory();
-            expect(history.length).toBe(2);
-            expect(history[0]).toEqual(userInput);
-            expect(history[1].role).toBe('model');
-            expect(history[1].parts).toEqual([{ text: 'Visible 1Visible 2' }]);
-        });
-        it('should handle thought part at the end of outputContents', () => {
-            const modelOutputThoughtAtEnd = [
-                { role: 'model', parts: [{ text: 'Visible text' }] },
-                { role: 'model', parts: [{ thought: true }] },
-            ];
-            // @ts-expect-error Accessing private method for testing purposes
-            chat.recordHistory(userInput, modelOutputThoughtAtEnd);
-            const history = chat.getHistory();
-            expect(history.length).toBe(2);
-            expect(history[0]).toEqual(userInput);
-            expect(history[1].role).toBe('model');
-            expect(history[1].parts).toEqual([{ text: 'Visible text' }]);
+            chat.setHistory(initialHistory);
+            // 2. Mock the API to return an empty/thought-only stream.
+            const emptyStreamResponse = (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { role: 'model', parts: [{ thought: true }] },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(emptyStreamResponse);
+            // 3. Action: Send the function response back to the model and consume the stream.
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, {
+                functionResponse: {
+                    name: 'find_restaurant',
+                    response: { name: 'Vesuvio' },
+                },
+            }, 'prompt-id-stream-1', new AbortController().signal);
+            // 4. Assert: The stream processing should throw an InvalidStreamError.
+            await expect((async () => {
+                for await (const _ of stream) {
+                    // This loop consumes the stream to trigger the internal logic.
+                }
+            })()).rejects.toThrow(InvalidStreamError);
         });
-    });
-    describe('addHistory', () => {
-        it('should add a new content item to the history', () => {
-            const newContent = {
-                role: 'user',
-                parts: [{ text: 'A new message' }],
-            };
-            chat.addHistory(newContent);
-            const history = chat.getHistory();
-            expect(history.length).toBe(1);
-            expect(history[0]).toEqual(newContent);
+        it('should succeed when there is a tool call without finish reason', async () => {
+            // Setup: Stream with tool call but no finish reason
+            const streamWithToolCall = (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: {
+                                role: 'model',
+                                parts: [
+                                    {
+                                        functionCall: {
+                                            name: 'test_function',
+                                            args: { param: 'value' },
+                                        },
+                                    },
+                                ],
+                            },
+                            // No finishReason
+                        },
+                    ],
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithToolCall);
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-1', new AbortController().signal);
+            // Should not throw an error
+            await expect((async () => {
+                for await (const _ of stream) {
+                    // consume stream
+                }
+            })()).resolves.not.toThrow();
         });
-        it('should add multiple items correctly', () => {
-            const content1 = {
-                role: 'user',
-                parts: [{ text: 'Message 1' }],
-            };
-            const content2 = {
-                role: 'model',
-                parts: [{ text: 'Message 2' }],
-            };
-            chat.addHistory(content1);
-            chat.addHistory(content2);
-            const history = chat.getHistory();
-            expect(history.length).toBe(2);
-            expect(history[0]).toEqual(content1);
-            expect(history[1]).toEqual(content2);
+        it('should throw InvalidStreamError when no tool call and no finish reason', async () => {
+            // Setup: Stream with text but no finish reason and no tool call
+            const streamWithoutFinishReason = (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: {
+                                role: 'model',
+                                parts: [{ text: 'some response' }],
+                            },
+                            // No finishReason
+                        },
+                    ],
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithoutFinishReason);
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-1', new AbortController().signal);
+            await expect((async () => {
+                for await (const _ of stream) {
+                    // consume stream
+                }
+            })()).rejects.toThrow(InvalidStreamError);
+        });
+        it('should throw InvalidStreamError when no tool call and empty response text', async () => {
+            // Setup: Stream with finish reason but empty response (only thoughts)
+            const streamWithEmptyResponse = (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: {
+                                role: 'model',
+                                parts: [{ thought: 'thinking...' }],
+                            },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithEmptyResponse);
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-1', new AbortController().signal);
+            await expect((async () => {
+                for await (const _ of stream) {
+                    // consume stream
+                }
+            })()).rejects.toThrow(InvalidStreamError);
+        });
+        it('should succeed when there is finish reason and response text', async () => {
+            // Setup: Stream with both finish reason and text content
+            const validStream = (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: {
+                                role: 'model',
+                                parts: [{ text: 'valid response' }],
+                            },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(validStream);
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-1', new AbortController().signal);
+            // Should not throw an error
+            await expect((async () => {
+                for await (const _ of stream) {
+                    // consume stream
+                }
+            })()).resolves.not.toThrow();
+        });
+        it('should throw InvalidStreamError when finishReason is MALFORMED_FUNCTION_CALL', async () => {
+            // Setup: Stream with MALFORMED_FUNCTION_CALL finish reason and empty response
+            const streamWithMalformedFunctionCall = (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: {
+                                role: 'model',
+                                parts: [], // Empty parts
+                            },
+                            finishReason: 'MALFORMED_FUNCTION_CALL',
+                        },
+                    ],
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithMalformedFunctionCall);
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.5-pro' }, 'test', 'prompt-id-malformed', new AbortController().signal);
+            // Should throw an error
+            await expect((async () => {
+                for await (const _ of stream) {
+                    // consume stream
+                }
+            })()).rejects.toThrow(InvalidStreamError);
+        });
+        it('should retry when finishReason is MALFORMED_FUNCTION_CALL', async () => {
+            // 1. Mock the API to fail once with MALFORMED_FUNCTION_CALL, then succeed.
+            vi.mocked(mockContentGenerator.generateContentStream)
+                .mockImplementationOnce(async () => (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { parts: [], role: 'model' },
+                            finishReason: 'MALFORMED_FUNCTION_CALL',
+                        },
+                    ],
+                };
+            })())
+                .mockImplementationOnce(async () => 
+            // Second attempt succeeds
+            (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { parts: [{ text: 'Success after retry' }] },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })());
+            // 2. Send a message
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.5-pro' }, 'test retry', 'prompt-id-retry-malformed', new AbortController().signal);
+            const events = [];
+            for await (const event of stream) {
+                events.push(event);
+            }
+            // 3. Assertions
+            // Should be called twice (initial + retry)
+            expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(2);
+            // Check for a retry event
+            expect(events.some((e) => e.type === StreamEventType.RETRY)).toBe(true);
+            // Check for the successful content chunk
+            expect(events.some((e) => e.type === StreamEventType.CHUNK &&
+                e.value.candidates?.[0]?.content?.parts?.[0]?.text ===
+                    'Success after retry')).toBe(true);
+        });
+        it('should call generateContentStream with the correct parameters', async () => {
+            const response = (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: {
+                                parts: [{ text: 'response' }],
+                                role: 'model',
+                            },
+                            finishReason: 'STOP',
+                            index: 0,
+                            safetyRatings: [],
+                        },
+                    ],
+                    text: () => 'response',
+                    usageMetadata: {
+                        promptTokenCount: 42,
+                        candidatesTokenCount: 15,
+                        totalTokenCount: 57,
+                    },
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(response);
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'hello', 'prompt-id-1', new AbortController().signal);
+            for await (const _ of stream) {
+                // consume stream
+            }
+            expect(mockContentGenerator.generateContentStream).toHaveBeenCalledWith({
+                model: 'test-model',
+                contents: [
+                    {
+                        role: 'user',
+                        parts: [{ text: 'hello' }],
+                    },
+                ],
+                config: {
+                    systemInstruction: '',
+                    tools: [],
+                    temperature: 0,
+                    thinkingConfig: {
+                        thinkingBudget: DEFAULT_THINKING_MODE,
+                    },
+                    abortSignal: expect.any(AbortSignal),
+                },
+            }, 'prompt-id-1');
+        });
+        it('should use thinkingLevel and remove thinkingBudget for gemini-3 models', async () => {
+            const response = (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { parts: [{ text: 'response' }], role: 'model' },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(response);
+            const stream = await chat.sendMessageStream({ model: 'gemini-3-test-only-model-string-for-testing' }, 'hello', 'prompt-id-thinking-level', new AbortController().signal);
+            for await (const _ of stream) {
+                // consume stream
+            }
+            expect(mockContentGenerator.generateContentStream).toHaveBeenCalledWith(expect.objectContaining({
+                model: 'gemini-3-test-only-model-string-for-testing',
+                config: expect.objectContaining({
+                    thinkingConfig: {
+                        thinkingBudget: undefined,
+                        thinkingLevel: ThinkingLevel.HIGH,
+                    },
+                }),
+            }), 'prompt-id-thinking-level');
+        });
+        it('should use thinkingBudget and remove thinkingLevel for non-gemini-3 models', async () => {
+            const response = (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { parts: [{ text: 'response' }], role: 'model' },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(response);
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'hello', 'prompt-id-thinking-budget', new AbortController().signal);
+            for await (const _ of stream) {
+                // consume stream
+            }
+            expect(mockContentGenerator.generateContentStream).toHaveBeenCalledWith(expect.objectContaining({
+                model: 'gemini-2.0-flash',
+                config: expect.objectContaining({
+                    thinkingConfig: {
+                        thinkingBudget: DEFAULT_THINKING_MODE,
+                        thinkingLevel: undefined,
+                    },
+                }),
+            }), 'prompt-id-thinking-budget');
+        });
+    });
+    describe('addHistory', () => {
+        it('should add a new content item to the history', () => {
+            const newContent = {
+                role: 'user',
+                parts: [{ text: 'A new message' }],
+            };
+            chat.addHistory(newContent);
+            const history = chat.getHistory();
+            expect(history.length).toBe(1);
+            expect(history[0]).toEqual(newContent);
+        });
+        it('should add multiple items correctly', () => {
+            const content1 = {
+                role: 'user',
+                parts: [{ text: 'Message 1' }],
+            };
+            const content2 = {
+                role: 'model',
+                parts: [{ text: 'Message 2' }],
+            };
+            chat.addHistory(content1);
+            chat.addHistory(content2);
+            const history = chat.getHistory();
+            expect(history.length).toBe(2);
+            expect(history[0]).toEqual(content1);
+            expect(history[1]).toEqual(content2);
+        });
+    });
+    describe('sendMessageStream with retries', () => {
+        it('should not retry on invalid content if model does not start with gemini-2', async () => {
+            // Mock the stream to fail.
+            vi.mocked(mockContentGenerator.generateContentStream).mockImplementation(async () => (async function* () {
+                yield {
+                    candidates: [{ content: { parts: [{ text: '' }] } }],
+                };
+            })());
+            const stream = await chat.sendMessageStream({ model: 'gemini-1.5-pro' }, 'test', 'prompt-id-no-retry', new AbortController().signal);
+            await expect((async () => {
+                for await (const _ of stream) {
+                    // Must loop to trigger the internal logic that throws.
+                }
+            })()).rejects.toThrow(InvalidStreamError);
+            // Should be called only 1 time (no retry)
+            expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(1);
+            expect(mockLogContentRetry).not.toHaveBeenCalled();
+        });
+        it('should yield a RETRY event when an invalid stream is encountered', async () => {
+            // ARRANGE: Mock the stream to fail once, then succeed.
+            vi.mocked(mockContentGenerator.generateContentStream)
+                .mockImplementationOnce(async () => 
+            // First attempt: An invalid stream with an empty text part.
+            (async function* () {
+                yield {
+                    candidates: [{ content: { parts: [{ text: '' }] } }],
+                };
+            })())
+                .mockImplementationOnce(async () => 
+            // Second attempt (the retry): A minimal valid stream.
+            (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { parts: [{ text: 'Success' }] },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })());
+            // ACT: Send a message and collect all events from the stream.
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-yield-retry', new AbortController().signal);
+            const events = [];
+            for await (const event of stream) {
+                events.push(event);
+            }
+            // ASSERT: Check that a RETRY event was present in the stream's output.
+            const retryEvent = events.find((e) => e.type === StreamEventType.RETRY);
+            expect(retryEvent).toBeDefined();
+            expect(retryEvent?.type).toBe(StreamEventType.RETRY);
+        });
+        it('should retry on invalid content, succeed, and report metrics', async () => {
+            // Use mockImplementationOnce to provide a fresh, promise-wrapped generator for each attempt.
+            vi.mocked(mockContentGenerator.generateContentStream)
+                .mockImplementationOnce(async () => 
+            // First call returns an invalid stream
+            (async function* () {
+                yield {
+                    candidates: [{ content: { parts: [{ text: '' }] } }], // Invalid empty text part
+                };
+            })())
+                .mockImplementationOnce(async () => 
+            // Second call returns a valid stream
+            (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { parts: [{ text: 'Successful response' }] },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })());
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test', 'prompt-id-retry-success', new AbortController().signal);
+            const chunks = [];
+            for await (const chunk of stream) {
+                chunks.push(chunk);
+            }
+            // Assertions
+            expect(mockLogContentRetry).toHaveBeenCalledTimes(1);
+            expect(mockLogContentRetryFailure).not.toHaveBeenCalled();
+            expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(2);
+            // Check for a retry event
+            expect(chunks.some((c) => c.type === StreamEventType.RETRY)).toBe(true);
+            // Check for the successful content chunk
+            expect(chunks.some((c) => c.type === StreamEventType.CHUNK &&
+                c.value.candidates?.[0]?.content?.parts?.[0]?.text ===
+                    'Successful response')).toBe(true);
+            // Check that history was recorded correctly once, with no duplicates.
+            const history = chat.getHistory();
+            expect(history.length).toBe(2);
+            expect(history[0]).toEqual({
+                role: 'user',
+                parts: [{ text: 'test' }],
+            });
+            expect(history[1]).toEqual({
+                role: 'model',
+                parts: [{ text: 'Successful response' }],
+            });
+            // Verify that token counting is not called when usageMetadata is missing
+            expect(uiTelemetryService.setLastPromptTokenCount).not.toHaveBeenCalled();
+        });
+        it('should set temperature to 1 on retry', async () => {
+            // Use mockImplementationOnce to provide a fresh, promise-wrapped generator for each attempt.
+            vi.mocked(mockContentGenerator.generateContentStream)
+                .mockImplementationOnce(async () => 
+            // First call returns an invalid stream
+            (async function* () {
+                yield {
+                    candidates: [{ content: { parts: [{ text: '' }] } }], // Invalid empty text part
+                };
+            })())
+                .mockImplementationOnce(async () => 
+            // Second call returns a valid stream
+            (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { parts: [{ text: 'Successful response' }] },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })());
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-retry-temperature', new AbortController().signal);
+            for await (const _ of stream) {
+                // consume stream
+            }
+            expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(2);
+            // First call should have original temperature
+            expect(mockContentGenerator.generateContentStream).toHaveBeenNthCalledWith(1, expect.objectContaining({
+                config: expect.objectContaining({
+                    temperature: 0,
+                }),
+            }), 'prompt-id-retry-temperature');
+            // Second call (retry) should have temperature 1
+            expect(mockContentGenerator.generateContentStream).toHaveBeenNthCalledWith(2, expect.objectContaining({
+                config: expect.objectContaining({
+                    temperature: 1,
+                }),
+            }), 'prompt-id-retry-temperature');
+        });
+        it('should fail after all retries on persistent invalid content and report metrics', async () => {
+            vi.mocked(mockContentGenerator.generateContentStream).mockImplementation(async () => (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: {
+                                parts: [{ text: '' }],
+                                role: 'model',
+                            },
+                        },
+                    ],
+                };
+            })());
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test', 'prompt-id-retry-fail', new AbortController().signal);
+            await expect(async () => {
+                for await (const _ of stream) {
+                    // Must loop to trigger the internal logic that throws.
+                }
+            }).rejects.toThrow(InvalidStreamError);
+            // Should be called 2 times (initial + 1 retry)
+            expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(2);
+            expect(mockLogContentRetry).toHaveBeenCalledTimes(1);
+            expect(mockLogContentRetryFailure).toHaveBeenCalledTimes(1);
+            // History should still contain the user message.
+            const history = chat.getHistory();
+            expect(history.length).toBe(1);
+            expect(history[0]).toEqual({
+                role: 'user',
+                parts: [{ text: 'test' }],
+            });
+        });
+        describe('API error retry behavior', () => {
+            beforeEach(() => {
+                // Use a more direct mock for retry testing
+                mockRetryWithBackoff.mockImplementation(async (apiCall) => {
+                    try {
+                        return await apiCall();
+                    }
+                    catch (error) {
+                        // Simulate the logic of defaultShouldRetry for ApiError
+                        let shouldRetry = false;
+                        if (error instanceof ApiError && error.message) {
+                            if (error.status === 429 ||
+                                (error.status >= 500 && error.status < 600)) {
+                                shouldRetry = true;
+                            }
+                            // Explicitly don't retry on these
+                            if (error.status === 400) {
+                                shouldRetry = false;
+                            }
+                        }
+                        if (shouldRetry) {
+                            // Try again
+                            return await apiCall();
+                        }
+                        throw error;
+                    }
+                });
+            });
+            it('should not retry on 400 Bad Request errors', async () => {
+                const error400 = new ApiError({ message: 'Bad Request', status: 400 });
+                vi.mocked(mockContentGenerator.generateContentStream).mockRejectedValue(error400);
+                const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-400', new AbortController().signal);
+                await expect((async () => {
+                    for await (const _ of stream) {
+                        /* consume stream */
+                    }
+                })()).rejects.toThrow(error400);
+                // Should only be called once (no retry)
+                expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(1);
+            });
+            it('should retry on 429 Rate Limit errors', async () => {
+                const error429 = new ApiError({ message: 'Rate Limited', status: 429 });
+                vi.mocked(mockContentGenerator.generateContentStream)
+                    .mockRejectedValueOnce(error429)
+                    .mockResolvedValueOnce((async function* () {
+                    yield {
+                        candidates: [
+                            {
+                                content: { parts: [{ text: 'Success after retry' }] },
+                                finishReason: 'STOP',
+                            },
+                        ],
+                    };
+                })());
+                const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-429-retry', new AbortController().signal);
+                const events = [];
+                for await (const event of stream) {
+                    events.push(event);
+                }
+                // Should be called twice (initial + retry)
+                expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(2);
+                // Should have successful content
+                expect(events.some((e) => e.type === StreamEventType.CHUNK &&
+                    e.value.candidates?.[0]?.content?.parts?.[0]?.text ===
+                        'Success after retry')).toBe(true);
+            });
+            it('should retry on 5xx server errors', async () => {
+                const error500 = new ApiError({
+                    message: 'Internal Server Error 500',
+                    status: 500,
+                });
+                vi.mocked(mockContentGenerator.generateContentStream)
+                    .mockRejectedValueOnce(error500)
+                    .mockResolvedValueOnce((async function* () {
+                    yield {
+                        candidates: [
+                            {
+                                content: { parts: [{ text: 'Recovered from 500' }] },
+                                finishReason: 'STOP',
+                            },
+                        ],
+                    };
+                })());
+                const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-500-retry', new AbortController().signal);
+                const events = [];
+                for await (const event of stream) {
+                    events.push(event);
+                }
+                // Should be called twice (initial + retry)
+                expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(2);
+            });
+            it('should retry on specific fetch errors when configured', async () => {
+                vi.mocked(mockConfig.getRetryFetchErrors).mockReturnValue(true);
+                const fetchError = new Error('exception TypeError: fetch failed sending request');
+                vi.mocked(mockContentGenerator.generateContentStream)
+                    .mockRejectedValueOnce(fetchError)
+                    .mockResolvedValueOnce((async function* () {
+                    yield {
+                        candidates: [
+                            {
+                                content: { parts: [{ text: 'Success after fetch error' }] },
+                                finishReason: 'STOP',
+                            },
+                        ],
+                    };
+                })());
+                mockRetryWithBackoff.mockImplementation(async (apiCall, options) => {
+                    try {
+                        return await apiCall();
+                    }
+                    catch (error) {
+                        if (options?.retryFetchErrors &&
+                            error instanceof Error &&
+                            error.message.includes('exception TypeError: fetch failed sending request')) {
+                            return await apiCall();
+                        }
+                        throw error;
+                    }
+                });
+                const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-fetch-error-retry', new AbortController().signal);
+                const events = [];
+                for await (const event of stream) {
+                    events.push(event);
+                }
+                expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(2);
+                expect(events.some((e) => e.type === StreamEventType.CHUNK &&
+                    e.value.candidates?.[0]?.content?.parts?.[0]?.text ===
+                        'Success after fetch error')).toBe(true);
+            });
+            afterEach(() => {
+                // Reset to default behavior
+                mockRetryWithBackoff.mockImplementation(async (apiCall) => apiCall());
+            });
+        });
+    });
+    it('should correctly retry and append to an existing history mid-conversation', async () => {
+        // 1. Setup
+        const initialHistory = [
+            { role: 'user', parts: [{ text: 'First question' }] },
+            { role: 'model', parts: [{ text: 'First answer' }] },
+        ];
+        chat.setHistory(initialHistory);
+        // 2. Mock the API to fail once with an empty stream, then succeed.
+        vi.mocked(mockContentGenerator.generateContentStream)
+            .mockImplementationOnce(async () => (async function* () {
+            yield {
+                candidates: [{ content: { parts: [{ text: '' }] } }],
+            };
+        })())
+            .mockImplementationOnce(async () => 
+        // Second attempt succeeds
+        (async function* () {
+            yield {
+                candidates: [
+                    {
+                        content: { parts: [{ text: 'Second answer' }] },
+                        finishReason: 'STOP',
+                    },
+                ],
+            };
+        })());
+        // 3. Send a new message
+        const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'Second question', 'prompt-id-retry-existing', new AbortController().signal);
+        for await (const _ of stream) {
+            // consume stream
+        }
+        // 4. Assert the final history and metrics
+        const history = chat.getHistory();
+        expect(history.length).toBe(4);
+        // Assert that the correct metrics were reported for one empty-stream retry
+        expect(mockLogContentRetry).toHaveBeenCalledTimes(1);
+        // Explicitly verify the structure of each part to satisfy TypeScript
+        const turn1 = history[0];
+        if (!turn1?.parts?.[0] || !('text' in turn1.parts[0])) {
+            throw new Error('Test setup error: First turn is not a valid text part.');
+        }
+        expect(turn1.parts[0].text).toBe('First question');
+        const turn2 = history[1];
+        if (!turn2?.parts?.[0] || !('text' in turn2.parts[0])) {
+            throw new Error('Test setup error: Second turn is not a valid text part.');
+        }
+        expect(turn2.parts[0].text).toBe('First answer');
+        const turn3 = history[2];
+        if (!turn3?.parts?.[0] || !('text' in turn3.parts[0])) {
+            throw new Error('Test setup error: Third turn is not a valid text part.');
+        }
+        expect(turn3.parts[0].text).toBe('Second question');
+        const turn4 = history[3];
+        if (!turn4?.parts?.[0] || !('text' in turn4.parts[0])) {
+            throw new Error('Test setup error: Fourth turn is not a valid text part.');
+        }
+        expect(turn4.parts[0].text).toBe('Second answer');
+    });
+    it('should retry if the model returns a completely empty stream (no chunks)', async () => {
+        // 1. Mock the API to return an empty stream first, then a valid one.
+        vi.mocked(mockContentGenerator.generateContentStream)
+            .mockImplementationOnce(
+        // First call resolves to an async generator that yields nothing.
+        async () => (async function* () { })())
+            .mockImplementationOnce(
+        // Second call returns a valid stream.
+        async () => (async function* () {
+            yield {
+                candidates: [
+                    {
+                        content: {
+                            parts: [{ text: 'Successful response after empty' }],
+                        },
+                        finishReason: 'STOP',
+                    },
+                ],
+            };
+        })());
+        // 2. Call the method and consume the stream.
+        const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test empty stream', 'prompt-id-empty-stream', new AbortController().signal);
+        const chunks = [];
+        for await (const chunk of stream) {
+            chunks.push(chunk);
+        }
+        // 3. Assert the results.
+        expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(2);
+        expect(chunks.some((c) => c.type === StreamEventType.CHUNK &&
+            c.value.candidates?.[0]?.content?.parts?.[0]?.text ===
+                'Successful response after empty')).toBe(true);
+        const history = chat.getHistory();
+        expect(history.length).toBe(2);
+        // Explicitly verify the structure of each part to satisfy TypeScript
+        const turn1 = history[0];
+        if (!turn1?.parts?.[0] || !('text' in turn1.parts[0])) {
+            throw new Error('Test setup error: First turn is not a valid text part.');
+        }
+        expect(turn1.parts[0].text).toBe('test empty stream');
+        const turn2 = history[1];
+        if (!turn2?.parts?.[0] || !('text' in turn2.parts[0])) {
+            throw new Error('Test setup error: Second turn is not a valid text part.');
+        }
+        expect(turn2.parts[0].text).toBe('Successful response after empty');
+    });
+    it('should queue a subsequent sendMessageStream call until the first stream is fully consumed', async () => {
+        // 1. Create a promise to manually control the stream's lifecycle
+        let continueFirstStream;
+        const firstStreamContinuePromise = new Promise((resolve) => {
+            continueFirstStream = resolve;
+        });
+        // 2. Mock the API to return controllable async generators
+        const firstStreamGenerator = (async function* () {
+            yield {
+                candidates: [
+                    { content: { parts: [{ text: 'first response part 1' }] } },
+                ],
+            };
+            await firstStreamContinuePromise; // Pause the stream
+            yield {
+                candidates: [
+                    {
+                        content: { parts: [{ text: ' part 2' }] },
+                        finishReason: 'STOP',
+                    },
+                ],
+            };
+        })();
+        const secondStreamGenerator = (async function* () {
+            yield {
+                candidates: [
+                    {
+                        content: { parts: [{ text: 'second response' }] },
+                        finishReason: 'STOP',
+                    },
+                ],
+            };
+        })();
+        vi.mocked(mockContentGenerator.generateContentStream)
+            .mockResolvedValueOnce(firstStreamGenerator)
+            .mockResolvedValueOnce(secondStreamGenerator);
+        // 3. Start the first stream and consume only the first chunk to pause it
+        const firstStream = await chat.sendMessageStream({ model: 'test-model' }, 'first', 'prompt-1', new AbortController().signal);
+        const firstStreamIterator = firstStream[Symbol.asyncIterator]();
+        await firstStreamIterator.next();
+        // 4. While the first stream is paused, start the second call. It will block.
+        const secondStreamPromise = chat.sendMessageStream({ model: 'test-model' }, 'second', 'prompt-2', new AbortController().signal);
+        // 5. Assert that only one API call has been made so far.
+        expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(1);
+        // 6. Unblock and fully consume the first stream to completion.
+        continueFirstStream();
+        await firstStreamIterator.next(); // Consume the rest of the stream
+        await firstStreamIterator.next(); // Finish the iterator
+        // 7. Now that the first stream is done, await the second promise to get its generator.
+        const secondStream = await secondStreamPromise;
+        // 8. Start consuming the second stream, which triggers its internal API call.
+        const secondStreamIterator = secondStream[Symbol.asyncIterator]();
+        await secondStreamIterator.next();
+        // 9. The second API call should now have been made.
+        expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(2);
+        // 10. FIX: Fully consume the second stream to ensure recordHistory is called.
+        await secondStreamIterator.next(); // This finishes the iterator.
+        // 11. Final check on history.
+        const history = chat.getHistory();
+        expect(history.length).toBe(4);
+        const turn4 = history[3];
+        if (!turn4?.parts?.[0] || !('text' in turn4.parts[0])) {
+            throw new Error('Test setup error: Fourth turn is not a valid text part.');
+        }
+        expect(turn4.parts[0].text).toBe('second response');
+    });
+    describe('Model Resolution', () => {
+        const mockResponse = {
+            candidates: [
+                {
+                    content: { parts: [{ text: 'response' }], role: 'model' },
+                    finishReason: 'STOP',
+                },
+            ],
+        };
+        it('should use the FLASH model when in fallback mode (sendMessageStream)', async () => {
+            vi.mocked(mockConfig.getModel).mockReturnValue('gemini-pro');
+            vi.mocked(mockConfig.isInFallbackMode).mockReturnValue(true);
+            vi.mocked(mockContentGenerator.generateContentStream).mockImplementation(async () => (async function* () {
+                yield mockResponse;
+            })());
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-res3', new AbortController().signal);
+            for await (const _ of stream) {
+                // consume stream
+            }
+            expect(mockContentGenerator.generateContentStream).toHaveBeenCalledWith(expect.objectContaining({
+                model: DEFAULT_GEMINI_FLASH_MODEL,
+            }), 'prompt-id-res3');
+        });
+    });
+    describe('Fallback Integration (Retries)', () => {
+        const error429 = new ApiError({
+            message: 'API Error 429: Quota exceeded',
+            status: 429,
+        });
+        // Define the simulated behavior for retryWithBackoff for these tests.
+        // This simulation tries the apiCall, if it fails, it calls the callback,
+        // and then tries the apiCall again if the callback returns true.
+        const simulateRetryBehavior = async (apiCall, options) => {
+            try {
+                return await apiCall();
+            }
+            catch (error) {
+                if (options.onPersistent429) {
+                    // We simulate the "persistent" trigger here for simplicity.
+                    const shouldRetry = await options.onPersistent429(options.authType, error);
+                    if (shouldRetry) {
+                        return await apiCall();
+                    }
+                }
+                throw error; // Stop if callback returns false/null or doesn't exist
+            }
+        };
+        beforeEach(() => {
+            mockRetryWithBackoff.mockImplementation(simulateRetryBehavior);
+        });
+        afterEach(() => {
+            mockRetryWithBackoff.mockImplementation(async (apiCall) => apiCall());
+        });
+        it('should call handleFallback with the specific failed model and retry if handler returns true', async () => {
+            const authType = AuthType.LOGIN_WITH_GOOGLE;
+            vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({
+                authType,
+            });
+            const isInFallbackModeSpy = vi.spyOn(mockConfig, 'isInFallbackMode');
+            isInFallbackModeSpy.mockReturnValue(false);
+            vi.mocked(mockContentGenerator.generateContentStream)
+                .mockRejectedValueOnce(error429) // Attempt 1 fails
+                .mockResolvedValueOnce(
+            // Attempt 2 succeeds
+            (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { parts: [{ text: 'Success on retry' }] },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })());
+            mockHandleFallback.mockImplementation(async () => {
+                isInFallbackModeSpy.mockReturnValue(true);
+                return true; // Signal retry
+            });
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'trigger 429', 'prompt-id-fb1', new AbortController().signal);
+            // Consume stream to trigger logic
+            for await (const _ of stream) {
+                // no-op
+            }
+            expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(2);
+            expect(mockHandleFallback).toHaveBeenCalledTimes(1);
+            expect(mockHandleFallback).toHaveBeenCalledWith(mockConfig, 'test-model', authType, error429);
+            const history = chat.getHistory();
+            const modelTurn = history[1];
+            expect(modelTurn.parts[0].text).toBe('Success on retry');
+        });
+        it('should switch to DEFAULT_GEMINI_FLASH_MODEL and use thinkingBudget when falling back from a gemini-3 model', async () => {
+            // ARRANGE
+            const authType = AuthType.LOGIN_WITH_GOOGLE;
+            vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({
+                authType,
+            });
+            // Initial state: Not in fallback mode
+            const isInFallbackModeSpy = vi.spyOn(mockConfig, 'isInFallbackMode');
+            isInFallbackModeSpy.mockReturnValue(false);
+            // Mock API calls:
+            // 1. Fails with 429 (simulating gemini-3 failure)
+            // 2. Succeeds (simulating fallback success)
+            vi.mocked(mockContentGenerator.generateContentStream)
+                .mockRejectedValueOnce(error429)
+                .mockResolvedValueOnce((async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { parts: [{ text: 'Fallback success' }] },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })());
+            // Mock handleFallback to enable fallback mode and signal retry
+            mockHandleFallback.mockImplementation(async () => {
+                isInFallbackModeSpy.mockReturnValue(true); // Next call will see fallback mode = true
+                return true;
+            });
+            // ACT
+            const stream = await chat.sendMessageStream({ model: 'gemini-3-test-model' }, // Start with a gemini-3 model
+            'test fallback thinking', 'prompt-id-fb3', new AbortController().signal);
+            for await (const _ of stream) {
+                // consume stream
+            }
+            // ASSERT
+            expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(2);
+            // First call: gemini-3 model, thinkingLevel set
+            expect(mockContentGenerator.generateContentStream).toHaveBeenNthCalledWith(1, expect.objectContaining({
+                model: 'gemini-3-test-model',
+                config: expect.objectContaining({
+                    thinkingConfig: {
+                        thinkingBudget: undefined,
+                        thinkingLevel: ThinkingLevel.HIGH,
+                    },
+                }),
+            }), 'prompt-id-fb3');
+            // Second call: DEFAULT_GEMINI_FLASH_MODEL (due to fallback), thinkingBudget set (due to fix)
+            expect(mockContentGenerator.generateContentStream).toHaveBeenNthCalledWith(2, expect.objectContaining({
+                model: DEFAULT_GEMINI_FLASH_MODEL,
+                config: expect.objectContaining({
+                    thinkingConfig: {
+                        thinkingBudget: DEFAULT_THINKING_MODE,
+                        thinkingLevel: undefined,
+                    },
+                }),
+            }), 'prompt-id-fb3');
+        });
+        it('should stop retrying if handleFallback returns false (e.g., auth intent)', async () => {
+            vi.mocked(mockConfig.getModel).mockReturnValue('gemini-pro');
+            vi.mocked(mockContentGenerator.generateContentStream).mockRejectedValue(error429);
+            mockHandleFallback.mockResolvedValue(false);
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test stop', 'prompt-id-fb2', new AbortController().signal);
+            await expect((async () => {
+                for await (const _ of stream) {
+                    /* consume stream */
+                }
+            })()).rejects.toThrow(error429);
+            expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(1);
+            expect(mockHandleFallback).toHaveBeenCalledTimes(1);
+        });
+    });
+    it('should discard valid partial content from a failed attempt upon retry', async () => {
+        // Mock the stream to fail on the first attempt after yielding some valid content.
+        vi.mocked(mockContentGenerator.generateContentStream)
+            .mockImplementationOnce(async () => 
+        // First attempt: yields one valid chunk, then one invalid chunk
+        (async function* () {
+            yield {
+                candidates: [
+                    {
+                        content: {
+                            parts: [{ text: 'This valid part should be discarded' }],
+                        },
+                    },
+                ],
+            };
+            yield {
+                candidates: [{ content: { parts: [{ text: '' }] } }], // Invalid chunk triggers retry
+            };
+        })())
+            .mockImplementationOnce(async () => 
+        // Second attempt (the retry): succeeds
+        (async function* () {
+            yield {
+                candidates: [
+                    {
+                        content: {
+                            parts: [{ text: 'Successful final response' }],
+                        },
+                        finishReason: 'STOP',
+                    },
+                ],
+            };
+        })());
+        // Send a message and consume the stream
+        const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-discard-test', new AbortController().signal);
+        const events = [];
+        for await (const event of stream) {
+            events.push(event);
+        }
+        // Check that a retry happened
+        expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(2);
+        expect(events.some((e) => e.type === StreamEventType.RETRY)).toBe(true);
+        // Check the final recorded history
+        const history = chat.getHistory();
+        expect(history.length).toBe(2); // user turn + final model turn
+        const modelTurn = history[1];
+        // The model turn should only contain the text from the successful attempt
+        expect(modelTurn.parts[0].text).toBe('Successful final response');
+        // It should NOT contain any text from the failed attempt
+        expect(modelTurn.parts[0].text).not.toContain('This valid part should be discarded');
+    });
+    describe('stripThoughtsFromHistory', () => {
+        it('should strip thought signatures', () => {
+            chat.setHistory([
+                {
+                    role: 'user',
+                    parts: [{ text: 'hello' }],
+                },
+                {
+                    role: 'model',
+                    parts: [
+                        { text: 'thinking...', thoughtSignature: 'thought-123' },
+                        {
+                            functionCall: { name: 'test', args: {} },
+                            thoughtSignature: 'thought-456',
+                        },
+                    ],
+                },
+            ]);
+            chat.stripThoughtsFromHistory();
+            expect(chat.getHistory()).toEqual([
+                {
+                    role: 'user',
+                    parts: [{ text: 'hello' }],
+                },
+                {
+                    role: 'model',
+                    parts: [
+                        { text: 'thinking...' },
+                        { functionCall: { name: 'test', args: {} } },
+                    ],
+                },
+            ]);
+        });
+    });
+    describe('Preview Model Fallback Logic', () => {
+        it('should reset previewModelBypassMode to false at the start of sendMessageStream', async () => {
+            const stream = (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { role: 'model', parts: [{ text: 'Success' }] },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(stream);
+            await chat.sendMessageStream({ model: 'test-model' }, 'test', 'prompt-id-preview-model-reset', new AbortController().signal);
+            expect(mockConfig.setPreviewModelBypassMode).toHaveBeenCalledWith(false);
+        });
+        it('should reset previewModelFallbackMode to false upon successful Preview Model usage', async () => {
+            const stream = (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { role: 'model', parts: [{ text: 'Success' }] },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(stream);
+            const resultStream = await chat.sendMessageStream({ model: PREVIEW_GEMINI_MODEL }, 'test', 'prompt-id-preview-model-healing', new AbortController().signal);
+            for await (const _ of resultStream) {
+                // consume stream
+            }
+            expect(mockConfig.setPreviewModelFallbackMode).toHaveBeenCalledWith(false);
+        });
+        it('should NOT reset previewModelFallbackMode if Preview Model was bypassed (downgraded)', async () => {
+            const stream = (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { role: 'model', parts: [{ text: 'Success' }] },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(stream);
+            // Simulate bypass mode being active (downgrade happened)
+            vi.mocked(mockConfig.isPreviewModelBypassMode).mockReturnValue(true);
+            const resultStream = await chat.sendMessageStream({ model: PREVIEW_GEMINI_MODEL }, 'test', 'prompt-id-bypass-no-healing', new AbortController().signal);
+            for await (const _ of resultStream) {
+                // consume stream
+            }
+            expect(mockConfig.setPreviewModelFallbackMode).not.toHaveBeenCalled();
+        });
+    });
+    describe('ensureActiveLoopHasThoughtSignatures', () => {
+        it('should add thoughtSignature to the first functionCall in each model turn of the active loop', () => {
+            const chat = new GeminiChat(mockConfig, '', [], []);
+            const history = [
+                { role: 'user', parts: [{ text: 'Old message' }] },
+                {
+                    role: 'model',
+                    parts: [{ functionCall: { name: 'old_tool', args: {} } }],
+                },
+                { role: 'user', parts: [{ text: 'Find a restaurant' }] }, // active loop starts here
+                {
+                    role: 'model',
+                    parts: [
+                        { functionCall: { name: 'find_restaurant', args: {} } }, // This one gets a signature
+                        { functionCall: { name: 'find_restaurant_2', args: {} } }, // This one does NOT
+                    ],
+                },
+                {
+                    role: 'user',
+                    parts: [
+                        { functionResponse: { name: 'find_restaurant', response: {} } },
+                    ],
+                },
+                {
+                    role: 'model',
+                    parts: [
+                        {
+                            functionCall: { name: 'tool_with_sig', args: {} },
+                            thoughtSignature: 'existing-sig',
+                        },
+                        { functionCall: { name: 'another_tool', args: {} } }, // This one does NOT get a signature
+                    ],
+                },
+            ];
+            const newContents = chat.ensureActiveLoopHasThoughtSignatures(history);
+            // Outside active loop - unchanged
+            expect(newContents[1]?.parts?.[0]).not.toHaveProperty('thoughtSignature');
+            // Inside active loop, first model turn
+            // First function call gets a signature
+            expect(newContents[3]?.parts?.[0]?.thoughtSignature).toBe(SYNTHETIC_THOUGHT_SIGNATURE);
+            // Second function call does NOT
+            expect(newContents[3]?.parts?.[1]).not.toHaveProperty('thoughtSignature');
+            // User functionResponse part - unchanged (this is not a model turn)
+            expect(newContents[4]?.parts?.[0]).not.toHaveProperty('thoughtSignature');
+            // Inside active loop, second model turn
+            // First function call already has a signature, so nothing changes
+            expect(newContents[5]?.parts?.[0]?.thoughtSignature).toBe('existing-sig');
+            // Second function call does NOT get a signature
+            expect(newContents[5]?.parts?.[1]).not.toHaveProperty('thoughtSignature');
+        });
+        it('should not modify contents if there is no user text message', () => {
+            const chat = new GeminiChat(mockConfig, '', [], []);
+            const history = [
+                {
+                    role: 'user',
+                    parts: [{ functionResponse: { name: 'tool1', response: {} } }],
+                },
+                {
+                    role: 'model',
+                    parts: [{ functionCall: { name: 'tool2', args: {} } }],
+                },
+            ];
+            const newContents = chat.ensureActiveLoopHasThoughtSignatures(history);
+            expect(newContents).toEqual(history);
+            expect(newContents[1]?.parts?.[0]).not.toHaveProperty('thoughtSignature');
+        });
+        it('should handle an empty history', () => {
+            const chat = new GeminiChat(mockConfig, '', []);
+            const history = [];
+            const newContents = chat.ensureActiveLoopHasThoughtSignatures(history);
+            expect(newContents).toEqual([]);
+        });
+        it('should handle history with only a user message', () => {
+            const chat = new GeminiChat(mockConfig, '', []);
+            const history = [{ role: 'user', parts: [{ text: 'Hello' }] }];
+            const newContents = chat.ensureActiveLoopHasThoughtSignatures(history);
+            expect(newContents).toEqual(history);
         });
     });
 });