jcarr
/
gemini-cli
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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
This commit is contained in:
Taylor Mullen 2025-05-25 14:16:08 -07:00 committed by N. Taylor Mullen
parent 48781272ee
commit c181fc1cf3
2 changed files with 189 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

137
packages/server/src/tools/edit.test.ts
View File

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

92
packages/server/src/tools/edit.ts
View File

@ -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 { params: correctedParams, occurrences: correctedOccurrences } =
await ensureCorrectEdit(currentContent, params, this.client);
const correctedEdit = await ensureCorrectEdit(
currentContent,
params,
this.client,
);
finalOldString = correctedEdit.params.old_string;
finalNewString = correctedEdit.params.new_string;
occurrences = correctedEdit.occurrences;
params.old_string = correctedParams.old_string;
params.new_string = correctedParams.new_string;
if (correctedOccurrences === 0 || correctedOccurrences !== 1) {
if (occurrences === 0 || occurrences !== 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,