@machina.ai/cell-cli-core 1.20.2-rc1 → 1.22.5-rc1

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

@@ -7,13 +7,16 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import { ApiError, ThinkingLevel } from '@google/genai';
 import { GeminiChat, InvalidStreamError, StreamEventType, SYNTHETIC_THOUGHT_SIGNATURE, } from './geminiChat.js';
 import { setSimulate429 } from '../utils/testUtils.js';
-import { DEFAULT_GEMINI_FLASH_MODEL, DEFAULT_GEMINI_MODEL, DEFAULT_THINKING_MODE, PREVIEW_GEMINI_MODEL, } from '../config/models.js';
+import { DEFAULT_GEMINI_FLASH_MODEL, DEFAULT_THINKING_MODE, PREVIEW_GEMINI_MODEL, PREVIEW_GEMINI_FLASH_MODEL, } from '../config/models.js';
 import { AuthType } from './contentGenerator.js';
 import { TerminalQuotaError } from '../utils/googleQuotaErrors.js';
-import { retryWithBackoff } from '../utils/retry.js';
+import {} from '../utils/retry.js';
 import { uiTelemetryService } from '../telemetry/uiTelemetry.js';
 import { HookSystem } from '../hooks/hookSystem.js';
 import { createMockMessageBus } from '../test-utils/mock-message-bus.js';
+import { createAvailabilityServiceMock } from '../availability/testUtils.js';
+import * as policyHelpers from '../availability/policyHelpers.js';
+import { makeResolvedModelConfig } from '../services/modelConfigServiceTestUtils.js';
 // Mock fs module to prevent actual file system operations during tests
 const mockFileSystem = new Map();
 vi.mock('node:fs', () => {
@@ -31,6 +34,10 @@ vi.mock('node:fs', () => {
             });
         }),
         existsSync: vi.fn((path) => mockFileSystem.has(path)),
+        createWriteStream: vi.fn(() => ({
+            write: vi.fn(),
+            on: vi.fn(),
+        })),
     };
     return {
         default: fsModule,
@@ -44,9 +51,13 @@ const { mockHandleFallback } = vi.hoisted(() => ({
 const { mockRetryWithBackoff } = vi.hoisted(() => ({
     mockRetryWithBackoff: vi.fn(),
 }));
-vi.mock('../utils/retry.js', () => ({
-    retryWithBackoff: mockRetryWithBackoff,
-}));
+vi.mock('../utils/retry.js', async (importOriginal) => {
+    const actual = await importOriginal();
+    return {
+        ...actual,
+        retryWithBackoff: mockRetryWithBackoff,
+    };
+});
 vi.mock('../fallback/handler.js', () => ({
     handleFallback: mockHandleFallback,
 }));
@@ -79,19 +90,32 @@ describe('GeminiChat', () => {
         };
         mockHandleFallback.mockClear();
         // Default mock implementation for tests that don't care about retry logic
-        mockRetryWithBackoff.mockImplementation(async (apiCall) => apiCall());
+        mockRetryWithBackoff.mockImplementation(async (apiCall, options) => {
+            const result = await apiCall();
+            const context = options?.getAvailabilityContext?.();
+            if (context) {
+                context.service.markHealthy(context.policy.model);
+            }
+            return result;
+        });
+        let currentModel = 'gemini-pro';
+        let currentActiveModel = 'gemini-pro';
         mockConfig = {
             getSessionId: () => 'test-session-id',
             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',
+            getContentGeneratorConfig: vi.fn().mockImplementation(() => ({
+                authType: 'oauth-personal',
+                model: currentModel,
+            })),
+            getModel: vi.fn().mockImplementation(() => currentModel),
+            setModel: vi.fn().mockImplementation((m) => {
+                currentModel = m;
+                // When model is explicitly set, active model usually resets or updates to it
+                currentActiveModel = m;
             }),
-            getModel: vi.fn().mockReturnValue('gemini-pro'),
-            setModel: vi.fn(),
             isInFallbackMode: vi.fn().mockReturnValue(false),
             getQuotaErrorOccurred: vi.fn().mockReturnValue(false),
             setQuotaErrorOccurred: vi.fn(),
@@ -105,9 +129,11 @@ describe('GeminiChat', () => {
             }),
             getContentGenerator: vi.fn().mockReturnValue(mockContentGenerator),
             getRetryFetchErrors: vi.fn().mockReturnValue(false),
+            getUserTier: vi.fn().mockReturnValue(undefined),
             modelConfigService: {
                 getResolvedConfig: vi.fn().mockImplementation((modelConfigKey) => {
-                    const thinkingConfig = modelConfigKey.model.startsWith('gemini-3')
+                    const model = modelConfigKey.model ?? mockConfig.getModel();
+                    const thinkingConfig = model.startsWith('gemini-3')
                         ? {
                             thinkingLevel: ThinkingLevel.HIGH,
                         }
@@ -115,7 +141,7 @@ describe('GeminiChat', () => {
                             thinkingBudget: DEFAULT_THINKING_MODE,
                         };
                     return {
-                        model: modelConfigKey.model,
+                        model,
                         generateContentConfig: {
                             temperature: 0,
                             thinkingConfig,
@@ -129,6 +155,12 @@ describe('GeminiChat', () => {
             setPreviewModelFallbackMode: vi.fn(),
             isInteractive: vi.fn().mockReturnValue(false),
             getEnableHooks: vi.fn().mockReturnValue(false),
+            isModelAvailabilityServiceEnabled: vi.fn().mockReturnValue(false),
+            getActiveModel: vi.fn().mockImplementation(() => currentActiveModel),
+            setActiveModel: vi
+                .fn()
+                .mockImplementation((m) => (currentActiveModel = m)),
+            getModelAvailabilityService: vi.fn(),
         };
         // Use proper MessageBus mocking for Phase 3 preparation
         const mockMessageBus = createMockMessageBus();
@@ -421,6 +453,26 @@ describe('GeminiChat', () => {
                 maxAttempts: 1,
             }));
         });
+        it('should use maxAttempts=1 for retryWithBackoff when in Preview Model Fallback Mode (Flash)', async () => {
+            vi.mocked(mockConfig.isPreviewModelFallbackMode).mockReturnValue(true);
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue((async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { parts: [{ text: 'Success' }] },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })());
+            const stream = await chat.sendMessageStream({ model: PREVIEW_GEMINI_FLASH_MODEL }, 'test', 'prompt-id-fast-retry-flash', new AbortController().signal);
+            for await (const _ of stream) {
+                // consume stream
+            }
+            expect(mockRetryWithBackoff).toHaveBeenCalledWith(expect.any(Function), expect.objectContaining({
+                maxAttempts: 1,
+            }));
+        });
         it('should NOT use maxAttempts=1 for other models even in Preview Model Fallback Mode', async () => {
             vi.mocked(mockConfig.isPreviewModelFallbackMode).mockReturnValue(true);
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue((async function* () {
@@ -441,44 +493,6 @@ describe('GeminiChat', () => {
                 maxAttempts: undefined, // Should use default
             }));
         });
-        it('should pass DEFAULT_GEMINI_MODEL to handleFallback when Preview Model is bypassed (downgraded)', async () => {
-            // ARRANGE
-            vi.mocked(mockConfig.isPreviewModelBypassMode).mockReturnValue(true);
-            // Mock retryWithBackoff to simulate catching the error and calling onPersistent429
-            vi.mocked(retryWithBackoff).mockImplementation(async (apiCall, options) => {
-                const onPersistent429 = options?.onPersistent429;
-                try {
-                    await apiCall();
-                }
-                catch (error) {
-                    if (onPersistent429) {
-                        await onPersistent429(AuthType.LOGIN_WITH_GOOGLE, error);
-                    }
-                    throw error;
-                }
-            });
-            // We need the API call to fail so retryWithBackoff calls the callback.
-            vi.mocked(mockContentGenerator.generateContentStream).mockRejectedValue(new TerminalQuotaError('Simulated Quota Error', {
-                code: 429,
-                message: 'Simulated Quota Error',
-                details: [],
-            }));
-            // ACT
-            const consumeStream = async () => {
-                const stream = await chat.sendMessageStream({ model: PREVIEW_GEMINI_MODEL }, 'test', 'prompt-id-bypass', new AbortController().signal);
-                // Consume the stream to trigger execution
-                for await (const _ of stream) {
-                    // do nothing
-                }
-            };
-            await expect(consumeStream()).rejects.toThrow('Simulated Quota Error');
-            expect(retryWithBackoff).toHaveBeenCalled();
-            // ASSERT
-            // handleFallback is called via onPersistent429Callback
-            // We verify it was called with DEFAULT_GEMINI_MODEL
-            expect(mockHandleFallback).toHaveBeenCalledWith(expect.anything(), DEFAULT_GEMINI_MODEL, // This is the key assertion
-            expect.anything(), expect.anything());
-        });
         it('should 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 = [
@@ -1295,30 +1309,6 @@ describe('GeminiChat', () => {
         }
         expect(turn4.parts[0].text).toBe('second response');
     });
-    describe('Model Resolution', () => {
-        const mockResponse = {
-            candidates: [
-                {
-                    content: { parts: [{ text: 'response' }], role: 'model' },
-                    finishReason: 'STOP',
-                },
-            ],
-        };
-        it('should use the FLASH model when in fallback mode (sendMessageStream)', async () => {
-            vi.mocked(mockConfig.getModel).mockReturnValue('gemini-pro');
-            vi.mocked(mockConfig.isInFallbackMode).mockReturnValue(true);
-            vi.mocked(mockContentGenerator.generateContentStream).mockImplementation(async () => (async function* () {
-                yield mockResponse;
-            })());
-            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-res3', new AbortController().signal);
-            for await (const _ of stream) {
-                // consume stream
-            }
-            expect(mockContentGenerator.generateContentStream).toHaveBeenCalledWith(expect.objectContaining({
-                model: DEFAULT_GEMINI_FLASH_MODEL,
-            }), 'prompt-id-res3');
-        });
-    });
     describe('Fallback Integration (Retries)', () => {
         const error429 = new ApiError({
             message: 'API Error 429: Quota exceeded',
@@ -1385,64 +1375,6 @@ describe('GeminiChat', () => {
             const modelTurn = history[1];
             expect(modelTurn.parts[0].text).toBe('Success on retry');
         });
-        it('should switch to DEFAULT_GEMINI_FLASH_MODEL and use thinkingBudget when falling back from a gemini-3 model', async () => {
-            // ARRANGE
-            const authType = AuthType.LOGIN_WITH_GOOGLE;
-            vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({
-                authType,
-            });
-            // Initial state: Not in fallback mode
-            const isInFallbackModeSpy = vi.spyOn(mockConfig, 'isInFallbackMode');
-            isInFallbackModeSpy.mockReturnValue(false);
-            // Mock API calls:
-            // 1. Fails with 429 (simulating gemini-3 failure)
-            // 2. Succeeds (simulating fallback success)
-            vi.mocked(mockContentGenerator.generateContentStream)
-                .mockRejectedValueOnce(error429)
-                .mockResolvedValueOnce((async function* () {
-                yield {
-                    candidates: [
-                        {
-                            content: { parts: [{ text: 'Fallback success' }] },
-                            finishReason: 'STOP',
-                        },
-                    ],
-                };
-            })());
-            // Mock handleFallback to enable fallback mode and signal retry
-            mockHandleFallback.mockImplementation(async () => {
-                isInFallbackModeSpy.mockReturnValue(true); // Next call will see fallback mode = true
-                return true;
-            });
-            // ACT
-            const stream = await chat.sendMessageStream({ model: 'gemini-3-test-model' }, // Start with a gemini-3 model
-            'test fallback thinking', 'prompt-id-fb3', new AbortController().signal);
-            for await (const _ of stream) {
-                // consume stream
-            }
-            // ASSERT
-            expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(2);
-            // First call: gemini-3 model, thinkingLevel set
-            expect(mockContentGenerator.generateContentStream).toHaveBeenNthCalledWith(1, expect.objectContaining({
-                model: 'gemini-3-test-model',
-                config: expect.objectContaining({
-                    thinkingConfig: {
-                        thinkingBudget: undefined,
-                        thinkingLevel: ThinkingLevel.HIGH,
-                    },
-                }),
-            }), 'prompt-id-fb3');
-            // Second call: DEFAULT_GEMINI_FLASH_MODEL (due to fallback), thinkingBudget set (due to fix)
-            expect(mockContentGenerator.generateContentStream).toHaveBeenNthCalledWith(2, expect.objectContaining({
-                model: DEFAULT_GEMINI_FLASH_MODEL,
-                config: expect.objectContaining({
-                    thinkingConfig: {
-                        thinkingBudget: DEFAULT_THINKING_MODE,
-                        thinkingLevel: undefined,
-                    },
-                }),
-            }), 'prompt-id-fb3');
-        });
         it('should stop retrying if handleFallback returns false (e.g., auth intent)', async () => {
             vi.mocked(mockConfig.getModel).mockReturnValue('gemini-pro');
             vi.mocked(mockContentGenerator.generateContentStream).mockRejectedValue(error429);
@@ -1542,61 +1474,6 @@ 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({ model: 'test-model' }, 'test', 'prompt-id-preview-model-reset', new AbortController().signal);
-            expect(mockConfig.setPreviewModelBypassMode).toHaveBeenCalledWith(false);
-        });
-        it('should reset previewModelFallbackMode to false upon successful Preview Model usage', async () => {
-            const stream = (async function* () {
-                yield {
-                    candidates: [
-                        {
-                            content: { role: 'model', parts: [{ text: 'Success' }] },
-                            finishReason: 'STOP',
-                        },
-                    ],
-                };
-            })();
-            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(stream);
-            const resultStream = await chat.sendMessageStream({ model: PREVIEW_GEMINI_MODEL }, 'test', 'prompt-id-preview-model-healing', new AbortController().signal);
-            for await (const _ of resultStream) {
-                // consume stream
-            }
-            expect(mockConfig.setPreviewModelFallbackMode).toHaveBeenCalledWith(false);
-        });
-        it('should NOT reset previewModelFallbackMode if Preview Model was bypassed (downgraded)', async () => {
-            const stream = (async function* () {
-                yield {
-                    candidates: [
-                        {
-                            content: { role: 'model', parts: [{ text: 'Success' }] },
-                            finishReason: 'STOP',
-                        },
-                    ],
-                };
-            })();
-            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(stream);
-            // Simulate bypass mode being active (downgrade happened)
-            vi.mocked(mockConfig.isPreviewModelBypassMode).mockReturnValue(true);
-            const resultStream = await chat.sendMessageStream({ model: PREVIEW_GEMINI_MODEL }, 'test', 'prompt-id-bypass-no-healing', new AbortController().signal);
-            for await (const _ of resultStream) {
-                // consume stream
-            }
-            expect(mockConfig.setPreviewModelFallbackMode).not.toHaveBeenCalled();
-        });
-    });
     describe('ensureActiveLoopHasThoughtSignatures', () => {
         it('should add thoughtSignature to the first functionCall in each model turn of the active loop', () => {
             const chat = new GeminiChat(mockConfig, '', [], []);
@@ -1676,5 +1553,176 @@ describe('GeminiChat', () => {
             expect(newContents).toEqual(history);
         });
     });
+    describe('Availability Service Integration', () => {
+        let mockAvailabilityService;
+        beforeEach(async () => {
+            mockAvailabilityService = createAvailabilityServiceMock();
+            vi.mocked(mockConfig.getModelAvailabilityService).mockReturnValue(mockAvailabilityService);
+            vi.mocked(mockConfig.isModelAvailabilityServiceEnabled).mockReturnValue(true);
+            // Stateful mock for activeModel
+            let activeModel = 'model-a';
+            vi.mocked(mockConfig.getActiveModel).mockImplementation(() => activeModel);
+            vi.mocked(mockConfig.setActiveModel).mockImplementation((model) => {
+                activeModel = model;
+            });
+            vi.spyOn(policyHelpers, 'resolvePolicyChain').mockReturnValue([
+                {
+                    model: 'model-a',
+                    isLastResort: false,
+                    actions: {},
+                    stateTransitions: {},
+                },
+                {
+                    model: 'model-b',
+                    isLastResort: false,
+                    actions: {},
+                    stateTransitions: {},
+                },
+                {
+                    model: 'model-c',
+                    isLastResort: true,
+                    actions: {},
+                    stateTransitions: {},
+                },
+            ]);
+        });
+        it('should mark healthy on successful stream', async () => {
+            vi.mocked(mockAvailabilityService.selectFirstAvailable).mockReturnValue({
+                selectedModel: 'model-b',
+                skipped: [],
+            });
+            // Simulate selection happening upstream
+            mockConfig.setActiveModel('model-b');
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue((async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { parts: [{ text: 'Response' }], role: 'model' },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })());
+            const stream = await chat.sendMessageStream({ model: 'gemini-pro' }, 'test', 'prompt-healthy', new AbortController().signal);
+            for await (const _ of stream) {
+                // consume
+            }
+            expect(mockAvailabilityService.markHealthy).toHaveBeenCalledWith('model-b');
+        });
+        it('caps retries to a single attempt when selection is sticky', async () => {
+            vi.mocked(mockAvailabilityService.selectFirstAvailable).mockReturnValue({
+                selectedModel: 'model-a',
+                attempts: 1,
+                skipped: [],
+            });
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue((async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { parts: [{ text: 'Response' }], role: 'model' },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })());
+            const stream = await chat.sendMessageStream({ model: 'gemini-pro' }, 'test', 'prompt-sticky-once', new AbortController().signal);
+            for await (const _ of stream) {
+                // consume
+            }
+            expect(mockRetryWithBackoff).toHaveBeenCalledWith(expect.any(Function), expect.objectContaining({ maxAttempts: 1 }));
+            expect(mockAvailabilityService.consumeStickyAttempt).toHaveBeenCalledWith('model-a');
+        });
+        it('should pass attempted model to onPersistent429 callback which calls handleFallback', async () => {
+            vi.mocked(mockAvailabilityService.selectFirstAvailable).mockReturnValue({
+                selectedModel: 'model-a',
+                skipped: [],
+            });
+            // Simulate selection happening upstream
+            mockConfig.setActiveModel('model-a');
+            // Simulate retry logic behavior: catch error, call onPersistent429
+            const error = new TerminalQuotaError('Quota', {
+                code: 429,
+                message: 'quota',
+                details: [],
+            });
+            vi.mocked(mockContentGenerator.generateContentStream).mockRejectedValue(error);
+            // We need retryWithBackoff to trigger the callback
+            mockRetryWithBackoff.mockImplementation(async (apiCall, options) => {
+                try {
+                    await apiCall();
+                }
+                catch (e) {
+                    if (options?.onPersistent429) {
+                        await options.onPersistent429(AuthType.LOGIN_WITH_GOOGLE, e);
+                    }
+                    throw e; // throw anyway to end test
+                }
+            });
+            const consume = async () => {
+                const stream = await chat.sendMessageStream({ model: 'gemini-pro' }, 'test', 'prompt-fallback-arg', new AbortController().signal);
+                for await (const _ of stream) {
+                    // consume
+                }
+            };
+            await expect(consume()).rejects.toThrow();
+            // handleFallback is called with the ATTEMPTED model (model-a), not the requested one (gemini-pro)
+            expect(mockHandleFallback).toHaveBeenCalledWith(expect.anything(), 'model-a', expect.anything(), error);
+        });
+        it('re-resolves generateContentConfig when active model changes between retries', async () => {
+            // Availability enabled with stateful active model
+            let activeModel = 'model-a';
+            vi.mocked(mockConfig.getActiveModel).mockImplementation(() => activeModel);
+            vi.mocked(mockConfig.setActiveModel).mockImplementation((model) => {
+                activeModel = model;
+            });
+            // Different configs per model
+            vi.mocked(mockConfig.modelConfigService.getResolvedConfig)
+                .mockReturnValueOnce(makeResolvedModelConfig('model-a', { temperature: 0.1 }))
+                .mockReturnValueOnce(makeResolvedModelConfig('model-b', { temperature: 0.9 }));
+            // First attempt uses model-a, then simulate availability switching to model-b
+            mockRetryWithBackoff.mockImplementation(async (apiCall) => {
+                await apiCall(); // first attempt
+                activeModel = 'model-b'; // simulate switch before retry
+                return apiCall(); // second attempt
+            });
+            // Generators for each attempt
+            const firstResponse = (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { parts: [{ text: 'first' }], role: 'model' },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })();
+            const secondResponse = (async function* () {
+                yield {
+                    candidates: [
+                        {
+                            content: { parts: [{ text: 'second' }], role: 'model' },
+                            finishReason: 'STOP',
+                        },
+                    ],
+                };
+            })();
+            vi.mocked(mockContentGenerator.generateContentStream)
+                .mockResolvedValueOnce(firstResponse)
+                .mockResolvedValueOnce(secondResponse);
+            const stream = await chat.sendMessageStream({ model: 'gemini-pro' }, 'test', 'prompt-config-refresh', new AbortController().signal);
+            // Consume to drive both attempts
+            for await (const _ of stream) {
+                // consume
+            }
+            expect(mockContentGenerator.generateContentStream).toHaveBeenNthCalledWith(1, expect.objectContaining({
+                model: 'model-a',
+                config: expect.objectContaining({ temperature: 0.1 }),
+            }), expect.any(String));
+            expect(mockContentGenerator.generateContentStream).toHaveBeenNthCalledWith(2, expect.objectContaining({
+                model: 'model-b',
+                config: expect.objectContaining({ temperature: 0.9 }),
+            }), expect.any(String));
+        });
+    });
 });
 //# sourceMappingURL=geminiChat.test.js.map