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

Co-authored-by: Colt McAnlis <colton@google.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Colt McAnlis 2025-07-07 10:28:56 -07:00 committed by GitHub
parent 229ae03631
commit 8f4046c71a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 249 additions and 21 deletions

View File

@ -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);

View File

@ -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,

View File

@ -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,

View File

@ -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

View File

@ -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', () => {

View File

@ -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 = {
@ -49,6 +58,96 @@ export interface CorrectedEditResult {
occurrences: number;
}
/**
* 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.
@ -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,

View File

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