From 2c6aae863af084fdd515c56a41ef01cd9d7b2c0b Mon Sep 17 00:00:00 2001 From: Leo <45218470+ngleo@users.noreply.github.com> Date: Sat, 14 Jun 2025 19:20:04 +0100 Subject: [PATCH] Enable "modify" in write tool (#1044) --- packages/core/src/core/coreToolScheduler.ts | 37 +- packages/core/src/tools/edit.test.ts | 137 ------- packages/core/src/tools/edit.ts | 172 ++------ .../core/src/tools/modifiable-tool.test.ts | 367 ++++++++++++++++++ packages/core/src/tools/modifiable-tool.ts | 163 ++++++++ packages/core/src/tools/write-file.ts | 37 +- 6 files changed, 626 insertions(+), 287 deletions(-) create mode 100644 packages/core/src/tools/modifiable-tool.test.ts create mode 100644 packages/core/src/tools/modifiable-tool.ts diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index d17e1511..14a39792 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -13,14 +13,17 @@ import { ToolResult, ToolRegistry, ApprovalMode, - EditTool, - EditToolParams, EditorType, Config, logToolCall, } from '../index.js'; import { Part, PartListUnion } from '@google/genai'; import { getResponseTextFromParts } from '../utils/generateContentResponseUtilities.js'; +import { + isModifiableTool, + ModifyContext, + modifyWithEditor, +} from '../tools/modifiable-tool.js'; export type ValidatingToolCall = { status: 'validating'; @@ -525,8 +528,8 @@ export class CoreToolScheduler { ); } else if (outcome === ToolConfirmationOutcome.ModifyWithEditor) { const waitingToolCall = toolCall as WaitingToolCall; - if (waitingToolCall?.confirmationDetails?.type === 'edit') { - const editTool = waitingToolCall.tool as EditTool; + if (isModifiableTool(waitingToolCall.tool)) { + const modifyContext = waitingToolCall.tool.getModifyContext(signal); const editorType = this.getPreferredEditor(); if (!editorType) { return; @@ -535,22 +538,22 @@ export class CoreToolScheduler { this.setStatusInternal(callId, 'awaiting_approval', { ...waitingToolCall.confirmationDetails, isModifying: true, - }); + } as ToolCallConfirmationDetails); - const modifyResults = await editTool.onModify( - waitingToolCall.request.args as unknown as EditToolParams, - signal, + const { updatedParams, updatedDiff } = await modifyWithEditor< + typeof waitingToolCall.request.args + >( + waitingToolCall.request.args, + modifyContext as ModifyContext, editorType, + signal, ); - - if (modifyResults) { - this.setArgsInternal(callId, modifyResults.updatedParams); - this.setStatusInternal(callId, 'awaiting_approval', { - ...waitingToolCall.confirmationDetails, - fileDiff: modifyResults.updatedDiff, - isModifying: false, - }); - } + this.setArgsInternal(callId, updatedParams); + this.setStatusInternal(callId, 'awaiting_approval', { + ...waitingToolCall.confirmationDetails, + fileDiff: updatedDiff, + isModifying: false, + } as ToolCallConfirmationDetails); } } else { this.setStatusInternal(callId, 'scheduled'); diff --git a/packages/core/src/tools/edit.test.ts b/packages/core/src/tools/edit.test.ts index dd3b481d..7dff5bf9 100644 --- a/packages/core/src/tools/edit.test.ts +++ b/packages/core/src/tools/edit.test.ts @@ -605,141 +605,4 @@ describe('EditTool', () => { ); }); }); - - 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, - old_string: originalContent, - new_string: 'modified content', - }; - - fs.writeFileSync(filePath, originalContent, 'utf8'); - - const result = await tool.onModify( - params, - new AbortController().signal, - 'vscode', - ); - - 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, - 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, - old_string: '', - new_string: 'new file content', - }; - - const result = await tool.onModify( - params, - new AbortController().signal, - 'vscode', - ); - - expect(mockOpenDiff).toHaveBeenCalledTimes(1); - - const [oldPath, newPath] = mockOpenDiff.mock.calls[0]; - - expect(result).toBeDefined(); - expect(result!.updatedParams).toEqual({ - file_path: filePath, - 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, - 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, - 'vscode', - ); - 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, - 'vscode', - ); - 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 39352121..3317d460 100644 --- a/packages/core/src/tools/edit.ts +++ b/packages/core/src/tools/edit.ts @@ -6,7 +6,6 @@ import * as fs from 'fs'; import * as path from 'path'; -import * as os from 'os'; import * as Diff from 'diff'; import { BaseTool, @@ -23,9 +22,8 @@ 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'; -import { EditorType } from '../utils/editor.js'; +import { ModifiableTool, ModifyContext } from './modifiable-tool.js'; /** * Parameters for the Edit tool @@ -64,13 +62,14 @@ interface CalculatedEdit { /** * Implementation of the Edit tool logic */ -export class EditTool extends BaseTool { +export class EditTool + extends BaseTool + implements ModifiableTool +{ static readonly Name = 'replace'; 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 @@ -432,132 +431,6 @@ 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, - editorType: EditorType, - ): Promise< - { updatedParams: EditToolParams; updatedDiff: string } | undefined - > { - const { oldPath, newPath } = this.createTempFiles(params); - this.tempOldDiffPath = oldPath; - this.tempNewDiffPath = newPath; - - await openDiff(this.tempOldDiffPath, this.tempNewDiffPath, editorType); - 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, - 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 { - 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; - proposedContent = this._applyReplacement( - proposedContent, - params.old_string, - params.new_string, - params.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 */ @@ -567,4 +440,39 @@ Expectation for required parameters: fs.mkdirSync(dirName, { recursive: true }); } } + + getModifyContext(_: AbortSignal): ModifyContext { + return { + getFilePath: (params: EditToolParams) => params.file_path, + getCurrentContent: async (params: EditToolParams): Promise => { + try { + return fs.readFileSync(params.file_path, 'utf8'); + } catch (err) { + if (!isNodeError(err) || err.code !== 'ENOENT') throw err; + return ''; + } + }, + getProposedContent: async (params: EditToolParams): Promise => { + try { + const currentContent = fs.readFileSync(params.file_path, 'utf8'); + return this._applyReplacement( + currentContent, + params.old_string, + params.new_string, + params.old_string === '' && currentContent === '', + ); + } catch (err) { + if (!isNodeError(err) || err.code !== 'ENOENT') throw err; + return ''; + } + }, + createUpdatedParams: ( + modifiedProposedContent: string, + originalParams: EditToolParams, + ): EditToolParams => ({ + ...originalParams, + new_string: modifiedProposedContent, + }), + }; + } } diff --git a/packages/core/src/tools/modifiable-tool.test.ts b/packages/core/src/tools/modifiable-tool.test.ts new file mode 100644 index 00000000..56c27fe0 --- /dev/null +++ b/packages/core/src/tools/modifiable-tool.test.ts @@ -0,0 +1,367 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + vi, + describe, + it, + expect, + beforeEach, + afterEach, + type Mock, +} from 'vitest'; +import { + modifyWithEditor, + ModifyContext, + ModifiableTool, + isModifiableTool, +} from './modifiable-tool.js'; +import { EditorType } from '../utils/editor.js'; +import fs from 'fs'; +import os from 'os'; +import * as path from 'path'; + +// Mock dependencies +const mockOpenDiff = vi.hoisted(() => vi.fn()); +const mockCreatePatch = vi.hoisted(() => vi.fn()); + +vi.mock('../utils/editor.js', () => ({ + openDiff: mockOpenDiff, +})); + +vi.mock('diff', () => ({ + createPatch: mockCreatePatch, +})); + +vi.mock('fs'); +vi.mock('os'); + +interface TestParams { + filePath: string; + someOtherParam: string; + modifiedContent?: string; +} + +describe('modifyWithEditor', () => { + let tempDir: string; + let mockModifyContext: ModifyContext; + let mockParams: TestParams; + let currentContent: string; + let proposedContent: string; + let modifiedContent: string; + let abortSignal: AbortSignal; + + beforeEach(() => { + vi.resetAllMocks(); + + tempDir = '/tmp/test-dir'; + abortSignal = new AbortController().signal; + + currentContent = 'original content\nline 2\nline 3'; + proposedContent = 'modified content\nline 2\nline 3'; + modifiedContent = 'user modified content\nline 2\nline 3\nnew line'; + mockParams = { + filePath: path.join(tempDir, 'test.txt'), + someOtherParam: 'value', + }; + + mockModifyContext = { + getFilePath: vi.fn().mockReturnValue(mockParams.filePath), + getCurrentContent: vi.fn().mockResolvedValue(currentContent), + getProposedContent: vi.fn().mockResolvedValue(proposedContent), + createUpdatedParams: vi + .fn() + .mockImplementation((modifiedContent, originalParams) => ({ + ...originalParams, + modifiedContent, + })), + }; + + (os.tmpdir as Mock).mockReturnValue(tempDir); + + (fs.existsSync as Mock).mockReturnValue(true); + (fs.mkdirSync as Mock).mockImplementation(() => undefined); + (fs.writeFileSync as Mock).mockImplementation(() => {}); + (fs.unlinkSync as Mock).mockImplementation(() => {}); + + (fs.readFileSync as Mock).mockImplementation((filePath: string) => { + if (filePath.includes('-new-')) { + return modifiedContent; + } + return currentContent; + }); + + mockCreatePatch.mockReturnValue('mock diff content'); + mockOpenDiff.mockResolvedValue(undefined); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + describe('successful modification', () => { + it('should successfully modify content with VSCode editor', async () => { + const result = await modifyWithEditor( + mockParams, + mockModifyContext, + 'vscode' as EditorType, + abortSignal, + ); + + expect(mockModifyContext.getCurrentContent).toHaveBeenCalledWith( + mockParams, + ); + expect(mockModifyContext.getProposedContent).toHaveBeenCalledWith( + mockParams, + ); + expect(mockModifyContext.getFilePath).toHaveBeenCalledWith(mockParams); + + expect(fs.writeFileSync).toHaveBeenCalledTimes(2); + expect(fs.writeFileSync).toHaveBeenNthCalledWith( + 1, + expect.stringContaining( + path.join(tempDir, 'gemini-cli-tool-modify-diffs'), + ), + currentContent, + 'utf8', + ); + expect(fs.writeFileSync).toHaveBeenNthCalledWith( + 2, + expect.stringContaining( + path.join(tempDir, 'gemini-cli-tool-modify-diffs'), + ), + proposedContent, + 'utf8', + ); + + expect(mockOpenDiff).toHaveBeenCalledWith( + expect.stringContaining('-old-'), + expect.stringContaining('-new-'), + 'vscode', + ); + + expect(fs.readFileSync).toHaveBeenCalledWith( + expect.stringContaining('-old-'), + 'utf8', + ); + expect(fs.readFileSync).toHaveBeenCalledWith( + expect.stringContaining('-new-'), + 'utf8', + ); + + expect(mockModifyContext.createUpdatedParams).toHaveBeenCalledWith( + modifiedContent, + mockParams, + ); + + expect(mockCreatePatch).toHaveBeenCalledWith( + path.basename(mockParams.filePath), + currentContent, + modifiedContent, + 'Current', + 'Proposed', + expect.objectContaining({ + context: 3, + ignoreWhitespace: true, + }), + ); + + expect(fs.unlinkSync).toHaveBeenCalledTimes(2); + expect(fs.unlinkSync).toHaveBeenNthCalledWith( + 1, + expect.stringContaining('-old-'), + ); + expect(fs.unlinkSync).toHaveBeenNthCalledWith( + 2, + expect.stringContaining('-new-'), + ); + + expect(result).toEqual({ + updatedParams: { + ...mockParams, + modifiedContent, + }, + updatedDiff: 'mock diff content', + }); + }); + + it('should create temp directory if it does not exist', async () => { + (fs.existsSync as Mock).mockReturnValue(false); + + await modifyWithEditor( + mockParams, + mockModifyContext, + 'vscode' as EditorType, + abortSignal, + ); + + expect(fs.mkdirSync).toHaveBeenCalledWith( + path.join(tempDir, 'gemini-cli-tool-modify-diffs'), + { recursive: true }, + ); + }); + + it('should not create temp directory if it already exists', async () => { + (fs.existsSync as Mock).mockReturnValue(true); + + await modifyWithEditor( + mockParams, + mockModifyContext, + 'vscode' as EditorType, + abortSignal, + ); + + expect(fs.mkdirSync).not.toHaveBeenCalled(); + }); + }); + + it('should handle missing old temp file gracefully', async () => { + (fs.readFileSync as Mock).mockImplementation((filePath: string) => { + if (filePath.includes('-old-')) { + const error = new Error('ENOENT: no such file or directory'); + (error as NodeJS.ErrnoException).code = 'ENOENT'; + throw error; + } + return modifiedContent; + }); + + const result = await modifyWithEditor( + mockParams, + mockModifyContext, + 'vscode' as EditorType, + abortSignal, + ); + + expect(mockCreatePatch).toHaveBeenCalledWith( + path.basename(mockParams.filePath), + '', + modifiedContent, + 'Current', + 'Proposed', + expect.objectContaining({ + context: 3, + ignoreWhitespace: true, + }), + ); + + expect(result.updatedParams).toBeDefined(); + expect(result.updatedDiff).toBe('mock diff content'); + }); + + it('should handle missing new temp file gracefully', async () => { + (fs.readFileSync as Mock).mockImplementation((filePath: string) => { + if (filePath.includes('-new-')) { + const error = new Error('ENOENT: no such file or directory'); + (error as NodeJS.ErrnoException).code = 'ENOENT'; + throw error; + } + return currentContent; + }); + + const result = await modifyWithEditor( + mockParams, + mockModifyContext, + 'vscode' as EditorType, + abortSignal, + ); + + expect(mockCreatePatch).toHaveBeenCalledWith( + path.basename(mockParams.filePath), + currentContent, + '', + 'Current', + 'Proposed', + expect.objectContaining({ + context: 3, + ignoreWhitespace: true, + }), + ); + + expect(result.updatedParams).toBeDefined(); + expect(result.updatedDiff).toBe('mock diff content'); + }); + + it('should clean up temp files even if editor fails', async () => { + const editorError = new Error('Editor failed to open'); + mockOpenDiff.mockRejectedValue(editorError); + + await expect( + modifyWithEditor( + mockParams, + mockModifyContext, + 'vscode' as EditorType, + abortSignal, + ), + ).rejects.toThrow('Editor failed to open'); + + expect(fs.unlinkSync).toHaveBeenCalledTimes(2); + }); + + it('should handle temp file cleanup errors gracefully', async () => { + const consoleErrorSpy = vi + .spyOn(console, 'error') + .mockImplementation(() => {}); + (fs.unlinkSync as Mock).mockImplementation((_filePath: string) => { + throw new Error('Failed to delete file'); + }); + + await modifyWithEditor( + mockParams, + mockModifyContext, + 'vscode' as EditorType, + abortSignal, + ); + + expect(consoleErrorSpy).toHaveBeenCalledTimes(2); + expect(consoleErrorSpy).toHaveBeenCalledWith( + expect.stringContaining('Error deleting temp diff file:'), + ); + + consoleErrorSpy.mockRestore(); + }); + + it('should create temp files with correct naming', async () => { + const testFilePath = path.join(tempDir, 'subfolder', 'test-file.txt'); + mockModifyContext.getFilePath = vi.fn().mockReturnValue(testFilePath); + + await modifyWithEditor( + mockParams, + mockModifyContext, + 'vscode' as EditorType, + abortSignal, + ); + + const writeFileCalls = (fs.writeFileSync as Mock).mock.calls; + expect(writeFileCalls).toHaveLength(2); + + const oldFilePath = writeFileCalls[0][0]; + const newFilePath = writeFileCalls[1][0]; + + expect(oldFilePath).toMatch(/gemini-cli-modify-test-file\.txt-old-\d+$/); + expect(newFilePath).toMatch(/gemini-cli-modify-test-file\.txt-new-\d+$/); + expect(oldFilePath).toContain(`${tempDir}/gemini-cli-tool-modify-diffs/`); + expect(newFilePath).toContain(`${tempDir}/gemini-cli-tool-modify-diffs/`); + }); +}); + +describe('isModifiableTool', () => { + it('should return true for objects with getModifyContext method', () => { + const mockTool = { + name: 'test-tool', + getModifyContext: vi.fn(), + } as unknown as ModifiableTool; + + expect(isModifiableTool(mockTool)).toBe(true); + }); + + it('should return false for objects without getModifyContext method', () => { + const mockTool = { + name: 'test-tool', + } as unknown as ModifiableTool; + + expect(isModifiableTool(mockTool)).toBe(false); + }); +}); diff --git a/packages/core/src/tools/modifiable-tool.ts b/packages/core/src/tools/modifiable-tool.ts new file mode 100644 index 00000000..96fe176c --- /dev/null +++ b/packages/core/src/tools/modifiable-tool.ts @@ -0,0 +1,163 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { EditorType } from '../utils/editor.js'; +import os from 'os'; +import path from 'path'; +import fs from 'fs'; +import * as Diff from 'diff'; +import { openDiff } from '../utils/editor.js'; +import { DEFAULT_DIFF_OPTIONS } from './diffOptions.js'; +import { isNodeError } from '../utils/errors.js'; +import { Tool } from './tools.js'; + +/** + * A tool that supports a modify operation. + */ +export interface ModifiableTool extends Tool { + getModifyContext(abortSignal: AbortSignal): ModifyContext; +} + +export interface ModifyContext { + getFilePath: (params: ToolParams) => string; + + getCurrentContent: (params: ToolParams) => Promise; + + getProposedContent: (params: ToolParams) => Promise; + + createUpdatedParams: ( + modifiedProposedContent: string, + originalParams: ToolParams, + ) => ToolParams; +} + +export interface ModifyResult { + updatedParams: ToolParams; + updatedDiff: string; +} + +export function isModifiableTool( + tool: Tool, +): tool is ModifiableTool { + return 'getModifyContext' in tool; +} + +function createTempFilesForModify( + currentContent: string, + proposedContent: string, + file_path: string, +): { oldPath: string; newPath: string } { + const tempDir = os.tmpdir(); + const diffDir = path.join(tempDir, 'gemini-cli-tool-modify-diffs'); + + if (!fs.existsSync(diffDir)) { + fs.mkdirSync(diffDir, { recursive: true }); + } + + const fileName = path.basename(file_path); + const timestamp = Date.now(); + const tempOldPath = path.join( + diffDir, + `gemini-cli-modify-${fileName}-old-${timestamp}`, + ); + const tempNewPath = path.join( + diffDir, + `gemini-cli-modify-${fileName}-new-${timestamp}`, + ); + + fs.writeFileSync(tempOldPath, currentContent, 'utf8'); + fs.writeFileSync(tempNewPath, proposedContent, 'utf8'); + + return { oldPath: tempOldPath, newPath: tempNewPath }; +} + +function getUpdatedParams( + tmpOldPath: string, + tempNewPath: string, + originalParams: ToolParams, + modifyContext: ModifyContext, +): { updatedParams: ToolParams; updatedDiff: string } { + let oldContent = ''; + let newContent = ''; + + try { + oldContent = fs.readFileSync(tmpOldPath, 'utf8'); + } catch (err) { + if (!isNodeError(err) || err.code !== 'ENOENT') throw err; + oldContent = ''; + } + + try { + newContent = fs.readFileSync(tempNewPath, 'utf8'); + } catch (err) { + if (!isNodeError(err) || err.code !== 'ENOENT') throw err; + newContent = ''; + } + + const updatedParams = modifyContext.createUpdatedParams( + newContent, + originalParams, + ); + const updatedDiff = Diff.createPatch( + path.basename(modifyContext.getFilePath(originalParams)), + oldContent, + newContent, + 'Current', + 'Proposed', + DEFAULT_DIFF_OPTIONS, + ); + + return { updatedParams, updatedDiff }; +} + +function deleteTempFiles(oldPath: string, newPath: string): void { + try { + fs.unlinkSync(oldPath); + } catch { + console.error(`Error deleting temp diff file: ${oldPath}`); + } + + try { + fs.unlinkSync(newPath); + } catch { + console.error(`Error deleting temp diff file: ${newPath}`); + } +} + +/** + * Triggers an external editor for the user to modify the proposed content, + * and returns the updated tool parameters and the diff after the user has modified the proposed content. + */ +export async function modifyWithEditor( + originalParams: ToolParams, + modifyContext: ModifyContext, + editorType: EditorType, + _abortSignal: AbortSignal, +): Promise> { + const currentContent = await modifyContext.getCurrentContent(originalParams); + const proposedContent = + await modifyContext.getProposedContent(originalParams); + + const { oldPath, newPath } = createTempFilesForModify( + currentContent, + proposedContent, + modifyContext.getFilePath(originalParams), + ); + + try { + await openDiff(oldPath, newPath, editorType); + const result = getUpdatedParams( + oldPath, + newPath, + originalParams, + modifyContext, + ); + + return result; + } finally { + deleteTempFiles(oldPath, newPath); + } +} diff --git a/packages/core/src/tools/write-file.ts b/packages/core/src/tools/write-file.ts index dc634cc8..b9e07034 100644 --- a/packages/core/src/tools/write-file.ts +++ b/packages/core/src/tools/write-file.ts @@ -25,6 +25,7 @@ import { } from '../utils/editCorrector.js'; import { GeminiClient } from '../core/client.js'; import { DEFAULT_DIFF_OPTIONS } from './diffOptions.js'; +import { ModifiableTool, ModifyContext } from './modifiable-tool.js'; /** * Parameters for the WriteFile tool @@ -51,7 +52,10 @@ interface GetCorrectedFileContentResult { /** * Implementation of the WriteFile tool logic */ -export class WriteFileTool extends BaseTool { +export class WriteFileTool + extends BaseTool + implements ModifiableTool +{ static readonly Name: string = 'write_file'; private readonly client: GeminiClient; @@ -336,4 +340,35 @@ export class WriteFileTool extends BaseTool { } return { originalContent, correctedContent, fileExists }; } + + getModifyContext( + abortSignal: AbortSignal, + ): ModifyContext { + return { + getFilePath: (params: WriteFileToolParams) => params.file_path, + getCurrentContent: async (params: WriteFileToolParams) => { + const correctedContentResult = await this._getCorrectedFileContent( + params.file_path, + params.content, + abortSignal, + ); + return correctedContentResult.originalContent; + }, + getProposedContent: async (params: WriteFileToolParams) => { + const correctedContentResult = await this._getCorrectedFileContent( + params.file_path, + params.content, + abortSignal, + ); + return correctedContentResult.correctedContent; + }, + createUpdatedParams: ( + modifiedProposedContent: string, + originalParams: WriteFileToolParams, + ) => ({ + ...originalParams, + content: modifiedProposedContent, + }), + }; + } }