From f2f2ecf9d83224778e5fc38cfcc4a1edddf9f7d4 Mon Sep 17 00:00:00 2001 From: Taylor Mullen Date: Tue, 27 May 2025 23:40:25 -0700 Subject: feat: Allow cancellation of in-progress Gemini requests and pre-execution checks - Implements cancellation for Gemini requests while they are actively being processed by the model. - Extends cancellation support to the logic within tools. This allows users to cancel operations during the phase where the system is determining if a tool execution requires user confirmation, which can include potentially long-running pre-flight checks or LLM-based corrections. - Underlying LLM calls for edit corrections (within and ) and next speaker checks can now also be cancelled. - Previously, cancellation of the main request was not possible until text started streaming, and pre-execution checks were not cancellable. - This change leverages the updated SDK's ability to accept an abort token and threads s throughout the request, tool execution, and pre-execution check lifecycle. Fixes https://github.com/google-gemini/gemini-cli/issues/531 --- packages/server/src/utils/editCorrector.test.ts | 33 ++++++++++++++++++++----- 1 file changed, 27 insertions(+), 6 deletions(-) (limited to 'packages/server/src/utils/editCorrector.test.ts') diff --git a/packages/server/src/utils/editCorrector.test.ts b/packages/server/src/utils/editCorrector.test.ts index 27c9ffe8..7d6f5a53 100644 --- a/packages/server/src/utils/editCorrector.test.ts +++ b/packages/server/src/utils/editCorrector.test.ts @@ -132,6 +132,7 @@ describe('editCorrector', () => { let mockGeminiClientInstance: Mocked; let mockToolRegistry: Mocked; let mockConfigInstance: Config; + const abortSignal = new AbortController().signal; beforeEach(() => { mockToolRegistry = new ToolRegistry({} as Config) as Mocked; @@ -187,12 +188,18 @@ describe('editCorrector', () => { callCount = 0; mockResponses.length = 0; - mockGenerateJson = vi.fn().mockImplementation(() => { - const response = mockResponses[callCount]; - callCount++; - if (response === undefined) return Promise.resolve({}); - return Promise.resolve(response); - }); + mockGenerateJson = vi + .fn() + .mockImplementation((_contents, _schema, signal) => { + // Check if the signal is aborted. If so, throw an error or return a specific response. + if (signal && signal.aborted) { + return Promise.reject(new Error('Aborted')); // Or some other specific error/response + } + const response = mockResponses[callCount]; + callCount++; + if (response === undefined) return Promise.resolve({}); + return Promise.resolve(response); + }); mockStartChat = vi.fn(); mockSendMessageStream = vi.fn(); @@ -217,6 +224,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(1); expect(result.params.new_string).toBe('replace with "this"'); @@ -234,6 +242,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(0); expect(result.params.new_string).toBe('replace with this'); @@ -254,6 +263,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(1); expect(result.params.new_string).toBe('replace with "this"'); @@ -271,6 +281,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(0); expect(result.params.new_string).toBe('replace with this'); @@ -292,6 +303,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(1); expect(result.params.new_string).toBe('replace with "this"'); @@ -309,6 +321,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(0); expect(result.params.new_string).toBe('replace with this'); @@ -329,6 +342,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(1); expect(result.params.new_string).toBe('replace with foobar'); @@ -351,6 +365,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(1); expect(result.params.new_string).toBe(llmNewString); @@ -372,6 +387,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(2); expect(result.params.new_string).toBe(llmNewString); @@ -391,6 +407,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(1); expect(result.params.new_string).toBe('replace with "this"'); @@ -412,6 +429,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(1); expect(result.params.new_string).toBe(newStringForLLMAndReturnedByLLM); @@ -432,6 +450,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(1); expect(result.params).toEqual(originalParams); @@ -449,6 +468,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(0); expect(result.params).toEqual(originalParams); @@ -471,6 +491,7 @@ describe('editCorrector', () => { currentContent, originalParams, mockGeminiClientInstance, + abortSignal, ); expect(mockGenerateJson).toHaveBeenCalledTimes(2); expect(result.params.old_string).toBe(currentContent); -- cgit v1.2.3