fix: Remove unreliable editCorrector that injects extra escape characters (#713)

This commit is contained in:
tanzhenxin
2025-09-25 16:46:58 +08:00
committed by GitHub
parent 4e7a7e2656
commit 673854b446
6 changed files with 56 additions and 1935 deletions

View File

@@ -6,22 +6,10 @@
/* eslint-disable @typescript-eslint/no-explicit-any */
const mockEnsureCorrectEdit = vi.hoisted(() => vi.fn());
const mockGenerateJson = vi.hoisted(() => vi.fn());
const mockOpenDiff = vi.hoisted(() => vi.fn());
import { IDEConnectionStatus } from '../ide/ide-client.js';
vi.mock('../utils/editCorrector.js', () => ({
ensureCorrectEdit: mockEnsureCorrectEdit,
}));
vi.mock('../core/client.js', () => ({
GeminiClient: vi.fn().mockImplementation(() => ({
generateJson: mockGenerateJson,
})),
}));
vi.mock('../utils/editor.js', () => ({
openDiff: mockOpenDiff,
}));
@@ -42,7 +30,6 @@ import fs from 'node:fs';
import os from 'node:os';
import type { Config } from '../config/config.js';
import { ApprovalMode } from '../config/config.js';
import type { Content, Part, SchemaUnion } from '@google/genai';
import { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.js';
import { StandardFileSystemService } from '../services/fileSystemService.js';
@@ -51,20 +38,13 @@ describe('EditTool', () => {
let tempDir: string;
let rootDir: string;
let mockConfig: Config;
let geminiClient: any;
beforeEach(() => {
vi.restoreAllMocks();
tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'edit-tool-test-'));
rootDir = path.join(tempDir, 'root');
fs.mkdirSync(rootDir);
geminiClient = {
generateJson: mockGenerateJson, // mockGenerateJson is already defined and hoisted
};
mockConfig = {
getGeminiClient: vi.fn().mockReturnValue(geminiClient),
getTargetDir: () => rootDir,
getApprovalMode: vi.fn(),
setApprovalMode: vi.fn(),
@@ -72,9 +52,6 @@ describe('EditTool', () => {
getFileSystemService: () => new StandardFileSystemService(),
getIdeClient: () => undefined,
getIdeMode: () => false,
// getGeminiConfig: () => ({ apiKey: 'test-api-key' }), // This was not a real Config method
// Add other properties/methods of Config if EditTool uses them
// Minimal other methods to satisfy Config type if needed by EditTool constructor or other direct uses:
getApiKey: () => 'test-api-key',
getModel: () => 'test-model',
getSandbox: () => false,
@@ -98,65 +75,6 @@ describe('EditTool', () => {
// Default to not skipping confirmation
(mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.DEFAULT);
// Reset mocks and set default implementation for ensureCorrectEdit
mockEnsureCorrectEdit.mockReset();
mockEnsureCorrectEdit.mockImplementation(
async (_, currentContent, params) => {
let occurrences = 0;
if (params.old_string && currentContent) {
// Simple string counting for the mock
let index = currentContent.indexOf(params.old_string);
while (index !== -1) {
occurrences++;
index = currentContent.indexOf(params.old_string, index + 1);
}
} else if (params.old_string === '') {
occurrences = 0; // Creating a new file
}
return Promise.resolve({ params, occurrences });
},
);
// Default mock for generateJson to return the snippet unchanged
mockGenerateJson.mockReset();
mockGenerateJson.mockImplementation(
async (contents: Content[], schema: SchemaUnion) => {
// The problematic_snippet is the last part of the user's content
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 as any)?.corrected_target_snippet) {
return Promise.resolve({
corrected_target_snippet: problematicSnippet,
});
}
if (((schema as any).properties as any)?.corrected_new_string) {
// For new_string correction, we might need more sophisticated logic,
// but for now, returning original is a safe default if not specified by a test.
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({}); // Default empty object if schema doesn't match
},
);
tool = new EditTool(mockConfig);
});
@@ -249,8 +167,6 @@ describe('EditTool', () => {
old_string: 'old',
new_string: 'new',
};
// ensureCorrectEdit will be called by shouldConfirmExecute
mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 1 });
const invocation = tool.build(params);
const confirmation = await invocation.shouldConfirmExecute(
new AbortController().signal,
@@ -264,14 +180,13 @@ describe('EditTool', () => {
);
});
it('should return false if old_string is not found (ensureCorrectEdit returns 0)', async () => {
it('should return false if old_string is not found', async () => {
fs.writeFileSync(filePath, 'some content here');
const params: EditToolParams = {
file_path: filePath,
old_string: 'not_found',
new_string: 'new',
};
mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 0 });
const invocation = tool.build(params);
const confirmation = await invocation.shouldConfirmExecute(
new AbortController().signal,
@@ -279,14 +194,13 @@ describe('EditTool', () => {
expect(confirmation).toBe(false);
});
it('should return false if multiple occurrences of old_string are found (ensureCorrectEdit returns > 1)', async () => {
it('should return false if multiple occurrences of old_string are found', async () => {
fs.writeFileSync(filePath, 'old old content here');
const params: EditToolParams = {
file_path: filePath,
old_string: 'old',
new_string: 'new',
};
mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 2 });
const invocation = tool.build(params);
const confirmation = await invocation.shouldConfirmExecute(
new AbortController().signal,
@@ -302,10 +216,6 @@ describe('EditTool', () => {
old_string: '',
new_string: 'new file content',
};
// ensureCorrectEdit might not be called if old_string is empty,
// as shouldConfirmExecute handles this for diff generation.
// If it is called, it should return 0 occurrences for a new file.
mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 0 });
const invocation = tool.build(params);
const confirmation = await invocation.shouldConfirmExecute(
new AbortController().signal,
@@ -319,65 +229,9 @@ describe('EditTool', () => {
);
});
it('should use corrected params from ensureCorrectEdit for diff generation', async () => {
const originalContent = 'This is the original string to be replaced.';
const originalOldString = 'original string';
const originalNewString = 'new string';
const correctedOldString = 'original string to be replaced'; // More specific
const correctedNewString = 'completely new string'; // Different replacement
const expectedFinalContent = 'This is the completely new string.';
fs.writeFileSync(filePath, originalContent);
const params: EditToolParams = {
file_path: filePath,
old_string: originalOldString,
new_string: originalNewString,
};
// The main beforeEach already calls mockEnsureCorrectEdit.mockReset()
// Set a specific mock for this test case
let mockCalled = false;
mockEnsureCorrectEdit.mockImplementationOnce(
async (_, content, p, client) => {
mockCalled = true;
expect(content).toBe(originalContent);
expect(p).toBe(params);
expect(client).toBe(geminiClient);
return {
params: {
file_path: filePath,
old_string: correctedOldString,
new_string: correctedNewString,
},
occurrences: 1,
};
},
);
const invocation = tool.build(params);
const confirmation = (await invocation.shouldConfirmExecute(
new AbortController().signal,
)) as FileDiff;
expect(mockCalled).toBe(true); // Check if the mock implementation was run
// expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(originalContent, params, expect.anything()); // Keep this commented for now
expect(confirmation).toEqual(
expect.objectContaining({
title: `Confirm Edit: ${testFile}`,
fileName: testFile,
}),
);
// Check that the diff is based on the corrected strings leading to the new state
expect(confirmation.fileDiff).toContain(`-${originalContent}`);
expect(confirmation.fileDiff).toContain(`+${expectedFinalContent}`);
// Verify that applying the correctedOldString and correctedNewString to originalContent
// indeed produces the expectedFinalContent, which is what the diff should reflect.
const patchedContent = originalContent.replace(
correctedOldString, // This was the string identified by ensureCorrectEdit for replacement
correctedNewString, // This was the string identified by ensureCorrectEdit as the replacement
);
expect(patchedContent).toBe(expectedFinalContent);
// This test is no longer relevant since editCorrector functionality was removed
it.skip('should use corrected params from ensureCorrectEdit for diff generation', async () => {
// Test skipped - editCorrector functionality removed
});
});
@@ -387,20 +241,6 @@ describe('EditTool', () => {
beforeEach(() => {
filePath = path.join(rootDir, testFile);
// Default for execute tests, can be overridden
mockEnsureCorrectEdit.mockImplementation(async (_, content, params) => {
let occurrences = 0;
if (params.old_string && content) {
let index = content.indexOf(params.old_string);
while (index !== -1) {
occurrences++;
index = content.indexOf(params.old_string, index + 1);
}
} else if (params.old_string === '') {
occurrences = 0;
}
return { params, occurrences };
});
});
it('should throw error if file path is not absolute', async () => {
@@ -433,10 +273,6 @@ describe('EditTool', () => {
new_string: 'new',
};
// Specific mock for this test's execution path in calculateEdit
// ensureCorrectEdit is NOT called by calculateEdit, only by shouldConfirmExecute
// So, the default mockEnsureCorrectEdit should correctly return 1 occurrence for 'old' in initialContent
const invocation = tool.build(params);
const result = await invocation.execute(new AbortController().signal);
@@ -477,7 +313,6 @@ describe('EditTool', () => {
old_string: 'nonexistent',
new_string: 'replacement',
};
// The default mockEnsureCorrectEdit will return 0 occurrences for 'nonexistent'
const invocation = tool.build(params);
const result = await invocation.execute(new AbortController().signal);
expect(result.llmContent).toMatch(
@@ -495,7 +330,6 @@ describe('EditTool', () => {
old_string: 'old',
new_string: 'new',
};
// The default mockEnsureCorrectEdit will return 2 occurrences for 'old'
const invocation = tool.build(params);
const result = await invocation.execute(new AbortController().signal);
expect(result.llmContent).toMatch(
@@ -638,8 +472,7 @@ describe('EditTool', () => {
});
it('should return EDIT_NO_CHANGE error if replacement results in identical content', async () => {
// This can happen if ensureCorrectEdit finds a fuzzy match, but the literal
// string replacement with `replaceAll` results in no change.
// This can happen if the literal string replacement with `replaceAll` results in no change.
const initialContent = 'line 1\nline 2\nline 3'; // Note the double space
fs.writeFileSync(filePath, initialContent, 'utf8');
const params: EditToolParams = {
@@ -649,16 +482,12 @@ describe('EditTool', () => {
new_string: 'line 1\nnew line 2\nline 3',
};
// Mock ensureCorrectEdit to simulate it finding a match (e.g., via fuzzy matching)
// but it doesn't correct the old_string to the literal content.
mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 1 });
const invocation = tool.build(params);
const result = await invocation.execute(new AbortController().signal);
expect(result.error?.type).toBe(ToolErrorType.EDIT_NO_CHANGE);
expect(result.error?.type).toBe(ToolErrorType.EDIT_NO_OCCURRENCE_FOUND);
expect(result.returnDisplay).toMatch(
/No changes to apply. The new content is identical to the current content./,
/Failed to edit, could not find the string to replace./,
);
// Ensure the file was not actually changed
expect(fs.readFileSync(filePath, 'utf8')).toBe(initialContent);
@@ -870,10 +699,6 @@ describe('EditTool', () => {
old_string: 'old',
new_string: 'new',
};
mockEnsureCorrectEdit.mockResolvedValueOnce({
params: { ...params, old_string: 'old', new_string: 'new' },
occurrences: 1,
});
ideClient.openDiff.mockResolvedValueOnce({
status: 'accepted',
content: modifiedContent,

View File

@@ -21,7 +21,6 @@ import { makeRelative, shortenPath } from '../utils/paths.js';
import { isNodeError } from '../utils/errors.js';
import type { Config } from '../config/config.js';
import { ApprovalMode } from '../config/config.js';
import { ensureCorrectEdit } from '../utils/editCorrector.js';
import { DEFAULT_DIFF_OPTIONS, getDiffStat } from './diffOptions.js';
import { ReadFileTool } from './read-file.js';
import { ToolNames } from './tool-names.js';
@@ -116,16 +115,13 @@ class EditToolInvocation implements ToolInvocation<EditToolParams, ToolResult> {
* @returns An object describing the potential edit outcome
* @throws File system errors if reading the file fails unexpectedly (e.g., permissions)
*/
private async calculateEdit(
params: EditToolParams,
abortSignal: AbortSignal,
): Promise<CalculatedEdit> {
private async calculateEdit(params: EditToolParams): Promise<CalculatedEdit> {
const expectedReplacements = params.expected_replacements ?? 1;
let currentContent: string | null = null;
let fileExists = false;
let isNewFile = false;
let finalNewString = params.new_string;
let finalOldString = params.old_string;
const finalNewString = params.new_string;
const finalOldString = params.old_string;
let occurrences = 0;
let error:
| { display: string; raw: string; type: ToolErrorType }
@@ -157,18 +153,7 @@ class EditToolInvocation implements ToolInvocation<EditToolParams, ToolResult> {
type: ToolErrorType.FILE_NOT_FOUND,
};
} else if (currentContent !== null) {
// Editing an existing file
const correctedEdit = await ensureCorrectEdit(
params.file_path,
currentContent,
params,
this.config.getGeminiClient(),
abortSignal,
);
finalOldString = correctedEdit.params.old_string;
finalNewString = correctedEdit.params.new_string;
occurrences = correctedEdit.occurrences;
occurrences = this.countOccurrences(currentContent, params.old_string);
if (params.old_string === '') {
// Error: Trying to create a file that already exists
error = {
@@ -234,12 +219,28 @@ class EditToolInvocation implements ToolInvocation<EditToolParams, ToolResult> {
};
}
/**
* Counts occurrences of a substring in a string
*/
private countOccurrences(str: string, substr: string): number {
if (substr === '') {
return 0;
}
let count = 0;
let pos = str.indexOf(substr);
while (pos !== -1) {
count++;
pos = str.indexOf(substr, pos + substr.length); // Start search after the current match
}
return count;
}
/**
* Handles the confirmation prompt for the Edit tool in the CLI.
* It needs to calculate the diff to show the user.
*/
async shouldConfirmExecute(
abortSignal: AbortSignal,
_abortSignal: AbortSignal,
): Promise<ToolCallConfirmationDetails | false> {
if (this.config.getApprovalMode() === ApprovalMode.AUTO_EDIT) {
return false;
@@ -247,7 +248,7 @@ class EditToolInvocation implements ToolInvocation<EditToolParams, ToolResult> {
let editData: CalculatedEdit;
try {
editData = await this.calculateEdit(this.params, abortSignal);
editData = await this.calculateEdit(this.params);
} catch (error) {
const errorMsg = error instanceof Error ? error.message : String(error);
console.log(`Error preparing edit: ${errorMsg}`);
@@ -330,10 +331,10 @@ class EditToolInvocation implements ToolInvocation<EditToolParams, ToolResult> {
* @param params Parameters for the edit operation
* @returns Result of the edit operation
*/
async execute(signal: AbortSignal): Promise<ToolResult> {
async execute(_signal: AbortSignal): Promise<ToolResult> {
let editData: CalculatedEdit;
try {
editData = await this.calculateEdit(this.params, signal);
editData = await this.calculateEdit(this.params);
} catch (error) {
const errorMsg = error instanceof Error ? error.message : String(error);
return {

View File

@@ -18,7 +18,6 @@ import { getCorrectedFileContent, WriteFileTool } from './write-file.js';
import { ToolErrorType } from './tool-error.js';
import type { FileDiff, ToolEditConfirmationDetails } from './tools.js';
import { ToolConfirmationOutcome } from './tools.js';
import { type EditToolParams } from './edit.js';
import type { Config } from '../config/config.js';
import { ApprovalMode } from '../config/config.js';
import type { ToolRegistry } from './tool-registry.js';
@@ -26,11 +25,6 @@ import path from 'node:path';
import fs from 'node:fs';
import os from 'node:os';
import { GeminiClient } from '../core/client.js';
import type { CorrectedEditResult } from '../utils/editCorrector.js';
import {
ensureCorrectEdit,
ensureCorrectFileContent,
} from '../utils/editCorrector.js';
import { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.js';
import { StandardFileSystemService } from '../services/fileSystemService.js';
@@ -38,17 +32,8 @@ const rootDir = path.resolve(os.tmpdir(), 'qwen-code-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 fsService = new StandardFileSystemService();
@@ -111,11 +96,6 @@ describe('WriteFileTool', () => {
) as Mocked<GeminiClient>;
vi.mocked(GeminiClient).mockImplementation(() => mockGeminiClientInstance);
vi.mocked(ensureCorrectEdit).mockImplementation(mockEnsureCorrectEdit);
vi.mocked(ensureCorrectFileContent).mockImplementation(
mockEnsureCorrectFileContent,
);
// Now that mockGeminiClientInstance is initialized, set the mock implementation for getGeminiClient
mockConfigInternal.getGeminiClient.mockReturnValue(
mockGeminiClientInstance,
@@ -134,40 +114,6 @@ describe('WriteFileTool', () => {
// Reset mocks before each test
mockConfigInternal.getApprovalMode.mockReturnValue(ApprovalMode.DEFAULT);
mockConfigInternal.setApprovalMode.mockClear();
mockEnsureCorrectEdit.mockReset();
mockEnsureCorrectFileContent.mockReset();
// Default mock implementations that return valid structures
mockEnsureCorrectEdit.mockImplementation(
async (
filePath: string,
_currentContent: string,
params: EditToolParams,
_client: GeminiClient,
signal?: AbortSignal, // Make AbortSignal optional to match usage
): Promise<CorrectedEditResult> => {
if (signal?.aborted) {
return Promise.reject(new Error('Aborted'));
}
return Promise.resolve({
params: { ...params, new_string: params.new_string ?? '' },
occurrences: 1,
});
},
);
mockEnsureCorrectFileContent.mockImplementation(
async (
content: string,
_client: GeminiClient,
signal?: AbortSignal,
): Promise<string> => {
// Make AbortSignal optional
if (signal?.aborted) {
return Promise.reject(new Error('Aborted'));
}
return Promise.resolve(content ?? '');
},
);
});
afterEach(() => {
@@ -240,71 +186,35 @@ describe('WriteFileTool', () => {
});
describe('getCorrectedFileContent', () => {
it('should call ensureCorrectFileContent for a new file', async () => {
it('should return proposed content unchanged for a new file', async () => {
const filePath = path.join(rootDir, 'new_corrected_file.txt');
const proposedContent = 'Proposed new content.';
const correctedContent = 'Corrected new content.';
const abortSignal = new AbortController().signal;
// Ensure the mock is set for this specific test case if needed, or rely on beforeEach
mockEnsureCorrectFileContent.mockResolvedValue(correctedContent);
const result = await getCorrectedFileContent(
mockConfig,
filePath,
proposedContent,
abortSignal,
);
expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith(
proposedContent,
mockGeminiClientInstance,
abortSignal,
);
expect(mockEnsureCorrectEdit).not.toHaveBeenCalled();
expect(result.correctedContent).toBe(correctedContent);
expect(result.correctedContent).toBe(proposedContent);
expect(result.originalContent).toBe('');
expect(result.fileExists).toBe(false);
expect(result.error).toBeUndefined();
});
it('should call ensureCorrectEdit for an existing file', async () => {
it('should return proposed content unchanged 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.';
const abortSignal = new AbortController().signal;
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);
const result = await getCorrectedFileContent(
mockConfig,
filePath,
proposedContent,
abortSignal,
);
expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
filePath,
originalContent,
{
old_string: originalContent,
new_string: proposedContent,
file_path: filePath,
},
mockGeminiClientInstance,
abortSignal,
);
expect(mockEnsureCorrectFileContent).not.toHaveBeenCalled();
expect(result.correctedContent).toBe(correctedProposedContent);
expect(result.correctedContent).toBe(proposedContent);
expect(result.originalContent).toBe(originalContent);
expect(result.fileExists).toBe(true);
expect(result.error).toBeUndefined();
@@ -313,7 +223,6 @@ describe('WriteFileTool', () => {
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';
const abortSignal = new AbortController().signal;
fs.writeFileSync(filePath, 'content', { mode: 0o000 });
const readError = new Error('Permission denied');
@@ -325,12 +234,9 @@ describe('WriteFileTool', () => {
mockConfig,
filePath,
proposedContent,
abortSignal,
);
expect(fsService.readTextFile).toHaveBeenCalledWith(filePath);
expect(mockEnsureCorrectEdit).not.toHaveBeenCalled();
expect(mockEnsureCorrectFileContent).not.toHaveBeenCalled();
expect(result.correctedContent).toBe(proposedContent);
expect(result.originalContent).toBe('');
expect(result.fileExists).toBe(true);
@@ -363,11 +269,9 @@ describe('WriteFileTool', () => {
fs.chmodSync(filePath, 0o600);
});
it('should request confirmation with diff for a new file (with corrected content)', async () => {
it('should request confirmation with diff for a new file', 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 invocation = tool.build(params);
@@ -375,16 +279,11 @@ describe('WriteFileTool', () => {
abortSignal,
)) as ToolEditConfirmationDetails;
expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith(
proposedContent,
mockGeminiClientInstance,
abortSignal,
);
expect(confirmation).toEqual(
expect.objectContaining({
title: `Confirm Write: ${path.basename(filePath)}`,
fileName: 'confirm_new_file.txt',
fileDiff: expect.stringContaining(correctedContent),
fileDiff: expect.stringContaining(proposedContent),
}),
);
expect(confirmation.fileDiff).toMatch(
@@ -395,45 +294,23 @@ describe('WriteFileTool', () => {
);
});
it('should request confirmation with diff for an existing file (with corrected content)', async () => {
it('should request confirmation with diff for an existing file', 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');
mockEnsureCorrectEdit.mockResolvedValue({
params: {
file_path: filePath,
old_string: originalContent,
new_string: correctedProposedContent,
},
occurrences: 1,
});
const params = { file_path: filePath, content: proposedContent };
const invocation = tool.build(params);
const confirmation = (await invocation.shouldConfirmExecute(
abortSignal,
)) as ToolEditConfirmationDetails;
expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
filePath,
originalContent,
{
old_string: originalContent,
new_string: proposedContent,
file_path: filePath,
},
mockGeminiClientInstance,
abortSignal,
);
expect(confirmation).toEqual(
expect.objectContaining({
title: `Confirm Write: ${path.basename(filePath)}`,
fileName: 'confirm_existing_file.txt',
fileDiff: expect.stringContaining(correctedProposedContent),
fileDiff: expect.stringContaining(proposedContent),
}),
);
expect(confirmation.fileDiff).toMatch(
@@ -470,11 +347,9 @@ describe('WriteFileTool', () => {
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');
it('should write a new file and return diff', async () => {
const filePath = path.join(rootDir, 'execute_new_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 invocation = tool.build(params);
@@ -490,49 +365,27 @@ describe('WriteFileTool', () => {
const result = await invocation.execute(abortSignal);
expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith(
proposedContent,
mockGeminiClientInstance,
abortSignal,
);
expect(result.llmContent).toMatch(
/Successfully created and wrote to new file/,
);
expect(fs.existsSync(filePath)).toBe(true);
const writtenContent = await fsService.readTextFile(filePath);
expect(writtenContent).toBe(correctedContent);
expect(writtenContent).toBe(proposedContent);
const display = result.returnDisplay as FileDiff;
expect(display.fileName).toBe('execute_new_corrected_file.txt');
expect(display.fileName).toBe('execute_new_file.txt');
expect(display.fileDiff).toMatch(/--- execute_new_file.txt\tOriginal/);
expect(display.fileDiff).toMatch(/\+\+\+ execute_new_file.txt\tWritten/);
expect(display.fileDiff).toMatch(
/--- execute_new_corrected_file.txt\tOriginal/,
);
expect(display.fileDiff).toMatch(
/\+\+\+ execute_new_corrected_file.txt\tWritten/,
);
expect(display.fileDiff).toMatch(
correctedContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'),
proposedContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'),
);
});
it('should overwrite an existing file with corrected content and return diff', async () => {
const filePath = path.join(
rootDir,
'execute_existing_corrected_file.txt',
);
it('should overwrite an existing file and return diff', async () => {
const filePath = path.join(rootDir, 'execute_existing_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');
mockEnsureCorrectEdit.mockResolvedValue({
params: {
file_path: filePath,
old_string: initialContent,
new_string: correctedProposedContent,
},
occurrences: 1,
});
const params = { file_path: filePath, content: proposedContent };
const invocation = tool.build(params);
@@ -547,27 +400,16 @@ describe('WriteFileTool', () => {
const result = await invocation.execute(abortSignal);
expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(
filePath,
initialContent,
{
old_string: initialContent,
new_string: proposedContent,
file_path: filePath,
},
mockGeminiClientInstance,
abortSignal,
);
expect(result.llmContent).toMatch(/Successfully overwrote file/);
const writtenContent = await fsService.readTextFile(filePath);
expect(writtenContent).toBe(correctedProposedContent);
expect(writtenContent).toBe(proposedContent);
const display = result.returnDisplay as FileDiff;
expect(display.fileName).toBe('execute_existing_corrected_file.txt');
expect(display.fileName).toBe('execute_existing_file.txt');
expect(display.fileDiff).toMatch(
initialContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'),
);
expect(display.fileDiff).toMatch(
correctedProposedContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'),
proposedContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'),
);
});
@@ -575,7 +417,6 @@ describe('WriteFileTool', () => {
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 };
const invocation = tool.build(params);
@@ -600,7 +441,6 @@ describe('WriteFileTool', () => {
it('should include modification message when proposed content is modified', async () => {
const filePath = path.join(rootDir, 'new_file_modified.txt');
const content = 'New file content modified by user';
mockEnsureCorrectFileContent.mockResolvedValue(content);
const params = {
file_path: filePath,
@@ -616,7 +456,6 @@ describe('WriteFileTool', () => {
it('should not include modification message when proposed content is not modified', async () => {
const filePath = path.join(rootDir, 'new_file_unmodified.txt');
const content = 'New file content not modified';
mockEnsureCorrectFileContent.mockResolvedValue(content);
const params = {
file_path: filePath,
@@ -632,7 +471,6 @@ describe('WriteFileTool', () => {
it('should not include modification message when modified_by_user is not provided', async () => {
const filePath = path.join(rootDir, 'new_file_unmodified.txt');
const content = 'New file content not modified';
mockEnsureCorrectFileContent.mockResolvedValue(content);
const params = {
file_path: filePath,

View File

@@ -26,10 +26,6 @@ import {
import { ToolErrorType } from './tool-error.js';
import { makeRelative, shortenPath } from '../utils/paths.js';
import { getErrorMessage, isNodeError } from '../utils/errors.js';
import {
ensureCorrectEdit,
ensureCorrectFileContent,
} from '../utils/editCorrector.js';
import { DEFAULT_DIFF_OPTIONS, getDiffStat } from './diffOptions.js';
import { ToolNames } from './tool-names.js';
import type {
@@ -79,11 +75,10 @@ export async function getCorrectedFileContent(
config: Config,
filePath: string,
proposedContent: string,
abortSignal: AbortSignal,
): Promise<GetCorrectedFileContentResult> {
let originalContent = '';
let fileExists = false;
let correctedContent = proposedContent;
const correctedContent = proposedContent;
try {
originalContent = await config
@@ -107,32 +102,6 @@ export async function getCorrectedFileContent(
}
}
// If readError is set, we have returned.
// So, file was either read successfully (fileExists=true, originalContent set)
// or it was ENOENT (fileExists=false, originalContent='').
if (fileExists) {
// This implies originalContent is available
const { params: correctedParams } = await ensureCorrectEdit(
filePath,
originalContent,
{
old_string: originalContent, // Treat entire current content as old_string
new_string: proposedContent,
file_path: filePath,
},
config.getGeminiClient(),
abortSignal,
);
correctedContent = correctedParams.new_string;
} else {
// This implies new file (ENOENT)
correctedContent = await ensureCorrectFileContent(
proposedContent,
config.getGeminiClient(),
abortSignal,
);
}
return { originalContent, correctedContent, fileExists };
}
@@ -160,7 +129,7 @@ class WriteFileToolInvocation extends BaseToolInvocation<
}
override async shouldConfirmExecute(
abortSignal: AbortSignal,
_abortSignal: AbortSignal,
): Promise<ToolCallConfirmationDetails | false> {
if (this.config.getApprovalMode() === ApprovalMode.AUTO_EDIT) {
return false;
@@ -170,7 +139,6 @@ class WriteFileToolInvocation extends BaseToolInvocation<
this.config,
this.params.file_path,
this.params.content,
abortSignal,
);
if (correctedContentResult.error) {
@@ -226,14 +194,13 @@ class WriteFileToolInvocation extends BaseToolInvocation<
return confirmationDetails;
}
async execute(abortSignal: AbortSignal): Promise<ToolResult> {
async execute(_abortSignal: AbortSignal): Promise<ToolResult> {
const { file_path, content, ai_proposed_content, modified_by_user } =
this.params;
const correctedContentResult = await getCorrectedFileContent(
this.config,
file_path,
content,
abortSignal,
);
if (correctedContentResult.error) {
@@ -476,7 +443,7 @@ export class WriteFileTool
}
getModifyContext(
abortSignal: AbortSignal,
_abortSignal: AbortSignal,
): ModifyContext<WriteFileToolParams> {
return {
getFilePath: (params: WriteFileToolParams) => params.file_path,
@@ -485,7 +452,6 @@ export class WriteFileTool
this.config,
params.file_path,
params.content,
abortSignal,
);
return correctedContentResult.originalContent;
},
@@ -494,7 +460,6 @@ export class WriteFileTool
this.config,
params.file_path,
params.content,
abortSignal,
);
return correctedContentResult.correctedContent;
},