@machina.ai/cell-cli-core 1.13.0-rc5 → 1.16.0-rc2

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

@@ -5,11 +5,12 @@
  */
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import { ApiError } from '@google/genai';
-import { GeminiChat, InvalidStreamError, StreamEventType, } from './geminiChat.js';
+import { GeminiChat, InvalidStreamError, StreamEventType, SYNTHETIC_THOUGHT_SIGNATURE, } from './geminiChat.js';
 import { setSimulate429 } from '../utils/testUtils.js';
-import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
+import { DEFAULT_GEMINI_FLASH_MODEL, DEFAULT_GEMINI_MODEL, PREVIEW_GEMINI_MODEL, } from '../config/models.js';
 import { AuthType } from './contentGenerator.js';
-import {} from '../utils/retry.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();
@@ -83,6 +84,7 @@ describe('GeminiChat', () => {
             getTelemetryLogPromptsEnabled: () => true,
             getUsageStatisticsEnabled: () => true,
             getDebugMode: () => false,
+            getPreviewFeatures: () => false,
             getContentGeneratorConfig: vi.fn().mockReturnValue({
                 authType: 'oauth-personal', // Ensure this is set for fallback tests
                 model: 'test-model',
@@ -102,6 +104,11 @@ describe('GeminiChat', () => {
             }),
             getContentGenerator: vi.fn().mockReturnValue(mockContentGenerator),
             getRetryFetchErrors: vi.fn().mockReturnValue(false),
+            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);
@@ -196,7 +203,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('test-model', { message: 'test message' }, 'prompt-id-no-finish-empty-end');
+            const stream = await chat.sendMessageStream('gemini-2.0-flash', { message: 'test message' }, 'prompt-id-no-finish-empty-end');
             await expect((async () => {
                 for await (const _ of stream) {
                     /* consume stream */
@@ -366,6 +373,84 @@ 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 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(PREVIEW_GEMINI_MODEL, { message: 'test' }, 'prompt-id-fast-retry');
+            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(DEFAULT_GEMINI_FLASH_MODEL, { message: 'test' }, 'prompt-id-normal-retry');
+            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(PREVIEW_GEMINI_MODEL, { message: 'test' }, 'prompt-id-bypass');
+                // 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 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 = [
@@ -399,7 +484,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('test-model', {
+            const stream = await chat.sendMessageStream('gemini-2.0-flash', {
                 message: {
                     functionResponse: {
                         name: 'find_restaurant',
@@ -461,7 +546,7 @@ describe('GeminiChat', () => {
                 };
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithoutFinishReason);
-            const stream = await chat.sendMessageStream('test-model', { message: 'test' }, 'prompt-id-1');
+            const stream = await chat.sendMessageStream('gemini-2.0-flash', { message: 'test' }, 'prompt-id-1');
             await expect((async () => {
                 for await (const _ of stream) {
                     // consume stream
@@ -484,7 +569,7 @@ describe('GeminiChat', () => {
                 };
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithEmptyResponse);
-            const stream = await chat.sendMessageStream('test-model', { message: 'test' }, 'prompt-id-1');
+            const stream = await chat.sendMessageStream('gemini-2.0-flash', { message: 'test' }, 'prompt-id-1');
             await expect((async () => {
                 for await (const _ of stream) {
                     // consume stream
@@ -515,6 +600,71 @@ describe('GeminiChat', () => {
                 }
             })()).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('gemini-2.5-pro', { message: 'test' }, 'prompt-id-malformed');
+            // 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('gemini-2.5-pro', { message: 'test retry' }, 'prompt-id-retry-malformed');
+            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 {
@@ -583,6 +733,23 @@ describe('GeminiChat', () => {
         });
     });
     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('gemini-1.5-pro', { message: 'test' }, 'prompt-id-no-retry');
+            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)
@@ -606,7 +773,7 @@ describe('GeminiChat', () => {
                 };
             })());
             // ACT: Send a message and collect all events from the stream.
-            const stream = await chat.sendMessageStream('test-model', { message: 'test' }, 'prompt-id-yield-retry');
+            const stream = await chat.sendMessageStream('gemini-2.0-flash', { message: 'test' }, 'prompt-id-yield-retry');
             const events = [];
             for await (const event of stream) {
                 events.push(event);
@@ -638,7 +805,7 @@ describe('GeminiChat', () => {
                     ],
                 };
             })());
-            const stream = await chat.sendMessageStream('test-model', { message: 'test' }, 'prompt-id-retry-success');
+            const stream = await chat.sendMessageStream('gemini-2.0-flash', { message: 'test' }, 'prompt-id-retry-success');
             const chunks = [];
             for await (const chunk of stream) {
                 chunks.push(chunk);
@@ -689,7 +856,7 @@ describe('GeminiChat', () => {
                     ],
                 };
             })());
-            const stream = await chat.sendMessageStream('test-model', { message: 'test', config: { temperature: 0.5 } }, 'prompt-id-retry-temperature');
+            const stream = await chat.sendMessageStream('gemini-2.0-flash', { message: 'test', config: { temperature: 0.5 } }, 'prompt-id-retry-temperature');
             for await (const _ of stream) {
                 // consume stream
             }
@@ -720,7 +887,7 @@ describe('GeminiChat', () => {
                     ],
                 };
             })());
-            const stream = await chat.sendMessageStream('test-model', { message: 'test' }, 'prompt-id-retry-fail');
+            const stream = await chat.sendMessageStream('gemini-2.0-flash', { message: 'test' }, 'prompt-id-retry-fail');
             await expect(async () => {
                 for await (const _ of stream) {
                     // Must loop to trigger the internal logic that throws.
@@ -769,7 +936,7 @@ describe('GeminiChat', () => {
             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');
+                const stream = await chat.sendMessageStream('gemini-2.0-flash', { message: 'test' }, 'prompt-id-400');
                 await expect((async () => {
                     for await (const _ of stream) {
                         /* consume stream */
@@ -900,7 +1067,7 @@ describe('GeminiChat', () => {
             };
         })());
         // 3. Send a new message
-        const stream = await chat.sendMessageStream('test-model', { message: 'Second question' }, 'prompt-id-retry-existing');
+        const stream = await chat.sendMessageStream('gemini-2.0-flash', { message: 'Second question' }, 'prompt-id-retry-existing');
         for await (const _ of stream) {
             // consume stream
         }
@@ -952,7 +1119,7 @@ describe('GeminiChat', () => {
             };
         })());
         // 2. Call the method and consume the stream.
-        const stream = await chat.sendMessageStream('test-model', { message: 'test empty stream' }, 'prompt-id-empty-stream');
+        const stream = await chat.sendMessageStream('gemini-2.0-flash', { message: 'test empty stream' }, 'prompt-id-empty-stream');
         const chunks = [];
         for await (const chunk of stream) {
             chunks.push(chunk);
@@ -1136,7 +1303,7 @@ describe('GeminiChat', () => {
             vi.mocked(mockConfig.getModel).mockReturnValue('gemini-pro');
             vi.mocked(mockContentGenerator.generateContentStream).mockRejectedValue(error429);
             mockHandleFallback.mockResolvedValue(false);
-            const stream = await chat.sendMessageStream('test-model', { message: 'test stop' }, 'prompt-id-fb2');
+            const stream = await chat.sendMessageStream('gemini-2.0-flash', { message: 'test stop' }, 'prompt-id-fb2');
             await expect((async () => {
                 for await (const _ of stream) {
                     /* consume stream */
@@ -1180,7 +1347,7 @@ describe('GeminiChat', () => {
             };
         })());
         // Send a message and consume the stream
-        const stream = await chat.sendMessageStream('test-model', { message: 'test' }, 'prompt-id-discard-test');
+        const stream = await chat.sendMessageStream('gemini-2.0-flash', { message: 'test' }, 'prompt-id-discard-test');
         const events = [];
         for await (const event of stream) {
             events.push(event);
@@ -1231,5 +1398,139 @@ describe('GeminiChat', () => {
             ]);
         });
     });
+    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('test-model', { message: 'test' }, 'prompt-id-preview-model-reset');
+            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(PREVIEW_GEMINI_MODEL, { message: 'test' }, 'prompt-id-preview-model-healing');
+            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(PREVIEW_GEMINI_MODEL, { message: 'test' }, 'prompt-id-bypass-no-healing');
+            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);
+        });
+    });
 });
 //# sourceMappingURL=geminiChat.test.js.map