fix(core-tools): bypass LLM correction for JSON and IPYNB files in write_file and replace (#28223)

Co-authored-by: David Pierce <davidapierce@google.com>

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

@@ -1331,6 +1331,44 @@ function doIt() {
 
       expect(mockFixLLMEditWithInstruction).toHaveBeenCalled();
     });
+
+    it('should NOT call FixLLMEditWithInstruction for .json files even when disableLLMCorrection is false', async () => {
+      const filePath = path.join(rootDir, 'test.json');
+      fs.writeFileSync(filePath, '{"key": "value"}', 'utf8');
+
+      (mockConfig.getDisableLLMCorrection as Mock).mockReturnValue(false);
+
+      const params: EditToolParams = {
+        file_path: filePath,
+        instruction: 'Replace value',
+        old_string: 'nonexistent',
+        new_string: 'replacement',
+      };
+
+      const invocation = tool.build(params);
+      await invocation.execute({ abortSignal: new AbortController().signal });
+
+      expect(mockFixLLMEditWithInstruction).not.toHaveBeenCalled();
+    });
+
+    it('should NOT call FixLLMEditWithInstruction for .ipynb files even when disableLLMCorrection is false', async () => {
+      const filePath = path.join(rootDir, 'notebook.ipynb');
+      fs.writeFileSync(filePath, '{"cells": []}', 'utf8');
+
+      (mockConfig.getDisableLLMCorrection as Mock).mockReturnValue(false);
+
+      const params: EditToolParams = {
+        file_path: filePath,
+        instruction: 'Replace cell',
+        old_string: 'nonexistent',
+        new_string: 'replacement',
+      };
+
+      const invocation = tool.build(params);
+      await invocation.execute({ abortSignal: new AbortController().signal });
+
+      expect(mockFixLLMEditWithInstruction).not.toHaveBeenCalled();
+    });
   });
 
   describe('JIT context discovery', () => {

packages/core/src/tools/edit.ts

@@ -767,7 +767,12 @@ class EditToolInvocation
       };
     }
 
-    if (this.config.getDisableLLMCorrection()) {
+    const fileExt = path.extname(this.resolvedPath).toLowerCase();
+    const isJsonOrIpynb = ['.json', '.ipynb', '.jsonc', '.json5'].includes(
+      fileExt,
+    );
+
+    if (this.config.getDisableLLMCorrection() || isJsonOrIpynb) {
       return {
         currentContent,
         newContent: currentContent,

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

@@ -428,6 +428,39 @@ describe('WriteFileTool', () => {
       expect(result.error).toBeUndefined();
     });
 
+    it('should not call ensureCorrectFileContent for .json files', async () => {
+      const filePath = path.join(rootDir, 'config.json');
+      const proposedContent = '{"key": "value\\nwith\\nescapes"}';
+      const abortSignal = new AbortController().signal;
+
+      const result = await getCorrectedFileContent(
+        mockConfig,
+        filePath,
+        proposedContent,
+        abortSignal,
+      );
+
+      expect(mockEnsureCorrectFileContent).not.toHaveBeenCalled();
+      expect(result.correctedContent).toBe(proposedContent);
+    });
+
+    it('should not call ensureCorrectFileContent for .ipynb files', async () => {
+      const filePath = path.join(rootDir, 'notebook.ipynb');
+      const proposedContent =
+        '{"cells": [{"source": ["print(\\"hello\\\\n\\")"]}]}';
+      const abortSignal = new AbortController().signal;
+
+      const result = await getCorrectedFileContent(
+        mockConfig,
+        filePath,
+        proposedContent,
+        abortSignal,
+      );
+
+      expect(mockEnsureCorrectFileContent).not.toHaveBeenCalled();
+      expect(result.correctedContent).toBe(proposedContent);
+    });
+
     it('should return error if reading an existing file fails (e.g. permissions)', async () => {
       const filePath = path.join(rootDir, 'unreadable_file.txt');
       const proposedContent = 'some content';

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

@@ -198,22 +198,29 @@ export async function getCorrectedFileContent(
     }
   }
 
-  const activeModel = config.getActiveModel();
-  const resolvedModel = resolveModel(activeModel, false, false, true, config);
-
-  const aggressiveUnescape =
-    !isGemini3Model(resolvedModel, config) &&
-    !isGemini2Model(resolvedModel) &&
-    !isCustomModel(resolvedModel, config);
-
-  correctedContent = await ensureCorrectFileContent(
-    proposedContent,
-    config.getBaseLlmClient(),
-    abortSignal,
-    config.getDisableLLMCorrection(),
-    aggressiveUnescape,
+  const fileExt = path.extname(filePath).toLowerCase();
+  const isJsonOrIpynb = ['.json', '.ipynb', '.jsonc', '.json5'].includes(
+    fileExt,
   );
 
+  if (!isJsonOrIpynb) {
+    const activeModel = config.getActiveModel();
+    const resolvedModel = resolveModel(activeModel, false, false, true, config);
+
+    const aggressiveUnescape =
+      !isGemini3Model(resolvedModel, config) &&
+      !isGemini2Model(resolvedModel) &&
+      !isCustomModel(resolvedModel, config);
+
+    correctedContent = await ensureCorrectFileContent(
+      proposedContent,
+      config.getBaseLlmClient(),
+      abortSignal,
+      config.getDisableLLMCorrection(),
+      aggressiveUnescape,
+    );
+  }
+
   return { originalContent, correctedContent, fileExists };
 }