From 93f8fe3671babbd3065d7a80b9e5ac50c42042da Mon Sep 17 00:00:00 2001 From: christine betts Date: Mon, 4 Aug 2025 21:36:23 +0000 Subject: [PATCH] [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> --- packages/vscode-ide-companion/package.json | 51 ++++ .../vscode-ide-companion/src/diff-manager.ts | 228 ++++++++++++++++++ .../vscode-ide-companion/src/extension.ts | 38 ++- .../vscode-ide-companion/src/ide-server.ts | 70 +++++- 4 files changed, 381 insertions(+), 6 deletions(-) create mode 100644 packages/vscode-ide-companion/src/diff-manager.ts 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,7 +31,22 @@ "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(); + private onDidChangeEmitter = new vscode.EventEmitter(); + + get onDidChange(): vscode.Event { + 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(); + readonly onDidChange = this.onDidChangeEmitter.event; + private diffDocuments = new Map(); + + 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(); const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {}; - const sessionsWithInitialNotification = new Set(); 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; };