jcarr
/
gemini-cli
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Minor refactoring of IDE companion server (#4331)

This commit is contained in:
Shreya Keshive 2025-07-16 21:03:56 -04:00 committed by GitHub
parent fbe09cd35e
commit 0c76affe6d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 173 additions and 119 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
packages/vscode-ide-companion/src/extension.ts
View File

@ -5,10 +5,32 @@
*/ */
import * as vscode from 'vscode'; import * as vscode from 'vscode';
import { startIDEServer } from './ide-server'; import { IDEServer } from './ide-server';
let ideServer: IDEServer;
let logger: vscode.OutputChannel;
export async function activate(context: vscode.ExtensionContext) { export async function activate(context: vscode.ExtensionContext) {
startIDEServer(context); logger = vscode.window.createOutputChannel('Gemini CLI IDE Companion');
logger.show();
logger.appendLine('Starting Gemini CLI IDE Companion server...');
ideServer = new IDEServer(logger);
try {
await ideServer.start(context);
} catch (err) {
const message = err instanceof Error ? err.message : String(err);
logger.appendLine(`Failed to start IDE server: ${message}`);
}
} }
export function deactivate() {} export function deactivate() {
if (ideServer) {
logger.appendLine('Deactivating Gemini CLI IDE Companion...');
return ideServer.stop().finally(() => {
logger.dispose();
});
}
if (logger) {
logger.dispose();
}
}

264
packages/vscode-ide-companion/src/ide-server.ts
View File

@ -13,14 +13,18 @@ import {
isInitializeRequest, isInitializeRequest,
type JSONRPCNotification, type JSONRPCNotification,
} from '@modelcontextprotocol/sdk/types.js'; } from '@modelcontextprotocol/sdk/types.js';
import { Server as HTTPServer } from 'node:http';
import { Server } from 'node:http'; const MCP_SESSION_ID_HEADER = 'mcp-session-id';
const IDE_SERVER_PORT_ENV_VAR = 'GEMINI_CLI_IDE_SERVER_PORT';
function sendActiveFileChangedNotification( function sendActiveFileChangedNotification(
transport: StreamableHTTPServerTransport, transport: StreamableHTTPServerTransport,
logger: vscode.OutputChannel,
) { ) {
const editor = vscode.window.activeTextEditor; const editor = vscode.window.activeTextEditor;
const filePath = editor ? editor.document.uri.fsPath : ''; const filePath = editor ? editor.document.uri.fsPath : '';
logger.appendLine(`Sending active file changed notification: ${filePath}`);
const notification: JSONRPCNotification = { const notification: JSONRPCNotification = {
jsonrpc: '2.0', jsonrpc: '2.0',
method: 'ide/activeFileChanged', method: 'ide/activeFileChanged',
@ -29,116 +33,157 @@ function sendActiveFileChangedNotification(
transport.send(notification); transport.send(notification);
} }
export async function startIDEServer(context: vscode.ExtensionContext) { export class IDEServer {
const app = express(); private server: HTTPServer | undefined;
app.use(express.json()); private context: vscode.ExtensionContext | undefined;
private logger: vscode.OutputChannel;
const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {}; constructor(logger: vscode.OutputChannel) {
const sessionsWithInitialNotification = new Set<string>(); this.logger = logger;
}
const disposable = vscode.window.onDidChangeActiveTextEditor((_editor) => { async start(context: vscode.ExtensionContext) {
for (const transport of Object.values(transports)) { this.context = context;
sendActiveFileChangedNotification(transport); const transports: { [sessionId: string]: StreamableHTTPServerTransport } =
} {};
}); const sessionsWithInitialNotification = new Set<string>();
context.subscriptions.push(disposable);
app.post('/mcp', async (req: Request, res: Response) => { const app = express();
const sessionId = req.headers['mcp-session-id'] as string | undefined; app.use(express.json());
let transport: StreamableHTTPServerTransport; const mcpServer = createMcpServer();
if (sessionId && transports[sessionId]) { const disposable = vscode.window.onDidChangeActiveTextEditor((_editor) => {
transport = transports[sessionId]; for (const transport of Object.values(transports)) {
} else if (!sessionId && isInitializeRequest(req.body)) { sendActiveFileChangedNotification(transport, this.logger);
transport = new StreamableHTTPServerTransport({ }
sessionIdGenerator: () => randomUUID(), });
onsessioninitialized: (newSessionId) => { context.subscriptions.push(disposable);
transports[newSessionId] = transport;
},
});
transport.onclose = () => { app.post('/mcp', async (req: Request, res: Response) => {
if (transport.sessionId) { const sessionId = req.headers[MCP_SESSION_ID_HEADER] as
sessionsWithInitialNotification.delete(transport.sessionId); | string
delete transports[transport.sessionId]; | undefined;
} let transport: StreamableHTTPServerTransport;
};
const server = createMcpServer(); if (sessionId && transports[sessionId]) {
server.connect(transport); transport = transports[sessionId];
} else { } else if (!sessionId && isInitializeRequest(req.body)) {
res.status(400).json({ transport = new StreamableHTTPServerTransport({
jsonrpc: '2.0', sessionIdGenerator: () => randomUUID(),
error: { onsessioninitialized: (newSessionId) => {
code: -32000, this.logger.appendLine(`New session initialized: ${newSessionId}`);
message: transports[newSessionId] = transport;
'Bad Request: No valid session ID provided for non-initialize request.', },
}, });
id: null, transport.onclose = () => {
}); if (transport.sessionId) {
return; this.logger.appendLine(`Session closed: ${transport.sessionId}`);
} sessionsWithInitialNotification.delete(transport.sessionId);
delete transports[transport.sessionId];
try { }
await transport.handleRequest(req, res, req.body); };
} catch (error) { mcpServer.connect(transport);
console.error('Error handling MCP request:', error); } else {
if (!res.headersSent) { this.logger.appendLine(
res.status(500).json({ 'Bad Request: No valid session ID provided for non-initialize request.',
jsonrpc: '2.0' as const, );
res.status(400).json({
jsonrpc: '2.0',
error: { error: {
code: -32603, code: -32000,
message: 'Internal server error', message:
'Bad Request: No valid session ID provided for non-initialize request.',
}, },
id: null, id: null,
}); });
return;
} }
}
});
const handleSessionRequest = async (req: Request, res: Response) => { try {
const sessionId = req.headers['mcp-session-id'] as string | undefined; await transport.handleRequest(req, res, req.body);
if (!sessionId || !transports[sessionId]) { } catch (error) {
res.status(400).send('Invalid or missing session ID'); const errorMessage =
return; error instanceof Error ? error.message : 'Unknown error';
} this.logger.appendLine(`Error handling MCP request: ${errorMessage}`);
if (!res.headersSent) {
const transport = transports[sessionId]; res.status(500).json({
jsonrpc: '2.0' as const,
try { error: {
await transport.handleRequest(req, res); code: -32603,
} catch (error) { message: 'Internal server error',
console.error('Error handling MCP GET request:', error); },
if (!res.headersSent) { id: null,
res.status(400).send('Bad Request'); });
}
} }
});
const handleSessionRequest = async (req: Request, res: Response) => {
const sessionId = req.headers[MCP_SESSION_ID_HEADER] as
| string
| undefined;
if (!sessionId || !transports[sessionId]) {
this.logger.appendLine('Invalid or missing session ID');
res.status(400).send('Invalid or missing session ID');
return;
}
const transport = transports[sessionId];
try {
await transport.handleRequest(req, res);
} catch (error) {
const errorMessage =
error instanceof Error ? error.message : 'Unknown error';
this.logger.appendLine(
`Error handling session request: ${errorMessage}`,
);
if (!res.headersSent) {
res.status(400).send('Bad Request');
}
}
if (!sessionsWithInitialNotification.has(sessionId)) {
sendActiveFileChangedNotification(transport, this.logger);
sessionsWithInitialNotification.add(sessionId);
}
};
app.get('/mcp', handleSessionRequest);
this.server = app.listen(0, () => {
const address = (this.server as HTTPServer).address();
if (address && typeof address !== 'string') {
const port = address.port;
context.environmentVariableCollection.replace(
IDE_SERVER_PORT_ENV_VAR,
port.toString(),
);
this.logger.appendLine(`IDE server listening on port ${port}`);
}
});
}
async stop(): Promise<void> {
if (this.server) {
await new Promise<void>((resolve, reject) => {
this.server!.close((err?: Error) => {
if (err) {
this.logger.appendLine(
`Error shutting down IDE server: ${err.message}`,
);
return reject(err);
}
this.logger.appendLine(`IDE server shut down`);
resolve();
});
});
this.server = undefined;
} }
if (!sessionsWithInitialNotification.has(sessionId)) { if (this.context) {
sendActiveFileChangedNotification(transport); this.context.environmentVariableCollection.clear();
sessionsWithInitialNotification.add(sessionId);
} }
}; }
app.get('/mcp', handleSessionRequest);
const server = app.listen(0, () => {
const address = (server as Server).address();
if (address && typeof address !== 'string') {
const port = address.port;
context.environmentVariableCollection.replace(
'GEMINI_CLI_IDE_SERVER_PORT',
port.toString(),
);
console.log(`MCP Streamable HTTP Server listening on port ${port}`);
} else {
const port = 0;
console.error('Failed to start server:', 'Unknown error');
vscode.window.showErrorMessage(
`Companion server failed to start on port ${port}: Unknown error`,
);
}
});
} }
const createMcpServer = () => { const createMcpServer = () => {
@ -157,31 +202,18 @@ const createMcpServer = () => {
inputSchema: {}, inputSchema: {},
}, },
async () => { async () => {
try { const activeEditor = vscode.window.activeTextEditor;
const activeEditor = vscode.window.activeTextEditor; const filePath = activeEditor ? activeEditor.document.uri.fsPath : '';
const filePath = activeEditor ? activeEditor.document.uri.fsPath : ''; if (filePath) {
if (filePath) { return {
return { content: [{ type: 'text', text: `Active file: ${filePath}` }],
content: [{ type: 'text', text: `Active file: ${filePath}` }], };
}; } else {
} else {
return {
content: [
{
type: 'text',
text: 'No file is currently active in the editor.',
},
],
};
}
} catch (error) {
return { return {
content: [ content: [
{ {
type: 'text', type: 'text',
text: `Failed to get active file: ${ text: 'No file is currently active in the editor.',
(error as Error).message || 'Unknown error'
}`,
}, },
], ],
}; };