[ide-mode] Add openDiff tool to IDE MCP server (#4519)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

4 files changed, 381 insertions, 6 deletions

diff --git a/packages/vscode-ide-companion/package.json b/packages/vscode-ide-companion/package.json
index 254d8ac2..263f1b18 100644
--- a/packages/vscode-ide-companion/package.json
+++ b/packages/vscode-ide-companion/package.json
@@ -31,8 +31,23 @@
     "onStartupFinished"
   ],
   "contributes": {
+    "languages": [
+      {
+        "id": "gemini-diff-editable"
+      }
+    ],
     "commands": [
       {
+        "command": "gemini.diff.accept",
+        "title": "Gemini CLI: Accept Current Diff",
+        "icon": "$(check)"
+      },
+      {
+        "command": "gemini.diff.cancel",
+        "title": "Cancel",
+        "icon": "$(close)"
+      },
+      {
         "command": "gemini-cli.runGeminiCLI",
         "title": "Gemini CLI: Run"
       },
@@ -40,6 +55,42 @@
         "command": "gemini-cli.showNotices",
         "title": "Gemini CLI: View Third-Party Notices"
       }
+    ],
+    "menus": {
+      "commandPalette": [
+        {
+          "command": "gemini.diff.accept",
+          "when": "gemini.diff.isVisible"
+        },
+        {
+          "command": "gemini.diff.cancel",
+          "when": "gemini.diff.isVisible"
+        }
+      ],
+      "editor/title": [
+        {
+          "command": "gemini.diff.accept",
+          "when": "gemini.diff.isVisible",
+          "group": "navigation"
+        },
+        {
+          "command": "gemini.diff.cancel",
+          "when": "gemini.diff.isVisible",
+          "group": "navigation"
+        }
+      ]
+    },
+    "keybindings": [
+      {
+        "command": "gemini.diff.accept",
+        "key": "ctrl+s",
+        "when": "gemini.diff.isVisible"
+      },
+      {
+        "command": "gemini.diff.accept",
+        "key": "cmd+s",
+        "when": "gemini.diff.isVisible"
+      }
     ]
   },
   "main": "./dist/extension.cjs",
diff --git a/packages/vscode-ide-companion/src/diff-manager.ts b/packages/vscode-ide-companion/src/diff-manager.ts
new file mode 100644
index 00000000..159a6101
--- /dev/null
+++ b/packages/vscode-ide-companion/src/diff-manager.ts
@@ -0,0 +1,228 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as vscode from 'vscode';
+import * as path from 'node:path';
+import { DIFF_SCHEME } from './extension.js';
+import { type JSONRPCNotification } from '@modelcontextprotocol/sdk/types.js';
+
+export class DiffContentProvider implements vscode.TextDocumentContentProvider {
+  private content = new Map<string, string>();
+  private onDidChangeEmitter = new vscode.EventEmitter<vscode.Uri>();
+
+  get onDidChange(): vscode.Event<vscode.Uri> {
+    return this.onDidChangeEmitter.event;
+  }
+
+  provideTextDocumentContent(uri: vscode.Uri): string {
+    return this.content.get(uri.toString()) ?? '';
+  }
+
+  setContent(uri: vscode.Uri, content: string): void {
+    this.content.set(uri.toString(), content);
+    this.onDidChangeEmitter.fire(uri);
+  }
+
+  deleteContent(uri: vscode.Uri): void {
+    this.content.delete(uri.toString());
+  }
+
+  getContent(uri: vscode.Uri): string | undefined {
+    return this.content.get(uri.toString());
+  }
+}
+
+// Information about a diff view that is currently open.
+interface DiffInfo {
+  originalFilePath: string;
+  newContent: string;
+  rightDocUri: vscode.Uri;
+}
+
+/**
+ * Manages the state and lifecycle of diff views within the IDE.
+ */
+export class DiffManager {
+  private readonly onDidChangeEmitter =
+    new vscode.EventEmitter<JSONRPCNotification>();
+  readonly onDidChange = this.onDidChangeEmitter.event;
+  private diffDocuments = new Map<string, DiffInfo>();
+
+  constructor(
+    private readonly logger: vscode.OutputChannel,
+    private readonly diffContentProvider: DiffContentProvider,
+  ) {}
+
+  /**
+   * Creates and shows a new diff view.
+   */
+  async showDiff(filePath: string, newContent: string) {
+    const fileUri = vscode.Uri.file(filePath);
+
+    const rightDocUri = vscode.Uri.from({
+      scheme: DIFF_SCHEME,
+      path: filePath,
+      // cache busting
+      query: `rand=${Math.random()}`,
+    });
+    this.diffContentProvider.setContent(rightDocUri, newContent);
+
+    this.addDiffDocument(rightDocUri, {
+      originalFilePath: filePath,
+      newContent,
+      rightDocUri,
+    });
+
+    const diffTitle = `${path.basename(filePath)} ↔ Modified`;
+    await vscode.commands.executeCommand(
+      'setContext',
+      'gemini.diff.isVisible',
+      true,
+    );
+
+    let leftDocUri;
+    try {
+      await vscode.workspace.fs.stat(fileUri);
+      leftDocUri = fileUri;
+    } catch {
+      // We need to provide an empty document to diff against.
+      // Using the 'untitled' scheme is one way to do this.
+      leftDocUri = vscode.Uri.from({
+        scheme: 'untitled',
+        path: filePath,
+      });
+    }
+
+    await vscode.commands.executeCommand(
+      'vscode.diff',
+      leftDocUri,
+      rightDocUri,
+      diffTitle,
+      {
+        preview: false,
+      },
+    );
+    await vscode.commands.executeCommand(
+      'workbench.action.files.setActiveEditorWriteableInSession',
+    );
+  }
+
+  /**
+   * Closes an open diff view for a specific file.
+   */
+  async closeDiff(filePath: string) {
+    let uriToClose: vscode.Uri | undefined;
+    for (const [uriString, diffInfo] of this.diffDocuments.entries()) {
+      if (diffInfo.originalFilePath === filePath) {
+        uriToClose = vscode.Uri.parse(uriString);
+        break;
+      }
+    }
+
+    if (uriToClose) {
+      const rightDoc = await vscode.workspace.openTextDocument(uriToClose);
+      const modifiedContent = rightDoc.getText();
+      await this.closeDiffEditor(uriToClose);
+      this.onDidChangeEmitter.fire({
+        jsonrpc: '2.0',
+        method: 'ide/diffClosed',
+        params: {
+          filePath,
+          content: modifiedContent,
+        },
+      });
+      vscode.window.showInformationMessage(`Diff for ${filePath} closed.`);
+    } else {
+      vscode.window.showWarningMessage(`No open diff found for ${filePath}.`);
+    }
+  }
+
+  /**
+   * User accepts the changes in a diff view. Does not apply changes.
+   */
+  async acceptDiff(rightDocUri: vscode.Uri) {
+    const diffInfo = this.diffDocuments.get(rightDocUri.toString());
+    if (!diffInfo) {
+      this.logger.appendLine(
+        `No diff info found for ${rightDocUri.toString()}`,
+      );
+      return;
+    }
+
+    const rightDoc = await vscode.workspace.openTextDocument(rightDocUri);
+    const modifiedContent = rightDoc.getText();
+    await this.closeDiffEditor(rightDocUri);
+
+    this.onDidChangeEmitter.fire({
+      jsonrpc: '2.0',
+      method: 'ide/diffAccepted',
+      params: {
+        filePath: diffInfo.originalFilePath,
+        content: modifiedContent,
+      },
+    });
+  }
+
+  /**
+   * Called when a user cancels a diff view.
+   */
+  async cancelDiff(rightDocUri: vscode.Uri) {
+    const diffInfo = this.diffDocuments.get(rightDocUri.toString());
+    if (!diffInfo) {
+      this.logger.appendLine(
+        `No diff info found for ${rightDocUri.toString()}`,
+      );
+      // Even if we don't have diff info, we should still close the editor.
+      await this.closeDiffEditor(rightDocUri);
+      return;
+    }
+
+    const rightDoc = await vscode.workspace.openTextDocument(rightDocUri);
+    const modifiedContent = rightDoc.getText();
+    await this.closeDiffEditor(rightDocUri);
+
+    this.onDidChangeEmitter.fire({
+      jsonrpc: '2.0',
+      method: 'ide/diffClosed',
+      params: {
+        filePath: diffInfo.originalFilePath,
+        content: modifiedContent,
+      },
+    });
+  }
+
+  private addDiffDocument(uri: vscode.Uri, diffInfo: DiffInfo) {
+    this.diffDocuments.set(uri.toString(), diffInfo);
+  }
+
+  private async closeDiffEditor(rightDocUri: vscode.Uri) {
+    const diffInfo = this.diffDocuments.get(rightDocUri.toString());
+    await vscode.commands.executeCommand(
+      'setContext',
+      'gemini.diff.isVisible',
+      false,
+    );
+
+    if (diffInfo) {
+      this.diffDocuments.delete(rightDocUri.toString());
+      this.diffContentProvider.deleteContent(rightDocUri);
+    }
+
+    // Find and close the tab corresponding to the diff view
+    for (const tabGroup of vscode.window.tabGroups.all) {
+      for (const tab of tabGroup.tabs) {
+        const input = tab.input as {
+          modified?: vscode.Uri;
+          original?: vscode.Uri;
+        };
+        if (input && input.modified?.toString() === rightDocUri.toString()) {
+          await vscode.window.tabGroups.close(tab);
+          return;
+        }
+      }
+    }
+  }
+}
diff --git a/packages/vscode-ide-companion/src/extension.ts b/packages/vscode-ide-companion/src/extension.ts
index 73090175..b31e15b8 100644
--- a/packages/vscode-ide-companion/src/extension.ts
+++ b/packages/vscode-ide-companion/src/extension.ts
@@ -6,12 +6,15 @@
 
 import * as vscode from 'vscode';
 import { IDEServer } from './ide-server.js';
+import { DiffContentProvider, DiffManager } from './diff-manager.js';
 import { createLogger } from './utils/logger.js';
 
 const IDE_WORKSPACE_PATH_ENV_VAR = 'GEMINI_CLI_IDE_WORKSPACE_PATH';
+export const DIFF_SCHEME = 'gemini-diff';
 
 let ideServer: IDEServer;
 let logger: vscode.OutputChannel;
+
 let log: (message: string) => void = () => {};
 
 function updateWorkspacePath(context: vscode.ExtensionContext) {
@@ -37,7 +40,40 @@ export async function activate(context: vscode.ExtensionContext) {
 
   updateWorkspacePath(context);
 
-  ideServer = new IDEServer(log);
+  const diffContentProvider = new DiffContentProvider();
+  const diffManager = new DiffManager(logger, diffContentProvider);
+
+  context.subscriptions.push(
+    vscode.workspace.onDidCloseTextDocument((doc) => {
+      if (doc.uri.scheme === DIFF_SCHEME) {
+        diffManager.cancelDiff(doc.uri);
+      }
+    }),
+    vscode.workspace.registerTextDocumentContentProvider(
+      DIFF_SCHEME,
+      diffContentProvider,
+    ),
+    vscode.commands.registerCommand(
+      'gemini.diff.accept',
+      (uri?: vscode.Uri) => {
+        const docUri = uri ?? vscode.window.activeTextEditor?.document.uri;
+        if (docUri && docUri.scheme === DIFF_SCHEME) {
+          diffManager.acceptDiff(docUri);
+        }
+      },
+    ),
+    vscode.commands.registerCommand(
+      'gemini.diff.cancel',
+      (uri?: vscode.Uri) => {
+        const docUri = uri ?? vscode.window.activeTextEditor?.document.uri;
+        if (docUri && docUri.scheme === DIFF_SCHEME) {
+          diffManager.cancelDiff(docUri);
+        }
+      },
+    ),
+  );
+
+  ideServer = new IDEServer(log, diffManager);
   try {
     await ideServer.start(context);
   } catch (err) {
diff --git a/packages/vscode-ide-companion/src/ide-server.ts b/packages/vscode-ide-companion/src/ide-server.ts
index 8296c64c..30215ccc 100644
--- a/packages/vscode-ide-companion/src/ide-server.ts
+++ b/packages/vscode-ide-companion/src/ide-server.ts
@@ -14,6 +14,8 @@ import {
   type JSONRPCNotification,
 } from '@modelcontextprotocol/sdk/types.js';
 import { Server as HTTPServer } from 'node:http';
+import { z } from 'zod';
+import { DiffManager } from './diff-manager.js';
 import { OpenFilesManager } from './open-files-manager.js';
 
 const MCP_SESSION_ID_HEADER = 'mcp-session-id';
@@ -45,20 +47,22 @@ export class IDEServer {
   private server: HTTPServer | undefined;
   private context: vscode.ExtensionContext | undefined;
   private log: (message: string) => void;
+  diffManager: DiffManager;
 
-  constructor(log: (message: string) => void) {
+  constructor(log: (message: string) => void, diffManager: DiffManager) {
     this.log = log;
+    this.diffManager = diffManager;
   }
 
   async start(context: vscode.ExtensionContext) {
     this.context = context;
+    const sessionsWithInitialNotification = new Set<string>();
     const transports: { [sessionId: string]: StreamableHTTPServerTransport } =
       {};
-    const sessionsWithInitialNotification = new Set<string>();
 
     const app = express();
     app.use(express.json());
-    const mcpServer = createMcpServer();
+    const mcpServer = createMcpServer(this.diffManager);
 
     const openFilesManager = new OpenFilesManager(context);
     const onDidChangeSubscription = openFilesManager.onDidChange(() => {
@@ -71,6 +75,14 @@ export class IDEServer {
       }
     });
     context.subscriptions.push(onDidChangeSubscription);
+    const onDidChangeDiffSubscription = this.diffManager.onDidChange(
+      (notification: JSONRPCNotification) => {
+        for (const transport of Object.values(transports)) {
+          transport.send(notification);
+        }
+      },
+    );
+    context.subscriptions.push(onDidChangeDiffSubscription);
 
     app.post('/mcp', async (req: Request, res: Response) => {
       const sessionId = req.headers[MCP_SESSION_ID_HEADER] as
@@ -88,7 +100,6 @@ export class IDEServer {
             transports[newSessionId] = transport;
           },
         });
-
         const keepAlive = setInterval(() => {
           try {
             transport.send({ jsonrpc: '2.0', method: 'ping' });
@@ -212,7 +223,7 @@ export class IDEServer {
   }
 }
 
-const createMcpServer = () => {
+const createMcpServer = (diffManager: DiffManager) => {
   const server = new McpServer(
     {
       name: 'gemini-cli-companion-mcp-server',
@@ -220,5 +231,54 @@ const createMcpServer = () => {
     },
     { capabilities: { logging: {} } },
   );
+  server.registerTool(
+    'openDiff',
+    {
+      description:
+        '(IDE Tool) Open a diff view to create or modify a file. Returns a notification once the diff has been accepted or rejcted.',
+      inputSchema: z.object({
+        filePath: z.string(),
+        // TODO(chrstn): determine if this should be required or not.
+        newContent: z.string().optional(),
+      }).shape,
+    },
+    async ({
+      filePath,
+      newContent,
+    }: {
+      filePath: string;
+      newContent?: string;
+    }) => {
+      await diffManager.showDiff(filePath, newContent ?? '');
+      return {
+        content: [
+          {
+            type: 'text',
+            text: `Showing diff for ${filePath}`,
+          },
+        ],
+      };
+    },
+  );
+  server.registerTool(
+    'closeDiff',
+    {
+      description: '(IDE Tool) Close an open diff view for a specific file.',
+      inputSchema: z.object({
+        filePath: z.string(),
+      }).shape,
+    },
+    async ({ filePath }: { filePath: string }) => {
+      await diffManager.closeDiff(filePath);
+      return {
+        content: [
+          {
+            type: 'text',
+            text: `Closed diff for ${filePath}`,
+          },
+        ],
+      };
+    },
+  );
   return server;
 };