MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-04-08 20:30:53 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat: Allow cancellation of in-progress Gemini requests and pre-execution checks

Browse Source 
- 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
This commit is contained in:
Taylor Mullen
2025-05-27 23:40:25 -07:00
committed by N. Taylor Mullen
parent bfeaac8441
commit f2f2ecf9d8
16 changed files with 260 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

76
packages/server/src/tools/write-file.test.ts
View File

@@ -110,18 +110,32 @@ describe('WriteFileTool', () => {
// Default mock implementations that return valid structures
mockEnsureCorrectEdit.mockImplementation(
async (
currentContent: string,
_currentContent: string,
params: EditToolParams,
_client: GeminiClient,
): Promise<CorrectedEditResult> =>
Promise.resolve({
signal?: AbortSignal, // Make AbortSignal optional to match usage
): Promise<CorrectedEditResult> => {
if (signal?.aborted) {
return Promise.reject(new Error('Aborted'));
}
return Promise.resolve({
params: { ...params, new_string: params.new_string ?? '' },
occurrences: 1,
}),
});
},
);
mockEnsureCorrectFileContent.mockImplementation(
async (content: string, _client: GeminiClient): Promise<string> =>
Promise.resolve(content ?? ''),
async (
content: string,
_client: GeminiClient,
signal?: AbortSignal,
): Promise<string> => {
// Make AbortSignal optional
if (signal?.aborted) {
return Promise.reject(new Error('Aborted'));
}
return Promise.resolve(content ?? '');
},
);
});
@@ -181,6 +195,7 @@ describe('WriteFileTool', () => {
const filePath = path.join(rootDir, 'new_corrected_file.txt');
const proposedContent = 'Proposed new content.';
const correctedContent = 'Corrected new content.';
const abortSignal = new AbortController().signal;
// Ensure the mock is set for this specific test case if needed, or rely on beforeEach
mockEnsureCorrectFileContent.mockResolvedValue(correctedContent);
@@ -188,11 +203,13 @@ describe('WriteFileTool', () => {
const result = await tool._getCorrectedFileContent(
filePath,
proposedContent,
abortSignal,
);
expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith(
proposedContent,
mockGeminiClientInstance,
abortSignal,
);
expect(mockEnsureCorrectEdit).not.toHaveBeenCalled();
expect(result.correctedContent).toBe(correctedContent);
@@ -206,6 +223,7 @@ describe('WriteFileTool', () => {
const originalContent = 'Original existing content.';
const proposedContent = 'Proposed replacement content.';
const correctedProposedContent = 'Corrected replacement content.';
const abortSignal = new AbortController().signal;
fs.writeFileSync(filePath, originalContent, 'utf8');
// Ensure this mock is active and returns the correct structure
@@ -222,6 +240,7 @@ describe('WriteFileTool', () => {
const result = await tool._getCorrectedFileContent(
filePath,
proposedContent,
abortSignal,
);
expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
@@ -232,6 +251,7 @@ describe('WriteFileTool', () => {
file_path: filePath,
},
mockGeminiClientInstance,
abortSignal,
);
expect(mockEnsureCorrectFileContent).not.toHaveBeenCalled();
expect(result.correctedContent).toBe(correctedProposedContent);
@@ -243,6 +263,7 @@ describe('WriteFileTool', () => {
it('should return error if reading an existing file fails (e.g. permissions)', async () => {
const filePath = path.join(rootDir, 'unreadable_file.txt');
const proposedContent = 'some content';
const abortSignal = new AbortController().signal;
fs.writeFileSync(filePath, 'content', { mode: 0o000 });
const readError = new Error('Permission denied');
@@ -255,6 +276,7 @@ describe('WriteFileTool', () => {
const result = await tool._getCorrectedFileContent(
filePath,
proposedContent,
abortSignal,
);
expect(fs.readFileSync).toHaveBeenCalledWith(filePath, 'utf8');
@@ -274,16 +296,17 @@ describe('WriteFileTool', () => {
});
describe('shouldConfirmExecute', () => {
const abortSignal = new AbortController().signal;
it('should return false if params are invalid (relative path)', async () => {
const params = { file_path: 'relative.txt', content: 'test' };
const confirmation = await tool.shouldConfirmExecute(params);
const confirmation = await tool.shouldConfirmExecute(params, abortSignal);
expect(confirmation).toBe(false);
});
it('should return false if params are invalid (outside root)', async () => {
const outsidePath = path.resolve(tempDir, 'outside-root.txt');
const params = { file_path: outsidePath, content: 'test' };
const confirmation = await tool.shouldConfirmExecute(params);
const confirmation = await tool.shouldConfirmExecute(params, abortSignal);
expect(confirmation).toBe(false);
});
@@ -298,7 +321,7 @@ describe('WriteFileTool', () => {
throw readError;
});
const confirmation = await tool.shouldConfirmExecute(params);
const confirmation = await tool.shouldConfirmExecute(params, abortSignal);
expect(confirmation).toBe(false);
vi.spyOn(fs, 'readFileSync').mockImplementation(originalReadFileSync);
@@ -314,11 +337,13 @@ describe('WriteFileTool', () => {
const params = { file_path: filePath, content: proposedContent };
const confirmation = (await tool.shouldConfirmExecute(
params,
abortSignal,
)) as ToolEditConfirmationDetails;
expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith(
proposedContent,
mockGeminiClientInstance,
abortSignal,
);
expect(confirmation).toEqual(
expect.objectContaining({
@@ -343,7 +368,6 @@ describe('WriteFileTool', () => {
'Corrected replacement for confirmation.';
fs.writeFileSync(filePath, originalContent, 'utf8');
// Ensure this mock is active and returns the correct structure
mockEnsureCorrectEdit.mockResolvedValue({
params: {
file_path: filePath,
@@ -356,6 +380,7 @@ describe('WriteFileTool', () => {
const params = { file_path: filePath, content: proposedContent };
const confirmation = (await tool.shouldConfirmExecute(
params,
abortSignal,
)) as ToolEditConfirmationDetails;
expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
@@ -366,6 +391,7 @@ describe('WriteFileTool', () => {
file_path: filePath,
},
mockGeminiClientInstance,
abortSignal,
);
expect(confirmation).toEqual(
expect.objectContaining({
@@ -381,9 +407,10 @@ describe('WriteFileTool', () => {
});
describe('execute', () => {
const abortSignal = new AbortController().signal;
it('should return error if params are invalid (relative path)', async () => {
const params = { file_path: 'relative.txt', content: 'test' };
const result = await tool.execute(params, new AbortController().signal);
const result = await tool.execute(params, abortSignal);
expect(result.llmContent).toMatch(/Error: Invalid parameters provided/);
expect(result.returnDisplay).toMatch(/Error: File path must be absolute/);
});
@@ -391,7 +418,7 @@ describe('WriteFileTool', () => {
it('should return error if params are invalid (path outside root)', async () => {
const outsidePath = path.resolve(tempDir, 'outside-root.txt');
const params = { file_path: outsidePath, content: 'test' };
const result = await tool.execute(params, new AbortController().signal);
const result = await tool.execute(params, abortSignal);
expect(result.llmContent).toMatch(/Error: Invalid parameters provided/);
expect(result.returnDisplay).toMatch(
/Error: File path must be within the root directory/,
@@ -409,7 +436,7 @@ describe('WriteFileTool', () => {
throw readError;
});
const result = await tool.execute(params, new AbortController().signal);
const result = await tool.execute(params, abortSignal);
expect(result.llmContent).toMatch(/Error checking existing file/);
expect(result.returnDisplay).toMatch(
/Error checking existing file: Simulated read error for execute/,
@@ -427,16 +454,20 @@ describe('WriteFileTool', () => {
const params = { file_path: filePath, content: proposedContent };
const confirmDetails = await tool.shouldConfirmExecute(params);
const confirmDetails = await tool.shouldConfirmExecute(
params,
abortSignal,
);
if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) {
await confirmDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce);
}
const result = await tool.execute(params, new AbortController().signal);
const result = await tool.execute(params, abortSignal);
expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith(
proposedContent,
mockGeminiClientInstance,
abortSignal,
);
expect(result.llmContent).toMatch(
/Successfully created and wrote to new file/,
@@ -477,12 +508,15 @@ describe('WriteFileTool', () => {
const params = { file_path: filePath, content: proposedContent };
const confirmDetails = await tool.shouldConfirmExecute(params);
const confirmDetails = await tool.shouldConfirmExecute(
params,
abortSignal,
);
if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) {
await confirmDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce);
}
const result = await tool.execute(params, new AbortController().signal);
const result = await tool.execute(params, abortSignal);
expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
initialContent,
@@ -492,6 +526,7 @@ describe('WriteFileTool', () => {
file_path: filePath,
},
mockGeminiClientInstance,
abortSignal,
);
expect(result.llmContent).toMatch(/Successfully overwrote file/);
expect(fs.readFileSync(filePath, 'utf8')).toBe(correctedProposedContent);
@@ -513,12 +548,15 @@ describe('WriteFileTool', () => {
const params = { file_path: filePath, content };
// Simulate confirmation if your logic requires it before execute, or remove if not needed for this path
const confirmDetails = await tool.shouldConfirmExecute(params);
const confirmDetails = await tool.shouldConfirmExecute(
params,
abortSignal,
);
if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) {
await confirmDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce);
}
await tool.execute(params, new AbortController().signal);
await tool.execute(params, abortSignal);
expect(fs.existsSync(dirPath)).toBe(true);
expect(fs.statSync(dirPath).isDirectory()).toBe(true);