diff options
Diffstat (limited to 'packages/core/src/tools')
| -rw-r--r-- | packages/core/src/tools/edit.test.ts | 142 | ||||
| -rw-r--r-- | packages/core/src/tools/edit.ts | 144 | ||||
| -rw-r--r-- | packages/core/src/tools/tools.ts | 3 |
3 files changed, 287 insertions, 2 deletions
diff --git a/packages/core/src/tools/edit.test.ts b/packages/core/src/tools/edit.test.ts index 8d7e2d14..2f6a6642 100644 --- a/packages/core/src/tools/edit.test.ts +++ b/packages/core/src/tools/edit.test.ts @@ -8,6 +8,7 @@ const mockEnsureCorrectEdit = vi.hoisted(() => vi.fn()); const mockGenerateJson = vi.hoisted(() => vi.fn()); +const mockOpenDiff = vi.hoisted(() => vi.fn()); vi.mock('../utils/editCorrector.js', () => ({ ensureCorrectEdit: mockEnsureCorrectEdit, @@ -19,6 +20,10 @@ vi.mock('../core/client.js', () => ({ })), })); +vi.mock('../utils/editor.js', () => ({ + openDiff: mockOpenDiff, +})); + import { describe, it, expect, beforeEach, afterEach, vi, Mock } from 'vitest'; import { EditTool, EditToolParams } from './edit.js'; import { FileDiff } from './tools.js'; @@ -27,6 +32,7 @@ import fs from 'fs'; import os from 'os'; import { ApprovalMode, Config } from '../config/config.js'; import { Content, Part, SchemaUnion } from '@google/genai'; +import { ToolConfirmationOutcome } from './tools.js'; describe('EditTool', () => { let tool: EditTool; @@ -715,4 +721,140 @@ function makeRequest() { ); }); }); + + describe('onModify', () => { + const testFile = 'some_file.txt'; + let filePath: string; + const diffDir = path.join(os.tmpdir(), 'gemini-cli-edit-tool-diffs'); + + beforeEach(() => { + filePath = path.join(rootDir, testFile); + mockOpenDiff.mockClear(); + }); + + afterEach(() => { + fs.rmSync(diffDir, { recursive: true, force: true }); + }); + + it('should create temporary files, call openDiff, and return updated params with diff', async () => { + const originalContent = 'original content'; + const params: EditToolParams = { + file_path: filePath, + edits: [ + { old_string: originalContent, new_string: 'modified content' }, + ], + }; + + fs.writeFileSync(filePath, originalContent, 'utf8'); + + const result = await tool.onModify( + params, + new AbortController().signal, + ToolConfirmationOutcome.ModifyVSCode, + ); + + expect(mockOpenDiff).toHaveBeenCalledTimes(1); + const [oldPath, newPath] = mockOpenDiff.mock.calls[0]; + expect(oldPath).toMatch( + /gemini-cli-edit-tool-diffs[/\\]gemini-cli-edit-some_file\.txt-old-\d+/, + ); + expect(newPath).toMatch( + /gemini-cli-edit-tool-diffs[/\\]gemini-cli-edit-some_file\.txt-new-\d+/, + ); + + expect(result).toBeDefined(); + expect(result!.updatedParams).toEqual({ + file_path: filePath, + edits: [ + { old_string: originalContent, new_string: 'modified content' }, + ], + }); + expect(result!.updatedDiff).toEqual(`Index: some_file.txt +=================================================================== +--- some_file.txt\tCurrent ++++ some_file.txt\tProposed +@@ -1,1 +1,1 @@ +-original content +\\ No newline at end of file ++modified content +\\ No newline at end of file +`); + + // Verify temp files are cleaned up + expect(fs.existsSync(oldPath)).toBe(false); + expect(fs.existsSync(newPath)).toBe(false); + }); + + it('should handle non-existent files and return updated params', async () => { + const params: EditToolParams = { + file_path: filePath, + edits: [{ old_string: '', new_string: 'new file content' }], + }; + + const result = await tool.onModify( + params, + new AbortController().signal, + ToolConfirmationOutcome.ModifyVSCode, + ); + + expect(mockOpenDiff).toHaveBeenCalledTimes(1); + + const [oldPath, newPath] = mockOpenDiff.mock.calls[0]; + + expect(result).toBeDefined(); + expect(result!.updatedParams).toEqual({ + file_path: filePath, + edits: [{ old_string: '', new_string: 'new file content' }], + }); + expect(result!.updatedDiff).toContain('new file content'); + + // Verify temp files are cleaned up + expect(fs.existsSync(oldPath)).toBe(false); + expect(fs.existsSync(newPath)).toBe(false); + }); + + it('should clean up previous temp files before creating new ones', async () => { + const params: EditToolParams = { + file_path: filePath, + edits: [{ old_string: 'old', new_string: 'new' }], + }; + + fs.writeFileSync(filePath, 'some old content', 'utf8'); + + // Call onModify first time + const result1 = await tool.onModify( + params, + new AbortController().signal, + ToolConfirmationOutcome.ModifyVSCode, + ); + const firstCall = mockOpenDiff.mock.calls[0]; + const firstOldPath = firstCall[0]; + const firstNewPath = firstCall[1]; + + expect(result1).toBeDefined(); + expect(fs.existsSync(firstOldPath)).toBe(false); + expect(fs.existsSync(firstNewPath)).toBe(false); + + // Ensure different timestamps so that the file names are different for testing. + await new Promise((resolve) => setTimeout(resolve, 2)); + + const result2 = await tool.onModify( + params, + new AbortController().signal, + ToolConfirmationOutcome.ModifyVSCode, + ); + const secondCall = mockOpenDiff.mock.calls[1]; + const secondOldPath = secondCall[0]; + const secondNewPath = secondCall[1]; + + // Call onModify second time + expect(result2).toBeDefined(); + expect(fs.existsSync(secondOldPath)).toBe(false); + expect(fs.existsSync(secondNewPath)).toBe(false); + + // Verify different file names were used + expect(firstOldPath).not.toBe(secondOldPath); + expect(firstNewPath).not.toBe(secondNewPath); + }); + }); }); diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts index eb22d0ac..f6a05ac1 100644 --- a/packages/core/src/tools/edit.ts +++ b/packages/core/src/tools/edit.ts @@ -4,8 +4,9 @@ * SPDX-License-Identifier: Apache-2.0 */ -import fs from 'fs'; -import path from 'path'; +import * as fs from 'fs'; +import * as path from 'path'; +import * as os from 'os'; import * as Diff from 'diff'; import { BaseTool, @@ -22,6 +23,7 @@ import { GeminiClient } from '../core/client.js'; import { Config, ApprovalMode } from '../config/config.js'; import { ensureCorrectEdit } from '../utils/editCorrector.js'; import { DEFAULT_DIFF_OPTIONS } from './diffOptions.js'; +import { openDiff } from '../utils/editor.js'; import { ReadFileTool } from './read-file.js'; /** @@ -75,6 +77,8 @@ export class EditTool extends BaseTool<EditToolParams, EditResult> { private readonly config: Config; private readonly rootDirectory: string; private readonly client: GeminiClient; + private tempOldDiffPath?: string; + private tempNewDiffPath?: string; /** * Creates a new instance of the EditLogic @@ -515,6 +519,142 @@ Expectation for required parameters: } /** + * Creates temp files for the current and proposed file contents and opens a diff tool. + * When the diff tool is closed, the tool will check if the file has been modified and provide the updated params. + * @returns Updated params and diff if the file has been modified, undefined otherwise. + */ + async onModify( + params: EditToolParams, + _abortSignal: AbortSignal, + outcome: ToolConfirmationOutcome, + ): Promise< + { updatedParams: EditToolParams; updatedDiff: string } | undefined + > { + const { oldPath, newPath } = this.createTempFiles(params); + this.tempOldDiffPath = oldPath; + this.tempNewDiffPath = newPath; + + await openDiff( + this.tempOldDiffPath, + this.tempNewDiffPath, + outcome === ToolConfirmationOutcome.ModifyVSCode ? 'vscode' : 'vim', + ); + return await this.getUpdatedParamsIfModified(params, _abortSignal); + } + + private async getUpdatedParamsIfModified( + params: EditToolParams, + _abortSignal: AbortSignal, + ): Promise< + { updatedParams: EditToolParams; updatedDiff: string } | undefined + > { + if (!this.tempOldDiffPath || !this.tempNewDiffPath) return undefined; + let oldContent = ''; + let newContent = ''; + try { + oldContent = fs.readFileSync(this.tempOldDiffPath, 'utf8'); + } catch (err) { + if (!isNodeError(err) || err.code !== 'ENOENT') throw err; + oldContent = ''; + } + try { + newContent = fs.readFileSync(this.tempNewDiffPath, 'utf8'); + } catch (err) { + if (!isNodeError(err) || err.code !== 'ENOENT') throw err; + newContent = ''; + } + + // Combine the edits into a single edit + const updatedParams: EditToolParams = { + ...params, + edits: [ + { + old_string: oldContent, + new_string: newContent, + }, + ], + }; + + const updatedDiff = Diff.createPatch( + path.basename(params.file_path), + oldContent, + newContent, + 'Current', + 'Proposed', + DEFAULT_DIFF_OPTIONS, + ); + + this.deleteTempFiles(); + return { updatedParams, updatedDiff }; + } + + private createTempFiles(params: EditToolParams): Record<string, string> { + this.deleteTempFiles(); + + const tempDir = os.tmpdir(); + const diffDir = path.join(tempDir, 'gemini-cli-edit-tool-diffs'); + + if (!fs.existsSync(diffDir)) { + fs.mkdirSync(diffDir, { recursive: true }); + } + + const fileName = path.basename(params.file_path); + const timestamp = Date.now(); + const tempOldPath = path.join( + diffDir, + `gemini-cli-edit-${fileName}-old-${timestamp}`, + ); + const tempNewPath = path.join( + diffDir, + `gemini-cli-edit-${fileName}-new-${timestamp}`, + ); + + let currentContent = ''; + try { + currentContent = fs.readFileSync(params.file_path, 'utf8'); + } catch (err) { + if (!isNodeError(err) || err.code !== 'ENOENT') throw err; + currentContent = ''; + } + + let proposedContent = currentContent; + for (const edit of params.edits) { + proposedContent = this._applyReplacement( + proposedContent, + edit.old_string, + edit.new_string, + edit.old_string === '' && currentContent === '', + ); + } + + fs.writeFileSync(tempOldPath, currentContent, 'utf8'); + fs.writeFileSync(tempNewPath, proposedContent, 'utf8'); + return { + oldPath: tempOldPath, + newPath: tempNewPath, + }; + } + + private deleteTempFiles(): void { + try { + if (this.tempOldDiffPath) { + fs.unlinkSync(this.tempOldDiffPath); + this.tempOldDiffPath = undefined; + } + } catch { + console.error(`Error deleting temp diff file: `, this.tempOldDiffPath); + } + try { + if (this.tempNewDiffPath) { + fs.unlinkSync(this.tempNewDiffPath); + this.tempNewDiffPath = undefined; + } + } catch { + console.error(`Error deleting temp diff file: `, this.tempNewDiffPath); + } + } + + /** * Creates parent directories if they don't exist */ private ensureParentDirectoriesExist(filePath: string): void { diff --git a/packages/core/src/tools/tools.ts b/packages/core/src/tools/tools.ts index 1b932229..6a1be9c9 100644 --- a/packages/core/src/tools/tools.ts +++ b/packages/core/src/tools/tools.ts @@ -202,6 +202,7 @@ export interface ToolEditConfirmationDetails { onConfirm: (outcome: ToolConfirmationOutcome) => Promise<void>; fileName: string; fileDiff: string; + isModifying?: boolean; } export interface ToolExecuteConfirmationDetails { @@ -231,5 +232,7 @@ export enum ToolConfirmationOutcome { ProceedAlways = 'proceed_always', ProceedAlwaysServer = 'proceed_always_server', ProceedAlwaysTool = 'proceed_always_tool', + ModifyVSCode = 'modify_vscode', + ModifyVim = 'modify_vim', Cancel = 'cancel', } |
