@office-ai/aioncli-core 0.18.6 → 0.24.0-preview.1

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

@@ -9,14 +9,17 @@ import { AuthType, } from './contentGenerator.js';
 import {} from './geminiChat.js';
 import { CompressionStatus, GeminiEventType, Turn, } from './turn.js';
 import { getCoreSystemPrompt } from './prompts.js';
-import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
+import { DEFAULT_GEMINI_MODEL_AUTO } from '../config/models.js';
 import { FileDiscoveryService } from '../services/fileDiscoveryService.js';
 import { setSimulate429 } from '../utils/testUtils.js';
 import { tokenLimit } from './tokenLimits.js';
 import { ideContextStore } from '../ide/ideContext.js';
 import { uiTelemetryService } from '../telemetry/uiTelemetry.js';
 import { ChatCompressionService } from '../services/chatCompressionService.js';
+import { createAvailabilityServiceMock } from '../availability/testUtils.js';
 import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js';
+import { HookSystem } from '../hooks/hookSystem.js';
+import * as policyCatalog from '../availability/policyCatalog.js';
 vi.mock('../services/chatCompressionService.js');
 // Mock fs module to prevent actual file system operations during tests
 const mockFileSystem = new Map();
@@ -35,13 +38,16 @@ vi.mock('node:fs', () => {
             });
         }),
         existsSync: vi.fn((path) => mockFileSystem.has(path)),
+        createWriteStream: vi.fn(() => ({
+            write: vi.fn(),
+            on: vi.fn(),
+        })),
     };
     return {
         default: fsModule,
         ...fsModule,
     };
 });
-// --- Mocks ---
 const mockTurnRunFn = vi.fn();
 vi.mock('./turn', async (importOriginal) => {
     const actual = await importOriginal();
@@ -53,6 +59,7 @@ vi.mock('./turn', async (importOriginal) => {
         constructor() {
             // The constructor can be empty or do some mock setup
         }
+        getResponseText = vi.fn().mockReturnValue('Mock Response');
     }
     // Export the mock class as 'Turn'
     return {
@@ -85,6 +92,16 @@ vi.mock('../telemetry/uiTelemetry.js', () => ({
         getLastPromptTokenCount: vi.fn(),
     },
 }));
+vi.mock('../hooks/hookSystem.js');
+vi.mock('./clientHookTriggers.js', () => ({
+    fireBeforeAgentHook: vi.fn(),
+    fireAfterAgentHook: vi.fn().mockResolvedValue({
+        decision: 'allow',
+        continue: false,
+        suppressOutput: false,
+        systemMessage: undefined,
+    }),
+}));
 /**
  * Array.fromAsync ponyfill, which will be available in es 2024.
  *
@@ -102,6 +119,7 @@ describe('Gemini Client (client.ts)', () => {
     let mockConfig;
     let client;
     let mockGenerateContentFn;
+    let mockRouterService;
     beforeEach(async () => {
         vi.resetAllMocks();
         ClearcutLogger.clearInstance();
@@ -119,10 +137,16 @@ describe('Gemini Client (client.ts)', () => {
         });
         // Disable 429 simulation for tests
         setSimulate429(false);
+        mockRouterService = {
+            route: vi
+                .fn()
+                .mockResolvedValue({ model: 'default-routed-model', reason: 'test' }),
+        };
         mockContentGenerator = {
             generateContent: mockGenerateContentFn,
             generateContentStream: vi.fn(),
             batchEmbedContents: vi.fn(),
+            countTokens: vi.fn().mockResolvedValue({ totalTokens: 100 }),
         };
         // Because the GeminiClient constructor kicks off an async process (startChat)
         // that depends on a fully-formed Config object, we need to mock the
@@ -143,11 +167,15 @@ describe('Gemini Client (client.ts)', () => {
                 .mockReturnValue(contentGeneratorConfig),
             getToolRegistry: vi.fn().mockReturnValue(mockToolRegistry),
             getModel: vi.fn().mockReturnValue('test-model'),
+            getUserTier: vi.fn().mockReturnValue(undefined),
             getEmbeddingModel: vi.fn().mockReturnValue('test-embedding-model'),
             getApiKey: vi.fn().mockReturnValue('test-key'),
             getVertexAI: vi.fn().mockReturnValue(false),
             getUserAgent: vi.fn().mockReturnValue('test-agent'),
             getUserMemory: vi.fn().mockReturnValue(''),
+            getGlobalMemory: vi.fn().mockReturnValue(''),
+            getEnvironmentMemory: vi.fn().mockReturnValue(''),
+            isJitContextEnabled: vi.fn().mockReturnValue(false),
             getSessionId: vi.fn().mockReturnValue('test-session-id'),
             getProxy: vi.fn().mockReturnValue(undefined),
             getWorkingDir: vi.fn().mockReturnValue('/test/dir'),
@@ -165,15 +193,13 @@ describe('Gemini Client (client.ts)', () => {
                 getDirectories: vi.fn().mockReturnValue(['/test/dir']),
             }),
             getGeminiClient: vi.fn(),
-            getModelRouterService: vi.fn().mockReturnValue({
-                route: vi.fn().mockResolvedValue({ model: 'default-routed-model' }),
-            }),
-            isInFallbackMode: vi.fn().mockReturnValue(false),
-            setFallbackMode: vi.fn(),
+            getModelRouterService: vi
+                .fn()
+                .mockReturnValue(mockRouterService),
+            getMessageBus: vi.fn().mockReturnValue(undefined),
+            getEnableHooks: vi.fn().mockReturnValue(false),
             getChatCompression: vi.fn().mockReturnValue(undefined),
             getSkipNextSpeakerCheck: vi.fn().mockReturnValue(false),
-            getUseSmartEdit: vi.fn().mockReturnValue(false),
-            getUseModelRouter: vi.fn().mockReturnValue(false),
             getShowModelInfoInChat: vi.fn().mockReturnValue(false),
             getContinueOnFailedApiCall: vi.fn(),
             getProjectRoot: vi.fn().mockReturnValue('/test/project/root'),
@@ -199,7 +225,17 @@ describe('Gemini Client (client.ts)', () => {
                 },
             },
             isInteractive: vi.fn().mockReturnValue(false),
+            getExperiments: () => { },
+            getActiveModel: vi.fn().mockReturnValue('test-model'),
+            setActiveModel: vi.fn(),
+            resetTurn: vi.fn(),
+            getModelAvailabilityService: vi
+                .fn()
+                .mockReturnValue(createAvailabilityServiceMock()),
         };
+        mockConfig.getHookSystem = vi
+            .fn()
+            .mockReturnValue(new HookSystem(mockConfig));
         client = new GeminiClient(mockConfig);
         await client.initialize();
         vi.mocked(mockConfig.getGeminiClient).mockReturnValue(client);
@@ -226,18 +262,18 @@ describe('Gemini Client (client.ts)', () => {
         it('should create a new chat session, clearing the old history', async () => {
             // 1. Get the initial chat instance and add some history.
             const initialChat = client.getChat();
-            const initialHistory = await client.getHistory();
+            const initialHistory = client.getHistory();
             await client.addHistory({
                 role: 'user',
                 parts: [{ text: 'some old message' }],
             });
-            const historyWithOldMessage = await client.getHistory();
+            const historyWithOldMessage = client.getHistory();
             expect(historyWithOldMessage.length).toBeGreaterThan(initialHistory.length);
             // 2. Call resetChat.
             await client.resetChat();
             // 3. Get the new chat instance and its history.
             const newChat = client.getChat();
-            const newHistory = await client.getHistory();
+            const newHistory = client.getHistory();
             // 4. Assert that the chat instance is new and the history is reset.
             expect(newChat).not.toBe(initialChat);
             expect(newHistory.length).toBe(initialHistory.length);
@@ -384,6 +420,27 @@ describe('Gemini Client (client.ts)', () => {
                 expect(ChatCompressionService.prototype.compress).toHaveBeenLastCalledWith(expect.anything(), 'prompt-id-5', false, expect.anything(), expect.anything(), true);
             });
         });
+        it('should correctly latch hasFailedCompressionAttempt flag', async () => {
+            // 1. Setup: Call setup() from this test file
+            // This helper function mocks the compression service for us.
+            const { client } = setup({
+                originalTokenCount: 100,
+                newTokenCount: 200, // Inflated
+                compressionStatus: CompressionStatus.COMPRESSION_FAILED_INFLATED_TOKEN_COUNT,
+            });
+            // 2. Test Step 1: Trigger a non-forced failure
+            await client.tryCompressChat('prompt-1', false); // force = false
+            // 3. Assert Step 1: Check that the flag became true
+            // 3. Assert Step 1: Check that the flag became true
+            expect(client
+                .hasFailedCompressionAttempt).toBe(true);
+            // 4. Test Step 2: Trigger a forced failure
+            await client.tryCompressChat('prompt-2', true); // force = true
+            // 5. Assert Step 2: Check that the flag REMAINS true
+            // 5. Assert Step 2: Check that the flag REMAINS true
+            expect(client
+                .hasFailedCompressionAttempt).toBe(true);
+        });
         it('should not trigger summarization if token count is below threshold', async () => {
             const MOCKED_TOKEN_LIMIT = 1000;
             const originalTokenCount = MOCKED_TOKEN_LIMIT * 0.699;
@@ -677,6 +734,50 @@ ${JSON.stringify({
                 parts: expectedRequest,
             });
         });
+        it('should use local estimation for text-only requests and NOT call countTokens', async () => {
+            const request = [{ text: 'Hello world' }];
+            const generator = client['getContentGeneratorOrFail']();
+            const countTokensSpy = vi.spyOn(generator, 'countTokens');
+            const stream = client.sendMessageStream(request, new AbortController().signal, 'test-prompt-id');
+            await stream.next(); // Trigger the generator
+            expect(countTokensSpy).not.toHaveBeenCalled();
+        });
+        it('should use countTokens API for requests with non-text parts', async () => {
+            const request = [
+                { text: 'Describe this image' },
+                { inlineData: { mimeType: 'image/png', data: 'base64...' } },
+            ];
+            const generator = client['getContentGeneratorOrFail']();
+            const countTokensSpy = vi
+                .spyOn(generator, 'countTokens')
+                .mockResolvedValue({ totalTokens: 123 });
+            const stream = client.sendMessageStream(request, new AbortController().signal, 'test-prompt-id');
+            await stream.next(); // Trigger the generator
+            expect(countTokensSpy).toHaveBeenCalledWith(expect.objectContaining({
+                contents: expect.arrayContaining([
+                    expect.objectContaining({
+                        parts: expect.arrayContaining([
+                            { text: 'Describe this image' },
+                            { inlineData: { mimeType: 'image/png', data: 'base64...' } },
+                        ]),
+                    }),
+                ]),
+            }));
+        });
+        it('should estimate CJK characters more conservatively (closer to 1 token/char)', async () => {
+            const request = [{ text: '你好世界' }]; // 4 chars
+            const generator = client['getContentGeneratorOrFail']();
+            const countTokensSpy = vi.spyOn(generator, 'countTokens');
+            // 4 chars.
+            // Old logic: 4/4 = 1.
+            // New logic (heuristic): 4 * 1 = 4. (Or at least > 1).
+            // Let's assert it's roughly accurate.
+            const stream = client.sendMessageStream(request, new AbortController().signal, 'test-prompt-id');
+            await stream.next();
+            // Should NOT call countTokens (it's text only)
+            expect(countTokensSpy).not.toHaveBeenCalled();
+            // The actual token calculation is unit tested in tokenCalculation.test.ts
+        });
         it('should return the turn instance after the stream is complete', async () => {
             // Arrange
             const mockStream = (async function* () {
@@ -859,7 +960,8 @@ ${JSON.stringify({
             // A string of length 400 is roughly 100 tokens.
             const longText = 'a'.repeat(400);
             const request = [{ text: longText }];
-            const estimatedRequestTokenCount = Math.floor(JSON.stringify(request).length / 4);
+            // estimateTextOnlyLength counts only text content (400 chars), not JSON structure
+            const estimatedRequestTokenCount = Math.floor(longText.length / 4);
             const remainingTokenCount = MOCKED_TOKEN_LIMIT - lastPromptTokenCount;
             // Mock tryCompressChat to not compress
             vi.spyOn(client, 'tryCompressChat').mockResolvedValue({
@@ -905,7 +1007,8 @@ ${JSON.stringify({
             // We need a request > 95 tokens.
             const longText = 'a'.repeat(400);
             const request = [{ text: longText }];
-            const estimatedRequestTokenCount = Math.floor(JSON.stringify(request).length / 4);
+            // estimateTextOnlyLength counts only text content (400 chars), not JSON structure
+            const estimatedRequestTokenCount = Math.floor(longText.length / 4);
             const remainingTokenCount = STICKY_MODEL_LIMIT - lastPromptTokenCount;
             vi.spyOn(client, 'tryCompressChat').mockResolvedValue({
                 originalTokenCount: lastPromptTokenCount,
@@ -927,6 +1030,51 @@ ${JSON.stringify({
             expect(tokenLimit).toHaveBeenCalledWith(STICKY_MODEL);
             expect(mockTurnRunFn).not.toHaveBeenCalled();
         });
+        it('should not trigger overflow warning for requests with large binary data (PDFs/images)', async () => {
+            // Arrange
+            const MOCKED_TOKEN_LIMIT = 1000000; // 1M tokens
+            vi.mocked(tokenLimit).mockReturnValue(MOCKED_TOKEN_LIMIT);
+            const lastPromptTokenCount = 10000;
+            const mockChat = {
+                getLastPromptTokenCount: vi.fn().mockReturnValue(lastPromptTokenCount),
+                getHistory: vi.fn().mockReturnValue([]),
+            };
+            client['chat'] = mockChat;
+            // Simulate a PDF file with large base64 data (11MB when encoded)
+            // In the old implementation, this would incorrectly estimate ~2.7M tokens
+            // In the new implementation, only the text part is counted
+            const largePdfBase64 = 'A'.repeat(11 * 1024 * 1024);
+            const request = [
+                { text: 'Please analyze this PDF document' }, // ~35 chars = ~8 tokens
+                {
+                    inlineData: {
+                        mimeType: 'application/pdf',
+                        data: largePdfBase64, // This should be ignored in token estimation
+                    },
+                },
+            ];
+            // Mock tryCompressChat to not compress
+            vi.spyOn(client, 'tryCompressChat').mockResolvedValue({
+                originalTokenCount: lastPromptTokenCount,
+                newTokenCount: lastPromptTokenCount,
+                compressionStatus: CompressionStatus.NOOP,
+            });
+            // Mock Turn.run to simulate successful processing
+            const mockStream = (async function* () {
+                yield { type: 'content', value: 'Analysis complete' };
+            })();
+            mockTurnRunFn.mockReturnValue(mockStream);
+            // Act
+            const stream = client.sendMessageStream(request, new AbortController().signal, 'prompt-id-pdf-test');
+            const events = await fromAsync(stream);
+            // Assert
+            // Should NOT contain overflow warning
+            expect(events).not.toContainEqual(expect.objectContaining({
+                type: GeminiEventType.ContextWindowWillOverflow,
+            }));
+            // Turn.run should be called (processing should continue)
+            expect(mockTurnRunFn).toHaveBeenCalled();
+        });
         describe('Model Routing', () => {
             let mockRouterService;
             beforeEach(() => {
@@ -985,37 +1133,23 @@ ${JSON.stringify({
                 // Should use the newly routed model
                 expect(mockTurnRunFn).toHaveBeenCalledWith({ model: 'new-routed-model' }, [{ text: 'A new topic' }], expect.any(AbortSignal));
             });
-            it('should use the fallback model and bypass routing when in fallback mode', async () => {
-                vi.mocked(mockConfig.isInFallbackMode).mockReturnValue(true);
-                mockRouterService.route.mockResolvedValue({
-                    model: DEFAULT_GEMINI_FLASH_MODEL,
-                    reason: 'fallback',
-                });
-                const stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-1');
-                await fromAsync(stream);
-                expect(mockTurnRunFn).toHaveBeenCalledWith({ model: DEFAULT_GEMINI_FLASH_MODEL }, [{ text: 'Hi' }], expect.any(AbortSignal));
-            });
-            it('should stick to the fallback model for the entire sequence even if fallback mode ends', async () => {
-                // Start the sequence in fallback mode
-                vi.mocked(mockConfig.isInFallbackMode).mockReturnValue(true);
-                mockRouterService.route.mockResolvedValue({
-                    model: DEFAULT_GEMINI_FLASH_MODEL,
-                    reason: 'fallback',
-                });
-                let stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-fallback-stickiness');
-                await fromAsync(stream);
-                // First call should use fallback model
-                expect(mockTurnRunFn).toHaveBeenCalledWith({ model: DEFAULT_GEMINI_FLASH_MODEL }, [{ text: 'Hi' }], expect.any(AbortSignal));
-                // End fallback mode
-                vi.mocked(mockConfig.isInFallbackMode).mockReturnValue(false);
-                // Second call in the same sequence
-                stream = client.sendMessageStream([{ text: 'Continue' }], new AbortController().signal, 'prompt-fallback-stickiness');
-                await fromAsync(stream);
-                // Router should still not be called, and it should stick to the fallback model
-                expect(mockTurnRunFn).toHaveBeenCalledTimes(2); // Ensure it was called again
-                expect(mockTurnRunFn).toHaveBeenLastCalledWith({ model: DEFAULT_GEMINI_FLASH_MODEL }, // Still the fallback model
-                [{ text: 'Continue' }], expect.any(AbortSignal));
-            });
+        });
+        it('should use getGlobalMemory for system instruction when JIT is enabled', async () => {
+            vi.mocked(mockConfig.isJitContextEnabled).mockReturnValue(true);
+            vi.mocked(mockConfig.getGlobalMemory).mockReturnValue('Global JIT Memory');
+            vi.mocked(mockConfig.getUserMemory).mockReturnValue('Full JIT Memory');
+            const { getCoreSystemPrompt } = await import('./prompts.js');
+            const mockGetCoreSystemPrompt = vi.mocked(getCoreSystemPrompt);
+            await client.updateSystemInstruction();
+            expect(mockGetCoreSystemPrompt).toHaveBeenCalledWith(mockConfig, 'Global JIT Memory');
+        });
+        it('should use getUserMemory for system instruction when JIT is disabled', async () => {
+            vi.mocked(mockConfig.isJitContextEnabled).mockReturnValue(false);
+            vi.mocked(mockConfig.getUserMemory).mockReturnValue('Legacy Memory');
+            const { getCoreSystemPrompt } = await import('./prompts.js');
+            const mockGetCoreSystemPrompt = vi.mocked(getCoreSystemPrompt);
+            await client.updateSystemInstruction();
+            expect(mockGetCoreSystemPrompt).toHaveBeenCalledWith(mockConfig, 'Legacy Memory');
         });
         it('should recursively call sendMessageStream with "Please continue." when InvalidStream event is received', async () => {
             vi.spyOn(client['config'], 'getContinueOnFailedApiCall').mockReturnValue(true);
@@ -1045,6 +1179,7 @@ ${JSON.stringify({
             expect(events).toEqual([
                 { type: GeminiEventType.ModelInfo, value: 'default-routed-model' },
                 { type: GeminiEventType.InvalidStream },
+                { type: GeminiEventType.ModelInfo, value: 'default-routed-model' },
                 { type: GeminiEventType.Content, value: 'Continued content' },
             ]);
             // Verify that turn.run was called twice
@@ -1101,8 +1236,8 @@ ${JSON.stringify({
             const stream = client.sendMessageStream(initialRequest, signal, promptId);
             const events = await fromAsync(stream);
             // Assert
-            // We expect 3 events (model_info + original + 1 retry)
-            expect(events.length).toBe(3);
+            // We expect 4 events (model_info + original + model_info + 1 retry)
+            expect(events.length).toBe(4);
             expect(events
                 .filter((e) => e.type !== GeminiEventType.ModelInfo)
                 .every((e) => e.type === GeminiEventType.InvalidStream)).toBe(true);
@@ -1328,6 +1463,94 @@ ${JSON.stringify({
                 expect(contextJson.activeFile.path).toBe('/path/to/active/file.ts');
             });
         });
+        describe('Availability Service Integration', () => {
+            let mockAvailabilityService;
+            beforeEach(() => {
+                mockAvailabilityService = createAvailabilityServiceMock();
+                vi.mocked(mockConfig.getModelAvailabilityService).mockReturnValue(mockAvailabilityService);
+                vi.mocked(mockConfig.setActiveModel).mockClear();
+                mockRouterService.route.mockResolvedValue({
+                    model: 'model-a',
+                    reason: 'test',
+                });
+                vi.mocked(mockConfig.getModelRouterService).mockReturnValue(mockRouterService);
+                vi.spyOn(policyCatalog, 'getModelPolicyChain').mockReturnValue([
+                    {
+                        model: 'model-a',
+                        isLastResort: false,
+                        actions: {},
+                        stateTransitions: {},
+                    },
+                    {
+                        model: 'model-b',
+                        isLastResort: true,
+                        actions: {},
+                        stateTransitions: {},
+                    },
+                ]);
+                mockTurnRunFn.mockReturnValue((async function* () {
+                    yield { type: 'content', value: 'Hello' };
+                })());
+            });
+            it('should select first available model, set active, and not consume sticky attempt (done lower in chain)', async () => {
+                vi.mocked(mockAvailabilityService.selectFirstAvailable).mockReturnValue({
+                    selectedModel: 'model-a',
+                    attempts: 1,
+                    skipped: [],
+                });
+                vi.mocked(mockConfig.getModel).mockReturnValue(DEFAULT_GEMINI_MODEL_AUTO);
+                const stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-avail');
+                await fromAsync(stream);
+                expect(mockAvailabilityService.selectFirstAvailable).toHaveBeenCalledWith(['model-a', 'model-b']);
+                expect(mockConfig.setActiveModel).toHaveBeenCalledWith('model-a');
+                expect(mockAvailabilityService.consumeStickyAttempt).not.toHaveBeenCalled();
+                // Ensure turn.run used the selected model
+                expect(mockTurnRunFn).toHaveBeenCalledWith(expect.objectContaining({ model: 'model-a' }), expect.anything(), expect.anything());
+            });
+            it('should default to last resort model if selection returns null', async () => {
+                vi.mocked(mockAvailabilityService.selectFirstAvailable).mockReturnValue({
+                    selectedModel: null,
+                    skipped: [],
+                });
+                vi.mocked(mockConfig.getModel).mockReturnValue(DEFAULT_GEMINI_MODEL_AUTO);
+                const stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-avail-fallback');
+                await fromAsync(stream);
+                expect(mockConfig.setActiveModel).toHaveBeenCalledWith('model-b'); // Last resort
+                expect(mockAvailabilityService.consumeStickyAttempt).not.toHaveBeenCalled();
+            });
+            it('should reset turn on new message stream', async () => {
+                vi.mocked(mockAvailabilityService.selectFirstAvailable).mockReturnValue({
+                    selectedModel: 'model-a',
+                    skipped: [],
+                });
+                const stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-reset');
+                await fromAsync(stream);
+                expect(mockConfig.resetTurn).toHaveBeenCalled();
+            });
+            it('should NOT reset turn on invalid stream retry', async () => {
+                vi.mocked(mockAvailabilityService.selectFirstAvailable).mockReturnValue({
+                    selectedModel: 'model-a',
+                    skipped: [],
+                });
+                // We simulate a retry by calling sendMessageStream with isInvalidStreamRetry=true
+                // But the public API doesn't expose that argument directly unless we use the private method or simulate the recursion.
+                // We can simulate recursion by mocking turn run to return invalid stream once.
+                vi.spyOn(client['config'], 'getContinueOnFailedApiCall').mockReturnValue(true);
+                const mockStream1 = (async function* () {
+                    yield { type: GeminiEventType.InvalidStream };
+                })();
+                const mockStream2 = (async function* () {
+                    yield { type: 'content', value: 'ok' };
+                })();
+                mockTurnRunFn
+                    .mockReturnValueOnce(mockStream1)
+                    .mockReturnValueOnce(mockStream2);
+                const stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-retry');
+                await fromAsync(stream);
+                // resetTurn should be called once (for the initial call) but NOT for the recursive call
+                expect(mockConfig.resetTurn).toHaveBeenCalledTimes(1);
+            });
+        });
         describe('IDE context with pending tool calls', () => {
             let mockChat;
             beforeEach(() => {
@@ -1680,9 +1903,9 @@ ${JSON.stringify({
         it('should call generateContent with the correct parameters', async () => {
             const contents = [{ role: 'user', parts: [{ text: 'hello' }] }];
             const abortSignal = new AbortController().signal;
-            await client.generateContent({ model: DEFAULT_GEMINI_FLASH_MODEL }, contents, abortSignal);
+            await client.generateContent({ model: 'test-model' }, contents, abortSignal);
             expect(mockContentGenerator.generateContent).toHaveBeenCalledWith({
-                model: DEFAULT_GEMINI_FLASH_MODEL,
+                model: 'test-model',
                 config: {
                     abortSignal,
                     systemInstruction: getCoreSystemPrompt({}, ''),
@@ -1693,33 +1916,200 @@ ${JSON.stringify({
             }, 'test-session-id');
         });
         it('should use current model from config for content generation', async () => {
-            const initialModel = client['config'].getModel();
+            const initialModel = 'test-model';
             const contents = [{ role: 'user', parts: [{ text: 'test' }] }];
-            const currentModel = initialModel + '-changed';
-            vi.spyOn(client['config'], 'getModel').mockReturnValueOnce(currentModel);
-            await client.generateContent({ model: DEFAULT_GEMINI_FLASH_MODEL }, contents, new AbortController().signal);
-            expect(mockContentGenerator.generateContent).not.toHaveBeenCalledWith({
+            await client.generateContent({ model: initialModel }, contents, new AbortController().signal);
+            expect(mockContentGenerator.generateContent).toHaveBeenCalledWith(expect.objectContaining({
                 model: initialModel,
-                config: expect.any(Object),
-                contents,
-            });
-            expect(mockContentGenerator.generateContent).toHaveBeenCalledWith({
-                model: DEFAULT_GEMINI_FLASH_MODEL,
-                config: expect.any(Object),
-                contents,
-            }, 'test-session-id');
-        });
-        it('should use the Flash model when fallback mode is active', async () => {
-            const contents = [{ role: 'user', parts: [{ text: 'hello' }] }];
-            const abortSignal = new AbortController().signal;
-            const requestedModel = 'gemini-2.5-pro'; // A non-flash model
-            // Mock config to be in fallback mode
-            vi.spyOn(client['config'], 'isInFallbackMode').mockReturnValue(true);
-            await client.generateContent({ model: requestedModel }, contents, abortSignal);
-            expect(mockGenerateContentFn).toHaveBeenCalledWith(expect.objectContaining({
-                model: DEFAULT_GEMINI_FLASH_MODEL,
             }), 'test-session-id');
         });
+        describe('Hook System', () => {
+            let mockMessageBus;
+            beforeEach(() => {
+                vi.clearAllMocks();
+                mockMessageBus = { publish: vi.fn(), subscribe: vi.fn() };
+                // Force override config methods on the client instance
+                client['config'].getEnableHooks = vi.fn().mockReturnValue(true);
+                client['config'].getMessageBus = vi
+                    .fn()
+                    .mockReturnValue(mockMessageBus);
+            });
+            it('should fire BeforeAgent and AfterAgent exactly once for a simple turn', async () => {
+                const promptId = 'test-prompt-hook-1';
+                const request = { text: 'Hello Hooks' };
+                const signal = new AbortController().signal;
+                const { fireBeforeAgentHook, fireAfterAgentHook } = await import('./clientHookTriggers.js');
+                mockTurnRunFn.mockImplementation(async function* () {
+                    this.getResponseText.mockReturnValue('Hook Response');
+                    yield { type: GeminiEventType.Content, value: 'Hook Response' };
+                });
+                const stream = client.sendMessageStream(request, signal, promptId);
+                while (!(await stream.next()).done)
+                    ;
+                expect(fireBeforeAgentHook).toHaveBeenCalledTimes(1);
+                expect(fireAfterAgentHook).toHaveBeenCalledTimes(1);
+                expect(fireAfterAgentHook).toHaveBeenCalledWith(expect.anything(), request, 'Hook Response');
+                // Map should be empty
+                expect(client['hookStateMap'].size).toBe(0);
+            });
+            it('should fire BeforeAgent once and AfterAgent once even with recursion', async () => {
+                const { checkNextSpeaker } = await import('../utils/nextSpeakerChecker.js');
+                vi.mocked(checkNextSpeaker)
+                    .mockResolvedValueOnce({ next_speaker: 'model', reasoning: 'more' })
+                    .mockResolvedValueOnce(null);
+                const promptId = 'test-prompt-hook-recursive';
+                const request = { text: 'Recursion Test' };
+                const signal = new AbortController().signal;
+                const { fireBeforeAgentHook, fireAfterAgentHook } = await import('./clientHookTriggers.js');
+                let callCount = 0;
+                mockTurnRunFn.mockImplementation(async function* () {
+                    callCount++;
+                    const response = `Response ${callCount}`;
+                    this.getResponseText.mockReturnValue(response);
+                    yield { type: GeminiEventType.Content, value: response };
+                });
+                const stream = client.sendMessageStream(request, signal, promptId);
+                while (!(await stream.next()).done)
+                    ;
+                // BeforeAgent should fire ONLY once despite multiple internal turns
+                expect(fireBeforeAgentHook).toHaveBeenCalledTimes(1);
+                // AfterAgent should fire ONLY when the stack unwinds
+                expect(fireAfterAgentHook).toHaveBeenCalledTimes(1);
+                // Check cumulative response (separated by newline)
+                expect(fireAfterAgentHook).toHaveBeenCalledWith(expect.anything(), request, 'Response 1\nResponse 2');
+                expect(client['hookStateMap'].size).toBe(0);
+            });
+            it('should use original request in AfterAgent hook even when continuation happened', async () => {
+                const { checkNextSpeaker } = await import('../utils/nextSpeakerChecker.js');
+                vi.mocked(checkNextSpeaker)
+                    .mockResolvedValueOnce({ next_speaker: 'model', reasoning: 'more' })
+                    .mockResolvedValueOnce(null);
+                const promptId = 'test-prompt-hook-original-req';
+                const request = { text: 'Do something' };
+                const signal = new AbortController().signal;
+                const { fireAfterAgentHook } = await import('./clientHookTriggers.js');
+                mockTurnRunFn.mockImplementation(async function* () {
+                    this.getResponseText.mockReturnValue('Ok');
+                    yield { type: GeminiEventType.Content, value: 'Ok' };
+                });
+                const stream = client.sendMessageStream(request, signal, promptId);
+                while (!(await stream.next()).done)
+                    ;
+                expect(fireAfterAgentHook).toHaveBeenCalledWith(expect.anything(), request, // Should be 'Do something'
+                expect.stringContaining('Ok'));
+            });
+            it('should cleanup state when prompt_id changes', async () => {
+                const signal = new AbortController().signal;
+                mockTurnRunFn.mockImplementation(async function* () {
+                    this.getResponseText.mockReturnValue('Ok');
+                    yield { type: GeminiEventType.Content, value: 'Ok' };
+                });
+                client['hookStateMap'].set('old-id', {
+                    hasFiredBeforeAgent: true,
+                    cumulativeResponse: 'Old',
+                    activeCalls: 0,
+                    originalRequest: { text: 'Old' },
+                });
+                client['lastPromptId'] = 'old-id';
+                const stream = client.sendMessageStream({ text: 'New' }, signal, 'new-id');
+                await stream.next();
+                expect(client['hookStateMap'].has('old-id')).toBe(false);
+                expect(client['hookStateMap'].has('new-id')).toBe(true);
+            });
+            it('should stop execution in BeforeAgent when hook returns continue: false', async () => {
+                const { fireBeforeAgentHook } = await import('./clientHookTriggers.js');
+                vi.mocked(fireBeforeAgentHook).mockResolvedValue({
+                    shouldStopExecution: () => true,
+                    getEffectiveReason: () => 'Stopped by hook',
+                });
+                const mockChat = {
+                    addHistory: vi.fn(),
+                    getHistory: vi.fn().mockReturnValue([]),
+                    getLastPromptTokenCount: vi.fn(),
+                };
+                client['chat'] = mockChat;
+                const request = [{ text: 'Hello' }];
+                const stream = client.sendMessageStream(request, new AbortController().signal, 'test-prompt');
+                const events = await fromAsync(stream);
+                expect(events).toContainEqual({
+                    type: GeminiEventType.AgentExecutionStopped,
+                    value: { reason: 'Stopped by hook' },
+                });
+                expect(mockChat.addHistory).toHaveBeenCalledWith({
+                    role: 'user',
+                    parts: request,
+                });
+                expect(mockTurnRunFn).not.toHaveBeenCalled();
+            });
+            it('should block execution in BeforeAgent when hook returns decision: block', async () => {
+                const { fireBeforeAgentHook } = await import('./clientHookTriggers.js');
+                vi.mocked(fireBeforeAgentHook).mockResolvedValue({
+                    shouldStopExecution: () => false,
+                    isBlockingDecision: () => true,
+                    getEffectiveReason: () => 'Blocked by hook',
+                });
+                const mockChat = {
+                    addHistory: vi.fn(),
+                    getHistory: vi.fn().mockReturnValue([]),
+                    getLastPromptTokenCount: vi.fn(),
+                };
+                client['chat'] = mockChat;
+                const request = [{ text: 'Hello' }];
+                const stream = client.sendMessageStream(request, new AbortController().signal, 'test-prompt');
+                const events = await fromAsync(stream);
+                expect(events).toContainEqual({
+                    type: GeminiEventType.AgentExecutionBlocked,
+                    value: {
+                        reason: 'Blocked by hook',
+                    },
+                });
+                expect(mockChat.addHistory).not.toHaveBeenCalled();
+                expect(mockTurnRunFn).not.toHaveBeenCalled();
+            });
+            it('should stop execution in AfterAgent when hook returns continue: false', async () => {
+                const { fireAfterAgentHook } = await import('./clientHookTriggers.js');
+                vi.mocked(fireAfterAgentHook).mockResolvedValue({
+                    shouldStopExecution: () => true,
+                    getEffectiveReason: () => 'Stopped after agent',
+                });
+                mockTurnRunFn.mockImplementation(async function* () {
+                    yield { type: GeminiEventType.Content, value: 'Hello' };
+                });
+                const stream = client.sendMessageStream({ text: 'Hi' }, new AbortController().signal, 'test-prompt');
+                const events = await fromAsync(stream);
+                expect(events).toContainEqual({
+                    type: GeminiEventType.AgentExecutionStopped,
+                    value: { reason: 'Stopped after agent' },
+                });
+                // sendMessageStream should not recurse
+                expect(mockTurnRunFn).toHaveBeenCalledTimes(1);
+            });
+            it('should yield AgentExecutionBlocked and recurse in AfterAgent when hook returns decision: block', async () => {
+                const { fireAfterAgentHook } = await import('./clientHookTriggers.js');
+                vi.mocked(fireAfterAgentHook)
+                    .mockResolvedValueOnce({
+                    shouldStopExecution: () => false,
+                    isBlockingDecision: () => true,
+                    getEffectiveReason: () => 'Please explain',
+                })
+                    .mockResolvedValueOnce({
+                    shouldStopExecution: () => false,
+                    isBlockingDecision: () => false,
+                });
+                mockTurnRunFn.mockImplementation(async function* () {
+                    yield { type: GeminiEventType.Content, value: 'Response' };
+                });
+                const stream = client.sendMessageStream({ text: 'Hi' }, new AbortController().signal, 'test-prompt');
+                const events = await fromAsync(stream);
+                expect(events).toContainEqual({
+                    type: GeminiEventType.AgentExecutionBlocked,
+                    value: { reason: 'Please explain' },
+                });
+                // Should have called turn run twice (original + re-prompt)
+                expect(mockTurnRunFn).toHaveBeenCalledTimes(2);
+                expect(mockTurnRunFn).toHaveBeenNthCalledWith(2, expect.anything(), [{ text: 'Please explain' }], expect.anything());
+            });
+        });
     });
 });
 //# sourceMappingURL=client.test.js.map