fix: EditTool can clobber human edits to the same file. (#3043)

Co-authored-by: Colt McAnlis <[email protected]>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

7 files changed, 249 insertions, 21 deletions

diff --git a/packages/core/src/tools/edit.test.ts b/packages/core/src/tools/edit.test.ts
index 1341fdd0..ab42450a 100644
--- a/packages/core/src/tools/edit.test.ts
+++ b/packages/core/src/tools/edit.test.ts
@@ -84,20 +84,22 @@ describe('EditTool', () => {
 
     // Reset mocks and set default implementation for ensureCorrectEdit
     mockEnsureCorrectEdit.mockReset();
-    mockEnsureCorrectEdit.mockImplementation(async (currentContent, params) => {
-      let occurrences = 0;
-      if (params.old_string && currentContent) {
-        // Simple string counting for the mock
-        let index = currentContent.indexOf(params.old_string);
-        while (index !== -1) {
-          occurrences++;
-          index = currentContent.indexOf(params.old_string, index + 1);
+    mockEnsureCorrectEdit.mockImplementation(
+      async (_, currentContent, params) => {
+        let occurrences = 0;
+        if (params.old_string && currentContent) {
+          // Simple string counting for the mock
+          let index = currentContent.indexOf(params.old_string);
+          while (index !== -1) {
+            occurrences++;
+            index = currentContent.indexOf(params.old_string, index + 1);
+          }
+        } else if (params.old_string === '') {
+          occurrences = 0; // Creating a new file
         }
-      } else if (params.old_string === '') {
-        occurrences = 0; // Creating a new file
-      }
-      return Promise.resolve({ params, occurrences });
-    });
+        return Promise.resolve({ params, occurrences });
+      },
+    );
 
     // Default mock for generateJson to return the snippet unchanged
     mockGenerateJson.mockReset();
@@ -333,7 +335,7 @@ describe('EditTool', () => {
       // Set a specific mock for this test case
       let mockCalled = false;
       mockEnsureCorrectEdit.mockImplementationOnce(
-        async (content, p, client) => {
+        async (_, content, p, client) => {
           mockCalled = true;
           expect(content).toBe(originalContent);
           expect(p).toBe(params);
@@ -383,7 +385,7 @@ describe('EditTool', () => {
     beforeEach(() => {
       filePath = path.join(rootDir, testFile);
       // Default for execute tests, can be overridden
-      mockEnsureCorrectEdit.mockImplementation(async (content, params) => {
+      mockEnsureCorrectEdit.mockImplementation(async (_, content, params) => {
         let occurrences = 0;
         if (params.old_string && content) {
           let index = content.indexOf(params.old_string);
diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts
index 1e04ab35..2df01a22 100644
--- a/packages/core/src/tools/edit.ts
+++ b/packages/core/src/tools/edit.ts
@@ -236,6 +236,7 @@ Expectation for required parameters:
     } else if (currentContent !== null) {
       // Editing an existing file
       const correctedEdit = await ensureCorrectEdit(
+        params.file_path,
         currentContent,
         params,
         this.client,
diff --git a/packages/core/src/tools/write-file.test.ts b/packages/core/src/tools/write-file.test.ts
index 3f9825c6..c33b5fa2 100644
--- a/packages/core/src/tools/write-file.test.ts
+++ b/packages/core/src/tools/write-file.test.ts
@@ -114,6 +114,7 @@ describe('WriteFileTool', () => {
     // Default mock implementations that return valid structures
     mockEnsureCorrectEdit.mockImplementation(
       async (
+        filePath: string,
         _currentContent: string,
         params: EditToolParams,
         _client: GeminiClient,
@@ -248,6 +249,7 @@ describe('WriteFileTool', () => {
       );
 
       expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
+        filePath,
         originalContent,
         {
           old_string: originalContent,
@@ -388,6 +390,7 @@ describe('WriteFileTool', () => {
       )) as ToolEditConfirmationDetails;
 
       expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
+        filePath,
         originalContent,
         {
           old_string: originalContent,
@@ -523,6 +526,7 @@ describe('WriteFileTool', () => {
       const result = await tool.execute(params, abortSignal);
 
       expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
+        filePath,
         initialContent,
         {
           old_string: initialContent,
diff --git a/packages/core/src/tools/write-file.ts b/packages/core/src/tools/write-file.ts
index c343cab8..37a1ba78 100644
--- a/packages/core/src/tools/write-file.ts
+++ b/packages/core/src/tools/write-file.ts
@@ -367,6 +367,7 @@ export class WriteFileTool
     if (fileExists) {
       // This implies originalContent is available
       const { params: correctedParams } = await ensureCorrectEdit(
+        filePath,
         originalContent,
         {
           old_string: originalContent, // Treat entire current content as old_string
diff --git a/packages/core/src/utils/editCorrector.test.ts b/packages/core/src/utils/editCorrector.test.ts
index cb305e28..bcf75dfe 100644
--- a/packages/core/src/utils/editCorrector.test.ts
+++ b/packages/core/src/utils/editCorrector.test.ts
@@ -5,7 +5,17 @@
  */
 
 /* eslint-disable @typescript-eslint/no-explicit-any */
-import { vi, describe, it, expect, beforeEach, type Mocked } from 'vitest';
+import {
+  vi,
+  describe,
+  it,
+  expect,
+  beforeEach,
+  Mock,
+  type Mocked,
+} from 'vitest';
+import * as fs from 'fs';
+import { EditTool } from '../tools/edit.js';
 
 // MOCKS
 let callCount = 0;
@@ -15,6 +25,10 @@ let mockGenerateJson: any;
 let mockStartChat: any;
 let mockSendMessageStream: any;
 
+vi.mock('fs', () => ({
+  statSync: vi.fn(),
+}));
+
 vi.mock('../core/client.js', () => ({
   GeminiClient: vi.fn().mockImplementation(function (
     this: any,
@@ -222,6 +236,7 @@ describe('editCorrector', () => {
       mockGeminiClientInstance = new GeminiClient(
         mockConfigInstance,
       ) as Mocked<GeminiClient>;
+      mockGeminiClientInstance.getHistory = vi.fn().mockResolvedValue([]);
       resetEditCorrectorCaches_TEST_ONLY();
     });
 
@@ -237,6 +252,7 @@ describe('editCorrector', () => {
           corrected_new_string_escaping: 'replace with "this"',
         });
         const result = await ensureCorrectEdit(
+          '/test/file.txt',
           currentContent,
           originalParams,
           mockGeminiClientInstance,
@@ -255,6 +271,7 @@ describe('editCorrector', () => {
           new_string: 'replace with this',
         };
         const result = await ensureCorrectEdit(
+          '/test/file.txt',
           currentContent,
           originalParams,
           mockGeminiClientInstance,
@@ -276,6 +293,7 @@ describe('editCorrector', () => {
           corrected_new_string_escaping: 'replace with "this"',
         });
         const result = await ensureCorrectEdit(
+          '/test/file.txt',
           currentContent,
           originalParams,
           mockGeminiClientInstance,
@@ -294,6 +312,7 @@ describe('editCorrector', () => {
           new_string: 'replace with this',
         };
         const result = await ensureCorrectEdit(
+          '/test/file.txt',
           currentContent,
           originalParams,
           mockGeminiClientInstance,
@@ -316,6 +335,7 @@ describe('editCorrector', () => {
         };
         mockResponses.push({ corrected_new_string: 'replace with "this"' });
         const result = await ensureCorrectEdit(
+          '/test/file.txt',
           currentContent,
           originalParams,
           mockGeminiClientInstance,
@@ -334,6 +354,7 @@ describe('editCorrector', () => {
           new_string: 'replace with this',
         };
         const result = await ensureCorrectEdit(
+          '/test/file.txt',
           currentContent,
           originalParams,
           mockGeminiClientInstance,
@@ -352,6 +373,7 @@ describe('editCorrector', () => {
           new_string: 'replace with foobar',
         };
         const result = await ensureCorrectEdit(
+          '/test/file.txt',
           currentContent,
           originalParams,
           mockGeminiClientInstance,
@@ -375,6 +397,7 @@ describe('editCorrector', () => {
         const llmNewString = 'LLM says replace with "that"';
         mockResponses.push({ corrected_new_string_escaping: llmNewString });
         const result = await ensureCorrectEdit(
+          '/test/file.txt',
           currentContent,
           originalParams,
           mockGeminiClientInstance,
@@ -397,6 +420,7 @@ describe('editCorrector', () => {
         mockResponses.push({ corrected_target_snippet: llmCorrectedOldString });
         mockResponses.push({ corrected_new_string: llmNewString });
         const result = await ensureCorrectEdit(
+          '/test/file.txt',
           currentContent,
           originalParams,
           mockGeminiClientInstance,
@@ -417,6 +441,7 @@ describe('editCorrector', () => {
         const llmCorrectedOldString = 'to be corrected';
         mockResponses.push({ corrected_target_snippet: llmCorrectedOldString });
         const result = await ensureCorrectEdit(
+          '/test/file.txt',
           currentContent,
           originalParams,
           mockGeminiClientInstance,
@@ -439,6 +464,7 @@ describe('editCorrector', () => {
           corrected_new_string_escaping: newStringForLLMAndReturnedByLLM,
         });
         const result = await ensureCorrectEdit(
+          '/test/file.txt',
           currentContent,
           originalParams,
           mockGeminiClientInstance,
@@ -460,6 +486,7 @@ describe('editCorrector', () => {
         };
         mockResponses.push({ corrected_target_snippet: 'still nonexistent' });
         const result = await ensureCorrectEdit(
+          '/test/file.txt',
           currentContent,
           originalParams,
           mockGeminiClientInstance,
@@ -478,6 +505,7 @@ describe('editCorrector', () => {
           new_string: 'some new string',
         };
         const result = await ensureCorrectEdit(
+          '/test/file.txt',
           currentContent,
           originalParams,
           mockGeminiClientInstance,
@@ -491,16 +519,17 @@ describe('editCorrector', () => {
 
     describe('Scenario Group 5: Specific unescapeStringForGeminiBug checks (integrated into ensureCorrectEdit)', () => {
       it('Test 5.1: old_string needs LLM to become currentContent, new_string also needs correction', async () => {
-        const currentContent = 'const x = "a\\nbc\\\\"def\\\\"';
+        const currentContent = 'const x = "a\nbc\\"def\\"';
         const originalParams = {
           file_path: '/test/file.txt',
-          old_string: 'const x = \\\\"a\\\\nbc\\\\\\\\"def\\\\\\\\"',
-          new_string: 'const y = \\\\"new\\\\nval\\\\\\\\"content\\\\\\\\"',
+          old_string: 'const x = \\"a\\nbc\\\\"def\\\\"',
+          new_string: 'const y = \\"new\\nval\\\\"content\\\\"',
         };
-        const expectedFinalNewString = 'const y = "new\\nval\\\\"content\\\\"';
+        const expectedFinalNewString = 'const y = "new\nval\\"content\\"';
         mockResponses.push({ corrected_target_snippet: currentContent });
         mockResponses.push({ corrected_new_string: expectedFinalNewString });
         const result = await ensureCorrectEdit(
+          '/test/file.txt',
           currentContent,
           originalParams,
           mockGeminiClientInstance,
@@ -512,6 +541,61 @@ describe('editCorrector', () => {
         expect(result.occurrences).toBe(1);
       });
     });
+
+    describe('Scenario Group 6: Concurrent Edits', () => {
+      it('Test 6.1: should return early if file was modified by another process', async () => {
+        const filePath = '/test/file.txt';
+        const currentContent =
+          'This content has been modified by someone else.';
+        const originalParams = {
+          file_path: filePath,
+          old_string: 'nonexistent string',
+          new_string: 'some new string',
+        };
+
+        const now = Date.now();
+        const lastEditTime = now - 5000; // 5 seconds ago
+
+        // Mock the file's modification time to be recent
+        vi.spyOn(fs, 'statSync').mockReturnValue({
+          mtimeMs: now,
+        } as fs.Stats);
+
+        // Mock the last edit timestamp from our history to be in the past
+        const history = [
+          {
+            role: 'model',
+            parts: [
+              {
+                functionResponse: {
+                  name: EditTool.Name,
+                  id: `${EditTool.Name}-${lastEditTime}-123`,
+                  response: {
+                    output: {
+                      llmContent: `Successfully modified file: ${filePath}`,
+                    },
+                  },
+                },
+              },
+            ],
+          },
+        ];
+        (mockGeminiClientInstance.getHistory as Mock).mockResolvedValue(
+          history,
+        );
+
+        const result = await ensureCorrectEdit(
+          filePath,
+          currentContent,
+          originalParams,
+          mockGeminiClientInstance,
+          abortSignal,
+        );
+
+        expect(result.occurrences).toBe(0);
+        expect(result.params).toEqual(originalParams);
+      });
+    });
   });
 
   describe('ensureCorrectFileContent', () => {
diff --git a/packages/core/src/utils/editCorrector.ts b/packages/core/src/utils/editCorrector.ts
index a12c7364..bff1c2f2 100644
--- a/packages/core/src/utils/editCorrector.ts
+++ b/packages/core/src/utils/editCorrector.ts
@@ -11,9 +11,18 @@ import {
   Type,
 } from '@google/genai';
 import { GeminiClient } from '../core/client.js';
-import { EditToolParams } from '../tools/edit.js';
+import { EditToolParams, EditTool } from '../tools/edit.js';
+import { WriteFileTool } from '../tools/write-file.js';
+import { ReadFileTool } from '../tools/read-file.js';
+import { ReadManyFilesTool } from '../tools/read-many-files.js';
+import { GrepTool } from '../tools/grep.js';
 import { LruCache } from './LruCache.js';
 import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
+import {
+  isFunctionResponse,
+  isFunctionCall,
+} from '../utils/messageInspectors.js';
+import * as fs from 'fs';
 
 const EditModel = DEFAULT_GEMINI_FLASH_MODEL;
 const EditConfig: GenerateContentConfig = {
@@ -50,6 +59,96 @@ export interface CorrectedEditResult {
 }
 
 /**
+ * Extracts the timestamp from the .id value, which is in format
+ * <tool.name>-<timestamp>-<uuid>
+ * @param fcnId the ID value of a functionCall or functionResponse object
+ * @returns -1 if the timestamp could not be extracted, else the timestamp (as a number)
+ */
+function getTimestampFromFunctionId(fcnId: string): number {
+  const idParts = fcnId.split('-');
+  if (idParts.length > 2) {
+    const timestamp = parseInt(idParts[1], 10);
+    if (!isNaN(timestamp)) {
+      return timestamp;
+    }
+  }
+  return -1;
+}
+
+/**
+ * Will look through the gemini client history and determine when the most recent
+ * edit to a target file occured. If no edit happened, it will return -1
+ * @param filePath the path to the file
+ * @param client the geminiClient, so that we can get the history
+ * @returns a DateTime (as a number) of when the last edit occured, or -1 if no edit was found.
+ */
+async function findLastEditTimestamp(
+  filePath: string,
+  client: GeminiClient,
+): Promise<number> {
+  const history = (await client.getHistory()) ?? [];
+
+  // Tools that may reference the file path in their FunctionResponse `output`.
+  const toolsInResp = new Set([
+    WriteFileTool.Name,
+    EditTool.Name,
+    ReadManyFilesTool.Name,
+    GrepTool.Name,
+  ]);
+  // Tools that may reference the file path in their FunctionCall `args`.
+  const toolsInCall = new Set([...toolsInResp, ReadFileTool.Name]);
+
+  // Iterate backwards to find the most recent relevant action.
+  for (const entry of history.slice().reverse()) {
+    if (!entry.parts) continue;
+
+    for (const part of entry.parts) {
+      let id: string | undefined;
+      let content: unknown;
+
+      // Check for a relevant FunctionCall with the file path in its arguments.
+      if (
+        isFunctionCall(entry) &&
+        part.functionCall?.name &&
+        toolsInCall.has(part.functionCall.name)
+      ) {
+        id = part.functionCall.id;
+        content = part.functionCall.args;
+      }
+      // Check for a relevant FunctionResponse with the file path in its output.
+      else if (
+        isFunctionResponse(entry) &&
+        part.functionResponse?.name &&
+        toolsInResp.has(part.functionResponse.name)
+      ) {
+        const { response } = part.functionResponse;
+        if (response && !('error' in response) && 'output' in response) {
+          id = part.functionResponse.id;
+          content = response.output;
+        }
+      }
+
+      if (!id || content === undefined) continue;
+
+      // Use the "blunt hammer" approach to find the file path in the content.
+      // Note that the tool response data is inconsistent in their formatting
+      // with successes and errors - so, we just check for the existance
+      // as the best guess to if error/failed occured with the response.
+      const stringified = JSON.stringify(content);
+      if (
+        !stringified.includes('Error') && // only applicable for functionResponse
+        !stringified.includes('Failed') && // only applicable for functionResponse
+        stringified.includes(filePath)
+      ) {
+        return getTimestampFromFunctionId(id);
+      }
+    }
+  }
+
+  return -1;
+}
+
+/**
  * Attempts to correct edit parameters if the original old_string is not found.
  * It tries unescaping, and then LLM-based correction.
  * Results are cached to avoid redundant processing.
@@ -61,6 +160,7 @@ export interface CorrectedEditResult {
  *          EditToolParams (as CorrectedEditParams) and the final occurrences count.
  */
 export async function ensureCorrectEdit(
+  filePath: string,
   currentContent: string,
   originalParams: EditToolParams, // This is the EditToolParams from edit.ts, without \'corrected\'
   client: GeminiClient,
@@ -140,6 +240,34 @@ export async function ensureCorrectEdit(
         );
       }
     } else if (occurrences === 0) {
+      if (filePath) {
+        // In order to keep from clobbering edits made outside our system,
+        // let's check if there was a more recent edit to the file than what
+        // our system has done
+        const lastEditedByUsTime = await findLastEditTimestamp(
+          filePath,
+          client,
+        );
+
+        // Add a 1-second buffer to account for timing inaccuracies. If the file
+        // was modified more than a second after the last edit tool was run, we
+        // can assume it was modified by something else.
+        if (lastEditedByUsTime > 0) {
+          const stats = fs.statSync(filePath);
+          const diff = stats.mtimeMs - lastEditedByUsTime;
+          if (diff > 2000) {
+            // Hard coded for 2 seconds
+            // This file was edited sooner
+            const result: CorrectedEditResult = {
+              params: { ...originalParams },
+              occurrences: 0, // Explicitly 0 as LLM failed
+            };
+            editCorrectionCache.set(cacheKey, result);
+            return result;
+          }
+        }
+      }
+
       const llmCorrectedOldString = await correctOldStringMismatch(
         client,
         currentContent,
diff --git a/packages/core/src/utils/messageInspectors.ts b/packages/core/src/utils/messageInspectors.ts
index b2c3cdce..5ef06d79 100644
--- a/packages/core/src/utils/messageInspectors.ts
+++ b/packages/core/src/utils/messageInspectors.ts
@@ -13,3 +13,11 @@ export function isFunctionResponse(content: Content): boolean {
     content.parts.every((part) => !!part.functionResponse)
   );
 }
+
+export function isFunctionCall(content: Content): boolean {
+  return (
+    content.role === 'model' &&
+    !!content.parts &&
+    content.parts.every((part) => !!part.functionCall)
+  );
+}