fix(cli): correctly handle tool invocation cancellation (#844)

This commit is contained in:
N. Taylor Mullen
2025-06-08 11:14:45 -07:00
committed by GitHub
parent 9efca40dae
commit 241c404573
6 changed files with 156 additions and 4 deletions

View File

@@ -30,6 +30,7 @@ const MockedGeminiClientClass = vi.hoisted(() =>
// _config
this.startChat = mockStartChat;
this.sendMessageStream = mockSendMessageStream;
this.addHistory = vi.fn();
}),
);
@@ -267,6 +268,7 @@ describe('useGeminiStream', () => {
() => ({ getToolSchemaList: vi.fn(() => []) }) as any,
),
getGeminiClient: mockGetGeminiClient,
addHistory: vi.fn(),
} as unknown as Config;
mockOnDebugMessage = vi.fn();
mockHandleSlashCommand = vi.fn().mockReturnValue(false);
@@ -294,7 +296,10 @@ describe('useGeminiStream', () => {
.mockReturnValue((async function* () {})());
});
const renderTestHook = (initialToolCalls: TrackedToolCall[] = []) => {
const renderTestHook = (
initialToolCalls: TrackedToolCall[] = [],
geminiClient?: any,
) => {
mockUseReactToolScheduler.mockReturnValue([
initialToolCalls,
mockScheduleToolCalls,
@@ -302,9 +307,11 @@ describe('useGeminiStream', () => {
mockMarkToolsAsSubmitted,
]);
const client = geminiClient || mockConfig.getGeminiClient();
const { result, rerender } = renderHook(() =>
useGeminiStream(
mockConfig.getGeminiClient(),
client,
mockAddItem as unknown as UseHistoryManagerReturn['addItem'],
mockSetShowHelp,
mockConfig,
@@ -318,6 +325,7 @@ describe('useGeminiStream', () => {
rerender,
mockMarkToolsAsSubmitted,
mockSendMessageStream,
client,
// mockFilter removed
};
};
@@ -444,4 +452,44 @@ describe('useGeminiStream', () => {
expect.any(AbortSignal),
);
});
it('should handle all tool calls being cancelled', async () => {
const toolCalls: TrackedToolCall[] = [
{
request: { callId: '1', name: 'testTool', args: {} },
status: 'cancelled',
response: {
callId: '1',
responseParts: [{ text: 'cancelled' }],
error: undefined,
resultDisplay: 'Tool 1 cancelled display',
},
responseSubmittedToGemini: false,
tool: {
name: 'testTool',
description: 'desc',
getDescription: vi.fn(),
} as any,
},
];
const client = new MockedGeminiClientClass(mockConfig);
const { mockMarkToolsAsSubmitted, rerender } = renderTestHook(
toolCalls,
client,
);
await act(async () => {
rerender({} as any);
});
await waitFor(() => {
expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['1']);
expect(client.addHistory).toHaveBeenCalledTimes(2);
expect(client.addHistory).toHaveBeenCalledWith({
role: 'user',
parts: [{ text: 'cancelled' }],
});
});
});
});

View File

@@ -19,7 +19,7 @@ import {
ToolCallRequestInfo,
logUserPrompt,
} from '@gemini-cli/core';
import { type PartListUnion } from '@google/genai';
import { type Part, type PartListUnion } from '@google/genai';
import {
StreamingState,
HistoryItemWithoutId,
@@ -531,6 +531,41 @@ export const useGeminiStream = (
completedAndReadyToSubmitTools.length > 0 &&
completedAndReadyToSubmitTools.length === toolCalls.length
) {
// If all the tools were cancelled, don't submit a response to Gemini.
const allToolsCancelled = completedAndReadyToSubmitTools.every(
(tc) => tc.status === 'cancelled',
);
if (allToolsCancelled) {
if (geminiClient) {
// We need to manually add the function responses to the history
// so the model knows the tools were cancelled.
const responsesToAdd = completedAndReadyToSubmitTools.flatMap(
(toolCall) => toolCall.response.responseParts,
);
for (const response of responsesToAdd) {
let parts: Part[];
if (Array.isArray(response)) {
parts = response;
} else if (typeof response === 'string') {
parts = [{ text: response }];
} else {
parts = [response];
}
geminiClient.addHistory({
role: 'user',
parts,
});
}
}
const callIdsToMarkAsSubmitted = completedAndReadyToSubmitTools.map(
(toolCall) => toolCall.request.callId,
);
markToolsAsSubmitted(callIdsToMarkAsSubmitted);
return;
}
const responsesToSend: PartListUnion[] =
completedAndReadyToSubmitTools.map(
(toolCall) => toolCall.response.responseParts,
@@ -542,7 +577,14 @@ export const useGeminiStream = (
markToolsAsSubmitted(callIdsToMarkAsSubmitted);
submitQuery(mergePartListUnions(responsesToSend));
}
}, [toolCalls, isResponding, submitQuery, markToolsAsSubmitted, addItem]);
}, [
toolCalls,
isResponding,
submitQuery,
markToolsAsSubmitted,
addItem,
geminiClient,
]);
const pendingHistoryItems = [
pendingHistoryItemRef.current,