From 84e882770b71039b82de2506e0a6cbc9908ea20a Mon Sep 17 00:00:00 2001 From: Vedant Mahajan Date: Sat, 24 Jan 2026 18:52:08 +0530 Subject: [PATCH] migrate fireToolNotificationHook to hookSystem (#17398) Co-authored-by: Tommaso Sciortino --- .../core/src/core/coreToolHookTriggers.ts | 145 +----------------- packages/core/src/core/coreToolScheduler.ts | 8 +- packages/core/src/hooks/hookSystem.ts | 93 ++++++++++- .../core/src/scheduler/confirmation.test.ts | 22 +-- packages/core/src/scheduler/confirmation.ts | 5 +- packages/core/src/scheduler/scheduler.test.ts | 4 - 6 files changed, 109 insertions(+), 168 deletions(-) diff --git a/packages/core/src/core/coreToolHookTriggers.ts b/packages/core/src/core/coreToolHookTriggers.ts index 873d344c30..551c6aef1f 100644 --- a/packages/core/src/core/coreToolHookTriggers.ts +++ b/packages/core/src/core/coreToolHookTriggers.ts @@ -4,23 +4,9 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { MessageBus } from '../confirmation-bus/message-bus.js'; -import { - MessageBusType, - type HookExecutionRequest, - type HookExecutionResponse, -} from '../confirmation-bus/types.js'; -import { - NotificationType, - type McpToolContext, - BeforeToolHookOutput, -} from '../hooks/types.js'; +import { type McpToolContext, BeforeToolHookOutput } from '../hooks/types.js'; import type { Config } from '../config/config.js'; -import type { - ToolCallConfirmationDetails, - ToolResult, - AnyDeclarativeTool, -} from '../tools/tools.js'; +import type { ToolResult, AnyDeclarativeTool } from '../tools/tools.js'; import { ToolErrorType } from '../tools/tool-error.js'; import { debugLogger } from '../utils/debugLogger.js'; import type { AnsiOutput, ShellExecutionConfig } from '../index.js'; @@ -28,133 +14,6 @@ import type { AnyToolInvocation } from '../tools/tools.js'; import { ShellToolInvocation } from '../tools/shell.js'; import { DiscoveredMCPToolInvocation } from '../tools/mcp-tool.js'; -/** - * Serializable representation of tool confirmation details for hooks. - * Excludes function properties like onConfirm that can't be serialized. - */ -interface SerializableConfirmationDetails { - type: 'edit' | 'exec' | 'mcp' | 'info'; - title: string; - // Edit-specific fields - fileName?: string; - filePath?: string; - fileDiff?: string; - originalContent?: string | null; - newContent?: string; - isModifying?: boolean; - // Exec-specific fields - command?: string; - rootCommand?: string; - // MCP-specific fields - serverName?: string; - toolName?: string; - toolDisplayName?: string; - // Info-specific fields - prompt?: string; - urls?: string[]; -} - -/** - * Converts ToolCallConfirmationDetails to a serializable format for hooks. - * Excludes function properties (onConfirm, ideConfirmation) that can't be serialized. - */ -function toSerializableDetails( - details: ToolCallConfirmationDetails, -): SerializableConfirmationDetails { - const base: SerializableConfirmationDetails = { - type: details.type, - title: details.title, - }; - - switch (details.type) { - case 'edit': - return { - ...base, - fileName: details.fileName, - filePath: details.filePath, - fileDiff: details.fileDiff, - originalContent: details.originalContent, - newContent: details.newContent, - isModifying: details.isModifying, - }; - case 'exec': - return { - ...base, - command: details.command, - rootCommand: details.rootCommand, - }; - case 'mcp': - return { - ...base, - serverName: details.serverName, - toolName: details.toolName, - toolDisplayName: details.toolDisplayName, - }; - case 'info': - return { - ...base, - prompt: details.prompt, - urls: details.urls, - }; - default: - return base; - } -} - -/** - * Gets the message to display in the notification hook for tool confirmation. - */ -function getNotificationMessage( - confirmationDetails: ToolCallConfirmationDetails, -): string { - switch (confirmationDetails.type) { - case 'edit': - return `Tool ${confirmationDetails.title} requires editing`; - case 'exec': - return `Tool ${confirmationDetails.title} requires execution`; - case 'mcp': - return `Tool ${confirmationDetails.title} requires MCP`; - case 'info': - return `Tool ${confirmationDetails.title} requires information`; - default: - return `Tool requires confirmation`; - } -} - -/** - * Fires the ToolPermission notification hook for a tool that needs confirmation. - * - * @param messageBus The message bus to use for hook communication - * @param confirmationDetails The tool confirmation details - */ -export async function fireToolNotificationHook( - messageBus: MessageBus, - confirmationDetails: ToolCallConfirmationDetails, -): Promise { - try { - const message = getNotificationMessage(confirmationDetails); - const serializedDetails = toSerializableDetails(confirmationDetails); - - await messageBus.request( - { - type: MessageBusType.HOOK_EXECUTION_REQUEST, - eventName: 'Notification', - input: { - notification_type: NotificationType.ToolPermission, - message, - details: serializedDetails, - }, - }, - MessageBusType.HOOK_EXECUTION_RESPONSE, - ); - } catch (error) { - debugLogger.debug( - `Notification hook failed for ${confirmationDetails.title}:`, - error, - ); - } -} - /** * Extracts MCP context from a tool invocation if it's an MCP tool. * diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index 585d7b9bf6..124cef32b9 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -24,7 +24,6 @@ import { getToolSuggestion } from '../utils/tool-utils.js'; import type { ToolConfirmationRequest } from '../confirmation-bus/types.js'; import { MessageBusType } from '../confirmation-bus/types.js'; import type { MessageBus } from '../confirmation-bus/message-bus.js'; -import { fireToolNotificationHook } from './coreToolHookTriggers.js'; import { type ToolCall, type ValidatingToolCall, @@ -651,10 +650,9 @@ export class CoreToolScheduler { } // Fire Notification hook before showing confirmation to user - const messageBus = this.config.getMessageBus(); - const hooksEnabled = this.config.getEnableHooks(); - if (hooksEnabled && messageBus) { - await fireToolNotificationHook(messageBus, confirmationDetails); + const hookSystem = this.config.getHookSystem(); + if (hookSystem) { + await hookSystem.fireToolNotificationEvent(confirmationDetails); } // Allow IDE to resolve confirmation diff --git a/packages/core/src/hooks/hookSystem.ts b/packages/core/src/hooks/hookSystem.ts index 0102d41b78..bfb855c5d5 100644 --- a/packages/core/src/hooks/hookSystem.ts +++ b/packages/core/src/hooks/hookSystem.ts @@ -24,6 +24,7 @@ import type { BeforeToolSelectionHookOutput, McpToolContext, } from './types.js'; +import { NotificationType } from './types.js'; import type { AggregatedHookResult } from './hookAggregator.js'; import type { GenerateContentParameters, @@ -33,6 +34,7 @@ import type { ToolConfig, ToolListUnion, } from '@google/genai'; +import type { ToolCallConfirmationDetails } from '../tools/tools.js'; /** * Main hook system that coordinates all hook-related functionality @@ -78,6 +80,73 @@ export interface AfterModelHookResult { reason?: string; } +/** + * Converts ToolCallConfirmationDetails to a serializable format for hooks. + * Excludes function properties (onConfirm, ideConfirmation) that can't be serialized. + */ +function toSerializableDetails( + details: ToolCallConfirmationDetails, +): Record { + const base: Record = { + type: details.type, + title: details.title, + }; + + switch (details.type) { + case 'edit': + return { + ...base, + fileName: details.fileName, + filePath: details.filePath, + fileDiff: details.fileDiff, + originalContent: details.originalContent, + newContent: details.newContent, + isModifying: details.isModifying, + }; + case 'exec': + return { + ...base, + command: details.command, + rootCommand: details.rootCommand, + }; + case 'mcp': + return { + ...base, + serverName: details.serverName, + toolName: details.toolName, + toolDisplayName: details.toolDisplayName, + }; + case 'info': + return { + ...base, + prompt: details.prompt, + urls: details.urls, + }; + default: + return base; + } +} + +/** + * Gets the message to display in the notification hook for tool confirmation. + */ +function getNotificationMessage( + confirmationDetails: ToolCallConfirmationDetails, +): string { + switch (confirmationDetails.type) { + case 'edit': + return `Tool ${confirmationDetails.title} requires editing`; + case 'exec': + return `Tool ${confirmationDetails.title} requires execution`; + case 'mcp': + return `Tool ${confirmationDetails.title} requires MCP`; + case 'info': + return `Tool ${confirmationDetails.title} requires information`; + default: + return `Tool requires confirmation`; + } +} + export class HookSystem { private readonly hookRegistry: HookRegistry; private readonly hookRunner: HookRunner; @@ -312,7 +381,7 @@ export class HookSystem { ); return result.finalOutput; } catch (error) { - debugLogger.debug(`BeforeTool hook failed for ${toolName}:`, error); + debugLogger.debug(`BeforeToolEvent failed for ${toolName}:`, error); return undefined; } } @@ -336,8 +405,28 @@ export class HookSystem { ); return result.finalOutput; } catch (error) { - debugLogger.debug(`AfterTool hook failed for ${toolName}:`, error); + debugLogger.debug(`AfterToolEvent failed for ${toolName}:`, error); return undefined; } } + + async fireToolNotificationEvent( + confirmationDetails: ToolCallConfirmationDetails, + ): Promise { + try { + const message = getNotificationMessage(confirmationDetails); + const serializedDetails = toSerializableDetails(confirmationDetails); + + await this.hookEventHandler.fireNotificationEvent( + NotificationType.ToolPermission, + message, + serializedDetails, + ); + } catch (error) { + debugLogger.debug( + `NotificationEvent failed for ${confirmationDetails.title}:`, + error, + ); + } + } } diff --git a/packages/core/src/scheduler/confirmation.test.ts b/packages/core/src/scheduler/confirmation.test.ts index 12243137cd..7162af9d46 100644 --- a/packages/core/src/scheduler/confirmation.test.ts +++ b/packages/core/src/scheduler/confirmation.test.ts @@ -32,17 +32,12 @@ import type { ValidatingToolCall, WaitingToolCall } from './types.js'; import type { Config } from '../config/config.js'; import type { EditorType } from '../utils/editor.js'; import { randomUUID } from 'node:crypto'; -import { fireToolNotificationHook } from '../core/coreToolHookTriggers.js'; // Mock Dependencies vi.mock('node:crypto', () => ({ randomUUID: vi.fn(), })); -vi.mock('../core/coreToolHookTriggers.js', () => ({ - fireToolNotificationHook: vi.fn(), -})); - describe('confirmation.ts', () => { let mockMessageBus: MessageBus; @@ -140,15 +135,19 @@ describe('confirmation.ts', () => { configurable: true, }); + const mockHookSystem = { + fireToolNotificationEvent: vi.fn().mockResolvedValue(undefined), + }; + mockConfig = { + getEnableHooks: vi.fn().mockReturnValue(true), + getHookSystem: vi.fn().mockReturnValue(mockHookSystem), + } as unknown as Mocked; + mockModifier = { handleModifyWithEditor: vi.fn(), applyInlineModify: vi.fn(), } as unknown as Mocked; - mockConfig = { - getEnableHooks: vi.fn().mockReturnValue(true), - } as unknown as Mocked; - getPreferredEditor = vi.fn().mockReturnValue('vim'); invocationMock = { @@ -263,8 +262,9 @@ describe('confirmation.ts', () => { }); await promise; - expect(fireToolNotificationHook).toHaveBeenCalledWith( - mockMessageBus, + expect( + mockConfig.getHookSystem()?.fireToolNotificationEvent, + ).toHaveBeenCalledWith( expect.objectContaining({ type: details.type, prompt: details.prompt, diff --git a/packages/core/src/scheduler/confirmation.ts b/packages/core/src/scheduler/confirmation.ts index f8d5f6b6b4..c6aa541508 100644 --- a/packages/core/src/scheduler/confirmation.ts +++ b/packages/core/src/scheduler/confirmation.ts @@ -23,7 +23,6 @@ import type { SchedulerStateManager } from './state-manager.js'; import type { ToolModificationHandler } from './tool-modifier.js'; import type { EditorType } from '../utils/editor.js'; import type { DiffUpdateResult } from '../ide/ide-client.js'; -import { fireToolNotificationHook } from '../core/coreToolHookTriggers.js'; import { debugLogger } from '../utils/debugLogger.js'; export interface ConfirmationResult { @@ -168,8 +167,8 @@ async function notifyHooks( deps: { config: Config; messageBus: MessageBus }, details: ToolCallConfirmationDetails, ): Promise { - if (deps.config.getEnableHooks()) { - await fireToolNotificationHook(deps.messageBus, { + if (deps.config.getHookSystem()) { + await deps.config.getHookSystem()?.fireToolNotificationEvent({ ...details, // Pass no-op onConfirm to satisfy type definition; side-effects via // callbacks are disallowed. diff --git a/packages/core/src/scheduler/scheduler.test.ts b/packages/core/src/scheduler/scheduler.test.ts index 25bdb34deb..96340e4d5e 100644 --- a/packages/core/src/scheduler/scheduler.test.ts +++ b/packages/core/src/scheduler/scheduler.test.ts @@ -35,10 +35,6 @@ vi.mock('../telemetry/types.js', () => ({ ToolCallEvent: vi.fn().mockImplementation((call) => ({ ...call })), })); -vi.mock('../core/coreToolHookTriggers.js', () => ({ - fireToolNotificationHook: vi.fn(), -})); - import { SchedulerStateManager } from './state-manager.js'; import { resolveConfirmation } from './confirmation.js'; import { checkPolicy, updatePolicy } from './policy.js';