[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:
parent
e7b468e122
commit
93f8fe3671
|
@ -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",
|
||||
|
|
|
@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
|
@ -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) {
|
||||
|
|
|
@ -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;
|
||||
};
|
||||
|
|
Loading…
Reference in New Issue