Correct edits even when auto-accept is enabled.

- Prior to this when a user would turn on auto-accept for edits we'd stop ensuring correct edits. This would result in a lot of back and forth by the model. This change also incoporates ensure correct edit into the normal execution flow.
- Added edit tests for this.

Part of https://github.com/google-gemini/gemini-cli/issues/484

packages/server/src/tools/edit.test.ts

@@ -6,6 +6,19 @@
 
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
+const mockEnsureCorrectEdit = vi.hoisted(() => vi.fn());
+const mockGenerateJson = vi.hoisted(() => vi.fn());
+
+vi.mock('../utils/editCorrector.js', () => ({
+  ensureCorrectEdit: mockEnsureCorrectEdit,
+}));
+
+vi.mock('../core/client.js', () => ({
+  GeminiClient: vi.fn().mockImplementation(() => ({
+    generateJson: mockGenerateJson,
+  })),
+}));
+
 import { describe, it, expect, beforeEach, afterEach, vi, Mock } from 'vitest';
 import { EditTool, EditToolParams } from './edit.js';
 import { FileDiff } from './tools.js';
@@ -15,16 +28,6 @@ import os from 'os';
 import { Config } from '../config/config.js';
 import { Content, Part, SchemaUnion } from '@google/genai';
 
-// Mock GeminiClient
-const mockEnsureCorrectEdit = vi.fn();
-const mockGenerateJson = vi.fn();
-vi.mock('../core/client.js', () => ({
-  GeminiClient: vi.fn().mockImplementation(() => ({
-    ensureCorrectEdit: mockEnsureCorrectEdit,
-    generateJson: mockGenerateJson,
-  })),
-}));
-
 describe('EditTool', () => {
   let tool: EditTool;
   let tempDir: string;
@@ -133,6 +136,48 @@ describe('EditTool', () => {
     fs.rmSync(tempDir, { recursive: true, force: true });
   });
 
+  describe('_applyReplacement', () => {
+    // Access private method for testing
+    // Note: `tool` is initialized in `beforeEach` of the parent describe block
+    it('should return newString if isNewFile is true', () => {
+      expect((tool as any)._applyReplacement(null, 'old', 'new', true)).toBe(
+        'new',
+      );
+      expect(
+        (tool as any)._applyReplacement('existing', 'old', 'new', true),
+      ).toBe('new');
+    });
+
+    it('should return newString if currentContent is null and oldString is empty (defensive)', () => {
+      expect((tool as any)._applyReplacement(null, '', 'new', false)).toBe(
+        'new',
+      );
+    });
+
+    it('should return empty string if currentContent is null and oldString is not empty (defensive)', () => {
+      expect((tool as any)._applyReplacement(null, 'old', 'new', false)).toBe(
+        '',
+      );
+    });
+
+    it('should replace oldString with newString in currentContent', () => {
+      expect(
+        (tool as any)._applyReplacement(
+          'hello old world old',
+          'old',
+          'new',
+          false,
+        ),
+      ).toBe('hello new world new');
+    });
+
+    it('should return currentContent if oldString is empty and not a new file', () => {
+      expect(
+        (tool as any)._applyReplacement('hello world', '', 'new', false),
+      ).toBe('hello world');
+    });
+  });
+
   describe('validateParams', () => {
     it('should return null for valid params', () => {
       const params: EditToolParams = {
@@ -243,6 +288,68 @@ describe('EditTool', () => {
         }),
       );
     });
+
+    it('should use corrected params from ensureCorrectEdit for diff generation', async () => {
+      const originalContent = 'This is the original string to be replaced.';
+      const originalOldString = 'original string';
+      const originalNewString = 'new string';
+
+      const correctedOldString = 'original string to be replaced'; // More specific
+      const correctedNewString = 'completely new string'; // Different replacement
+      const expectedFinalContent = 'This is the completely new string.';
+
+      fs.writeFileSync(filePath, originalContent);
+      const params: EditToolParams = {
+        file_path: filePath,
+        old_string: originalOldString,
+        new_string: originalNewString,
+      };
+
+      // The main beforeEach already calls mockEnsureCorrectEdit.mockReset()
+      // Set a specific mock for this test case
+      let mockCalled = false;
+      mockEnsureCorrectEdit.mockImplementationOnce(
+        async (content, p, client) => {
+          console.log('mockEnsureCorrectEdit CALLED IN TEST');
+          mockCalled = true;
+          expect(content).toBe(originalContent);
+          expect(p).toBe(params);
+          expect(client).toBe((tool as any).client);
+          return {
+            params: {
+              file_path: filePath,
+              old_string: correctedOldString,
+              new_string: correctedNewString,
+            },
+            occurrences: 1,
+          };
+        },
+      );
+
+      const confirmation = (await tool.shouldConfirmExecute(
+        params,
+      )) as FileDiff;
+
+      expect(mockCalled).toBe(true); // Check if the mock implementation was run
+      // expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(originalContent, params, expect.anything()); // Keep this commented for now
+      expect(confirmation).toEqual(
+        expect.objectContaining({
+          title: `Confirm Edit: ${testFile}`,
+          fileName: testFile,
+        }),
+      );
+      // Check that the diff is based on the corrected strings leading to the new state
+      expect(confirmation.fileDiff).toContain(`-${originalContent}`);
+      expect(confirmation.fileDiff).toContain(`+${expectedFinalContent}`);
+
+      // Verify that applying the correctedOldString and correctedNewString to originalContent
+      // indeed produces the expectedFinalContent, which is what the diff should reflect.
+      const patchedContent = originalContent.replace(
+        correctedOldString, // This was the string identified by ensureCorrectEdit for replacement
+        correctedNewString, // This was the string identified by ensureCorrectEdit as the replacement
+      );
+      expect(patchedContent).toBe(expectedFinalContent);
+    });
   });
 
   describe('execute', () => {
@@ -337,9 +444,11 @@ describe('EditTool', () => {
       };
       // The default mockEnsureCorrectEdit will return 0 occurrences for 'nonexistent'
       const result = await tool.execute(params, new AbortController().signal);
-      expect(result.llmContent).toMatch(/0 occurrences found/);
+      expect(result.llmContent).toMatch(
+        /0 occurrences found for old_string in/,
+      );
       expect(result.returnDisplay).toMatch(
-        /Failed to edit, could not find the string to replace/,
+        /Failed to edit, could not find the string to replace./,
       );
     });
 
@@ -352,7 +461,9 @@ describe('EditTool', () => {
       };
       // The default mockEnsureCorrectEdit will return 2 occurrences for 'old'
       const result = await tool.execute(params, new AbortController().signal);
-      expect(result.llmContent).toMatch(/Expected 1 occurrences but found 2/);
+      expect(result.llmContent).toMatch(
+        /Expected 1 occurrences but found 2 for old_string in file/,
+      );
       expect(result.returnDisplay).toMatch(
         /Failed to edit, expected 1 occurrence\(s\) but found 2/,
       );

packages/server/src/tools/edit.ts

@@ -21,7 +21,7 @@ import { isNodeError } from '../utils/errors.js';
 import { ReadFileTool } from './read-file.js';
 import { GeminiClient } from '../core/client.js';
 import { Config } from '../config/config.js';
-import { countOccurrences, ensureCorrectEdit } from '../utils/editCorrector.js';
+import { ensureCorrectEdit } from '../utils/editCorrector.js';
 
 /**
  * Parameters for the Edit tool
@@ -147,6 +147,26 @@ Expectation for parameters:
     return null;
   }
 
+  private _applyReplacement(
+    currentContent: string | null,
+    oldString: string,
+    newString: string,
+    isNewFile: boolean,
+  ): string {
+    if (isNewFile) {
+      return newString;
+    }
+    if (currentContent === null) {
+      // Should not happen if not a new file, but defensively return empty or newString if oldString is also empty
+      return oldString === '' ? newString : '';
+    }
+    // If oldString is empty and it's not a new file, do not modify the content.
+    if (oldString === '' && !isNewFile) {
+      return currentContent;
+    }
+    return currentContent.replaceAll(oldString, newString);
+  }
+
   /**
    * Calculates the potential outcome of an edit operation.
    * @param params Parameters for the edit operation
@@ -158,7 +178,8 @@ Expectation for parameters:
     let currentContent: string | null = null;
     let fileExists = false;
     let isNewFile = false;
-    let newContent = '';
+    let finalNewString = params.new_string;
+    let finalOldString = params.old_string;
     let occurrences = 0;
     let error: { display: string; raw: string } | undefined = undefined;
 
@@ -176,8 +197,6 @@ Expectation for parameters:
     if (params.old_string === '' && !fileExists) {
       // Creating a new file
       isNewFile = true;
-      newContent = params.new_string;
-      occurrences = 0;
     } else if (!fileExists) {
       // Trying to edit a non-existent file (and old_string is not empty)
       error = {
@@ -186,7 +205,14 @@ Expectation for parameters:
       };
     } else if (currentContent !== null) {
       // Editing an existing file
-      occurrences = countOccurrences(currentContent, params.old_string);
+      const correctedEdit = await ensureCorrectEdit(
+        currentContent,
+        params,
+        this.client,
+      );
+      finalOldString = correctedEdit.params.old_string;
+      finalNewString = correctedEdit.params.new_string;
+      occurrences = correctedEdit.occurrences;
 
       if (params.old_string === '') {
         // Error: Trying to create a file that already exists
@@ -197,19 +223,13 @@ Expectation for parameters:
       } else if (occurrences === 0) {
         error = {
           display: `Failed to edit, could not find the string to replace.`,
-          raw: `Failed to edit, 0 occurrences found for old_string in ${params.file_path}. No edits made. The exact text in old_string was not found. Ensure you're not escaping content incorrectly and check whitespace, indentation, and context. Use ReadFile tool to verify.`,
+          raw: `Failed to edit, 0 occurrences found for old_string in ${params.file_path}. No edits made. The exact text in old_string was not found. Ensure you're not escaping content incorrectly and check whitespace, indentation, and context. Use ${ReadFileTool.Name} tool to verify.`,
         };
       } else if (occurrences !== expectedReplacements) {
         error = {
           display: `Failed to edit, expected ${expectedReplacements} occurrence(s) but found ${occurrences}.`,
           raw: `Failed to edit, Expected ${expectedReplacements} occurrences but found ${occurrences} for old_string in file: ${params.file_path}`,
         };
-      } else {
-        // Successful edit calculation
-        newContent = currentContent.replaceAll(
-          params.old_string,
-          params.new_string,
-        );
       }
     } else {
       // Should not happen if fileExists and no exception was thrown, but defensively:
@@ -219,6 +239,13 @@ Expectation for parameters:
       };
     }
 
+    const newContent = this._applyReplacement(
+      currentContent,
+      finalOldString,
+      finalNewString,
+      isNewFile,
+    );
+
     return {
       currentContent,
       newContent,
@@ -247,7 +274,10 @@ Expectation for parameters:
     }
     let currentContent: string | null = null;
     let fileExists = false;
-    let newContent = '';
+    let finalNewString = params.new_string;
+    let finalOldString = params.old_string;
+    let occurrences = 0;
+
     try {
       currentContent = fs.readFileSync(params.file_path, 'utf8');
       fileExists = true;
@@ -259,28 +289,36 @@ Expectation for parameters:
         return false;
       }
     }
+
     if (params.old_string === '' && !fileExists) {
-      newContent = params.new_string;
+      // Creating new file, newContent is just params.new_string
     } else if (!fileExists) {
-      return false;
+      return false; // Cannot edit non-existent file if old_string is not empty
     } else if (currentContent !== null) {
-      // Use the correctEdit utility to potentially correct params and get occurrences
+      const correctedEdit = await ensureCorrectEdit(
-      const { params: correctedParams, occurrences: correctedOccurrences } =
+        currentContent,
-        await ensureCorrectEdit(currentContent, params, this.client);
+        params,
+        this.client,
+      );
+      finalOldString = correctedEdit.params.old_string;
+      finalNewString = correctedEdit.params.new_string;
+      occurrences = correctedEdit.occurrences;
 
-      params.old_string = correctedParams.old_string;
+      if (occurrences === 0 || occurrences !== 1) {
-      params.new_string = correctedParams.new_string;
-
-      if (correctedOccurrences === 0 || correctedOccurrences !== 1) {
         return false;
       }
-      newContent = currentContent.replaceAll(
-        params.old_string,
-        params.new_string,
-      );
     } else {
-      return false;
+      return false; // Should not happen
     }
+
+    const isNewFileScenario = params.old_string === '' && !fileExists;
+    const newContent = this._applyReplacement(
+      currentContent,
+      finalOldString,
+      finalNewString,
+      isNewFileScenario,
+    );
+
     const fileName = path.basename(params.file_path);
     const fileDiff = Diff.createPatch(
       fileName,