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

packages/core/src/config/config.ts

@@ -358,6 +358,7 @@ export interface ConfigParameters {
   skillsSupport?: boolean;
   disabledSkills?: string[];
   experimentalJitContext?: boolean;
+  disableLLMCorrection?: boolean;
   onModelChange?: (model: string) => void;
   mcpEnabled?: boolean;
   extensionsEnabled?: boolean;
@@ -497,6 +498,7 @@ export class Config {
   private disabledSkills: string[];
 
   private readonly experimentalJitContext: boolean;
+  private readonly disableLLMCorrection: boolean;
   private contextManager?: ContextManager;
   private terminalBackground: string | undefined = undefined;
   private remoteAdminSettings: GeminiCodeAssistSetting | undefined;
@@ -566,6 +568,7 @@ export class Config {
     this.model = params.model;
     this._activeModel = params.model;
     this.enableAgents = params.enableAgents ?? false;
+    this.disableLLMCorrection = params.disableLLMCorrection ?? false;
     this.skillsSupport = params.skillsSupport ?? false;
     this.disabledSkills = params.disabledSkills ?? [];
     this.modelAvailabilityService = new ModelAvailabilityService();
@@ -1426,6 +1429,10 @@ export class Config {
     return this.enableExtensionReloading;
   }
 
+  getDisableLLMCorrection(): boolean {
+    return this.disableLLMCorrection;
+  }
+
   isAgentsEnabled(): boolean {
     return this.enableAgents;
   }

packages/core/src/tools/confirmation-policy.test.ts

@@ -64,6 +64,7 @@ describe('Tool Confirmation Policy Updates', () => {
       getFileFilteringOptions: () => ({}),
       getGeminiClient: () => ({}),
       getBaseLlmClient: () => ({}),
+      getDisableLLMCorrection: () => false,
       getIdeMode: () => false,
       getWorkspaceContext: () => ({
         isPathWithinWorkspace: () => true,

packages/core/src/tools/edit.test.ts

@@ -120,6 +120,7 @@ describe('EditTool', () => {
       setGeminiMdFileCount: vi.fn(),
       getToolRegistry: () => ({}) as any,
       isInteractive: () => false,
+      getDisableLLMCorrection: vi.fn(() => false),
       getExperiments: () => {},
     } as unknown as Config;
 
@@ -858,4 +859,47 @@ describe('EditTool', () => {
       expect(totalActualRemoved).toBe(totalExpectedRemoved);
     });
   });
+
+  describe('disableLLMCorrection', () => {
+    it('should NOT call FixLLMEditWithInstruction when disableLLMCorrection is true', async () => {
+      const filePath = path.join(rootDir, 'disable_llm_test.txt');
+      fs.writeFileSync(filePath, 'Some content.', 'utf8');
+
+      // Enable the setting
+      (mockConfig.getDisableLLMCorrection as Mock).mockReturnValue(true);
+
+      const params: EditToolParams = {
+        file_path: filePath,
+        instruction: 'Replace non-existent text',
+        old_string: 'nonexistent',
+        new_string: 'replacement',
+      };
+
+      const invocation = tool.build(params);
+      const result = await invocation.execute(new AbortController().signal);
+
+      expect(result.error?.type).toBe(ToolErrorType.EDIT_NO_OCCURRENCE_FOUND);
+      expect(mockFixLLMEditWithInstruction).not.toHaveBeenCalled();
+    });
+
+    it('should call FixLLMEditWithInstruction when disableLLMCorrection is false (default)', async () => {
+      const filePath = path.join(rootDir, 'enable_llm_test.txt');
+      fs.writeFileSync(filePath, 'Some content.', 'utf8');
+
+      // Default is false, but being explicit
+      (mockConfig.getDisableLLMCorrection as Mock).mockReturnValue(false);
+
+      const params: EditToolParams = {
+        file_path: filePath,
+        instruction: 'Replace non-existent text',
+        old_string: 'nonexistent',
+        new_string: 'replacement',
+      };
+
+      const invocation = tool.build(params);
+      await invocation.execute(new AbortController().signal);
+
+      expect(mockFixLLMEditWithInstruction).toHaveBeenCalled();
+    });
+  });
 });

packages/core/src/tools/edit.ts

@@ -639,6 +639,17 @@ class EditToolInvocation
       };
     }
 
+    if (this.config.getDisableLLMCorrection()) {
+      return {
+        currentContent,
+        newContent: currentContent,
+        occurrences: replacementResult.occurrences,
+        isNewFile: false,
+        error: initialError,
+        originalLineEnding,
+      };
+    }
+
     // If there was an error, try to self-correct.
     return this.attemptSelfCorrection(
       params,

packages/core/src/tools/write-file.test.ts

@@ -106,6 +106,7 @@ const mockConfigInternal = {
       discoverTools: vi.fn(),
     }) as unknown as ToolRegistry,
   isInteractive: () => false,
+  getDisableLLMCorrection: vi.fn(() => false),
 };
 const mockConfig = mockConfigInternal as unknown as Config;
 
@@ -293,6 +294,7 @@ describe('WriteFileTool', () => {
         proposedContent,
         mockBaseLlmClientInstance,
         abortSignal,
+        false,
       );
       expect(mockEnsureCorrectEdit).not.toHaveBeenCalled();
       expect(result.correctedContent).toBe(correctedContent);
@@ -337,6 +339,7 @@ describe('WriteFileTool', () => {
         mockGeminiClientInstance,
         mockBaseLlmClientInstance,
         abortSignal,
+        false,
       );
       expect(mockEnsureCorrectFileContent).not.toHaveBeenCalled();
       expect(result.correctedContent).toBe(correctedProposedContent);
@@ -414,6 +417,7 @@ describe('WriteFileTool', () => {
         proposedContent,
         mockBaseLlmClientInstance,
         abortSignal,
+        false,
       );
       expect(confirmation).toEqual(
         expect.objectContaining({
@@ -464,6 +468,7 @@ describe('WriteFileTool', () => {
         mockGeminiClientInstance,
         mockBaseLlmClientInstance,
         abortSignal,
+        false,
       );
       expect(confirmation).toEqual(
         expect.objectContaining({
@@ -658,6 +663,7 @@ describe('WriteFileTool', () => {
         proposedContent,
         mockBaseLlmClientInstance,
         abortSignal,
+        false,
       );
       expect(result.llmContent).toMatch(
         /Successfully created and wrote to new file/,
@@ -715,6 +721,7 @@ describe('WriteFileTool', () => {
         mockGeminiClientInstance,
         mockBaseLlmClientInstance,
         abortSignal,
+        false,
       );
       expect(result.llmContent).toMatch(/Successfully overwrote file/);
       const writtenContent = await fsService.readTextFile(filePath);
@@ -899,4 +906,73 @@ describe('WriteFileTool', () => {
       },
     );
   });
+
+  describe('disableLLMCorrection', () => {
+    const abortSignal = new AbortController().signal;
+
+    it('should call ensureCorrectFileContent with disableLLMCorrection=true for a new file when disabled', async () => {
+      const filePath = path.join(rootDir, 'new_file_no_correction.txt');
+      const proposedContent = 'Proposed content.';
+
+      mockConfigInternal.getDisableLLMCorrection.mockReturnValue(true);
+      // Ensure the mock returns the content passed to it (simulating no change or unescaped change)
+      mockEnsureCorrectFileContent.mockResolvedValue(proposedContent);
+
+      const result = await getCorrectedFileContent(
+        mockConfig,
+        filePath,
+        proposedContent,
+        abortSignal,
+      );
+
+      expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith(
+        proposedContent,
+        mockBaseLlmClientInstance,
+        abortSignal,
+        true,
+      );
+      expect(mockEnsureCorrectEdit).not.toHaveBeenCalled();
+      expect(result.correctedContent).toBe(proposedContent);
+      expect(result.fileExists).toBe(false);
+    });
+
+    it('should call ensureCorrectEdit with disableLLMCorrection=true for an existing file when disabled', async () => {
+      const filePath = path.join(rootDir, 'existing_file_no_correction.txt');
+      const originalContent = 'Original content.';
+      const proposedContent = 'Proposed content.';
+      fs.writeFileSync(filePath, originalContent, 'utf8');
+
+      mockConfigInternal.getDisableLLMCorrection.mockReturnValue(true);
+      // Ensure the mock returns the content passed to it
+      mockEnsureCorrectEdit.mockResolvedValue({
+        params: {
+          file_path: filePath,
+          old_string: originalContent,
+          new_string: proposedContent,
+        },
+        occurrences: 1,
+      });
+
+      const result = await getCorrectedFileContent(
+        mockConfig,
+        filePath,
+        proposedContent,
+        abortSignal,
+      );
+
+      expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
+        filePath,
+        originalContent,
+        expect.anything(), // params object
+        mockGeminiClientInstance,
+        mockBaseLlmClientInstance,
+        abortSignal,
+        true,
+      );
+      expect(mockEnsureCorrectFileContent).not.toHaveBeenCalled();
+      expect(result.correctedContent).toBe(proposedContent);
+      expect(result.originalContent).toBe(originalContent);
+      expect(result.fileExists).toBe(true);
+    });
+  });
 });

packages/core/src/tools/write-file.ts

@@ -127,6 +127,7 @@ export async function getCorrectedFileContent(
       config.getGeminiClient(),
       config.getBaseLlmClient(),
       abortSignal,
+      config.getDisableLLMCorrection(),
     );
     correctedContent = correctedParams.new_string;
   } else {
@@ -135,6 +136,7 @@ export async function getCorrectedFileContent(
       proposedContent,
       config.getBaseLlmClient(),
       abortSignal,
+      config.getDisableLLMCorrection(),
     );
   }
   return { originalContent, correctedContent, fileExists };

packages/core/src/utils/editCorrector.test.ts

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

packages/core/src/utils/editCorrector.ts

@@ -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,