Ensure edit correction isn't re-done after confirm.

- Edit corretion leans on LLM-isms to ensure we properly fix poorly escaped content. Beacues of this we need to ensure that we don't re-run edit correction in many cases.
  - To ensure this an `LruCache` has been added to capture intermediate steps of edit correction to avoid re-computations.
  - Max cache size is 50 currently. This means a user can have a muti-confirmation flow of 25 items without recomputing anything (assuming they all break edit correction).
- Laid some groundwork for future testing.

Part of https://github.com/google-gemini/gemini-cli/issues/484

1 files changed, 22 insertions, 0 deletions

diff --git a/packages/server/src/utils/editCorrector.ts b/packages/server/src/utils/editCorrector.ts
index a17de6d4..91aa85a3 100644
--- a/packages/server/src/utils/editCorrector.ts
+++ b/packages/server/src/utils/editCorrector.ts
@@ -12,6 +12,7 @@ import {
 } from '@google/genai';
 import { GeminiClient } from '../core/client.js';
 import { EditToolParams } from '../tools/edit.js';
+import { LruCache } from './LruCache.js';
 
 const EditModel = 'gemini-2.5-flash-preview-04-17';
 const EditConfig: GenerateContentConfig = {
@@ -20,6 +21,13 @@ const EditConfig: GenerateContentConfig = {
   },
 };
 
+const MAX_CACHE_SIZE = 50;
+
+// Cache for ensureCorrectEdit results
+const editCorrectionCache = new LruCache<string, CorrectedEditResult>(
+  MAX_CACHE_SIZE,
+);
+
 /**
  * Defines the structure of the parameters within CorrectedEditResult
  */
@@ -53,6 +61,12 @@ export async function ensureCorrectEdit(
   originalParams: EditToolParams, // This is the EditToolParams from edit.ts, without \'corrected\'
   client: GeminiClient,
 ): Promise<CorrectedEditResult> {
+  const cacheKey = `${currentContent}---${originalParams.old_string}---${originalParams.new_string}`;
+  const cachedResult = editCorrectionCache.get(cacheKey);
+  if (cachedResult) {
+    return cachedResult;
+  }
+
   let finalNewString = originalParams.new_string;
   const newStringPotentiallyEscaped =
     unescapeStringForGeminiBug(originalParams.new_string) !==
@@ -74,6 +88,7 @@ export async function ensureCorrectEdit(
       params: { ...originalParams },
       occurrences,
     };
+    editCorrectionCache.set(cacheKey, result);
     return result;
   } else {
     // occurrences is 0 or some other unexpected state initially
@@ -124,6 +139,7 @@ export async function ensureCorrectEdit(
           params: { ...originalParams },
           occurrences: 0, // Explicitly 0 as LLM failed
         };
+        editCorrectionCache.set(cacheKey, result);
         return result;
       }
     } else {
@@ -132,6 +148,7 @@ export async function ensureCorrectEdit(
         params: { ...originalParams },
         occurrences, // This will be > 1
       };
+      editCorrectionCache.set(cacheKey, result);
       return result;
     }
   }
@@ -153,6 +170,7 @@ export async function ensureCorrectEdit(
     },
     occurrences: countOccurrences(currentContent, finalOldString), // Recalculate occurrences with the final old_string
   };
+  editCorrectionCache.set(cacheKey, result);
   return result;
 }
 
@@ -449,3 +467,7 @@ export function countOccurrences(str: string, substr: string): number {
   }
   return count;
 }
+
+export function resetEditCorrectorCaches_TEST_ONLY() {
+  editCorrectionCache.clear();
+}