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:
parent
229ae03631
commit
8f4046c71a
|
@ -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);
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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', () => {
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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)
|
||||
);
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue