@machina.ai/cell-cli-core 1.6.1-rc2 → 1.10.0-rc1

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

@@ -4,12 +4,13 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
-import { GeminiChat, EmptyStreamError, StreamEventType, } from './geminiChat.js';
+import { ApiError } from '@google/genai';
+import { GeminiChat, InvalidStreamError, StreamEventType, } from './geminiChat.js';
 import { setSimulate429 } from '../utils/testUtils.js';
 import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
 import { AuthType } from './contentGenerator.js';
 import {} from '../utils/retry.js';
-import { Kind } from '../tools/tools.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', () => {
@@ -46,16 +47,19 @@ vi.mock('../utils/retry.js', () => ({
 vi.mock('../fallback/handler.js', () => ({
     handleFallback: mockHandleFallback,
 }));
-const { mockLogInvalidChunk, mockLogContentRetry, mockLogContentRetryFailure } = vi.hoisted(() => ({
-    mockLogInvalidChunk: vi.fn(),
+const { mockLogContentRetry, mockLogContentRetryFailure } = vi.hoisted(() => ({
     mockLogContentRetry: vi.fn(),
     mockLogContentRetryFailure: vi.fn(),
 }));
 vi.mock('../telemetry/loggers.js', () => ({
-    logInvalidChunk: mockLogInvalidChunk,
     logContentRetry: mockLogContentRetry,
     logContentRetryFailure: mockLogContentRetryFailure,
 }));
+vi.mock('../telemetry/uiTelemetry.js', () => ({
+    uiTelemetryService: {
+        setLastPromptTokenCount: vi.fn(),
+    },
+}));
 describe('GeminiChat', () => {
     let mockContentGenerator;
     let chat;
@@ -63,6 +67,7 @@ describe('GeminiChat', () => {
     const config = {};
     beforeEach(() => {
         vi.clearAllMocks();
+        vi.mocked(uiTelemetryService.setLastPromptTokenCount).mockClear();
         mockContentGenerator = {
             generateContent: vi.fn(),
             generateContentStream: vi.fn(),
@@ -96,6 +101,7 @@ describe('GeminiChat', () => {
                 getTool: vi.fn(),
             }),
             getContentGenerator: vi.fn().mockReturnValue(mockContentGenerator),
+            getRetryFetchErrors: vi.fn().mockReturnValue(false),
         };
         // Disable 429 simulation for tests
         setSimulate429(false);
@@ -180,7 +186,7 @@ describe('GeminiChat', () => {
                 for await (const _ of stream) {
                     /* consume stream */
                 }
-            })()).rejects.toThrow(EmptyStreamError);
+            })()).rejects.toThrow(InvalidStreamError);
         });
         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.
@@ -345,7 +351,7 @@ describe('GeminiChat', () => {
             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 stream response', async () => {
+        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 = [
                 {
@@ -386,20 +392,113 @@ describe('GeminiChat', () => {
                     },
                 },
             }, 'prompt-id-stream-1');
-            for await (const _ of stream) {
-                // This loop consumes the stream to trigger the internal logic.
-            }
-            // 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();
+            // 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);
+        });
+        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('test-model', { message: 'test' }, 'prompt-id-1');
+            // Should not throw an error
+            await expect((async () => {
+                for await (const _ of stream) {
+                    // consume stream
+                }
+            })()).resolves.not.toThrow();
+        });
+        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('test-model', { message: 'test' }, 'prompt-id-1');
+            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('test-model', { message: 'test' }, 'prompt-id-1');
+            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('test-model', { message: 'test' }, 'prompt-id-1');
+            // Should not throw an error
+            await expect((async () => {
+                for await (const _ of stream) {
+                    // consume stream
+                }
+            })()).resolves.not.toThrow();
         });
         it('should call generateContentStream with the correct parameters', async () => {
             const response = (async function* () {
@@ -416,6 +515,11 @@ describe('GeminiChat', () => {
                         },
                     ],
                     text: () => 'response',
+                    usageMetadata: {
+                        promptTokenCount: 42,
+                        candidatesTokenCount: 15,
+                        totalTokenCount: 57,
+                    },
                 };
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(response);
@@ -433,6 +537,9 @@ describe('GeminiChat', () => {
                 ],
                 config: {},
             }, 'prompt-id-1');
+            // Verify that token counting is called when usageMetadata is present
+            expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalledWith(42);
+            expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalledTimes(1);
         });
     });
     describe('addHistory', () => {
@@ -525,7 +632,6 @@ describe('GeminiChat', () => {
                 chunks.push(chunk);
             }
             // Assertions
-            expect(mockLogInvalidChunk).toHaveBeenCalledTimes(1);
             expect(mockLogContentRetry).toHaveBeenCalledTimes(1);
             expect(mockLogContentRetryFailure).not.toHaveBeenCalled();
             expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(2);
@@ -546,6 +652,48 @@ describe('GeminiChat', () => {
                 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('test-model', { message: 'test', config: { temperature: 0.5 } }, 'prompt-id-retry-temperature');
+            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.5,
+                }),
+            }), '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* () {
@@ -565,15 +713,152 @@ describe('GeminiChat', () => {
                 for await (const _ of stream) {
                     // Must loop to trigger the internal logic that throws.
                 }
-            }).rejects.toThrow(EmptyStreamError);
-            // Should be called 3 times (initial + 2 retries)
-            expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(3);
-            expect(mockLogInvalidChunk).toHaveBeenCalledTimes(3);
-            expect(mockLogContentRetry).toHaveBeenCalledTimes(2);
+            }).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 be clean, as if the failed turn never happened.
+            // History should still contain the user message.
             const history = chat.getHistory();
-            expect(history.length).toBe(0);
+            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('test-model', { message: 'test' }, 'prompt-id-400');
+                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('test-model', { message: 'test' }, 'prompt-id-429-retry');
+                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('test-model', { message: 'test' }, 'prompt-id-500-retry');
+                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('test-model', { message: 'test' }, 'prompt-id-fetch-error-retry');
+                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 () => {
@@ -745,215 +1030,6 @@ describe('GeminiChat', () => {
         }
         expect(turn4.parts[0].text).toBe('second response');
     });
-    describe('stopBeforeSecondMutator', () => {
-        beforeEach(() => {
-            // Common setup for these tests: mock the tool registry.
-            const mockToolRegistry = {
-                getTool: vi.fn((toolName) => {
-                    if (toolName === 'edit') {
-                        return { kind: Kind.Edit };
-                    }
-                    return { kind: Kind.Other };
-                }),
-            };
-            vi.mocked(mockConfig.getToolRegistry).mockReturnValue(mockToolRegistry);
-        });
-        it('should stop streaming before a second mutator tool call', async () => {
-            const responses = [
-                {
-                    candidates: [
-                        { content: { role: 'model', parts: [{ text: 'First part. ' }] } },
-                    ],
-                },
-                {
-                    candidates: [
-                        {
-                            content: {
-                                role: 'model',
-                                parts: [{ functionCall: { name: 'edit', args: {} } }],
-                            },
-                        },
-                    ],
-                },
-                {
-                    candidates: [
-                        {
-                            content: {
-                                role: 'model',
-                                parts: [{ functionCall: { name: 'fetch', args: {} } }],
-                            },
-                        },
-                    ],
-                },
-                // This chunk contains the second mutator and should be clipped.
-                {
-                    candidates: [
-                        {
-                            content: {
-                                role: 'model',
-                                parts: [
-                                    { functionCall: { name: 'edit', args: {} } },
-                                    { text: 'some trailing text' },
-                                ],
-                            },
-                        },
-                    ],
-                },
-                // This chunk should never be reached.
-                {
-                    candidates: [
-                        {
-                            content: {
-                                role: 'model',
-                                parts: [{ text: 'This should not appear.' }],
-                            },
-                        },
-                    ],
-                },
-            ];
-            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue((async function* () {
-                for (const response of responses) {
-                    yield response;
-                }
-            })());
-            const stream = await chat.sendMessageStream('test-model', { message: 'test message' }, 'prompt-id-mutator-test');
-            for await (const _ of stream) {
-                // Consume the stream to trigger history recording.
-            }
-            const history = chat.getHistory();
-            expect(history.length).toBe(2);
-            const modelTurn = history[1];
-            expect(modelTurn.role).toBe('model');
-            expect(modelTurn?.parts?.length).toBe(3);
-            expect(modelTurn?.parts[0].text).toBe('First part. ');
-            expect(modelTurn.parts[1].functionCall?.name).toBe('edit');
-            expect(modelTurn.parts[2].functionCall?.name).toBe('fetch');
-        });
-        it('should not stop streaming if only one mutator is present', async () => {
-            const responses = [
-                {
-                    candidates: [
-                        { content: { role: 'model', parts: [{ text: 'Part 1. ' }] } },
-                    ],
-                },
-                {
-                    candidates: [
-                        {
-                            content: {
-                                role: 'model',
-                                parts: [{ functionCall: { name: 'edit', args: {} } }],
-                            },
-                        },
-                    ],
-                },
-                {
-                    candidates: [
-                        {
-                            content: {
-                                role: 'model',
-                                parts: [{ text: 'Part 2.' }],
-                            },
-                            finishReason: 'STOP',
-                        },
-                    ],
-                },
-            ];
-            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue((async function* () {
-                for (const response of responses) {
-                    yield response;
-                }
-            })());
-            const stream = await chat.sendMessageStream('test-model', { message: 'test message' }, 'prompt-id-one-mutator');
-            for await (const _ of stream) {
-                /* consume */
-            }
-            const history = chat.getHistory();
-            const modelTurn = history[1];
-            expect(modelTurn?.parts?.length).toBe(3);
-            expect(modelTurn.parts[1].functionCall?.name).toBe('edit');
-            expect(modelTurn.parts[2].text).toBe('Part 2.');
-        });
-        it('should clip the chunk containing the second mutator, preserving prior parts', async () => {
-            const responses = [
-                {
-                    candidates: [
-                        {
-                            content: {
-                                role: 'model',
-                                parts: [{ functionCall: { name: 'edit', args: {} } }],
-                            },
-                        },
-                    ],
-                },
-                // This chunk has a valid part before the second mutator.
-                // The valid part should be kept, the rest of the chunk discarded.
-                {
-                    candidates: [
-                        {
-                            content: {
-                                role: 'model',
-                                parts: [
-                                    { text: 'Keep this text. ' },
-                                    { functionCall: { name: 'edit', args: {} } },
-                                    { text: 'Discard this text.' },
-                                ],
-                            },
-                            finishReason: 'STOP',
-                        },
-                    ],
-                },
-            ];
-            const stream = (async function* () {
-                for (const response of responses) {
-                    yield response;
-                }
-            })();
-            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(stream);
-            const resultStream = await chat.sendMessageStream('test-model', { message: 'test' }, 'prompt-id-clip-chunk');
-            for await (const _ of resultStream) {
-                /* consume */
-            }
-            const history = chat.getHistory();
-            const modelTurn = history[1];
-            expect(modelTurn?.parts?.length).toBe(2);
-            expect(modelTurn.parts[0].functionCall?.name).toBe('edit');
-            expect(modelTurn.parts[1].text).toBe('Keep this text. ');
-        });
-        it('should handle two mutators in the same chunk (parallel call scenario)', async () => {
-            const responses = [
-                {
-                    candidates: [
-                        {
-                            content: {
-                                role: 'model',
-                                parts: [
-                                    { text: 'Some text. ' },
-                                    { functionCall: { name: 'edit', args: {} } },
-                                    { functionCall: { name: 'edit', args: {} } },
-                                ],
-                            },
-                            finishReason: 'STOP',
-                        },
-                    ],
-                },
-            ];
-            const stream = (async function* () {
-                for (const response of responses) {
-                    yield response;
-                }
-            })();
-            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(stream);
-            const resultStream = await chat.sendMessageStream('test-model', { message: 'test' }, 'prompt-id-parallel-mutators');
-            for await (const _ of resultStream) {
-                /* consume */
-            }
-            const history = chat.getHistory();
-            const modelTurn = history[1];
-            expect(modelTurn?.parts?.length).toBe(2);
-            expect(modelTurn.parts[0].text).toBe('Some text. ');
-            expect(modelTurn.parts[1].functionCall?.name).toBe('edit');
-        });
-    });
     describe('Model Resolution', () => {
         const mockResponse = {
             candidates: [
@@ -979,7 +1055,8 @@ describe('GeminiChat', () => {
         });
     });
     describe('Fallback Integration (Retries)', () => {
-        const error429 = Object.assign(new Error('API Error 429: Quota exceeded'), {
+        const error429 = new ApiError({
+            message: 'API Error 429: Quota exceeded',
             status: 429,
         });
         // Define the simulated behavior for retryWithBackoff for these tests.