Fix(test): Improve write-file and editCorrector test suites

- Enhanced  by:
  - Mocking  and  utilities (, ) to allow for more focused unit testing of .
  - Adding comprehensive tests for the private  method, covering new and existing file scenarios, as well as error handling.
  - Expanding tests for  and  to verify behavior with the new content correction logic, including diff generation and directory creation.
- Refined  by:
  - Introducing robust mocking for  and its methods (, , ) to simulate LLM interactions accurately.
  - Adding extensive test scenarios for , categorized by how  matches and how  is processed, including direct matches, unescaping, and LLM-based corrections.
  - Including tests for edge cases like no matches or multiple matches.
  - Adding a  utility for better test isolation.

Final fix for https://github.com/google-gemini/gemini-cli/issues/484

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

@@ -4,17 +4,50 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
+import {
+  describe,
+  it,
+  expect,
+  beforeEach,
+  afterEach,
+  vi,
+  type Mocked,
+} from 'vitest';
 import { WriteFileTool } from './write-file.js';
-import { FileDiff, ToolConfirmationOutcome } from './tools.js';
+import {
+  FileDiff,
+  ToolConfirmationOutcome,
+  ToolEditConfirmationDetails,
+} from './tools.js';
+import { type EditToolParams } from './edit.js';
 import { Config } from '../config/config.js';
-import { ToolRegistry } from './tool-registry.js'; // Added import
+import { ToolRegistry } from './tool-registry.js';
 import path from 'path';
 import fs from 'fs';
 import os from 'os';
+import { GeminiClient } from '../core/client.js';
+import {
+  ensureCorrectEdit,
+  ensureCorrectFileContent,
+  CorrectedEditResult,
+} from '../utils/editCorrector.js';
 
 const rootDir = path.resolve(os.tmpdir(), 'gemini-cli-test-root');
 
+// --- MOCKS ---
+vi.mock('../core/client.js');
+vi.mock('../utils/editCorrector.js');
+
+let mockGeminiClientInstance: Mocked<GeminiClient>;
+const mockEnsureCorrectEdit = vi.fn<typeof ensureCorrectEdit>();
+const mockEnsureCorrectFileContent = vi.fn<typeof ensureCorrectFileContent>();
+
+// Wire up the mocked functions to be used by the actual module imports
+vi.mocked(ensureCorrectEdit).mockImplementation(mockEnsureCorrectEdit);
+vi.mocked(ensureCorrectFileContent).mockImplementation(
+  mockEnsureCorrectFileContent,
+);
+
 // Mock Config
 const mockConfigInternal = {
   getTargetDir: () => rootDir,
@@ -42,13 +75,14 @@ const mockConfigInternal = {
     }) as unknown as ToolRegistry,
 };
 const mockConfig = mockConfigInternal as unknown as Config;
+// --- END MOCKS ---
 
 describe('WriteFileTool', () => {
   let tool: WriteFileTool;
   let tempDir: string;
 
   beforeEach(() => {
-    // Create a unique temporary directory for files created outside the root (for testing boundary conditions)
+    // Create a unique temporary directory for files created outside the root
     tempDir = fs.mkdtempSync(
       path.join(os.tmpdir(), 'write-file-test-external-'),
     );
@@ -56,12 +90,39 @@ describe('WriteFileTool', () => {
     if (!fs.existsSync(rootDir)) {
       fs.mkdirSync(rootDir, { recursive: true });
     }
+
+    // Setup GeminiClient mock
+    mockGeminiClientInstance = new (vi.mocked(GeminiClient))(
+      mockConfig,
+    ) as Mocked<GeminiClient>;
+    vi.mocked(GeminiClient).mockImplementation(() => mockGeminiClientInstance);
+
     tool = new WriteFileTool(mockConfig);
-    // Reset mocks before each test that might use them for confirmation logic
+
+    // Reset mocks before each test
     mockConfigInternal.getAlwaysSkipModificationConfirmation.mockReturnValue(
       false,
     );
     mockConfigInternal.setAlwaysSkipModificationConfirmation.mockClear();
+    mockEnsureCorrectEdit.mockReset();
+    mockEnsureCorrectFileContent.mockReset();
+
+    // Default mock implementations that return valid structures
+    mockEnsureCorrectEdit.mockImplementation(
+      async (
+        currentContent: string,
+        params: EditToolParams,
+        _client: GeminiClient,
+      ): Promise<CorrectedEditResult> =>
+        Promise.resolve({
+          params: { ...params, new_string: params.new_string ?? '' },
+          occurrences: 1,
+        }),
+    );
+    mockEnsureCorrectFileContent.mockImplementation(
+      async (content: string, _client: GeminiClient): Promise<string> =>
+        Promise.resolve(content ?? ''),
+    );
   });
 
   afterEach(() => {
@@ -72,6 +133,7 @@ describe('WriteFileTool', () => {
     if (fs.existsSync(rootDir)) {
       fs.rmSync(rootDir, { recursive: true, force: true });
     }
+    vi.clearAllMocks();
   });
 
   describe('validateToolParams', () => {
@@ -101,27 +163,113 @@ describe('WriteFileTool', () => {
       );
     });
 
-    it('should return error for path that is the root itself', () => {
+    it('should return error if path is a directory', () => {
+      const dirAsFilePath = path.join(rootDir, 'a_directory');
+      fs.mkdirSync(dirAsFilePath);
       const params = {
-        file_path: rootDir, // Attempting to write to the root directory itself
+        file_path: dirAsFilePath,
         content: 'hello',
       };
-      // With the new validation, this should now return an error as rootDir is a directory.
       expect(tool.validateToolParams(params)).toMatch(
-        `Path is a directory, not a file: ${rootDir}`,
+        `Path is a directory, not a file: ${dirAsFilePath}`,
       );
     });
+  });
 
-    it('should return error for path that is just / and root is not /', () => {
+  describe('_getCorrectedFileContent', () => {
-      const params = { file_path: path.resolve('/'), content: 'hello' };
+    it('should call ensureCorrectFileContent for a new file', async () => {
-      if (rootDir === path.resolve('/')) {
+      const filePath = path.join(rootDir, 'new_corrected_file.txt');
-        // This case would only occur if the test runner somehow sets rootDir to actual '/', which is highly unlikely and unsafe.
+      const proposedContent = 'Proposed new content.';
-        expect(tool.validateToolParams(params)).toBeNull();
+      const correctedContent = 'Corrected new content.';
-      } else {
+      // Ensure the mock is set for this specific test case if needed, or rely on beforeEach
-        expect(tool.validateToolParams(params)).toMatch(
+      mockEnsureCorrectFileContent.mockResolvedValue(correctedContent);
-          /File path must be within the root directory/,
+
-        );
+      // @ts-expect-error _getCorrectedFileContent is private
-      }
+      const result = await tool._getCorrectedFileContent(
+        filePath,
+        proposedContent,
+      );
+
+      expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith(
+        proposedContent,
+        mockGeminiClientInstance,
+      );
+      expect(mockEnsureCorrectEdit).not.toHaveBeenCalled();
+      expect(result.correctedContent).toBe(correctedContent);
+      expect(result.originalContent).toBe('');
+      expect(result.fileExists).toBe(false);
+      expect(result.error).toBeUndefined();
+    });
+
+    it('should call ensureCorrectEdit for an existing file', async () => {
+      const filePath = path.join(rootDir, 'existing_corrected_file.txt');
+      const originalContent = 'Original existing content.';
+      const proposedContent = 'Proposed replacement content.';
+      const correctedProposedContent = 'Corrected replacement content.';
+      fs.writeFileSync(filePath, originalContent, 'utf8');
+
+      // Ensure this mock is active and returns the correct structure
+      mockEnsureCorrectEdit.mockResolvedValue({
+        params: {
+          file_path: filePath,
+          old_string: originalContent,
+          new_string: correctedProposedContent,
+        },
+        occurrences: 1,
+      } as CorrectedEditResult);
+
+      // @ts-expect-error _getCorrectedFileContent is private
+      const result = await tool._getCorrectedFileContent(
+        filePath,
+        proposedContent,
+      );
+
+      expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
+        originalContent,
+        {
+          old_string: originalContent,
+          new_string: proposedContent,
+          file_path: filePath,
+        },
+        mockGeminiClientInstance,
+      );
+      expect(mockEnsureCorrectFileContent).not.toHaveBeenCalled();
+      expect(result.correctedContent).toBe(correctedProposedContent);
+      expect(result.originalContent).toBe(originalContent);
+      expect(result.fileExists).toBe(true);
+      expect(result.error).toBeUndefined();
+    });
+
+    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';
+      fs.writeFileSync(filePath, 'content', { mode: 0o000 });
+
+      const readError = new Error('Permission denied');
+      const originalReadFileSync = fs.readFileSync;
+      vi.spyOn(fs, 'readFileSync').mockImplementationOnce(() => {
+        throw readError;
+      });
+
+      // @ts-expect-error _getCorrectedFileContent is private
+      const result = await tool._getCorrectedFileContent(
+        filePath,
+        proposedContent,
+      );
+
+      expect(fs.readFileSync).toHaveBeenCalledWith(filePath, 'utf8');
+      expect(mockEnsureCorrectEdit).not.toHaveBeenCalled();
+      expect(mockEnsureCorrectFileContent).not.toHaveBeenCalled();
+      expect(result.correctedContent).toBe(proposedContent);
+      expect(result.originalContent).toBe('');
+      expect(result.fileExists).toBe(true);
+      expect(result.error).toEqual({
+        message: 'Permission denied',
+        code: undefined,
+      });
+
+      vi.spyOn(fs, 'readFileSync').mockImplementation(originalReadFileSync);
+      fs.chmodSync(filePath, 0o600);
     });
   });
 
@@ -139,17 +287,96 @@ describe('WriteFileTool', () => {
       expect(confirmation).toBe(false);
     });
 
-    it('should request confirmation for valid params if file does not exist', async () => {
+    it('should return false if _getCorrectedFileContent returns an error', async () => {
-      const filePath = path.join(rootDir, 'new_file.txt');
+      const filePath = path.join(rootDir, 'confirm_error_file.txt');
-      const params = { file_path: filePath, content: 'new content' };
+      const params = { file_path: filePath, content: 'test content' };
+      fs.writeFileSync(filePath, 'original', { mode: 0o000 });
+
+      const readError = new Error('Simulated read error for confirmation');
+      const originalReadFileSync = fs.readFileSync;
+      vi.spyOn(fs, 'readFileSync').mockImplementationOnce(() => {
+        throw readError;
+      });
+
       const confirmation = await tool.shouldConfirmExecute(params);
+      expect(confirmation).toBe(false);
+
+      vi.spyOn(fs, 'readFileSync').mockImplementation(originalReadFileSync);
+      fs.chmodSync(filePath, 0o600);
+    });
+
+    it('should request confirmation with diff for a new file (with corrected content)', async () => {
+      const filePath = path.join(rootDir, 'confirm_new_file.txt');
+      const proposedContent = 'Proposed new content for confirmation.';
+      const correctedContent = 'Corrected new content for confirmation.';
+      mockEnsureCorrectFileContent.mockResolvedValue(correctedContent); // Ensure this mock is active
+
+      const params = { file_path: filePath, content: proposedContent };
+      const confirmation = (await tool.shouldConfirmExecute(
+        params,
+      )) as ToolEditConfirmationDetails;
+
+      expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith(
+        proposedContent,
+        mockGeminiClientInstance,
+      );
       expect(confirmation).toEqual(
         expect.objectContaining({
           title: `Confirm Write: ${path.basename(filePath)}`,
-          fileName: 'new_file.txt',
+          fileName: 'confirm_new_file.txt',
-          fileDiff: expect.any(String),
+          fileDiff: expect.stringContaining(correctedContent),
         }),
       );
+      expect(confirmation.fileDiff).toMatch(
+        /--- confirm_new_file.txt\tCurrent/,
+      );
+      expect(confirmation.fileDiff).toMatch(
+        /\+\+\+ confirm_new_file.txt\tProposed/,
+      );
+    });
+
+    it('should request confirmation with diff for an existing file (with corrected content)', async () => {
+      const filePath = path.join(rootDir, 'confirm_existing_file.txt');
+      const originalContent = 'Original content for confirmation.';
+      const proposedContent = 'Proposed replacement for confirmation.';
+      const correctedProposedContent =
+        'Corrected replacement for confirmation.';
+      fs.writeFileSync(filePath, originalContent, 'utf8');
+
+      // Ensure this mock is active and returns the correct structure
+      mockEnsureCorrectEdit.mockResolvedValue({
+        params: {
+          file_path: filePath,
+          old_string: originalContent,
+          new_string: correctedProposedContent,
+        },
+        occurrences: 1,
+      });
+
+      const params = { file_path: filePath, content: proposedContent };
+      const confirmation = (await tool.shouldConfirmExecute(
+        params,
+      )) as ToolEditConfirmationDetails;
+
+      expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
+        originalContent,
+        {
+          old_string: originalContent,
+          new_string: proposedContent,
+          file_path: filePath,
+        },
+        mockGeminiClientInstance,
+      );
+      expect(confirmation).toEqual(
+        expect.objectContaining({
+          title: `Confirm Write: ${path.basename(filePath)}`,
+          fileName: 'confirm_existing_file.txt',
+          fileDiff: expect.stringContaining(correctedProposedContent),
+        }),
+      );
+      expect(confirmation.fileDiff).toMatch(
+        originalContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'),
+      );
     });
   });
 
@@ -171,10 +398,34 @@ describe('WriteFileTool', () => {
       );
     });
 
-    it('should write a new file and return diff', async () => {
+    it('should return error if _getCorrectedFileContent returns an error during execute', async () => {
-      const filePath = path.join(rootDir, 'execute_new_file.txt');
+      const filePath = path.join(rootDir, 'execute_error_file.txt');
-      const content = 'Hello from execute!';
+      const params = { file_path: filePath, content: 'test content' };
-      const params = { file_path: filePath, content };
+      fs.writeFileSync(filePath, 'original', { mode: 0o000 });
+
+      const readError = new Error('Simulated read error for execute');
+      const originalReadFileSync = fs.readFileSync;
+      vi.spyOn(fs, 'readFileSync').mockImplementationOnce(() => {
+        throw readError;
+      });
+
+      const result = await tool.execute(params, new AbortController().signal);
+      expect(result.llmContent).toMatch(/Error checking existing file/);
+      expect(result.returnDisplay).toMatch(
+        /Error checking existing file: Simulated read error for execute/,
+      );
+
+      vi.spyOn(fs, 'readFileSync').mockImplementation(originalReadFileSync);
+      fs.chmodSync(filePath, 0o600);
+    });
+
+    it('should write a new file with corrected content and return diff', async () => {
+      const filePath = path.join(rootDir, 'execute_new_corrected_file.txt');
+      const proposedContent = 'Proposed new content for execute.';
+      const correctedContent = 'Corrected new content for execute.';
+      mockEnsureCorrectFileContent.mockResolvedValue(correctedContent);
+
+      const params = { file_path: filePath, content: proposedContent };
 
       const confirmDetails = await tool.shouldConfirmExecute(params);
       if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) {
@@ -183,26 +434,48 @@ describe('WriteFileTool', () => {
 
       const result = await tool.execute(params, new AbortController().signal);
 
+      expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith(
+        proposedContent,
+        mockGeminiClientInstance,
+      );
       expect(result.llmContent).toMatch(
         /Successfully created and wrote to new file/,
       );
       expect(fs.existsSync(filePath)).toBe(true);
-      expect(fs.readFileSync(filePath, 'utf8')).toBe(content);
+      expect(fs.readFileSync(filePath, 'utf8')).toBe(correctedContent);
-      const display = result.returnDisplay as FileDiff; // Type assertion
+      const display = result.returnDisplay as FileDiff;
-      expect(display.fileName).toBe('execute_new_file.txt');
+      expect(display.fileName).toBe('execute_new_corrected_file.txt');
-      // For new files, the diff will include the filename in the "Original" header
+      expect(display.fileDiff).toMatch(
-      expect(display.fileDiff).toMatch(/--- execute_new_file.txt\tOriginal/);
+        /--- execute_new_corrected_file.txt\tOriginal/,
-      expect(display.fileDiff).toMatch(/\+\+\+ execute_new_file.txt\tWritten/);
+      );
-      expect(display.fileDiff).toMatch(content);
+      expect(display.fileDiff).toMatch(
+        /\+\+\+ execute_new_corrected_file.txt\tWritten/,
+      );
+      expect(display.fileDiff).toMatch(
+        correctedContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'),
+      );
     });
 
-    it('should overwrite an existing file and return diff', async () => {
+    it('should overwrite an existing file with corrected content and return diff', async () => {
-      const filePath = path.join(rootDir, 'execute_existing_file.txt');
+      const filePath = path.join(
-      const initialContent = 'Initial content.';
+        rootDir,
-      const newContent = 'Overwritten content!';
+        'execute_existing_corrected_file.txt',
+      );
+      const initialContent = 'Initial content for execute.';
+      const proposedContent = 'Proposed overwrite for execute.';
+      const correctedProposedContent = 'Corrected overwrite for execute.';
       fs.writeFileSync(filePath, initialContent, 'utf8');
 
-      const params = { file_path: filePath, content: newContent };
+      mockEnsureCorrectEdit.mockResolvedValue({
+        params: {
+          file_path: filePath,
+          old_string: initialContent,
+          new_string: correctedProposedContent,
+        },
+        occurrences: 1,
+      });
+
+      const params = { file_path: filePath, content: proposedContent };
 
       const confirmDetails = await tool.shouldConfirmExecute(params);
       if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) {
@@ -211,12 +484,46 @@ describe('WriteFileTool', () => {
 
       const result = await tool.execute(params, new AbortController().signal);
 
+      expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
+        initialContent,
+        {
+          old_string: initialContent,
+          new_string: proposedContent,
+          file_path: filePath,
+        },
+        mockGeminiClientInstance,
+      );
       expect(result.llmContent).toMatch(/Successfully overwrote file/);
-      expect(fs.readFileSync(filePath, 'utf8')).toBe(newContent);
+      expect(fs.readFileSync(filePath, 'utf8')).toBe(correctedProposedContent);
-      const display = result.returnDisplay as FileDiff; // Type assertion
+      const display = result.returnDisplay as FileDiff;
-      expect(display.fileName).toBe('execute_existing_file.txt');
+      expect(display.fileName).toBe('execute_existing_corrected_file.txt');
-      expect(display.fileDiff).toMatch(initialContent);
+      expect(display.fileDiff).toMatch(
-      expect(display.fileDiff).toMatch(newContent);
+        initialContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'),
+      );
+      expect(display.fileDiff).toMatch(
+        correctedProposedContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'),
+      );
+    });
+
+    it('should create directory if it does not exist', async () => {
+      const dirPath = path.join(rootDir, 'new_dir_for_write');
+      const filePath = path.join(dirPath, 'file_in_new_dir.txt');
+      const content = 'Content in new directory';
+      mockEnsureCorrectFileContent.mockResolvedValue(content); // Ensure this mock is active
+
+      const params = { file_path: filePath, content };
+      // Simulate confirmation if your logic requires it before execute, or remove if not needed for this path
+      const confirmDetails = await tool.shouldConfirmExecute(params);
+      if (typeof confirmDetails === 'object' && confirmDetails.onConfirm) {
+        await confirmDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce);
+      }
+
+      await tool.execute(params, new AbortController().signal);
+
+      expect(fs.existsSync(dirPath)).toBe(true);
+      expect(fs.statSync(dirPath).isDirectory()).toBe(true);
+      expect(fs.existsSync(filePath)).toBe(true);
+      expect(fs.readFileSync(filePath, 'utf8')).toBe(content);
     });
   });
 });

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

@@ -4,44 +4,68 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect } from 'vitest';
+/* eslint-disable @typescript-eslint/no-explicit-any */
+import { vi, describe, it, expect, beforeEach, type Mocked } from 'vitest';
+
+// MOCKS
+let callCount = 0;
+const mockResponses: any[] = [];
+
+let mockGenerateJson: any;
+let mockStartChat: any;
+let mockSendMessageStream: any;
+
+vi.mock('../core/client.js', () => ({
+  GeminiClient: vi.fn().mockImplementation(function (
+    this: any,
+    _config: Config,
+  ) {
+    this.generateJson = (...params: any[]) => mockGenerateJson(...params); // Corrected: use mockGenerateJson
+    this.startChat = (...params: any[]) => mockStartChat(...params); // Corrected: use mockStartChat
+    this.sendMessageStream = (...params: any[]) =>
+      mockSendMessageStream(...params); // Corrected: use mockSendMessageStream
+    return this;
+  }),
+}));
+// END MOCKS
+
 import {
   countOccurrences,
+  ensureCorrectEdit,
   unescapeStringForGeminiBug,
+  resetEditCorrectorCaches_TEST_ONLY,
 } from './editCorrector.js';
+import { GeminiClient } from '../core/client.js';
+import type { Config } from '../config/config.js';
+import { ToolRegistry } from '../tools/tool-registry.js';
+
+vi.mock('../tools/tool-registry.js');
 
 describe('editCorrector', () => {
   describe('countOccurrences', () => {
     it('should return 0 for empty string', () => {
       expect(countOccurrences('', 'a')).toBe(0);
     });
-
     it('should return 0 for empty substring', () => {
       expect(countOccurrences('abc', '')).toBe(0);
     });
-
     it('should return 0 if substring is not found', () => {
       expect(countOccurrences('abc', 'd')).toBe(0);
     });
-
     it('should return 1 if substring is found once', () => {
       expect(countOccurrences('abc', 'b')).toBe(1);
     });
-
     it('should return correct count for multiple occurrences', () => {
       expect(countOccurrences('ababa', 'a')).toBe(3);
       expect(countOccurrences('ababab', 'ab')).toBe(3);
     });
-
     it('should count non-overlapping occurrences', () => {
-      expect(countOccurrences('aaaaa', 'aa')).toBe(2); // Non-overlapping: aa_aa_
+      expect(countOccurrences('aaaaa', 'aa')).toBe(2);
-      expect(countOccurrences('ababab', 'aba')).toBe(1); // Non-overlapping: aba_ab -> 1
+      expect(countOccurrences('ababab', 'aba')).toBe(1);
     });
-
     it('should correctly count occurrences when substring is longer', () => {
       expect(countOccurrences('abc', 'abcdef')).toBe(0);
     });
-
     it('should be case sensitive', () => {
       expect(countOccurrences('abcABC', 'a')).toBe(1);
       expect(countOccurrences('abcABC', 'A')).toBe(1);
@@ -56,62 +80,403 @@ describe('editCorrector', () => {
       expect(unescapeStringForGeminiBug('\\"')).toBe('"');
       expect(unescapeStringForGeminiBug('\\`')).toBe('`');
     });
-
     it('should handle multiple escaped sequences', () => {
       expect(unescapeStringForGeminiBug('Hello\\nWorld\\tTest')).toBe(
         'Hello\nWorld\tTest',
       );
     });
-
     it('should not alter already correct sequences', () => {
       expect(unescapeStringForGeminiBug('\n')).toBe('\n');
       expect(unescapeStringForGeminiBug('Correct string')).toBe(
         'Correct string',
       );
     });
-
     it('should handle mixed correct and incorrect sequences', () => {
       expect(unescapeStringForGeminiBug('\\nCorrect\t\\`')).toBe(
         '\nCorrect\t`',
       );
     });
-
     it('should handle backslash followed by actual newline character', () => {
       expect(unescapeStringForGeminiBug('\\\n')).toBe('\n');
       expect(unescapeStringForGeminiBug('First line\\\nSecond line')).toBe(
         'First line\nSecond line',
       );
     });
-
     it('should handle multiple backslashes before an escapable character', () => {
-      expect(unescapeStringForGeminiBug('\\\\n')).toBe('\n'); // \\n -> \n
+      expect(unescapeStringForGeminiBug('\\\\n')).toBe('\n');
-      expect(unescapeStringForGeminiBug('\\\\\\t')).toBe('\t'); // \\\t -> \t
+      expect(unescapeStringForGeminiBug('\\\\\\t')).toBe('\t');
-      expect(unescapeStringForGeminiBug('\\\\\\\\`')).toBe('`'); // \\\\` -> `
+      expect(unescapeStringForGeminiBug('\\\\\\\\`')).toBe('`');
     });
-
     it('should return empty string for empty input', () => {
       expect(unescapeStringForGeminiBug('')).toBe('');
     });
-
     it('should not alter strings with no targeted escape sequences', () => {
       expect(unescapeStringForGeminiBug('abc def')).toBe('abc def');
-      // \\F and \\S are not targeted escapes, so they should remain as \\F and \\S
       expect(unescapeStringForGeminiBug('C:\\Folder\\File')).toBe(
         'C:\\Folder\\File',
       );
     });
-
     it('should correctly process strings with some targeted escapes', () => {
-      // \\U is not targeted, \\n is.
       expect(unescapeStringForGeminiBug('C:\\Users\\name')).toBe(
         'C:\\Users\name',
       );
     });
-
     it('should handle complex cases with mixed slashes and characters', () => {
       expect(
-        unescapeStringForGeminiBug('\\\\\\nLine1\\\nLine2\\tTab\\\\`Tick\\"'),
+        unescapeStringForGeminiBug('\\\\\\\nLine1\\\nLine2\\tTab\\\\`Tick\\"'),
       ).toBe('\nLine1\nLine2\tTab`Tick"');
     });
   });
+
+  describe('ensureCorrectEdit', () => {
+    let mockGeminiClientInstance: Mocked<GeminiClient>;
+    let mockToolRegistry: Mocked<ToolRegistry>;
+    let mockConfigInstance: Config;
+
+    beforeEach(() => {
+      mockToolRegistry = new ToolRegistry({} as Config) as Mocked<ToolRegistry>;
+      const configParams = {
+        apiKey: 'test-api-key',
+        model: 'test-model',
+        sandbox: false as boolean | string,
+        targetDir: '/test',
+        debugMode: false,
+        question: undefined as string | undefined,
+        fullContext: false,
+        coreTools: undefined as string[] | undefined,
+        toolDiscoveryCommand: undefined as string | undefined,
+        toolCallCommand: undefined as string | undefined,
+        mcpServerCommand: undefined as string | undefined,
+        mcpServers: undefined as Record<string, any> | undefined,
+        userAgent: 'test-agent',
+        userMemory: '',
+        geminiMdFileCount: 0,
+        alwaysSkipModificationConfirmation: false,
+      };
+      mockConfigInstance = {
+        ...configParams,
+        getApiKey: vi.fn(() => configParams.apiKey),
+        getModel: vi.fn(() => configParams.model),
+        getSandbox: vi.fn(() => configParams.sandbox),
+        getTargetDir: vi.fn(() => configParams.targetDir),
+        getToolRegistry: vi.fn(() => mockToolRegistry),
+        getDebugMode: vi.fn(() => configParams.debugMode),
+        getQuestion: vi.fn(() => configParams.question),
+        getFullContext: vi.fn(() => configParams.fullContext),
+        getCoreTools: vi.fn(() => configParams.coreTools),
+        getToolDiscoveryCommand: vi.fn(() => configParams.toolDiscoveryCommand),
+        getToolCallCommand: vi.fn(() => configParams.toolCallCommand),
+        getMcpServerCommand: vi.fn(() => configParams.mcpServerCommand),
+        getMcpServers: vi.fn(() => configParams.mcpServers),
+        getUserAgent: vi.fn(() => configParams.userAgent),
+        getUserMemory: vi.fn(() => configParams.userMemory),
+        setUserMemory: vi.fn((mem: string) => {
+          configParams.userMemory = mem;
+        }),
+        getGeminiMdFileCount: vi.fn(() => configParams.geminiMdFileCount),
+        setGeminiMdFileCount: vi.fn((count: number) => {
+          configParams.geminiMdFileCount = count;
+        }),
+        getAlwaysSkipModificationConfirmation: vi.fn(
+          () => configParams.alwaysSkipModificationConfirmation,
+        ),
+        setAlwaysSkipModificationConfirmation: vi.fn((skip: boolean) => {
+          configParams.alwaysSkipModificationConfirmation = skip;
+        }),
+      } as unknown as Config;
+
+      callCount = 0;
+      mockResponses.length = 0;
+      mockGenerateJson = vi.fn().mockImplementation(() => {
+        const response = mockResponses[callCount];
+        callCount++;
+        if (response === undefined) return Promise.resolve({});
+        return Promise.resolve(response);
+      });
+      mockStartChat = vi.fn();
+      mockSendMessageStream = vi.fn();
+
+      mockGeminiClientInstance = new GeminiClient(
+        mockConfigInstance,
+      ) as Mocked<GeminiClient>;
+      resetEditCorrectorCaches_TEST_ONLY();
+    });
+
+    describe('Scenario Group 1: originalParams.old_string matches currentContent directly', () => {
+      it('Test 1.1: old_string (no literal \\), new_string (escaped by Gemini) -> new_string unescaped', async () => {
+        const currentContent = 'This is a test string to find me.';
+        const originalParams = {
+          file_path: '/test/file.txt',
+          old_string: 'find me',
+          new_string: 'replace with \\"this\\"',
+        };
+        mockResponses.push({
+          corrected_new_string_escaping: 'replace with "this"',
+        });
+        const result = await ensureCorrectEdit(
+          currentContent,
+          originalParams,
+          mockGeminiClientInstance,
+        );
+        expect(mockGenerateJson).toHaveBeenCalledTimes(1);
+        expect(result.params.new_string).toBe('replace with "this"');
+        expect(result.params.old_string).toBe('find me');
+        expect(result.occurrences).toBe(1);
+      });
+      it('Test 1.2: old_string (no literal \\), new_string (correctly formatted) -> new_string unchanged', async () => {
+        const currentContent = 'This is a test string to find me.';
+        const originalParams = {
+          file_path: '/test/file.txt',
+          old_string: 'find me',
+          new_string: 'replace with this',
+        };
+        const result = await ensureCorrectEdit(
+          currentContent,
+          originalParams,
+          mockGeminiClientInstance,
+        );
+        expect(mockGenerateJson).toHaveBeenCalledTimes(0);
+        expect(result.params.new_string).toBe('replace with this');
+        expect(result.params.old_string).toBe('find me');
+        expect(result.occurrences).toBe(1);
+      });
+      it('Test 1.3: old_string (with literal \\), new_string (escaped by Gemini) -> new_string unchanged (still escaped)', async () => {
+        const currentContent = 'This is a test string to find\\me.';
+        const originalParams = {
+          file_path: '/test/file.txt',
+          old_string: 'find\\me',
+          new_string: 'replace with \\"this\\"',
+        };
+        mockResponses.push({
+          corrected_new_string_escaping: 'replace with "this"',
+        });
+        const result = await ensureCorrectEdit(
+          currentContent,
+          originalParams,
+          mockGeminiClientInstance,
+        );
+        expect(mockGenerateJson).toHaveBeenCalledTimes(1);
+        expect(result.params.new_string).toBe('replace with "this"');
+        expect(result.params.old_string).toBe('find\\me');
+        expect(result.occurrences).toBe(1);
+      });
+      it('Test 1.4: old_string (with literal \\), new_string (correctly formatted) -> new_string unchanged', async () => {
+        const currentContent = 'This is a test string to find\\me.';
+        const originalParams = {
+          file_path: '/test/file.txt',
+          old_string: 'find\\me',
+          new_string: 'replace with this',
+        };
+        const result = await ensureCorrectEdit(
+          currentContent,
+          originalParams,
+          mockGeminiClientInstance,
+        );
+        expect(mockGenerateJson).toHaveBeenCalledTimes(0);
+        expect(result.params.new_string).toBe('replace with this');
+        expect(result.params.old_string).toBe('find\\me');
+        expect(result.occurrences).toBe(1);
+      });
+    });
+
+    describe('Scenario Group 2: originalParams.old_string does NOT match, but unescapeStringForGeminiBug(originalParams.old_string) DOES match', () => {
+      it('Test 2.1: old_string (over-escaped, no intended literal \\), new_string (escaped by Gemini) -> new_string unescaped', async () => {
+        const currentContent = 'This is a test string to find "me".';
+        const originalParams = {
+          file_path: '/test/file.txt',
+          old_string: 'find \\"me\\"',
+          new_string: 'replace with \\"this\\"',
+        };
+        mockResponses.push({ corrected_new_string: 'replace with "this"' });
+        const result = await ensureCorrectEdit(
+          currentContent,
+          originalParams,
+          mockGeminiClientInstance,
+        );
+        expect(mockGenerateJson).toHaveBeenCalledTimes(1);
+        expect(result.params.new_string).toBe('replace with "this"');
+        expect(result.params.old_string).toBe('find "me"');
+        expect(result.occurrences).toBe(1);
+      });
+      it('Test 2.2: old_string (over-escaped, no intended literal \\), new_string (correctly formatted) -> new_string unescaped (harmlessly)', async () => {
+        const currentContent = 'This is a test string to find "me".';
+        const originalParams = {
+          file_path: '/test/file.txt',
+          old_string: 'find \\"me\\"',
+          new_string: 'replace with this',
+        };
+        const result = await ensureCorrectEdit(
+          currentContent,
+          originalParams,
+          mockGeminiClientInstance,
+        );
+        expect(mockGenerateJson).toHaveBeenCalledTimes(0);
+        expect(result.params.new_string).toBe('replace with this');
+        expect(result.params.old_string).toBe('find "me"');
+        expect(result.occurrences).toBe(1);
+      });
+      it('Test 2.3: old_string (over-escaped, with intended literal \\), new_string (simple) -> new_string corrected', async () => {
+        const currentContent = 'This is a test string to find \\me.';
+        const originalParams = {
+          file_path: '/test/file.txt',
+          old_string: 'find \\\\me',
+          new_string: 'replace with foobar',
+        };
+        mockResponses.push({
+          corrected_target_snippet: 'find \\me',
+        });
+        const result = await ensureCorrectEdit(
+          currentContent,
+          originalParams,
+          mockGeminiClientInstance,
+        );
+        expect(mockGenerateJson).toHaveBeenCalledTimes(1);
+        expect(result.params.new_string).toBe('replace with foobar');
+        expect(result.params.old_string).toBe('find \\me');
+        expect(result.occurrences).toBe(1);
+      });
+    });
+
+    describe('Scenario Group 3: LLM Correction Path', () => {
+      it('Test 3.1: old_string (no literal \\), new_string (escaped by Gemini), LLM re-escapes new_string -> final new_string is double unescaped', async () => {
+        const currentContent = 'This is a test string to corrected find me.';
+        const originalParams = {
+          file_path: '/test/file.txt',
+          old_string: 'find me',
+          new_string: 'replace with \\\\"this\\\\"',
+        };
+        const llmNewString = 'LLM says replace with "that"';
+        mockResponses.push({ corrected_new_string_escaping: llmNewString });
+        const result = await ensureCorrectEdit(
+          currentContent,
+          originalParams,
+          mockGeminiClientInstance,
+        );
+        expect(mockGenerateJson).toHaveBeenCalledTimes(1);
+        expect(result.params.new_string).toBe(llmNewString);
+        expect(result.params.old_string).toBe('find me');
+        expect(result.occurrences).toBe(1);
+      });
+      it('Test 3.2: old_string (with literal \\), new_string (escaped by Gemini), LLM re-escapes new_string -> final new_string is unescaped once', async () => {
+        const currentContent = 'This is a test string to corrected find me.';
+        const originalParams = {
+          file_path: '/test/file.txt',
+          old_string: 'find\\me',
+          new_string: 'replace with \\\\"this\\\\"',
+        };
+        const llmCorrectedOldString = 'corrected find me';
+        const llmNewString = 'LLM says replace with "that"';
+        mockResponses.push({ corrected_target_snippet: llmCorrectedOldString });
+        mockResponses.push({ corrected_new_string: llmNewString });
+        const result = await ensureCorrectEdit(
+          currentContent,
+          originalParams,
+          mockGeminiClientInstance,
+        );
+        expect(mockGenerateJson).toHaveBeenCalledTimes(2);
+        expect(result.params.new_string).toBe(llmNewString);
+        expect(result.params.old_string).toBe(llmCorrectedOldString);
+        expect(result.occurrences).toBe(1);
+      });
+      it('Test 3.3: old_string needs LLM, new_string is fine -> old_string corrected, new_string original', async () => {
+        const currentContent = 'This is a test string to be corrected.';
+        const originalParams = {
+          file_path: '/test/file.txt',
+          old_string: 'fiiind me',
+          new_string: 'replace with "this"',
+        };
+        const llmCorrectedOldString = 'to be corrected';
+        mockResponses.push({ corrected_target_snippet: llmCorrectedOldString });
+        const result = await ensureCorrectEdit(
+          currentContent,
+          originalParams,
+          mockGeminiClientInstance,
+        );
+        expect(mockGenerateJson).toHaveBeenCalledTimes(1);
+        expect(result.params.new_string).toBe('replace with "this"');
+        expect(result.params.old_string).toBe(llmCorrectedOldString);
+        expect(result.occurrences).toBe(1);
+      });
+      it('Test 3.4: LLM correction path, correctNewString returns the originalNewString it was passed (which was unescaped) -> final new_string is unescaped', async () => {
+        const currentContent = 'This is a test string to corrected find me.';
+        const originalParams = {
+          file_path: '/test/file.txt',
+          old_string: 'find me',
+          new_string: 'replace with \\\\"this\\\\"',
+        };
+        const newStringForLLMAndReturnedByLLM = 'replace with "this"';
+        mockResponses.push({
+          corrected_new_string_escaping: newStringForLLMAndReturnedByLLM,
+        });
+        const result = await ensureCorrectEdit(
+          currentContent,
+          originalParams,
+          mockGeminiClientInstance,
+        );
+        expect(mockGenerateJson).toHaveBeenCalledTimes(1);
+        expect(result.params.new_string).toBe(newStringForLLMAndReturnedByLLM);
+        expect(result.occurrences).toBe(1);
+      });
+    });
+
+    describe('Scenario Group 4: No Match Found / Multiple Matches', () => {
+      it('Test 4.1: No version of old_string (original, unescaped, LLM-corrected) matches -> returns original params, 0 occurrences', async () => {
+        const currentContent = 'This content has nothing to find.';
+        const originalParams = {
+          file_path: '/test/file.txt',
+          old_string: 'nonexistent string',
+          new_string: 'some new string',
+        };
+        mockResponses.push({ corrected_target_snippet: 'still nonexistent' });
+        const result = await ensureCorrectEdit(
+          currentContent,
+          originalParams,
+          mockGeminiClientInstance,
+        );
+        expect(mockGenerateJson).toHaveBeenCalledTimes(1);
+        expect(result.params).toEqual(originalParams);
+        expect(result.occurrences).toBe(0);
+      });
+      it('Test 4.2: unescapedOldStringAttempt results in >1 occurrences -> returns original params, count occurrences', async () => {
+        const currentContent =
+          'This content has find "me" and also find "me" again.';
+        const originalParams = {
+          file_path: '/test/file.txt',
+          old_string: 'find "me"',
+          new_string: 'some new string',
+        };
+        const result = await ensureCorrectEdit(
+          currentContent,
+          originalParams,
+          mockGeminiClientInstance,
+        );
+        expect(mockGenerateJson).toHaveBeenCalledTimes(0);
+        expect(result.params).toEqual(originalParams);
+        expect(result.occurrences).toBe(2);
+      });
+    });
+
+    describe('Scenario Group 5: Specific unescapeStringForGeminiBug checks (integrated into ensureCorrectEdit)', () => {
+      it('Test 5.1: old_string needs LLM to become currentContent, new_string also needs correction', async () => {
+        const currentContent = 'const x = "a\\nbc\\\\"def\\\\"';
+        const originalParams = {
+          file_path: '/test/file.txt',
+          old_string: 'const x = \\\\"a\\\\nbc\\\\\\\\"def\\\\\\\\"',
+          new_string: 'const y = \\\\"new\\\\nval\\\\\\\\"content\\\\\\\\"',
+        };
+        const expectedFinalNewString = 'const y = "new\\nval\\\\"content\\\\"';
+        mockResponses.push({ corrected_target_snippet: currentContent });
+        mockResponses.push({ corrected_new_string: expectedFinalNewString });
+        const result = await ensureCorrectEdit(
+          currentContent,
+          originalParams,
+          mockGeminiClientInstance,
+        );
+        expect(mockGenerateJson).toHaveBeenCalledTimes(2);
+        expect(result.params.old_string).toBe(currentContent);
+        expect(result.params.new_string).toBe(expectedFinalNewString);
+        expect(result.occurrences).toBe(1);
+      });
+    });
+  });
 });