1 files changed, 89 insertions, 19 deletions

diff --git a/packages/vscode-ide-companion/src/ide-server.ts b/packages/vscode-ide-companion/src/ide-server.ts
index 6d25367c..43c55e7d 100644
--- a/packages/vscode-ide-companion/src/ide-server.ts
+++ b/packages/vscode-ide-companion/src/ide-server.ts
@@ -8,28 +8,81 @@ import * as vscode from 'vscode';
 import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
 import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
 import express, { Request, Response } from 'express';
+import { randomUUID } from 'node:crypto';
+import {
+  isInitializeRequest,
+  type JSONRPCNotification,
+} from '@modelcontextprotocol/sdk/types.js';
 
-export async function startIDEServer(_context: vscode.ExtensionContext) {
+export async function startIDEServer(context: vscode.ExtensionContext) {
   const app = express();
   app.use(express.json());
 
-  const mcpServer = createMcpServer();
-  const transport = new StreamableHTTPServerTransport({
-    sessionIdGenerator: undefined,
-    enableJsonResponse: true,
-  });
+  const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {};
 
-  mcpServer.connect(transport);
+  const disposable = vscode.window.onDidChangeActiveTextEditor((editor) => {
+    const filePath = editor ? editor.document.uri.fsPath : null;
+    const notification: JSONRPCNotification = {
+      jsonrpc: '2.0',
+      method: 'ide/activeFileChanged',
+      params: { filePath },
+    };
+    for (const transport of Object.values(transports)) {
+      transport.send(notification);
+    }
+  });
+  context.subscriptions.push(disposable);
 
   app.post('/mcp', async (req: Request, res: Response) => {
-    console.log('Received MCP request:', req.body);
+    const sessionId = req.headers['mcp-session-id'] as string | undefined;
+    let transport: StreamableHTTPServerTransport;
+
+    if (sessionId && transports[sessionId]) {
+      transport = transports[sessionId];
+    } else if (!sessionId && isInitializeRequest(req.body)) {
+      transport = new StreamableHTTPServerTransport({
+        sessionIdGenerator: () => randomUUID(),
+        onsessioninitialized: (newSessionId) => {
+          transports[newSessionId] = transport;
+          const editor = vscode.window.activeTextEditor;
+          const filePath = editor ? editor.document.uri.fsPath : null;
+          const notification: JSONRPCNotification = {
+            jsonrpc: '2.0',
+            method: 'ide/activeFileChanged',
+            params: { filePath },
+          };
+          transport.send(notification);
+        },
+      });
+
+      transport.onclose = () => {
+        if (transport.sessionId) {
+          delete transports[transport.sessionId];
+        }
+      };
+
+      const server = createMcpServer();
+      server.connect(transport);
+    } else {
+      res.status(400).json({
+        jsonrpc: '2.0',
+        error: {
+          code: -32000,
+          message:
+            'Bad Request: No valid session ID provided for non-initialize request.',
+        },
+        id: null,
+      });
+      return;
+    }
+
     try {
       await transport.handleRequest(req, res, req.body);
     } catch (error) {
       console.error('Error handling MCP request:', error);
       if (!res.headersSent) {
         res.status(500).json({
-          jsonrpc: '2.0',
+          jsonrpc: '2.0' as const,
           error: {
             code: -32603,
             message: 'Internal server error',
@@ -40,15 +93,29 @@ export async function startIDEServer(_context: vscode.ExtensionContext) {
     }
   });
 
-  // Handle GET requests for SSE streams
-  app.get('/mcp', async (req: Request, res: Response) => {
-    res.status(405).set('Allow', 'POST').send('Method Not Allowed');
-  });
+  const handleSessionRequest = async (req: Request, res: Response) => {
+    const sessionId = req.headers['mcp-session-id'] as string | undefined;
+    if (!sessionId || !transports[sessionId]) {
+      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);
+      if (!res.headersSent) {
+        res.status(400).send('Bad Request');
+      }
+    }
+  };
+
+  app.get('/mcp', handleSessionRequest);
 
-  // Start the server
   // TODO(#3918): Generate dynamically and write to env variable
   const PORT = 3000;
-  app.listen(PORT, (error) => {
+  app.listen(PORT, (error?: Error) => {
     if (error) {
       console.error('Failed to start server:', error);
       vscode.window.showErrorMessage(
@@ -60,10 +127,13 @@ export async function startIDEServer(_context: vscode.ExtensionContext) {
 }
 
 const createMcpServer = () => {
-  const server = new McpServer({
-    name: 'vscode-ide-server',
-    version: '1.0.0',
-  });
+  const server = new McpServer(
+    {
+      name: 'gemini-cli-companion-mcp-server',
+      version: '1.0.0',
+    },
+    { capabilities: { logging: {} } },
+  );
   server.registerTool(
     'getActiveFile',
     {