@office-ai/aioncli-core 0.26.0 → 0.30.0

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

@@ -12,11 +12,11 @@ import { AuthType } from './contentGenerator.js';
 import { TerminalQuotaError } from '../utils/googleQuotaErrors.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';
+import { LlmRole } from '../telemetry/types.js';
 // Mock fs module to prevent actual file system operations during tests
 const mockFileSystem = new Map();
 vi.mock('node:fs', () => {
@@ -105,7 +105,6 @@ describe('GeminiChat', () => {
             getTelemetryLogPromptsEnabled: () => true,
             getUsageStatisticsEnabled: () => true,
             getDebugMode: () => false,
-            getPreviewFeatures: () => false,
             getContentGeneratorConfig: vi.fn().mockImplementation(() => ({
                 authType: 'oauth-personal',
                 model: currentModel,
@@ -165,9 +164,7 @@ describe('GeminiChat', () => {
         setSimulate429(false);
         // Reset history for each test by creating a new instance
         chat = new GeminiChat(mockConfig);
-        mockConfig.getHookSystem = vi
-            .fn()
-            .mockReturnValue(new HookSystem(mockConfig));
+        mockConfig.getHookSystem = vi.fn().mockReturnValue(undefined);
     });
     afterEach(() => {
         vi.restoreAllMocks();
@@ -190,6 +187,27 @@ describe('GeminiChat', () => {
             expect(chatEmpty.getLastPromptTokenCount()).toBe(0);
         });
     });
+    describe('setHistory', () => {
+        it('should recalculate lastPromptTokenCount when history is updated', () => {
+            const initialHistory = [
+                { role: 'user', parts: [{ text: 'Hello' }] },
+            ];
+            const chatWithHistory = new GeminiChat(mockConfig, '', [], initialHistory);
+            const initialCount = chatWithHistory.getLastPromptTokenCount();
+            const newHistory = [
+                {
+                    role: 'user',
+                    parts: [
+                        {
+                            text: 'This is a much longer history item that should result in more tokens than just hello.',
+                        },
+                    ],
+                },
+            ];
+            chatWithHistory.setHistory(newHistory);
+            expect(chatWithHistory.getLastPromptTokenCount()).toBeGreaterThan(initialCount);
+        });
+    });
     describe('sendMessageStream', () => {
         it('should succeed if a tool call is followed by an empty part', async () => {
             // 1. Mock a stream that contains a tool call, then an invalid (empty) part.
@@ -219,7 +237,7 @@ describe('GeminiChat', () => {
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithToolCall);
             // 2. Action & Assert: The stream processing should complete without throwing an error
             // because the presence of a tool call makes the empty final chunk acceptable.
-            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-tool-call-empty-end', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-tool-call-empty-end', new AbortController().signal, LlmRole.MAIN);
             await expect((async () => {
                 for await (const _ of stream) {
                     /* consume stream */
@@ -259,7 +277,7 @@ describe('GeminiChat', () => {
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithNoFinish);
             // 2. Action & Assert: The stream should fail because there's no finish reason.
-            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-no-finish-empty-end', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-no-finish-empty-end', new AbortController().signal, LlmRole.MAIN);
             await expect((async () => {
                 for await (const _ of stream) {
                     /* consume stream */
@@ -294,7 +312,7 @@ describe('GeminiChat', () => {
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithInvalidEnd);
             // 2. Action & Assert: The stream should complete without throwing an error.
-            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-valid-then-invalid-end', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-valid-then-invalid-end', new AbortController().signal, LlmRole.MAIN);
             await expect((async () => {
                 for await (const _ of stream) {
                     /* consume stream */
@@ -329,7 +347,7 @@ describe('GeminiChat', () => {
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(multiChunkStream);
             // 2. Action: Send a message and consume the stream.
-            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-empty-chunk-consolidation', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-empty-chunk-consolidation', new AbortController().signal, LlmRole.MAIN);
             for await (const _ of stream) {
                 // Consume the stream
             }
@@ -377,7 +395,7 @@ describe('GeminiChat', () => {
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(multiChunkStream);
             // 2. Action: Send a message and consume the stream.
-            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-multi-chunk', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-multi-chunk', new AbortController().signal, LlmRole.MAIN);
             for await (const _ of stream) {
                 // Consume the stream to trigger history recording.
             }
@@ -413,7 +431,7 @@ describe('GeminiChat', () => {
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(mixedContentStream);
             // 2. Action: Send a message and fully consume the stream to trigger history recording.
-            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-mixed-chunk', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-mixed-chunk', new AbortController().signal, LlmRole.MAIN);
             for await (const _ of stream) {
                 // This loop consumes the stream.
             }
@@ -467,7 +485,7 @@ describe('GeminiChat', () => {
                     name: 'find_restaurant',
                     response: { name: 'Vesuvio' },
                 },
-            }, 'prompt-id-stream-1', new AbortController().signal);
+            }, 'prompt-id-stream-1', new AbortController().signal, LlmRole.MAIN);
             // 4. Assert: The stream processing should throw an InvalidStreamError.
             await expect((async () => {
                 for await (const _ of stream) {
@@ -498,7 +516,7 @@ describe('GeminiChat', () => {
                 };
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithToolCall);
-            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-1', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-1', new AbortController().signal, LlmRole.MAIN);
             // Should not throw an error
             await expect((async () => {
                 for await (const _ of stream) {
@@ -522,7 +540,7 @@ describe('GeminiChat', () => {
                 };
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithoutFinishReason);
-            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-1', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-1', new AbortController().signal, LlmRole.MAIN);
             await expect((async () => {
                 for await (const _ of stream) {
                     // consume stream
@@ -545,7 +563,7 @@ describe('GeminiChat', () => {
                 };
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithEmptyResponse);
-            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-1', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-1', new AbortController().signal, LlmRole.MAIN);
             await expect((async () => {
                 for await (const _ of stream) {
                     // consume stream
@@ -568,7 +586,7 @@ describe('GeminiChat', () => {
                 };
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(validStream);
-            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-1', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-1', new AbortController().signal, LlmRole.MAIN);
             // Should not throw an error
             await expect((async () => {
                 for await (const _ of stream) {
@@ -592,7 +610,7 @@ describe('GeminiChat', () => {
                 };
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(streamWithMalformedFunctionCall);
-            const stream = await chat.sendMessageStream({ model: 'gemini-2.5-pro' }, 'test', 'prompt-id-malformed', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.5-pro' }, 'test', 'prompt-id-malformed', new AbortController().signal, LlmRole.MAIN);
             // Should throw an error
             await expect((async () => {
                 for await (const _ of stream) {
@@ -626,7 +644,7 @@ describe('GeminiChat', () => {
                 };
             })());
             // 2. Send a message
-            const stream = await chat.sendMessageStream({ model: 'gemini-2.5-pro' }, 'test retry', 'prompt-id-retry-malformed', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.5-pro' }, 'test retry', 'prompt-id-retry-malformed', new AbortController().signal, LlmRole.MAIN);
             const events = [];
             for await (const event of stream) {
                 events.push(event);
@@ -664,7 +682,7 @@ describe('GeminiChat', () => {
                 };
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(response);
-            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'hello', 'prompt-id-1', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'hello', 'prompt-id-1', new AbortController().signal, LlmRole.MAIN);
             for await (const _ of stream) {
                 // consume stream
             }
@@ -685,7 +703,7 @@ describe('GeminiChat', () => {
                     },
                     abortSignal: expect.any(AbortSignal),
                 },
-            }, 'prompt-id-1');
+            }, 'prompt-id-1', LlmRole.MAIN);
         });
         it('should use thinkingLevel and remove thinkingBudget for gemini-3 models', async () => {
             const response = (async function* () {
@@ -699,7 +717,7 @@ describe('GeminiChat', () => {
                 };
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(response);
-            const stream = await chat.sendMessageStream({ model: 'gemini-3-test-only-model-string-for-testing' }, 'hello', 'prompt-id-thinking-level', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'gemini-3-test-only-model-string-for-testing' }, 'hello', 'prompt-id-thinking-level', new AbortController().signal, LlmRole.MAIN);
             for await (const _ of stream) {
                 // consume stream
             }
@@ -711,7 +729,7 @@ describe('GeminiChat', () => {
                         thinkingLevel: ThinkingLevel.HIGH,
                     },
                 }),
-            }), 'prompt-id-thinking-level');
+            }), 'prompt-id-thinking-level', LlmRole.MAIN);
         });
         it('should use thinkingBudget and remove thinkingLevel for non-gemini-3 models', async () => {
             const response = (async function* () {
@@ -725,7 +743,7 @@ describe('GeminiChat', () => {
                 };
             })();
             vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue(response);
-            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'hello', 'prompt-id-thinking-budget', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'hello', 'prompt-id-thinking-budget', new AbortController().signal, LlmRole.MAIN);
             for await (const _ of stream) {
                 // consume stream
             }
@@ -733,11 +751,11 @@ describe('GeminiChat', () => {
                 model: 'gemini-2.0-flash',
                 config: expect.objectContaining({
                     thinkingConfig: {
-                        thinkingBudget: DEFAULT_THINKING_MODE,
+                        thinkingBudget: 8192,
                         thinkingLevel: undefined,
                     },
                 }),
-            }), 'prompt-id-thinking-budget');
+            }), 'prompt-id-thinking-budget', LlmRole.MAIN);
         });
     });
     describe('addHistory', () => {
@@ -776,7 +794,7 @@ describe('GeminiChat', () => {
                     candidates: [{ content: { parts: [{ text: '' }] } }],
                 };
             })());
-            const stream = await chat.sendMessageStream({ model: 'gemini-1.5-pro' }, 'test', 'prompt-id-no-retry', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'gemini-1.5-pro' }, 'test', 'prompt-id-no-retry', new AbortController().signal, LlmRole.MAIN);
             await expect((async () => {
                 for await (const _ of stream) {
                     // Must loop to trigger the internal logic that throws.
@@ -809,7 +827,7 @@ describe('GeminiChat', () => {
                 };
             })());
             // ACT: Send a message and collect all events from the stream.
-            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-yield-retry', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-yield-retry', new AbortController().signal, LlmRole.MAIN);
             const events = [];
             for await (const event of stream) {
                 events.push(event);
@@ -841,7 +859,7 @@ describe('GeminiChat', () => {
                     ],
                 };
             })());
-            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test', 'prompt-id-retry-success', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test', 'prompt-id-retry-success', new AbortController().signal, LlmRole.MAIN);
             const chunks = [];
             for await (const chunk of stream) {
                 chunks.push(chunk);
@@ -892,7 +910,7 @@ describe('GeminiChat', () => {
                     ],
                 };
             })());
-            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-retry-temperature', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-retry-temperature', new AbortController().signal, LlmRole.MAIN);
             for await (const _ of stream) {
                 // consume stream
             }
@@ -902,13 +920,13 @@ describe('GeminiChat', () => {
                 config: expect.objectContaining({
                     temperature: 0,
                 }),
-            }), 'prompt-id-retry-temperature');
+            }), 'prompt-id-retry-temperature', LlmRole.MAIN);
             // Second call (retry) should have temperature 1
             expect(mockContentGenerator.generateContentStream).toHaveBeenNthCalledWith(2, expect.objectContaining({
                 config: expect.objectContaining({
                     temperature: 1,
                 }),
-            }), 'prompt-id-retry-temperature');
+            }), 'prompt-id-retry-temperature', LlmRole.MAIN);
         });
         it('should fail after all retries on persistent invalid content and report metrics', async () => {
             vi.mocked(mockContentGenerator.generateContentStream).mockImplementation(async () => (async function* () {
@@ -923,7 +941,7 @@ describe('GeminiChat', () => {
                     ],
                 };
             })());
-            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test', 'prompt-id-retry-fail', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test', 'prompt-id-retry-fail', new AbortController().signal, LlmRole.MAIN);
             await expect(async () => {
                 for await (const _ of stream) {
                     // Must loop to trigger the internal logic that throws.
@@ -972,7 +990,7 @@ describe('GeminiChat', () => {
             it('should not retry on 400 Bad Request errors', async () => {
                 const error400 = new ApiError({ message: 'Bad Request', status: 400 });
                 vi.mocked(mockContentGenerator.generateContentStream).mockRejectedValue(error400);
-                const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-400', new AbortController().signal);
+                const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-400', new AbortController().signal, LlmRole.MAIN);
                 await expect((async () => {
                     for await (const _ of stream) {
                         /* consume stream */
@@ -995,7 +1013,7 @@ describe('GeminiChat', () => {
                         ],
                     };
                 })());
-                const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-429-retry', new AbortController().signal);
+                const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-429-retry', new AbortController().signal, LlmRole.MAIN);
                 const events = [];
                 for await (const event of stream) {
                     events.push(event);
@@ -1024,7 +1042,7 @@ describe('GeminiChat', () => {
                         ],
                     };
                 })());
-                const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-500-retry', new AbortController().signal);
+                const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-500-retry', new AbortController().signal, LlmRole.MAIN);
                 const events = [];
                 for await (const event of stream) {
                     events.push(event);
@@ -1060,7 +1078,7 @@ describe('GeminiChat', () => {
                         throw error;
                     }
                 });
-                const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-fetch-error-retry', new AbortController().signal);
+                const stream = await chat.sendMessageStream({ model: 'test-model' }, 'test message', 'prompt-id-fetch-error-retry', new AbortController().signal, LlmRole.MAIN);
                 const events = [];
                 for await (const event of stream) {
                     events.push(event);
@@ -1103,7 +1121,7 @@ describe('GeminiChat', () => {
             };
         })());
         // 3. Send a new message
-        const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'Second question', 'prompt-id-retry-existing', new AbortController().signal);
+        const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'Second question', 'prompt-id-retry-existing', new AbortController().signal, LlmRole.MAIN);
         for await (const _ of stream) {
             // consume stream
         }
@@ -1155,7 +1173,7 @@ describe('GeminiChat', () => {
             };
         })());
         // 2. Call the method and consume the stream.
-        const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test empty stream', 'prompt-id-empty-stream', new AbortController().signal);
+        const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test empty stream', 'prompt-id-empty-stream', new AbortController().signal, LlmRole.MAIN);
         const chunks = [];
         for await (const chunk of stream) {
             chunks.push(chunk);
@@ -1216,11 +1234,11 @@ describe('GeminiChat', () => {
             .mockResolvedValueOnce(firstStreamGenerator)
             .mockResolvedValueOnce(secondStreamGenerator);
         // 3. Start the first stream and consume only the first chunk to pause it
-        const firstStream = await chat.sendMessageStream({ model: 'test-model' }, 'first', 'prompt-1', new AbortController().signal);
+        const firstStream = await chat.sendMessageStream({ model: 'test-model' }, 'first', 'prompt-1', new AbortController().signal, LlmRole.MAIN);
         const firstStreamIterator = firstStream[Symbol.asyncIterator]();
         await firstStreamIterator.next();
         // 4. While the first stream is paused, start the second call. It will block.
-        const secondStreamPromise = chat.sendMessageStream({ model: 'test-model' }, 'second', 'prompt-2', new AbortController().signal);
+        const secondStreamPromise = chat.sendMessageStream({ model: 'test-model' }, 'second', 'prompt-2', new AbortController().signal, LlmRole.MAIN);
         // 5. Assert that only one API call has been made so far.
         expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(1);
         // 6. Unblock and fully consume the first stream to completion.
@@ -1294,7 +1312,7 @@ describe('GeminiChat', () => {
                 };
             })());
             mockHandleFallback.mockImplementation(async () => true);
-            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'trigger 429', 'prompt-id-fb1', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'test-model' }, 'trigger 429', 'prompt-id-fb1', new AbortController().signal, LlmRole.MAIN);
             // Consume stream to trigger logic
             for await (const _ of stream) {
                 // no-op
@@ -1341,7 +1359,7 @@ describe('GeminiChat', () => {
             };
         })());
         // Send a message and consume the stream
-        const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-discard-test', new AbortController().signal);
+        const stream = await chat.sendMessageStream({ model: 'gemini-2.0-flash' }, 'test message', 'prompt-id-discard-test', new AbortController().signal, LlmRole.MAIN);
         const events = [];
         for await (const event of stream) {
             events.push(event);
@@ -1520,7 +1538,7 @@ describe('GeminiChat', () => {
                     ],
                 };
             })());
-            const stream = await chat.sendMessageStream({ model: 'gemini-pro' }, 'test', 'prompt-healthy', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'gemini-pro' }, 'test', 'prompt-healthy', new AbortController().signal, LlmRole.MAIN);
             for await (const _ of stream) {
                 // consume
             }
@@ -1542,7 +1560,7 @@ describe('GeminiChat', () => {
                     ],
                 };
             })());
-            const stream = await chat.sendMessageStream({ model: 'gemini-pro' }, 'test', 'prompt-sticky-once', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'gemini-pro' }, 'test', 'prompt-sticky-once', new AbortController().signal, LlmRole.MAIN);
             for await (const _ of stream) {
                 // consume
             }
@@ -1576,7 +1594,7 @@ describe('GeminiChat', () => {
                 }
             });
             const consume = async () => {
-                const stream = await chat.sendMessageStream({ model: 'gemini-pro' }, 'test', 'prompt-fallback-arg', new AbortController().signal);
+                const stream = await chat.sendMessageStream({ model: 'gemini-pro' }, 'test', 'prompt-fallback-arg', new AbortController().signal, LlmRole.MAIN);
                 for await (const _ of stream) {
                     // consume
                 }
@@ -1633,19 +1651,126 @@ describe('GeminiChat', () => {
             vi.mocked(mockContentGenerator.generateContentStream)
                 .mockResolvedValueOnce(firstResponse)
                 .mockResolvedValueOnce(secondResponse);
-            const stream = await chat.sendMessageStream({ model: 'gemini-pro' }, 'test', 'prompt-config-refresh', new AbortController().signal);
+            const stream = await chat.sendMessageStream({ model: 'gemini-pro' }, 'test', 'prompt-config-refresh', new AbortController().signal, LlmRole.MAIN);
             // 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));
+                config: expect.objectContaining({
+                    temperature: 0.1,
+                }),
+            }), expect.any(String), LlmRole.MAIN);
             expect(mockContentGenerator.generateContentStream).toHaveBeenNthCalledWith(2, expect.objectContaining({
                 model: 'model-b',
-                config: expect.objectContaining({ temperature: 0.9 }),
-            }), expect.any(String));
+                config: expect.objectContaining({
+                    temperature: 0.9,
+                }),
+            }), expect.any(String), LlmRole.MAIN);
+        });
+    });
+    describe('Hook execution control', () => {
+        let mockHookSystem;
+        beforeEach(() => {
+            vi.mocked(mockConfig.getEnableHooks).mockReturnValue(true);
+            mockHookSystem = {
+                fireBeforeModelEvent: vi.fn().mockResolvedValue({ blocked: false }),
+                fireAfterModelEvent: vi.fn().mockResolvedValue({ response: {} }),
+                fireBeforeToolSelectionEvent: vi.fn().mockResolvedValue({}),
+            };
+            mockConfig.getHookSystem = vi.fn().mockReturnValue(mockHookSystem);
+        });
+        it('should yield AGENT_EXECUTION_STOPPED when BeforeModel hook stops execution', async () => {
+            vi.mocked(mockHookSystem.fireBeforeModelEvent).mockResolvedValue({
+                blocked: true,
+                stopped: true,
+                reason: 'stopped by hook',
+            });
+            const stream = await chat.sendMessageStream({ model: 'gemini-pro' }, 'test', 'prompt-id', new AbortController().signal, LlmRole.MAIN);
+            const events = [];
+            for await (const event of stream) {
+                events.push(event);
+            }
+            expect(events).toHaveLength(1);
+            expect(events[0]).toEqual({
+                type: StreamEventType.AGENT_EXECUTION_STOPPED,
+                reason: 'stopped by hook',
+            });
+        });
+        it('should yield AGENT_EXECUTION_BLOCKED and synthetic response when BeforeModel hook blocks execution', async () => {
+            const syntheticResponse = {
+                candidates: [{ content: { parts: [{ text: 'blocked' }] } }],
+            };
+            vi.mocked(mockHookSystem.fireBeforeModelEvent).mockResolvedValue({
+                blocked: true,
+                reason: 'blocked by hook',
+                syntheticResponse,
+            });
+            const stream = await chat.sendMessageStream({ model: 'gemini-pro' }, 'test', 'prompt-id', new AbortController().signal, LlmRole.MAIN);
+            const events = [];
+            for await (const event of stream) {
+                events.push(event);
+            }
+            expect(events).toHaveLength(2);
+            expect(events[0]).toEqual({
+                type: StreamEventType.AGENT_EXECUTION_BLOCKED,
+                reason: 'blocked by hook',
+            });
+            expect(events[1]).toEqual({
+                type: StreamEventType.CHUNK,
+                value: syntheticResponse,
+            });
+        });
+        it('should yield AGENT_EXECUTION_STOPPED when AfterModel hook stops execution', async () => {
+            // Mock content generator to return a stream
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue((async function* () {
+                yield {
+                    candidates: [{ content: { parts: [{ text: 'response' }] } }],
+                };
+            })());
+            vi.mocked(mockHookSystem.fireAfterModelEvent).mockResolvedValue({
+                response: {},
+                stopped: true,
+                reason: 'stopped by after hook',
+            });
+            const stream = await chat.sendMessageStream({ model: 'gemini-pro' }, 'test', 'prompt-id', new AbortController().signal, LlmRole.MAIN);
+            const events = [];
+            for await (const event of stream) {
+                events.push(event);
+            }
+            expect(events).toContainEqual({
+                type: StreamEventType.AGENT_EXECUTION_STOPPED,
+                reason: 'stopped by after hook',
+            });
+        });
+        it('should yield AGENT_EXECUTION_BLOCKED and response when AfterModel hook blocks execution', async () => {
+            const response = {
+                candidates: [{ content: { parts: [{ text: 'response' }] } }],
+            };
+            // Mock content generator to return a stream
+            vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue((async function* () {
+                yield response;
+            })());
+            vi.mocked(mockHookSystem.fireAfterModelEvent).mockResolvedValue({
+                response,
+                blocked: true,
+                reason: 'blocked by after hook',
+            });
+            const stream = await chat.sendMessageStream({ model: 'gemini-pro' }, 'test', 'prompt-id', new AbortController().signal, LlmRole.MAIN);
+            const events = [];
+            for await (const event of stream) {
+                events.push(event);
+            }
+            expect(events).toContainEqual({
+                type: StreamEventType.AGENT_EXECUTION_BLOCKED,
+                reason: 'blocked by after hook',
+            });
+            // Should also contain the chunk (hook response)
+            expect(events).toContainEqual({
+                type: StreamEventType.CHUNK,
+                value: response,
+            });
         });
     });
 });