From a8b115caa1765261087d011530ffd7199d6b96c5 Mon Sep 17 00:00:00 2001 From: amelidev Date: Thu, 9 Jul 2026 13:31:50 -0600 Subject: [PATCH] fix(core-tools): bypass LLM correction for JSON and IPYNB files in write_file and replace (#28223) Co-authored-by: David Pierce --- packages/core/src/tools/edit.test.ts | 38 ++++++++++++++++++++++ packages/core/src/tools/edit.ts | 7 +++- packages/core/src/tools/write-file.test.ts | 33 +++++++++++++++++++ packages/core/src/tools/write-file.ts | 35 ++++++++++++-------- 4 files changed, 98 insertions(+), 15 deletions(-) diff --git a/packages/core/src/tools/edit.test.ts b/packages/core/src/tools/edit.test.ts index 8af188ba97..c164490eee 100644 --- a/packages/core/src/tools/edit.test.ts +++ b/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', () => { diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts index 19d017f642..47a5721028 100644 --- a/packages/core/src/tools/edit.ts +++ b/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, diff --git a/packages/core/src/tools/write-file.test.ts b/packages/core/src/tools/write-file.test.ts index 038e1b6ead..7e2c3f86c8 100644 --- a/packages/core/src/tools/write-file.test.ts +++ b/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'; diff --git a/packages/core/src/tools/write-file.ts b/packages/core/src/tools/write-file.ts index ad35e72c8d..92e005b3bb 100644 --- a/packages/core/src/tools/write-file.ts +++ b/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 }; }