diff --git a/packages/core/src/tools/edit.test.ts b/packages/core/src/tools/edit.test.ts index bf003dc2..430b163a 100644 --- a/packages/core/src/tools/edit.test.ts +++ b/packages/core/src/tools/edit.test.ts @@ -6,22 +6,10 @@ /* eslint-disable @typescript-eslint/no-explicit-any */ -const mockEnsureCorrectEdit = vi.hoisted(() => vi.fn()); -const mockGenerateJson = vi.hoisted(() => vi.fn()); const mockOpenDiff = vi.hoisted(() => vi.fn()); import { IDEConnectionStatus } from '../ide/ide-client.js'; -vi.mock('../utils/editCorrector.js', () => ({ - ensureCorrectEdit: mockEnsureCorrectEdit, -})); - -vi.mock('../core/client.js', () => ({ - GeminiClient: vi.fn().mockImplementation(() => ({ - generateJson: mockGenerateJson, - })), -})); - vi.mock('../utils/editor.js', () => ({ openDiff: mockOpenDiff, })); @@ -42,7 +30,6 @@ import fs from 'node:fs'; import os from 'node:os'; import type { Config } from '../config/config.js'; import { ApprovalMode } from '../config/config.js'; -import type { Content, Part, SchemaUnion } from '@google/genai'; import { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.js'; import { StandardFileSystemService } from '../services/fileSystemService.js'; @@ -51,20 +38,13 @@ describe('EditTool', () => { let tempDir: string; let rootDir: string; let mockConfig: Config; - let geminiClient: any; - beforeEach(() => { vi.restoreAllMocks(); tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'edit-tool-test-')); rootDir = path.join(tempDir, 'root'); fs.mkdirSync(rootDir); - geminiClient = { - generateJson: mockGenerateJson, // mockGenerateJson is already defined and hoisted - }; - mockConfig = { - getGeminiClient: vi.fn().mockReturnValue(geminiClient), getTargetDir: () => rootDir, getApprovalMode: vi.fn(), setApprovalMode: vi.fn(), @@ -72,9 +52,6 @@ describe('EditTool', () => { getFileSystemService: () => new StandardFileSystemService(), getIdeClient: () => undefined, getIdeMode: () => false, - // getGeminiConfig: () => ({ apiKey: 'test-api-key' }), // This was not a real Config method - // Add other properties/methods of Config if EditTool uses them - // Minimal other methods to satisfy Config type if needed by EditTool constructor or other direct uses: getApiKey: () => 'test-api-key', getModel: () => 'test-model', getSandbox: () => false, @@ -98,65 +75,6 @@ describe('EditTool', () => { // Default to not skipping confirmation (mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.DEFAULT); - // Reset mocks and set default implementation for ensureCorrectEdit - mockEnsureCorrectEdit.mockReset(); - mockEnsureCorrectEdit.mockImplementation( - async (_, currentContent, params) => { - let occurrences = 0; - if (params.old_string && currentContent) { - // Simple string counting for the mock - let index = currentContent.indexOf(params.old_string); - while (index !== -1) { - occurrences++; - index = currentContent.indexOf(params.old_string, index + 1); - } - } else if (params.old_string === '') { - occurrences = 0; // Creating a new file - } - return Promise.resolve({ params, occurrences }); - }, - ); - - // Default mock for generateJson to return the snippet unchanged - mockGenerateJson.mockReset(); - mockGenerateJson.mockImplementation( - async (contents: Content[], schema: SchemaUnion) => { - // The problematic_snippet is the last part of the user's content - const userContent = contents.find((c: Content) => c.role === 'user'); - let promptText = ''; - if (userContent && userContent.parts) { - promptText = userContent.parts - .filter((p: Part) => typeof (p as any).text === 'string') - .map((p: Part) => (p as any).text) - .join('\n'); - } - const snippetMatch = promptText.match( - /Problematic target snippet:\n```\n([\s\S]*?)\n```/, - ); - const problematicSnippet = - snippetMatch && snippetMatch[1] ? snippetMatch[1] : ''; - - if (((schema as any).properties as any)?.corrected_target_snippet) { - return Promise.resolve({ - corrected_target_snippet: problematicSnippet, - }); - } - if (((schema as any).properties as any)?.corrected_new_string) { - // For new_string correction, we might need more sophisticated logic, - // but for now, returning original is a safe default if not specified by a test. - const originalNewStringMatch = promptText.match( - /original_new_string \(what was intended to replace original_old_string\):\n```\n([\s\S]*?)\n```/, - ); - const originalNewString = - originalNewStringMatch && originalNewStringMatch[1] - ? originalNewStringMatch[1] - : ''; - return Promise.resolve({ corrected_new_string: originalNewString }); - } - return Promise.resolve({}); // Default empty object if schema doesn't match - }, - ); - tool = new EditTool(mockConfig); }); @@ -249,8 +167,6 @@ describe('EditTool', () => { old_string: 'old', new_string: 'new', }; - // ensureCorrectEdit will be called by shouldConfirmExecute - mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 1 }); const invocation = tool.build(params); const confirmation = await invocation.shouldConfirmExecute( new AbortController().signal, @@ -264,14 +180,13 @@ describe('EditTool', () => { ); }); - it('should return false if old_string is not found (ensureCorrectEdit returns 0)', async () => { + it('should return false if old_string is not found', async () => { fs.writeFileSync(filePath, 'some content here'); const params: EditToolParams = { file_path: filePath, old_string: 'not_found', new_string: 'new', }; - mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 0 }); const invocation = tool.build(params); const confirmation = await invocation.shouldConfirmExecute( new AbortController().signal, @@ -279,14 +194,13 @@ describe('EditTool', () => { expect(confirmation).toBe(false); }); - it('should return false if multiple occurrences of old_string are found (ensureCorrectEdit returns > 1)', async () => { + it('should return false if multiple occurrences of old_string are found', async () => { fs.writeFileSync(filePath, 'old old content here'); const params: EditToolParams = { file_path: filePath, old_string: 'old', new_string: 'new', }; - mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 2 }); const invocation = tool.build(params); const confirmation = await invocation.shouldConfirmExecute( new AbortController().signal, @@ -302,10 +216,6 @@ describe('EditTool', () => { old_string: '', new_string: 'new file content', }; - // ensureCorrectEdit might not be called if old_string is empty, - // as shouldConfirmExecute handles this for diff generation. - // If it is called, it should return 0 occurrences for a new file. - mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 0 }); const invocation = tool.build(params); const confirmation = await invocation.shouldConfirmExecute( new AbortController().signal, @@ -319,65 +229,9 @@ describe('EditTool', () => { ); }); - it('should use corrected params from ensureCorrectEdit for diff generation', async () => { - const originalContent = 'This is the original string to be replaced.'; - const originalOldString = 'original string'; - const originalNewString = 'new string'; - - const correctedOldString = 'original string to be replaced'; // More specific - const correctedNewString = 'completely new string'; // Different replacement - const expectedFinalContent = 'This is the completely new string.'; - - fs.writeFileSync(filePath, originalContent); - const params: EditToolParams = { - file_path: filePath, - old_string: originalOldString, - new_string: originalNewString, - }; - - // The main beforeEach already calls mockEnsureCorrectEdit.mockReset() - // Set a specific mock for this test case - let mockCalled = false; - mockEnsureCorrectEdit.mockImplementationOnce( - async (_, content, p, client) => { - mockCalled = true; - expect(content).toBe(originalContent); - expect(p).toBe(params); - expect(client).toBe(geminiClient); - return { - params: { - file_path: filePath, - old_string: correctedOldString, - new_string: correctedNewString, - }, - occurrences: 1, - }; - }, - ); - const invocation = tool.build(params); - const confirmation = (await invocation.shouldConfirmExecute( - new AbortController().signal, - )) as FileDiff; - - expect(mockCalled).toBe(true); // Check if the mock implementation was run - // expect(mockEnsureCorrectEdit).toHaveBeenCalledWith(originalContent, params, expect.anything()); // Keep this commented for now - expect(confirmation).toEqual( - expect.objectContaining({ - title: `Confirm Edit: ${testFile}`, - fileName: testFile, - }), - ); - // Check that the diff is based on the corrected strings leading to the new state - expect(confirmation.fileDiff).toContain(`-${originalContent}`); - expect(confirmation.fileDiff).toContain(`+${expectedFinalContent}`); - - // Verify that applying the correctedOldString and correctedNewString to originalContent - // indeed produces the expectedFinalContent, which is what the diff should reflect. - const patchedContent = originalContent.replace( - correctedOldString, // This was the string identified by ensureCorrectEdit for replacement - correctedNewString, // This was the string identified by ensureCorrectEdit as the replacement - ); - expect(patchedContent).toBe(expectedFinalContent); + // This test is no longer relevant since editCorrector functionality was removed + it.skip('should use corrected params from ensureCorrectEdit for diff generation', async () => { + // Test skipped - editCorrector functionality removed }); }); @@ -387,20 +241,6 @@ describe('EditTool', () => { beforeEach(() => { filePath = path.join(rootDir, testFile); - // Default for execute tests, can be overridden - mockEnsureCorrectEdit.mockImplementation(async (_, content, params) => { - let occurrences = 0; - if (params.old_string && content) { - let index = content.indexOf(params.old_string); - while (index !== -1) { - occurrences++; - index = content.indexOf(params.old_string, index + 1); - } - } else if (params.old_string === '') { - occurrences = 0; - } - return { params, occurrences }; - }); }); it('should throw error if file path is not absolute', async () => { @@ -433,10 +273,6 @@ describe('EditTool', () => { new_string: 'new', }; - // Specific mock for this test's execution path in calculateEdit - // ensureCorrectEdit is NOT called by calculateEdit, only by shouldConfirmExecute - // So, the default mockEnsureCorrectEdit should correctly return 1 occurrence for 'old' in initialContent - const invocation = tool.build(params); const result = await invocation.execute(new AbortController().signal); @@ -477,7 +313,6 @@ describe('EditTool', () => { old_string: 'nonexistent', new_string: 'replacement', }; - // The default mockEnsureCorrectEdit will return 0 occurrences for 'nonexistent' const invocation = tool.build(params); const result = await invocation.execute(new AbortController().signal); expect(result.llmContent).toMatch( @@ -495,7 +330,6 @@ describe('EditTool', () => { old_string: 'old', new_string: 'new', }; - // The default mockEnsureCorrectEdit will return 2 occurrences for 'old' const invocation = tool.build(params); const result = await invocation.execute(new AbortController().signal); expect(result.llmContent).toMatch( @@ -638,8 +472,7 @@ describe('EditTool', () => { }); it('should return EDIT_NO_CHANGE error if replacement results in identical content', async () => { - // This can happen if ensureCorrectEdit finds a fuzzy match, but the literal - // string replacement with `replaceAll` results in no change. + // This can happen if the literal string replacement with `replaceAll` results in no change. const initialContent = 'line 1\nline 2\nline 3'; // Note the double space fs.writeFileSync(filePath, initialContent, 'utf8'); const params: EditToolParams = { @@ -649,16 +482,12 @@ describe('EditTool', () => { new_string: 'line 1\nnew line 2\nline 3', }; - // Mock ensureCorrectEdit to simulate it finding a match (e.g., via fuzzy matching) - // but it doesn't correct the old_string to the literal content. - mockEnsureCorrectEdit.mockResolvedValueOnce({ params, occurrences: 1 }); - const invocation = tool.build(params); const result = await invocation.execute(new AbortController().signal); - expect(result.error?.type).toBe(ToolErrorType.EDIT_NO_CHANGE); + expect(result.error?.type).toBe(ToolErrorType.EDIT_NO_OCCURRENCE_FOUND); expect(result.returnDisplay).toMatch( - /No changes to apply. The new content is identical to the current content./, + /Failed to edit, could not find the string to replace./, ); // Ensure the file was not actually changed expect(fs.readFileSync(filePath, 'utf8')).toBe(initialContent); @@ -870,10 +699,6 @@ describe('EditTool', () => { old_string: 'old', new_string: 'new', }; - mockEnsureCorrectEdit.mockResolvedValueOnce({ - params: { ...params, old_string: 'old', new_string: 'new' }, - occurrences: 1, - }); ideClient.openDiff.mockResolvedValueOnce({ status: 'accepted', content: modifiedContent, diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts index e40a6cd1..f05fdbd8 100644 --- a/packages/core/src/tools/edit.ts +++ b/packages/core/src/tools/edit.ts @@ -21,7 +21,6 @@ import { makeRelative, shortenPath } from '../utils/paths.js'; import { isNodeError } from '../utils/errors.js'; import type { Config } from '../config/config.js'; import { ApprovalMode } from '../config/config.js'; -import { ensureCorrectEdit } from '../utils/editCorrector.js'; import { DEFAULT_DIFF_OPTIONS, getDiffStat } from './diffOptions.js'; import { ReadFileTool } from './read-file.js'; import { ToolNames } from './tool-names.js'; @@ -116,16 +115,13 @@ class EditToolInvocation implements ToolInvocation { * @returns An object describing the potential edit outcome * @throws File system errors if reading the file fails unexpectedly (e.g., permissions) */ - private async calculateEdit( - params: EditToolParams, - abortSignal: AbortSignal, - ): Promise { + private async calculateEdit(params: EditToolParams): Promise { const expectedReplacements = params.expected_replacements ?? 1; let currentContent: string | null = null; let fileExists = false; let isNewFile = false; - let finalNewString = params.new_string; - let finalOldString = params.old_string; + const finalNewString = params.new_string; + const finalOldString = params.old_string; let occurrences = 0; let error: | { display: string; raw: string; type: ToolErrorType } @@ -157,18 +153,7 @@ class EditToolInvocation implements ToolInvocation { type: ToolErrorType.FILE_NOT_FOUND, }; } else if (currentContent !== null) { - // Editing an existing file - const correctedEdit = await ensureCorrectEdit( - params.file_path, - currentContent, - params, - this.config.getGeminiClient(), - abortSignal, - ); - finalOldString = correctedEdit.params.old_string; - finalNewString = correctedEdit.params.new_string; - occurrences = correctedEdit.occurrences; - + occurrences = this.countOccurrences(currentContent, params.old_string); if (params.old_string === '') { // Error: Trying to create a file that already exists error = { @@ -234,12 +219,28 @@ class EditToolInvocation implements ToolInvocation { }; } + /** + * Counts occurrences of a substring in a string + */ + private countOccurrences(str: string, substr: string): number { + if (substr === '') { + return 0; + } + let count = 0; + let pos = str.indexOf(substr); + while (pos !== -1) { + count++; + pos = str.indexOf(substr, pos + substr.length); // Start search after the current match + } + return count; + } + /** * Handles the confirmation prompt for the Edit tool in the CLI. * It needs to calculate the diff to show the user. */ async shouldConfirmExecute( - abortSignal: AbortSignal, + _abortSignal: AbortSignal, ): Promise { if (this.config.getApprovalMode() === ApprovalMode.AUTO_EDIT) { return false; @@ -247,7 +248,7 @@ class EditToolInvocation implements ToolInvocation { let editData: CalculatedEdit; try { - editData = await this.calculateEdit(this.params, abortSignal); + editData = await this.calculateEdit(this.params); } catch (error) { const errorMsg = error instanceof Error ? error.message : String(error); console.log(`Error preparing edit: ${errorMsg}`); @@ -330,10 +331,10 @@ class EditToolInvocation implements ToolInvocation { * @param params Parameters for the edit operation * @returns Result of the edit operation */ - async execute(signal: AbortSignal): Promise { + async execute(_signal: AbortSignal): Promise { let editData: CalculatedEdit; try { - editData = await this.calculateEdit(this.params, signal); + editData = await this.calculateEdit(this.params); } catch (error) { const errorMsg = error instanceof Error ? error.message : String(error); return { diff --git a/packages/core/src/tools/write-file.test.ts b/packages/core/src/tools/write-file.test.ts index c4ed3755..0782c6ba 100644 --- a/packages/core/src/tools/write-file.test.ts +++ b/packages/core/src/tools/write-file.test.ts @@ -18,7 +18,6 @@ import { getCorrectedFileContent, WriteFileTool } from './write-file.js'; import { ToolErrorType } from './tool-error.js'; import type { FileDiff, ToolEditConfirmationDetails } from './tools.js'; import { ToolConfirmationOutcome } from './tools.js'; -import { type EditToolParams } from './edit.js'; import type { Config } from '../config/config.js'; import { ApprovalMode } from '../config/config.js'; import type { ToolRegistry } from './tool-registry.js'; @@ -26,11 +25,6 @@ import path from 'node:path'; import fs from 'node:fs'; import os from 'node:os'; import { GeminiClient } from '../core/client.js'; -import type { CorrectedEditResult } from '../utils/editCorrector.js'; -import { - ensureCorrectEdit, - ensureCorrectFileContent, -} from '../utils/editCorrector.js'; import { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.js'; import { StandardFileSystemService } from '../services/fileSystemService.js'; @@ -38,17 +32,8 @@ const rootDir = path.resolve(os.tmpdir(), 'qwen-code-test-root'); // --- MOCKS --- vi.mock('../core/client.js'); -vi.mock('../utils/editCorrector.js'); let mockGeminiClientInstance: Mocked; -const mockEnsureCorrectEdit = vi.fn(); -const mockEnsureCorrectFileContent = vi.fn(); - -// Wire up the mocked functions to be used by the actual module imports -vi.mocked(ensureCorrectEdit).mockImplementation(mockEnsureCorrectEdit); -vi.mocked(ensureCorrectFileContent).mockImplementation( - mockEnsureCorrectFileContent, -); // Mock Config const fsService = new StandardFileSystemService(); @@ -111,11 +96,6 @@ describe('WriteFileTool', () => { ) as Mocked; vi.mocked(GeminiClient).mockImplementation(() => mockGeminiClientInstance); - vi.mocked(ensureCorrectEdit).mockImplementation(mockEnsureCorrectEdit); - vi.mocked(ensureCorrectFileContent).mockImplementation( - mockEnsureCorrectFileContent, - ); - // Now that mockGeminiClientInstance is initialized, set the mock implementation for getGeminiClient mockConfigInternal.getGeminiClient.mockReturnValue( mockGeminiClientInstance, @@ -134,40 +114,6 @@ describe('WriteFileTool', () => { // Reset mocks before each test mockConfigInternal.getApprovalMode.mockReturnValue(ApprovalMode.DEFAULT); mockConfigInternal.setApprovalMode.mockClear(); - mockEnsureCorrectEdit.mockReset(); - mockEnsureCorrectFileContent.mockReset(); - - // Default mock implementations that return valid structures - mockEnsureCorrectEdit.mockImplementation( - async ( - filePath: string, - _currentContent: string, - params: EditToolParams, - _client: GeminiClient, - signal?: AbortSignal, // Make AbortSignal optional to match usage - ): Promise => { - 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 => { - // Make AbortSignal optional - if (signal?.aborted) { - return Promise.reject(new Error('Aborted')); - } - return Promise.resolve(content ?? ''); - }, - ); }); afterEach(() => { @@ -240,71 +186,35 @@ describe('WriteFileTool', () => { }); describe('getCorrectedFileContent', () => { - it('should call ensureCorrectFileContent for a new file', async () => { + it('should return proposed content unchanged for a new file', async () => { const filePath = path.join(rootDir, 'new_corrected_file.txt'); const proposedContent = 'Proposed new content.'; - const correctedContent = 'Corrected new content.'; - const abortSignal = new AbortController().signal; - // Ensure the mock is set for this specific test case if needed, or rely on beforeEach - mockEnsureCorrectFileContent.mockResolvedValue(correctedContent); const result = await getCorrectedFileContent( mockConfig, filePath, proposedContent, - abortSignal, ); - expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith( - proposedContent, - mockGeminiClientInstance, - abortSignal, - ); - expect(mockEnsureCorrectEdit).not.toHaveBeenCalled(); - expect(result.correctedContent).toBe(correctedContent); + expect(result.correctedContent).toBe(proposedContent); expect(result.originalContent).toBe(''); expect(result.fileExists).toBe(false); expect(result.error).toBeUndefined(); }); - it('should call ensureCorrectEdit for an existing file', async () => { + it('should return proposed content unchanged for an existing file', async () => { const filePath = path.join(rootDir, 'existing_corrected_file.txt'); const originalContent = 'Original existing content.'; const proposedContent = 'Proposed replacement content.'; - const correctedProposedContent = 'Corrected replacement content.'; - const abortSignal = new AbortController().signal; fs.writeFileSync(filePath, originalContent, 'utf8'); - // Ensure this mock is active and returns the correct structure - mockEnsureCorrectEdit.mockResolvedValue({ - params: { - file_path: filePath, - old_string: originalContent, - new_string: correctedProposedContent, - }, - occurrences: 1, - } as CorrectedEditResult); - const result = await getCorrectedFileContent( mockConfig, filePath, proposedContent, - abortSignal, ); - expect(mockEnsureCorrectEdit).toHaveBeenCalledWith( - filePath, - originalContent, - { - old_string: originalContent, - new_string: proposedContent, - file_path: filePath, - }, - mockGeminiClientInstance, - abortSignal, - ); - expect(mockEnsureCorrectFileContent).not.toHaveBeenCalled(); - expect(result.correctedContent).toBe(correctedProposedContent); + expect(result.correctedContent).toBe(proposedContent); expect(result.originalContent).toBe(originalContent); expect(result.fileExists).toBe(true); expect(result.error).toBeUndefined(); @@ -313,7 +223,6 @@ describe('WriteFileTool', () => { it('should return error if reading an existing file fails (e.g. permissions)', async () => { const filePath = path.join(rootDir, 'unreadable_file.txt'); const proposedContent = 'some content'; - const abortSignal = new AbortController().signal; fs.writeFileSync(filePath, 'content', { mode: 0o000 }); const readError = new Error('Permission denied'); @@ -325,12 +234,9 @@ describe('WriteFileTool', () => { mockConfig, filePath, proposedContent, - abortSignal, ); expect(fsService.readTextFile).toHaveBeenCalledWith(filePath); - expect(mockEnsureCorrectEdit).not.toHaveBeenCalled(); - expect(mockEnsureCorrectFileContent).not.toHaveBeenCalled(); expect(result.correctedContent).toBe(proposedContent); expect(result.originalContent).toBe(''); expect(result.fileExists).toBe(true); @@ -363,11 +269,9 @@ describe('WriteFileTool', () => { fs.chmodSync(filePath, 0o600); }); - it('should request confirmation with diff for a new file (with corrected content)', async () => { + it('should request confirmation with diff for a new file', async () => { const filePath = path.join(rootDir, 'confirm_new_file.txt'); const proposedContent = 'Proposed new content for confirmation.'; - const correctedContent = 'Corrected new content for confirmation.'; - mockEnsureCorrectFileContent.mockResolvedValue(correctedContent); // Ensure this mock is active const params = { file_path: filePath, content: proposedContent }; const invocation = tool.build(params); @@ -375,16 +279,11 @@ describe('WriteFileTool', () => { abortSignal, )) as ToolEditConfirmationDetails; - expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith( - proposedContent, - mockGeminiClientInstance, - abortSignal, - ); expect(confirmation).toEqual( expect.objectContaining({ title: `Confirm Write: ${path.basename(filePath)}`, fileName: 'confirm_new_file.txt', - fileDiff: expect.stringContaining(correctedContent), + fileDiff: expect.stringContaining(proposedContent), }), ); expect(confirmation.fileDiff).toMatch( @@ -395,45 +294,23 @@ describe('WriteFileTool', () => { ); }); - it('should request confirmation with diff for an existing file (with corrected content)', async () => { + it('should request confirmation with diff for an existing file', async () => { const filePath = path.join(rootDir, 'confirm_existing_file.txt'); const originalContent = 'Original content for confirmation.'; const proposedContent = 'Proposed replacement for confirmation.'; - const correctedProposedContent = - 'Corrected replacement for confirmation.'; fs.writeFileSync(filePath, originalContent, 'utf8'); - mockEnsureCorrectEdit.mockResolvedValue({ - params: { - file_path: filePath, - old_string: originalContent, - new_string: correctedProposedContent, - }, - occurrences: 1, - }); - const params = { file_path: filePath, content: proposedContent }; const invocation = tool.build(params); const confirmation = (await invocation.shouldConfirmExecute( abortSignal, )) as ToolEditConfirmationDetails; - expect(mockEnsureCorrectEdit).toHaveBeenCalledWith( - filePath, - originalContent, - { - old_string: originalContent, - new_string: proposedContent, - file_path: filePath, - }, - mockGeminiClientInstance, - abortSignal, - ); expect(confirmation).toEqual( expect.objectContaining({ title: `Confirm Write: ${path.basename(filePath)}`, fileName: 'confirm_existing_file.txt', - fileDiff: expect.stringContaining(correctedProposedContent), + fileDiff: expect.stringContaining(proposedContent), }), ); expect(confirmation.fileDiff).toMatch( @@ -470,11 +347,9 @@ describe('WriteFileTool', () => { fs.chmodSync(filePath, 0o600); }); - it('should write a new file with corrected content and return diff', async () => { - const filePath = path.join(rootDir, 'execute_new_corrected_file.txt'); + it('should write a new file and return diff', async () => { + const filePath = path.join(rootDir, 'execute_new_file.txt'); const proposedContent = 'Proposed new content for execute.'; - const correctedContent = 'Corrected new content for execute.'; - mockEnsureCorrectFileContent.mockResolvedValue(correctedContent); const params = { file_path: filePath, content: proposedContent }; const invocation = tool.build(params); @@ -490,49 +365,27 @@ describe('WriteFileTool', () => { const result = await invocation.execute(abortSignal); - expect(mockEnsureCorrectFileContent).toHaveBeenCalledWith( - proposedContent, - mockGeminiClientInstance, - abortSignal, - ); expect(result.llmContent).toMatch( /Successfully created and wrote to new file/, ); expect(fs.existsSync(filePath)).toBe(true); const writtenContent = await fsService.readTextFile(filePath); - expect(writtenContent).toBe(correctedContent); + expect(writtenContent).toBe(proposedContent); const display = result.returnDisplay as FileDiff; - expect(display.fileName).toBe('execute_new_corrected_file.txt'); + expect(display.fileName).toBe('execute_new_file.txt'); + expect(display.fileDiff).toMatch(/--- execute_new_file.txt\tOriginal/); + expect(display.fileDiff).toMatch(/\+\+\+ execute_new_file.txt\tWritten/); expect(display.fileDiff).toMatch( - /--- execute_new_corrected_file.txt\tOriginal/, - ); - expect(display.fileDiff).toMatch( - /\+\+\+ execute_new_corrected_file.txt\tWritten/, - ); - expect(display.fileDiff).toMatch( - correctedContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'), + proposedContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'), ); }); - it('should overwrite an existing file with corrected content and return diff', async () => { - const filePath = path.join( - rootDir, - 'execute_existing_corrected_file.txt', - ); + it('should overwrite an existing file and return diff', async () => { + const filePath = path.join(rootDir, 'execute_existing_file.txt'); const initialContent = 'Initial content for execute.'; const proposedContent = 'Proposed overwrite for execute.'; - const correctedProposedContent = 'Corrected overwrite for execute.'; fs.writeFileSync(filePath, initialContent, 'utf8'); - mockEnsureCorrectEdit.mockResolvedValue({ - params: { - file_path: filePath, - old_string: initialContent, - new_string: correctedProposedContent, - }, - occurrences: 1, - }); - const params = { file_path: filePath, content: proposedContent }; const invocation = tool.build(params); @@ -547,27 +400,16 @@ describe('WriteFileTool', () => { const result = await invocation.execute(abortSignal); - expect(mockEnsureCorrectEdit).toHaveBeenCalledWith( - filePath, - initialContent, - { - old_string: initialContent, - new_string: proposedContent, - file_path: filePath, - }, - mockGeminiClientInstance, - abortSignal, - ); expect(result.llmContent).toMatch(/Successfully overwrote file/); const writtenContent = await fsService.readTextFile(filePath); - expect(writtenContent).toBe(correctedProposedContent); + expect(writtenContent).toBe(proposedContent); const display = result.returnDisplay as FileDiff; - expect(display.fileName).toBe('execute_existing_corrected_file.txt'); + expect(display.fileName).toBe('execute_existing_file.txt'); expect(display.fileDiff).toMatch( initialContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'), ); expect(display.fileDiff).toMatch( - correctedProposedContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'), + proposedContent.replace(/[.*+?^${}()|[\\]\\]/g, '\\$&'), ); }); @@ -575,7 +417,6 @@ describe('WriteFileTool', () => { const dirPath = path.join(rootDir, 'new_dir_for_write'); const filePath = path.join(dirPath, 'file_in_new_dir.txt'); const content = 'Content in new directory'; - mockEnsureCorrectFileContent.mockResolvedValue(content); // Ensure this mock is active const params = { file_path: filePath, content }; const invocation = tool.build(params); @@ -600,7 +441,6 @@ describe('WriteFileTool', () => { it('should include modification message when proposed content is modified', async () => { const filePath = path.join(rootDir, 'new_file_modified.txt'); const content = 'New file content modified by user'; - mockEnsureCorrectFileContent.mockResolvedValue(content); const params = { file_path: filePath, @@ -616,7 +456,6 @@ describe('WriteFileTool', () => { it('should not include modification message when proposed content is not modified', async () => { const filePath = path.join(rootDir, 'new_file_unmodified.txt'); const content = 'New file content not modified'; - mockEnsureCorrectFileContent.mockResolvedValue(content); const params = { file_path: filePath, @@ -632,7 +471,6 @@ describe('WriteFileTool', () => { it('should not include modification message when modified_by_user is not provided', async () => { const filePath = path.join(rootDir, 'new_file_unmodified.txt'); const content = 'New file content not modified'; - mockEnsureCorrectFileContent.mockResolvedValue(content); const params = { file_path: filePath, diff --git a/packages/core/src/tools/write-file.ts b/packages/core/src/tools/write-file.ts index dc32bcc8..4c92bef8 100644 --- a/packages/core/src/tools/write-file.ts +++ b/packages/core/src/tools/write-file.ts @@ -26,10 +26,6 @@ import { import { ToolErrorType } from './tool-error.js'; import { makeRelative, shortenPath } from '../utils/paths.js'; import { getErrorMessage, isNodeError } from '../utils/errors.js'; -import { - ensureCorrectEdit, - ensureCorrectFileContent, -} from '../utils/editCorrector.js'; import { DEFAULT_DIFF_OPTIONS, getDiffStat } from './diffOptions.js'; import { ToolNames } from './tool-names.js'; import type { @@ -79,11 +75,10 @@ export async function getCorrectedFileContent( config: Config, filePath: string, proposedContent: string, - abortSignal: AbortSignal, ): Promise { let originalContent = ''; let fileExists = false; - let correctedContent = proposedContent; + const correctedContent = proposedContent; try { originalContent = await config @@ -107,32 +102,6 @@ export async function getCorrectedFileContent( } } - // If readError is set, we have returned. - // So, file was either read successfully (fileExists=true, originalContent set) - // or it was ENOENT (fileExists=false, originalContent=''). - - if (fileExists) { - // This implies originalContent is available - const { params: correctedParams } = await ensureCorrectEdit( - filePath, - originalContent, - { - old_string: originalContent, // Treat entire current content as old_string - new_string: proposedContent, - file_path: filePath, - }, - config.getGeminiClient(), - abortSignal, - ); - correctedContent = correctedParams.new_string; - } else { - // This implies new file (ENOENT) - correctedContent = await ensureCorrectFileContent( - proposedContent, - config.getGeminiClient(), - abortSignal, - ); - } return { originalContent, correctedContent, fileExists }; } @@ -160,7 +129,7 @@ class WriteFileToolInvocation extends BaseToolInvocation< } override async shouldConfirmExecute( - abortSignal: AbortSignal, + _abortSignal: AbortSignal, ): Promise { if (this.config.getApprovalMode() === ApprovalMode.AUTO_EDIT) { return false; @@ -170,7 +139,6 @@ class WriteFileToolInvocation extends BaseToolInvocation< this.config, this.params.file_path, this.params.content, - abortSignal, ); if (correctedContentResult.error) { @@ -226,14 +194,13 @@ class WriteFileToolInvocation extends BaseToolInvocation< return confirmationDetails; } - async execute(abortSignal: AbortSignal): Promise { + async execute(_abortSignal: AbortSignal): Promise { const { file_path, content, ai_proposed_content, modified_by_user } = this.params; const correctedContentResult = await getCorrectedFileContent( this.config, file_path, content, - abortSignal, ); if (correctedContentResult.error) { @@ -476,7 +443,7 @@ export class WriteFileTool } getModifyContext( - abortSignal: AbortSignal, + _abortSignal: AbortSignal, ): ModifyContext { return { getFilePath: (params: WriteFileToolParams) => params.file_path, @@ -485,7 +452,6 @@ export class WriteFileTool this.config, params.file_path, params.content, - abortSignal, ); return correctedContentResult.originalContent; }, @@ -494,7 +460,6 @@ export class WriteFileTool this.config, params.file_path, params.content, - abortSignal, ); return correctedContentResult.correctedContent; }, diff --git a/packages/core/src/utils/editCorrector.test.ts b/packages/core/src/utils/editCorrector.test.ts deleted file mode 100644 index 7f76bed4..00000000 --- a/packages/core/src/utils/editCorrector.test.ts +++ /dev/null @@ -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; - let mockToolRegistry: Mocked; - let mockConfigInstance: Config; - const abortSignal = new AbortController().signal; - - beforeEach(() => { - mockToolRegistry = new ToolRegistry({} as Config) as Mocked; - 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 | 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; - 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; - let mockToolRegistry: Mocked; - let mockConfigInstance: Config; - const abortSignal = new AbortController().signal; - - beforeEach(() => { - mockToolRegistry = new ToolRegistry({} as Config) as Mocked; - 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 | 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; - 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); - }); - }); -}); diff --git a/packages/core/src/utils/editCorrector.ts b/packages/core/src/utils/editCorrector.ts deleted file mode 100644 index 381108e4..00000000 --- a/packages/core/src/utils/editCorrector.ts +++ /dev/null @@ -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( - MAX_CACHE_SIZE, -); - -// Cache for ensureCorrectFileContent results -const fileContentCorrectionCache = new LruCache(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 - * -- - * @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 { - const history = (await client.getHistory()) ?? []; - - // Tools that may reference the file path in their FunctionResponse `output`. - const toolsInResp = new Set([ - 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([...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 { - 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 { - 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 = { - 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 { - 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 = { - 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 { - 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 = { - 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 { - 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 = { - 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 { - 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(); -}