From 05049b5abfaea58dd70427391c21002fe0bc5dc2 Mon Sep 17 00:00:00 2001 From: Sandy Tao Date: Wed, 31 Dec 2025 07:22:53 +0800 Subject: [PATCH] feat(hooks): implement STOP_EXECUTION and enhance hook decision handling (#15685) --- .../hooks-system.before-tool-stop.responses | 3 +- integration-tests/hooks-system.test.ts | 2 +- packages/cli/src/nonInteractiveCli.test.ts | 52 +++++++ packages/cli/src/nonInteractiveCli.ts | 38 +++++ .../cli/src/ui/hooks/useGeminiStream.test.tsx | 93 +++++++++++ packages/cli/src/ui/hooks/useGeminiStream.ts | 23 +++ .../src/core/coreToolHookTriggers.test.ts | 144 +++++++++++++++++- .../core/src/core/coreToolHookTriggers.ts | 41 +++-- packages/core/src/hooks/types.ts | 2 +- packages/core/src/tools/tool-error.ts | 3 + 10 files changed, 379 insertions(+), 22 deletions(-) diff --git a/integration-tests/hooks-system.before-tool-stop.responses b/integration-tests/hooks-system.before-tool-stop.responses index 3a45bc756e..e27a393013 100644 --- a/integration-tests/hooks-system.before-tool-stop.responses +++ b/integration-tests/hooks-system.before-tool-stop.responses @@ -1,2 +1 @@ -{"method":"generateContentStream","response":[{"candidates":[{"content":{"parts":[{"functionCall":{"name":"write_file","args":{"file_path":"test.txt","content":"hello"}}}],"role":"model"},"finishReason":"STOP","index":0}]}]} -{"method":"generateContentStream","response":[{"candidates":[{"content":{"parts":[{"text":"Okay, stopping."}],"role":"model"},"finishReason":"STOP","index":0}]}]} \ No newline at end of file +{"method":"generateContentStream","response":[{"candidates":[{"content":{"parts":[{"text":"**Initializing File Creation**\n\nI'm starting to think about how to make a new file called `test.txt`. My plan is to use a `write_file` tool. I'll need to specify the location and what the file should contain. For now, it will be empty.\n\n\n","thought":true}],"role":"model"},"index":0}],"usageMetadata":{"promptTokenCount":13216,"totalTokenCount":13269,"promptTokensDetails":[{"modality":"TEXT","tokenCount":13216}],"thoughtsTokenCount":53}},{"candidates":[{"content":{"parts":[{"functionCall":{"name":"write_file","args":{"file_path":"test.txt","content":""}},"thoughtSignature":"CiQBcsjafJ20Qbx0YvING6aZ0wYoGWJh3eqornOG4E4AfBLiVsQKXwFyyNp8UlwYs/pv9IRQQGhDlrmlOJF2hfQijryyUYLI+qjDYTpZ6KKIfZF4+vS0soL2BJ3eTXA6gaadFEfNQem3WQVeQoKLFoW4Hv4mbasXqQc0K3p15DuSAtZZENTbCnsBcsjafGK+BJyF/Npnd7gyU0TL5PXePT0nuDFjhJDxlSRUJHDP315TewD3PUYsXd10oWsfhy4B5AngyUiBPUoajdsxg8WxaxnOZYqcp8EIuwtGZrCTev6IihT5nE5jj7u0P9vtnCmkAc6p+4O7Q7Jku1uVGqeJChgzI4YKSAFyyNp8EXSdbttV4xzX+NLKkc276L8Y63tnKU6/Y7fc9/58tU29DSdrgwfe9qmvwtTsO0piFXSLazqHJt8h2bgR7A7GnKDiIA=="}],"role":"model"},"index":0}],"usageMetadata":{"promptTokenCount":13216,"candidatesTokenCount":21,"totalTokenCount":13290,"promptTokensDetails":[{"modality":"TEXT","tokenCount":13216}],"thoughtsTokenCount":53}},{"candidates":[{"content":{"role":"model"},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":13216,"candidatesTokenCount":21,"totalTokenCount":13290,"promptTokensDetails":[{"modality":"TEXT","tokenCount":13216}],"thoughtsTokenCount":53}}]} diff --git a/integration-tests/hooks-system.test.ts b/integration-tests/hooks-system.test.ts index 6ad81bad2d..bb7d2cd565 100644 --- a/integration-tests/hooks-system.test.ts +++ b/integration-tests/hooks-system.test.ts @@ -1608,7 +1608,7 @@ console.log(JSON.stringify({decision: "block", systemMessage: "Disabled hook sho }); const result = await rig.run({ - args: 'Run tool', + args: 'Use write_file to create test.txt', }); // The hook should have stopped execution message (returned from tool) diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts index 7c2b61bb55..f2a2a43592 100644 --- a/packages/cli/src/nonInteractiveCli.test.ts +++ b/packages/cli/src/nonInteractiveCli.test.ts @@ -1745,4 +1745,56 @@ describe('runNonInteractive', () => { ); expect(getWrittenOutput()).toContain('Done'); }); + + it('should stop agent execution immediately when a tool call returns STOP_EXECUTION error', async () => { + const toolCallEvent: ServerGeminiStreamEvent = { + type: GeminiEventType.ToolCallRequest, + value: { + callId: 'stop-call', + name: 'stopTool', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-id-stop', + }, + }; + + // Mock tool execution returning STOP_EXECUTION + mockCoreExecuteToolCall.mockResolvedValue({ + status: 'error', + request: toolCallEvent.value, + tool: {} as AnyDeclarativeTool, + invocation: {} as AnyToolInvocation, + response: { + callId: 'stop-call', + responseParts: [{ text: 'error occurred' }], + errorType: ToolErrorType.STOP_EXECUTION, + error: new Error('Stop reason from hook'), + resultDisplay: undefined, + }, + }); + + const firstCallEvents: ServerGeminiStreamEvent[] = [ + { type: GeminiEventType.Content, value: 'Executing tool...' }, + toolCallEvent, + ]; + + // Setup the mock to return events for the first call. + // We expect the loop to terminate after the tool execution. + // If it doesn't, it might call sendMessageStream again, which we'll assert against. + mockGeminiClient.sendMessageStream + .mockReturnValueOnce(createStreamFromEvents(firstCallEvents)) + .mockReturnValueOnce(createStreamFromEvents([])); + + await runNonInteractive({ + config: mockConfig, + settings: mockSettings, + input: 'Run stop tool', + prompt_id: 'prompt-id-stop', + }); + + expect(mockCoreExecuteToolCall).toHaveBeenCalled(); + + // The key assertion: sendMessageStream should have been called ONLY ONCE (initial user input). + expect(mockGeminiClient.sendMessageStream).toHaveBeenCalledTimes(1); + }); }); diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index d76503822e..c81efd72f5 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -28,6 +28,7 @@ import { CoreEvent, createWorkingStdio, recordToolCallInteractions, + ToolErrorType, } from '@google/gemini-cli-core'; import type { Content, Part } from '@google/genai'; @@ -416,6 +417,43 @@ export async function runNonInteractive({ ); } + // Check if any tool requested to stop execution immediately + const stopExecutionTool = completedToolCalls.find( + (tc) => tc.response.errorType === ToolErrorType.STOP_EXECUTION, + ); + + if (stopExecutionTool && stopExecutionTool.response.error) { + const stopMessage = `Agent execution stopped: ${stopExecutionTool.response.error.message}`; + + if (config.getOutputFormat() === OutputFormat.TEXT) { + process.stderr.write(`${stopMessage}\n`); + } + + // Emit final result event for streaming JSON + if (streamFormatter) { + const metrics = uiTelemetryService.getMetrics(); + const durationMs = Date.now() - startTime; + streamFormatter.emitEvent({ + type: JsonStreamEventType.RESULT, + timestamp: new Date().toISOString(), + status: 'success', + stats: streamFormatter.convertToStreamStats( + metrics, + durationMs, + ), + }); + } else if (config.getOutputFormat() === OutputFormat.JSON) { + const formatter = new JsonFormatter(); + const stats = uiTelemetryService.getMetrics(); + textOutput.write( + formatter.format(config.getSessionId(), responseText, stats), + ); + } else { + textOutput.ensureTrailingNewline(); // Ensure a final newline + } + return; + } + currentMessages = [{ role: 'user', parts: toolResponseParts }]; } else { // Emit final result event for streaming JSON diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index a813adc722..8847995568 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -694,6 +694,99 @@ describe('useGeminiStream', () => { }); }); + it('should stop agent execution immediately when a tool call returns STOP_EXECUTION error', async () => { + const stopExecutionToolCalls: TrackedToolCall[] = [ + { + request: { + callId: 'stop-call', + name: 'stopTool', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-id-stop', + }, + status: 'error', + response: { + callId: 'stop-call', + responseParts: [{ text: 'error occurred' }], + errorType: ToolErrorType.STOP_EXECUTION, + error: new Error('Stop reason from hook'), + resultDisplay: undefined, + }, + responseSubmittedToGemini: false, + tool: { + displayName: 'stop tool', + }, + invocation: { + getDescription: () => `Mock description`, + } as unknown as AnyToolInvocation, + } as unknown as TrackedCompletedToolCall, + ]; + const client = new MockedGeminiClientClass(mockConfig); + + // Capture the onComplete callback + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise) + | null = null; + + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [ + [], + mockScheduleToolCalls, + mockMarkToolsAsSubmitted, + vi.fn(), + mockCancelAllToolCalls, + ]; + }); + + const { result } = renderHook(() => + useGeminiStream( + client, + [], + mockAddItem, + mockConfig, + mockLoadedSettings, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + () => {}, + () => {}, + 80, + 24, + ), + ); + + // Trigger the onComplete callback with STOP_EXECUTION tool + await act(async () => { + if (capturedOnComplete) { + await (capturedOnComplete as any)(stopExecutionToolCalls); + } + }); + + await waitFor(() => { + expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['stop-call']); + // Should add an info message to history + expect(mockAddItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: MessageType.INFO, + text: expect.stringContaining( + 'Agent execution stopped: Stop reason from hook', + ), + }), + expect.any(Number), + ); + // Ensure we do NOT call back to the API + expect(mockSendMessageStream).not.toHaveBeenCalled(); + // Streaming state should be Idle + expect(result.current.streamingState).toBe(StreamingState.Idle); + }); + }); + it('should group multiple cancelled tool call responses into a single history entry', async () => { const cancelledToolCall1: TrackedCancelledToolCall = { request: { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 3ab2cbce5f..d36d9f57ed 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -39,6 +39,7 @@ import { EDIT_TOOL_NAMES, processRestorableToolCalls, recordToolCallInteractions, + ToolErrorType, } from '@google/gemini-cli-core'; import { type Part, type PartListUnion, FinishReason } from '@google/genai'; import type { @@ -1153,6 +1154,28 @@ export const useGeminiStream = ( return; } + // Check if any tool requested to stop execution immediately + const stopExecutionTool = geminiTools.find( + (tc) => tc.response.errorType === ToolErrorType.STOP_EXECUTION, + ); + + if (stopExecutionTool && stopExecutionTool.response.error) { + addItem( + { + type: MessageType.INFO, + text: `Agent execution stopped: ${stopExecutionTool.response.error.message}`, + }, + Date.now(), + ); + setIsResponding(false); + + const callIdsToMarkAsSubmitted = geminiTools.map( + (toolCall) => toolCall.request.callId, + ); + markToolsAsSubmitted(callIdsToMarkAsSubmitted); + return; + } + // If all the tools were cancelled, don't submit a response to Gemini. const allToolsCancelled = geminiTools.every( (tc) => tc.status === 'cancelled', diff --git a/packages/core/src/core/coreToolHookTriggers.test.ts b/packages/core/src/core/coreToolHookTriggers.test.ts index e49b48e4ae..68a4357523 100644 --- a/packages/core/src/core/coreToolHookTriggers.test.ts +++ b/packages/core/src/core/coreToolHookTriggers.test.ts @@ -6,6 +6,7 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'; import { executeToolWithHooks } from './coreToolHookTriggers.js'; +import { ToolErrorType } from '../tools/tool-error.js'; import { BaseToolInvocation, type ToolResult, @@ -17,8 +18,8 @@ import { type HookExecutionResponse, } from '../confirmation-bus/types.js'; -class MockInvocation extends BaseToolInvocation<{ key: string }, ToolResult> { - constructor(params: { key: string }) { +class MockInvocation extends BaseToolInvocation<{ key?: string }, ToolResult> { + constructor(params: { key?: string }) { super(params); } getDescription() { @@ -26,8 +27,10 @@ class MockInvocation extends BaseToolInvocation<{ key: string }, ToolResult> { } async execute() { return { - llmContent: `key: ${this.params.key}`, - returnDisplay: `key: ${this.params.key}`, + llmContent: this.params.key ? `key: ${this.params.key}` : 'success', + returnDisplay: this.params.key + ? `key: ${this.params.key}` + : 'success display', }; } } @@ -39,12 +42,145 @@ describe('executeToolWithHooks', () => { beforeEach(() => { messageBus = { request: vi.fn(), + publish: vi.fn(), + subscribe: vi.fn(), + unsubscribe: vi.fn(), } as unknown as MessageBus; mockTool = { build: vi.fn().mockImplementation((params) => new MockInvocation(params)), } as unknown as AnyDeclarativeTool; }); + it('should prioritize continue: false over decision: block in BeforeTool', async () => { + const invocation = new MockInvocation({}); + const abortSignal = new AbortController().signal; + + vi.mocked(messageBus.request).mockResolvedValue({ + type: MessageBusType.HOOK_EXECUTION_RESPONSE, + correlationId: 'test-id', + success: true, + output: { + continue: false, + stopReason: 'Stop immediately', + decision: 'block', + reason: 'Should be ignored because continue is false', + }, + } as HookExecutionResponse); + + const result = await executeToolWithHooks( + invocation, + 'test_tool', + abortSignal, + messageBus, + true, + mockTool, + ); + + expect(result.error?.type).toBe(ToolErrorType.STOP_EXECUTION); + expect(result.error?.message).toBe('Stop immediately'); + }); + + it('should block execution in BeforeTool if decision is block', async () => { + const invocation = new MockInvocation({}); + const abortSignal = new AbortController().signal; + + vi.mocked(messageBus.request).mockResolvedValue({ + type: MessageBusType.HOOK_EXECUTION_RESPONSE, + correlationId: 'test-id', + success: true, + output: { + decision: 'block', + reason: 'Execution blocked', + }, + } as HookExecutionResponse); + + const result = await executeToolWithHooks( + invocation, + 'test_tool', + abortSignal, + messageBus, + true, + mockTool, + ); + + expect(result.error?.type).toBe(ToolErrorType.EXECUTION_FAILED); + expect(result.error?.message).toBe('Execution blocked'); + }); + + it('should handle continue: false in AfterTool', async () => { + const invocation = new MockInvocation({}); + const abortSignal = new AbortController().signal; + const spy = vi.spyOn(invocation, 'execute'); + + // BeforeTool allow + vi.mocked(messageBus.request) + .mockResolvedValueOnce({ + type: MessageBusType.HOOK_EXECUTION_RESPONSE, + correlationId: 'test-id', + success: true, + output: { decision: 'allow' }, + } as HookExecutionResponse) + // AfterTool stop + .mockResolvedValueOnce({ + type: MessageBusType.HOOK_EXECUTION_RESPONSE, + correlationId: 'test-id', + success: true, + output: { + continue: false, + stopReason: 'Stop after execution', + }, + } as HookExecutionResponse); + + const result = await executeToolWithHooks( + invocation, + 'test_tool', + abortSignal, + messageBus, + true, + mockTool, + ); + + expect(result.error?.type).toBe(ToolErrorType.STOP_EXECUTION); + expect(result.error?.message).toBe('Stop after execution'); + expect(spy).toHaveBeenCalled(); + }); + + it('should block result in AfterTool if decision is deny', async () => { + const invocation = new MockInvocation({}); + const abortSignal = new AbortController().signal; + + // BeforeTool allow + vi.mocked(messageBus.request) + .mockResolvedValueOnce({ + type: MessageBusType.HOOK_EXECUTION_RESPONSE, + correlationId: 'test-id', + success: true, + output: { decision: 'allow' }, + } as HookExecutionResponse) + // AfterTool deny + .mockResolvedValueOnce({ + type: MessageBusType.HOOK_EXECUTION_RESPONSE, + correlationId: 'test-id', + success: true, + output: { + decision: 'deny', + reason: 'Result denied', + }, + } as HookExecutionResponse); + + const result = await executeToolWithHooks( + invocation, + 'test_tool', + abortSignal, + messageBus, + true, + mockTool, + ); + + expect(result.error?.type).toBe(ToolErrorType.EXECUTION_FAILED); + expect(result.error?.message).toBe('Result denied'); + }); + it('should apply modified tool input from BeforeTool hook', async () => { const params = { key: 'original' }; const invocation = new MockInvocation(params); diff --git a/packages/core/src/core/coreToolHookTriggers.ts b/packages/core/src/core/coreToolHookTriggers.ts index 24fd113ee6..70f9e93c1d 100644 --- a/packages/core/src/core/coreToolHookTriggers.ts +++ b/packages/core/src/core/coreToolHookTriggers.ts @@ -273,6 +273,19 @@ export async function executeToolWithHooks( toolInput, ); + // Check if hook requested to stop entire agent execution + if (beforeOutput?.shouldStopExecution()) { + const reason = beforeOutput.getEffectiveReason(); + return { + llmContent: `Agent execution stopped by hook: ${reason}`, + returnDisplay: `Agent execution stopped by hook: ${reason}`, + error: { + type: ToolErrorType.STOP_EXECUTION, + message: reason, + }, + }; + } + // Check if hook blocked the tool execution const blockingError = beforeOutput?.getBlockingError(); if (blockingError?.blocked) { @@ -286,19 +299,6 @@ export async function executeToolWithHooks( }; } - // Check if hook requested to stop entire agent execution - if (beforeOutput?.shouldStopExecution()) { - const reason = beforeOutput.getEffectiveReason(); - return { - llmContent: `Agent execution stopped by hook: ${reason}`, - returnDisplay: `Agent execution stopped by hook: ${reason}`, - error: { - type: ToolErrorType.EXECUTION_FAILED, - message: `Agent execution stopped: ${reason}`, - }, - }; - } - // Check if hook requested to update tool input if (beforeOutput instanceof BeforeToolHookOutput) { const modifiedInput = beforeOutput.getModifiedToolInput(); @@ -386,9 +386,22 @@ export async function executeToolWithHooks( return { llmContent: `Agent execution stopped by hook: ${reason}`, returnDisplay: `Agent execution stopped by hook: ${reason}`, + error: { + type: ToolErrorType.STOP_EXECUTION, + message: reason, + }, + }; + } + + // Check if hook blocked the tool result + const blockingError = afterOutput?.getBlockingError(); + if (blockingError?.blocked) { + return { + llmContent: `Tool result blocked: ${blockingError.reason}`, + returnDisplay: `Tool result blocked: ${blockingError.reason}`, error: { type: ToolErrorType.EXECUTION_FAILED, - message: `Agent execution stopped: ${reason}`, + message: blockingError.reason, }, }; } diff --git a/packages/core/src/hooks/types.ts b/packages/core/src/hooks/types.ts index 08fec0c5dc..e54a03f840 100644 --- a/packages/core/src/hooks/types.ts +++ b/packages/core/src/hooks/types.ts @@ -180,7 +180,7 @@ export class DefaultHookOutput implements HookOutput { * Get the effective reason for blocking or stopping */ getEffectiveReason(): string { - return this.reason || this.stopReason || 'No reason provided'; + return this.stopReason || this.reason || 'No reason provided'; } /** diff --git a/packages/core/src/tools/tool-error.ts b/packages/core/src/tools/tool-error.ts index 50ce99f2e8..4102f1a490 100644 --- a/packages/core/src/tools/tool-error.ts +++ b/packages/core/src/tools/tool-error.ts @@ -71,6 +71,9 @@ export enum ToolErrorType { // WebSearch-specific Errors WEB_SEARCH_FAILED = 'web_search_failed', + + // Hook-specific Errors + STOP_EXECUTION = 'stop_execution', } /**