From ca78a0f1771ff4520e59e4876f4d7a86f1b0f9b8 Mon Sep 17 00:00:00 2001 From: "N. Taylor Mullen" Date: Wed, 1 Apr 2026 16:16:34 -0700 Subject: [PATCH] fix(cli): ensure agent stops when all declinable tools are cancelled (#24479) --- .../cli/src/ui/hooks/useGeminiStream.test.tsx | 126 +++++++++++++++++- packages/cli/src/ui/hooks/useGeminiStream.ts | 15 ++- 2 files changed, 134 insertions(+), 7 deletions(-) diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index e7d9949124..d246d06a77 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -52,6 +52,7 @@ import { MCPDiscoveryState, GeminiCliOperation, getPlanModeExitMessage, + UPDATE_TOPIC_TOOL_NAME, } from '@google/gemini-cli-core'; import type { Part, PartListUnion } from '@google/genai'; import type { UseHistoryManagerReturn } from './useHistoryManager.js'; @@ -904,6 +905,30 @@ describe('useGeminiStream', () => { it('should handle all tool calls being cancelled', async () => { const cancelledToolCalls: TrackedToolCall[] = [ + { + request: { + callId: 'topic1', + name: UPDATE_TOPIC_TOOL_NAME, + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-id-3', + }, + status: CoreToolCallStatus.Success, + response: { + callId: 'topic1', + responseParts: [ + { + functionResponse: { + name: UPDATE_TOPIC_TOOL_NAME, + id: 'topic1', + response: {}, + }, + }, + ], + }, + tool: { displayName: 'Update Topic Context' }, + invocation: { getDescription: () => 'Updating topic' }, + } as any, { request: { callId: '1', @@ -924,8 +949,8 @@ describe('useGeminiStream', () => { }, invocation: { getDescription: () => `Mock description`, - } as unknown as AnyToolInvocation, - } as TrackedCancelledToolCall, + }, + } as any, ]; const client = new MockedGeminiClientClass(mockConfig); @@ -978,16 +1003,109 @@ describe('useGeminiStream', () => { }); await waitFor(() => { - expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['1']); + expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['topic1', '1']); expect(client.addHistory).toHaveBeenCalledWith({ role: 'user', - parts: [{ text: CoreToolCallStatus.Cancelled }], + parts: [ + { + functionResponse: { + name: UPDATE_TOPIC_TOOL_NAME, + id: 'topic1', + response: {}, + }, + }, + { text: CoreToolCallStatus.Cancelled }, + ], }); // Ensure we do NOT call back to the API expect(mockSendMessageStream).not.toHaveBeenCalled(); }); }); + it('should NOT stop responding when only update_topic is called', async () => { + const topicToolCalls: TrackedToolCall[] = [ + { + request: { + callId: 'topic1', + name: UPDATE_TOPIC_TOOL_NAME, + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-id-3', + }, + status: CoreToolCallStatus.Success, + response: { + callId: 'topic1', + responseParts: [ + { + functionResponse: { + name: UPDATE_TOPIC_TOOL_NAME, + id: 'topic1', + response: {}, + }, + }, + ], + }, + tool: { displayName: 'Update Topic Context' }, + invocation: { getDescription: () => 'Updating topic' }, + } as any, + ]; + const client = new MockedGeminiClientClass(mockConfig); + + // Capture the onComplete callback + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise) + | null = null; + + mockUseToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [ + topicToolCalls, + vi.fn(), + mockMarkToolsAsSubmitted, + vi.fn(), + vi.fn(), + 0, + ]; + }); + + await renderHookWithProviders(() => + useGeminiStream( + client, + [], + mockAddItem, + mockConfig, + mockLoadedSettings, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + () => {}, + () => {}, + 80, + 24, + ), + ); + + // Trigger the onComplete callback with the topic tool + await act(async () => { + if (capturedOnComplete) { + await capturedOnComplete(topicToolCalls); + } + }); + + await waitFor(() => { + // The streaming state should still be Responding because we didn't cancel anything important + // and we expect a continuation. + expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['topic1']); + // Should HAVE called back to the API for continuation + expect(mockSendMessageStream).toHaveBeenCalled(); + }); + }); + it('should stop agent execution immediately when a tool call returns STOP_EXECUTION error', async () => { const stopExecutionToolCalls: TrackedCompletedToolCall[] = [ { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index fb975a4429..a27334391a 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -1968,11 +1968,20 @@ export const useGeminiStream = ( } // If all the tools were cancelled, don't submit a response to Gemini. - const allToolsCancelled = geminiTools.every( - (tc) => tc.status === CoreToolCallStatus.Cancelled, + // Note: we ignore the topic tool because the user doesn't have a chance to decline it. + const declinableTools = geminiTools.filter( + (tc) => !isTopicTool(tc.request.name), ); + const allDeclinableToolsCancelled = + declinableTools.length > 0 && + declinableTools.every( + (tc) => tc.status === CoreToolCallStatus.Cancelled, + ); + const allToolsCancelled = + geminiTools.length > 0 && + geminiTools.every((tc) => tc.status === CoreToolCallStatus.Cancelled); - if (allToolsCancelled) { + if (allDeclinableToolsCancelled || allToolsCancelled) { // If the turn was cancelled via the imperative escape key flow, // the cancellation message is added there. We check the ref to avoid duplication. if (!turnCancelledRef.current) {