diff --git a/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx b/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx index 8b7f93d1..b4fe4167 100644 --- a/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx +++ b/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx @@ -138,8 +138,6 @@ export const ToolConfirmationMessage: React.FC< value: ToolConfirmationOutcome.Cancel, }); } else { - // TODO(chrstnb): support edit tool in IDE mode. - options.push({ label: 'Modify with external editor', value: ToolConfirmationOutcome.ModifyWithEditor, diff --git a/packages/core/src/tools/edit.test.ts b/packages/core/src/tools/edit.test.ts index 029d3a3c..3bfa023e 100644 --- a/packages/core/src/tools/edit.test.ts +++ b/packages/core/src/tools/edit.test.ts @@ -10,6 +10,8 @@ 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, })); @@ -26,7 +28,7 @@ vi.mock('../utils/editor.js', () => ({ import { describe, it, expect, beforeEach, afterEach, vi, Mock } from 'vitest'; import { EditTool, EditToolParams } from './edit.js'; -import { FileDiff } from './tools.js'; +import { FileDiff, ToolConfirmationOutcome } from './tools.js'; import { ToolErrorType } from './tool-error.js'; import path from 'path'; import fs from 'fs'; @@ -58,6 +60,9 @@ describe('EditTool', () => { getApprovalMode: vi.fn(), setApprovalMode: vi.fn(), getWorkspaceContext: () => createMockWorkspaceContext(rootDir), + getIdeClient: () => undefined, + getIdeMode: () => false, + getIdeModeFeature: () => 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: @@ -796,4 +801,57 @@ describe('EditTool', () => { expect(error).toContain(rootDir); }); }); + + 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).getIdeModeFeature = () => 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, + 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, + }); + + const confirmation = await tool.shouldConfirmExecute( + params, + 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/edit.ts b/packages/core/src/tools/edit.ts index 853ad4c1..43505182 100644 --- a/packages/core/src/tools/edit.ts +++ b/packages/core/src/tools/edit.ts @@ -27,6 +27,7 @@ import { ensureCorrectEdit } from '../utils/editCorrector.js'; import { DEFAULT_DIFF_OPTIONS } from './diffOptions.js'; import { ReadFileTool } from './read-file.js'; import { ModifiableDeclarativeTool, ModifyContext } from './modifiable-tool.js'; +import { IDEConnectionStatus } from '../ide/ide-client.js'; /** * Parameters for the Edit tool @@ -328,6 +329,14 @@ Expectation for required parameters: 'Proposed', DEFAULT_DIFF_OPTIONS, ); + const ideClient = this.config.getIdeClient(); + const ideConfirmation = + this.config.getIdeModeFeature() && + this.config.getIdeMode() && + ideClient?.getConnectionStatus().status === IDEConnectionStatus.Connected + ? ideClient.openDiff(params.file_path, editData.newContent) + : undefined; + const confirmationDetails: ToolEditConfirmationDetails = { type: 'edit', title: `Confirm Edit: ${shortenPath(makeRelative(params.file_path, this.config.getTargetDir()))}`, @@ -340,7 +349,18 @@ Expectation for required parameters: 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. + params.old_string = editData.currentContent ?? ''; + params.new_string = result.content; + } + } }, + ideConfirmation, }; return confirmationDetails; }