@google/gemini-cli-core 0.6.0-nightly.20250910.a31830a3 → 0.6.0-preview.0

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

@@ -105,206 +105,6 @@ describe('GeminiChat', () => {
         vi.restoreAllMocks();
         vi.resetAllMocks();
     });
-    describe('sendMessage', () => {
-        it('should retain the initial user message when an automatic function call occurs', async () => {
-            // 1. Define the user's initial text message. This is the turn that gets dropped by the buggy logic.
-            const userInitialMessage = {
-                role: 'user',
-                parts: [{ text: 'How is the weather in Boston?' }],
-            };
-            // 2. Mock the full API response, including the automaticFunctionCallingHistory.
-            // This history represents the full turn: user asks, model calls tool, tool responds, model answers.
-            const mockAfcResponse = {
-                candidates: [
-                    {
-                        content: {
-                            role: 'model',
-                            parts: [
-                                { text: 'The weather in Boston is 72 degrees and sunny.' },
-                            ],
-                        },
-                    },
-                ],
-                automaticFunctionCallingHistory: [
-                    userInitialMessage, // The user's turn
-                    {
-                        // The model's first response: a tool call
-                        role: 'model',
-                        parts: [
-                            {
-                                functionCall: {
-                                    name: 'get_weather',
-                                    args: { location: 'Boston' },
-                                },
-                            },
-                        ],
-                    },
-                    {
-                        // The tool's response, which has a 'user' role
-                        role: 'user',
-                        parts: [
-                            {
-                                functionResponse: {
-                                    name: 'get_weather',
-                                    response: { temperature: 72, condition: 'sunny' },
-                                },
-                            },
-                        ],
-                    },
-                ],
-            };
-            vi.mocked(mockContentGenerator.generateContent).mockResolvedValue(mockAfcResponse);
-            // 3. Action: Send the initial message.
-            await chat.sendMessage({ message: 'How is the weather in Boston?' }, 'prompt-id-afc-bug');
-            // 4. Assert: Check the final state of the history.
-            const history = chat.getHistory();
-            // With the bug, history.length will be 3, because the first user message is dropped.
-            // The correct behavior is for the history to contain all 4 turns.
-            expect(history.length).toBe(4);
-            // Crucially, assert that the very first turn in the history matches the user's initial message.
-            // This is the assertion that will fail.
-            const firstTurn = history[0];
-            expect(firstTurn.role).toBe('user');
-            expect(firstTurn?.parts[0].text).toBe('How is the weather in Boston?');
-            // Verify the rest of the history is also correct.
-            const secondTurn = history[1];
-            expect(secondTurn.role).toBe('model');
-            expect(secondTurn?.parts[0].functionCall).toBeDefined();
-            const thirdTurn = history[2];
-            expect(thirdTurn.role).toBe('user');
-            expect(thirdTurn?.parts[0].functionResponse).toBeDefined();
-            const fourthTurn = history[3];
-            expect(fourthTurn.role).toBe('model');
-            expect(fourthTurn?.parts[0].text).toContain('72 degrees and sunny');
-        });
-        it('should throw an error when attempting to add a user turn after another user turn', async () => {
-            // 1. Setup: Create a history that already ends with a user turn (a functionResponse).
-            const initialHistory = [
-                { role: 'user', parts: [{ text: 'Initial prompt' }] },
-                {
-                    role: 'model',
-                    parts: [{ functionCall: { name: 'test_tool', args: {} } }],
-                },
-                {
-                    role: 'user',
-                    parts: [{ functionResponse: { name: 'test_tool', response: {} } }],
-                },
-            ];
-            chat.setHistory(initialHistory);
-            // 2. Mock a valid model response so the call doesn't fail for other reasons.
-            const mockResponse = {
-                candidates: [
-                    { content: { role: 'model', parts: [{ text: 'some response' }] } },
-                ],
-            };
-            vi.mocked(mockContentGenerator.generateContent).mockResolvedValue(mockResponse);
-            // 3. Action & Assert: Expect that sending another user message immediately
-            //    after a user-role turn throws the specific error.
-            await expect(chat.sendMessage({ message: 'This is an invalid consecutive user message' }, 'prompt-id-1')).rejects.toThrow('Cannot add a user turn after another user turn.');
-        });
-        it('should preserve text parts that are in the same response as a thought', async () => {
-            // 1. Mock the API to return a single response containing both a thought and visible text.
-            const mixedContentResponse = {
-                candidates: [
-                    {
-                        content: {
-                            role: 'model',
-                            parts: [
-                                { thought: 'This is a thought.' },
-                                { text: 'This is the visible text that should not be lost.' },
-                            ],
-                        },
-                    },
-                ],
-            };
-            vi.mocked(mockContentGenerator.generateContent).mockResolvedValue(mixedContentResponse);
-            // 2. Action: Send a standard, non-streaming message.
-            await chat.sendMessage({ message: 'test message' }, 'prompt-id-mixed-response');
-            // 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);
-            const modelTurn = history[1];
-            expect(modelTurn.role).toBe('model');
-            // CRUCIAL ASSERTION:
-            // Buggy code would discard the entire response because a "thought" was present,
-            // resulting in an empty placeholder turn with 0 parts.
-            // 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 add a placeholder model turn when a tool call is followed by an empty model 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: [
-                        {
-                            functionCall: {
-                                name: 'find_restaurant',
-                                args: { cuisine: 'Italian' },
-                            },
-                        },
-                    ],
-                },
-            ];
-            chat.setHistory(initialHistory);
-            // 2. Mock the API to return an empty/thought-only response.
-            const emptyModelResponse = {
-                candidates: [
-                    { content: { role: 'model', parts: [{ thought: true }] } },
-                ],
-            };
-            vi.mocked(mockContentGenerator.generateContent).mockResolvedValue(emptyModelResponse);
-            // 3. Action: Send the function response back to the model.
-            await chat.sendMessage({
-                message: {
-                    functionResponse: {
-                        name: 'find_restaurant',
-                        response: { name: 'Vesuvio' },
-                    },
-                },
-            }, 'prompt-id-1');
-            // 4. Assert: The history should now have four valid, alternating turns.
-            const history = chat.getHistory();
-            expect(history.length).toBe(4);
-            // The final turn must be the empty model placeholder.
-            const lastTurn = history[3];
-            expect(lastTurn.role).toBe('model');
-            expect(lastTurn?.parts?.length).toBe(0);
-            // The second-to-last turn must be the function response we sent.
-            const secondToLastTurn = history[2];
-            expect(secondToLastTurn.role).toBe('user');
-            expect(secondToLastTurn?.parts[0].functionResponse).toBeDefined();
-        });
-        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(mockContentGenerator.generateContent).mockResolvedValue(response);
-            await chat.sendMessage({ message: 'hello' }, 'prompt-id-1');
-            expect(mockContentGenerator.generateContent).toHaveBeenCalledWith({
-                model: 'gemini-pro',
-                contents: [{ role: 'user', parts: [{ text: 'hello' }] }],
-                config: {},
-            }, 'prompt-id-1');
-        });
-    });
     describe('sendMessageStream', () => {
         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.
@@ -334,7 +134,7 @@ describe('GeminiChat', () => {
             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({ message: 'test message' }, 'prompt-id-tool-call-empty-end');
+            const stream = await chat.sendMessageStream('test-model', { message: 'test message' }, 'prompt-id-tool-call-empty-end');
             await expect((async () => {
                 for await (const _ of stream) {
                     /* consume stream */
@@ -374,7 +174,7 @@ describe('GeminiChat', () => {
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithNoFinish);
             // 2. Action & Assert: The stream should fail because there's no finish reason.
-            const stream = await chat.sendMessageStream({ message: 'test message' }, 'prompt-id-no-finish-empty-end');
+            const stream = await chat.sendMessageStream('test-model', { message: 'test message' }, 'prompt-id-no-finish-empty-end');
             await expect((async () => {
                 for await (const _ of stream) {
                     /* consume stream */
@@ -409,7 +209,7 @@ describe('GeminiChat', () => {
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithInvalidEnd);
             // 2. Action & Assert: The stream should complete without throwing an error.
-            const stream = await chat.sendMessageStream({ message: 'test message' }, 'prompt-id-valid-then-invalid-end');
+            const stream = await chat.sendMessageStream('test-model', { message: 'test message' }, 'prompt-id-valid-then-invalid-end');
             await expect((async () => {
                 for await (const _ of stream) {
                     /* consume stream */
@@ -422,55 +222,6 @@ describe('GeminiChat', () => {
             expect(modelTurn?.parts?.length).toBe(1);
             expect(modelTurn?.parts[0].text).toBe('Initial valid content...');
         });
-        it('should not consolidate text into a part that also contains a functionCall', async () => {
-            // 1. Mock the API to stream a malformed part followed by a valid text part.
-            const multiChunkStream = (async function* () {
-                // This malformed part has both text and a functionCall.
-                yield {
-                    candidates: [
-                        {
-                            content: {
-                                role: 'model',
-                                parts: [
-                                    {
-                                        text: 'Some text',
-                                        functionCall: { name: 'do_stuff', args: {} },
-                                    },
-                                ],
-                            },
-                        },
-                    ],
-                };
-                // This valid text part should NOT be merged into the malformed one.
-                yield {
-                    candidates: [
-                        {
-                            content: {
-                                role: 'model',
-                                parts: [{ text: ' that should not be merged.' }],
-                            },
-                        },
-                    ],
-                };
-            })();
-            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(multiChunkStream);
-            // 2. Action: Send a message and consume the stream.
-            const stream = await chat.sendMessageStream({ message: 'test message' }, 'prompt-id-malformed-chunk');
-            for await (const _ of stream) {
-                // Consume the stream to trigger history recording.
-            }
-            // 3. Assert: Check that the final history was not incorrectly consolidated.
-            const history = chat.getHistory();
-            expect(history.length).toBe(2);
-            const modelTurn = history[1];
-            // CRUCIAL ASSERTION: There should be two separate parts.
-            // The old, non-strict logic would incorrectly merge them, resulting in one part.
-            expect(modelTurn?.parts?.length).toBe(2);
-            // Verify the contents of each part.
-            expect(modelTurn?.parts[0].text).toBe('Some text');
-            expect(modelTurn?.parts[0].functionCall).toBeDefined();
-            expect(modelTurn?.parts[1].text).toBe(' that should not be merged.');
-        });
         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* () {
@@ -493,7 +244,7 @@ describe('GeminiChat', () => {
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(multiChunkStream);
             // 2. Action: Send a message and consume the stream.
-            const stream = await chat.sendMessageStream({ message: 'test message' }, 'prompt-id-empty-chunk-consolidation');
+            const stream = await chat.sendMessageStream('test-model', { message: 'test message' }, 'prompt-id-empty-chunk-consolidation');
             for await (const _ of stream) {
                 // Consume the stream
             }
@@ -541,7 +292,7 @@ describe('GeminiChat', () => {
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(multiChunkStream);
             // 2. Action: Send a message and consume the stream.
-            const stream = await chat.sendMessageStream({ message: 'test message' }, 'prompt-id-multi-chunk');
+            const stream = await chat.sendMessageStream('test-model', { message: 'test message' }, 'prompt-id-multi-chunk');
             for await (const _ of stream) {
                 // Consume the stream to trigger history recording.
             }
@@ -577,7 +328,7 @@ describe('GeminiChat', () => {
             })();
             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({ message: 'test message' }, 'prompt-id-mixed-chunk');
+            const stream = await chat.sendMessageStream('test-model', { message: 'test message' }, 'prompt-id-mixed-chunk');
             for await (const _ of stream) {
                 // This loop consumes the stream.
             }
@@ -626,7 +377,7 @@ describe('GeminiChat', () => {
             })();
             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({
+            const stream = await chat.sendMessageStream('test-model', {
                 message: {
                     functionResponse: {
                         name: 'find_restaurant',
@@ -667,146 +418,22 @@ describe('GeminiChat', () => {
                 };
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(response);
-            const stream = await chat.sendMessageStream({ message: 'hello' }, 'prompt-id-1');
+            const stream = await chat.sendMessageStream('test-model', { message: 'hello' }, 'prompt-id-1');
             for await (const _ of stream) {
-                // consume stream to trigger internal logic
+                // consume stream
             }
             expect(mockContentGenerator.generateContentStream).toHaveBeenCalledWith({
-                model: 'gemini-pro',
-                contents: [{ role: 'user', parts: [{ text: 'hello' }] }],
+                model: 'test-model',
+                contents: [
+                    {
+                        role: 'user',
+                        parts: [{ text: 'hello' }],
+                    },
+                ],
                 config: {},
             }, 'prompt-id-1');
         });
     });
-    describe('recordHistory', () => {
-        const userInput = {
-            role: 'user',
-            parts: [{ text: 'User input' }],
-        };
-        it('should consolidate all consecutive model turns into a single turn', () => {
-            const userInput = {
-                role: 'user',
-                parts: [{ text: 'User input' }],
-            };
-            // This simulates a multi-part model response with different part types.
-            const modelOutput = [
-                { role: 'model', parts: [{ text: 'Thinking...' }] },
-                {
-                    role: 'model',
-                    parts: [{ functionCall: { name: 'do_stuff', args: {} } }],
-                },
-            ];
-            // @ts-expect-error Accessing private method for testing
-            chat.recordHistory(userInput, modelOutput);
-            const history = chat.getHistory();
-            // The history should contain the user's turn and ONE consolidated model turn.
-            // The old code would fail here, resulting in a length of 3.
-            //expect(history).toBe([]);
-            expect(history.length).toBe(2);
-            const modelTurn = history[1];
-            expect(modelTurn.role).toBe('model');
-            // The consolidated turn should contain both the text part and the functionCall part.
-            expect(modelTurn?.parts?.length).toBe(2);
-            expect(modelTurn?.parts[0].text).toBe('Thinking...');
-            expect(modelTurn?.parts[1].functionCall).toBeDefined();
-        });
-        it('should add a placeholder model turn when a tool call is followed by an empty response', () => {
-            // 1. Setup: A history where the model has just made a function call.
-            const initialHistory = [
-                { role: 'user', parts: [{ text: 'Initial prompt' }] },
-                {
-                    role: 'model',
-                    parts: [{ functionCall: { name: 'test_tool', args: {} } }],
-                },
-            ];
-            chat.setHistory(initialHistory);
-            // 2. Action: The user provides the tool's response, and the model's
-            // final output is empty (e.g., just a thought, which gets filtered out).
-            const functionResponse = {
-                role: 'user',
-                parts: [{ functionResponse: { name: 'test_tool', response: {} } }],
-            };
-            const emptyModelOutput = [];
-            // @ts-expect-error Accessing private method for testing
-            chat.recordHistory(functionResponse, emptyModelOutput, [
-                functionResponse,
-            ]);
-            // 3. Assert: The history should now have four valid, alternating turns.
-            const history = chat.getHistory();
-            expect(history.length).toBe(4);
-            // The final turn must be the empty model placeholder.
-            const lastTurn = history[3];
-            expect(lastTurn.role).toBe('model');
-            expect(lastTurn?.parts?.length).toBe(0);
-            // The second-to-last turn must be the function response we provided.
-            const secondToLastTurn = history[2];
-            expect(secondToLastTurn.role).toBe('user');
-            expect(secondToLastTurn?.parts[0].functionResponse).toBeDefined();
-        });
-        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
-            chat.recordHistory(userInput, modelOutput);
-            const history = chat.getHistory();
-            expect(history.length).toBe(2);
-            expect(history[0]).toEqual(userInput);
-            expect(history[1]).toEqual(modelOutput[0]);
-        });
-        it('should consolidate adjacent text parts from multiple content objects', () => {
-            const modelOutput = [
-                { role: 'model', parts: [{ text: 'Part 1.' }] },
-                { role: 'model', parts: [{ text: ' Part 2.' }] },
-                { role: 'model', parts: [{ text: ' Part 3.' }] },
-            ];
-            // @ts-expect-error Accessing private method for testing
-            chat.recordHistory(userInput, modelOutput);
-            const history = chat.getHistory();
-            expect(history.length).toBe(2);
-            expect(history[1].role).toBe('model');
-            expect(history[1].parts).toEqual([{ text: 'Part 1. Part 2. Part 3.' }]);
-        });
-        it('should add an empty placeholder turn if modelOutput is empty', () => {
-            // This simulates receiving a pre-filtered, thought-only response.
-            const emptyModelOutput = [];
-            // @ts-expect-error Accessing private method for testing
-            chat.recordHistory(userInput, emptyModelOutput);
-            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([]);
-        });
-        it('should preserve model outputs with undefined or empty parts arrays', () => {
-            const malformedOutput = [
-                { role: 'model', parts: [{ text: 'Text part' }] },
-                { role: 'model', parts: undefined },
-                { role: 'model', parts: [] },
-            ];
-            // @ts-expect-error Accessing private method for testing
-            chat.recordHistory(userInput, malformedOutput);
-            const history = chat.getHistory();
-            expect(history.length).toBe(4); // userInput + 3 model turns
-            expect(history[1].parts).toEqual([{ text: 'Text part' }]);
-            expect(history[2].parts).toBeUndefined();
-            expect(history[3].parts).toEqual([]);
-        });
-        it('should not consolidate content with different roles', () => {
-            const mixedOutput = [
-                { role: 'model', parts: [{ text: 'Model 1' }] },
-                { role: 'user', parts: [{ text: 'Unexpected User' }] },
-                { role: 'model', parts: [{ text: 'Model 2' }] },
-            ];
-            // @ts-expect-error Accessing private method for testing
-            chat.recordHistory(userInput, mixedOutput);
-            const history = chat.getHistory();
-            expect(history.length).toBe(4); // userInput, model1, unexpected_user, model2
-            expect(history[1]).toEqual(mixedOutput[0]);
-            expect(history[2]).toEqual(mixedOutput[1]);
-            expect(history[3]).toEqual(mixedOutput[2]);
-        });
-    });
     describe('addHistory', () => {
         it('should add a new content item to the history', () => {
             const newContent = {
@@ -859,7 +486,7 @@ describe('GeminiChat', () => {
                 };
             })());
             // ACT: Send a message and collect all events from the stream.
-            const stream = await chat.sendMessageStream({ message: 'test' }, 'prompt-id-yield-retry');
+            const stream = await chat.sendMessageStream('test-model', { message: 'test' }, 'prompt-id-yield-retry');
             const events = [];
             for await (const event of stream) {
                 events.push(event);
@@ -891,7 +518,7 @@ describe('GeminiChat', () => {
                     ],
                 };
             })());
-            const stream = await chat.sendMessageStream({ message: 'test' }, 'prompt-id-retry-success');
+            const stream = await chat.sendMessageStream('test-model', { message: 'test' }, 'prompt-id-retry-success');
             const chunks = [];
             for await (const chunk of stream) {
                 chunks.push(chunk);
@@ -932,14 +559,12 @@ describe('GeminiChat', () => {
                     ],
                 };
             })());
-            // This helper function consumes the stream and allows us to test for rejection.
-            async function consumeStreamAndExpectError() {
-                const stream = await chat.sendMessageStream({ message: 'test' }, 'prompt-id-retry-fail');
+            const stream = await chat.sendMessageStream('test-model', { message: 'test' }, 'prompt-id-retry-fail');
+            await expect(async () => {
                 for await (const _ of stream) {
                     // Must loop to trigger the internal logic that throws.
                 }
-            }
-            await expect(consumeStreamAndExpectError()).rejects.toThrow(EmptyStreamError);
+            }).rejects.toThrow(EmptyStreamError);
             // Should be called 3 times (initial + 2 retries)
             expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(3);
             expect(mockLogInvalidChunk).toHaveBeenCalledTimes(3);
@@ -977,7 +602,7 @@ describe('GeminiChat', () => {
             };
         })());
         // 3. Send a new message
-        const stream = await chat.sendMessageStream({ message: 'Second question' }, 'prompt-id-retry-existing');
+        const stream = await chat.sendMessageStream('test-model', { message: 'Second question' }, 'prompt-id-retry-existing');
         for await (const _ of stream) {
             // consume stream
         }
@@ -1008,60 +633,6 @@ describe('GeminiChat', () => {
         }
         expect(turn4.parts[0].text).toBe('Second answer');
     });
-    describe('concurrency control', () => {
-        it('should queue a subsequent sendMessage call until the first one completes', async () => {
-            // 1. Create promises to manually control when the API calls resolve
-            let firstCallResolver;
-            const firstCallPromise = new Promise((resolve) => {
-                firstCallResolver = resolve;
-            });
-            let secondCallResolver;
-            const secondCallPromise = new Promise((resolve) => {
-                secondCallResolver = resolve;
-            });
-            // A standard response body for the mock
-            const mockResponse = {
-                candidates: [
-                    {
-                        content: { parts: [{ text: 'response' }], role: 'model' },
-                    },
-                ],
-            };
-            // 2. Mock the API to return our controllable promises in order
-            vi.mocked(mockContentGenerator.generateContent)
-                .mockReturnValueOnce(firstCallPromise)
-                .mockReturnValueOnce(secondCallPromise);
-            // 3. Start the first message call. Do not await it yet.
-            const firstMessagePromise = chat.sendMessage({ message: 'first' }, 'prompt-1');
-            // Give the event loop a chance to run the async call up to the `await`
-            await new Promise(process.nextTick);
-            // 4. While the first call is "in-flight", start the second message call.
-            const secondMessagePromise = chat.sendMessage({ message: 'second' }, 'prompt-2');
-            // 5. CRUCIAL CHECK: At this point, only the first API call should have been made.
-            // The second call should be waiting on `sendPromise`.
-            expect(mockContentGenerator.generateContent).toHaveBeenCalledTimes(1);
-            expect(mockContentGenerator.generateContent).toHaveBeenCalledWith(expect.objectContaining({
-                contents: expect.arrayContaining([
-                    expect.objectContaining({ parts: [{ text: 'first' }] }),
-                ]),
-            }), 'prompt-1');
-            // 6. Unblock the first API call and wait for the first message to fully complete.
-            firstCallResolver(mockResponse);
-            await firstMessagePromise;
-            // Give the event loop a chance to unblock and run the second call.
-            await new Promise(process.nextTick);
-            // 7. CRUCIAL CHECK: Now, the second API call should have been made.
-            expect(mockContentGenerator.generateContent).toHaveBeenCalledTimes(2);
-            expect(mockContentGenerator.generateContent).toHaveBeenCalledWith(expect.objectContaining({
-                contents: expect.arrayContaining([
-                    expect.objectContaining({ parts: [{ text: 'second' }] }),
-                ]),
-            }), 'prompt-2');
-            // 8. Clean up by resolving the second call.
-            secondCallResolver(mockResponse);
-            await secondMessagePromise;
-        });
-    });
     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)
@@ -1083,7 +654,7 @@ describe('GeminiChat', () => {
             };
         })());
         // 2. Call the method and consume the stream.
-        const stream = await chat.sendMessageStream({ message: 'test empty stream' }, 'prompt-id-empty-stream');
+        const stream = await chat.sendMessageStream('test-model', { message: 'test empty stream' }, 'prompt-id-empty-stream');
         const chunks = [];
         for await (const chunk of stream) {
             chunks.push(chunk);
@@ -1144,11 +715,11 @@ describe('GeminiChat', () => {
             .mockResolvedValueOnce(firstStreamGenerator)
             .mockResolvedValueOnce(secondStreamGenerator);
         // 3. Start the first stream and consume only the first chunk to pause it
-        const firstStream = await chat.sendMessageStream({ message: 'first' }, 'prompt-1');
+        const firstStream = await chat.sendMessageStream('test-model', { message: 'first' }, 'prompt-1');
         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({ message: 'second' }, 'prompt-2');
+        const secondStreamPromise = chat.sendMessageStream('test-model', { message: 'second' }, 'prompt-2');
         // 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.
@@ -1182,31 +753,13 @@ describe('GeminiChat', () => {
                 },
             ],
         };
-        it('should use the configured model when not in fallback mode (sendMessage)', async () => {
-            vi.mocked(mockConfig.getModel).mockReturnValue('gemini-2.5-pro');
-            vi.mocked(mockConfig.isInFallbackMode).mockReturnValue(false);
-            vi.mocked(mockContentGenerator.generateContent).mockResolvedValue(mockResponse);
-            await chat.sendMessage({ message: 'test' }, 'prompt-id-res1');
-            expect(mockContentGenerator.generateContent).toHaveBeenCalledWith(expect.objectContaining({
-                model: 'gemini-2.5-pro',
-            }), 'prompt-id-res1');
-        });
-        it('should use the FLASH model when in fallback mode (sendMessage)', async () => {
-            vi.mocked(mockConfig.getModel).mockReturnValue('gemini-2.5-pro');
-            vi.mocked(mockConfig.isInFallbackMode).mockReturnValue(true);
-            vi.mocked(mockContentGenerator.generateContent).mockResolvedValue(mockResponse);
-            await chat.sendMessage({ message: 'test' }, 'prompt-id-res2');
-            expect(mockContentGenerator.generateContent).toHaveBeenCalledWith(expect.objectContaining({
-                model: DEFAULT_GEMINI_FLASH_MODEL,
-            }), 'prompt-id-res2');
-        });
         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({ message: 'test' }, 'prompt-id-res3');
+            const stream = await chat.sendMessageStream('test-model', { message: 'test' }, 'prompt-id-res3');
             for await (const _ of stream) {
                 // consume stream
             }
@@ -1244,37 +797,53 @@ describe('GeminiChat', () => {
             mockRetryWithBackoff.mockImplementation(async (apiCall) => apiCall());
         });
         it('should call handleFallback with the specific failed model and retry if handler returns true', async () => {
-            const FAILED_MODEL = 'gemini-2.5-pro';
-            vi.mocked(mockConfig.getModel).mockReturnValue(FAILED_MODEL);
             const authType = AuthType.LOGIN_WITH_GOOGLE;
             vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({
                 authType,
-                model: FAILED_MODEL,
             });
             const isInFallbackModeSpy = vi.spyOn(mockConfig, 'isInFallbackMode');
             isInFallbackModeSpy.mockReturnValue(false);
-            vi.mocked(mockContentGenerator.generateContent)
+            vi.mocked(mockContentGenerator.generateContentStream)
                 .mockRejectedValueOnce(error429) // Attempt 1 fails
-                .mockResolvedValueOnce({
-                candidates: [{ content: { parts: [{ text: 'Success on retry' }] } }],
-            }); // Attempt 2 succeeds
+                .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 result = await chat.sendMessage({ message: 'trigger 429' }, 'prompt-id-fb1');
-            expect(mockRetryWithBackoff).toHaveBeenCalledTimes(1);
-            expect(mockContentGenerator.generateContent).toHaveBeenCalledTimes(2);
+            const stream = await chat.sendMessageStream('test-model', { message: 'trigger 429' }, 'prompt-id-fb1');
+            // 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, FAILED_MODEL, authType, error429);
-            expect(result.candidates?.[0]?.content?.parts?.[0]?.text).toBe('Success on retry');
+            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 stop retrying if handleFallback returns false (e.g., auth intent)', async () => {
             vi.mocked(mockConfig.getModel).mockReturnValue('gemini-pro');
-            vi.mocked(mockContentGenerator.generateContent).mockRejectedValue(error429);
+            vi.mocked(mockContentGenerator.generateContentStream).mockRejectedValue(error429);
             mockHandleFallback.mockResolvedValue(false);
-            await expect(chat.sendMessage({ message: 'test stop' }, 'prompt-id-fb2')).rejects.toThrow(error429);
-            expect(mockContentGenerator.generateContent).toHaveBeenCalledTimes(1);
+            const stream = await chat.sendMessageStream('test-model', { message: 'test stop' }, 'prompt-id-fb2');
+            await expect((async () => {
+                for await (const _ of stream) {
+                    /* consume stream */
+                }
+            })()).rejects.toThrow(error429);
+            expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(1);
             expect(mockHandleFallback).toHaveBeenCalledTimes(1);
         });
     });
@@ -1312,7 +881,7 @@ describe('GeminiChat', () => {
             };
         })());
         // Send a message and consume the stream
-        const stream = await chat.sendMessageStream({ message: 'test' }, 'prompt-id-discard-test');
+        const stream = await chat.sendMessageStream('test-model', { message: 'test' }, 'prompt-id-discard-test');
         const events = [];
         for await (const event of stream) {
             events.push(event);