feat(ide): Check for IDE diffing capabilities before opening diffs (google-gemini#8266)

skeshive · web-flow · commit 8969a232eca1 · 2025-09-11T20:17:57.000Z
diff --git 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<IdeClient | null>(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 = (
       <DiffRenderer
         diffContent={confirmationDetails.fileDiff}
diff --git a/packages/core/src/ide/ide-client.test.ts b/packages/core/src/ide/ide-client.test.ts
@@ -83,6 +83,7 @@ describe('IdeClient', () => {
       close: vi.fn(),
       setNotificationHandler: vi.fn(),
       callTool: vi.fn(),
+      request: vi.fn(),
     } as unknown as Mocked<Client>;
     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<string[]>
+        >
+      ).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<string[]>
+        >
+      ).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<string[]>
+        >
+      ).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<string[]>
+        >
+      ).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
@@ -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<string, (result: DiffUpdateResult) => 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<void> {
+    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
@@ -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
@@ -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<EditToolParams, ToolResult> {
     );
     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
@@ -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
@@ -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<EditToolParams, ToolResult> {
     );
     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
diff --git a/packages/core/src/tools/write-file.ts b/packages/core/src/tools/write-file.ts
