Added replace tool ability to replace more than 1 occurrence (#669)

Co-authored-by: N. Taylor Mullen <ntaylormullen@google.com>
This commit is contained in:
Bryan Morgan 2025-06-01 17:49:48 -04:00 committed by GitHub
parent f2a8d39f42
commit f7a2442fac
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 94 additions and 10 deletions

View File

@ -484,6 +484,49 @@ describe('EditTool', () => {
);
});
it('should successfully replace multiple occurrences when expected_replacements specified', async () => {
fs.writeFileSync(filePath, 'old text old text old text', 'utf8');
const params: EditToolParams = {
file_path: filePath,
old_string: 'old',
new_string: 'new',
expected_replacements: 3,
};
// Simulate confirmation by setting shouldAlwaysEdit
(tool as any).shouldAlwaysEdit = true;
const result = await tool.execute(params, new AbortController().signal);
(tool as any).shouldAlwaysEdit = false; // Reset for other tests
expect(result.llmContent).toMatch(/Successfully modified file/);
expect(fs.readFileSync(filePath, 'utf8')).toBe(
'new text new text new text',
);
const display = result.returnDisplay as FileDiff;
expect(display.fileDiff).toMatch(/old text old text old text/);
expect(display.fileDiff).toMatch(/new text new text new text/);
expect(display.fileName).toBe(testFile);
});
it('should return error if expected_replacements does not match actual occurrences', async () => {
fs.writeFileSync(filePath, 'old text old text', 'utf8');
const params: EditToolParams = {
file_path: filePath,
old_string: 'old',
new_string: 'new',
expected_replacements: 3, // Expecting 3 but only 2 exist
};
const result = await tool.execute(params, new AbortController().signal);
expect(result.llmContent).toMatch(
/Expected 3 occurrences but found 2 for old_string in file/,
);
expect(result.returnDisplay).toMatch(
/Failed to edit, expected 3 occurrence\(s\) but found 2/,
);
});
it('should return error if trying to create a file that already exists (empty old_string)', async () => {
fs.writeFileSync(filePath, 'Existing content', 'utf8');
const params: EditToolParams = {

View File

@ -42,6 +42,12 @@ export interface EditToolParams {
* The text to replace it with
*/
new_string: string;
/**
* Number of replacements expected. Defaults to 1 if not specified.
* Use when you want to replace multiple occurrences.
*/
expected_replacements?: number;
}
interface CalculatedEdit {
@ -69,14 +75,15 @@ export class EditTool extends BaseTool<EditToolParams, ToolResult> {
super(
EditTool.Name,
'Edit',
`Replaces a single, unique occurrence of text within a file. This tool requires providing significant context around the change to ensure uniqueness and precise targeting. Always use the ${ReadFileTool} tool to examine the file's current content before attempting a text replacement.
`Replaces text within a file. By default, replaces a single occurrence, but can replace multiple occurrences when \`expected_replacements\` is specified. This tool requires providing significant context around the change to ensure precise targeting. Always use the ${ReadFileTool} tool to examine the file's current content before attempting a text replacement.
Expectation for parameters:
Expectation for required parameters:
1. \`file_path\` MUST be an absolute path; otherwise an error will be thrown.
2. \`old_string\` MUST be the exact literal text to replace (including all whitespace, indentation, newlines, and surrounding code etc.).
3. \`new_string\` MUST be the exact literal text to replace \`old_string\` with (also including all whitespace, indentation, newlines, and surrounding code etc.). Ensure the resulting code is correct and idiomatic.
4. NEVER escape \`old_string\` or \`new_string\`, that would break the exact literal text requirement.
**Important:** If ANY of the above are not satisfied, the tool will fail. CRITICAL for \`old_string\`: Must uniquely identify the single instance to change. Include at least 3 lines of context BEFORE and AFTER the target text, matching whitespace and indentation precisely. If this string matches multiple locations, or does not match exactly, the tool will fail.`,
**Important:** If ANY of the above are not satisfied, the tool will fail. CRITICAL for \`old_string\`: Must uniquely identify the single instance to change. Include at least 3 lines of context BEFORE and AFTER the target text, matching whitespace and indentation precisely. If this string matches multiple locations, or does not match exactly, the tool will fail.,
**Multiple replacements:** Set \`expected_replacements\` to the number of occurrences you want to replace. The tool will replace ALL occurrences that match \`old_string\` exactly. Ensure the number of replacements matches your expectation.`,
{
properties: {
file_path: {
@ -86,7 +93,7 @@ Expectation for parameters:
},
old_string: {
description:
'The exact literal text to replace, preferably unescaped. CRITICAL: Must uniquely identify the single instance to change. Include at least 3 lines of context BEFORE and AFTER the target text, matching whitespace and indentation precisely. If this string is not the exact literal text (i.e. you escaped it), matches multiple locations, or does not match exactly, the tool will fail.',
'The exact literal text to replace, preferably unescaped. For single replacements (default), include at least 3 lines of context BEFORE and AFTER the target text, matching whitespace and indentation precisely. For multiple replacements, specify expected_replacements parameter. If this string is not the exact literal text (i.e. you escaped it) or does not match exactly, the tool will fail.',
type: 'string',
},
new_string: {
@ -94,6 +101,12 @@ Expectation for parameters:
'The exact literal text to replace `old_string` with, preferably unescaped. Provide the EXACT text. Ensure the resulting code is correct and idiomatic.',
type: 'string',
},
expected_replacements: {
type: 'number',
description:
'Number of replacements expected. Defaults to 1 if not specified. Use when you want to replace multiple occurrences.',
minimum: 1,
},
},
required: ['file_path', 'old_string', 'new_string'],
type: 'object',
@ -178,7 +191,7 @@ Expectation for parameters:
params: EditToolParams,
abortSignal: AbortSignal,
): Promise<CalculatedEdit> {
const expectedReplacements = 1;
const expectedReplacements = params.expected_replacements ?? 1;
let currentContent: string | null = null;
let fileExists = false;
let isNewFile = false;
@ -311,7 +324,8 @@ Expectation for parameters:
finalNewString = correctedEdit.params.new_string;
occurrences = correctedEdit.occurrences;
if (occurrences === 0 || occurrences !== 1) {
const expectedReplacements = params.expected_replacements ?? 1;
if (occurrences === 0 || occurrences !== expectedReplacements) {
return false;
}
} else {

View File

@ -79,7 +79,7 @@ export async function ensureCorrectEdit(
let finalOldString = originalParams.old_string;
let occurrences = countOccurrences(currentContent, finalOldString);
if (occurrences === 1) {
if (occurrences === (originalParams.expected_replacements ?? 1)) {
if (newStringPotentiallyEscaped) {
finalNewString = await correctNewStringEscaping(
client,
@ -89,6 +89,29 @@ export async function ensureCorrectEdit(
);
}
} else if (occurrences > 1) {
const expectedReplacements = originalParams.expected_replacements ?? 1;
// If user expects multiple replacements, return as-is
if (occurrences === expectedReplacements) {
const result: CorrectedEditResult = {
params: { ...originalParams },
occurrences,
};
editCorrectionCache.set(cacheKey, result);
return result;
}
// If user expects 1 but found multiple, try to correct (existing behavior)
if (expectedReplacements === 1) {
const result: CorrectedEditResult = {
params: { ...originalParams },
occurrences,
};
editCorrectionCache.set(cacheKey, result);
return result;
}
// If occurrences don't match expected, return as-is (will fail validation later)
const result: CorrectedEditResult = {
params: { ...originalParams },
occurrences,
@ -102,7 +125,7 @@ export async function ensureCorrectEdit(
);
occurrences = countOccurrences(currentContent, unescapedOldStringAttempt);
if (occurrences === 1) {
if (occurrences === (originalParams.expected_replacements ?? 1)) {
finalOldString = unescapedOldStringAttempt;
if (newStringPotentiallyEscaped) {
finalNewString = await correctNewString(
@ -125,7 +148,7 @@ export async function ensureCorrectEdit(
llmCorrectedOldString,
);
if (llmOldOccurrences === 1) {
if (llmOldOccurrences === (originalParams.expected_replacements ?? 1)) {
finalOldString = llmCorrectedOldString;
occurrences = llmOldOccurrences;
@ -165,6 +188,7 @@ export async function ensureCorrectEdit(
finalOldString,
finalNewString,
currentContent,
originalParams,
);
finalOldString = targetString;
finalNewString = pair;
@ -508,6 +532,7 @@ function trimPairIfPossible(
target: string,
trimIfTargetTrims: string,
currentContent: string,
originalParams: EditToolParams,
) {
const trimmedTargetString = target.trim();
if (target.length !== trimmedTargetString.length) {
@ -516,7 +541,9 @@ function trimPairIfPossible(
trimmedTargetString,
);
if (trimmedTargetOccurrences === 1) {
if (
trimmedTargetOccurrences === (originalParams.expected_replacements ?? 1)
) {
const trimmedReactiveString = trimIfTargetTrims.trim();
return {
targetString: trimmedTargetString,