diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index 85942e66ee5..489431d2ff7 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -2262,6 +2262,18 @@ const SETTINGS_SCHEMA = { mergeStrategy: MergeStrategy.CONCAT, items: HOOK_DEFINITION_ITEMS, }, + PostToolBatch: { + type: 'array', + label: 'Post Tool Batch Hooks', + category: 'Advanced', + requiresRestart: false, + default: [], + description: + 'Hooks that execute once after all tool calls in a batch resolve.', + showInDialog: false, + mergeStrategy: MergeStrategy.CONCAT, + items: HOOK_DEFINITION_ITEMS, + }, SessionStart: { type: 'array', label: 'Session Start Hooks', diff --git a/packages/cli/src/ui/components/hooks/HooksManagementDialog.test.tsx b/packages/cli/src/ui/components/hooks/HooksManagementDialog.test.tsx index e4fd15235f3..3cd5dd7ff45 100644 --- a/packages/cli/src/ui/components/hooks/HooksManagementDialog.test.tsx +++ b/packages/cli/src/ui/components/hooks/HooksManagementDialog.test.tsx @@ -6,12 +6,14 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; import { cleanup } from 'ink-testing-library'; +import { HookEventName } from '@qwen-code/qwen-code-core'; import { HooksManagementDialog } from './HooksManagementDialog.js'; import { renderWithProviders } from '../../../test-utils/render.js'; import { useKeypress } from '../../hooks/useKeypress.js'; import { useConfig } from '../../contexts/ConfigContext.js'; import { loadSettings, SettingScope } from '../../../config/settings.js'; import type { Key } from '../../contexts/KeypressContext.js'; +import { DISPLAY_HOOK_EVENTS } from './constants.js'; vi.mock('../../hooks/useKeypress.js', () => ({ useKeypress: vi.fn(), @@ -338,14 +340,15 @@ describe('HooksManagementDialog', () => { expect(lastFrame()).toContain('Hooks'); }); - for (let i = 0; i < 6; i++) { + const stopEventIndex = DISPLAY_HOOK_EVENTS.indexOf(HookEventName.Stop); + for (let i = 0; i < stopEventIndex; i++) { pressKey('down'); await vi.waitFor(() => { expect(lastFrame()).toContain(`❯ ${i + 2}.`); }); } await vi.waitFor(() => { - expect(lastFrame()).toContain('❯ 7. Stop'); + expect(lastFrame()).toContain(`❯ ${stopEventIndex + 1}. Stop`); }); pressKey('return'); await vi.waitFor(() => { diff --git a/packages/cli/src/ui/components/hooks/constants.test.ts b/packages/cli/src/ui/components/hooks/constants.test.ts index f9100667ae3..25d149cb520 100644 --- a/packages/cli/src/ui/components/hooks/constants.test.ts +++ b/packages/cli/src/ui/components/hooks/constants.test.ts @@ -231,6 +231,7 @@ describe('hooks constants', () => { expect(DISPLAY_HOOK_EVENTS).toContain(HookEventName.PreToolUse); expect(DISPLAY_HOOK_EVENTS).toContain(HookEventName.PostToolUse); expect(DISPLAY_HOOK_EVENTS).toContain(HookEventName.PostToolUseFailure); + expect(DISPLAY_HOOK_EVENTS).toContain(HookEventName.PostToolBatch); expect(DISPLAY_HOOK_EVENTS).toContain(HookEventName.Notification); expect(DISPLAY_HOOK_EVENTS).toContain(HookEventName.UserPromptSubmit); expect(DISPLAY_HOOK_EVENTS).toContain(HookEventName.SessionStart); @@ -245,8 +246,8 @@ describe('hooks constants', () => { expect(DISPLAY_HOOK_EVENTS).toContain(HookEventName.TodoCompleted); }); - it('should have 17 events', () => { - expect(DISPLAY_HOOK_EVENTS).toHaveLength(17); + it('should have 18 events', () => { + expect(DISPLAY_HOOK_EVENTS).toHaveLength(18); }); }); @@ -268,6 +269,7 @@ describe('hooks constants', () => { it('returns false for events without matchers', () => { expect(supportsMatchers(HookEventName.Stop)).toBe(false); + expect(supportsMatchers(HookEventName.PostToolBatch)).toBe(false); expect(supportsMatchers(HookEventName.UserPromptSubmit)).toBe(false); expect(supportsMatchers(HookEventName.TodoCreated)).toBe(false); expect(supportsMatchers(HookEventName.TodoCompleted)).toBe(false); diff --git a/packages/cli/src/ui/components/hooks/constants.ts b/packages/cli/src/ui/components/hooks/constants.ts index bc86d34a2f7..af259984194 100644 --- a/packages/cli/src/ui/components/hooks/constants.ts +++ b/packages/cli/src/ui/components/hooks/constants.ts @@ -43,6 +43,11 @@ export function getHookExitCodes(eventName: string): HookExitCode[] { { code: 2, description: t('show stderr to model immediately') }, { code: 'Other', description: t('show stderr to user only') }, ], + [HookEventName.PostToolBatch]: [ + { code: 0, description: t('stdout shown in transcript mode (ctrl+o)') }, + { code: 2, description: t('show stderr to model immediately') }, + { code: 'Other', description: t('show stderr to user only') }, + ], [HookEventName.Notification]: [ { code: 0, description: t('stdout/stderr not shown') }, { code: 'Other', description: t('show stderr to user only') }, @@ -144,6 +149,9 @@ export function getHookShortDescription(eventName: string): string { [HookEventName.PreToolUse]: t('Before tool execution'), [HookEventName.PostToolUse]: t('After tool execution'), [HookEventName.PostToolUseFailure]: t('After tool execution fails'), + [HookEventName.PostToolBatch]: t( + 'After all tool calls in a batch resolve', + ), [HookEventName.Notification]: t('When notifications are sent'), [HookEventName.UserPromptSubmit]: t('When the user submits a prompt'), [HookEventName.SessionStart]: t('When a new session is started'), @@ -187,6 +195,9 @@ export function getHookDescription(eventName: string): string { [HookEventName.PostToolUseFailure]: t( 'Input to command is JSON with tool_name, tool_input, tool_use_id, error, error_type, is_interrupt, and is_timeout.', ), + [HookEventName.PostToolBatch]: t( + 'Input to command is JSON with tool_calls, an array of resolved tool calls containing tool_name, tool_input, tool_use_id, and tool_response.', + ), [HookEventName.Notification]: t( 'Input to command is JSON with notification message and type.', ), diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 72afe807be0..fdf84a51cf2 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -119,6 +119,7 @@ import { type PermissionSuggestion, type HookEventName, type HookDefinition, + type PostToolBatchToolCall, } from '../hooks/types.js'; import { fireNotificationHook } from '../core/toolHookTriggers.js'; @@ -1485,6 +1486,13 @@ export class Config { signal, ); break; + case 'PostToolBatch': + result = await hookSystem.firePostToolBatchEvent( + (input['tool_calls'] as PostToolBatchToolCall[]) || [], + (input['permission_mode'] as PermissionMode) || 'default', + signal, + ); + break; case 'Notification': result = await hookSystem.fireNotificationEvent( (input['message'] as string) || '', diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index f3eda272d8e..32909a5d4d5 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/packages/core/src/core/coreToolScheduler.test.ts @@ -76,6 +76,8 @@ type ToolSpanRecord = { shouldStop?: boolean; blockType?: string; hasAdditionalContext?: boolean; + postBatchStop?: boolean; + postBatchStopReason?: string; error?: string; }; }; @@ -84,6 +86,7 @@ const toolSpanRecords = vi.hoisted((): ToolSpanRecord[] => []); const shouldThrowToolSpanSetAttribute = vi.hoisted(() => ({ value: false })); const shouldThrowToolSpanSetStatus = vi.hoisted(() => ({ value: false })); const debugLoggerWarnSpy = vi.hoisted(() => vi.fn()); +const debugLoggerInfoSpy = vi.hoisted(() => vi.fn()); const runSideQueryMock = vi.hoisted(() => vi.fn()); vi.mock('../utils/debugLogger.js', async (importOriginal) => { @@ -93,7 +96,7 @@ vi.mock('../utils/debugLogger.js', async (importOriginal) => { ...actual, createDebugLogger: () => ({ debug: vi.fn(), - info: vi.fn(), + info: debugLoggerInfoSpy, warn: debugLoggerWarnSpy, error: vi.fn(), }), @@ -500,6 +503,7 @@ async function waitForStatus( describe('CoreToolScheduler', () => { beforeEach(() => { + debugLoggerInfoSpy.mockClear(); runSideQueryMock.mockReset(); }); @@ -699,6 +703,7 @@ describe('CoreToolScheduler', () => { getChatRecordingService: () => undefined, getMemoryPressureMonitor: () => options.memoryMonitor, getMessageBus: vi.fn().mockReturnValue(options.messageBus), + hasHooksForEvent: vi.fn().mockReturnValue(!options.disableHooks), getHookSystem: vi.fn().mockReturnValue(options.hookSystem), getDisableAllHooks: vi .fn() @@ -1263,6 +1268,673 @@ describe('CoreToolScheduler', () => { expect(execute).toHaveBeenCalledOnce(); }); + it('fires PostToolBatch once after a resolved tool batch before completion callback', async () => { + const executeA = vi.fn().mockResolvedValue({ + llmContent: [ + { + inlineData: { + mimeType: 'image/png', + data: 'raw-binary-payload', + }, + }, + ], + returnDisplay: 'alpha output', + }); + const executeB = vi.fn().mockResolvedValue({ + llmContent: 'beta output', + returnDisplay: 'beta output', + }); + const toolsByName = new Map([ + [ + 'alpha', + new MockTool({ + name: 'alpha', + kind: Kind.Read, + execute: executeA, + }), + ], + [ + 'beta', + new MockTool({ + name: 'beta', + kind: Kind.Read, + execute: executeB, + }), + ], + ]); + const callOrder: string[] = []; + const messageBus = { + request: vi + .fn() + .mockImplementation( + async (request: { + eventName: string; + }): Promise => { + callOrder.push(request.eventName); + return { + type: MessageBusType.HOOK_EXECUTION_RESPONSE, + correlationId: `${request.eventName}-hook`, + success: true, + output: + request.eventName === 'PostToolBatch' + ? { + hookSpecificOutput: { + hookEventName: 'PostToolBatch', + additionalContext: 'batch context', + }, + } + : { decision: 'allow' }, + }; + }, + ), + }; + const onAllToolCallsComplete = vi.fn(() => { + callOrder.push('complete'); + }); + const { scheduler } = createSchedulerForLegacyToolTests({ + toolsByName, + messageBus, + disableHooks: false, + onAllToolCallsComplete, + }); + + const abortController = new AbortController(); + await scheduler.schedule( + [ + { + callId: 'call-alpha', + name: 'alpha', + args: { value: 'a' }, + isClientInitiated: false, + prompt_id: 'prompt-batch', + }, + { + callId: 'call-beta', + name: 'beta', + args: { value: 'b' }, + isClientInitiated: false, + prompt_id: 'prompt-batch', + }, + ], + abortController.signal, + ); + + await vi.waitFor(() => { + expect(onAllToolCallsComplete).toHaveBeenCalled(); + }); + + const batchRequests = messageBus.request.mock.calls.filter( + ([request]) => request.eventName === 'PostToolBatch', + ); + expect(batchRequests).toHaveLength(1); + expect(batchRequests[0][0]).toEqual( + expect.objectContaining({ + eventName: 'PostToolBatch', + signal: abortController.signal, + input: { + permission_mode: 'yolo', + tool_calls: [ + expect.objectContaining({ + tool_name: 'alpha', + tool_input: { value: 'a' }, + tool_use_id: 'call-alpha', + status: 'success', + tool_response: expect.objectContaining({ + error: undefined, + response_parts: [ + expect.objectContaining({ + functionResponse: expect.objectContaining({ + parts: [ + { + inlineData: { + mimeType: 'image/png', + data: '', + }, + }, + ], + }), + }), + ], + }), + }), + expect.objectContaining({ + tool_name: 'beta', + tool_input: { value: 'b' }, + tool_use_id: 'call-beta', + status: 'success', + tool_response: expect.objectContaining({ + error: undefined, + }), + }), + ], + }, + }), + ); + expect(callOrder.indexOf('PostToolBatch')).toBeLessThan( + callOrder.indexOf('complete'), + ); + + const completionCalls = onAllToolCallsComplete.mock + .calls as unknown as Array<[ToolCall[]]>; + const completedCalls = completionCalls[0]?.[0]; + const lastCompletedCall = completedCalls?.at(-1); + const lastResponse = + lastCompletedCall && 'response' in lastCompletedCall + ? lastCompletedCall.response.responseParts.at(-1) + : undefined; + expect(lastResponse?.functionResponse?.response?.['output']).toContain( + 'batch context', + ); + expect( + ( + scheduler as unknown as { + callIdToPostToolBatchSignal: Map; + } + ).callIdToPostToolBatchSignal.size, + ).toBe(0); + }); + + it('includes failed tool responses in PostToolBatch payloads', async () => { + const executeA = vi.fn().mockResolvedValue({ + llmContent: 'alpha output', + returnDisplay: 'alpha output', + }); + const executeB = vi.fn().mockRejectedValue(new Error('beta failed')); + const toolsByName = new Map([ + [ + 'alpha', + new MockTool({ + name: 'alpha', + kind: Kind.Read, + execute: executeA, + }), + ], + [ + 'beta', + new MockTool({ + name: 'beta', + kind: Kind.Read, + execute: executeB, + }), + ], + ]); + const messageBus = { + request: vi.fn().mockImplementation( + async (request: { + eventName: string; + }): Promise => ({ + type: MessageBusType.HOOK_EXECUTION_RESPONSE, + correlationId: `${request.eventName}-hook`, + success: true, + output: { decision: 'allow' }, + }), + ), + }; + const onAllToolCallsComplete = vi.fn(); + const { scheduler } = createSchedulerForLegacyToolTests({ + toolsByName, + messageBus, + disableHooks: false, + onAllToolCallsComplete, + }); + + await scheduler.schedule( + [ + { + callId: 'call-alpha', + name: 'alpha', + args: { value: 'a' }, + isClientInitiated: false, + prompt_id: 'prompt-batch-failure', + }, + { + callId: 'call-beta', + name: 'beta', + args: { value: 'b' }, + isClientInitiated: false, + prompt_id: 'prompt-batch-failure', + }, + ], + new AbortController().signal, + ); + + await vi.waitFor(() => { + expect(onAllToolCallsComplete).toHaveBeenCalled(); + }); + + const batchRequest = messageBus.request.mock.calls.find( + ([request]) => request.eventName === 'PostToolBatch', + )?.[0]; + expect(batchRequest).toEqual( + expect.objectContaining({ + input: { + permission_mode: 'yolo', + tool_calls: [ + expect.objectContaining({ + tool_name: 'alpha', + status: 'success', + tool_response: expect.objectContaining({ + error: undefined, + error_type: undefined, + }), + }), + expect.objectContaining({ + tool_name: 'beta', + status: 'error', + tool_response: expect.objectContaining({ + error: 'beta failed', + error_type: ToolErrorType.UNHANDLED_EXCEPTION, + }), + }), + ], + }, + }), + ); + }); + + it('queues new tool calls while a PostToolBatch hook is still running', async () => { + const executeA = vi.fn().mockResolvedValue({ + llmContent: 'alpha output', + returnDisplay: 'alpha output', + }); + const executeB = vi.fn().mockResolvedValue({ + llmContent: 'beta output', + returnDisplay: 'beta output', + }); + const toolsByName = new Map([ + [ + 'alpha', + new MockTool({ + name: 'alpha', + kind: Kind.Read, + execute: executeA, + }), + ], + [ + 'beta', + new MockTool({ + name: 'beta', + kind: Kind.Read, + execute: executeB, + }), + ], + ]); + let resolveBatchHookStarted!: () => void; + const batchHookStarted = new Promise((resolve) => { + resolveBatchHookStarted = resolve; + }); + let releaseBatchHook!: () => void; + const batchHookRelease = new Promise((resolve) => { + releaseBatchHook = resolve; + }); + const messageBus = { + request: vi + .fn() + .mockImplementation( + async (request: { + eventName: string; + }): Promise => { + if (request.eventName === 'PostToolBatch') { + resolveBatchHookStarted(); + await batchHookRelease; + } + return { + type: MessageBusType.HOOK_EXECUTION_RESPONSE, + correlationId: `${request.eventName}-hook`, + success: true, + output: { decision: 'allow' }, + }; + }, + ), + }; + const onAllToolCallsComplete = vi.fn(); + const { scheduler } = createSchedulerForLegacyToolTests({ + toolsByName, + messageBus, + disableHooks: false, + onAllToolCallsComplete, + }); + + const firstSchedule = scheduler.schedule( + [ + { + callId: 'call-alpha', + name: 'alpha', + args: { value: 'a' }, + isClientInitiated: false, + prompt_id: 'prompt-batch-pending', + }, + ], + new AbortController().signal, + ); + + await batchHookStarted; + const secondSchedule = scheduler.schedule( + [ + { + callId: 'call-beta', + name: 'beta', + args: { value: 'b' }, + isClientInitiated: false, + prompt_id: 'prompt-batch-queued', + }, + ], + new AbortController().signal, + ); + + await Promise.resolve(); + expect(executeB).not.toHaveBeenCalled(); + + releaseBatchHook(); + await firstSchedule; + await secondSchedule; + + await vi.waitFor(() => { + expect(onAllToolCallsComplete).toHaveBeenCalledTimes(2); + }); + }); + + it('drains queued tool calls when completion finalization throws', async () => { + const executeA = vi.fn().mockResolvedValue({ + llmContent: 'alpha output', + returnDisplay: 'alpha output', + }); + const executeB = vi.fn().mockResolvedValue({ + llmContent: 'beta output', + returnDisplay: 'beta output', + }); + const toolsByName = new Map([ + [ + 'alpha', + new MockTool({ + name: 'alpha', + kind: Kind.Read, + execute: executeA, + }), + ], + [ + 'beta', + new MockTool({ + name: 'beta', + kind: Kind.Read, + execute: executeB, + }), + ], + ]); + let resolveBatchHookStarted!: () => void; + const batchHookStarted = new Promise((resolve) => { + resolveBatchHookStarted = resolve; + }); + let releaseBatchHook!: () => void; + const batchHookRelease = new Promise((resolve) => { + releaseBatchHook = resolve; + }); + const messageBus = { + request: vi + .fn() + .mockImplementation( + async (request: { + eventName: string; + }): Promise => { + if (request.eventName === 'PostToolBatch') { + resolveBatchHookStarted(); + await batchHookRelease; + } + return { + type: MessageBusType.HOOK_EXECUTION_RESPONSE, + correlationId: `${request.eventName}-hook`, + success: true, + output: { decision: 'allow' }, + }; + }, + ), + }; + const onAllToolCallsComplete = vi + .fn() + .mockRejectedValueOnce(new Error('completion failed')) + .mockResolvedValue(undefined); + const { scheduler } = createSchedulerForLegacyToolTests({ + toolsByName, + messageBus, + disableHooks: false, + onAllToolCallsComplete, + }); + + const firstSchedule = scheduler.schedule( + [ + { + callId: 'call-alpha', + name: 'alpha', + args: { value: 'a' }, + isClientInitiated: false, + prompt_id: 'prompt-batch-throws', + }, + ], + new AbortController().signal, + ); + + await batchHookStarted; + const secondSchedule = scheduler.schedule( + [ + { + callId: 'call-beta', + name: 'beta', + args: { value: 'b' }, + isClientInitiated: false, + prompt_id: 'prompt-batch-after-throw', + }, + ], + new AbortController().signal, + ); + + await Promise.resolve(); + expect(executeB).not.toHaveBeenCalled(); + + releaseBatchHook(); + await firstSchedule; + await secondSchedule; + + await vi.waitFor(() => { + expect(executeB).toHaveBeenCalled(); + expect(onAllToolCallsComplete).toHaveBeenCalledTimes(2); + }); + }); + + it('applies PostToolBatch stop decisions and preserves additional context', async () => { + const executeA = vi.fn().mockResolvedValue({ + llmContent: 'alpha output', + returnDisplay: 'alpha output', + }); + const executeB = vi.fn().mockResolvedValue({ + llmContent: 'beta output', + returnDisplay: 'beta output', + }); + const toolsByName = new Map([ + [ + 'alpha', + new MockTool({ + name: 'alpha', + kind: Kind.Read, + execute: executeA, + }), + ], + [ + 'beta', + new MockTool({ + name: 'beta', + kind: Kind.Read, + execute: executeB, + }), + ], + ]); + const messageBus = { + request: vi.fn().mockImplementation( + async (request: { + eventName: string; + }): Promise => ({ + type: MessageBusType.HOOK_EXECUTION_RESPONSE, + correlationId: `${request.eventName}-hook`, + success: true, + output: + request.eventName === 'PostToolBatch' + ? { + continue: false, + stopReason: 'halt', + hookSpecificOutput: { + hookEventName: 'PostToolBatch', + additionalContext: 'batch context', + }, + } + : { decision: 'allow' }, + }), + ), + }; + const onAllToolCallsComplete = vi.fn(); + const { scheduler } = createSchedulerForLegacyToolTests({ + toolsByName, + messageBus, + disableHooks: false, + onAllToolCallsComplete, + }); + + await scheduler.schedule( + [ + { + callId: 'call-alpha', + name: 'alpha', + args: { value: 'a' }, + isClientInitiated: false, + prompt_id: 'prompt-batch-stop', + }, + { + callId: 'call-beta', + name: 'beta', + args: { value: 'b' }, + isClientInitiated: false, + prompt_id: 'prompt-batch-stop', + }, + ], + new AbortController().signal, + ); + + await vi.waitFor(() => { + expect(onAllToolCallsComplete).toHaveBeenCalled(); + }); + + const completionCalls = onAllToolCallsComplete.mock + .calls as unknown as Array<[ToolCall[]]>; + const completedCalls = completionCalls[0]?.[0]; + const lastCompletedCall = completedCalls?.at(-1); + expect(completedCalls?.some((call) => call.status === 'success')).toBe( + true, + ); + expect(lastCompletedCall?.status).toBe('error'); + if (lastCompletedCall?.status === 'error') { + expect(lastCompletedCall.response.errorType).toBe( + ToolErrorType.EXECUTION_DENIED, + ); + expect(lastCompletedCall.response.error?.message).toContain('halt'); + const lastResponse = + lastCompletedCall.response.responseParts.at(-1)?.functionResponse + ?.response; + expect(lastResponse?.['error']).toContain('halt'); + expect(lastResponse?.['error']).toContain('batch context'); + expect(lastCompletedCall.response.contentLength).toBe( + 'halt'.length + 'batch context'.length + 2, + ); + expect(lastCompletedCall.outcome).toBeUndefined(); + } + expect(debugLoggerInfoSpy).toHaveBeenCalledWith( + 'PostToolBatch hook stopped batch (2 calls): halt', + ); + const batchHookSpan = toolSpanRecords.findLast( + (record) => + record.name === 'hook' && + record.attributes['hook_event'] === 'PostToolBatch', + ); + expect(batchHookSpan?.hookMetadata?.postBatchStop).toBe(true); + expect(batchHookSpan?.hookMetadata?.postBatchStopReason).toBe('halt'); + }); + + it('passes through completed calls when PostToolBatch returns hookError', async () => { + const execute = vi.fn().mockResolvedValue({ + llmContent: 'alpha output', + returnDisplay: 'alpha output', + }); + const toolsByName = new Map([ + [ + 'alpha', + new MockTool({ + name: 'alpha', + kind: Kind.Read, + execute, + }), + ], + ]); + const messageBus = { + request: vi.fn().mockImplementation( + async (request: { + eventName: string; + }): Promise => ({ + type: MessageBusType.HOOK_EXECUTION_RESPONSE, + correlationId: `${request.eventName}-hook`, + success: request.eventName !== 'PostToolBatch', + output: + request.eventName === 'PostToolBatch' + ? undefined + : { decision: 'allow' }, + error: + request.eventName === 'PostToolBatch' + ? new Error('bus timeout') + : undefined, + }), + ), + }; + const onAllToolCallsComplete = vi.fn(); + const { scheduler } = createSchedulerForLegacyToolTests({ + toolsByName, + messageBus, + disableHooks: false, + onAllToolCallsComplete, + }); + + await scheduler.schedule( + [ + { + callId: 'call-alpha', + name: 'alpha', + args: { value: 'a' }, + isClientInitiated: false, + prompt_id: 'prompt-batch-hook-error', + }, + ], + new AbortController().signal, + ); + + await vi.waitFor(() => { + expect(onAllToolCallsComplete).toHaveBeenCalled(); + }); + + const completionCalls = onAllToolCallsComplete.mock + .calls as unknown as Array<[ToolCall[]]>; + const completedCalls = completionCalls[0]?.[0]; + expect(completedCalls).toHaveLength(1); + expect(completedCalls?.[0]?.status).toBe('success'); + const batchHookSpan = toolSpanRecords.findLast( + (record) => + record.name === 'hook' && + record.attributes['hook_event'] === 'PostToolBatch', + ); + expect(batchHookSpan?.hookMetadata?.postBatchStop).toBe(false); + expect( + ( + scheduler as unknown as { + callIdToPostToolBatchSignal: Map; + } + ).callIdToPostToolBatchSignal.size, + ).toBe(0); + }); + it('should cancel a tool call if the signal is aborted before confirmation', async () => { const mockTool = new MockTool({ name: 'mockTool', @@ -4623,13 +5295,13 @@ describe('CoreToolScheduler telemetry spans', () => { }; await runSingleTool({ messageBus, disableHooks: false }); - // The PreToolUse hook span is the only one fired in this path. - const hookSpans = getHookSpans(); - expect(hookSpans).toHaveLength(1); - expect(hookSpans[0].attributes['hook_event']).toBe('PreToolUse'); - expect(hookSpans[0].hookMetadata?.success).toBe(true); - expect(hookSpans[0].hookMetadata?.shouldProceed).toBe(false); - expect(hookSpans[0].hookMetadata?.blockType).toBe('denied'); + const preToolUseSpan = getHookSpans().find( + (span) => span.attributes['hook_event'] === 'PreToolUse', + ); + expect(preToolUseSpan).toBeDefined(); + expect(preToolUseSpan?.hookMetadata?.success).toBe(true); + expect(preToolUseSpan?.hookMetadata?.shouldProceed).toBe(false); + expect(preToolUseSpan?.hookMetadata?.blockType).toBe('denied'); }); it('hook span records error when underlying hook helper surfaces hookError (#4321)', async () => { diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index e7b3623b516..16d8898647e 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -24,11 +24,13 @@ import { firePreToolUseHook, firePostToolUseHook, firePostToolUseFailureHook, + firePostToolBatchHook, fireNotificationHook, firePermissionRequestHook, appendAdditionalContext, } from './toolHookTriggers.js'; import { NotificationType } from '../hooks/types.js'; +import type { PostToolBatchToolCall } from '../hooks/types.js'; import type { MessageBus } from '../confirmation-bus/message-bus.js'; const debugLogger = createDebugLogger('TOOL_SCHEDULER'); @@ -712,6 +714,168 @@ const createErrorResponse = ( contentLength: error.message.length, }); +function serializeToolResponse( + response: ToolCallResponseInfo, +): Record { + // Keep this payload aligned with the persisted ToolCallResponseInfo fields + // hook authors need for batch-level auditing. + return { + response_parts: response.responseParts.map(summarizeBatchResponsePart), + result_display: response.resultDisplay, + error: response.error?.message, + error_type: response.errorType, + content_length: response.contentLength, + }; +} + +function summarizeBatchResponsePart(part: Part): Part { + const summarized = part.inlineData + ? { + ...part, + inlineData: { + mimeType: part.inlineData.mimeType, + data: '', + }, + } + : part; + + if (!summarized.functionResponse?.parts) { + return summarized; + } + + return { + ...summarized, + functionResponse: { + ...summarized.functionResponse, + parts: summarized.functionResponse.parts.map(summarizeBatchResponsePart), + }, + }; +} + +function toPostToolBatchToolCall( + call: CompletedToolCall, +): PostToolBatchToolCall { + return { + tool_name: call.request.name, + tool_input: call.request.args, + tool_use_id: call.request.callId, + status: call.status, + tool_response: serializeToolResponse(call.response), + }; +} + +function appendContextToResponsePart( + part: Part, + additionalContext: string, +): Part { + if (!part.functionResponse) { + debugLogger.warn( + 'appendContextToResponsePart: no functionResponse on part, additionalContext dropped', + ); + return part; + } + + const response = part.functionResponse.response ?? {}; + const output = response['output']; + const error = response['error']; + const hasOutput = Object.prototype.hasOwnProperty.call(response, 'output'); + const useOutputKey = + typeof output === 'string' || (hasOutput && typeof error !== 'string'); + const key = useOutputKey ? 'output' : 'error'; + const currentText = useOutputKey + ? typeof output === 'string' + ? output + : JSON.stringify(output) + : typeof error === 'string' + ? error + : JSON.stringify(response); + + return { + ...part, + functionResponse: { + ...part.functionResponse, + response: { + ...response, + [key]: `${currentText}\n\n${additionalContext}`, + }, + }, + }; +} + +function appendContextToToolResponse( + response: ToolCallResponseInfo, + additionalContext: string | undefined, +): ToolCallResponseInfo { + if (!additionalContext || response.responseParts.length === 0) { + return response; + } + + const responseParts = [...response.responseParts]; + const lastIndex = responseParts.length - 1; + const appendedPart = appendContextToResponsePart( + responseParts[lastIndex], + additionalContext, + ); + if (appendedPart === responseParts[lastIndex]) { + return response; + } + responseParts[lastIndex] = appendedPart; + + return { + ...response, + responseParts, + contentLength: + response.contentLength !== undefined + ? response.contentLength + additionalContext.length + 2 + : undefined, + }; +} + +function withPostToolBatchAdditionalContext( + completedCalls: CompletedToolCall[], + additionalContext: string | undefined, +): CompletedToolCall[] { + if (!additionalContext || completedCalls.length === 0) { + return completedCalls; + } + + const calls = [...completedCalls]; + const lastIndex = calls.length - 1; + calls[lastIndex] = { + ...calls[lastIndex], + response: appendContextToToolResponse( + calls[lastIndex].response, + additionalContext, + ), + } as CompletedToolCall; + return calls; +} + +function withPostToolBatchStop( + completedCalls: CompletedToolCall[], + stopReason: string, +): CompletedToolCall[] { + if (completedCalls.length === 0) { + return completedCalls; + } + + const calls = [...completedCalls]; + const lastCall = calls[calls.length - 1]; + calls[calls.length - 1] = { + status: 'error', + request: lastCall.request, + tool: lastCall.tool, + response: createErrorResponse( + lastCall.request, + new Error(stopReason), + ToolErrorType.EXECUTION_DENIED, + ), + durationMs: lastCall.durationMs, + outcome: undefined, + } as ErroredToolCall; + return calls; +} + interface CoreToolSchedulerOptions { config: Config; outputUpdateHandler?: OutputUpdateHandler; @@ -826,6 +990,10 @@ export class CoreToolScheduler { // sessions reusing the same AbortSignal don't accumulate listeners // and trip Node's MaxListenersExceededWarning (#4321 review-3). private callIdToBatch = new Map(); + // Keep the scheduling signal until the all-calls-complete hook fires. + // callIdToBatch is drained earlier when spans end, so it cannot be used + // to recover the PostToolBatch AbortSignal reliably. + private callIdToPostToolBatchSignal = new Map(); private requestQueue: Array<{ request: ToolCallRequestInfo | ToolCallRequestInfo[]; signal: AbortSignal; @@ -1032,7 +1200,13 @@ export class CoreToolScheduler { } }); this.notifyToolCallsUpdate(); - this.checkAndNotifyCompletion(); + void this.checkAndNotifyCompletion().catch((error: unknown) => { + debugLogger.warn( + `setStatusInternal completion notification failed: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + }); } private setArgsInternal(targetCallId: string, args: unknown): void { @@ -1187,6 +1361,7 @@ export class CoreToolScheduler { setToolSpanCancelled(span); this.finalizeToolSpan(callId); } + this.callIdToPostToolBatchSignal.delete(callId); } catch (e) { debugLogger.warn( `drainSpansForBatch: failed to drain ${callId}: ${e instanceof Error ? e.message : String(e)}`, @@ -1629,6 +1804,7 @@ export class CoreToolScheduler { this.toolSpans.set(reqInfo.callId, toolSpan); batchState.callIds.add(reqInfo.callId); this.callIdToBatch.set(reqInfo.callId, batchState); + this.callIdToPostToolBatchSignal.set(reqInfo.callId, signal); try { if (signal.aborted) { @@ -2112,7 +2288,13 @@ export class CoreToolScheduler { } } await this.attemptExecutionOfScheduledCalls(signal); - void this.checkAndNotifyCompletion(); + void this.checkAndNotifyCompletion().catch((error: unknown) => { + debugLogger.warn( + `_schedule completion notification failed: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + }); // Listener removal happens inside `finalizeToolSpan` → // `releaseBatchListenerIfDrained` for every callId, so we don't // need a duplicate cleanup here. That path also covers the @@ -3239,28 +3421,107 @@ export class CoreToolScheduler { ); if (this.toolCalls.length > 0 && allCallsAreTerminal) { - const completedCalls = [...this.toolCalls] as CompletedToolCall[]; + let completedCalls = [...this.toolCalls] as CompletedToolCall[]; this.toolCalls = []; - + this.isFinalizingToolCalls = true; + const batchSignal = completedCalls + .map((call) => + this.callIdToPostToolBatchSignal.get(call.request.callId), + ) + .find((candidate): candidate is AbortSignal => !!candidate); for (const call of completedCalls) { - logToolCall(this.config, new ToolCallEvent(call)); + this.callIdToPostToolBatchSignal.delete(call.request.callId); } - // Record tool results before notifying completion - this.recordToolResults(completedCalls); + let messageBus: MessageBus | undefined; + try { + const shouldFirePostToolBatch = + !this.config.getDisableAllHooks() && + (this.config.hasHooksForEvent?.('PostToolBatch') ?? false); + messageBus = shouldFirePostToolBatch + ? this.config.getMessageBus() + : undefined; + } catch (error) { + debugLogger.warn( + `PostToolBatch hook setup failed: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + } + try { + if (messageBus) { + const batchToolCalls = completedCalls.map(toPostToolBatchToolCall); + const permissionMode = this.config.getApprovalMode(); + const batchHookResult = await this.withHookSpan( + { hookEvent: 'PostToolBatch', toolName: 'batch' }, + () => + firePostToolBatchHook( + messageBus, + batchToolCalls, + permissionMode, + batchSignal, + ), + (r) => + r.hookError + ? { + success: false, + error: r.hookError, + shouldStop: false, + postBatchStop: false, + } + : { + success: true, + shouldStop: r.shouldStop, + hasAdditionalContext: !!r.additionalContext, + blockType: r.shouldStop ? 'stop' : undefined, + postBatchStop: r.shouldStop, + postBatchStopReason: r.shouldStop + ? r.stopReason || 'no reason given' + : undefined, + }, + ); + + // Order matters: stop replaces the last response, so append + // additionalContext only after the stop decision is applied. + if (batchHookResult.shouldStop) { + debugLogger.info( + `PostToolBatch hook stopped batch (${completedCalls.length} calls): ${ + batchHookResult.stopReason || 'no reason given' + }`, + ); + completedCalls = withPostToolBatchStop( + completedCalls, + batchHookResult.stopReason || + 'Execution stopped by PostToolBatch hook', + ); + } + + completedCalls = withPostToolBatchAdditionalContext( + completedCalls, + batchHookResult.additionalContext, + ); + } + + for (const call of completedCalls) { + logToolCall(this.config, new ToolCallEvent(call)); + } - if (this.onAllToolCallsComplete) { - this.isFinalizingToolCalls = true; - await this.onAllToolCallsComplete(completedCalls); + // Record tool results before notifying completion + this.recordToolResults(completedCalls); + + if (this.onAllToolCallsComplete) { + await this.onAllToolCallsComplete(completedCalls); + } + this.notifyToolCallsUpdate(); + } finally { this.isFinalizingToolCalls = false; - } - this.notifyToolCallsUpdate(); - // After completion, process the next item in the queue. - if (this.requestQueue.length > 0) { - const next = this.requestQueue.shift()!; - this._schedule(next.request, next.signal) - .then(next.resolve) - .catch(next.reject); + // Always drain the queue, even if completion callbacks throw. + if (this.requestQueue.length > 0) { + const next = this.requestQueue.shift()!; + this._schedule(next.request, next.signal) + .then(next.resolve) + .catch(next.reject); + } } } } diff --git a/packages/core/src/core/toolHookTriggers.test.ts b/packages/core/src/core/toolHookTriggers.test.ts index 165dfe8a942..cf08e3ba9fe 100644 --- a/packages/core/src/core/toolHookTriggers.test.ts +++ b/packages/core/src/core/toolHookTriggers.test.ts @@ -10,6 +10,7 @@ import { firePreToolUseHook, firePostToolUseHook, firePostToolUseFailureHook, + firePostToolBatchHook, fireNotificationHook, appendAdditionalContext, firePermissionRequestHook, @@ -18,6 +19,22 @@ import type { MessageBus } from '../confirmation-bus/message-bus.js'; import { NotificationType } from '../hooks/types.js'; import { MessageBusType } from '../confirmation-bus/types.js'; +const debugLoggerWarnSpy = vi.hoisted(() => vi.fn()); + +vi.mock('../utils/debugLogger.js', async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + createDebugLogger: () => ({ + debug: vi.fn(), + info: vi.fn(), + warn: debugLoggerWarnSpy, + error: vi.fn(), + }), + }; +}); + // Mock the MessageBus const createMockMessageBus = () => ({ @@ -405,6 +422,146 @@ describe('toolHookTriggers', () => { }); }); + describe('firePostToolBatchHook', () => { + it('should return shouldStop: false when no messageBus is provided', async () => { + const result = await firePostToolBatchHook(undefined, []); + + expect(result).toEqual({ shouldStop: false }); + }); + + it('should send resolved tool calls and return additional context', async () => { + const mockMessageBus = createMockMessageBus(); + (mockMessageBus.request as ReturnType).mockResolvedValue({ + success: true, + output: { + hookSpecificOutput: { + hookEventName: 'PostToolBatch', + additionalContext: 'batch note', + }, + }, + }); + + const result = await firePostToolBatchHook( + mockMessageBus, + [ + { + tool_name: 'read_file', + tool_input: { path: 'README.md' }, + tool_use_id: 'call-1', + status: 'success', + tool_response: { output: 'contents' }, + }, + ], + 'auto', + ); + + expect(mockMessageBus.request).toHaveBeenCalledWith( + { + type: MessageBusType.HOOK_EXECUTION_REQUEST, + eventName: 'PostToolBatch', + input: { + permission_mode: 'auto', + tool_calls: [ + { + tool_name: 'read_file', + tool_input: { path: 'README.md' }, + tool_use_id: 'call-1', + status: 'success', + tool_response: { output: 'contents' }, + }, + ], + }, + signal: undefined, + }, + MessageBusType.HOOK_EXECUTION_RESPONSE, + 15_000, + undefined, + ); + expect(result).toEqual({ + shouldStop: false, + additionalContext: 'batch note', + }); + }); + + it('should surface stop decisions', async () => { + const mockMessageBus = createMockMessageBus(); + (mockMessageBus.request as ReturnType).mockResolvedValue({ + success: true, + output: { + continue: false, + stopReason: 'stop after batch', + }, + }); + + const result = await firePostToolBatchHook(mockMessageBus, []); + + expect(result).toEqual({ + shouldStop: true, + stopReason: 'stop after batch', + additionalContext: undefined, + }); + }); + + it('should stop on deny decisions', async () => { + const mockMessageBus = createMockMessageBus(); + (mockMessageBus.request as ReturnType).mockResolvedValue({ + success: true, + output: { + decision: 'deny', + reason: 'blocked after batch', + }, + }); + + const result = await firePostToolBatchHook(mockMessageBus, []); + + expect(result).toEqual({ + shouldStop: true, + stopReason: 'blocked after batch', + additionalContext: undefined, + }); + }); + + it('should return hookError when hook execution fails without an error message', async () => { + const mockMessageBus = createMockMessageBus(); + (mockMessageBus.request as ReturnType).mockResolvedValue({ + success: false, + }); + + const result = await firePostToolBatchHook(mockMessageBus, []); + + expect(result.shouldStop).toBe(false); + expect(result.hookError).toMatch(/success: false/); + expect(debugLoggerWarnSpy).toHaveBeenCalledWith( + expect.stringContaining('PostToolBatch hook returned failure'), + ); + }); + + it('should return hookError when hook returns success without output', async () => { + const mockMessageBus = createMockMessageBus(); + (mockMessageBus.request as ReturnType).mockResolvedValue({ + success: true, + output: undefined, + }); + + const result = await firePostToolBatchHook(mockMessageBus, []); + + expect(result.shouldStop).toBe(false); + expect(result.hookError).toMatch(/no output/); + }); + + it('should return hookError when messageBus.request throws', async () => { + const mockMessageBus = createMockMessageBus(); + (mockMessageBus.request as ReturnType).mockRejectedValue( + new Error('bus timeout'), + ); + + const result = await firePostToolBatchHook(mockMessageBus, []); + + expect(result.shouldStop).toBe(false); + expect(result.hookError).toContain('bus timeout'); + }); + }); + describe('firePostToolUseFailureHook', () => { it('should return empty object when no messageBus is provided', async () => { const result = await firePostToolUseFailureHook( diff --git a/packages/core/src/core/toolHookTriggers.ts b/packages/core/src/core/toolHookTriggers.ts index 85981eaf557..7745c620c4e 100644 --- a/packages/core/src/core/toolHookTriggers.ts +++ b/packages/core/src/core/toolHookTriggers.ts @@ -18,11 +18,13 @@ import { type NotificationType, type PermissionRequestHookOutput, type PermissionSuggestion, + type PostToolBatchToolCall, } from '../hooks/types.js'; import { createDebugLogger } from '../utils/debugLogger.js'; import type { Part, PartListUnion } from '@google/genai'; const debugLogger = createDebugLogger('TOOL_HOOKS'); +const POST_TOOL_BATCH_HOOK_TIMEOUT_MS = 15_000; /** * Generate a unique tool_use_id for tracking tool executions @@ -77,6 +79,20 @@ export interface PostToolUseFailureHookResult { hookError?: string; } +/** + * Result of PostToolBatch hook execution + */ +export interface PostToolBatchHookResult { + /** Whether execution should stop before the next model request */ + shouldStop: boolean; + /** Stop reason if applicable */ + stopReason?: string; + /** Additional context to append once for the whole batch */ + additionalContext?: string; + /** See PreToolUseHookResult.hookError. */ + hookError?: string; +} + /** * Fire PreToolUse hook via MessageBus and process the result * @@ -361,6 +377,65 @@ export async function firePostToolUseFailureHook( } } +/** + * Fire PostToolBatch hook via MessageBus and process the result + * + * @param messageBus - The message bus instance + * @param toolCalls - Resolved tool calls in the batch + * @returns PostToolBatchHookResult with stop/additional-context decisions + */ +export async function firePostToolBatchHook( + messageBus: MessageBus | undefined, + toolCalls: PostToolBatchToolCall[], + permissionMode = 'default', + signal?: AbortSignal, +): Promise { + if (!messageBus) { + return { shouldStop: false }; + } + + try { + const response = await messageBus.request< + HookExecutionRequest, + HookExecutionResponse + >( + { + type: MessageBusType.HOOK_EXECUTION_REQUEST, + eventName: 'PostToolBatch', + input: { + permission_mode: permissionMode, + tool_calls: toolCalls, + }, + signal, + }, + MessageBusType.HOOK_EXECUTION_RESPONSE, + POST_TOOL_BATCH_HOOK_TIMEOUT_MS, + signal, + ); + + if (!response.success || !response.output) { + const message = + response.error?.message || + `hook runner returned ${response.success ? 'no output' : 'success: false'} without error detail`; + debugLogger.warn(`PostToolBatch hook returned failure: ${message}`); + return { shouldStop: false, hookError: message }; + } + + const batchOutput = createHookOutput('PostToolBatch', response.output); + const shouldStop = batchOutput.shouldStopExecution(); + + return { + shouldStop, + stopReason: shouldStop ? batchOutput.getEffectiveReason() : undefined, + additionalContext: batchOutput.getAdditionalContext(), + }; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + debugLogger.warn(`PostToolBatch hook error: ${message}`); + return { shouldStop: false, hookError: message }; + } +} + /** * Result of Notification hook execution */ diff --git a/packages/core/src/hooks/hookAggregator.test.ts b/packages/core/src/hooks/hookAggregator.test.ts index 7a4a6c169ba..427f9a00453 100644 --- a/packages/core/src/hooks/hookAggregator.test.ts +++ b/packages/core/src/hooks/hookAggregator.test.ts @@ -11,6 +11,7 @@ import type { HookExecutionResult, HookOutput, PermissionRequestHookOutput, + PostToolBatchHookOutput, } from './types.js'; describe('HookAggregator', () => { @@ -207,6 +208,60 @@ describe('HookAggregator', () => { result.finalOutput?.hookSpecificOutput?.['additionalContext'], ).toBe('ctx\nctx2'); }); + + it('should preserve PostToolBatch stop decisions across multiple hooks', () => { + const outputs: HookOutput[] = [ + { continue: false, stopReason: 'first hook stopped' }, + { continue: true }, + ]; + + const results: HookExecutionResult[] = outputs.map((output) => ({ + hookConfig: { type: HookType.Command, command: 'echo test' }, + eventName: HookEventName.PostToolBatch, + success: true, + output, + duration: 100, + })); + + const result = aggregator.aggregateResults( + results, + HookEventName.PostToolBatch, + ); + + const hookOutput = createHookOutput( + HookEventName.PostToolBatch, + result.finalOutput ?? {}, + ) as PostToolBatchHookOutput; + expect(hookOutput.shouldStopExecution()).toBe(true); + expect(hookOutput.getEffectiveReason()).toBe('first hook stopped'); + }); + + it('should preserve PostToolBatch deny decisions after aggregation', () => { + const outputs: HookOutput[] = [ + { decision: 'deny', reason: 'blocked' }, + { decision: 'allow' }, + ]; + + const results: HookExecutionResult[] = outputs.map((output) => ({ + hookConfig: { type: HookType.Command, command: 'echo test' }, + eventName: HookEventName.PostToolBatch, + success: true, + output, + duration: 100, + })); + + const result = aggregator.aggregateResults( + results, + HookEventName.PostToolBatch, + ); + + const hookOutput = createHookOutput( + HookEventName.PostToolBatch, + result.finalOutput ?? {}, + ) as PostToolBatchHookOutput; + expect(hookOutput.shouldStopExecution()).toBe(true); + expect(hookOutput.getEffectiveReason()).toBe('blocked'); + }); }); describe('mergePermissionRequestOutputs', () => { diff --git a/packages/core/src/hooks/hookAggregator.ts b/packages/core/src/hooks/hookAggregator.ts index e5eef22251b..37cf11fc6bb 100644 --- a/packages/core/src/hooks/hookAggregator.ts +++ b/packages/core/src/hooks/hookAggregator.ts @@ -10,6 +10,7 @@ import { PreToolUseHookOutput, PostToolUseHookOutput, PostToolUseFailureHookOutput, + PostToolBatchHookOutput, StopHookOutput, PermissionRequestHookOutput, } from './types.js'; @@ -104,6 +105,7 @@ export class HookAggregator { case HookEventName.PreToolUse: case HookEventName.PostToolUse: case HookEventName.PostToolUseFailure: + case HookEventName.PostToolBatch: case HookEventName.Stop: case HookEventName.UserPromptSubmit: case HookEventName.SubagentStop: @@ -365,6 +367,8 @@ export class HookAggregator { return new PostToolUseHookOutput(output); case HookEventName.PostToolUseFailure: return new PostToolUseFailureHookOutput(output); + case HookEventName.PostToolBatch: + return new PostToolBatchHookOutput(output); case HookEventName.Stop: case HookEventName.SubagentStop: return new StopHookOutput(output); diff --git a/packages/core/src/hooks/hookEventHandler.test.ts b/packages/core/src/hooks/hookEventHandler.test.ts index 77bb57e80ff..89f261638a0 100644 --- a/packages/core/src/hooks/hookEventHandler.test.ts +++ b/packages/core/src/hooks/hookEventHandler.test.ts @@ -773,6 +773,85 @@ describe('HookEventHandler', () => { }); }); + describe('firePostToolBatchEvent', () => { + it('should execute hooks for PostToolBatch without matcher context', async () => { + const mockPlan = createMockExecutionPlan([]); + const mockAggregated = createMockAggregatedResult(true); + + vi.mocked(mockHookPlanner.createExecutionPlan).mockReturnValue(mockPlan); + vi.mocked(mockHookRunner.executeHooksParallel).mockResolvedValue([]); + vi.mocked(mockHookAggregator.aggregateResults).mockReturnValue( + mockAggregated, + ); + + const result = await hookEventHandler.firePostToolBatchEvent([ + { + tool_name: 'read_file', + tool_input: { path: 'README.md' }, + tool_use_id: 'call-1', + status: 'success', + tool_response: { output: 'contents' }, + }, + ]); + + expect(mockHookPlanner.createExecutionPlan).toHaveBeenCalledWith( + HookEventName.PostToolBatch, + undefined, + ); + expect(result.success).toBe(true); + }); + + it('should include tool_calls in hook input', async () => { + const mockPlan = createMockExecutionPlan([ + { + type: HookType.Command, + command: 'echo test', + source: HooksConfigSource.Project, + }, + ]); + vi.mocked(mockHookPlanner.createExecutionPlan).mockReturnValue(mockPlan); + vi.mocked(mockHookRunner.executeHooksParallel).mockResolvedValue([]); + vi.mocked(mockHookAggregator.aggregateResults).mockReturnValue( + createMockAggregatedResult(true), + ); + + await hookEventHandler.firePostToolBatchEvent([ + { + tool_name: 'shell', + tool_input: { command: 'pwd' }, + tool_use_id: 'call-2', + status: 'success', + tool_response: { output: '/tmp/project' }, + }, + ]); + + const mockCalls = (mockHookRunner.executeHooksParallel as Mock).mock + .calls; + const input = mockCalls[0][2] as { + hook_event_name: string; + permission_mode: string; + tool_calls: Array<{ + tool_name: string; + tool_input: Record; + tool_use_id: string; + tool_response?: Record; + }>; + }; + + expect(input.hook_event_name).toBe(HookEventName.PostToolBatch); + expect(input.permission_mode).toBe(PermissionMode.Default); + expect(input.tool_calls).toEqual([ + { + tool_name: 'shell', + tool_input: { command: 'pwd' }, + tool_use_id: 'call-2', + status: 'success', + tool_response: { output: '/tmp/project' }, + }, + ]); + }); + }); + describe('firePostToolUseFailureEvent', () => { it('should execute hooks for PostToolUseFailure event', async () => { const mockPlan = createMockExecutionPlan([]); diff --git a/packages/core/src/hooks/hookEventHandler.ts b/packages/core/src/hooks/hookEventHandler.ts index 92f1f92adf8..91d111dfa5f 100644 --- a/packages/core/src/hooks/hookEventHandler.ts +++ b/packages/core/src/hooks/hookEventHandler.ts @@ -25,6 +25,8 @@ import type { PreToolUseInput, PostToolUseInput, PostToolUseFailureInput, + PostToolBatchInput, + PostToolBatchToolCall, PreCompactInput, PreCompactTrigger, PostCompactInput, @@ -310,6 +312,29 @@ export class HookEventHandler { ); } + /** + * Fire a PostToolBatch event + * Called once after every tool call in a batch has resolved + */ + async firePostToolBatchEvent( + toolCalls: PostToolBatchToolCall[], + permissionMode: PermissionMode = PermissionMode.Default, + signal?: AbortSignal, + ): Promise { + const input: PostToolBatchInput = { + ...this.createBaseInput(HookEventName.PostToolBatch), + permission_mode: permissionMode, + tool_calls: toolCalls, + }; + + return this.executeHooks( + HookEventName.PostToolBatch, + input, + undefined, + signal, + ); + } + /** * Fire a Notification event */ diff --git a/packages/core/src/hooks/hookPlanner.test.ts b/packages/core/src/hooks/hookPlanner.test.ts index 7fae0a350d3..22f19970872 100644 --- a/packages/core/src/hooks/hookPlanner.test.ts +++ b/packages/core/src/hooks/hookPlanner.test.ts @@ -87,6 +87,7 @@ describe('HookPlanner', () => { expect(getHookMatcherTarget(HookEventName.UserPromptSubmit)).toBe( undefined, ); + expect(getHookMatcherTarget(HookEventName.PostToolBatch)).toBe(undefined); }); }); diff --git a/packages/core/src/hooks/hookPlanner.ts b/packages/core/src/hooks/hookPlanner.ts index 68bc95bbc94..6f21d58d105 100644 --- a/packages/core/src/hooks/hookPlanner.ts +++ b/packages/core/src/hooks/hookPlanner.ts @@ -59,6 +59,7 @@ export function getHookMatcherTarget( case HookEventName.UserPromptSubmit: case HookEventName.Stop: + case HookEventName.PostToolBatch: case HookEventName.TodoCreated: case HookEventName.TodoCompleted: return undefined; diff --git a/packages/core/src/hooks/hookSystem.test.ts b/packages/core/src/hooks/hookSystem.test.ts index 6f4cbc98658..e668d69ac1e 100644 --- a/packages/core/src/hooks/hookSystem.test.ts +++ b/packages/core/src/hooks/hookSystem.test.ts @@ -93,6 +93,7 @@ describe('HookSystem', () => { firePreToolUseEvent: vi.fn(), firePostToolUseEvent: vi.fn(), firePostToolUseFailureEvent: vi.fn(), + firePostToolBatchEvent: vi.fn(), firePreCompactEvent: vi.fn(), fireNotificationEvent: vi.fn(), firePermissionRequestEvent: vi.fn(), @@ -845,6 +846,59 @@ describe('HookSystem', () => { }); }); + describe('firePostToolBatchEvent', () => { + it('should fire PostToolBatch event and return output', async () => { + const mockResult = { + success: true, + allOutputs: [], + errors: [], + totalDuration: 50, + finalOutput: { + hookSpecificOutput: { + hookEventName: 'PostToolBatch', + additionalContext: 'batch context', + }, + }, + }; + vi.mocked(mockHookEventHandler.firePostToolBatchEvent).mockResolvedValue( + mockResult, + ); + + const toolCalls = [ + { + tool_name: 'read_file', + tool_input: { path: 'README.md' }, + tool_use_id: 'call-1', + status: 'success' as const, + tool_response: { output: 'contents' }, + }, + ]; + const result = await hookSystem.firePostToolBatchEvent(toolCalls); + + expect(mockHookEventHandler.firePostToolBatchEvent).toHaveBeenCalledWith( + toolCalls, + PermissionMode.Default, + undefined, + ); + expect(result).toBeDefined(); + expect(result?.getAdditionalContext()).toBe('batch context'); + }); + + it('should return undefined when no final output', async () => { + vi.mocked(mockHookEventHandler.firePostToolBatchEvent).mockResolvedValue({ + success: true, + allOutputs: [], + errors: [], + totalDuration: 0, + finalOutput: undefined, + }); + + const result = await hookSystem.firePostToolBatchEvent([]); + + expect(result).toBeUndefined(); + }); + }); + describe('firePostToolUseFailureEvent', () => { it('should fire PostToolUseFailure event and return output', async () => { const mockResult = { diff --git a/packages/core/src/hooks/hookSystem.ts b/packages/core/src/hooks/hookSystem.ts index 5d4122f914c..601f681eb37 100644 --- a/packages/core/src/hooks/hookSystem.ts +++ b/packages/core/src/hooks/hookSystem.ts @@ -13,12 +13,12 @@ import { HookEventHandler } from './hookEventHandler.js'; import type { HookRegistryEntry } from './hookRegistry.js'; import { createDebugLogger } from '../utils/debugLogger.js'; import type { DefaultHookOutput, HookPhase } from './types.js'; -import { createHookOutput } from './types.js'; +import { createHookOutput, PermissionMode } from './types.js'; import type { SessionStartSource, SessionEndReason, AgentType, - PermissionMode, + PostToolBatchToolCall, PreCompactTrigger, PostCompactTrigger, NotificationType, @@ -268,6 +268,24 @@ export class HookSystem { : undefined; } + /** + * Fire a PostToolBatch event - called once after a tool-call batch resolves + */ + async firePostToolBatchEvent( + toolCalls: PostToolBatchToolCall[], + permissionMode: PermissionMode = PermissionMode.Default, + signal?: AbortSignal, + ): Promise { + const result = await this.hookEventHandler.firePostToolBatchEvent( + toolCalls, + permissionMode, + signal, + ); + return result.finalOutput + ? createHookOutput('PostToolBatch', result.finalOutput) + : undefined; + } + /** * Fire a PreCompact event - called before conversation compaction */ diff --git a/packages/core/src/hooks/types.ts b/packages/core/src/hooks/types.ts index 15bccb1c7bb..30afed80f57 100644 --- a/packages/core/src/hooks/types.ts +++ b/packages/core/src/hooks/types.ts @@ -26,6 +26,8 @@ export enum HookEventName { PostToolUse = 'PostToolUse', // PostToolUseFailure - After tool execution fails PostToolUseFailure = 'PostToolUseFailure', + // PostToolBatch - After a batch of tool calls resolves + PostToolBatch = 'PostToolBatch', // Notification - When notifications are sent Notification = 'Notification', // UserPromptSubmit - When the user submits a prompt @@ -281,6 +283,8 @@ export function createHookOutput( return new PostToolUseHookOutput(data); case HookEventName.PostToolUseFailure: return new PostToolUseFailureHookOutput(data); + case HookEventName.PostToolBatch: + return new PostToolBatchHookOutput(data); case HookEventName.Stop: case HookEventName.SubagentStop: return new StopHookOutput(data); @@ -485,6 +489,18 @@ export class PostToolUseFailureHookOutput extends DefaultHookOutput { } } +/** + * Specific hook output class for PostToolBatch events. + */ +export class PostToolBatchHookOutput extends DefaultHookOutput { + /** + * Check if batch processing should stop after the resolved tool calls. + */ + override shouldStopExecution(): boolean { + return super.shouldStopExecution() || this.isBlockingDecision(); + } +} + /** * Specific hook output class for Stop events. */ @@ -683,6 +699,40 @@ export interface PostToolUseFailureOutput extends HookOutput { }; } +/** + * Tool call summary for PostToolBatch hook input + */ +export interface PostToolBatchToolCall { + tool_name: string; + tool_input: Record; + tool_use_id: string; + status: 'success' | 'error' | 'cancelled'; + /** + * Serialized ToolCallResponseInfo fields for the resolved call: + * response_parts, result_display, error, error_type, and content_length. + */ + tool_response?: Record; +} + +/** + * PostToolBatch hook input + * Fired once after all tool calls in a batch have resolved. + */ +export interface PostToolBatchInput extends HookInput { + permission_mode: PermissionMode; + tool_calls: PostToolBatchToolCall[]; +} + +/** + * PostToolBatch hook output + */ +export interface PostToolBatchOutput extends HookOutput { + hookSpecificOutput?: { + hookEventName: 'PostToolBatch'; + additionalContext?: string; + }; +} + /** * UserPromptSubmit hook input */ diff --git a/packages/core/src/telemetry/session-tracing.test.ts b/packages/core/src/telemetry/session-tracing.test.ts index c0119982010..7a40c304532 100644 --- a/packages/core/src/telemetry/session-tracing.test.ts +++ b/packages/core/src/telemetry/session-tracing.test.ts @@ -985,6 +985,30 @@ describe('session-tracing', () => { endToolSpan(toolSpan, { success: true }); }); + it('records shouldStop/hasAdditionalContext on PostToolBatch', () => { + const hookSpan = startHookSpan({ + hookEvent: 'PostToolBatch', + toolName: 'batch', + }); + endHookSpan(hookSpan, { + success: true, + shouldStop: true, + hasAdditionalContext: true, + postBatchStop: true, + postBatchStopReason: 'policy halt', + }); + + const hookRecord = mockSpans.find((s) => s.name === 'qwen-code.hook'); + expect(hookRecord?.attributes['hook_event']).toBe('PostToolBatch'); + expect(hookRecord?.attributes['should_stop']).toBe(true); + expect(hookRecord?.attributes['has_additional_context']).toBe(true); + expect(hookRecord?.attributes['post_batch_stop']).toBe(true); + expect(hookRecord?.attributes['post_batch_stop_reason']).toBe( + 'policy halt', + ); + expect(hookRecord?.statuses).toHaveLength(0); + }); + it('marks status ERROR only when the hook itself threw', () => { const toolSpan = startToolSpan('Bash'); let hookSpan!: ReturnType; diff --git a/packages/core/src/telemetry/session-tracing.ts b/packages/core/src/telemetry/session-tracing.ts index 118192bea91..ee53a56091d 100644 --- a/packages/core/src/telemetry/session-tracing.ts +++ b/packages/core/src/telemetry/session-tracing.ts @@ -820,7 +820,11 @@ export function endToolBlockedOnUserSpan( // --- Hook Spans --- -export type HookEvent = 'PreToolUse' | 'PostToolUse' | 'PostToolUseFailure'; +export type HookEvent = + | 'PreToolUse' + | 'PostToolUse' + | 'PostToolUseFailure' + | 'PostToolBatch'; export interface StartHookSpanOptions { hookEvent: HookEvent; @@ -840,6 +844,10 @@ export interface HookSpanMetadata { /** Discriminator for blocking decision when applicable. */ blockType?: 'denied' | 'ask' | 'stop'; hasAdditionalContext?: boolean; + /** PostToolBatch only: true when the batch hook stopped before the next turn. */ + postBatchStop?: boolean; + /** PostToolBatch only: reason attached to a stop decision. */ + postBatchStopReason?: string; /** Hook threw — span ends as ERROR with this message. */ error?: string; } @@ -915,6 +923,12 @@ export function endHookSpan(span: Span, metadata?: HookSpanMetadata): void { endAttributes['block_type'] = metadata.blockType; if (metadata.hasAdditionalContext !== undefined) endAttributes['has_additional_context'] = metadata.hasAdditionalContext; + if (metadata.postBatchStop !== undefined) + endAttributes['post_batch_stop'] = metadata.postBatchStop; + if (metadata.postBatchStopReason !== undefined) + endAttributes['post_batch_stop_reason'] = truncateSpanError( + metadata.postBatchStopReason, + ); if (metadata.error !== undefined) endAttributes['error'] = truncateSpanError(metadata.error); } diff --git a/packages/vscode-ide-companion/schemas/settings.schema.json b/packages/vscode-ide-companion/schemas/settings.schema.json index 3ad72e458de..30c04d391ff 100644 --- a/packages/vscode-ide-companion/schemas/settings.schema.json +++ b/packages/vscode-ide-companion/schemas/settings.schema.json @@ -1646,6 +1646,109 @@ ] } }, + "PostToolBatch": { + "description": "Hooks that execute once after all tool calls in a batch resolve.", + "type": "array", + "items": { + "description": "A hook definition with an optional matcher and a list of hook configurations.", + "type": "object", + "properties": { + "matcher": { + "description": "An optional matcher pattern to filter when this hook definition applies.", + "type": "string" + }, + "sequential": { + "description": "Whether the hooks should be executed sequentially instead of in parallel.", + "type": "boolean" + }, + "hooks": { + "description": "The list of hook configurations to execute.", + "type": "array", + "items": { + "description": "A hook configuration entry that defines a hook to execute.", + "type": "object", + "properties": { + "type": { + "description": "The type of hook. Note: \"function\" type is only available via SDK registration, not settings.json.", + "type": "string", + "enum": [ + "command", + "http" + ] + }, + "command": { + "description": "The command to execute when the hook is triggered. Required for \"command\" type.", + "type": "string" + }, + "url": { + "description": "The URL to send the POST request to. Required for \"http\" type.", + "type": "string" + }, + "headers": { + "description": "HTTP headers to include in the request. Supports env var interpolation ($VAR, ${VAR}).", + "type": "object", + "additionalProperties": { + "type": "string" + } + }, + "allowedEnvVars": { + "description": "List of environment variables allowed for interpolation in headers and URL.", + "type": "array", + "items": { + "type": "string" + } + }, + "name": { + "description": "An optional name for the hook.", + "type": "string" + }, + "description": { + "description": "An optional description of what the hook does.", + "type": "string" + }, + "timeout": { + "description": "Timeout in seconds for the hook execution.", + "type": "number" + }, + "env": { + "description": "Environment variables to set when executing the hook command.", + "type": "object", + "additionalProperties": { + "type": "string" + } + }, + "async": { + "description": "Whether to execute the hook asynchronously (non-blocking, for \"command\" type only).", + "type": "boolean" + }, + "once": { + "description": "Whether to execute the hook only once per session (for \"http\" type).", + "type": "boolean" + }, + "statusMessage": { + "description": "A message to display while the hook is executing.", + "type": "string" + }, + "shell": { + "description": "The shell to use for command execution.", + "type": "string", + "enum": [ + "bash", + "powershell" + ] + } + }, + "required": [ + "type" + ] + } + } + }, + "required": [ + "hooks" + ] + } + }, "SessionStart": { "description": "Hooks that execute when a new session starts or resumes.", "type": "array",