diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 2703517819..0673fecb6d 100755 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -78,6 +78,7 @@ export interface CliArgs { proxy: string | undefined; includeDirectories: string[] | undefined; screenReader: boolean | undefined; + useSmartEdit: boolean | undefined; sessionSummary: string | undefined; } @@ -621,6 +622,7 @@ export async function loadCliConfig( skipNextSpeakerCheck: settings.model?.skipNextSpeakerCheck, enablePromptCompletion: settings.general?.enablePromptCompletion ?? false, eventEmitter: appEvents, + useSmartEdit: argv.useSmartEdit ?? settings.useSmartEdit, }); } diff --git a/packages/cli/src/config/settingsSchema.ts b/packages/cli/src/config/settingsSchema.ts index 9210f0b5c7..73e6f7e425 100644 --- a/packages/cli/src/config/settingsSchema.ts +++ b/packages/cli/src/config/settingsSchema.ts @@ -613,7 +613,15 @@ export const SETTINGS_SCHEMA = { }, }, }, - + useSmartEdit: { + type: 'boolean', + label: 'Use Smart Edit', + category: 'Advanced', + requiresRestart: false, + default: false, + description: 'Enable the smart-edit tool instead of the replace tool.', + showInDialog: false, + }, security: { type: 'object', label: 'Security', diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 2a127ef0c3..d1ebd718d2 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -19,6 +19,7 @@ import { GrepTool } from '../tools/grep.js'; import { RipGrepTool } from '../tools/ripGrep.js'; import { GlobTool } from '../tools/glob.js'; import { EditTool } from '../tools/edit.js'; +import { SmartEditTool } from '../tools/smart-edit.js'; import { ShellTool } from '../tools/shell.js'; import { WriteFileTool } from '../tools/write-file.js'; import { WebFetchTool } from '../tools/web-fetch.js'; @@ -209,6 +210,7 @@ export interface ConfigParameters { extensionManagement?: boolean; enablePromptCompletion?: boolean; eventEmitter?: EventEmitter; + useSmartEdit?: boolean; } export class Config { @@ -285,6 +287,7 @@ export class Config { readonly storage: Storage; private readonly fileExclusions: FileExclusions; private readonly eventEmitter?: EventEmitter; + private readonly useSmartEdit: boolean; constructor(params: ConfigParameters) { this.sessionId = params.sessionId; @@ -355,6 +358,7 @@ export class Config { this.useRipgrep = params.useRipgrep ?? false; this.shouldUseNodePtyShell = params.shouldUseNodePtyShell ?? false; this.skipNextSpeakerCheck = params.skipNextSpeakerCheck ?? false; + this.useSmartEdit = params.useSmartEdit ?? false; this.extensionManagement = params.extensionManagement ?? false; this.storage = new Storage(this.targetDir); this.enablePromptCompletion = params.enablePromptCompletion ?? false; @@ -808,6 +812,10 @@ export class Config { return this.enablePromptCompletion; } + getUseSmartEdit(): boolean { + return this.useSmartEdit; + } + async getGitService(): Promise { if (!this.gitService) { this.gitService = new GitService(this.targetDir, this.storage); @@ -865,7 +873,11 @@ export class Config { } registerCoreTool(GlobTool, this); - registerCoreTool(EditTool, this); + if (this.getUseSmartEdit()) { + registerCoreTool(SmartEditTool, this); + } else { + registerCoreTool(EditTool, this); + } registerCoreTool(WriteFileTool, this); registerCoreTool(WebFetchTool, this); registerCoreTool(ReadManyFilesTool, this); diff --git a/packages/core/src/tools/smart-edit.test.ts b/packages/core/src/tools/smart-edit.test.ts new file mode 100644 index 0000000000..51ab29a57f --- /dev/null +++ b/packages/core/src/tools/smart-edit.test.ts @@ -0,0 +1,488 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/* eslint-disable @typescript-eslint/no-explicit-any */ + +const mockFixLLMEditWithInstruction = vi.hoisted(() => vi.fn()); +const mockGenerateJson = vi.hoisted(() => vi.fn()); +const mockOpenDiff = vi.hoisted(() => vi.fn()); + +import { IDEConnectionStatus } from '../ide/ide-client.js'; + +vi.mock('../utils/llm-edit-fixer.js', () => ({ + FixLLMEditWithInstruction: mockFixLLMEditWithInstruction, +})); + +vi.mock('../core/client.js', () => ({ + GeminiClient: vi.fn().mockImplementation(() => ({ + generateJson: mockGenerateJson, + })), +})); + +vi.mock('../utils/editor.js', () => ({ + openDiff: mockOpenDiff, +})); + +import { + describe, + it, + expect, + beforeEach, + afterEach, + vi, + type Mock, +} from 'vitest'; +import { + applyReplacement, + SmartEditTool, + type EditToolParams, + calculateReplacement, +} from './smart-edit.js'; +import { type FileDiff, ToolConfirmationOutcome } from './tools.js'; +import { ToolErrorType } from './tool-error.js'; +import path from 'node:path'; +import fs from 'node:fs'; +import os from 'node:os'; +import { ApprovalMode, type Config } from '../config/config.js'; +import { type Content, type Part, type SchemaUnion } from '@google/genai'; +import { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.js'; +import { StandardFileSystemService } from '../services/fileSystemService.js'; + +describe('SmartEditTool', () => { + let tool: SmartEditTool; + let tempDir: string; + let rootDir: string; + let mockConfig: Config; + let geminiClient: any; + + beforeEach(() => { + vi.restoreAllMocks(); + tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'smart-edit-tool-test-')); + rootDir = path.join(tempDir, 'root'); + fs.mkdirSync(rootDir); + + geminiClient = { + generateJson: mockGenerateJson, + }; + + mockConfig = { + getGeminiClient: vi.fn().mockReturnValue(geminiClient), + getTargetDir: () => rootDir, + getApprovalMode: vi.fn(), + setApprovalMode: vi.fn(), + getWorkspaceContext: () => createMockWorkspaceContext(rootDir), + getFileSystemService: () => new StandardFileSystemService(), + getIdeClient: () => undefined, + getIdeMode: () => false, + getApiKey: () => 'test-api-key', + getModel: () => 'test-model', + getSandbox: () => false, + getDebugMode: () => false, + getQuestion: () => undefined, + getFullContext: () => false, + getToolDiscoveryCommand: () => undefined, + getToolCallCommand: () => undefined, + getMcpServerCommand: () => undefined, + getMcpServers: () => undefined, + getUserAgent: () => 'test-agent', + getUserMemory: () => '', + setUserMemory: vi.fn(), + getGeminiMdFileCount: () => 0, + setGeminiMdFileCount: vi.fn(), + getToolRegistry: () => ({}) as any, + } as unknown as Config; + + (mockConfig.getApprovalMode as Mock).mockClear(); + (mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.DEFAULT); + + mockFixLLMEditWithInstruction.mockReset(); + mockFixLLMEditWithInstruction.mockResolvedValue({ + noChangesRequired: false, + search: '', + replace: '', + explanation: 'LLM fix failed', + }); + + mockGenerateJson.mockReset(); + mockGenerateJson.mockImplementation( + async (contents: Content[], schema: SchemaUnion) => { + const userContent = contents.find((c: Content) => c.role === 'user'); + let promptText = ''; + if (userContent && userContent.parts) { + promptText = userContent.parts + .filter((p: Part) => typeof (p as any).text === 'string') + .map((p: Part) => (p as any).text) + .join('\n'); + } + const snippetMatch = promptText.match( + /Problematic target snippet:\n```\n([\s\S]*?)\n```/, + ); + const problematicSnippet = + snippetMatch && snippetMatch[1] ? snippetMatch[1] : ''; + + if (((schema as any).properties as any)?.corrected_target_snippet) { + return Promise.resolve({ + corrected_target_snippet: problematicSnippet, + }); + } + if (((schema as any).properties as any)?.corrected_new_string) { + const originalNewStringMatch = promptText.match( + /original_new_string \(what was intended to replace original_old_string\):\n```\n([\s\S]*?)\n```/, + ); + const originalNewString = + originalNewStringMatch && originalNewStringMatch[1] + ? originalNewStringMatch[1] + : ''; + return Promise.resolve({ corrected_new_string: originalNewString }); + } + return Promise.resolve({}); + }, + ); + + tool = new SmartEditTool(mockConfig); + }); + + afterEach(() => { + fs.rmSync(tempDir, { recursive: true, force: true }); + }); + + describe('applyReplacement', () => { + it('should return newString if isNewFile is true', () => { + expect(applyReplacement(null, 'old', 'new', true)).toBe('new'); + expect(applyReplacement('existing', 'old', 'new', true)).toBe('new'); + }); + + it('should replace oldString with newString in currentContent', () => { + expect(applyReplacement('hello old world old', 'old', 'new', false)).toBe( + 'hello new world new', + ); + }); + }); + + describe('calculateReplacement', () => { + const abortSignal = new AbortController().signal; + + it('should perform an exact replacement', async () => { + const content = 'hello world'; + const result = await calculateReplacement({ + params: { + file_path: 'test.txt', + instruction: 'test', + old_string: 'world', + new_string: 'moon', + }, + currentContent: content, + abortSignal, + }); + expect(result.newContent).toBe('hello moon'); + expect(result.occurrences).toBe(1); + }); + + it('should perform a flexible, whitespace-insensitive replacement', async () => { + const content = ' hello\n world\n'; + const result = await calculateReplacement({ + params: { + file_path: 'test.txt', + instruction: 'test', + old_string: 'hello\nworld', + new_string: 'goodbye\nmoon', + }, + currentContent: content, + abortSignal, + }); + expect(result.newContent).toBe(' goodbye\n moon\n'); + expect(result.occurrences).toBe(1); + }); + + it('should return 0 occurrences if no match is found', async () => { + const content = 'hello world'; + const result = await calculateReplacement({ + params: { + file_path: 'test.txt', + instruction: 'test', + old_string: 'nomatch', + new_string: 'moon', + }, + currentContent: content, + abortSignal, + }); + expect(result.newContent).toBe(content); + expect(result.occurrences).toBe(0); + }); + }); + + describe('validateToolParams', () => { + it('should return null for valid params', () => { + const params: EditToolParams = { + file_path: path.join(rootDir, 'test.txt'), + instruction: 'An instruction', + old_string: 'old', + new_string: 'new', + }; + expect(tool.validateToolParams(params)).toBeNull(); + }); + + it('should return error for relative path', () => { + const params: EditToolParams = { + file_path: 'test.txt', + instruction: 'An instruction', + old_string: 'old', + new_string: 'new', + }; + expect(tool.validateToolParams(params)).toMatch( + /File path must be absolute/, + ); + }); + }); + + describe('execute', () => { + const testFile = 'execute_me.txt'; + let filePath: string; + + beforeEach(() => { + filePath = path.join(rootDir, testFile); + }); + + it('should edit an existing file and return diff with fileName', async () => { + const initialContent = 'This is some old text.'; + const newContent = 'This is some new text.'; + fs.writeFileSync(filePath, initialContent, 'utf8'); + const params: EditToolParams = { + file_path: filePath, + instruction: 'Replace old with new', + old_string: 'old', + new_string: 'new', + }; + + const invocation = tool.build(params); + const result = await invocation.execute(new AbortController().signal); + + expect(result.llmContent).toMatch(/Successfully modified file/); + expect(fs.readFileSync(filePath, 'utf8')).toBe(newContent); + const display = result.returnDisplay as FileDiff; + expect(display.fileDiff).toMatch(initialContent); + expect(display.fileDiff).toMatch(newContent); + expect(display.fileName).toBe(testFile); + }); + + it('should return error if old_string is not found in file', async () => { + fs.writeFileSync(filePath, 'Some content.', 'utf8'); + const params: EditToolParams = { + file_path: filePath, + instruction: 'Replace non-existent text', + old_string: 'nonexistent', + new_string: 'replacement', + }; + const invocation = tool.build(params); + const result = await invocation.execute(new AbortController().signal); + expect(result.llmContent).toMatch(/0 occurrences found for old_string/); + expect(result.returnDisplay).toMatch( + /Failed to edit, could not find the string to replace./, + ); + expect(mockFixLLMEditWithInstruction).toHaveBeenCalled(); + }); + + it('should succeed if FixLLMEditWithInstruction corrects the params', async () => { + const initialContent = 'This is some original text.'; + const finalContent = 'This is some brand new text.'; + fs.writeFileSync(filePath, initialContent, 'utf8'); + const params: EditToolParams = { + file_path: filePath, + instruction: 'Replace original with brand new', + old_string: 'original text that is slightly wrong', // This will fail first + new_string: 'brand new text', + }; + + mockFixLLMEditWithInstruction.mockResolvedValueOnce({ + noChangesRequired: false, + search: 'original text', // The corrected search string + replace: 'brand new text', + explanation: 'Corrected the search string to match the file content.', + }); + + const invocation = tool.build(params); + const result = await invocation.execute(new AbortController().signal); + + expect(result.error).toBeUndefined(); + expect(result.llmContent).toMatch(/Successfully modified file/); + expect(fs.readFileSync(filePath, 'utf8')).toBe(finalContent); + expect(mockFixLLMEditWithInstruction).toHaveBeenCalledTimes(1); + }); + + it('should return NO_CHANGE if FixLLMEditWithInstruction determines no changes are needed', async () => { + const initialContent = 'The price is $100.'; + fs.writeFileSync(filePath, initialContent, 'utf8'); + const params: EditToolParams = { + file_path: filePath, + instruction: 'Ensure the price is $100', + old_string: 'price is $50', // Incorrect old string + new_string: 'price is $100', + }; + + mockFixLLMEditWithInstruction.mockResolvedValueOnce({ + noChangesRequired: true, + search: '', + replace: '', + explanation: 'The price is already correctly set to $100.', + }); + + const invocation = tool.build(params); + const result = await invocation.execute(new AbortController().signal); + + expect(result.error?.type).toBe(ToolErrorType.EDIT_NO_CHANGE); + expect(result.llmContent).toMatch(/A secondary check determined/); + expect(fs.readFileSync(filePath, 'utf8')).toBe(initialContent); // File is unchanged + }); + + it('should preserve CRLF line endings when editing a file', async () => { + const initialContent = 'line one\r\nline two\r\n'; + const newContent = 'line one\r\nline three\r\n'; + fs.writeFileSync(filePath, initialContent, 'utf8'); + const params: EditToolParams = { + file_path: filePath, + instruction: 'Replace two with three', + old_string: 'line two', + new_string: 'line three', + }; + + const invocation = tool.build(params); + await invocation.execute(new AbortController().signal); + + const finalContent = fs.readFileSync(filePath, 'utf8'); + expect(finalContent).toBe(newContent); + }); + + it('should create a new file with CRLF line endings if new_string has them', async () => { + const newContentWithCRLF = 'new line one\r\nnew line two\r\n'; + const params: EditToolParams = { + file_path: filePath, + instruction: 'Create a new file', + old_string: '', + new_string: newContentWithCRLF, + }; + + const invocation = tool.build(params); + await invocation.execute(new AbortController().signal); + + const finalContent = fs.readFileSync(filePath, 'utf8'); + expect(finalContent).toBe(newContentWithCRLF); + }); + }); + + describe('Error Scenarios', () => { + const testFile = 'error_test.txt'; + let filePath: string; + + beforeEach(() => { + filePath = path.join(rootDir, testFile); + }); + + it('should return FILE_NOT_FOUND error', async () => { + const params: EditToolParams = { + file_path: filePath, + instruction: 'test', + old_string: 'any', + new_string: 'new', + }; + const invocation = tool.build(params); + const result = await invocation.execute(new AbortController().signal); + expect(result.error?.type).toBe(ToolErrorType.FILE_NOT_FOUND); + }); + + it('should return ATTEMPT_TO_CREATE_EXISTING_FILE error', async () => { + fs.writeFileSync(filePath, 'existing content', 'utf8'); + const params: EditToolParams = { + file_path: filePath, + instruction: 'test', + old_string: '', + new_string: 'new content', + }; + const invocation = tool.build(params); + const result = await invocation.execute(new AbortController().signal); + expect(result.error?.type).toBe( + ToolErrorType.ATTEMPT_TO_CREATE_EXISTING_FILE, + ); + }); + + it('should return NO_OCCURRENCE_FOUND error', async () => { + fs.writeFileSync(filePath, 'content', 'utf8'); + const params: EditToolParams = { + file_path: filePath, + instruction: 'test', + old_string: 'not-found', + new_string: 'new', + }; + const invocation = tool.build(params); + const result = await invocation.execute(new AbortController().signal); + expect(result.error?.type).toBe(ToolErrorType.EDIT_NO_OCCURRENCE_FOUND); + }); + + it('should return EXPECTED_OCCURRENCE_MISMATCH error', async () => { + fs.writeFileSync(filePath, 'one one two', 'utf8'); + const params: EditToolParams = { + file_path: filePath, + instruction: 'test', + old_string: 'one', + new_string: 'new', + }; + const invocation = tool.build(params); + const result = await invocation.execute(new AbortController().signal); + expect(result.error?.type).toBe( + ToolErrorType.EDIT_EXPECTED_OCCURRENCE_MISMATCH, + ); + }); + }); + + describe('IDE mode', () => { + const testFile = 'edit_me.txt'; + let filePath: string; + let ideClient: any; + + beforeEach(() => { + filePath = path.join(rootDir, testFile); + ideClient = { + openDiff: vi.fn(), + getConnectionStatus: vi.fn().mockReturnValue({ + status: IDEConnectionStatus.Connected, + }), + }; + (mockConfig as any).getIdeMode = () => true; + (mockConfig as any).getIdeClient = () => ideClient; + }); + + it('should call ideClient.openDiff and update params on confirmation', async () => { + const initialContent = 'some old content here'; + const newContent = 'some new content here'; + const modifiedContent = 'some modified content here'; + fs.writeFileSync(filePath, initialContent); + const params: EditToolParams = { + file_path: filePath, + instruction: 'test', + old_string: 'old', + new_string: 'new', + }; + + ideClient.openDiff.mockResolvedValueOnce({ + status: 'accepted', + content: modifiedContent, + }); + + const invocation = tool.build(params); + const confirmation = await invocation.shouldConfirmExecute( + new AbortController().signal, + ); + + expect(ideClient.openDiff).toHaveBeenCalledWith(filePath, newContent); + + if (confirmation && 'onConfirm' in confirmation) { + await confirmation.onConfirm(ToolConfirmationOutcome.ProceedOnce); + } + + expect(params.old_string).toBe(initialContent); + expect(params.new_string).toBe(modifiedContent); + }); + }); +}); diff --git a/packages/core/src/tools/smart-edit.ts b/packages/core/src/tools/smart-edit.ts new file mode 100644 index 0000000000..a084da5059 --- /dev/null +++ b/packages/core/src/tools/smart-edit.ts @@ -0,0 +1,846 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as fs from 'node:fs'; +import * as path from 'node:path'; +import * as Diff from 'diff'; +import { + BaseDeclarativeTool, + Kind, + type ToolCallConfirmationDetails, + ToolConfirmationOutcome, + type ToolEditConfirmationDetails, + type ToolInvocation, + type ToolLocation, + type ToolResult, + type ToolResultDisplay, +} from './tools.js'; +import { ToolErrorType } from './tool-error.js'; +import { makeRelative, shortenPath } from '../utils/paths.js'; +import { isNodeError } from '../utils/errors.js'; +import { type Config, ApprovalMode } from '../config/config.js'; +import { DEFAULT_DIFF_OPTIONS, getDiffStat } from './diffOptions.js'; +import { ReadFileTool } from './read-file.js'; +import { + type ModifiableDeclarativeTool, + type ModifyContext, +} from './modifiable-tool.js'; +import { IDEConnectionStatus } from '../ide/ide-client.js'; +import { FixLLMEditWithInstruction } from '../utils/llm-edit-fixer.js'; + +export function 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); +} + +interface ReplacementContext { + params: EditToolParams; + currentContent: string; + abortSignal: AbortSignal; +} + +interface ReplacementResult { + newContent: string; + occurrences: number; + finalOldString: string; + finalNewString: string; +} + +function restoreTrailingNewline( + originalContent: string, + modifiedContent: string, +): string { + const hadTrailingNewline = originalContent.endsWith('\n'); + if (hadTrailingNewline && !modifiedContent.endsWith('\n')) { + return modifiedContent + '\n'; + } else if (!hadTrailingNewline && modifiedContent.endsWith('\n')) { + return modifiedContent.replace(/\n$/, ''); + } + return modifiedContent; +} + +async function calculateExactReplacement( + context: ReplacementContext, +): Promise { + const { currentContent, params } = context; + const { old_string, new_string } = params; + + const normalizedCode = currentContent; + const normalizedSearch = old_string.replace(/\r\n/g, '\n'); + const normalizedReplace = new_string.replace(/\r\n/g, '\n'); + + const exactOccurrences = normalizedCode.split(normalizedSearch).length - 1; + if (exactOccurrences > 0) { + let modifiedCode = normalizedCode.replaceAll( + normalizedSearch, + normalizedReplace, + ); + modifiedCode = restoreTrailingNewline(currentContent, modifiedCode); + return { + newContent: modifiedCode, + occurrences: exactOccurrences, + finalOldString: normalizedSearch, + finalNewString: normalizedReplace, + }; + } + + return null; +} + +async function calculateFlexibleReplacement( + context: ReplacementContext, +): Promise { + const { currentContent, params } = context; + const { old_string, new_string } = params; + + const normalizedCode = currentContent; + const normalizedSearch = old_string.replace(/\r\n/g, '\n'); + const normalizedReplace = new_string.replace(/\r\n/g, '\n'); + + const sourceLines = normalizedCode.match(/.*(?:\n|$)/g)?.slice(0, -1) ?? []; + const searchLinesStripped = normalizedSearch + .split('\n') + .map((line: string) => line.trim()); + const replaceLines = normalizedReplace.split('\n'); + + let flexibleOccurrences = 0; + let i = 0; + while (i <= sourceLines.length - searchLinesStripped.length) { + const window = sourceLines.slice(i, i + searchLinesStripped.length); + const windowStripped = window.map((line: string) => line.trim()); + const isMatch = windowStripped.every( + (line: string, index: number) => line === searchLinesStripped[index], + ); + + if (isMatch) { + flexibleOccurrences++; + const firstLineInMatch = window[0]; + const indentationMatch = firstLineInMatch.match(/^(\s*)/); + const indentation = indentationMatch ? indentationMatch[1] : ''; + const newBlockWithIndent = replaceLines.map( + (line: string) => `${indentation}${line}`, + ); + sourceLines.splice( + i, + searchLinesStripped.length, + newBlockWithIndent.join('\n'), + ); + i += replaceLines.length; + } else { + i++; + } + } + + if (flexibleOccurrences > 0) { + let modifiedCode = sourceLines.join(''); + modifiedCode = restoreTrailingNewline(currentContent, modifiedCode); + return { + newContent: modifiedCode, + occurrences: flexibleOccurrences, + finalOldString: normalizedSearch, + finalNewString: normalizedReplace, + }; + } + + return null; +} + +/** + * Detects the line ending style of a string. + * @param content The string content to analyze. + * @returns '\r\n' for Windows-style, '\n' for Unix-style. + */ +function detectLineEnding(content: string): '\r\n' | '\n' { + // If a Carriage Return is found, assume Windows-style endings. + // This is a simple but effective heuristic. + return content.includes('\r\n') ? '\r\n' : '\n'; +} + +export async function calculateReplacement( + context: ReplacementContext, +): Promise { + const { currentContent, params } = context; + const { old_string, new_string } = params; + const normalizedSearch = old_string.replace(/\r\n/g, '\n'); + const normalizedReplace = new_string.replace(/\r\n/g, '\n'); + + if (normalizedSearch === '') { + return { + newContent: currentContent, + occurrences: 0, + finalOldString: normalizedSearch, + finalNewString: normalizedReplace, + }; + } + + const exactResult = await calculateExactReplacement(context); + if (exactResult) { + return exactResult; + } + + const flexibleResult = await calculateFlexibleReplacement(context); + if (flexibleResult) { + return flexibleResult; + } + + return { + newContent: currentContent, + occurrences: 0, + finalOldString: normalizedSearch, + finalNewString: normalizedReplace, + }; +} + +export function getErrorReplaceResult( + params: EditToolParams, + occurrences: number, + expectedReplacements: number, + finalOldString: string, + finalNewString: string, +) { + let error: { display: string; raw: string; type: ToolErrorType } | undefined = + undefined; + 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 (${finalOldString}). Original old_string was (${params.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.`, + type: ToolErrorType.EDIT_NO_OCCURRENCE_FOUND, + }; + } else if (occurrences !== expectedReplacements) { + const occurrenceTerm = + expectedReplacements === 1 ? 'occurrence' : 'occurrences'; + + error = { + display: `Failed to edit, expected ${expectedReplacements} ${occurrenceTerm} but found ${occurrences}.`, + raw: `Failed to edit, Expected ${expectedReplacements} ${occurrenceTerm} but found ${occurrences} for old_string in file: ${params.file_path}`, + type: ToolErrorType.EDIT_EXPECTED_OCCURRENCE_MISMATCH, + }; + } else if (finalOldString === finalNewString) { + error = { + display: `No changes to apply. The old_string and new_string are identical.`, + raw: `No changes to apply. The old_string and new_string are identical in file: ${params.file_path}`, + type: ToolErrorType.EDIT_NO_CHANGE, + }; + } + return error; +} + +/** + * Parameters for the Edit tool + */ +export interface EditToolParams { + /** + * The absolute path to the file to modify + */ + file_path: string; + + /** + * The text to replace + */ + old_string: string; + + /** + * The text to replace it with + */ + new_string: string; + + /** + * The instruction for what needs to be done. + */ + instruction: string; + + /** + * Whether the edit was modified manually by the user. + */ + modified_by_user?: boolean; + + /** + * Initially proposed string. + */ + ai_proposed_string?: string; +} + +interface CalculatedEdit { + currentContent: string | null; + newContent: string; + occurrences: number; + error?: { display: string; raw: string; type: ToolErrorType }; + isNewFile: boolean; + originalLineEnding: '\r\n' | '\n'; +} + +class EditToolInvocation implements ToolInvocation { + constructor( + private readonly config: Config, + public params: EditToolParams, + ) {} + + toolLocations(): ToolLocation[] { + return [{ path: this.params.file_path }]; + } + + private async attemptSelfCorrection( + params: EditToolParams, + currentContent: string, + initialError: { display: string; raw: string; type: ToolErrorType }, + abortSignal: AbortSignal, + originalLineEnding: '\r\n' | '\n', + ): Promise { + const fixedEdit = await FixLLMEditWithInstruction( + params.instruction, + params.old_string, + params.new_string, + initialError.raw, + currentContent, + this.config.getGeminiClient(), + abortSignal, + ); + + if (fixedEdit.noChangesRequired) { + return { + currentContent, + newContent: currentContent, + occurrences: 0, + isNewFile: false, + error: { + display: `No changes required. The file already meets the specified conditions.`, + raw: `A secondary check determined that no changes were necessary to fulfill the instruction. Explanation: ${fixedEdit.explanation}. Original error with the parameters given: ${initialError.raw}`, + type: ToolErrorType.EDIT_NO_CHANGE, + }, + originalLineEnding, + }; + } + + const secondAttemptResult = await calculateReplacement({ + params: { + ...params, + old_string: fixedEdit.search, + new_string: fixedEdit.replace, + }, + currentContent, + abortSignal, + }); + + const secondError = getErrorReplaceResult( + params, + secondAttemptResult.occurrences, + 1, // expectedReplacements is always 1 for smart_edit + secondAttemptResult.finalOldString, + secondAttemptResult.finalNewString, + ); + + if (secondError) { + // The fix failed, return the original error + return { + currentContent, + newContent: currentContent, + occurrences: 0, + isNewFile: false, + error: initialError, + originalLineEnding, + }; + } + + return { + currentContent, + newContent: secondAttemptResult.newContent, + occurrences: secondAttemptResult.occurrences, + isNewFile: false, + error: undefined, + originalLineEnding, + }; + } + + /** + * Calculates the potential outcome of an edit operation. + * @param params Parameters for the edit operation + * @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 { + const expectedReplacements = 1; + let currentContent: string | null = null; + let fileExists = false; + let originalLineEnding: '\r\n' | '\n' = '\n'; // Default for new files + + try { + currentContent = await this.config + .getFileSystemService() + .readTextFile(params.file_path); + originalLineEnding = detectLineEnding(currentContent); + currentContent = currentContent.replace(/\r\n/g, '\n'); + fileExists = true; + } catch (err: unknown) { + if (!isNodeError(err) || err.code !== 'ENOENT') { + throw err; + } + fileExists = false; + } + + const isNewFile = params.old_string === '' && !fileExists; + + if (isNewFile) { + return { + currentContent, + newContent: params.new_string, + occurrences: 1, + isNewFile: true, + error: undefined, + originalLineEnding, + }; + } + + // after this point, it's not a new file/edit + if (!fileExists) { + return { + currentContent, + newContent: '', + occurrences: 0, + isNewFile: false, + error: { + display: `File not found. Cannot apply edit. Use an empty old_string to create a new file.`, + raw: `File not found: ${params.file_path}`, + type: ToolErrorType.FILE_NOT_FOUND, + }, + originalLineEnding, + }; + } + + if (currentContent === null) { + return { + currentContent, + newContent: '', + occurrences: 0, + isNewFile: false, + error: { + display: `Failed to read content of file.`, + raw: `Failed to read content of existing file: ${params.file_path}`, + type: ToolErrorType.READ_CONTENT_FAILURE, + }, + originalLineEnding, + }; + } + + if (params.old_string === '') { + return { + currentContent, + newContent: currentContent, + occurrences: 0, + isNewFile: false, + error: { + display: `Failed to edit. Attempted to create a file that already exists.`, + raw: `File already exists, cannot create: ${params.file_path}`, + type: ToolErrorType.ATTEMPT_TO_CREATE_EXISTING_FILE, + }, + originalLineEnding, + }; + } + + const replacementResult = await calculateReplacement({ + params, + currentContent, + abortSignal, + }); + + const initialError = getErrorReplaceResult( + params, + replacementResult.occurrences, + expectedReplacements, + replacementResult.finalOldString, + replacementResult.finalNewString, + ); + + if (!initialError) { + return { + currentContent, + newContent: replacementResult.newContent, + occurrences: replacementResult.occurrences, + isNewFile: false, + error: undefined, + originalLineEnding, + }; + } + + // If there was an error, try to self-correct. + return this.attemptSelfCorrection( + params, + currentContent, + initialError, + abortSignal, + originalLineEnding, + ); + } + + /** + * 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, + ): Promise { + if (this.config.getApprovalMode() === ApprovalMode.AUTO_EDIT) { + return false; + } + + let editData: CalculatedEdit; + try { + editData = await this.calculateEdit(this.params, abortSignal); + } catch (error) { + const errorMsg = error instanceof Error ? error.message : String(error); + console.log(`Error preparing edit: ${errorMsg}`); + return false; + } + + if (editData.error) { + console.log(`Error: ${editData.error.display}`); + return false; + } + + const fileName = path.basename(this.params.file_path); + const fileDiff = Diff.createPatch( + fileName, + editData.currentContent ?? '', + editData.newContent, + 'Current', + 'Proposed', + DEFAULT_DIFF_OPTIONS, + ); + const ideClient = this.config.getIdeClient(); + const ideConfirmation = + this.config.getIdeMode() && + ideClient?.getConnectionStatus().status === IDEConnectionStatus.Connected + ? ideClient.openDiff(this.params.file_path, editData.newContent) + : undefined; + + const confirmationDetails: ToolEditConfirmationDetails = { + type: 'edit', + title: `Confirm Edit: ${shortenPath(makeRelative(this.params.file_path, this.config.getTargetDir()))}`, + fileName, + filePath: this.params.file_path, + fileDiff, + originalContent: editData.currentContent, + newContent: editData.newContent, + onConfirm: async (outcome: ToolConfirmationOutcome) => { + if (outcome === ToolConfirmationOutcome.ProceedAlways) { + this.config.setApprovalMode(ApprovalMode.AUTO_EDIT); + } + + if (ideConfirmation) { + const result = await ideConfirmation; + if (result.status === 'accepted' && result.content) { + // TODO(chrstn): See https://github.com/google-gemini/gemini-cli/pull/5618#discussion_r2255413084 + // for info on a possible race condition where the file is modified on disk while being edited. + this.params.old_string = editData.currentContent ?? ''; + this.params.new_string = result.content; + } + } + }, + ideConfirmation, + }; + return confirmationDetails; + } + + getDescription(): string { + const relativePath = makeRelative( + this.params.file_path, + this.config.getTargetDir(), + ); + if (this.params.old_string === '') { + return `Create ${shortenPath(relativePath)}`; + } + + const oldStringSnippet = + this.params.old_string.split('\n')[0].substring(0, 30) + + (this.params.old_string.length > 30 ? '...' : ''); + const newStringSnippet = + this.params.new_string.split('\n')[0].substring(0, 30) + + (this.params.new_string.length > 30 ? '...' : ''); + + if (this.params.old_string === this.params.new_string) { + return `No file changes to ${shortenPath(relativePath)}`; + } + return `${shortenPath(relativePath)}: ${oldStringSnippet} => ${newStringSnippet}`; + } + + /** + * Executes the edit operation with the given parameters. + * @param params Parameters for the edit operation + * @returns Result of the edit operation + */ + async execute(signal: AbortSignal): Promise { + let editData: CalculatedEdit; + try { + editData = await this.calculateEdit(this.params, signal); + } catch (error) { + const errorMsg = error instanceof Error ? error.message : String(error); + return { + llmContent: `Error preparing edit: ${errorMsg}`, + returnDisplay: `Error preparing edit: ${errorMsg}`, + error: { + message: errorMsg, + type: ToolErrorType.EDIT_PREPARATION_FAILURE, + }, + }; + } + + if (editData.error) { + return { + llmContent: editData.error.raw, + returnDisplay: `Error: ${editData.error.display}`, + error: { + message: editData.error.raw, + type: editData.error.type, + }, + }; + } + + try { + this.ensureParentDirectoriesExist(this.params.file_path); + let finalContent = editData.newContent; + + // Restore original line endings if they were CRLF + if (!editData.isNewFile && editData.originalLineEnding === '\r\n') { + finalContent = finalContent.replace(/\n/g, '\r\n'); + } + await this.config + .getFileSystemService() + .writeTextFile(this.params.file_path, finalContent); + + let displayResult: ToolResultDisplay; + if (editData.isNewFile) { + displayResult = `Created ${shortenPath(makeRelative(this.params.file_path, this.config.getTargetDir()))}`; + } else { + // Generate diff for display, even though core logic doesn't technically need it + // The CLI wrapper will use this part of the ToolResult + const fileName = path.basename(this.params.file_path); + const fileDiff = Diff.createPatch( + fileName, + editData.currentContent ?? '', // Should not be null here if not isNewFile + editData.newContent, + 'Current', + 'Proposed', + DEFAULT_DIFF_OPTIONS, + ); + const originallyProposedContent = + this.params.ai_proposed_string || this.params.new_string; + const diffStat = getDiffStat( + fileName, + editData.currentContent ?? '', + originallyProposedContent, + this.params.new_string, + ); + displayResult = { + fileDiff, + fileName, + originalContent: editData.currentContent, + newContent: editData.newContent, + diffStat, + }; + } + + const llmSuccessMessageParts = [ + editData.isNewFile + ? `Created new file: ${this.params.file_path} with provided content.` + : `Successfully modified file: ${this.params.file_path} (${editData.occurrences} replacements).`, + ]; + if (this.params.modified_by_user) { + llmSuccessMessageParts.push( + `User modified the \`new_string\` content to be: ${this.params.new_string}.`, + ); + } + + return { + llmContent: llmSuccessMessageParts.join(' '), + returnDisplay: displayResult, + }; + } catch (error) { + const errorMsg = error instanceof Error ? error.message : String(error); + return { + llmContent: `Error executing edit: ${errorMsg}`, + returnDisplay: `Error writing file: ${errorMsg}`, + error: { + message: errorMsg, + type: ToolErrorType.FILE_WRITE_FAILURE, + }, + }; + } + } + + /** + * Creates parent directories if they don't exist + */ + private ensureParentDirectoriesExist(filePath: string): void { + const dirName = path.dirname(filePath); + if (!fs.existsSync(dirName)) { + fs.mkdirSync(dirName, { recursive: true }); + } + } +} + +/** + * Implementation of the Edit tool logic + */ +export class SmartEditTool + extends BaseDeclarativeTool + implements ModifiableDeclarativeTool +{ + static readonly Name = 'smart_edit'; + + constructor(private readonly config: Config) { + super( + SmartEditTool.Name, + 'Edit', + `Replaces text within a file. Replaces a single occurrence. This tool requires providing significant context around the change to ensure precise targeting. Always use the ${ReadFileTool.Name} tool to examine the file's current content before attempting a text replacement. + + The user has the ability to modify the \`new_string\` content. If modified, this will be stated in the response. + + Expectation for required parameters: + 1. \`file_path\` MUST be an absolute path; otherwise an error will be thrown. + 2. \`old_string\` MUST be the exact literal text to replace (including all whitespace, indentation, newlines, and surrounding code etc.). + 3. \`new_string\` MUST be the exact literal text to replace \`old_string\` with (also including all whitespace, indentation, newlines, and surrounding code etc.). Ensure the resulting code is correct and idiomatic and that \`old_string\` and \`new_string\` are different. + 4. \`instruction\` is the detailed instruction of what needs to be changed. It is important to Make it specific and detailed so developers or large language models can understand what needs to be changed and perform the changes on their own if necessary. + 5. NEVER escape \`old_string\` or \`new_string\`, that would break the exact literal text requirement. + **Important:** If ANY of the above are not satisfied, the tool will fail. CRITICAL for \`old_string\`: Must uniquely identify the single instance to change. Include at least 3 lines of context BEFORE and AFTER the target text, matching whitespace and indentation precisely. If this string matches multiple locations, or does not match exactly, the tool will fail. + 6. Prefer to break down complex and long changes into multiple smaller atomic calls to this tool. Always check the content of the file after changes or not finding a string to match. + **Multiple replacements:** If there are multiple and ambiguous occurences of the \`old_string\` in the file, the tool will also fail.`, + Kind.Edit, + { + properties: { + file_path: { + description: + "The absolute path to the file to modify. Must start with '/'.", + type: 'string', + }, + instruction: { + description: `A clear, semantic instruction for the code change, acting as a high-quality prompt for an expert LLM assistant. It must be self-contained and explain the goal of the change. + +A good instruction should concisely answer: +1. WHY is the change needed? (e.g., "To fix a bug where users can be null...") +2. WHERE should the change happen? (e.g., "...in the 'renderUserProfile' function...") +3. WHAT is the high-level change? (e.g., "...add a null check for the 'user' object...") +4. WHAT is the desired outcome? (e.g., "...so that it displays a loading spinner instead of crashing.") + +**GOOD Example:** "In the 'calculateTotal' function, correct the sales tax calculation by updating the 'taxRate' constant from 0.05 to 0.075 to reflect the new regional tax laws." + +**BAD Examples:** +- "Change the text." (Too vague) +- "Fix the bug." (Doesn't explain the bug or the fix) +- "Replace the line with this new line." (Brittle, just repeats the other parameters) +`, + type: 'string', + }, + old_string: { + description: + 'The exact literal text to replace, preferably unescaped. Include at least 3 lines of context BEFORE and AFTER the target text, matching whitespace and indentation precisely. If this string is not the exact literal text (i.e. you escaped it) or does not match exactly, the tool will fail.', + type: 'string', + }, + new_string: { + description: + 'The exact literal text to replace `old_string` with, preferably unescaped. Provide the EXACT text. Ensure the resulting code is correct and idiomatic.', + type: 'string', + }, + }, + required: ['file_path', 'instruction', 'old_string', 'new_string'], + type: 'object', + }, + ); + } + + /** + * Validates the parameters for the Edit tool + * @param params Parameters to validate + * @returns Error message string or null if valid + */ + protected override validateToolParamValues( + params: EditToolParams, + ): string | null { + if (!params.file_path) { + return "The 'file_path' parameter must be non-empty."; + } + + if (!path.isAbsolute(params.file_path)) { + return `File path must be absolute: ${params.file_path}`; + } + + const workspaceContext = this.config.getWorkspaceContext(); + if (!workspaceContext.isPathWithinWorkspace(params.file_path)) { + const directories = workspaceContext.getDirectories(); + return `File path must be within one of the workspace directories: ${directories.join(', ')}`; + } + + return null; + } + + protected createInvocation( + params: EditToolParams, + ): ToolInvocation { + return new EditToolInvocation(this.config, params); + } + + getModifyContext(_: AbortSignal): ModifyContext { + return { + getFilePath: (params: EditToolParams) => params.file_path, + getCurrentContent: async (params: EditToolParams): Promise => { + try { + return this.config + .getFileSystemService() + .readTextFile(params.file_path); + } catch (err) { + if (!isNodeError(err) || err.code !== 'ENOENT') throw err; + return ''; + } + }, + getProposedContent: async (params: EditToolParams): Promise => { + try { + const currentContent = await this.config + .getFileSystemService() + .readTextFile(params.file_path); + return applyReplacement( + currentContent, + params.old_string, + params.new_string, + params.old_string === '' && currentContent === '', + ); + } catch (err) { + if (!isNodeError(err) || err.code !== 'ENOENT') throw err; + return ''; + } + }, + createUpdatedParams: ( + oldContent: string, + modifiedProposedContent: string, + originalParams: EditToolParams, + ): EditToolParams => { + const content = originalParams.new_string; + return { + ...originalParams, + ai_proposed_string: content, + old_string: oldContent, + new_string: modifiedProposedContent, + modified_by_user: true, + }; + }, + }; + } +} diff --git a/packages/core/src/utils/llm-edit-fixer.ts b/packages/core/src/utils/llm-edit-fixer.ts new file mode 100644 index 0000000000..95496d4779 --- /dev/null +++ b/packages/core/src/utils/llm-edit-fixer.ts @@ -0,0 +1,145 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { type Content, Type } from '@google/genai'; +import { type GeminiClient } from '../core/client.js'; +import { LruCache } from './LruCache.js'; +import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js'; + +const MAX_CACHE_SIZE = 50; + +const EDIT_SYS_PROMPT = ` +You are an expert code-editing assistant specializing in debugging and correcting failed search-and-replace operations. + +# Primary Goal +Your task is to analyze a failed edit attempt and provide a corrected \`search\` string that will match the text in the file precisely. The correction should be as minimal as possible, staying very close to the original, failed \`search\` string. Do NOT invent a completely new edit based on the instruction; your job is to fix the provided parameters. + +It is important that you do no try to figure out if the instruction is correct. DO NOT GIVE ADVICE. Your only goal here is to do your best to perform the search and replace task! + +# Input Context +You will be given: +1. The high-level instruction for the original edit. +2. The exact \`search\` and \`replace\` strings that failed. +3. The error message that was produced. +4. The full content of the source file. + +# Rules for Correction +1. **Minimal Correction:** Your new \`search\` string must be a close variation of the original. Focus on fixing issues like whitespace, indentation, line endings, or small contextual differences. +2. **Explain the Fix:** Your \`explanation\` MUST state exactly why the original \`search\` failed and how your new \`search\` string resolves that specific failure. (e.g., "The original search failed due to incorrect indentation; the new search corrects the indentation to match the source file."). +3. **Preserve the \`replace\` String:** Do NOT modify the \`replace\` string unless the instruction explicitly requires it and it was the source of the error. Your primary focus is fixing the \`search\` string. +4. **No Changes Case:** CRUCIAL: if the change is already present in the file, set \`noChangesRequired\` to True and explain why in the \`explanation\`. It is crucial that you only do this if the changes outline in \`replace\` are alredy in the file and suits the instruction!! +5. **Exactness:** The final \`search\` field must be the EXACT literal text from the file. Do not escape characters. +`; + +const EDIT_USER_PROMPT = ` +# Goal of the Original Edit + +{instruction} + + +# Failed Attempt Details +- **Original \`search\` parameter (failed):** + +{old_string} + +- **Original \`replace\` parameter:** + +{new_string} + +- **Error Encountered:** + +{error} + + +# Full File Content + +{current_content} + + +# Your Task +Based on the error and the file content, provide a corrected \`search\` string that will succeed. Remember to keep your correction minimal and explain the precise reason for the failure in your \`explanation\`. +`; + +export interface SearchReplaceEdit { + search: string; + replace: string; + noChangesRequired: boolean; + explanation: string; +} + +const SearchReplaceEditSchema = { + type: Type.OBJECT, + properties: { + explanation: { type: Type.STRING }, + search: { type: Type.STRING }, + replace: { type: Type.STRING }, + noChangesRequired: { type: Type.BOOLEAN }, + }, + required: ['search', 'replace', 'explanation'], +}; + +const editCorrectionWithInstructionCache = new LruCache< + string, + SearchReplaceEdit +>(MAX_CACHE_SIZE); + +/** + * Attempts to fix a failed edit by using an LLM to generate a new search and replace pair. + * @param instruction The instruction for what needs to be done. + * @param old_string The original string to be replaced. + * @param new_string The original replacement string. + * @param error The error that occurred during the initial edit. + * @param current_content The current content of the file. + * @param geminiClient The Gemini client to use for the LLM call. + * @param abortSignal An abort signal to cancel the operation. + * @returns A new search and replace pair. + */ +export async function FixLLMEditWithInstruction( + instruction: string, + old_string: string, + new_string: string, + error: string, + current_content: string, + geminiClient: GeminiClient, + abortSignal: AbortSignal, +): Promise { + const cacheKey = `${instruction}---${old_string}---${new_string}--${current_content}--${error}`; + const cachedResult = editCorrectionWithInstructionCache.get(cacheKey); + if (cachedResult) { + return cachedResult; + } + const userPrompt = EDIT_USER_PROMPT.replace('{instruction}', instruction) + .replace('{old_string}', old_string) + .replace('{new_string}', new_string) + .replace('{error}', error) + .replace('{current_content}', current_content); + + const contents: Content[] = [ + { + role: 'user', + parts: [ + { + text: `${EDIT_SYS_PROMPT} +${userPrompt}`, + }, + ], + }, + ]; + + const result = (await geminiClient.generateJson( + contents, + SearchReplaceEditSchema, + abortSignal, + DEFAULT_GEMINI_FLASH_MODEL, + )) as unknown as SearchReplaceEdit; + + editCorrectionWithInstructionCache.set(cacheKey, result); + return result; +} + +export function resetLlmEditFixerCaches_TEST_ONLY() { + editCorrectionWithInstructionCache.clear(); +}