[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>
This commit is contained in:
christine betts 2025-08-04 21:36:23 +00:00 committed by GitHub
parent e7b468e122
commit 93f8fe3671
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 381 additions and 6 deletions

View File

@ -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",

View File

@ -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;
}
}
}
}
}

View File

@ -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) {

View File

@ -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;
};