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 { startIDEServer } from './ide-server';
import { IDEServer } from './ide-server';
let ideServer: IDEServer;
let logger: vscode.OutputChannel;
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();
}
}

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

@ -13,14 +13,18 @@ import {
isInitializeRequest,
type JSONRPCNotification,
} 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(
transport: StreamableHTTPServerTransport,
logger: vscode.OutputChannel,
) {
const editor = vscode.window.activeTextEditor;
const filePath = editor ? editor.document.uri.fsPath : '';
logger.appendLine(`Sending active file changed notification: ${filePath}`);
const notification: JSONRPCNotification = {
jsonrpc: '2.0',
method: 'ide/activeFileChanged',
@ -29,22 +33,36 @@ function sendActiveFileChangedNotification(
transport.send(notification);
}
export async function startIDEServer(context: vscode.ExtensionContext) {
export class IDEServer {
private server: HTTPServer | undefined;
private context: vscode.ExtensionContext | undefined;
private logger: vscode.OutputChannel;
constructor(logger: vscode.OutputChannel) {
this.logger = logger;
}
async start(context: vscode.ExtensionContext) {
this.context = context;
const transports: { [sessionId: string]: StreamableHTTPServerTransport } =
{};
const sessionsWithInitialNotification = new Set<string>();
const app = express();
app.use(express.json());
const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {};
const sessionsWithInitialNotification = new Set<string>();
const mcpServer = createMcpServer();
const disposable = vscode.window.onDidChangeActiveTextEditor((_editor) => {
for (const transport of Object.values(transports)) {
sendActiveFileChangedNotification(transport);
sendActiveFileChangedNotification(transport, this.logger);
}
});
context.subscriptions.push(disposable);
app.post('/mcp', async (req: Request, res: Response) => {
const sessionId = req.headers['mcp-session-id'] as string | undefined;
const sessionId = req.headers[MCP_SESSION_ID_HEADER] as
| string
| undefined;
let transport: StreamableHTTPServerTransport;
if (sessionId && transports[sessionId]) {
@ -53,20 +71,22 @@ export async function startIDEServer(context: vscode.ExtensionContext) {
transport = new StreamableHTTPServerTransport({
sessionIdGenerator: () => randomUUID(),
onsessioninitialized: (newSessionId) => {
this.logger.appendLine(`New session initialized: ${newSessionId}`);
transports[newSessionId] = transport;
},
});
transport.onclose = () => {
if (transport.sessionId) {
this.logger.appendLine(`Session closed: ${transport.sessionId}`);
sessionsWithInitialNotification.delete(transport.sessionId);
delete transports[transport.sessionId];
}
};
const server = createMcpServer();
server.connect(transport);
mcpServer.connect(transport);
} else {
this.logger.appendLine(
'Bad Request: No valid session ID provided for non-initialize request.',
);
res.status(400).json({
jsonrpc: '2.0',
error: {
@ -82,7 +102,9 @@ export async function startIDEServer(context: vscode.ExtensionContext) {
try {
await transport.handleRequest(req, res, req.body);
} catch (error) {
console.error('Error handling MCP request:', error);
const errorMessage =
error instanceof Error ? error.message : 'Unknown error';
this.logger.appendLine(`Error handling MCP request: ${errorMessage}`);
if (!res.headersSent) {
res.status(500).json({
jsonrpc: '2.0' as const,
@ -97,50 +119,73 @@ export async function startIDEServer(context: vscode.ExtensionContext) {
});
const handleSessionRequest = async (req: Request, res: Response) => {
const sessionId = req.headers['mcp-session-id'] as string | undefined;
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) {
console.error('Error handling MCP GET request:', 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);
sendActiveFileChangedNotification(transport, this.logger);
sessionsWithInitialNotification.add(sessionId);
}
};
app.get('/mcp', handleSessionRequest);
const server = app.listen(0, () => {
const address = (server as Server).address();
this.server = app.listen(0, () => {
const address = (this.server as HTTPServer).address();
if (address && typeof address !== 'string') {
const port = address.port;
context.environmentVariableCollection.replace(
'GEMINI_CLI_IDE_SERVER_PORT',
IDE_SERVER_PORT_ENV_VAR,
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`,
);
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 (this.context) {
this.context.environmentVariableCollection.clear();
}
}
}
const createMcpServer = () => {
const server = new McpServer(
{
@ -157,7 +202,6 @@ const createMcpServer = () => {
inputSchema: {},
},
async () => {
try {
const activeEditor = vscode.window.activeTextEditor;
const filePath = activeEditor ? activeEditor.document.uri.fsPath : '';
if (filePath) {
@ -174,18 +218,6 @@ const createMcpServer = () => {
],
};
}
} catch (error) {
return {
content: [
{
type: 'text',
text: `Failed to get active file: ${
(error as Error).message || 'Unknown error'
}`,
},
],
};
}
},
);
return server;