gemini-cli/packages/core/src/tools/edit.test.ts

/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

/* eslint-disable @typescript-eslint/no-explicit-any */

const mockFixLLMEditWithInstruction = vi.hoisted(() => vi.fn());
const mockGenerateJson = vi.hoisted(() => vi.fn());
const mockOpenDiff = vi.hoisted(() => vi.fn());

import { IdeClient } from '../ide/ide-client.js';

vi.mock('../ide/ide-client.js', () => ({
  IdeClient: {
    getInstance: vi.fn(),
  },
}));

vi.mock('../utils/llm-edit-fixer.js', () => ({
  FixLLMEditWithInstruction: mockFixLLMEditWithInstruction,
}));

vi.mock('../core/client.js', () => ({
  GeminiClient: vi.fn().mockImplementation(() => ({
    generateJson: mockGenerateJson,
    getHistory: vi.fn().mockResolvedValue([]),
  })),
}));

vi.mock('../utils/editor.js', () => ({
  openDiff: mockOpenDiff,
}));

import {
  describe,
  it,
  expect,
  beforeEach,
  afterEach,
  vi,
  type Mock,
} from 'vitest';
import {
  EditTool,
  type EditToolParams,
  applyReplacement,
  calculateReplacement,
} from './edit.js';
import { type FileDiff, ToolConfirmationOutcome } from './tools.js';
import { ToolErrorType } from './tool-error.js';
import {
  createMockMessageBus,
  getMockMessageBusInstance,
} from '../test-utils/mock-message-bus.js';
import path from 'node:path';
import { isSubpath } from '../utils/paths.js';
import fs from 'node:fs';
import os from 'node:os';
import { ApprovalMode } from '../policy/types.js';
import { type Config } from '../config/config.js';
import { type Content, type Part, type SchemaUnion } from '@google/genai';
import { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.js';
import { StandardFileSystemService } from '../services/fileSystemService.js';
import type { BaseLlmClient } from '../core/baseLlmClient.js';

describe('EditTool', () => {
  let tool: EditTool;
  let tempDir: string;
  let rootDir: string;
  let mockConfig: Config;
  let geminiClient: any;
  let fileSystemService: StandardFileSystemService;
  let baseLlmClient: BaseLlmClient;

  beforeEach(() => {
    vi.restoreAllMocks();
    tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'edit-tool-test-'));
    rootDir = path.join(tempDir, 'root');
    fs.mkdirSync(rootDir);

    geminiClient = {
      generateJson: mockGenerateJson,
      getHistory: vi.fn().mockResolvedValue([]),
    };

    baseLlmClient = {
      generateJson: mockGenerateJson,
    } as unknown as BaseLlmClient;

    fileSystemService = new StandardFileSystemService();

    mockConfig = {
      getUsageStatisticsEnabled: vi.fn(() => true),
      getSessionId: vi.fn(() => 'mock-session-id'),
      getContentGeneratorConfig: vi.fn(() => ({ authType: 'mock' })),
      getProxy: vi.fn(() => undefined),
      getGeminiClient: vi.fn().mockReturnValue(geminiClient),
      getBaseLlmClient: vi.fn().mockReturnValue(baseLlmClient),
      getTargetDir: () => rootDir,
      getApprovalMode: vi.fn(),
      setApprovalMode: vi.fn(),
      getWorkspaceContext: () => createMockWorkspaceContext(rootDir),
      getFileSystemService: () => fileSystemService,
      getIdeMode: () => false,
      getApiKey: () => 'test-api-key',
      getModel: () => 'test-model',
      getSandbox: () => false,
      getDebugMode: () => false,
      getQuestion: () => undefined,

      getToolDiscoveryCommand: () => undefined,
      getToolCallCommand: () => undefined,
      getMcpServerCommand: () => undefined,
      getMcpServers: () => undefined,
      getUserAgent: () => 'test-agent',
      getUserMemory: () => '',
      setUserMemory: vi.fn(),
      getGeminiMdFileCount: () => 0,
      setGeminiMdFileCount: vi.fn(),
      getToolRegistry: () => ({}) as any,
      isInteractive: () => false,
      getDisableLLMCorrection: vi.fn(() => true),
      getExperiments: () => {},
      storage: {
        getProjectTempDir: vi.fn().mockReturnValue('/tmp/project'),
      },
      isPathAllowed(this: Config, absolutePath: string): boolean {
        const workspaceContext = this.getWorkspaceContext();
        if (workspaceContext.isPathWithinWorkspace(absolutePath)) {
          return true;
        }

        const projectTempDir = this.storage.getProjectTempDir();
        return isSubpath(path.resolve(projectTempDir), absolutePath);
      },
      validatePathAccess(this: Config, absolutePath: string): string | null {
        if (this.isPathAllowed(absolutePath)) {
          return null;
        }

        const workspaceDirs = this.getWorkspaceContext().getDirectories();
        const projectTempDir = this.storage.getProjectTempDir();
        return `Path not in workspace: Attempted path "${absolutePath}" resolves outside the allowed workspace directories: ${workspaceDirs.join(', ')} or the project temp directory: ${projectTempDir}`;
      },
    } as unknown as Config;

    (mockConfig.getApprovalMode as Mock).mockClear();
    (mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.DEFAULT);

    mockFixLLMEditWithInstruction.mockReset();
    mockFixLLMEditWithInstruction.mockResolvedValue({
      noChangesRequired: false,
      search: '',
      replace: '',
      explanation: 'LLM fix failed',
    });

    mockGenerateJson.mockReset();
    mockGenerateJson.mockImplementation(
      async (contents: Content[], schema: SchemaUnion) => {
        const userContent = contents.find((c: Content) => c.role === 'user');
        let promptText = '';
        if (userContent && userContent.parts) {
          promptText = userContent.parts
            .filter((p: Part) => typeof (p as any).text === 'string')
            .map((p: Part) => (p as any).text)
            .join('\n');
        }
        const snippetMatch = promptText.match(
          /Problematic target snippet:\n```\n([\s\S]*?)\n```/,
        );
        const problematicSnippet =
          snippetMatch && snippetMatch[1] ? snippetMatch[1] : '';

        if ((schema as any).properties?.corrected_target_snippet) {
          return Promise.resolve({
            corrected_target_snippet: problematicSnippet,
          });
        }
        if ((schema as any).properties?.corrected_new_string) {
          const originalNewStringMatch = promptText.match(
            /original_new_string \(what was intended to replace original_old_string\):\n```\n([\s\S]*?)\n```/,
          );
          const originalNewString =
            originalNewStringMatch && originalNewStringMatch[1]
              ? originalNewStringMatch[1]
              : '';
          return Promise.resolve({ corrected_new_string: originalNewString });
        }
        return Promise.resolve({});
      },
    );

    const bus = createMockMessageBus();
    getMockMessageBusInstance(bus).defaultToolDecision = 'ask_user';
    tool = new EditTool(mockConfig, bus);
  });

  afterEach(() => {
    fs.rmSync(tempDir, { recursive: true, force: true });
  });

  describe('applyReplacement', () => {
    it('should return newString if isNewFile is true', () => {
      expect(applyReplacement(null, 'old', 'new', true)).toBe('new');
      expect(applyReplacement('existing', 'old', 'new', true)).toBe('new');
    });

    it('should return newString if currentContent is null and oldString is empty (defensive)', () => {
      expect(applyReplacement(null, '', 'new', false)).toBe('new');
    });

    it('should return empty string if currentContent is null and oldString is not empty (defensive)', () => {
      expect(applyReplacement(null, 'old', 'new', false)).toBe('');
    });

    it('should replace oldString with newString in currentContent', () => {
      expect(applyReplacement('hello old world old', 'old', 'new', false)).toBe(
        'hello new world new',
      );
    });

    it('should return currentContent if oldString is empty and not a new file', () => {
      expect(applyReplacement('hello world', '', 'new', false)).toBe(
        'hello world',
      );
    });

    it.each([
      {
        name: '$ literal',
        current: "price is $100 and pattern end is ' '",
        oldStr: 'price is $100',
        newStr: 'price is $200',
        expected: "price is $200 and pattern end is ' '",
      },
      {
        name: "$' literal",
        current: 'foo',
        oldStr: 'foo',
        newStr: "bar$'baz",
        expected: "bar$'baz",
      },
      {
        name: '$& literal',
        current: 'hello world',
        oldStr: 'hello',
        newStr: '$&-replacement',
        expected: '$&-replacement world',
      },
      {
        name: '$` literal',
        current: 'prefix-middle-suffix',
        oldStr: 'middle',
        newStr: 'new$`content',
        expected: 'prefix-new$`content-suffix',
      },
      {
        name: '$1, $2 capture groups literal',
        current: 'test string',
        oldStr: 'test',
        newStr: '$1$2replacement',
        expected: '$1$2replacement string',
      },
      {
        name: 'normal strings without problematic $',
        current: 'normal text replacement',
        oldStr: 'text',
        newStr: 'string',
        expected: 'normal string replacement',
      },
      {
        name: 'multiple occurrences with $ sequences',
        current: 'foo bar foo baz',
        oldStr: 'foo',
        newStr: "test$'end",
        expected: "test$'end bar test$'end baz",
      },
      {
        name: 'complex regex patterns with $ at end',
        current: "| select('match', '^[sv]d[a-z]$')",
        oldStr: "'^[sv]d[a-z]$'",
        newStr: "'^[sv]d[a-z]$' # updated",
        expected: "| select('match', '^[sv]d[a-z]$' # updated)",
      },
      {
        name: 'empty replacement with problematic $',
        current: 'test content',
        oldStr: 'nothing',
        newStr: "replacement$'text",
        expected: 'test content',
      },
      {
        name: '$$ (escaped dollar)',
        current: 'price value',
        oldStr: 'value',
        newStr: '$$100',
        expected: 'price $$100',
      },
    ])('should handle $name', ({ current, oldStr, newStr, expected }) => {
      const result = applyReplacement(current, oldStr, newStr, false);
      expect(result).toBe(expected);
    });
  });

  describe('calculateReplacement', () => {
    const abortSignal = new AbortController().signal;

    it.each([
      {
        name: 'perform an exact replacement',
        content: 'hello world',
        old_string: 'world',
        new_string: 'moon',
        expected: 'hello moon',
        occurrences: 1,
      },
      {
        name: 'perform a flexible, whitespace-insensitive replacement',
        content: '  hello\n    world\n',
        old_string: 'hello\nworld',
        new_string: 'goodbye\nmoon',
        expected: '  goodbye\n  moon\n',
        occurrences: 1,
      },
      {
        name: 'return 0 occurrences if no match is found',
        content: 'hello world',
        old_string: 'nomatch',
        new_string: 'moon',
        expected: 'hello world',
        occurrences: 0,
      },
    ])(
      'should $name',
      async ({ content, old_string, new_string, expected, occurrences }) => {
        const result = await calculateReplacement(mockConfig, {
          params: {
            file_path: 'test.txt',
            instruction: 'test',
            old_string,
            new_string,
          },
          currentContent: content,
          abortSignal,
        });
        expect(result.newContent).toBe(expected);
        expect(result.occurrences).toBe(occurrences);
      },
    );

    it('should perform a regex-based replacement for flexible intra-line whitespace', async () => {
      // This case would fail with the previous exact and line-trimming flexible logic
      // because the whitespace *within* the line is different.
      const content = '  function  myFunc( a, b ) {\n    return a + b;\n  }';
      const result = await calculateReplacement(mockConfig, {
        params: {
          file_path: 'test.js',
          instruction: 'test',
          old_string: 'function myFunc(a, b) {', // Note the normalized whitespace
          new_string: 'const yourFunc = (a, b) => {',
        },
        currentContent: content,
        abortSignal,
      });

      // The indentation from the original line should be preserved and applied to the new string.
      const expectedContent =
        '  const yourFunc = (a, b) => {\n    return a + b;\n  }';
      expect(result.newContent).toBe(expectedContent);
      expect(result.occurrences).toBe(1);
    });

    it('should perform a fuzzy replacement when exact match fails but similarity is high', async () => {
      const content =
        'const myConfig = {\n  enableFeature: true,\n  retries: 3\n};';
      // Typo: missing comma after true
      const oldString =
        'const myConfig = {\n  enableFeature: true\n  retries: 3\n};';
      const newString =
        'const myConfig = {\n  enableFeature: false,\n  retries: 5\n};';

      const result = await calculateReplacement(mockConfig, {
        params: {
          file_path: 'config.ts',
          instruction: 'update config',
          old_string: oldString,
          new_string: newString,
        },
        currentContent: content,
        abortSignal,
      });

      expect(result.occurrences).toBe(1);
      expect(result.newContent).toBe(newString);
    });

    it('should NOT perform a fuzzy replacement when similarity is below threshold', async () => {
      const content =
        'const myConfig = {\n  enableFeature: true,\n  retries: 3\n};';
      // Completely different string
      const oldString = 'function somethingElse() {\n  return false;\n}';
      const newString =
        'const myConfig = {\n  enableFeature: false,\n  retries: 5\n};';

      const result = await calculateReplacement(mockConfig, {
        params: {
          file_path: 'config.ts',
          instruction: 'update config',
          old_string: oldString,
          new_string: newString,
        },
        currentContent: content,
        abortSignal,
      });

      expect(result.occurrences).toBe(0);
      expect(result.newContent).toBe(content);
    });

    it('should perform multiple fuzzy replacements if multiple valid matches are found', async () => {
      const content = `
function doIt() {
  console.log("hello");
}

function doIt() {
  console.log("hello");
}
`;
      // old_string uses single quotes, file uses double.
      // This is a fuzzy match (quote difference).
      const oldString = `
function doIt() {
  console.log('hello');
}
`.trim();

      const newString = `
function doIt() {
  console.log("bye");
}
`.trim();

      const result = await calculateReplacement(mockConfig, {
        params: {
          file_path: 'test.ts',
          instruction: 'update',
          old_string: oldString,
          new_string: newString,
        },
        currentContent: content,
        abortSignal,
      });

      expect(result.occurrences).toBe(2);
      const expectedContent = `
function doIt() {
  console.log("bye");
}

function doIt() {
  console.log("bye");
}
`;
      expect(result.newContent).toBe(expectedContent);
    });

    it('should NOT insert extra newlines when replacing a block preceded by a blank line (regression)', async () => {
      const content = '\n  function oldFunc() {\n    // some code\n  }';
      const result = await calculateReplacement(mockConfig, {
        params: {
          file_path: 'test.js',
          instruction: 'test',
          old_string: 'function  oldFunc() {\n    // some code\n  }', // Two spaces after function to trigger regex
          new_string: 'function newFunc() {\n  // new code\n}', // Unindented
        },
        currentContent: content,
        abortSignal,
      });

      // The blank line at the start should be preserved as-is,
      // and the discovered indentation (2 spaces) should be applied to each line.
      const expectedContent = '\n  function newFunc() {\n    // new code\n  }';
      expect(result.newContent).toBe(expectedContent);
    });

    it('should NOT insert extra newlines in flexible replacement when old_string starts with a blank line (regression)', async () => {
      const content = '  // some comment\n\n  function oldFunc() {}';
      const result = await calculateReplacement(mockConfig, {
        params: {
          file_path: 'test.js',
          instruction: 'test',
          old_string: '\nfunction oldFunc() {}',
          new_string: '\n  function newFunc() {}', // Include desired indentation
        },
        currentContent: content,
        abortSignal,
      });

      // The blank line at the start is preserved, and the new block is inserted.
      const expectedContent = '  // some comment\n\n  function newFunc() {}';
      expect(result.newContent).toBe(expectedContent);
    });
  });

  describe('validateToolParams', () => {
    it('should return null for valid params', () => {
      const params: EditToolParams = {
        file_path: path.join(rootDir, 'test.txt'),
        instruction: 'An instruction',
        old_string: 'old',
        new_string: 'new',
      };
      expect(tool.validateToolParams(params)).toBeNull();
    });

    it('should return an error if path is outside the workspace', () => {
      const params: EditToolParams = {
        file_path: path.join(os.tmpdir(), 'outside.txt'),
        instruction: 'An instruction',
        old_string: 'old',
        new_string: 'new',
      };
      expect(tool.validateToolParams(params)).toMatch(/Path not in workspace/);
    });
  });

  describe('execute', () => {
    const testFile = 'execute_me.txt';
    let filePath: string;

    beforeEach(() => {
      filePath = path.join(rootDir, testFile);
    });

    it('should reject when calculateEdit fails after an abort signal', async () => {
      const params: EditToolParams = {
        file_path: path.join(rootDir, 'abort-execute.txt'),
        instruction: 'Abort during execute',
        old_string: 'old',
        new_string: 'new',
      };

      const invocation = tool.build(params);
      const abortController = new AbortController();
      const abortError = new Error('Abort requested during edit execution');

      const calculateSpy = vi
        .spyOn(invocation as any, 'calculateEdit')
        .mockImplementation(async () => {
          if (!abortController.signal.aborted) {
            abortController.abort();
          }
          throw abortError;
        });

      await expect(invocation.execute(abortController.signal)).rejects.toBe(
        abortError,
      );

      calculateSpy.mockRestore();
    });

    it('should edit an existing file and return diff with fileName', async () => {
      const initialContent = 'This is some old text.';
      const newContent = 'This is some new text.';
      fs.writeFileSync(filePath, initialContent, 'utf8');
      const params: EditToolParams = {
        file_path: filePath,
        instruction: 'Replace old with new',
        old_string: 'old',
        new_string: 'new',
      };

      const invocation = tool.build(params);
      const result = await invocation.execute(new AbortController().signal);

      expect(result.llmContent).toMatch(/Successfully modified file/);
      expect(fs.readFileSync(filePath, 'utf8')).toBe(newContent);
      const display = result.returnDisplay as FileDiff;
      expect(display.fileDiff).toMatch(initialContent);
      expect(display.fileDiff).toMatch(newContent);
      expect(display.fileName).toBe(testFile);
    });

    it('should return error if old_string is not found in file', async () => {
      fs.writeFileSync(filePath, 'Some content.', 'utf8');

      // Enable LLM correction for this test
      (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);
      const result = await invocation.execute(new AbortController().signal);
      expect(result.llmContent).toMatch(/0 occurrences found for old_string/);
      expect(result.returnDisplay).toMatch(
        /Failed to edit, could not find the string to replace./,
      );
      expect(mockFixLLMEditWithInstruction).toHaveBeenCalled();
    });

    it('should succeed if FixLLMEditWithInstruction corrects the params', async () => {
      const initialContent = 'This is some original text.';
      const finalContent = 'This is some brand new text.';
      fs.writeFileSync(filePath, initialContent, 'utf8');

      // Enable LLM correction for this test
      (mockConfig.getDisableLLMCorrection as Mock).mockReturnValue(false);

      const params: EditToolParams = {
        file_path: filePath,
        instruction: 'Replace original with brand new',
        old_string: 'wrong text', // This will fail first
        new_string: 'brand new text',
      };

      mockFixLLMEditWithInstruction.mockResolvedValueOnce({
        noChangesRequired: false,
        search: 'original text', // The corrected search string
        replace: 'brand new text',
        explanation: 'Corrected the search string to match the file content.',
      });

      const invocation = tool.build(params);
      const result = await invocation.execute(new AbortController().signal);

      expect(result.error).toBeUndefined();
      expect(result.llmContent).toMatch(/Successfully modified file/);
      expect(fs.readFileSync(filePath, 'utf8')).toBe(finalContent);
      expect(mockFixLLMEditWithInstruction).toHaveBeenCalledTimes(1);
    });

    it('should preserve CRLF line endings when editing a file', async () => {
      const initialContent = 'line one\r\nline two\r\n';
      const newContent = 'line one\r\nline three\r\n';
      fs.writeFileSync(filePath, initialContent, 'utf8');
      const params: EditToolParams = {
        file_path: filePath,
        instruction: 'Replace two with three',
        old_string: 'line two',
        new_string: 'line three',
      };

      const invocation = tool.build(params);
      await invocation.execute(new AbortController().signal);

      const finalContent = fs.readFileSync(filePath, 'utf8');
      expect(finalContent).toBe(newContent);
    });

    it('should create a new file with CRLF line endings if new_string has them', async () => {
      const newContentWithCRLF = 'new line one\r\nnew line two\r\n';
      const params: EditToolParams = {
        file_path: filePath,
        instruction: 'Create a new file',
        old_string: '',
        new_string: newContentWithCRLF,
      };

      const invocation = tool.build(params);
      await invocation.execute(new AbortController().signal);

      const finalContent = fs.readFileSync(filePath, 'utf8');
      expect(finalContent).toBe(newContentWithCRLF);
    });

    it('should return NO_CHANGE if FixLLMEditWithInstruction determines no changes are needed', async () => {
      const initialContent = 'The price is $100.';
      fs.writeFileSync(filePath, initialContent, 'utf8');

      // Enable LLM correction for this test
      (mockConfig.getDisableLLMCorrection as Mock).mockReturnValue(false);

      const params: EditToolParams = {
        file_path: filePath,
        instruction: 'Ensure the price is $100',
        old_string: 'price is $50', // Incorrect old string
        new_string: 'price is $100',
      };

      mockFixLLMEditWithInstruction.mockResolvedValueOnce({
        noChangesRequired: true,
        search: '',
        replace: '',
        explanation: 'The price is already correctly set to $100.',
      });

      const invocation = tool.build(params);
      const result = await invocation.execute(new AbortController().signal);

      expect(result.error?.type).toBe(
        ToolErrorType.EDIT_NO_CHANGE_LLM_JUDGEMENT,
      );
      expect(result.llmContent).toMatch(
        /A secondary check by an LLM determined/,
      );
      expect(fs.readFileSync(filePath, 'utf8')).toBe(initialContent); // File is unchanged
    });
  });

  describe('self-correction with content refresh to pull in external edits', () => {
    const testFile = 'test.txt';
    let filePath: string;

    beforeEach(() => {
      filePath = path.join(rootDir, testFile);
    });

    it('should use refreshed file content for self-correction if file was modified externally', async () => {
      const initialContent = 'This is the original content.';
      const externallyModifiedContent =
        'This is the externally modified content.';
      fs.writeFileSync(filePath, initialContent, 'utf8');

      // Enable LLM correction for this test
      (mockConfig.getDisableLLMCorrection as Mock).mockReturnValue(false);

      const params: EditToolParams = {
        file_path: filePath,
        instruction:
          'Replace "externally modified content" with "externally modified string"',
        old_string: 'externally modified content', // This will fail the first attempt, triggering self-correction.
        new_string: 'externally modified string',
      };

      // Spy on `readTextFile` to simulate an external file change between reads.
      const readTextFileSpy = vi
        .spyOn(fileSystemService, 'readTextFile')
        .mockResolvedValueOnce(initialContent) // First call in `calculateEdit`
        .mockResolvedValueOnce(externallyModifiedContent); // Second call in `attemptSelfCorrection`

      const invocation = tool.build(params);
      await invocation.execute(new AbortController().signal);

      // Assert that the file was read twice (initial read, then re-read for hash comparison).
      expect(readTextFileSpy).toHaveBeenCalledTimes(2);

      // Assert that the self-correction LLM was called with the updated content and a specific message.
      expect(mockFixLLMEditWithInstruction).toHaveBeenCalledWith(
        expect.any(String), // instruction
        params.old_string,
        params.new_string,
        expect.stringContaining(
          'However, the file has been modified by either the user or an external process',
        ), // errorForLlmEditFixer
        externallyModifiedContent, // The new content for correction
        expect.any(Object), // baseLlmClient
        expect.any(Object), // abortSignal
      );
    });
  });

  describe('Error Scenarios', () => {
    const testFile = 'error_test.txt';
    let filePath: string;

    beforeEach(() => {
      filePath = path.join(rootDir, testFile);
    });

    it.each([
      {
        name: 'FILE_NOT_FOUND',
        setup: () => {}, // no file created
        params: { old_string: 'any', new_string: 'new' },
        expectedError: ToolErrorType.FILE_NOT_FOUND,
      },
      {
        name: 'ATTEMPT_TO_CREATE_EXISTING_FILE',
        setup: (fp: string) => fs.writeFileSync(fp, 'existing content', 'utf8'),
        params: { old_string: '', new_string: 'new content' },
        expectedError: ToolErrorType.ATTEMPT_TO_CREATE_EXISTING_FILE,
      },
      {
        name: 'NO_OCCURRENCE_FOUND',
        setup: (fp: string) => fs.writeFileSync(fp, 'content', 'utf8'),
        params: { old_string: 'not-found', new_string: 'new' },
        expectedError: ToolErrorType.EDIT_NO_OCCURRENCE_FOUND,
      },
      {
        name: 'EXPECTED_OCCURRENCE_MISMATCH',
        setup: (fp: string) => fs.writeFileSync(fp, 'one one two', 'utf8'),
        params: { old_string: 'one', new_string: 'new' },
        expectedError: ToolErrorType.EDIT_EXPECTED_OCCURRENCE_MISMATCH,
      },
    ])(
      'should return $name error',
      async ({ setup, params, expectedError }) => {
        setup(filePath);
        const invocation = tool.build({
          file_path: filePath,
          instruction: 'test',
          ...params,
        });
        const result = await invocation.execute(new AbortController().signal);
        expect(result.error?.type).toBe(expectedError);
      },
    );
  });

  describe('expected_replacements', () => {
    const testFile = 'replacements_test.txt';
    let filePath: string;

    beforeEach(() => {
      filePath = path.join(rootDir, testFile);
    });

    it.each([
      {
        name: 'succeed when occurrences match expected_replacements',
        content: 'foo foo foo',
        expected: 3,
        shouldSucceed: true,
        finalContent: 'bar bar bar',
      },
      {
        name: 'fail when occurrences do not match expected_replacements',
        content: 'foo foo foo',
        expected: 2,
        shouldSucceed: false,
      },
      {
        name: 'default to 1 expected replacement if not specified',
        content: 'foo foo',
        expected: undefined,
        shouldSucceed: false,
      },
    ])(
      'should $name',
      async ({ content, expected, shouldSucceed, finalContent }) => {
        fs.writeFileSync(filePath, content, 'utf8');
        const params: EditToolParams = {
          file_path: filePath,
          instruction: 'Replace all foo with bar',
          old_string: 'foo',
          new_string: 'bar',
          ...(expected !== undefined && { expected_replacements: expected }),
        };
        const invocation = tool.build(params);
        const result = await invocation.execute(new AbortController().signal);

        if (shouldSucceed) {
          expect(result.error).toBeUndefined();
          if (finalContent)
            expect(fs.readFileSync(filePath, 'utf8')).toBe(finalContent);
        } else {
          expect(result.error?.type).toBe(
            ToolErrorType.EDIT_EXPECTED_OCCURRENCE_MISMATCH,
          );
        }
      },
    );
  });

  describe('IDE mode', () => {
    const testFile = 'edit_me.txt';
    let filePath: string;
    let ideClient: any;

    beforeEach(() => {
      filePath = path.join(rootDir, testFile);
      ideClient = {
        openDiff: vi.fn(),
        isDiffingEnabled: vi.fn().mockReturnValue(true),
      };
      vi.mocked(IdeClient.getInstance).mockResolvedValue(ideClient);
      (mockConfig as any).getIdeMode = () => true;
    });

    it('should call ideClient.openDiff and update params on confirmation', async () => {
      const initialContent = 'some old content here';
      const newContent = 'some new content here';
      const modifiedContent = 'some modified content here';
      fs.writeFileSync(filePath, initialContent);
      const params: EditToolParams = {
        file_path: filePath,
        instruction: 'test',
        old_string: 'old',
        new_string: 'new',
      };

      ideClient.openDiff.mockResolvedValueOnce({
        status: 'accepted',
        content: modifiedContent,
      });

      const invocation = tool.build(params);
      const confirmation = await invocation.shouldConfirmExecute(
        new AbortController().signal,
      );

      expect(ideClient.openDiff).toHaveBeenCalledWith(filePath, newContent);

      if (confirmation && 'onConfirm' in confirmation) {
        await confirmation.onConfirm(ToolConfirmationOutcome.ProceedOnce);
      }

      expect(params.old_string).toBe(initialContent);
      expect(params.new_string).toBe(modifiedContent);
    });
  });

  describe('shouldConfirmExecute', () => {
    it('should rethrow calculateEdit errors when the abort signal is triggered', async () => {
      const filePath = path.join(rootDir, 'abort-confirmation.txt');
      const params: EditToolParams = {
        file_path: filePath,
        instruction: 'Abort during confirmation',
        old_string: 'old',
        new_string: 'new',
      };

      const invocation = tool.build(params);
      const abortController = new AbortController();
      const abortError = new Error('Abort requested during edit confirmation');

      const calculateSpy = vi
        .spyOn(invocation as any, 'calculateEdit')
        .mockImplementation(async () => {
          if (!abortController.signal.aborted) {
            abortController.abort();
          }
          throw abortError;
        });

      await expect(
        invocation.shouldConfirmExecute(abortController.signal),
      ).rejects.toBe(abortError);

      calculateSpy.mockRestore();
    });
  });

  describe('multiple file edits', () => {
    it('should perform multiple removals and report correct diff stats', async () => {
      const numFiles = 10;
      const files: Array<{
        path: string;
        initialContent: string;
        toRemove: string;
      }> = [];
      const expectedLinesRemoved: number[] = [];
      const actualLinesRemoved: number[] = [];

      // 1. Create 10 files with 5-10 lines each
      for (let i = 0; i < numFiles; i++) {
        const fileName = `test-file-${i}.txt`;
        const filePath = path.join(rootDir, fileName);
        const numLines = Math.floor(Math.random() * 6) + 5; // 5 to 10 lines
        const lines = Array.from(
          { length: numLines },
          (_, j) => `File ${i}, Line ${j + 1}`,
        );
        const content = lines.join('\n') + '\n';

        // Determine which lines to remove (2 or 3 lines)
        const numLinesToRemove = Math.floor(Math.random() * 2) + 2; // 2 or 3
        expectedLinesRemoved.push(numLinesToRemove);
        const startLineToRemove = 1; // Start removing from the second line
        const linesToRemove = lines.slice(
          startLineToRemove,
          startLineToRemove + numLinesToRemove,
        );
        const toRemove = linesToRemove.join('\n') + '\n';

        fs.writeFileSync(filePath, content, 'utf8');
        files.push({
          path: filePath,
          initialContent: content,
          toRemove,
        });
      }

      // 2. Create and execute 10 tool calls for removal
      for (const file of files) {
        const params: EditToolParams = {
          file_path: file.path,
          instruction: `Remove lines from the file`,
          old_string: file.toRemove,
          new_string: '', // Removing the content
          ai_proposed_content: '',
        };
        const invocation = tool.build(params);
        const result = await invocation.execute(new AbortController().signal);

        if (
          result.returnDisplay &&
          typeof result.returnDisplay === 'object' &&
          'diffStat' in result.returnDisplay &&
          result.returnDisplay.diffStat
        ) {
          actualLinesRemoved.push(
            result.returnDisplay.diffStat?.model_removed_lines,
          );
        } else if (result.error) {
          throw result.error;
        }
      }

      // 3. Assert that the content was removed from each file
      for (const file of files) {
        const finalContent = fs.readFileSync(file.path, 'utf8');
        const expectedContent = file.initialContent.replace(file.toRemove, '');
        expect(finalContent).toBe(expectedContent);
        expect(finalContent).not.toContain(file.toRemove);
      }

      // 4. Assert that the total number of removed lines matches the diffStat total
      const totalExpectedRemoved = expectedLinesRemoved.reduce(
        (sum, current) => sum + current,
        0,
      );
      const totalActualRemoved = actualLinesRemoved.reduce(
        (sum, current) => sum + current,
        0,
      );
      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', async () => {
      const filePath = path.join(rootDir, 'enable_llm_test.txt');
      fs.writeFileSync(filePath, 'Some content.', 'utf8');

      // Now explicit as it's not the default anymore
      (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();
    });
  });
});