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
This commit is contained in:
Taylor Mullen 2025-05-25 14:21:27 -07:00 committed by N. Taylor Mullen
parent c181fc1cf3
commit 1a5fe16b22
2 changed files with 63 additions and 0 deletions

View File

@ -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();
}
}

View File

@ -12,6 +12,7 @@ import {
} from '@google/genai'; } from '@google/genai';
import { GeminiClient } from '../core/client.js'; import { GeminiClient } from '../core/client.js';
import { EditToolParams } from '../tools/edit.js'; import { EditToolParams } from '../tools/edit.js';
import { LruCache } from './LruCache.js';
const EditModel = 'gemini-2.5-flash-preview-04-17'; const EditModel = 'gemini-2.5-flash-preview-04-17';
const EditConfig: GenerateContentConfig = { 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 * 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\' originalParams: EditToolParams, // This is the EditToolParams from edit.ts, without \'corrected\'
client: GeminiClient, client: GeminiClient,
): Promise<CorrectedEditResult> { ): 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; let finalNewString = originalParams.new_string;
const newStringPotentiallyEscaped = const newStringPotentiallyEscaped =
unescapeStringForGeminiBug(originalParams.new_string) !== unescapeStringForGeminiBug(originalParams.new_string) !==
@ -74,6 +88,7 @@ export async function ensureCorrectEdit(
params: { ...originalParams }, params: { ...originalParams },
occurrences, occurrences,
}; };
editCorrectionCache.set(cacheKey, result);
return result; return result;
} else { } else {
// occurrences is 0 or some other unexpected state initially // occurrences is 0 or some other unexpected state initially
@ -124,6 +139,7 @@ export async function ensureCorrectEdit(
params: { ...originalParams }, params: { ...originalParams },
occurrences: 0, // Explicitly 0 as LLM failed occurrences: 0, // Explicitly 0 as LLM failed
}; };
editCorrectionCache.set(cacheKey, result);
return result; return result;
} }
} else { } else {
@ -132,6 +148,7 @@ export async function ensureCorrectEdit(
params: { ...originalParams }, params: { ...originalParams },
occurrences, // This will be > 1 occurrences, // This will be > 1
}; };
editCorrectionCache.set(cacheKey, result);
return result; return result;
} }
} }
@ -153,6 +170,7 @@ export async function ensureCorrectEdit(
}, },
occurrences: countOccurrences(currentContent, finalOldString), // Recalculate occurrences with the final old_string occurrences: countOccurrences(currentContent, finalOldString), // Recalculate occurrences with the final old_string
}; };
editCorrectionCache.set(cacheKey, result);
return result; return result;
} }
@ -449,3 +467,7 @@ export function countOccurrences(str: string, substr: string): number {
} }
return count; return count;
} }
export function resetEditCorrectorCaches_TEST_ONLY() {
editCorrectionCache.clear();
}