feat(core): add disableLLMCorrection setting to skip auto-correction in edit tools (#16000)

This commit is contained in:
Sandy Tao
2026-01-13 09:26:53 +08:00
committed by GitHub
parent 2fc61685a3
commit b81fe68325
13 changed files with 221 additions and 12 deletions
@@ -273,6 +273,7 @@ describe('editCorrector', () => {
mockGeminiClientInstance,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(mockGenerateJson).toHaveBeenCalledTimes(1);
expect(result.params.new_string).toBe('replace with "this"');
@@ -293,6 +294,7 @@ describe('editCorrector', () => {
mockGeminiClientInstance,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(mockGenerateJson).toHaveBeenCalledTimes(0);
expect(result.params.new_string).toBe('replace with this');
@@ -316,6 +318,7 @@ describe('editCorrector', () => {
mockGeminiClientInstance,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(mockGenerateJson).toHaveBeenCalledTimes(1);
expect(result.params.new_string).toBe('replace with "this"');
@@ -336,6 +339,7 @@ describe('editCorrector', () => {
mockGeminiClientInstance,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(mockGenerateJson).toHaveBeenCalledTimes(0);
expect(result.params.new_string).toBe('replace with this');
@@ -360,6 +364,7 @@ describe('editCorrector', () => {
mockGeminiClientInstance,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(mockGenerateJson).toHaveBeenCalledTimes(1);
expect(result.params.new_string).toBe('replace with "this"');
@@ -380,6 +385,7 @@ describe('editCorrector', () => {
mockGeminiClientInstance,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(mockGenerateJson).toHaveBeenCalledTimes(0);
expect(result.params.new_string).toBe('replace with this');
@@ -400,6 +406,7 @@ describe('editCorrector', () => {
mockGeminiClientInstance,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(mockGenerateJson).toHaveBeenCalledTimes(0);
expect(result.params.new_string).toBe('replace with foobar');
@@ -425,6 +432,7 @@ describe('editCorrector', () => {
mockGeminiClientInstance,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(mockGenerateJson).toHaveBeenCalledTimes(1);
expect(result.params.new_string).toBe(llmNewString);
@@ -449,6 +457,7 @@ describe('editCorrector', () => {
mockGeminiClientInstance,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(mockGenerateJson).toHaveBeenCalledTimes(2);
expect(result.params.new_string).toBe(llmNewString);
@@ -471,6 +480,7 @@ describe('editCorrector', () => {
mockGeminiClientInstance,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(mockGenerateJson).toHaveBeenCalledTimes(1);
expect(result.params.new_string).toBe('replace with "this"');
@@ -495,6 +505,7 @@ describe('editCorrector', () => {
mockGeminiClientInstance,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(mockGenerateJson).toHaveBeenCalledTimes(1);
expect(result.params.new_string).toBe(newStringForLLMAndReturnedByLLM);
@@ -518,6 +529,7 @@ describe('editCorrector', () => {
mockGeminiClientInstance,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(mockGenerateJson).toHaveBeenCalledTimes(1);
expect(result.params).toEqual(originalParams);
@@ -538,6 +550,7 @@ describe('editCorrector', () => {
mockGeminiClientInstance,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(mockGenerateJson).toHaveBeenCalledTimes(0);
expect(result.params).toEqual(originalParams);
@@ -563,6 +576,7 @@ describe('editCorrector', () => {
mockGeminiClientInstance,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(mockGenerateJson).toHaveBeenCalledTimes(2);
expect(result.params.old_string).toBe(currentContent);
@@ -618,6 +632,7 @@ describe('editCorrector', () => {
mockGeminiClientInstance,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(result.occurrences).toBe(0);
@@ -665,6 +680,7 @@ describe('editCorrector', () => {
content,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(result).toBe(content);
expect(mockGenerateJson).toHaveBeenCalledTimes(0);
@@ -681,6 +697,7 @@ describe('editCorrector', () => {
content,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(result).toBe(correctedContent);
@@ -701,6 +718,7 @@ describe('editCorrector', () => {
content,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(result).toBe(correctedContent);
@@ -716,6 +734,7 @@ describe('editCorrector', () => {
content,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(result).toBe(content);
@@ -736,6 +755,7 @@ describe('editCorrector', () => {
content,
mockBaseLlmClientInstance,
abortSignal,
false,
);
expect(result).toBe(correctedContent);
+23 -3
View File
@@ -167,10 +167,11 @@ async function findLastEditTimestamp(
export async function ensureCorrectEdit(
filePath: string,
currentContent: string,
originalParams: EditToolParams, // This is the EditToolParams from edit.ts, without \'corrected\'
originalParams: EditToolParams, // This is the EditToolParams from edit.ts, without 'corrected'
geminiClient: GeminiClient,
baseLlmClient: BaseLlmClient,
abortSignal: AbortSignal,
disableLLMCorrection: boolean,
): Promise<CorrectedEditResult> {
const cacheKey = `${currentContent}---${originalParams.old_string}---${originalParams.new_string}`;
const cachedResult = editCorrectionCache.get(cacheKey);
@@ -189,7 +190,7 @@ export async function ensureCorrectEdit(
let occurrences = countOccurrences(currentContent, finalOldString);
if (occurrences === expectedReplacements) {
if (newStringPotentiallyEscaped) {
if (newStringPotentiallyEscaped && !disableLLMCorrection) {
finalNewString = await correctNewStringEscaping(
baseLlmClient,
finalOldString,
@@ -236,7 +237,7 @@ export async function ensureCorrectEdit(
if (occurrences === expectedReplacements) {
finalOldString = unescapedOldStringAttempt;
if (newStringPotentiallyEscaped) {
if (newStringPotentiallyEscaped && !disableLLMCorrection) {
finalNewString = await correctNewString(
baseLlmClient,
originalParams.old_string, // original old
@@ -274,6 +275,15 @@ export async function ensureCorrectEdit(
}
}
if (disableLLMCorrection) {
const result: CorrectedEditResult = {
params: { ...originalParams },
occurrences: 0,
};
editCorrectionCache.set(cacheKey, result);
return result;
}
const llmCorrectedOldString = await correctOldStringMismatch(
baseLlmClient,
currentContent,
@@ -347,6 +357,7 @@ export async function ensureCorrectFileContent(
content: string,
baseLlmClient: BaseLlmClient,
abortSignal: AbortSignal,
disableLLMCorrection: boolean = false,
): Promise<string> {
const cachedResult = fileContentCorrectionCache.get(content);
if (cachedResult) {
@@ -360,6 +371,15 @@ export async function ensureCorrectFileContent(
return content;
}
if (disableLLMCorrection) {
// If we can't use LLM, we should at least use the unescaped content
// as it's likely better than the original if it was detected as potentially escaped.
// unescapeStringForGeminiBug is a heuristic, not an LLM call.
const unescaped = unescapeStringForGeminiBug(content);
fileContentCorrectionCache.set(content, unescaped);
return unescaped;
}
const correctedContent = await correctStringEscaping(
content,
baseLlmClient,