Add expected_replacements to smart-edit tool (#12885)

packages/core/src/tools/smart-edit.test.ts

@@ -47,7 +47,6 @@ import {
   type EditToolParams,
   calculateReplacement,
 } from './smart-edit.js';
-import { applyReplacement } from './edit.js';
 import { type FileDiff, ToolConfirmationOutcome } from './tools.js';
 import { ToolErrorType } from './tool-error.js';
 import path from 'node:path';
@@ -174,35 +173,6 @@ describe('SmartEditTool', () => {
     fs.rmSync(tempDir, { recursive: true, force: true });
   });
 
-  describe('applyReplacement', () => {
-    it('should return newString if isNewFile is true', () => {
-      expect(applyReplacement(null, 'old', 'new', true)).toBe('new');
-      expect(applyReplacement('existing', 'old', 'new', true)).toBe('new');
-    });
-
-    it('should replace oldString with newString in currentContent', () => {
-      expect(applyReplacement('hello old world old', 'old', 'new', false)).toBe(
-        'hello new world new',
-      );
-    });
-
-    it('should treat $ literally and not as replacement pattern', () => {
-      const current = 'regex end is $ and more';
-      const oldStr = 'regex end is $';
-      const newStr = 'regex end is $ and correct';
-      const result = applyReplacement(current, oldStr, newStr, false);
-      expect(result).toBe('regex end is $ and correct and more');
-    });
-
-    it("should treat $' literally and not as a replacement pattern", () => {
-      const current = 'foo';
-      const oldStr = 'foo';
-      const newStr = "bar$'baz";
-      const result = applyReplacement(current, oldStr, newStr, false);
-      expect(result).toBe("bar$'baz");
-    });
-  });
-
   describe('calculateReplacement', () => {
     const abortSignal = new AbortController().signal;
 
@@ -582,6 +552,63 @@ describe('SmartEditTool', () => {
     });
   });
 
+  describe('expected_replacements', () => {
+    const testFile = 'replacements_test.txt';
+    let filePath: string;
+
+    beforeEach(() => {
+      filePath = path.join(rootDir, testFile);
+    });
+
+    it('should succeed when occurrences match expected_replacements', async () => {
+      fs.writeFileSync(filePath, 'foo foo foo', 'utf8');
+      const params: EditToolParams = {
+        file_path: filePath,
+        instruction: 'Replace all foo with bar',
+        old_string: 'foo',
+        new_string: 'bar',
+        expected_replacements: 3,
+      };
+      const invocation = tool.build(params);
+      const result = await invocation.execute(new AbortController().signal);
+      expect(result.error).toBeUndefined();
+      expect(fs.readFileSync(filePath, 'utf8')).toBe('bar bar bar');
+    });
+
+    it('should fail when occurrences do not match expected_replacements', async () => {
+      fs.writeFileSync(filePath, 'foo foo foo', 'utf8');
+      const params: EditToolParams = {
+        file_path: filePath,
+        instruction: 'Replace all foo with bar',
+        old_string: 'foo',
+        new_string: 'bar',
+        expected_replacements: 2,
+      };
+      const invocation = tool.build(params);
+      const result = await invocation.execute(new AbortController().signal);
+      expect(result.error?.type).toBe(
+        ToolErrorType.EDIT_EXPECTED_OCCURRENCE_MISMATCH,
+      );
+    });
+
+    it('should default to 1 expected replacement if not specified', async () => {
+      fs.writeFileSync(filePath, 'foo foo', 'utf8');
+      const params: EditToolParams = {
+        file_path: filePath,
+        instruction: 'Replace foo with bar',
+        old_string: 'foo',
+        new_string: 'bar',
+        // expected_replacements is undefined, defaults to 1
+      };
+      const invocation = tool.build(params);
+      const result = await invocation.execute(new AbortController().signal);
+      // Should fail because there are 2 occurrences but default expectation is 1
+      expect(result.error?.type).toBe(
+        ToolErrorType.EDIT_EXPECTED_OCCURRENCE_MISMATCH,
+      );
+    });
+  });
+
   describe('IDE mode', () => {
     const testFile = 'edit_me.txt';
     let filePath: string;

packages/core/src/tools/smart-edit.ts

@@ -348,6 +348,12 @@ export interface EditToolParams {
    */
   new_string: string;
 
+  /**
+   * Number of replacements expected. Defaults to 1 if not specified.
+   * Use when you want to replace multiple occurrences.
+   */
+  expected_replacements?: number;
+
   /**
    * The instruction for what needs to be done.
    */
@@ -466,7 +472,7 @@ class EditToolInvocation
     const secondError = getErrorReplaceResult(
       params,
       secondAttemptResult.occurrences,
-      1, // expectedReplacements is always 1 for smart_edit
+      params.expected_replacements ?? 1,
       secondAttemptResult.finalOldString,
       secondAttemptResult.finalNewString,
     );
@@ -509,7 +515,7 @@ class EditToolInvocation
     params: EditToolParams,
     abortSignal: AbortSignal,
   ): Promise<CalculatedEdit> {
-    const expectedReplacements = 1;
+    const expectedReplacements = params.expected_replacements ?? 1;
     let currentContent: string | null = null;
     let fileExists = false;
     let originalLineEnding: '\r\n' | '\n' = '\n'; // Default for new files
@@ -848,7 +854,7 @@ export class SmartEditTool
     super(
       SmartEditTool.Name,
       'Edit',
-      `Replaces text within a file. Replaces a single occurrence. This tool requires providing significant context around the change to ensure precise targeting. Always use the ${READ_FILE_TOOL_NAME} 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 ${READ_FILE_TOOL_NAME} tool to examine the file's current content before attempting a text replacement.
       
       The user has the ability to modify the \`new_string\` content. If modified, this will be stated in the response.
       
@@ -859,7 +865,7 @@ export class SmartEditTool
       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.
       5. Prefer to break down complex and long changes into multiple smaller atomic calls to this tool. Always check the content of the file after changes or not finding a string to match.
-      **Multiple replacements:** If there are multiple and ambiguous occurences of the \`old_string\` in the file, the tool will also 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.`,
       Kind.Edit,
       {
         properties: {
@@ -887,7 +893,7 @@ A good instruction should concisely answer:
           },
           old_string: {
             description:
-              'The exact literal text to replace, preferably unescaped. 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) 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. 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: {
@@ -895,6 +901,12 @@ A good instruction should concisely answer:
               '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', 'instruction', 'old_string', 'new_string'],
         type: 'object',