From 8969a232eca176b1d596c103435c3e6b0523ce79 Mon Sep 17 00:00:00 2001 From: Shreya Keshive Date: Thu, 11 Sep 2025 16:17:57 -0400 Subject: [PATCH] feat(ide): Check for IDE diffing capabilities before opening diffs (#8266) --- .../messages/ToolConfirmationMessage.tsx | 40 ++++++--- packages/core/src/ide/ide-client.test.ts | 90 +++++++++++++++++++ packages/core/src/ide/ide-client.ts | 51 +++++++++++ packages/core/src/tools/edit.test.ts | 10 +-- packages/core/src/tools/edit.ts | 5 +- packages/core/src/tools/smart-edit.test.ts | 10 +-- packages/core/src/tools/smart-edit.ts | 5 +- packages/core/src/tools/write-file.test.ts | 17 +--- packages/core/src/tools/write-file.ts | 5 +- 9 files changed, 183 insertions(+), 50 deletions(-) diff --git a/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx b/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx index 427e3fb38c..dfcaddfedf 100644 --- a/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx +++ b/packages/cli/src/ui/components/messages/ToolConfirmationMessage.tsx @@ -5,6 +5,7 @@ */ import type React from 'react'; +import { useEffect, useState } from 'react'; import { Box, Text } from 'ink'; import { DiffRenderer } from './DiffRenderer.js'; import { RenderInline } from '../../utils/InlineMarkdownRenderer.js'; @@ -41,12 +42,31 @@ export const ToolConfirmationMessage: React.FC< const { onConfirm } = confirmationDetails; const childWidth = terminalWidth - 2; // 2 for padding + const [ideClient, setIdeClient] = useState(null); + const [isDiffingEnabled, setIsDiffingEnabled] = useState(false); + + useEffect(() => { + let isMounted = true; + if (config.getIdeMode()) { + const getIdeClient = async () => { + const client = await IdeClient.getInstance(); + if (isMounted) { + setIdeClient(client); + setIsDiffingEnabled(client?.isDiffingEnabled() ?? false); + } + }; + getIdeClient(); + } + return () => { + isMounted = false; + }; + }, [config]); + const handleConfirm = async (outcome: ToolConfirmationOutcome) => { if (confirmationDetails.type === 'edit') { - if (config.getIdeMode()) { + if (config.getIdeMode() && isDiffingEnabled) { const cliOutcome = outcome === ToolConfirmationOutcome.Cancel ? 'rejected' : 'accepted'; - const ideClient = await IdeClient.getInstance(); await ideClient?.resolveDiffFromCli( confirmationDetails.filePath, cliOutcome, @@ -137,22 +157,18 @@ export const ToolConfirmationMessage: React.FC< value: ToolConfirmationOutcome.ProceedAlways, }); } - if (config.getIdeMode()) { - options.push({ - label: 'No (esc)', - value: ToolConfirmationOutcome.Cancel, - }); - } else { + if (!config.getIdeMode() || !isDiffingEnabled) { options.push({ label: 'Modify with external editor', value: ToolConfirmationOutcome.ModifyWithEditor, }); - options.push({ - label: 'No, suggest changes (esc)', - value: ToolConfirmationOutcome.Cancel, - }); } + options.push({ + label: 'No, suggest changes (esc)', + value: ToolConfirmationOutcome.Cancel, + }); + bodyContent = ( { close: vi.fn(), setNotificationHandler: vi.fn(), callTool: vi.fn(), + request: vi.fn(), } as unknown as Mocked; mockHttpTransport = { close: vi.fn(), @@ -534,4 +535,93 @@ describe('IdeClient', () => { ); }); }); + + describe('isDiffingEnabled', () => { + it('should return false if not connected', async () => { + const ideClient = await IdeClient.getInstance(); + expect(ideClient.isDiffingEnabled()).toBe(false); + }); + + it('should return false if tool discovery fails', async () => { + const config = { port: '8080' }; + vi.mocked(fs.promises.readFile).mockResolvedValue(JSON.stringify(config)); + ( + vi.mocked(fs.promises.readdir) as Mock< + (path: fs.PathLike) => Promise + > + ).mockResolvedValue([]); + mockClient.request.mockRejectedValue(new Error('Method not found')); + + const ideClient = await IdeClient.getInstance(); + await ideClient.connect(); + + expect(ideClient.getConnectionStatus().status).toBe( + IDEConnectionStatus.Connected, + ); + expect(ideClient.isDiffingEnabled()).toBe(false); + }); + + it('should return false if diffing tools are not available', async () => { + const config = { port: '8080' }; + vi.mocked(fs.promises.readFile).mockResolvedValue(JSON.stringify(config)); + ( + vi.mocked(fs.promises.readdir) as Mock< + (path: fs.PathLike) => Promise + > + ).mockResolvedValue([]); + mockClient.request.mockResolvedValue({ + tools: [{ name: 'someOtherTool' }], + }); + + const ideClient = await IdeClient.getInstance(); + await ideClient.connect(); + + expect(ideClient.getConnectionStatus().status).toBe( + IDEConnectionStatus.Connected, + ); + expect(ideClient.isDiffingEnabled()).toBe(false); + }); + + it('should return false if only openDiff tool is available', async () => { + const config = { port: '8080' }; + vi.mocked(fs.promises.readFile).mockResolvedValue(JSON.stringify(config)); + ( + vi.mocked(fs.promises.readdir) as Mock< + (path: fs.PathLike) => Promise + > + ).mockResolvedValue([]); + mockClient.request.mockResolvedValue({ + tools: [{ name: 'openDiff' }], + }); + + const ideClient = await IdeClient.getInstance(); + await ideClient.connect(); + + expect(ideClient.getConnectionStatus().status).toBe( + IDEConnectionStatus.Connected, + ); + expect(ideClient.isDiffingEnabled()).toBe(false); + }); + + it('should return true if connected and diffing tools are available', async () => { + const config = { port: '8080' }; + vi.mocked(fs.promises.readFile).mockResolvedValue(JSON.stringify(config)); + ( + vi.mocked(fs.promises.readdir) as Mock< + (path: fs.PathLike) => Promise + > + ).mockResolvedValue([]); + mockClient.request.mockResolvedValue({ + tools: [{ name: 'openDiff' }, { name: 'closeDiff' }], + }); + + const ideClient = await IdeClient.getInstance(); + await ideClient.connect(); + + expect(ideClient.getConnectionStatus().status).toBe( + IDEConnectionStatus.Connected, + ); + expect(ideClient.isDiffingEnabled()).toBe(true); + }); + }); }); diff --git a/packages/core/src/ide/ide-client.ts b/packages/core/src/ide/ide-client.ts index bcba4b81d5..7c4b50d8ba 100644 --- a/packages/core/src/ide/ide-client.ts +++ b/packages/core/src/ide/ide-client.ts @@ -22,6 +22,7 @@ import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js' import * as os from 'node:os'; import * as path from 'node:path'; import { EnvHttpProxyAgent } from 'undici'; +import { ListToolsResultSchema } from '@modelcontextprotocol/sdk/types.js'; const logger = { // eslint-disable-next-line @typescript-eslint/no-explicit-any @@ -78,6 +79,7 @@ export class IdeClient { private diffResponses = new Map void>(); private statusListeners = new Set<(state: IDEConnectionState) => void>(); private trustChangeListeners = new Set<(isTrusted: boolean) => void>(); + private availableTools: string[] = []; /** * A mutex to ensure that only one diff view is open in the IDE at a time. * This prevents race conditions and UI issues in IDEs like VSCode that @@ -334,6 +336,53 @@ export class IdeClient { return this.currentIdeDisplayName; } + isDiffingEnabled(): boolean { + return ( + !!this.client && + this.state.status === IDEConnectionStatus.Connected && + this.availableTools.includes('openDiff') && + this.availableTools.includes('closeDiff') + ); + } + + private async discoverTools(): Promise { + if (!this.client) { + return; + } + try { + logger.debug('Discovering tools from IDE...'); + const response = await this.client.request( + { method: 'tools/list', params: {} }, + ListToolsResultSchema, + ); + + // Map the array of tool objects to an array of tool names (strings) + this.availableTools = response.tools.map((tool) => tool.name); + + if (this.availableTools.length > 0) { + logger.debug( + `Discovered ${this.availableTools.length} tools from IDE: ${this.availableTools.join(', ')}`, + ); + } else { + logger.debug( + 'IDE supports tool discovery, but no tools are available.', + ); + } + } catch (error) { + // It's okay if this fails, the IDE might not support it. + // Don't log an error if the method is not found, which is a common case. + if ( + error instanceof Error && + !error.message?.includes('Method not found') + ) { + logger.error(`Error discovering tools from IDE: ${error.message}`); + } else { + logger.debug('IDE does not support tool discovery.'); + } + this.availableTools = []; + } + } + private setState( status: IDEConnectionStatus, details?: string, @@ -631,6 +680,7 @@ export class IdeClient { ); await this.client.connect(transport); this.registerClientHandlers(); + await this.discoverTools(); this.setState(IDEConnectionStatus.Connected); return true; } catch (_error) { @@ -664,6 +714,7 @@ export class IdeClient { }); await this.client.connect(transport); this.registerClientHandlers(); + await this.discoverTools(); this.setState(IDEConnectionStatus.Connected); return true; } catch (_error) { diff --git a/packages/core/src/tools/edit.test.ts b/packages/core/src/tools/edit.test.ts index 7b98083e31..49ed12e5a6 100644 --- a/packages/core/src/tools/edit.test.ts +++ b/packages/core/src/tools/edit.test.ts @@ -10,16 +10,12 @@ const mockEnsureCorrectEdit = vi.hoisted(() => vi.fn()); const mockGenerateJson = vi.hoisted(() => vi.fn()); const mockOpenDiff = vi.hoisted(() => vi.fn()); -import { IdeClient, IDEConnectionStatus } from '../ide/ide-client.js'; +import { IdeClient } from '../ide/ide-client.js'; vi.mock('../ide/ide-client.js', () => ({ IdeClient: { getInstance: vi.fn(), }, - IDEConnectionStatus: { - Connected: 'connected', - Disconnected: 'disconnected', - }, })); vi.mock('../utils/editCorrector.js', () => ({ @@ -896,9 +892,7 @@ describe('EditTool', () => { filePath = path.join(rootDir, testFile); ideClient = { openDiff: vi.fn(), - getConnectionStatus: vi.fn().mockReturnValue({ - status: IDEConnectionStatus.Connected, - }), + isDiffingEnabled: vi.fn().mockReturnValue(true), }; vi.mocked(IdeClient.getInstance).mockResolvedValue(ideClient); (mockConfig as any).getIdeMode = () => true; diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts index cb0418b1e4..2c17fe86f0 100644 --- a/packages/core/src/tools/edit.ts +++ b/packages/core/src/tools/edit.ts @@ -32,7 +32,7 @@ import type { ModifiableDeclarativeTool, ModifyContext, } from './modifiable-tool.js'; -import { IdeClient, IDEConnectionStatus } from '../ide/ide-client.js'; +import { IdeClient } from '../ide/ide-client.js'; export function applyReplacement( currentContent: string | null, @@ -268,8 +268,7 @@ class EditToolInvocation implements ToolInvocation { ); const ideClient = await IdeClient.getInstance(); const ideConfirmation = - this.config.getIdeMode() && - ideClient?.getConnectionStatus().status === IDEConnectionStatus.Connected + this.config.getIdeMode() && ideClient.isDiffingEnabled() ? ideClient.openDiff(this.params.file_path, editData.newContent) : undefined; diff --git a/packages/core/src/tools/smart-edit.test.ts b/packages/core/src/tools/smart-edit.test.ts index c72fcb48df..9ce42c506e 100644 --- a/packages/core/src/tools/smart-edit.test.ts +++ b/packages/core/src/tools/smart-edit.test.ts @@ -10,16 +10,12 @@ const mockFixLLMEditWithInstruction = vi.hoisted(() => vi.fn()); const mockGenerateJson = vi.hoisted(() => vi.fn()); const mockOpenDiff = vi.hoisted(() => vi.fn()); -import { IdeClient, IDEConnectionStatus } from '../ide/ide-client.js'; +import { IdeClient } from '../ide/ide-client.js'; vi.mock('../ide/ide-client.js', () => ({ IdeClient: { getInstance: vi.fn(), }, - IDEConnectionStatus: { - Connected: 'connected', - Disconnected: 'disconnected', - }, })); vi.mock('../utils/llm-edit-fixer.js', () => ({ @@ -461,9 +457,7 @@ describe('SmartEditTool', () => { filePath = path.join(rootDir, testFile); ideClient = { openDiff: vi.fn(), - getConnectionStatus: vi.fn().mockReturnValue({ - status: IDEConnectionStatus.Connected, - }), + isDiffingEnabled: vi.fn().mockReturnValue(true), }; vi.mocked(IdeClient.getInstance).mockResolvedValue(ideClient); (mockConfig as any).getIdeMode = () => true; diff --git a/packages/core/src/tools/smart-edit.ts b/packages/core/src/tools/smart-edit.ts index 6291296f13..a39d906d2e 100644 --- a/packages/core/src/tools/smart-edit.ts +++ b/packages/core/src/tools/smart-edit.ts @@ -28,7 +28,7 @@ import { type ModifiableDeclarativeTool, type ModifyContext, } from './modifiable-tool.js'; -import { IdeClient, IDEConnectionStatus } from '../ide/ide-client.js'; +import { IdeClient } from '../ide/ide-client.js'; import { FixLLMEditWithInstruction } from '../utils/llm-edit-fixer.js'; export function applyReplacement( @@ -528,8 +528,7 @@ class EditToolInvocation implements ToolInvocation { ); const ideClient = await IdeClient.getInstance(); const ideConfirmation = - this.config.getIdeMode() && - ideClient?.getConnectionStatus().status === IDEConnectionStatus.Connected + this.config.getIdeMode() && ideClient.isDiffingEnabled() ? ideClient.openDiff(this.params.file_path, editData.newContent) : undefined; diff --git a/packages/core/src/tools/write-file.test.ts b/packages/core/src/tools/write-file.test.ts index 5dcafe6249..adb6ebf751 100644 --- a/packages/core/src/tools/write-file.test.ts +++ b/packages/core/src/tools/write-file.test.ts @@ -33,7 +33,7 @@ import { } from '../utils/editCorrector.js'; import { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.js'; import { StandardFileSystemService } from '../services/fileSystemService.js'; -import { IdeClient, IDEConnectionStatus } from '../ide/ide-client.js'; +import { IdeClient } from '../ide/ide-client.js'; import type { DiffUpdateResult } from '../ide/ideContext.js'; const rootDir = path.resolve(os.tmpdir(), 'gemini-cli-test-root'); @@ -45,18 +45,13 @@ vi.mock('../ide/ide-client.js', () => ({ IdeClient: { getInstance: vi.fn(), }, - IDEConnectionStatus: { - Connected: 'connected', - Disconnected: 'disconnected', - Connecting: 'connecting', - }, })); let mockGeminiClientInstance: Mocked; const mockEnsureCorrectEdit = vi.fn(); const mockEnsureCorrectFileContent = vi.fn(); const mockIdeClient = { - getConnectionStatus: vi.fn(), openDiff: vi.fn(), + isDiffingEnabled: vi.fn(), }; // Wire up the mocked functions to be used by the actual module imports @@ -456,9 +451,7 @@ describe('WriteFileTool', () => { beforeEach(() => { // Enable IDE mode and set connection status for these tests mockConfigInternal.getIdeMode.mockReturnValue(true); - mockIdeClient.getConnectionStatus.mockReturnValue({ - status: IDEConnectionStatus.Connected, - }); + mockIdeClient.isDiffingEnabled.mockReturnValue(true); mockIdeClient.openDiff.mockResolvedValue({ status: 'accepted', content: 'ide-modified-content', @@ -495,9 +488,7 @@ describe('WriteFileTool', () => { }); it('should not call openDiff if IDE is not connected', async () => { - mockIdeClient.getConnectionStatus.mockReturnValue({ - status: IDEConnectionStatus.Disconnected, - }); + mockIdeClient.isDiffingEnabled.mockReturnValue(false); const filePath = path.join(rootDir, 'ide_disconnected_file.txt'); const params = { file_path: filePath, content: 'test' }; const invocation = tool.build(params); diff --git a/packages/core/src/tools/write-file.ts b/packages/core/src/tools/write-file.ts index 597043f05b..3253fc2d6e 100644 --- a/packages/core/src/tools/write-file.ts +++ b/packages/core/src/tools/write-file.ts @@ -35,7 +35,7 @@ import type { ModifiableDeclarativeTool, ModifyContext, } from './modifiable-tool.js'; -import { IdeClient, IDEConnectionStatus } from '../ide/ide-client.js'; +import { IdeClient } from '../ide/ide-client.js'; import { logFileOperation } from '../telemetry/loggers.js'; import { FileOperationEvent } from '../telemetry/types.js'; import { FileOperation } from '../telemetry/metrics.js'; @@ -195,8 +195,7 @@ class WriteFileToolInvocation extends BaseToolInvocation< const ideClient = await IdeClient.getInstance(); const ideConfirmation = - this.config.getIdeMode() && - ideClient.getConnectionStatus().status === IDEConnectionStatus.Connected + this.config.getIdeMode() && ideClient.isDiffingEnabled() ? ideClient.openDiff(this.params.file_path, correctedContent) : undefined;