abbyswag/qwen-code
1
0
Fork 0
mirror of https://github.com/QwenLM/qwen-code.git synced 2025-12-19 09:33:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

191
packages/core/src/tools/edit.test.ts
View File

@@ -6,22 +6,10 @@
/* eslint-disable @typescript-eslint/no-explicit-any */ /* 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()); const mockOpenDiff = vi.hoisted(() => vi.fn());
import { IDEConnectionStatus } from '../ide/ide-client.js'; 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', () => ({ vi.mock('../utils/editor.js', () => ({
openDiff: mockOpenDiff, openDiff: mockOpenDiff,
})); }));
@@ -42,7 +30,6 @@ import fs from 'node:fs';
import os from 'node:os'; import os from 'node:os';
import type { Config } from '../config/config.js'; import type { Config } from '../config/config.js';
import { ApprovalMode } 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 { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.js';
import { StandardFileSystemService } from '../services/fileSystemService.js'; import { StandardFileSystemService } from '../services/fileSystemService.js';
@@ -51,20 +38,13 @@ describe('EditTool', () => {
let tempDir: string; let tempDir: string;
let rootDir: string; let rootDir: string;
let mockConfig: Config; let mockConfig: Config;
let geminiClient: any;
beforeEach(() => { beforeEach(() => {
vi.restoreAllMocks(); vi.restoreAllMocks();
tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'edit-tool-test-')); tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'edit-tool-test-'));
rootDir = path.join(tempDir, 'root'); rootDir = path.join(tempDir, 'root');
fs.mkdirSync(rootDir); fs.mkdirSync(rootDir);
geminiClient = {
generateJson: mockGenerateJson, // mockGenerateJson is already defined and hoisted
};
mockConfig = { mockConfig = {
getGeminiClient: vi.fn().mockReturnValue(geminiClient),
getTargetDir: () => rootDir, getTargetDir: () => rootDir,
getApprovalMode: vi.fn(), getApprovalMode: vi.fn(),
setApprovalMode: vi.fn(), setApprovalMode: vi.fn(),
@@ -72,9 +52,6 @@ describe('EditTool', () => {
getFileSystemService: () => new StandardFileSystemService(), getFileSystemService: () => new StandardFileSystemService(),
getIdeClient: () => undefined, getIdeClient: () => undefined,
getIdeMode: () => false, 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', getApiKey: () => 'test-api-key',
getModel: () => 'test-model', getModel: () => 'test-model',
getSandbox: () => false, getSandbox: () => false,
@@ -98,65 +75,6 @@ describe('EditTool', () => {
// Default to not skipping confirmation // Default to not skipping confirmation
(mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.DEFAULT); (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); tool = new EditTool(mockConfig);
}); });
@@ -249,8 +167,6 @@ describe('EditTool', () => {
old_string: 'old', old_string: 'old',
new_string: 'new', new_string: 'new',
}; };
// ensureCorrectEdit will be called by shouldConfirmExecute
mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 1 });
const invocation = tool.build(params); const invocation = tool.build(params);
const confirmation = await invocation.shouldConfirmExecute( const confirmation = await invocation.shouldConfirmExecute(
new AbortController().signal, 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'); fs.writeFileSync(filePath, 'some content here');
const params: EditToolParams = { const params: EditToolParams = {
file_path: filePath, file_path: filePath,
old_string: 'not_found', old_string: 'not_found',
new_string: 'new', new_string: 'new',
}; };
mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 0 });
const invocation = tool.build(params); const invocation = tool.build(params);
const confirmation = await invocation.shouldConfirmExecute( const confirmation = await invocation.shouldConfirmExecute(
new AbortController().signal, new AbortController().signal,
@@ -279,14 +194,13 @@ describe('EditTool', () => {
expect(confirmation).toBe(false); 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'); fs.writeFileSync(filePath, 'old old content here');
const params: EditToolParams = { const params: EditToolParams = {
file_path: filePath, file_path: filePath,
old_string: 'old', old_string: 'old',
new_string: 'new', new_string: 'new',
}; };
mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 2 });
const invocation = tool.build(params); const invocation = tool.build(params);
const confirmation = await invocation.shouldConfirmExecute( const confirmation = await invocation.shouldConfirmExecute(
new AbortController().signal, new AbortController().signal,
@@ -302,10 +216,6 @@ describe('EditTool', () => {
old_string: '', old_string: '',
new_string: 'new file content', 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 invocation = tool.build(params);
const confirmation = await invocation.shouldConfirmExecute( const confirmation = await invocation.shouldConfirmExecute(
new AbortController().signal, new AbortController().signal,
@@ -319,65 +229,9 @@ describe('EditTool', () => {
); );
}); });
it('should use corrected params from ensureCorrectEdit for diff generation', async () => { // This test is no longer relevant since editCorrector functionality was removed
const originalContent = 'This is the original string to be replaced.'; it.skip('should use corrected params from ensureCorrectEdit for diff generation', async () => {
const originalOldString = 'original string'; // Test skipped - editCorrector functionality removed
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);
}); });
}); });
@@ -387,20 +241,6 @@ describe('EditTool', () => {
beforeEach(() => { beforeEach(() => {
filePath = path.join(rootDir, testFile); 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 () => { it('should throw error if file path is not absolute', async () => {
@@ -433,10 +273,6 @@ describe('EditTool', () => {
new_string: 'new', 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 invocation = tool.build(params);
const result = await invocation.execute(new AbortController().signal); const result = await invocation.execute(new AbortController().signal);
@@ -477,7 +313,6 @@ describe('EditTool', () => {
old_string: 'nonexistent', old_string: 'nonexistent',
new_string: 'replacement', new_string: 'replacement',
}; };
// The default mockEnsureCorrectEdit will return 0 occurrences for 'nonexistent'
const invocation = tool.build(params); const invocation = tool.build(params);
const result = await invocation.execute(new AbortController().signal); const result = await invocation.execute(new AbortController().signal);
expect(result.llmContent).toMatch( expect(result.llmContent).toMatch(
@@ -495,7 +330,6 @@ describe('EditTool', () => {
old_string: 'old', old_string: 'old',
new_string: 'new', new_string: 'new',
}; };
// The default mockEnsureCorrectEdit will return 2 occurrences for 'old'
const invocation = tool.build(params); const invocation = tool.build(params);
const result = await invocation.execute(new AbortController().signal); const result = await invocation.execute(new AbortController().signal);
expect(result.llmContent).toMatch( expect(result.llmContent).toMatch(
@@ -638,8 +472,7 @@ describe('EditTool', () => {
}); });
it('should return EDIT_NO_CHANGE error if replacement results in identical content', async () => { 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 // This can happen if the literal string replacement with `replaceAll` results in no change.
// string replacement with `replaceAll` results in no change.
const initialContent = 'line 1\nline 2\nline 3'; // Note the double space const initialContent = 'line 1\nline 2\nline 3'; // Note the double space
fs.writeFileSync(filePath, initialContent, 'utf8'); fs.writeFileSync(filePath, initialContent, 'utf8');
const params: EditToolParams = { const params: EditToolParams = {
@@ -649,16 +482,12 @@ describe('EditTool', () => {
new_string: 'line 1\nnew line 2\nline 3', 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 invocation = tool.build(params);
const result = await invocation.execute(new AbortController().signal); 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( 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 // Ensure the file was not actually changed
expect(fs.readFileSync(filePath, 'utf8')).toBe(initialContent); expect(fs.readFileSync(filePath, 'utf8')).toBe(initialContent);
@@ -870,10 +699,6 @@ describe('EditTool', () => {
old_string: 'old', old_string: 'old',
new_string: 'new', new_string: 'new',
}; };
mockEnsureCorrectEdit.mockResolvedValueOnce({
params: { ...params, old_string: 'old', new_string: 'new' },
occurrences: 1,
});
ideClient.openDiff.mockResolvedValueOnce({ ideClient.openDiff.mockResolvedValueOnce({
status: 'accepted', status: 'accepted',
content: modifiedContent, content: modifiedContent,

47
packages/core/src/tools/edit.ts
View File

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

202
packages/core/src/tools/write-file.test.ts
View File

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

43
packages/core/src/tools/write-file.ts
View File

@@ -26,10 +26,6 @@ import {
import { ToolErrorType } from './tool-error.js'; import { ToolErrorType } from './tool-error.js';
import { makeRelative, shortenPath } from '../utils/paths.js'; import { makeRelative, shortenPath } from '../utils/paths.js';
import { getErrorMessage, isNodeError } from '../utils/errors.js'; import { getErrorMessage, isNodeError } from '../utils/errors.js';
import {
ensureCorrectEdit,
ensureCorrectFileContent,
} from '../utils/editCorrector.js';
import { DEFAULT_DIFF_OPTIONS, getDiffStat } from './diffOptions.js'; import { DEFAULT_DIFF_OPTIONS, getDiffStat } from './diffOptions.js';
import { ToolNames } from './tool-names.js'; import { ToolNames } from './tool-names.js';
import type { import type {
@@ -79,11 +75,10 @@ export async function getCorrectedFileContent(
config: Config, config: Config,
filePath: string, filePath: string,
proposedContent: string, proposedContent: string,
abortSignal: AbortSignal,
): Promise<GetCorrectedFileContentResult> { ): Promise<GetCorrectedFileContentResult> {
let originalContent = ''; let originalContent = '';
let fileExists = false; let fileExists = false;
let correctedContent = proposedContent; const correctedContent = proposedContent;
try { try {
originalContent = await config 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 }; return { originalContent, correctedContent, fileExists };
} }
@@ -160,7 +129,7 @@ class WriteFileToolInvocation extends BaseToolInvocation<
} }
override async shouldConfirmExecute( override async shouldConfirmExecute(
abortSignal: AbortSignal, _abortSignal: AbortSignal,
): Promise<ToolCallConfirmationDetails | false> { ): Promise<ToolCallConfirmationDetails | false> {
if (this.config.getApprovalMode() === ApprovalMode.AUTO_EDIT) { if (this.config.getApprovalMode() === ApprovalMode.AUTO_EDIT) {
return false; return false;
@@ -170,7 +139,6 @@ class WriteFileToolInvocation extends BaseToolInvocation<
this.config, this.config,
this.params.file_path, this.params.file_path,
this.params.content, this.params.content,
abortSignal,
); );
if (correctedContentResult.error) { if (correctedContentResult.error) {
@@ -226,14 +194,13 @@ class WriteFileToolInvocation extends BaseToolInvocation<
return confirmationDetails; return confirmationDetails;
} }
async execute(abortSignal: AbortSignal): Promise<ToolResult> { async execute(_abortSignal: AbortSignal): Promise<ToolResult> {
const { file_path, content, ai_proposed_content, modified_by_user } = const { file_path, content, ai_proposed_content, modified_by_user } =
this.params; this.params;
const correctedContentResult = await getCorrectedFileContent( const correctedContentResult = await getCorrectedFileContent(
this.config, this.config,
file_path, file_path,
content, content,
abortSignal,
); );
if (correctedContentResult.error) { if (correctedContentResult.error) {
@@ -476,7 +443,7 @@ export class WriteFileTool
} }
getModifyContext( getModifyContext(
abortSignal: AbortSignal, _abortSignal: AbortSignal,
): ModifyContext<WriteFileToolParams> { ): ModifyContext<WriteFileToolParams> {
return { return {
getFilePath: (params: WriteFileToolParams) => params.file_path, getFilePath: (params: WriteFileToolParams) => params.file_path,
@@ -485,7 +452,6 @@ export class WriteFileTool
this.config, this.config,
params.file_path, params.file_path,
params.content, params.content,
abortSignal,
); );
return correctedContentResult.originalContent; return correctedContentResult.originalContent;
}, },
@@ -494,7 +460,6 @@ export class WriteFileTool
this.config, this.config,
params.file_path, params.file_path,
params.content, params.content,
abortSignal,
); );
return correctedContentResult.correctedContent; return correctedContentResult.correctedContent;
}, },

761
packages/core/src/utils/editCorrector.test.ts
View File

@@ -1,761 +0,0 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
/* eslint-disable @typescript-eslint/no-explicit-any */
import type { Mock } from 'vitest';
import { vi, describe, it, expect, beforeEach, type Mocked } from 'vitest';
import * as fs from 'node:fs';
import { EditTool } from '../tools/edit.js';
// MOCKS
let callCount = 0;
const mockResponses: any[] = [];
let mockGenerateJson: any;
let mockStartChat: any;
let mockSendMessageStream: any;
vi.mock('fs', () => ({
statSync: vi.fn(),
mkdirSync: vi.fn(),
}));
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,
ensureCorrectFileContent,
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);
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);
});
});
describe('unescapeStringForGeminiBug', () => {
it('should unescape common sequences', () => {
expect(unescapeStringForGeminiBug('\\n')).toBe('\n');
expect(unescapeStringForGeminiBug('\\t')).toBe('\t');
expect(unescapeStringForGeminiBug("\\'")).toBe("'");
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 (aggressive unescaping)', () => {
expect(unescapeStringForGeminiBug('\\\\n')).toBe('\n');
expect(unescapeStringForGeminiBug('\\\\\\t')).toBe('\t');
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');
expect(unescapeStringForGeminiBug('C:\\Folder\\File')).toBe(
'C:\\Folder\\File',
);
});
it('should correctly process strings with some targeted escapes', () => {
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\\"'),
).toBe('\nLine1\nLine2\tTab`Tick"');
});
it('should handle escaped backslashes', () => {
expect(unescapeStringForGeminiBug('\\\\')).toBe('\\');
expect(unescapeStringForGeminiBug('C:\\\\Users')).toBe('C:\\Users');
expect(unescapeStringForGeminiBug('path\\\\to\\\\file')).toBe(
'path\to\\file',
);
});
it('should handle escaped backslashes mixed with other escapes (aggressive unescaping)', () => {
expect(unescapeStringForGeminiBug('line1\\\\\\nline2')).toBe(
'line1\nline2',
);
expect(unescapeStringForGeminiBug('quote\\\\"text\\\\nline')).toBe(
'quote"text\nline',
);
});
});
describe('ensureCorrectEdit', () => {
let mockGeminiClientInstance: Mocked<GeminiClient>;
let mockToolRegistry: Mocked<ToolRegistry>;
let mockConfigInstance: Config;
const abortSignal = new AbortController().signal;
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;
}),
getQuotaErrorOccurred: vi.fn().mockReturnValue(false),
setQuotaErrorOccurred: vi.fn(),
} as unknown as Config;
callCount = 0;
mockResponses.length = 0;
mockGenerateJson = vi
.fn()
.mockImplementation((_contents, _schema, signal) => {
// Check if the signal is aborted. If so, throw an error or return a specific response.
if (signal && signal.aborted) {
return Promise.reject(new Error('Aborted')); // Or some other specific error/response
}
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>;
mockGeminiClientInstance.getHistory = vi.fn().mockResolvedValue([]);
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(
'/test/file.txt',
currentContent,
originalParams,
mockGeminiClientInstance,
abortSignal,
);
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(
'/test/file.txt',
currentContent,
originalParams,
mockGeminiClientInstance,
abortSignal,
);
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(
'/test/file.txt',
currentContent,
originalParams,
mockGeminiClientInstance,
abortSignal,
);
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(
'/test/file.txt',
currentContent,
originalParams,
mockGeminiClientInstance,
abortSignal,
);
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(
'/test/file.txt',
currentContent,
originalParams,
mockGeminiClientInstance,
abortSignal,
);
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(
'/test/file.txt',
currentContent,
originalParams,
mockGeminiClientInstance,
abortSignal,
);
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',
};
const result = await ensureCorrectEdit(
'/test/file.txt',
currentContent,
originalParams,
mockGeminiClientInstance,
abortSignal,
);
expect(mockGenerateJson).toHaveBeenCalledTimes(0);
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(
'/test/file.txt',
currentContent,
originalParams,
mockGeminiClientInstance,
abortSignal,
);
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(
'/test/file.txt',
currentContent,
originalParams,
mockGeminiClientInstance,
abortSignal,
);
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(
'/test/file.txt',
currentContent,
originalParams,
mockGeminiClientInstance,
abortSignal,
);
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(
'/test/file.txt',
currentContent,
originalParams,
mockGeminiClientInstance,
abortSignal,
);
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(
'/test/file.txt',
currentContent,
originalParams,
mockGeminiClientInstance,
abortSignal,
);
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(
'/test/file.txt',
currentContent,
originalParams,
mockGeminiClientInstance,
abortSignal,
);
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(
'/test/file.txt',
currentContent,
originalParams,
mockGeminiClientInstance,
abortSignal,
);
expect(mockGenerateJson).toHaveBeenCalledTimes(2);
expect(result.params.old_string).toBe(currentContent);
expect(result.params.new_string).toBe(expectedFinalNewString);
expect(result.occurrences).toBe(1);
});
});
describe('Scenario Group 6: Concurrent Edits', () => {
it('Test 6.1: should return early if file was modified by another process', async () => {
const filePath = '/test/file.txt';
const currentContent =
'This content has been modified by someone else.';
const originalParams = {
file_path: filePath,
old_string: 'nonexistent string',
new_string: 'some new string',
};
const now = Date.now();
const lastEditTime = now - 5000; // 5 seconds ago
// Mock the file's modification time to be recent
vi.spyOn(fs, 'statSync').mockReturnValue({
mtimeMs: now,
} as fs.Stats);
// Mock the last edit timestamp from our history to be in the past
const history = [
{
role: 'model',
parts: [
{
functionResponse: {
name: EditTool.Name,
id: `${EditTool.Name}-${lastEditTime}-123`,
response: {
output: {
llmContent: `Successfully modified file: ${filePath}`,
},
},
},
},
],
},
];
(mockGeminiClientInstance.getHistory as Mock).mockResolvedValue(
history,
);
const result = await ensureCorrectEdit(
filePath,
currentContent,
originalParams,
mockGeminiClientInstance,
abortSignal,
);
expect(result.occurrences).toBe(0);
expect(result.params).toEqual(originalParams);
});
});
});
describe('ensureCorrectFileContent', () => {
let mockGeminiClientInstance: Mocked<GeminiClient>;
let mockToolRegistry: Mocked<ToolRegistry>;
let mockConfigInstance: Config;
const abortSignal = new AbortController().signal;
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;
}),
getQuotaErrorOccurred: vi.fn().mockReturnValue(false),
setQuotaErrorOccurred: vi.fn(),
} as unknown as Config;
callCount = 0;
mockResponses.length = 0;
mockGenerateJson = vi
.fn()
.mockImplementation((_contents, _schema, signal) => {
if (signal && signal.aborted) {
return Promise.reject(new Error('Aborted'));
}
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();
});
it('should return content unchanged if no escaping issues detected', async () => {
const content = 'This is normal content without escaping issues';
const result = await ensureCorrectFileContent(
content,
mockGeminiClientInstance,
abortSignal,
);
expect(result).toBe(content);
expect(mockGenerateJson).toHaveBeenCalledTimes(0);
});
it('should call correctStringEscaping for potentially escaped content', async () => {
const content = 'console.log(\\"Hello World\\");';
const correctedContent = 'console.log("Hello World");';
mockResponses.push({
corrected_string_escaping: correctedContent,
});
const result = await ensureCorrectFileContent(
content,
mockGeminiClientInstance,
abortSignal,
);
expect(result).toBe(correctedContent);
expect(mockGenerateJson).toHaveBeenCalledTimes(1);
});
it('should handle correctStringEscaping returning corrected content via correct property name', async () => {
// This test specifically verifies the property name fix
const content = 'const message = \\"Hello\\nWorld\\";';
const correctedContent = 'const message = "Hello\nWorld";';
// Mock the response with the correct property name
mockResponses.push({
corrected_string_escaping: correctedContent,
});
const result = await ensureCorrectFileContent(
content,
mockGeminiClientInstance,
abortSignal,
);
expect(result).toBe(correctedContent);
expect(mockGenerateJson).toHaveBeenCalledTimes(1);
});
it('should return original content if LLM correction fails', async () => {
const content = 'console.log(\\"Hello World\\");';
// Mock empty response to simulate LLM failure
mockResponses.push({});
const result = await ensureCorrectFileContent(
content,
mockGeminiClientInstance,
abortSignal,
);
expect(result).toBe(content);
expect(mockGenerateJson).toHaveBeenCalledTimes(1);
});
it('should handle various escape sequences that need correction', async () => {
const content =
'const obj = { name: \\"John\\", age: 30, bio: \\"Developer\\nEngineer\\" };';
const correctedContent =
'const obj = { name: "John", age: 30, bio: "Developer\nEngineer" };';
mockResponses.push({
corrected_string_escaping: correctedContent,
});
const result = await ensureCorrectFileContent(
content,
mockGeminiClientInstance,
abortSignal,
);
expect(result).toBe(correctedContent);
});
});
});

747
packages/core/src/utils/editCorrector.ts
View File

@@ -1,747 +0,0 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Content, GenerateContentConfig } from '@google/genai';
import type { GeminiClient } from '../core/client.js';
import type { EditToolParams } from '../tools/edit.js';
import { ToolNames } from '../tools/tool-names.js';
import { LruCache } from './LruCache.js';
import { DEFAULT_QWEN_FLASH_MODEL } from '../config/models.js';
import {
isFunctionResponse,
isFunctionCall,
} from '../utils/messageInspectors.js';
import * as fs from 'node:fs';
const EditModel = DEFAULT_QWEN_FLASH_MODEL;
const EditConfig: GenerateContentConfig = {
thinkingConfig: {
thinkingBudget: 0,
},
};
const MAX_CACHE_SIZE = 50;
// Cache for ensureCorrectEdit results
const editCorrectionCache = new LruCache<string, CorrectedEditResult>(
MAX_CACHE_SIZE,
);
// Cache for ensureCorrectFileContent results
const fileContentCorrectionCache = new LruCache<string, string>(MAX_CACHE_SIZE);
/**
* Defines the structure of the parameters within CorrectedEditResult
*/
interface CorrectedEditParams {
file_path: string;
old_string: string;
new_string: string;
}
/**
* Defines the result structure for ensureCorrectEdit.
*/
export interface CorrectedEditResult {
params: CorrectedEditParams;
occurrences: number;
}
/**
* Extracts the timestamp from the .id value, which is in format
* <tool.name>-<timestamp>-<uuid>
* @param fcnId the ID value of a functionCall or functionResponse object
* @returns -1 if the timestamp could not be extracted, else the timestamp (as a number)
*/
function getTimestampFromFunctionId(fcnId: string): number {
const idParts = fcnId.split('-');
if (idParts.length > 2) {
const timestamp = parseInt(idParts[1], 10);
if (!isNaN(timestamp)) {
return timestamp;
}
}
return -1;
}
/**
* Will look through the gemini client history and determine when the most recent
* edit to a target file occurred. If no edit happened, it will return -1
* @param filePath the path to the file
* @param client the geminiClient, so that we can get the history
* @returns a DateTime (as a number) of when the last edit occurred, or -1 if no edit was found.
*/
async function findLastEditTimestamp(
filePath: string,
client: GeminiClient,
): Promise<number> {
const history = (await client.getHistory()) ?? [];
// Tools that may reference the file path in their FunctionResponse `output`.
const toolsInResp = new Set<string>([
ToolNames.WRITE_FILE,
ToolNames.EDIT,
ToolNames.READ_MANY_FILES,
ToolNames.GREP,
]);
// Tools that may reference the file path in their FunctionCall `args`.
const toolsInCall = new Set<string>([...toolsInResp, ToolNames.READ_FILE]);
// Iterate backwards to find the most recent relevant action.
for (const entry of history.slice().reverse()) {
if (!entry.parts) continue;
for (const part of entry.parts) {
let id: string | undefined;
let content: unknown;
// Check for a relevant FunctionCall with the file path in its arguments.
if (
isFunctionCall(entry) &&
part.functionCall?.name &&
toolsInCall.has(part.functionCall.name)
) {
id = part.functionCall.id;
content = part.functionCall.args;
}
// Check for a relevant FunctionResponse with the file path in its output.
else if (
isFunctionResponse(entry) &&
part.functionResponse?.name &&
toolsInResp.has(part.functionResponse.name)
) {
const { response } = part.functionResponse;
if (response && !('error' in response) && 'output' in response) {
id = part.functionResponse.id;
content = response['output'];
}
}
if (!id || content === undefined) continue;
// Use the "blunt hammer" approach to find the file path in the content.
// Note that the tool response data is inconsistent in their formatting
// with successes and errors - so, we just check for the existence
// as the best guess to if error/failed occurred with the response.
const stringified = JSON.stringify(content);
if (
!stringified.includes('Error') && // only applicable for functionResponse
!stringified.includes('Failed') && // only applicable for functionResponse
stringified.includes(filePath)
) {
return getTimestampFromFunctionId(id);
}
}
}
return -1;
}
/**
* Attempts to correct edit parameters if the original old_string is not found.
* It tries unescaping, and then LLM-based correction.
* Results are cached to avoid redundant processing.
*
* @param currentContent The current content of the file.
* @param originalParams The original EditToolParams
* @param client The GeminiClient for LLM calls.
* @returns A promise resolving to an object containing the (potentially corrected)
* EditToolParams (as CorrectedEditParams) and the final occurrences count.
*/
export async function ensureCorrectEdit(
filePath: string,
currentContent: string,
originalParams: EditToolParams, // This is the EditToolParams from edit.ts, without \'corrected\'
client: GeminiClient,
abortSignal: AbortSignal,
): Promise<CorrectedEditResult> {
const cacheKey = `${currentContent}---${originalParams.old_string}---${originalParams.new_string}`;
const cachedResult = editCorrectionCache.get(cacheKey);
if (cachedResult) {
return cachedResult;
}
let finalNewString = originalParams.new_string;
const newStringPotentiallyEscaped =
unescapeStringForGeminiBug(originalParams.new_string) !==
originalParams.new_string;
const expectedReplacements = originalParams.expected_replacements ?? 1;
let finalOldString = originalParams.old_string;
let occurrences = countOccurrences(currentContent, finalOldString);
if (occurrences === expectedReplacements) {
if (newStringPotentiallyEscaped) {
finalNewString = await correctNewStringEscaping(
client,
finalOldString,
originalParams.new_string,
abortSignal,
);
}
} else if (occurrences > expectedReplacements) {
const expectedReplacements = originalParams.expected_replacements ?? 1;
// If user expects multiple replacements, return as-is
if (occurrences === expectedReplacements) {
const result: CorrectedEditResult = {
params: { ...originalParams },
occurrences,
};
editCorrectionCache.set(cacheKey, result);
return result;
}
// If user expects 1 but found multiple, try to correct (existing behavior)
if (expectedReplacements === 1) {
const result: CorrectedEditResult = {
params: { ...originalParams },
occurrences,
};
editCorrectionCache.set(cacheKey, result);
return result;
}
// If occurrences don't match expected, return as-is (will fail validation later)
const result: CorrectedEditResult = {
params: { ...originalParams },
occurrences,
};
editCorrectionCache.set(cacheKey, result);
return result;
} else {
// occurrences is 0 or some other unexpected state initially
const unescapedOldStringAttempt = unescapeStringForGeminiBug(
originalParams.old_string,
);
occurrences = countOccurrences(currentContent, unescapedOldStringAttempt);
if (occurrences === expectedReplacements) {
finalOldString = unescapedOldStringAttempt;
if (newStringPotentiallyEscaped) {
finalNewString = await correctNewString(
client,
originalParams.old_string, // original old
unescapedOldStringAttempt, // corrected old
originalParams.new_string, // original new (which is potentially escaped)
abortSignal,
);
}
} else if (occurrences === 0) {
if (filePath) {
// In order to keep from clobbering edits made outside our system,
// let's check if there was a more recent edit to the file than what
// our system has done
const lastEditedByUsTime = await findLastEditTimestamp(
filePath,
client,
);
// Add a 1-second buffer to account for timing inaccuracies. If the file
// was modified more than a second after the last edit tool was run, we
// can assume it was modified by something else.
if (lastEditedByUsTime > 0) {
const stats = fs.statSync(filePath);
const diff = stats.mtimeMs - lastEditedByUsTime;
if (diff > 2000) {
// Hard coded for 2 seconds
// This file was edited sooner
const result: CorrectedEditResult = {
params: { ...originalParams },
occurrences: 0, // Explicitly 0 as LLM failed
};
editCorrectionCache.set(cacheKey, result);
return result;
}
}
}
const llmCorrectedOldString = await correctOldStringMismatch(
client,
currentContent,
unescapedOldStringAttempt,
abortSignal,
);
const llmOldOccurrences = countOccurrences(
currentContent,
llmCorrectedOldString,
);
if (llmOldOccurrences === expectedReplacements) {
finalOldString = llmCorrectedOldString;
occurrences = llmOldOccurrences;
if (newStringPotentiallyEscaped) {
const baseNewStringForLLMCorrection = unescapeStringForGeminiBug(
originalParams.new_string,
);
finalNewString = await correctNewString(
client,
originalParams.old_string, // original old
llmCorrectedOldString, // corrected old
baseNewStringForLLMCorrection, // base new for correction
abortSignal,
);
}
} else {
// LLM correction also failed for old_string
const result: CorrectedEditResult = {
params: { ...originalParams },
occurrences: 0, // Explicitly 0 as LLM failed
};
editCorrectionCache.set(cacheKey, result);
return result;
}
} else {
// Unescaping old_string resulted in > 1 occurrence
const result: CorrectedEditResult = {
params: { ...originalParams },
occurrences, // This will be > 1
};
editCorrectionCache.set(cacheKey, result);
return result;
}
}
const { targetString, pair } = trimPairIfPossible(
finalOldString,
finalNewString,
currentContent,
expectedReplacements,
);
finalOldString = targetString;
finalNewString = pair;
// Final result construction
const result: CorrectedEditResult = {
params: {
file_path: originalParams.file_path,
old_string: finalOldString,
new_string: finalNewString,
},
occurrences: countOccurrences(currentContent, finalOldString), // Recalculate occurrences with the final old_string
};
editCorrectionCache.set(cacheKey, result);
return result;
}
export async function ensureCorrectFileContent(
content: string,
client: GeminiClient,
abortSignal: AbortSignal,
): Promise<string> {
const cachedResult = fileContentCorrectionCache.get(content);
if (cachedResult) {
return cachedResult;
}
const contentPotentiallyEscaped =
unescapeStringForGeminiBug(content) !== content;
if (!contentPotentiallyEscaped) {
fileContentCorrectionCache.set(content, content);
return content;
}
const correctedContent = await correctStringEscaping(
content,
client,
abortSignal,
);
fileContentCorrectionCache.set(content, correctedContent);
return correctedContent;
}
// Define the expected JSON schema for the LLM response for old_string correction
const OLD_STRING_CORRECTION_SCHEMA: Record<string, unknown> = {
type: 'object',
properties: {
corrected_target_snippet: {
type: 'string',
description:
'The corrected version of the target snippet that exactly and uniquely matches a segment within the provided file content.',
},
},
required: ['corrected_target_snippet'],
};
export async function correctOldStringMismatch(
geminiClient: GeminiClient,
fileContent: string,
problematicSnippet: string,
abortSignal: AbortSignal,
): Promise<string> {
const prompt = `
Context: A process needs to find an exact literal, unique match for a specific text snippet within a file's content. The provided snippet failed to match exactly. This is most likely because it has been overly escaped.
Task: Analyze the provided file content and the problematic target snippet. Identify the segment in the file content that the snippet was *most likely* intended to match. Output the *exact*, literal text of that segment from the file content. Focus *only* on removing extra escape characters and correcting formatting, whitespace, or minor differences to achieve a PERFECT literal match. The output must be the exact literal text as it appears in the file.
Problematic target snippet:
\`\`\`
${problematicSnippet}
\`\`\`
File Content:
\`\`\`
${fileContent}
\`\`\`
For example, if the problematic target snippet was "\\\\\\nconst greeting = \`Hello \\\\\`\${name}\\\\\`\`;" and the file content had content that looked like "\nconst greeting = \`Hello ${'\\`'}\${name}${'\\`'}\`;", then corrected_target_snippet should likely be "\nconst greeting = \`Hello ${'\\`'}\${name}${'\\`'}\`;" to fix the incorrect escaping to match the original file content.
If the differences are only in whitespace or formatting, apply similar whitespace/formatting changes to the corrected_target_snippet.
Return ONLY the corrected target snippet in the specified JSON format with the key 'corrected_target_snippet'. If no clear, unique match can be found, return an empty string for 'corrected_target_snippet'.
`.trim();
const contents: Content[] = [{ role: 'user', parts: [{ text: prompt }] }];
try {
const result = await geminiClient.generateJson(
contents,
OLD_STRING_CORRECTION_SCHEMA,
abortSignal,
EditModel,
EditConfig,
);
if (
result &&
typeof result['corrected_target_snippet'] === 'string' &&
result['corrected_target_snippet'].length > 0
) {
return result['corrected_target_snippet'];
} else {
return problematicSnippet;
}
} catch (error) {
if (abortSignal.aborted) {
throw error;
}
console.error(
'Error during LLM call for old string snippet correction:',
error,
);
return problematicSnippet;
}
}
// Define the expected JSON schema for the new_string correction LLM response
const NEW_STRING_CORRECTION_SCHEMA: Record<string, unknown> = {
type: 'object',
properties: {
corrected_new_string: {
type: 'string',
description:
'The original_new_string adjusted to be a suitable replacement for the corrected_old_string, while maintaining the original intent of the change.',
},
},
required: ['corrected_new_string'],
};
/**
* Adjusts the new_string to align with a corrected old_string, maintaining the original intent.
*/
export async function correctNewString(
geminiClient: GeminiClient,
originalOldString: string,
correctedOldString: string,
originalNewString: string,
abortSignal: AbortSignal,
): Promise<string> {
if (originalOldString === correctedOldString) {
return originalNewString;
}
const prompt = `
Context: A text replacement operation was planned. The original text to be replaced (original_old_string) was slightly different from the actual text in the file (corrected_old_string). The original_old_string has now been corrected to match the file content.
We now need to adjust the replacement text (original_new_string) so that it makes sense as a replacement for the corrected_old_string, while preserving the original intent of the change.
original_old_string (what was initially intended to be found):
\`\`\`
${originalOldString}
\`\`\`
corrected_old_string (what was actually found in the file and will be replaced):
\`\`\`
${correctedOldString}
\`\`\`
original_new_string (what was intended to replace original_old_string):
\`\`\`
${originalNewString}
\`\`\`
Task: Based on the differences between original_old_string and corrected_old_string, and the content of original_new_string, generate a corrected_new_string. This corrected_new_string should be what original_new_string would have been if it was designed to replace corrected_old_string directly, while maintaining the spirit of the original transformation.
For example, if original_old_string was "\\\\\\nconst greeting = \`Hello \\\\\`\${name}\\\\\`\`;" and corrected_old_string is "\nconst greeting = \`Hello ${'\\`'}\${name}${'\\`'}\`;", and original_new_string was "\\\\\\nconst greeting = \`Hello \\\\\`\${name} \${lastName}\\\\\`\`;", then corrected_new_string should likely be "\nconst greeting = \`Hello ${'\\`'}\${name} \${lastName}${'\\`'}\`;" to fix the incorrect escaping.
If the differences are only in whitespace or formatting, apply similar whitespace/formatting changes to the corrected_new_string.
Return ONLY the corrected string in the specified JSON format with the key 'corrected_new_string'. If no adjustment is deemed necessary or possible, return the original_new_string.
`.trim();
const contents: Content[] = [{ role: 'user', parts: [{ text: prompt }] }];
try {
const result = await geminiClient.generateJson(
contents,
NEW_STRING_CORRECTION_SCHEMA,
abortSignal,
EditModel,
EditConfig,
);
if (
result &&
typeof result['corrected_new_string'] === 'string' &&
result['corrected_new_string'].length > 0
) {
return result['corrected_new_string'];
} else {
return originalNewString;
}
} catch (error) {
if (abortSignal.aborted) {
throw error;
}
console.error('Error during LLM call for new_string correction:', error);
return originalNewString;
}
}
const CORRECT_NEW_STRING_ESCAPING_SCHEMA: Record<string, unknown> = {
type: 'object',
properties: {
corrected_new_string_escaping: {
type: 'string',
description:
'The new_string with corrected escaping, ensuring it is a proper replacement for the old_string, especially considering potential over-escaping issues from previous LLM generations.',
},
},
required: ['corrected_new_string_escaping'],
};
export async function correctNewStringEscaping(
geminiClient: GeminiClient,
oldString: string,
potentiallyProblematicNewString: string,
abortSignal: AbortSignal,
): Promise<string> {
const prompt = `
Context: A text replacement operation is planned. The text to be replaced (old_string) has been correctly identified in the file. However, the replacement text (new_string) might have been improperly escaped by a previous LLM generation (e.g. too many backslashes for newlines like \\n instead of \n, or unnecessarily quotes like \\"Hello\\" instead of "Hello").
old_string (this is the exact text that will be replaced):
\`\`\`
${oldString}
\`\`\`
potentially_problematic_new_string (this is the text that should replace old_string, but MIGHT have bad escaping, or might be entirely correct):
\`\`\`
${potentiallyProblematicNewString}
\`\`\`
Task: Analyze the potentially_problematic_new_string. If it's syntactically invalid due to incorrect escaping (e.g., "\n", "\t", "\\", "\\'", "\\""), correct the invalid syntax. The goal is to ensure the new_string, when inserted into the code, will be a valid and correctly interpreted.
For example, if old_string is "foo" and potentially_problematic_new_string is "bar\\nbaz", the corrected_new_string_escaping should be "bar\nbaz".
If potentially_problematic_new_string is console.log(\\"Hello World\\"), it should be console.log("Hello World").
Return ONLY the corrected string in the specified JSON format with the key 'corrected_new_string_escaping'. If no escaping correction is needed, return the original potentially_problematic_new_string.
`.trim();
const contents: Content[] = [{ role: 'user', parts: [{ text: prompt }] }];
try {
const result = await geminiClient.generateJson(
contents,
CORRECT_NEW_STRING_ESCAPING_SCHEMA,
abortSignal,
EditModel,
EditConfig,
);
if (
result &&
typeof result['corrected_new_string_escaping'] === 'string' &&
result['corrected_new_string_escaping'].length > 0
) {
return result['corrected_new_string_escaping'];
} else {
return potentiallyProblematicNewString;
}
} catch (error) {
if (abortSignal.aborted) {
throw error;
}
console.error(
'Error during LLM call for new_string escaping correction:',
error,
);
return potentiallyProblematicNewString;
}
}
const CORRECT_STRING_ESCAPING_SCHEMA: Record<string, unknown> = {
type: 'object',
properties: {
corrected_string_escaping: {
type: 'string',
description:
'The string with corrected escaping, ensuring it is valid, specially considering potential over-escaping issues from previous LLM generations.',
},
},
required: ['corrected_string_escaping'],
};
export async function correctStringEscaping(
potentiallyProblematicString: string,
client: GeminiClient,
abortSignal: AbortSignal,
): Promise<string> {
const prompt = `
Context: An LLM has just generated potentially_problematic_string and the text might have been improperly escaped (e.g. too many backslashes for newlines like \\n instead of \n, or unnecessarily quotes like \\"Hello\\" instead of "Hello").
potentially_problematic_string (this text MIGHT have bad escaping, or might be entirely correct):
\`\`\`
${potentiallyProblematicString}
\`\`\`
Task: Analyze the potentially_problematic_string. If it's syntactically invalid due to incorrect escaping (e.g., "\n", "\t", "\\", "\\'", "\\""), correct the invalid syntax. The goal is to ensure the text will be a valid and correctly interpreted.
For example, if potentially_problematic_string is "bar\\nbaz", the corrected_new_string_escaping should be "bar\nbaz".
If potentially_problematic_string is console.log(\\"Hello World\\"), it should be console.log("Hello World").
Return ONLY the corrected string in the specified JSON format with the key 'corrected_string_escaping'. If no escaping correction is needed, return the original potentially_problematic_string.
`.trim();
const contents: Content[] = [{ role: 'user', parts: [{ text: prompt }] }];
try {
const result = await client.generateJson(
contents,
CORRECT_STRING_ESCAPING_SCHEMA,
abortSignal,
EditModel,
EditConfig,
);
if (
result &&
typeof result['corrected_string_escaping'] === 'string' &&
result['corrected_string_escaping'].length > 0
) {
return result['corrected_string_escaping'];
} else {
return potentiallyProblematicString;
}
} catch (error) {
if (abortSignal.aborted) {
throw error;
}
console.error(
'Error during LLM call for string escaping correction:',
error,
);
return potentiallyProblematicString;
}
}
function trimPairIfPossible(
target: string,
trimIfTargetTrims: string,
currentContent: string,
expectedReplacements: number,
) {
const trimmedTargetString = target.trim();
if (target.length !== trimmedTargetString.length) {
const trimmedTargetOccurrences = countOccurrences(
currentContent,
trimmedTargetString,
);
if (trimmedTargetOccurrences === expectedReplacements) {
const trimmedReactiveString = trimIfTargetTrims.trim();
return {
targetString: trimmedTargetString,
pair: trimmedReactiveString,
};
}
}
return {
targetString: target,
pair: trimIfTargetTrims,
};
}
/**
* Unescapes a string that might have been overly escaped by an LLM.
*/
export function unescapeStringForGeminiBug(inputString: string): string {
// Regex explanation:
// \\ : Matches exactly one literal backslash character.
// (n|t|r|'|"|`|\\|\n) : This is a capturing group. It matches one of the following:
// n, t, r, ', ", ` : These match the literal characters 'n', 't', 'r', single quote, double quote, or backtick.
// This handles cases like "\\n", "\\`", etc.
// \\ : This matches a literal backslash. This handles cases like "\\\\" (escaped backslash).
// \n : This matches an actual newline character. This handles cases where the input
// string might have something like "\\\n" (a literal backslash followed by a newline).
// g : Global flag, to replace all occurrences.
return inputString.replace(
/\\+(n|t|r|'|"|`|\\|\n)/g,
(match, capturedChar) => {
// 'match' is the entire erroneous sequence, e.g., if the input (in memory) was "\\\\`", match is "\\\\`".
// 'capturedChar' is the character that determines the true meaning, e.g., '`'.
switch (capturedChar) {
case 'n':
return '\n'; // Correctly escaped: \n (newline character)
case 't':
return '\t'; // Correctly escaped: \t (tab character)
case 'r':
return '\r'; // Correctly escaped: \r (carriage return character)
case "'":
return "'"; // Correctly escaped: ' (apostrophe character)
case '"':
return '"'; // Correctly escaped: " (quotation mark character)
case '`':
return '`'; // Correctly escaped: ` (backtick character)
case '\\': // This handles when 'capturedChar' is a literal backslash
return '\\'; // Replace escaped backslash (e.g., "\\\\") with single backslash
case '\n': // This handles when 'capturedChar' is an actual newline
return '\n'; // Replace the whole erroneous sequence (e.g., "\\\n" in memory) with a clean newline
default:
// This fallback should ideally not be reached if the regex captures correctly.
// It would return the original matched sequence if an unexpected character was captured.
return match;
}
},
);
}
/**
* Counts occurrences of a substring in a string
*/
export function 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;
}
export function resetEditCorrectorCaches_TEST_ONLY() {
editCorrectionCache.clear();
fileContentCorrectionCache.clear();
}