diff options
Diffstat (limited to 'packages/server/src')
| -rw-r--r-- | packages/server/src/tools/edit.test.ts | 137 | ||||
| -rw-r--r-- | packages/server/src/tools/edit.ts | 92 |
2 files changed, 189 insertions, 40 deletions
diff --git a/packages/server/src/tools/edit.test.ts b/packages/server/src/tools/edit.test.ts index a552cf53..016e31bf 100644 --- a/packages/server/src/tools/edit.test.ts +++ b/packages/server/src/tools/edit.test.ts @@ -6,6 +6,19 @@ /* eslint-disable @typescript-eslint/no-explicit-any */ +const mockEnsureCorrectEdit = vi.hoisted(() => vi.fn()); +const mockGenerateJson = vi.hoisted(() => vi.fn()); + +vi.mock('../utils/editCorrector.js', () => ({ + ensureCorrectEdit: mockEnsureCorrectEdit, +})); + +vi.mock('../core/client.js', () => ({ + GeminiClient: vi.fn().mockImplementation(() => ({ + generateJson: mockGenerateJson, + })), +})); + import { describe, it, expect, beforeEach, afterEach, vi, Mock } from 'vitest'; import { EditTool, EditToolParams } from './edit.js'; import { FileDiff } from './tools.js'; @@ -15,16 +28,6 @@ import os from 'os'; import { Config } from '../config/config.js'; import { Content, Part, SchemaUnion } from '@google/genai'; -// Mock GeminiClient -const mockEnsureCorrectEdit = vi.fn(); -const mockGenerateJson = vi.fn(); -vi.mock('../core/client.js', () => ({ - GeminiClient: vi.fn().mockImplementation(() => ({ - ensureCorrectEdit: mockEnsureCorrectEdit, - generateJson: mockGenerateJson, - })), -})); - describe('EditTool', () => { let tool: EditTool; let tempDir: string; @@ -133,6 +136,48 @@ describe('EditTool', () => { fs.rmSync(tempDir, { recursive: true, force: true }); }); + describe('_applyReplacement', () => { + // Access private method for testing + // Note: `tool` is initialized in `beforeEach` of the parent describe block + it('should return newString if isNewFile is true', () => { + expect((tool as any)._applyReplacement(null, 'old', 'new', true)).toBe( + 'new', + ); + expect( + (tool as any)._applyReplacement('existing', 'old', 'new', true), + ).toBe('new'); + }); + + it('should return newString if currentContent is null and oldString is empty (defensive)', () => { + expect((tool as any)._applyReplacement(null, '', 'new', false)).toBe( + 'new', + ); + }); + + it('should return empty string if currentContent is null and oldString is not empty (defensive)', () => { + expect((tool as any)._applyReplacement(null, 'old', 'new', false)).toBe( + '', + ); + }); + + it('should replace oldString with newString in currentContent', () => { + expect( + (tool as any)._applyReplacement( + 'hello old world old', + 'old', + 'new', + false, + ), + ).toBe('hello new world new'); + }); + + it('should return currentContent if oldString is empty and not a new file', () => { + expect( + (tool as any)._applyReplacement('hello world', '', 'new', false), + ).toBe('hello world'); + }); + }); + describe('validateParams', () => { it('should return null for valid params', () => { const params: EditToolParams = { @@ -243,6 +288,68 @@ 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) => { + console.log('mockEnsureCorrectEdit CALLED IN TEST'); + mockCalled = true; + expect(content).toBe(originalContent); + expect(p).toBe(params); + expect(client).toBe((tool as any).client); + return { + params: { + file_path: filePath, + old_string: correctedOldString, + new_string: correctedNewString, + }, + occurrences: 1, + }; + }, + ); + + const confirmation = (await tool.shouldConfirmExecute( + params, + )) 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); + }); }); describe('execute', () => { @@ -337,9 +444,11 @@ describe('EditTool', () => { }; // The default mockEnsureCorrectEdit will return 0 occurrences for 'nonexistent' const result = await tool.execute(params, new AbortController().signal); - expect(result.llmContent).toMatch(/0 occurrences found/); + expect(result.llmContent).toMatch( + /0 occurrences found for old_string in/, + ); expect(result.returnDisplay).toMatch( - /Failed to edit, could not find the string to replace/, + /Failed to edit, could not find the string to replace./, ); }); @@ -352,7 +461,9 @@ describe('EditTool', () => { }; // The default mockEnsureCorrectEdit will return 2 occurrences for 'old' const result = await tool.execute(params, new AbortController().signal); - expect(result.llmContent).toMatch(/Expected 1 occurrences but found 2/); + expect(result.llmContent).toMatch( + /Expected 1 occurrences but found 2 for old_string in file/, + ); expect(result.returnDisplay).toMatch( /Failed to edit, expected 1 occurrence\(s\) but found 2/, ); diff --git a/packages/server/src/tools/edit.ts b/packages/server/src/tools/edit.ts index 9301d5e8..bd070d65 100644 --- a/packages/server/src/tools/edit.ts +++ b/packages/server/src/tools/edit.ts @@ -21,7 +21,7 @@ import { isNodeError } from '../utils/errors.js'; import { ReadFileTool } from './read-file.js'; import { GeminiClient } from '../core/client.js'; import { Config } from '../config/config.js'; -import { countOccurrences, ensureCorrectEdit } from '../utils/editCorrector.js'; +import { ensureCorrectEdit } from '../utils/editCorrector.js'; /** * Parameters for the Edit tool @@ -147,6 +147,26 @@ Expectation for parameters: return null; } + private _applyReplacement( + currentContent: string | null, + oldString: string, + newString: string, + isNewFile: boolean, + ): string { + if (isNewFile) { + return newString; + } + if (currentContent === null) { + // Should not happen if not a new file, but defensively return empty or newString if oldString is also empty + return oldString === '' ? newString : ''; + } + // If oldString is empty and it's not a new file, do not modify the content. + if (oldString === '' && !isNewFile) { + return currentContent; + } + return currentContent.replaceAll(oldString, newString); + } + /** * Calculates the potential outcome of an edit operation. * @param params Parameters for the edit operation @@ -158,7 +178,8 @@ Expectation for parameters: let currentContent: string | null = null; let fileExists = false; let isNewFile = false; - let newContent = ''; + let finalNewString = params.new_string; + let finalOldString = params.old_string; let occurrences = 0; let error: { display: string; raw: string } | undefined = undefined; @@ -176,8 +197,6 @@ Expectation for parameters: if (params.old_string === '' && !fileExists) { // Creating a new file isNewFile = true; - newContent = params.new_string; - occurrences = 0; } else if (!fileExists) { // Trying to edit a non-existent file (and old_string is not empty) error = { @@ -186,7 +205,14 @@ Expectation for parameters: }; } else if (currentContent !== null) { // Editing an existing file - occurrences = countOccurrences(currentContent, params.old_string); + const correctedEdit = await ensureCorrectEdit( + currentContent, + params, + this.client, + ); + finalOldString = correctedEdit.params.old_string; + finalNewString = correctedEdit.params.new_string; + occurrences = correctedEdit.occurrences; if (params.old_string === '') { // Error: Trying to create a file that already exists @@ -197,19 +223,13 @@ Expectation for parameters: } else if (occurrences === 0) { error = { display: `Failed to edit, could not find the string to replace.`, - raw: `Failed to edit, 0 occurrences found for old_string in ${params.file_path}. No edits made. The exact text in old_string was not found. Ensure you're not escaping content incorrectly and check whitespace, indentation, and context. Use ReadFile tool to verify.`, + raw: `Failed to edit, 0 occurrences found for old_string in ${params.file_path}. No edits made. The exact text in old_string was not found. Ensure you're not escaping content incorrectly and check whitespace, indentation, and context. Use ${ReadFileTool.Name} tool to verify.`, }; } else if (occurrences !== expectedReplacements) { error = { display: `Failed to edit, expected ${expectedReplacements} occurrence(s) but found ${occurrences}.`, raw: `Failed to edit, Expected ${expectedReplacements} occurrences but found ${occurrences} for old_string in file: ${params.file_path}`, }; - } else { - // Successful edit calculation - newContent = currentContent.replaceAll( - params.old_string, - params.new_string, - ); } } else { // Should not happen if fileExists and no exception was thrown, but defensively: @@ -219,6 +239,13 @@ Expectation for parameters: }; } + const newContent = this._applyReplacement( + currentContent, + finalOldString, + finalNewString, + isNewFile, + ); + return { currentContent, newContent, @@ -247,7 +274,10 @@ Expectation for parameters: } let currentContent: string | null = null; let fileExists = false; - let newContent = ''; + let finalNewString = params.new_string; + let finalOldString = params.old_string; + let occurrences = 0; + try { currentContent = fs.readFileSync(params.file_path, 'utf8'); fileExists = true; @@ -259,28 +289,36 @@ Expectation for parameters: return false; } } + if (params.old_string === '' && !fileExists) { - newContent = params.new_string; + // Creating new file, newContent is just params.new_string } else if (!fileExists) { - return false; + return false; // Cannot edit non-existent file if old_string is not empty } else if (currentContent !== null) { - // Use the correctEdit utility to potentially correct params and get occurrences - const { params: correctedParams, occurrences: correctedOccurrences } = - await ensureCorrectEdit(currentContent, params, this.client); - - params.old_string = correctedParams.old_string; - params.new_string = correctedParams.new_string; + const correctedEdit = await ensureCorrectEdit( + currentContent, + params, + this.client, + ); + finalOldString = correctedEdit.params.old_string; + finalNewString = correctedEdit.params.new_string; + occurrences = correctedEdit.occurrences; - if (correctedOccurrences === 0 || correctedOccurrences !== 1) { + if (occurrences === 0 || occurrences !== 1) { return false; } - newContent = currentContent.replaceAll( - params.old_string, - params.new_string, - ); } else { - return false; + return false; // Should not happen } + + const isNewFileScenario = params.old_string === '' && !fileExists; + const newContent = this._applyReplacement( + currentContent, + finalOldString, + finalNewString, + isNewFileScenario, + ); + const fileName = path.basename(params.file_path); const fileDiff = Diff.createPatch( fileName, |