feat: enhance MCP server functionality with new ServerRuntime interface and improved graceful shutdown process

Author: regenrek

packages/mcp/.browser-echo-mcp.json

@@ -1 +1 @@
-{"url":"http://127.0.0.1:5181","routeLogs":"/__client-logs","timestamp":1755771840799,"pid":51019}
+{"url":"http://127.0.0.1:5179","routeLogs":"/__client-logs","timestamp":1755775962617,"pid":14791}

packages/mcp/src/server.ts

@@ -1,4 +1,4 @@
-import { createServer as createNodeServer } from 'node:http';
+import { createServer as createNodeServer, type Server as NodeHttpServer } from 'node:http';
 import { randomUUID } from 'node:crypto';
 
 import { createApp, createRouter, defineEventHandler, getQuery, readRawBody, setResponseStatus, toNodeListener } from 'h3';
@@ -25,13 +25,21 @@ interface HttpOptions {
 
 export type StartOptions = HttpOptions;
 
+/** New: a runtime you can close without exiting the process */
+export interface ServerRuntime {
+  mcp: McpServer;
+  transport: StreamableHTTPServerTransport;
+  http: NodeHttpServer;
+  close(): Promise<void>;
+}
+
 /**
  * Starts the given MCP server with HTTP transport.
  */
 export async function startServer(
   server: McpServer,
   options: StartOptions,
-): Promise<void> {
+): Promise<ServerRuntime> {
   const { port, host, endpoint, logsRoute } = options;
 
   // Create store
@@ -43,29 +51,31 @@ export async function startServer(
   registerTools(context);
   registerResources(context);
 
-  // Start HTTP server
-  await startHttpServer(server, store, { host, port, endpoint, logsRoute });
+  // Start HTTP server & return a runtime handle
+  return await startHttpServer(server, store, { host, port, endpoint, logsRoute });
 }
 
-export async function stopServer(server: McpServer) {
-  try {
-    await server.close();
-  }
-  catch (error) {
-    console.error('Error occurred during server stop:', error);
+export async function stopServer(runtime: ServerRuntime | McpServer) {
+  // Back-compat: if someone passes the raw MCP server,
+  // close it but don't exit; prefer passing the runtime.
+  if ((runtime as ServerRuntime).close) {
+    await (runtime as ServerRuntime).close();
+    return;
   }
-  finally {
-    process.exit(0);
+  try {
+    await (runtime as McpServer).close();
+  } catch (error) {
+    console.error('Error occurred during MCP server stop:', error);
   }
 }
 
-/** Start a standalone H3 HTTP server exposing:
+/** Start an H3 HTTP server exposing:
  *  - MCP endpoint (POST for requests, GET for SSE) at {endpoint}
  *  - Log ingest/diagnostics at {logsRoute} (POST to append, GET to view)
  */
 async function startHttpServer(mcp: McpServer, store: LogStore, opts: {
   host: string; port: number; endpoint: `/${string}`; logsRoute: `/${string}`;
-}): Promise<void> {
+}): Promise<ServerRuntime> {
   const app = createApp();
   const router = createRouter();
 
@@ -100,45 +110,52 @@ async function startHttpServer(mcp: McpServer, store: LogStore, opts: {
         return '';
       }
 
-      // For GET/DELETE: require session id to prevent ambiguous streams without init
-      if (method === 'GET' || method === 'DELETE') {
-        const sidHeader = (event.node.req.headers['mcp-session-id'] as string | undefined) || '';
-        if (!sidHeader) {
-          setResponseStatus(event, 405);
-          try {
-            event.node.res.setHeader('Allow', 'POST');
-            event.node.res.setHeader('content-type', 'application/json');
-          } catch {}
-          return JSON.stringify({
-            jsonrpc: '2.0',
-            error: {
-              code: -32000,
-              message: 'Method not allowed. First POST InitializeRequest, read Mcp-Session-Id, and include it (plus MCP-Protocol-Version) on subsequent GET/DELETE/POST.'
-            },
-            id: null
-          });
-        }
-      }
-
-      // Normalize body for POST
+      // Read/parse body for POST-like early so we can detect "initialize"
       let bodyBuf: Buffer | undefined;
+      let parsed: any | undefined;
       if (method === 'POST' || method === 'PUT' || method === 'PATCH') {
         const raw = await readRawBody(event);
         if (raw && typeof raw === 'string') bodyBuf = Buffer.from(raw);
         else if (raw && (raw as any) instanceof Uint8Array) bodyBuf = Buffer.from(raw as any);
+        if (bodyBuf) {
+          try { parsed = JSON.parse(bodyBuf.toString('utf-8')); } catch { /* ignore */ }
+        }
+      }
 
-        // Allow tool calls without explicit session id by assigning a default
+      const isInitialize = method === 'POST' && parsed && parsed.method === 'initialize';
+      const headers = event.node.req.headers as Record<string, string | string[] | undefined>;
+      const sidHeader = (headers['mcp-session-id'] as string | undefined) || '';
+
+      // For GET/DELETE always require session id.
+      // For POST require session id unless it's an initialize request (first handshake).
+      if ((method === 'GET' || method === 'DELETE' || (method === 'POST' && !isInitialize)) && !sidHeader) {
+        setResponseStatus(event, 405);
         try {
-          const reqHeaders: Record<string, string | string[] | undefined> = event.node.req.headers as any;
-          const existingSid = reqHeaders['mcp-session-id'] || reqHeaders['Mcp-Session-Id'];
-          if (!existingSid) {
-            (event.node.req.headers as any)['mcp-session-id'] = 'default-session';
-          }
+          event.node.res.setHeader('Allow', 'POST');
+          event.node.res.setHeader('content-type', 'application/json');
         } catch {}
+        // If we know the JSON-RPC id from the body, echo it; otherwise null.
+        const id = parsed?.id ?? null;
+        return JSON.stringify({
+          jsonrpc: '2.0',
+          error: {
+            code: -32000,
+            message: 'Method not allowed. First POST initialize (no Mcp-Session-Id), read Mcp-Session-Id from the response, and include it (plus MCP-Protocol-Version) on subsequent GET/DELETE/POST.'
+          },
+          id
+        });
       }
 
+      // ❌ Removed: never force a "default-session" header.
+      // Let the transport generate a session id on initialize and return it.
+
       // Delegate to the SDK transport (writes to the Node response directly)
-      await transport.handleRequest(event.node.req as any, event.node.res as any, bodyBuf ? JSON.parse(bodyBuf.toString('utf-8')) : undefined);
+      await transport.handleRequest(
+        event.node.req as any,
+        event.node.res as any,
+        // Pass already-parsed JSON to avoid double parse
+        parsed ?? undefined
+      );
       // h3 will consider the response handled.
       return undefined as any;
     } catch (err) {
@@ -179,6 +196,27 @@ async function startHttpServer(mcp: McpServer, store: LogStore, opts: {
   console.log(`MCP (Streamable HTTP) listening → http://${opts.host}:${opts.port}${opts.endpoint}`);
   // eslint-disable-next-line no-console
   console.log(`Log ingest endpoint        → http://${opts.host}:${opts.port}${opts.logsRoute}`);
+
+  const close = async () => {
+    // Close in order: transport → MCP sessions → HTTP server
+    try {
+      if (typeof (transport as any).close === 'function') {
+        await (transport as any).close();
+      }
+    } catch (e) {
+      console.warn('Transport close failed:', e);
+    }
+    try {
+      await mcp.close();
+    } catch (e) {
+      console.warn('MCP close failed:', e);
+    }
+    await new Promise<void>((resolve) => {
+      try { nodeServer.close(() => resolve()); } catch { resolve(); }
+    });
+  };
+
+  return { mcp, transport, http: nodeServer, close };
 }
 
 async function advertiseDiscovery(host: string, port: number, logsRoute: `/${string}`) {
@@ -255,3 +293,24 @@ function createLogIngestRoutes(store: LogStore, logsRoute: `/${string}`) {
 
   return router;
 }
+
+// Example usage with the new ServerRuntime:
+/*
+```ts
+import { createServer, startServer, stopServer } from './server';
+
+async function example() {
+  const mcp = createServer({ name: 'Browser Echo (Frontend Logs)', version: '1.0.0' });
+  const runtime = await startServer(mcp, {
+    type: 'http',
+    host: 'localhost',
+    port: 3001,
+    endpoint: '/mcp',
+    logsRoute: '/logs'
+  });
+
+  // Later, to stop gracefully:
+  await stopServer(runtime); // No process.exit, just graceful shutdown
+}
+```
+*/