feat: Add flow to allow modifying edits during edit tool call (#808)

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);
+    });
+  });
 });

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
@@ -514,6 +518,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
    */

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',
 }