diff --git a/packages/core/src/core/geminiChatHookTriggers.test.ts b/packages/core/src/core/geminiChatHookTriggers.test.ts deleted file mode 100644 index 0bc1501386..0000000000 --- a/packages/core/src/core/geminiChatHookTriggers.test.ts +++ /dev/null @@ -1,204 +0,0 @@ -/** - * @license - * Copyright 2025 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { describe, it, expect, vi, beforeEach } from 'vitest'; -import { - fireBeforeModelHook, - fireAfterModelHook, -} from './geminiChatHookTriggers.js'; -import type { MessageBus } from '../confirmation-bus/message-bus.js'; -import type { - GenerateContentParameters, - GenerateContentResponse, -} from '@google/genai'; - -// Mock dependencies -const mockRequest = vi.fn(); -const mockMessageBus = { - request: mockRequest, -} as unknown as MessageBus; - -// Mock hook types -vi.mock('../hooks/types.js', async () => { - const actual = await vi.importActual('../hooks/types.js'); - return { - ...actual, - createHookOutput: vi.fn(), - }; -}); - -import { createHookOutput } from '../hooks/types.js'; - -describe('Gemini Chat Hook Triggers', () => { - beforeEach(() => { - vi.clearAllMocks(); - }); - - describe('fireBeforeModelHook', () => { - const llmRequest = { - model: 'gemini-pro', - contents: [{ parts: [{ text: 'test' }] }], - } as GenerateContentParameters; - - it('should return stopped: true when hook requests stop execution', async () => { - mockRequest.mockResolvedValue({ - output: { continue: false, stopReason: 'stopped by hook' }, - }); - vi.mocked(createHookOutput).mockReturnValue({ - shouldStopExecution: () => true, - getEffectiveReason: () => 'stopped by hook', - getBlockingError: () => ({ blocked: false, reason: '' }), - } as unknown as ReturnType); - - const result = await fireBeforeModelHook(mockMessageBus, llmRequest); - - expect(result).toEqual({ - blocked: true, - stopped: true, - reason: 'stopped by hook', - }); - }); - - it('should return blocked: true when hook blocks execution', async () => { - mockRequest.mockResolvedValue({ - output: { decision: 'block', reason: 'blocked by hook' }, - }); - vi.mocked(createHookOutput).mockReturnValue({ - shouldStopExecution: () => false, - getBlockingError: () => ({ blocked: true, reason: 'blocked by hook' }), - getEffectiveReason: () => 'blocked by hook', - getSyntheticResponse: () => undefined, - } as unknown as ReturnType); - - const result = await fireBeforeModelHook(mockMessageBus, llmRequest); - - expect(result).toEqual({ - blocked: true, - reason: 'blocked by hook', - syntheticResponse: undefined, - }); - }); - - it('should return modifications when hook allows execution', async () => { - mockRequest.mockResolvedValue({ - output: { decision: 'allow' }, - }); - vi.mocked(createHookOutput).mockReturnValue({ - shouldStopExecution: () => false, - getBlockingError: () => ({ blocked: false, reason: '' }), - applyLLMRequestModifications: (req: GenerateContentParameters) => req, - } as unknown as ReturnType); - - const result = await fireBeforeModelHook(mockMessageBus, llmRequest); - - expect(result).toEqual({ - blocked: false, - modifiedConfig: undefined, - modifiedContents: llmRequest.contents, - }); - }); - }); - - describe('fireAfterModelHook', () => { - const llmRequest = { - model: 'gemini-pro', - contents: [], - } as GenerateContentParameters; - const llmResponse = { - candidates: [ - { content: { role: 'model', parts: [{ text: 'response' }] } }, - ], - } as GenerateContentResponse; - - it('should return stopped: true when hook requests stop execution', async () => { - mockRequest.mockResolvedValue({ - output: { continue: false, stopReason: 'stopped by hook' }, - }); - vi.mocked(createHookOutput).mockReturnValue({ - shouldStopExecution: () => true, - getEffectiveReason: () => 'stopped by hook', - } as unknown as ReturnType); - - const result = await fireAfterModelHook( - mockMessageBus, - llmRequest, - llmResponse, - ); - - expect(result).toEqual({ - response: llmResponse, - stopped: true, - reason: 'stopped by hook', - }); - }); - - it('should return blocked: true when hook blocks execution', async () => { - mockRequest.mockResolvedValue({ - output: { decision: 'block', reason: 'blocked by hook' }, - }); - vi.mocked(createHookOutput).mockReturnValue({ - shouldStopExecution: () => false, - getBlockingError: () => ({ blocked: true, reason: 'blocked by hook' }), - getEffectiveReason: () => 'blocked by hook', - } as unknown as ReturnType); - - const result = await fireAfterModelHook( - mockMessageBus, - llmRequest, - llmResponse, - ); - - expect(result).toEqual({ - response: llmResponse, - blocked: true, - reason: 'blocked by hook', - }); - }); - - it('should return modified response when hook modifies response', async () => { - const modifiedResponse = { ...llmResponse, text: 'modified' }; - mockRequest.mockResolvedValue({ - output: { hookSpecificOutput: { llm_response: {} } }, - }); - vi.mocked(createHookOutput).mockReturnValue({ - shouldStopExecution: () => false, - getBlockingError: () => ({ blocked: false, reason: '' }), - getModifiedResponse: () => modifiedResponse, - } as unknown as ReturnType); - - const result = await fireAfterModelHook( - mockMessageBus, - llmRequest, - llmResponse, - ); - - expect(result).toEqual({ - response: modifiedResponse, - }); - }); - - it('should return original response when hook has no effect', async () => { - mockRequest.mockResolvedValue({ - output: {}, - }); - vi.mocked(createHookOutput).mockReturnValue({ - shouldStopExecution: () => false, - getBlockingError: () => ({ blocked: false, reason: '' }), - getModifiedResponse: () => undefined, - } as unknown as ReturnType); - - const result = await fireAfterModelHook( - mockMessageBus, - llmRequest, - llmResponse, - ); - - expect(result).toEqual({ - response: llmResponse, - }); - }); - }); -}); diff --git a/packages/core/src/core/geminiChatHookTriggers.ts b/packages/core/src/core/geminiChatHookTriggers.ts deleted file mode 100644 index e0632105de..0000000000 --- a/packages/core/src/core/geminiChatHookTriggers.ts +++ /dev/null @@ -1,269 +0,0 @@ -/** - * @license - * Copyright 2025 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import type { - GenerateContentResponse, - GenerateContentParameters, - GenerateContentConfig, - ContentListUnion, - ToolConfig, - ToolListUnion, -} from '@google/genai'; -import type { MessageBus } from '../confirmation-bus/message-bus.js'; -import { - MessageBusType, - type HookExecutionRequest, - type HookExecutionResponse, -} from '../confirmation-bus/types.js'; -import { - createHookOutput, - type BeforeModelHookOutput, - type BeforeToolSelectionHookOutput, - type AfterModelHookOutput, -} from '../hooks/types.js'; -import { debugLogger } from '../utils/debugLogger.js'; - -/** - * Result from firing the BeforeModel hook. - */ -export interface BeforeModelHookResult { - /** Whether the model call was blocked */ - blocked: boolean; - /** Whether the execution should be stopped entirely */ - stopped?: boolean; - /** Reason for blocking (if blocked) */ - reason?: string; - /** Synthetic response to return instead of calling the model (if blocked) */ - syntheticResponse?: GenerateContentResponse; - /** Modified config (if not blocked) */ - modifiedConfig?: GenerateContentConfig; - /** Modified contents (if not blocked) */ - modifiedContents?: ContentListUnion; -} - -/** - * Result from firing the BeforeToolSelection hook. - */ -export interface BeforeToolSelectionHookResult { - /** Modified tool config */ - toolConfig?: ToolConfig; - /** Modified tools */ - tools?: ToolListUnion; -} - -/** - * Result from firing the AfterModel hook. - * Contains either a modified response or indicates to use the original chunk. - */ -export interface AfterModelHookResult { - /** The response to yield (either modified or original) */ - response: GenerateContentResponse; - /** Whether the execution should be stopped entirely */ - stopped?: boolean; - /** Whether the model call was blocked */ - blocked?: boolean; - /** Reason for blocking or stopping */ - reason?: string; -} - -/** - * Fires the BeforeModel hook and returns the result. - */ -export async function fireBeforeModelHook( - messageBus: MessageBus, - llmRequest: GenerateContentParameters, -): Promise { - try { - const response = await messageBus.request< - HookExecutionRequest, - HookExecutionResponse - >( - { - type: MessageBusType.HOOK_EXECUTION_REQUEST, - eventName: 'BeforeModel', - input: { - llm_request: llmRequest as unknown as Record, - }, - }, - MessageBusType.HOOK_EXECUTION_RESPONSE, - ); - - // Reconstruct result from response - const beforeResultFinalOutput = response.output - ? createHookOutput('BeforeModel', response.output) - : undefined; - - const hookOutput = beforeResultFinalOutput; - - // Check if hook requested to stop execution - if (hookOutput?.shouldStopExecution()) { - return { - blocked: true, - stopped: true, - reason: hookOutput.getEffectiveReason(), - }; - } - - // Check if hook blocked the model call - const blockingError = hookOutput?.getBlockingError(); - if (blockingError?.blocked) { - const beforeModelOutput = hookOutput as BeforeModelHookOutput; - const syntheticResponse = beforeModelOutput.getSyntheticResponse(); - const reason = - hookOutput?.getEffectiveReason() || 'Model call blocked by hook'; - - return { - blocked: true, - reason, - syntheticResponse, - }; - } - - // Apply modifications from hook - if (hookOutput) { - const beforeModelOutput = hookOutput as BeforeModelHookOutput; - const modifiedRequest = - beforeModelOutput.applyLLMRequestModifications(llmRequest); - - return { - blocked: false, - modifiedConfig: modifiedRequest.config, - modifiedContents: modifiedRequest.contents, - }; - } - - return { blocked: false }; - } catch (error) { - debugLogger.debug(`BeforeModel hook failed:`, error); - return { blocked: false }; - } -} - -/** - * Fires the BeforeToolSelection hook and returns the result. - * - * @param messageBus The message bus to use for hook communication - * @param llmRequest The LLM request parameters - * @returns The hook result with tool configuration modifications - */ -export async function fireBeforeToolSelectionHook( - messageBus: MessageBus, - llmRequest: GenerateContentParameters, -): Promise { - try { - const response = await messageBus.request< - HookExecutionRequest, - HookExecutionResponse - >( - { - type: MessageBusType.HOOK_EXECUTION_REQUEST, - eventName: 'BeforeToolSelection', - input: { - llm_request: llmRequest as unknown as Record, - }, - }, - MessageBusType.HOOK_EXECUTION_RESPONSE, - ); - - // Reconstruct result from response - const toolSelectionResultFinalOutput = response.output - ? createHookOutput('BeforeToolSelection', response.output) - : undefined; - - // Apply tool configuration modifications - if (toolSelectionResultFinalOutput) { - const beforeToolSelectionOutput = - toolSelectionResultFinalOutput as BeforeToolSelectionHookOutput; - const modifiedConfig = - beforeToolSelectionOutput.applyToolConfigModifications({ - toolConfig: llmRequest.config?.toolConfig, - tools: llmRequest.config?.tools, - }); - - return { - toolConfig: modifiedConfig.toolConfig, - tools: modifiedConfig.tools, - }; - } - - return {}; - } catch (error) { - debugLogger.debug(`BeforeToolSelection hook failed:`, error); - return {}; - } -} - -/** - * Fires the AfterModel hook and returns the result. - * - * @param messageBus The message bus to use for hook communication - * @param originalRequest The original LLM request parameters - * @param chunk The current response chunk from the model - * @returns The hook result containing the response to yield - */ -export async function fireAfterModelHook( - messageBus: MessageBus, - originalRequest: GenerateContentParameters, - chunk: GenerateContentResponse, -): Promise { - try { - const response = await messageBus.request< - HookExecutionRequest, - HookExecutionResponse - >( - { - type: MessageBusType.HOOK_EXECUTION_REQUEST, - eventName: 'AfterModel', - input: { - llm_request: originalRequest as unknown as Record, - llm_response: chunk as unknown as Record, - }, - }, - MessageBusType.HOOK_EXECUTION_RESPONSE, - ); - - // Reconstruct result from response - const afterResultFinalOutput = response.output - ? createHookOutput('AfterModel', response.output) - : undefined; - - const hookOutput = afterResultFinalOutput; - - // Check if hook requested to stop execution - if (hookOutput?.shouldStopExecution()) { - return { - response: chunk, - stopped: true, - reason: hookOutput.getEffectiveReason(), - }; - } - - // Check if hook blocked the model call - const blockingError = hookOutput?.getBlockingError(); - if (blockingError?.blocked) { - return { - response: chunk, - blocked: true, - reason: hookOutput?.getEffectiveReason(), - }; - } - - // Apply modifications from hook - if (hookOutput) { - const afterModelOutput = hookOutput as AfterModelHookOutput; - const modifiedResponse = afterModelOutput.getModifiedResponse(); - if (modifiedResponse) { - return { response: modifiedResponse }; - } - } - - return { response: chunk }; - } catch (error) { - debugLogger.debug(`AfterModel hook failed:`, error); - // On error, return original chunk to avoid interrupting the stream. - return { response: chunk }; - } -} diff --git a/packages/core/src/hooks/hookSystem.ts b/packages/core/src/hooks/hookSystem.ts index c10380be8e..839fa47675 100644 --- a/packages/core/src/hooks/hookSystem.ts +++ b/packages/core/src/hooks/hookSystem.ts @@ -27,15 +27,56 @@ import type { AggregatedHookResult } from './hookAggregator.js'; import type { GenerateContentParameters, GenerateContentResponse, + GenerateContentConfig, + ContentListUnion, + ToolConfig, + ToolListUnion, } from '@google/genai'; -import type { - AfterModelHookResult, - BeforeModelHookResult, - BeforeToolSelectionHookResult, -} from '../core/geminiChatHookTriggers.js'; + /** * Main hook system that coordinates all hook-related functionality */ + +export interface BeforeModelHookResult { + /** Whether the model call was blocked */ + blocked: boolean; + /** Whether the execution should be stopped entirely */ + stopped?: boolean; + /** Reason for blocking (if blocked) */ + reason?: string; + /** Synthetic response to return instead of calling the model (if blocked) */ + syntheticResponse?: GenerateContentResponse; + /** Modified config (if not blocked) */ + modifiedConfig?: GenerateContentConfig; + /** Modified contents (if not blocked) */ + modifiedContents?: ContentListUnion; +} + +/** + * Result from firing the BeforeToolSelection hook. + */ +export interface BeforeToolSelectionHookResult { + /** Modified tool config */ + toolConfig?: ToolConfig; + /** Modified tools */ + tools?: ToolListUnion; +} + +/** + * Result from firing the AfterModel hook. + * Contains either a modified response or indicates to use the original chunk. + */ +export interface AfterModelHookResult { + /** The response to yield (either modified or original) */ + response: GenerateContentResponse; + /** Whether the execution should be stopped entirely */ + stopped?: boolean; + /** Whether the model call was blocked */ + blocked?: boolean; + /** Reason for blocking or stopping */ + reason?: string; +} + export class HookSystem { private readonly hookRegistry: HookRegistry; private readonly hookRunner: HookRunner;