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

2 files changed, 63 insertions, 0 deletions

diff --git a/packages/server/src/utils/LruCache.ts b/packages/server/src/utils/LruCache.ts
new file mode 100644
index 00000000..076828c4
--- /dev/null
+++ b/packages/server/src/utils/LruCache.ts
@@ -0,0 +1,41 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export class LruCache<K, V> {
+  private cache: Map<K, V>;
+  private maxSize: number;
+
+  constructor(maxSize: number) {
+    this.cache = new Map<K, V>();
+    this.maxSize = maxSize;
+  }
+
+  get(key: K): V | undefined {
+    const value = this.cache.get(key);
+    if (value) {
+      // Move to end to mark as recently used
+      this.cache.delete(key);
+      this.cache.set(key, value);
+    }
+    return value;
+  }
+
+  set(key: K, value: V): void {
+    if (this.cache.has(key)) {
+      this.cache.delete(key);
+    } else if (this.cache.size >= this.maxSize) {
+      const firstKey = this.cache.keys().next().value;
+      if (firstKey !== undefined) {
+        this.cache.delete(firstKey);
+      }
+    }
+    this.cache.set(key, value);
+  }
+
+  clear(): void {
+    this.cache.clear();
+  }
+}
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();
+}