@office-ai/aioncli-core 0.26.0 → 0.30.0

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

@@ -6,7 +6,7 @@
 import { describe, it, expect, vi, beforeEach, afterEach, } from 'vitest';
 import { GeminiClient } from './client.js';
 import { AuthType, } from './contentGenerator.js';
-import {} from './geminiChat.js';
+import { GeminiChat } from './geminiChat.js';
 import { CompressionStatus, GeminiEventType, Turn, } from './turn.js';
 import { getCoreSystemPrompt } from './prompts.js';
 import { DEFAULT_GEMINI_MODEL_AUTO } from '../config/models.js';
@@ -18,9 +18,10 @@ 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');
+import { LlmRole } from '../telemetry/types.js';
+import { partToString } from '../utils/partUtils.js';
+import { coreEvents } from '../utils/events.js';
 // Mock fs module to prevent actual file system operations during tests
 const mockFileSystem = new Map();
 vi.mock('node:fs', () => {
@@ -93,15 +94,11 @@ vi.mock('../telemetry/uiTelemetry.js', () => ({
     },
 }));
 vi.mock('../hooks/hookSystem.js');
-vi.mock('./clientHookTriggers.js', () => ({
-    fireBeforeAgentHook: vi.fn(),
-    fireAfterAgentHook: vi.fn().mockResolvedValue({
-        decision: 'allow',
-        continue: false,
-        suppressOutput: false,
-        systemMessage: undefined,
-    }),
-}));
+const mockHookSystem = {
+    fireBeforeAgentEvent: vi.fn().mockResolvedValue(undefined),
+    fireAfterAgentEvent: vi.fn().mockResolvedValue(undefined),
+    firePreCompressEvent: vi.fn().mockResolvedValue(undefined),
+};
 /**
  * Array.fromAsync ponyfill, which will be available in es 2024.
  *
@@ -124,14 +121,6 @@ describe('Gemini Client (client.ts)', () => {
         vi.resetAllMocks();
         ClearcutLogger.clearInstance();
         vi.mocked(uiTelemetryService.setLastPromptTokenCount).mockClear();
-        vi.mocked(ChatCompressionService.prototype.compress).mockResolvedValue({
-            newHistory: null,
-            info: {
-                originalTokenCount: 0,
-                newTokenCount: 0,
-                compressionStatus: CompressionStatus.NOOP,
-            },
-        });
         mockGenerateContentFn = vi.fn().mockResolvedValue({
             candidates: [{ content: { parts: [{ text: '{"key": "value"}' }] } }],
         });
@@ -176,6 +165,8 @@ describe('Gemini Client (client.ts)', () => {
             getGlobalMemory: vi.fn().mockReturnValue(''),
             getEnvironmentMemory: vi.fn().mockReturnValue(''),
             isJitContextEnabled: vi.fn().mockReturnValue(false),
+            getToolOutputMaskingEnabled: vi.fn().mockReturnValue(false),
+            getDisableLoopDetection: vi.fn().mockReturnValue(false),
             getSessionId: vi.fn().mockReturnValue('test-session-id'),
             getProxy: vi.fn().mockReturnValue(undefined),
             getWorkingDir: vi.fn().mockReturnValue('/test/dir'),
@@ -188,7 +179,6 @@ describe('Gemini Client (client.ts)', () => {
             getIdeModeFeature: vi.fn().mockReturnValue(false),
             getIdeMode: vi.fn().mockReturnValue(true),
             getDebugMode: vi.fn().mockReturnValue(false),
-            getPreviewFeatures: vi.fn().mockReturnValue(false),
             getWorkspaceContext: vi.fn().mockReturnValue({
                 getDirectories: vi.fn().mockReturnValue(['/test/dir']),
             }),
@@ -199,10 +189,12 @@ describe('Gemini Client (client.ts)', () => {
             getMessageBus: vi.fn().mockReturnValue(undefined),
             getEnableHooks: vi.fn().mockReturnValue(false),
             getChatCompression: vi.fn().mockReturnValue(undefined),
+            getCompressionThreshold: vi.fn().mockReturnValue(undefined),
             getSkipNextSpeakerCheck: vi.fn().mockReturnValue(false),
             getShowModelInfoInChat: vi.fn().mockReturnValue(false),
             getContinueOnFailedApiCall: vi.fn(),
             getProjectRoot: vi.fn().mockReturnValue('/test/project/root'),
+            getIncludeDirectoryTree: vi.fn().mockReturnValue(true),
             storage: {
                 getProjectTempDir: vi.fn().mockReturnValue('/test/temp'),
             },
@@ -233,21 +225,21 @@ describe('Gemini Client (client.ts)', () => {
                 .fn()
                 .mockReturnValue(createAvailabilityServiceMock()),
         };
-        mockConfig.getHookSystem = vi
-            .fn()
-            .mockReturnValue(new HookSystem(mockConfig));
+        mockConfig.getHookSystem = vi.fn().mockReturnValue(mockHookSystem);
         client = new GeminiClient(mockConfig);
         await client.initialize();
         vi.mocked(mockConfig.getGeminiClient).mockReturnValue(client);
         vi.mocked(uiTelemetryService.setLastPromptTokenCount).mockClear();
     });
     afterEach(() => {
+        client.dispose();
         vi.restoreAllMocks();
     });
     describe('addHistory', () => {
         it('should call chat.addHistory with the provided content', async () => {
             const mockChat = {
                 addHistory: vi.fn(),
+                setTools: vi.fn(),
             };
             client['chat'] = mockChat;
             const newContent = {
@@ -258,6 +250,24 @@ describe('Gemini Client (client.ts)', () => {
             expect(mockChat.addHistory).toHaveBeenCalledWith(newContent);
         });
     });
+    describe('setHistory', () => {
+        it('should update telemetry token count when history is set', () => {
+            const history = [
+                { role: 'user', parts: [{ text: 'some message' }] },
+            ];
+            client.setHistory(history);
+            expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalled();
+        });
+    });
+    describe('resumeChat', () => {
+        it('should update telemetry token count when a chat is resumed', async () => {
+            const history = [
+                { role: 'user', parts: [{ text: 'resumed message' }] },
+            ];
+            await client.resumeChat(history);
+            expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalled();
+        });
+    });
     describe('resetChat', () => {
         it('should create a new chat session, clearing the old history', async () => {
             // 1. Get the initial chat instance and add some history.
@@ -307,6 +317,7 @@ describe('Gemini Client (client.ts)', () => {
                 getHistory: mockGetHistory,
                 addHistory: vi.fn(),
                 setHistory: vi.fn(),
+                setTools: vi.fn(),
                 getLastPromptTokenCount: vi.fn(),
             };
         });
@@ -318,6 +329,10 @@ describe('Gemini Client (client.ts)', () => {
                 getHistory: vi.fn((_curated) => chatHistory),
                 setHistory: vi.fn(),
                 getLastPromptTokenCount: vi.fn().mockReturnValue(originalTokenCount),
+                getChatRecordingService: vi.fn().mockReturnValue({
+                    getConversation: vi.fn().mockReturnValue(null),
+                    getConversationFilePath: vi.fn().mockReturnValue(null),
+                }),
             };
             client['chat'] = mockOriginalChat;
             vi.mocked(uiTelemetryService.getLastPromptTokenCount).mockReturnValue(originalTokenCount);
@@ -325,7 +340,7 @@ describe('Gemini Client (client.ts)', () => {
                 { role: 'user', parts: [{ text: 'Summary' }] },
                 { role: 'model', parts: [{ text: 'Got it' }] },
             ];
-            vi.mocked(ChatCompressionService.prototype.compress).mockResolvedValue({
+            vi.spyOn(ChatCompressionService.prototype, 'compress').mockResolvedValue({
                 newHistory: compressionStatus === CompressionStatus.COMPRESSED
                     ? newHistory
                     : null,
@@ -444,7 +459,7 @@ describe('Gemini Client (client.ts)', () => {
         it('should not trigger summarization if token count is below threshold', async () => {
             const MOCKED_TOKEN_LIMIT = 1000;
             const originalTokenCount = MOCKED_TOKEN_LIMIT * 0.699;
-            vi.mocked(ChatCompressionService.prototype.compress).mockResolvedValue({
+            vi.spyOn(ChatCompressionService.prototype, 'compress').mockResolvedValue({
                 newHistory: null,
                 info: {
                     originalTokenCount,
@@ -476,6 +491,25 @@ describe('Gemini Client (client.ts)', () => {
                 newTokenCount: 50,
             });
         });
+        it('should resume the session file when compression succeeds', async () => {
+            const { client, mockOriginalChat } = setup({
+                compressionStatus: CompressionStatus.COMPRESSED,
+            });
+            const mockConversation = { some: 'conversation' };
+            const mockFilePath = '/tmp/session.json';
+            // Override the mock to return values
+            const mockRecordingService = {
+                getConversation: vi.fn().mockReturnValue(mockConversation),
+                getConversationFilePath: vi.fn().mockReturnValue(mockFilePath),
+            };
+            vi.mocked(mockOriginalChat.getChatRecordingService).mockReturnValue(mockRecordingService);
+            await client.tryCompressChat('prompt-id', false);
+            expect(client['startChat']).toHaveBeenCalledWith(expect.anything(), // newHistory
+            {
+                conversation: mockConversation,
+                filePath: mockFilePath,
+            });
+        });
     });
     describe('sendMessageStream', () => {
         it('emits a compression event when the context was automatically compressed', async () => {
@@ -498,6 +532,21 @@ describe('Gemini Client (client.ts)', () => {
                 value: compressionInfo,
             });
         });
+        it('does not emit ModelInfo event if signal is aborted', async () => {
+            // Arrange
+            mockTurnRunFn.mockReturnValue((async function* () {
+                yield { type: 'content', value: 'Hello' };
+            })());
+            const controller = new AbortController();
+            controller.abort();
+            // Act
+            const stream = client.sendMessageStream([{ text: 'Hi' }], controller.signal, 'prompt-id-1');
+            const events = await fromAsync(stream);
+            // Assert
+            expect(events).not.toContainEqual(expect.objectContaining({
+                type: GeminiEventType.ModelInfo,
+            }));
+        });
         it.each([
             {
                 compressionStatus: CompressionStatus.COMPRESSION_FAILED_INFLATED_TOKEN_COUNT,
@@ -558,6 +607,7 @@ describe('Gemini Client (client.ts)', () => {
             })());
             const mockChat = {
                 addHistory: vi.fn(),
+                setTools: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
                 getLastPromptTokenCount: vi.fn(),
             };
@@ -606,6 +656,7 @@ ${JSON.stringify({
             mockTurnRunFn.mockReturnValue(mockStream);
             const mockChat = {
                 addHistory: vi.fn(),
+                setTools: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
                 getLastPromptTokenCount: vi.fn(),
             };
@@ -618,7 +669,7 @@ ${JSON.stringify({
             }
             // Assert
             expect(ideContextStore.get).toHaveBeenCalled();
-            expect(mockTurnRunFn).toHaveBeenCalledWith({ model: 'default-routed-model' }, initialRequest, expect.any(AbortSignal));
+            expect(mockTurnRunFn).toHaveBeenCalledWith({ model: 'default-routed-model', isChatModel: true }, initialRequest, expect.any(AbortSignal), undefined);
         });
         it('should add context if ideMode is enabled and there is one active file', async () => {
             // Arrange
@@ -647,6 +698,7 @@ ${JSON.stringify({
             mockTurnRunFn.mockReturnValue(mockStream);
             const mockChat = {
                 addHistory: vi.fn(),
+                setTools: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
                 getLastPromptTokenCount: vi.fn(),
             };
@@ -708,6 +760,7 @@ ${JSON.stringify({
             mockTurnRunFn.mockReturnValue(mockStream);
             const mockChat = {
                 addHistory: vi.fn(),
+                setTools: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
                 getLastPromptTokenCount: vi.fn(),
             };
@@ -786,6 +839,7 @@ ${JSON.stringify({
             mockTurnRunFn.mockReturnValue(mockStream);
             const mockChat = {
                 addHistory: vi.fn(),
+                setTools: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
                 getLastPromptTokenCount: vi.fn(),
             };
@@ -820,6 +874,7 @@ ${JSON.stringify({
             mockTurnRunFn.mockReturnValue(mockStream);
             const mockChat = {
                 addHistory: vi.fn(),
+                setTools: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
                 getLastPromptTokenCount: vi.fn(),
             };
@@ -866,6 +921,7 @@ ${JSON.stringify({
             mockTurnRunFn.mockReturnValue(mockStream);
             const mockChat = {
                 addHistory: vi.fn(),
+                setTools: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
                 getLastPromptTokenCount: vi.fn(),
             };
@@ -905,6 +961,7 @@ ${JSON.stringify({
             mockTurnRunFn.mockReturnValue(mockStream);
             const mockChat = {
                 addHistory: vi.fn(),
+                setTools: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
                 getLastPromptTokenCount: vi.fn(),
             };
@@ -952,13 +1009,14 @@ ${JSON.stringify({
             const lastPromptTokenCount = 900;
             const mockChat = {
                 getLastPromptTokenCount: vi.fn().mockReturnValue(lastPromptTokenCount),
+                setTools: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
             };
             client['chat'] = mockChat;
-            // Remaining = 100. Threshold (95%) = 95.
-            // We need a request > 95 tokens.
-            // A string of length 400 is roughly 100 tokens.
-            const longText = 'a'.repeat(400);
+            // Remaining = 100.
+            // We need a request > 100 tokens.
+            // A string of length 404 is roughly 101 tokens.
+            const longText = 'a'.repeat(404);
             const request = [{ text: longText }];
             // estimateTextOnlyLength counts only text content (400 chars), not JSON structure
             const estimatedRequestTokenCount = Math.floor(longText.length / 4);
@@ -1000,12 +1058,13 @@ ${JSON.stringify({
             const lastPromptTokenCount = 900;
             const mockChat = {
                 getLastPromptTokenCount: vi.fn().mockReturnValue(lastPromptTokenCount),
+                setTools: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
             };
             client['chat'] = mockChat;
-            // Remaining (sticky) = 100. Threshold (95%) = 95.
-            // We need a request > 95 tokens.
-            const longText = 'a'.repeat(400);
+            // Remaining (sticky) = 100.
+            // We need a request > 100 tokens.
+            const longText = 'a'.repeat(404);
             const request = [{ text: longText }];
             // estimateTextOnlyLength counts only text content (400 chars), not JSON structure
             const estimatedRequestTokenCount = Math.floor(longText.length / 4);
@@ -1030,6 +1089,131 @@ ${JSON.stringify({
             expect(tokenLimit).toHaveBeenCalledWith(STICKY_MODEL);
             expect(mockTurnRunFn).not.toHaveBeenCalled();
         });
+        it('should attempt compression before overflow check and proceed if compression frees space', async () => {
+            // Arrange
+            const MOCKED_TOKEN_LIMIT = 1000;
+            vi.mocked(tokenLimit).mockReturnValue(MOCKED_TOKEN_LIMIT);
+            // Initial state: 950 tokens used, 50 remaining.
+            const initialTokenCount = 950;
+            // Request: 60 tokens. (950 + 60 = 1010 > 1000) -> Would overflow without compression.
+            const longText = 'a'.repeat(240); // 240 / 4 = 60 tokens
+            const request = [{ text: longText }];
+            // Use the real GeminiChat to manage state and token counts more realistically
+            const mockChatCompressed = {
+                getLastPromptTokenCount: vi.fn().mockReturnValue(400),
+                getHistory: vi
+                    .fn()
+                    .mockReturnValue([{ role: 'user', parts: [{ text: 'old' }] }]),
+                addHistory: vi.fn(),
+                setTools: vi.fn(),
+                getChatRecordingService: vi.fn().mockReturnValue({
+                    getConversation: vi.fn(),
+                    getConversationFilePath: vi.fn(),
+                }),
+            };
+            const mockChatInitial = {
+                getLastPromptTokenCount: vi.fn().mockReturnValue(initialTokenCount),
+                getHistory: vi
+                    .fn()
+                    .mockReturnValue([{ role: 'user', parts: [{ text: 'old' }] }]),
+                addHistory: vi.fn(),
+                setTools: vi.fn(),
+                getChatRecordingService: vi.fn().mockReturnValue({
+                    getConversation: vi.fn(),
+                    getConversationFilePath: vi.fn(),
+                }),
+            };
+            client['chat'] = mockChatInitial;
+            // Mock tryCompressChat to simulate successful compression
+            const tryCompressSpy = vi
+                .spyOn(client, 'tryCompressChat')
+                .mockImplementation(async () => {
+                // In reality, tryCompressChat replaces this.chat
+                client['chat'] = mockChatCompressed;
+                return {
+                    originalTokenCount: initialTokenCount,
+                    newTokenCount: 400,
+                    compressionStatus: CompressionStatus.COMPRESSED,
+                };
+            });
+            // Use a manual spy on Turn.prototype.run since Turn is a real class in this test context
+            // but mocked at the top of the file
+            mockTurnRunFn.mockImplementation(async function* () {
+                yield { type: 'content', value: 'Success after compression' };
+            });
+            // Act
+            const stream = client.sendMessageStream(request, new AbortController().signal, 'prompt-id-compression-test');
+            const events = await fromAsync(stream);
+            // Assert
+            // 1. Should NOT contain overflow warning
+            expect(events).not.toContainEqual(expect.objectContaining({
+                type: GeminiEventType.ContextWindowWillOverflow,
+            }));
+            // 2. Should contain compression event
+            expect(events).toContainEqual(expect.objectContaining({
+                type: GeminiEventType.ChatCompressed,
+            }));
+            // 3. Should have called tryCompressChat
+            expect(tryCompressSpy).toHaveBeenCalled();
+            // 4. Should have called Turn.run (proceeded with the request)
+            expect(mockTurnRunFn).toHaveBeenCalled();
+        });
+        it('should handle massive function responses by truncating them and then yielding overflow warning', async () => {
+            // Arrange
+            const MOCKED_TOKEN_LIMIT = 1000;
+            vi.mocked(tokenLimit).mockReturnValue(MOCKED_TOKEN_LIMIT);
+            // History has a large compressible part and a massive function response at the end.
+            const massiveText = 'a'.repeat(200000);
+            const history = [
+                { role: 'user', parts: [{ text: 'a'.repeat(100000) }] }, // compressible part
+                { role: 'model', parts: [{ text: 'ok' }] },
+                {
+                    role: 'model',
+                    parts: [{ functionCall: { name: 'huge_tool', args: {} } }],
+                },
+                {
+                    role: 'user',
+                    parts: [
+                        {
+                            functionResponse: {
+                                name: 'huge_tool',
+                                response: { data: massiveText },
+                            },
+                        },
+                    ],
+                },
+            ];
+            const realChat = new GeminiChat(mockConfig, '', [], history);
+            client['chat'] = realChat;
+            // Use a realistic mock for compression that simulates the 40k truncation effect.
+            // We spy on the instance directly to ensure it intercepts correctly.
+            const compressSpy = vi
+                .spyOn(client['compressionService'], 'compress')
+                .mockResolvedValue({
+                newHistory: history, // Keep history large for the overflow check
+                info: {
+                    originalTokenCount: 50000,
+                    newTokenCount: 10000, // Reduced from 50k but still > 1000 limit
+                    compressionStatus: CompressionStatus.COMPRESSED,
+                },
+            });
+            // The new request
+            const request = [{ text: 'next question' }];
+            // Act
+            const stream = client.sendMessageStream(request, new AbortController().signal, 'prompt-id-massive-test');
+            const events = await fromAsync(stream);
+            // Assert
+            // 1. Should have attempted compression
+            expect(compressSpy).toHaveBeenCalled();
+            // 2. Should yield overflow warning because 10000 > 1000 limit.
+            expect(events).toContainEqual(expect.objectContaining({
+                type: GeminiEventType.ContextWindowWillOverflow,
+                value: expect.objectContaining({
+                    estimatedRequestTokenCount: expect.any(Number),
+                    remainingTokenCount: expect.any(Number),
+                }),
+            }));
+        });
         it('should not trigger overflow warning for requests with large binary data (PDFs/images)', async () => {
             // Arrange
             const MOCKED_TOKEN_LIMIT = 1000000; // 1M tokens
@@ -1037,6 +1221,7 @@ ${JSON.stringify({
             const lastPromptTokenCount = 10000;
             const mockChat = {
                 getLastPromptTokenCount: vi.fn().mockReturnValue(lastPromptTokenCount),
+                setTools: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
             };
             client['chat'] = mockChat;
@@ -1089,6 +1274,7 @@ ${JSON.stringify({
                 })());
                 const mockChat = {
                     addHistory: vi.fn(),
+                    setTools: vi.fn(),
                     getHistory: vi.fn().mockReturnValue([]),
                     getLastPromptTokenCount: vi.fn(),
                 };
@@ -1099,28 +1285,28 @@ ${JSON.stringify({
                 await fromAsync(stream); // consume stream
                 expect(mockConfig.getModelRouterService).toHaveBeenCalled();
                 expect(mockRouterService.route).toHaveBeenCalled();
-                expect(mockTurnRunFn).toHaveBeenCalledWith({ model: 'routed-model' }, [{ text: 'Hi' }], expect.any(AbortSignal));
+                expect(mockTurnRunFn).toHaveBeenCalledWith({ model: 'routed-model', isChatModel: true }, [{ text: 'Hi' }], expect.any(AbortSignal), undefined);
             });
             it('should use the same model for subsequent turns in the same prompt (stickiness)', async () => {
                 // First turn
                 let stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-1');
                 await fromAsync(stream);
                 expect(mockRouterService.route).toHaveBeenCalledTimes(1);
-                expect(mockTurnRunFn).toHaveBeenCalledWith({ model: 'routed-model' }, [{ text: 'Hi' }], expect.any(AbortSignal));
+                expect(mockTurnRunFn).toHaveBeenCalledWith({ model: 'routed-model', isChatModel: true }, [{ text: 'Hi' }], expect.any(AbortSignal), undefined);
                 // Second turn
                 stream = client.sendMessageStream([{ text: 'Continue' }], new AbortController().signal, 'prompt-1');
                 await fromAsync(stream);
                 // Router should not be called again
                 expect(mockRouterService.route).toHaveBeenCalledTimes(1);
                 // Should stick to the first model
-                expect(mockTurnRunFn).toHaveBeenCalledWith({ model: 'routed-model' }, [{ text: 'Continue' }], expect.any(AbortSignal));
+                expect(mockTurnRunFn).toHaveBeenCalledWith({ model: 'routed-model', isChatModel: true }, [{ text: 'Continue' }], expect.any(AbortSignal), undefined);
             });
             it('should reset the sticky model and re-route when the prompt_id changes', async () => {
                 // First prompt
                 let stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-1');
                 await fromAsync(stream);
                 expect(mockRouterService.route).toHaveBeenCalledTimes(1);
-                expect(mockTurnRunFn).toHaveBeenCalledWith({ model: 'routed-model' }, [{ text: 'Hi' }], expect.any(AbortSignal));
+                expect(mockTurnRunFn).toHaveBeenCalledWith({ model: 'routed-model', isChatModel: true }, [{ text: 'Hi' }], expect.any(AbortSignal), undefined);
                 // New prompt
                 mockRouterService.route.mockResolvedValue({
                     model: 'new-routed-model',
@@ -1131,7 +1317,31 @@ ${JSON.stringify({
                 // Router should be called again for the new prompt
                 expect(mockRouterService.route).toHaveBeenCalledTimes(2);
                 // Should use the newly routed model
-                expect(mockTurnRunFn).toHaveBeenCalledWith({ model: 'new-routed-model' }, [{ text: 'A new topic' }], expect.any(AbortSignal));
+                expect(mockTurnRunFn).toHaveBeenCalledWith({ model: 'new-routed-model', isChatModel: true }, [{ text: 'A new topic' }], expect.any(AbortSignal), undefined);
+            });
+            it('should re-route within the same prompt when the configured model changes', async () => {
+                mockTurnRunFn.mockClear();
+                mockTurnRunFn.mockImplementation(async function* () {
+                    yield { type: 'content', value: 'Hello' };
+                });
+                mockRouterService.route.mockResolvedValueOnce({
+                    model: 'original-model',
+                    reason: 'test',
+                });
+                let stream = client.sendMessageStream([{ text: 'Hi' }], new AbortController().signal, 'prompt-1');
+                await fromAsync(stream);
+                expect(mockRouterService.route).toHaveBeenCalledTimes(1);
+                expect(mockTurnRunFn).toHaveBeenNthCalledWith(1, { model: 'original-model', isChatModel: true }, [{ text: 'Hi' }], expect.any(AbortSignal), undefined);
+                mockRouterService.route.mockResolvedValue({
+                    model: 'fallback-model',
+                    reason: 'test',
+                });
+                vi.mocked(mockConfig.getModel).mockReturnValue('gemini-2.5-flash');
+                coreEvents.emitModelChanged('gemini-2.5-flash');
+                stream = client.sendMessageStream([{ text: 'Continue' }], new AbortController().signal, 'prompt-1');
+                await fromAsync(stream);
+                expect(mockRouterService.route).toHaveBeenCalledTimes(2);
+                expect(mockTurnRunFn).toHaveBeenNthCalledWith(2, { model: 'fallback-model', isChatModel: true }, [{ text: 'Continue' }], expect.any(AbortSignal), undefined);
             });
         });
         it('should use getGlobalMemory for system instruction when JIT is enabled', async () => {
@@ -1140,15 +1350,15 @@ ${JSON.stringify({
             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');
+            client.updateSystemInstruction();
+            expect(mockGetCoreSystemPrompt).toHaveBeenCalledWith(mockConfig, 'Full 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();
+            client.updateSystemInstruction();
             expect(mockGetCoreSystemPrompt).toHaveBeenCalledWith(mockConfig, 'Legacy Memory');
         });
         it('should recursively call sendMessageStream with "Please continue." when InvalidStream event is received', async () => {
@@ -1165,6 +1375,7 @@ ${JSON.stringify({
                 .mockReturnValueOnce(mockStream2);
             const mockChat = {
                 addHistory: vi.fn(),
+                setTools: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
                 getLastPromptTokenCount: vi.fn(),
             };
@@ -1179,15 +1390,14 @@ ${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
             expect(mockTurnRunFn).toHaveBeenCalledTimes(2);
             // First call with original request
-            expect(mockTurnRunFn).toHaveBeenNthCalledWith(1, { model: 'default-routed-model' }, initialRequest, expect.any(AbortSignal));
+            expect(mockTurnRunFn).toHaveBeenNthCalledWith(1, { model: 'default-routed-model', isChatModel: true }, initialRequest, expect.any(AbortSignal), undefined);
             // Second call with "Please continue."
-            expect(mockTurnRunFn).toHaveBeenNthCalledWith(2, { model: 'default-routed-model' }, [{ text: 'System: Please continue.' }], expect.any(AbortSignal));
+            expect(mockTurnRunFn).toHaveBeenNthCalledWith(2, { model: 'default-routed-model', isChatModel: true }, [{ text: 'System: Please continue.' }], expect.any(AbortSignal), undefined);
         });
         it('should not recursively call sendMessageStream with "Please continue." when InvalidStream event is received and flag is false', async () => {
             vi.spyOn(client['config'], 'getContinueOnFailedApiCall').mockReturnValue(false);
@@ -1198,6 +1408,7 @@ ${JSON.stringify({
             mockTurnRunFn.mockReturnValueOnce(mockStream1);
             const mockChat = {
                 addHistory: vi.fn(),
+                setTools: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
                 getLastPromptTokenCount: vi.fn(),
             };
@@ -1225,6 +1436,7 @@ ${JSON.stringify({
             })());
             const mockChat = {
                 addHistory: vi.fn(),
+                setTools: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
                 getLastPromptTokenCount: vi.fn(),
             };
@@ -1236,11 +1448,11 @@ ${JSON.stringify({
             const stream = client.sendMessageStream(initialRequest, signal, promptId);
             const events = await fromAsync(stream);
             // Assert
-            // We expect 4 events (model_info + original + model_info + 1 retry)
-            expect(events.length).toBe(4);
+            // We expect 3 events (model_info + original + 1 retry)
+            expect(events.length).toBe(3);
             expect(events
-                .filter((e) => e.type !== GeminiEventType.ModelInfo)
-                .every((e) => e.type === GeminiEventType.InvalidStream)).toBe(true);
+                .filter((e) => e.type === GeminiEventType.ModelInfo)
+                .map((e) => e.value)).toEqual(['default-routed-model']);
             // Verify that turn.run was called twice
             expect(mockTurnRunFn).toHaveBeenCalledTimes(2);
         });
@@ -1260,6 +1472,7 @@ ${JSON.stringify({
                 const mockChat = {
                     addHistory: vi.fn(),
                     setHistory: vi.fn(),
+                    setTools: vi.fn(),
                     // Assume history is not empty for delta checks
                     getHistory: vi
                         .fn()
@@ -1505,7 +1718,7 @@ ${JSON.stringify({
                 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());
+                expect(mockTurnRunFn).toHaveBeenCalledWith(expect.objectContaining({ model: 'model-a' }), expect.anything(), expect.anything(), undefined);
             });
             it('should default to last resort model if selection returns null', async () => {
                 vi.mocked(mockAvailabilityService.selectFirstAvailable).mockReturnValue({
@@ -1567,6 +1780,7 @@ ${JSON.stringify({
                     addHistory: vi.fn(),
                     getHistory: vi.fn().mockReturnValue([]), // Default empty history
                     setHistory: vi.fn(),
+                    setTools: vi.fn(),
                     getLastPromptTokenCount: vi.fn(),
                 };
                 client['chat'] = mockChat;
@@ -1830,6 +2044,7 @@ ${JSON.stringify({
             mockTurnRunFn.mockReturnValue(mockStream);
             const mockChat = {
                 addHistory: vi.fn(),
+                setTools: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
                 getLastPromptTokenCount: vi.fn(),
             };
@@ -1856,6 +2071,7 @@ ${JSON.stringify({
             mockTurnRunFn.mockReturnValue(mockStream);
             const mockChat = {
                 addHistory: vi.fn(),
+                setTools: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
                 getLastPromptTokenCount: vi.fn(),
             };
@@ -1884,6 +2100,7 @@ ${JSON.stringify({
             });
             const mockChat = {
                 addHistory: vi.fn(),
+                setTools: vi.fn(),
                 getHistory: vi.fn().mockReturnValue([]),
                 getLastPromptTokenCount: vi.fn(),
             };
@@ -1903,7 +2120,7 @@ ${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: 'test-model' }, contents, abortSignal);
+            await client.generateContent({ model: 'test-model' }, contents, abortSignal, LlmRole.MAIN);
             expect(mockContentGenerator.generateContent).toHaveBeenCalledWith({
                 model: 'test-model',
                 config: {
@@ -1913,15 +2130,15 @@ ${JSON.stringify({
                     topP: 1,
                 },
                 contents,
-            }, 'test-session-id');
+            }, 'test-session-id', LlmRole.MAIN);
         });
         it('should use current model from config for content generation', async () => {
             const initialModel = 'test-model';
             const contents = [{ role: 'user', parts: [{ text: 'test' }] }];
-            await client.generateContent({ model: initialModel }, contents, new AbortController().signal);
+            await client.generateContent({ model: initialModel }, contents, new AbortController().signal, LlmRole.MAIN);
             expect(mockContentGenerator.generateContent).toHaveBeenCalledWith(expect.objectContaining({
                 model: initialModel,
-            }), 'test-session-id');
+            }), 'test-session-id', LlmRole.MAIN);
         });
         describe('Hook System', () => {
             let mockMessageBus;
@@ -1938,7 +2155,6 @@ ${JSON.stringify({
                 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' };
@@ -1946,9 +2162,9 @@ ${JSON.stringify({
                 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');
+                expect(mockHookSystem.fireBeforeAgentEvent).toHaveBeenCalledTimes(1);
+                expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledTimes(1);
+                expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledWith(partToString(request), 'Hook Response');
                 // Map should be empty
                 expect(client['hookStateMap'].size).toBe(0);
             });
@@ -1960,7 +2176,6 @@ ${JSON.stringify({
                 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++;
@@ -1972,11 +2187,11 @@ ${JSON.stringify({
                 while (!(await stream.next()).done)
                     ;
                 // BeforeAgent should fire ONLY once despite multiple internal turns
-                expect(fireBeforeAgentHook).toHaveBeenCalledTimes(1);
+                expect(mockHookSystem.fireBeforeAgentEvent).toHaveBeenCalledTimes(1);
                 // AfterAgent should fire ONLY when the stack unwinds
-                expect(fireAfterAgentHook).toHaveBeenCalledTimes(1);
+                expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledTimes(1);
                 // Check cumulative response (separated by newline)
-                expect(fireAfterAgentHook).toHaveBeenCalledWith(expect.anything(), request, 'Response 1\nResponse 2');
+                expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledWith(partToString(request), 'Response 1\nResponse 2');
                 expect(client['hookStateMap'].size).toBe(0);
             });
             it('should use original request in AfterAgent hook even when continuation happened', async () => {
@@ -1987,7 +2202,6 @@ ${JSON.stringify({
                 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' };
@@ -1995,7 +2209,7 @@ ${JSON.stringify({
                 const stream = client.sendMessageStream(request, signal, promptId);
                 while (!(await stream.next()).done)
                     ;
-                expect(fireAfterAgentHook).toHaveBeenCalledWith(expect.anything(), request, // Should be 'Do something'
+                expect(mockHookSystem.fireAfterAgentEvent).toHaveBeenCalledWith(partToString(request), // Should be 'Do something'
                 expect.stringContaining('Ok'));
             });
             it('should cleanup state when prompt_id changes', async () => {
@@ -2017,13 +2231,14 @@ ${JSON.stringify({
                 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({
+                mockHookSystem.fireBeforeAgentEvent.mockResolvedValue({
                     shouldStopExecution: () => true,
                     getEffectiveReason: () => 'Stopped by hook',
+                    systemMessage: undefined,
                 });
                 const mockChat = {
                     addHistory: vi.fn(),
+                    setTools: vi.fn(),
                     getHistory: vi.fn().mockReturnValue([]),
                     getLastPromptTokenCount: vi.fn(),
                 };
@@ -2042,14 +2257,15 @@ ${JSON.stringify({
                 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({
+                mockHookSystem.fireBeforeAgentEvent.mockResolvedValue({
                     shouldStopExecution: () => false,
                     isBlockingDecision: () => true,
                     getEffectiveReason: () => 'Blocked by hook',
+                    systemMessage: undefined,
                 });
                 const mockChat = {
                     addHistory: vi.fn(),
+                    setTools: vi.fn(),
                     getHistory: vi.fn().mockReturnValue([]),
                     getLastPromptTokenCount: vi.fn(),
                 };
@@ -2067,47 +2283,85 @@ ${JSON.stringify({
                 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({
+                mockHookSystem.fireAfterAgentEvent.mockResolvedValue({
                     shouldStopExecution: () => true,
                     getEffectiveReason: () => 'Stopped after agent',
+                    shouldClearContext: () => false,
+                    systemMessage: undefined,
                 });
                 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({
+                expect(events).toContainEqual(expect.objectContaining({
                     type: GeminiEventType.AgentExecutionStopped,
-                    value: { reason: 'Stopped after agent' },
-                });
+                    value: expect.objectContaining({ 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)
+                mockHookSystem.fireAfterAgentEvent
                     .mockResolvedValueOnce({
                     shouldStopExecution: () => false,
                     isBlockingDecision: () => true,
                     getEffectiveReason: () => 'Please explain',
+                    shouldClearContext: () => false,
+                    systemMessage: undefined,
                 })
                     .mockResolvedValueOnce({
                     shouldStopExecution: () => false,
                     isBlockingDecision: () => false,
+                    shouldClearContext: () => false,
+                    systemMessage: undefined,
                 });
                 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({
+                expect(events).toContainEqual(expect.objectContaining({
                     type: GeminiEventType.AgentExecutionBlocked,
-                    value: { reason: 'Please explain' },
-                });
+                    value: expect.objectContaining({ 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());
+                expect(mockTurnRunFn).toHaveBeenNthCalledWith(2, expect.anything(), [{ text: 'Please explain' }], expect.anything(), undefined);
+            });
+            it('should call resetChat when AfterAgent hook returns shouldClearContext: true', async () => {
+                const resetChatSpy = vi
+                    .spyOn(client, 'resetChat')
+                    .mockResolvedValue(undefined);
+                mockHookSystem.fireAfterAgentEvent
+                    .mockResolvedValueOnce({
+                    shouldStopExecution: () => false,
+                    isBlockingDecision: () => true,
+                    getEffectiveReason: () => 'Blocked and clearing context',
+                    shouldClearContext: () => true,
+                    systemMessage: undefined,
+                })
+                    .mockResolvedValueOnce({
+                    shouldStopExecution: () => false,
+                    isBlockingDecision: () => false,
+                    shouldClearContext: () => false,
+                    systemMessage: undefined,
+                });
+                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: 'Blocked and clearing context',
+                        systemMessage: undefined,
+                        contextCleared: true,
+                    },
+                });
+                expect(resetChatSpy).toHaveBeenCalledTimes(1);
+                resetChatSpy.mockRestore();
             });
         });
     });