From f7a2442faca5e3c51bab753672130968456a4c23 Mon Sep 17 00:00:00 2001
From: Bryan Morgan <bryanmorgan@google.com>
Date: Sun, 1 Jun 2025 17:49:48 -0400
Subject: Added replace tool ability to replace more than 1 occurrence (#669)

Co-authored-by: N. Taylor Mullen <ntaylormullen@google.com>
---
 packages/core/src/utils/editCorrector.ts | 35 ++++++++++++++++++++++++++++----
 1 file changed, 31 insertions(+), 4 deletions(-)

(limited to 'packages/core/src/utils')

diff --git a/packages/core/src/utils/editCorrector.ts b/packages/core/src/utils/editCorrector.ts
index 78663954..5e56bac8 100644
--- a/packages/core/src/utils/editCorrector.ts
+++ b/packages/core/src/utils/editCorrector.ts
@@ -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,
-- 
cgit v1.2.3