@testdriverai/mcp 7.11.6-test → 7.11.7-canary

package/mcp-server/src/server.ts

@@ -18,6 +18,7 @@ import * as http from "http";
 import type { RequestHandlerExtra } from "@modelcontextprotocol/sdk/shared/protocol.js";
 import type { Variables } from "@modelcontextprotocol/sdk/shared/uriTemplate.js";
 import type { CallToolResult, ReadResourceResult, ServerNotification, ServerRequest } from "@modelcontextprotocol/sdk/types.js";
+import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
 import * as Sentry from "@sentry/node";
 import * as fs from "fs";
 import * as os from "os";
@@ -29,7 +30,7 @@ import * as core from "./core/actions.js";
 import { NoActiveSessionError, type ActionResult } from "./core/actions.js";
 import { resolveE2bTemplateId, resolveOs } from "./env-utils.js";
 import { SessionStartInputSchema, type SessionStartInput } from "./provision-types.js";
-import { sessionManager, type SessionState } from "./session.js";
+import { type SessionState } from "./session.js";
 
 // =============================================================================
 // Sentry
@@ -225,21 +226,41 @@ interface StoredImage {
   timestamp: number;
 }
 
-// Map of image ID -> image data
-const imageStore = new Map<string, StoredImage>();
-
-// Counter for generating unique image IDs
-let imageIdCounter = 0;
-
 // Maximum number of images to store (to prevent memory leaks)
 const MAX_STORED_IMAGES = 100;
 
+/**
+ * Per-connection image store. It lives in the active {@link core.CoreContext}'s
+ * adapter scratch space (not a module global), so concurrent HTTP MCP clients
+ * don't share — or evict — each other's screenshots. The stdio/eve hosts have a
+ * single context, so this behaves exactly like the old module-level Map for them.
+ * Resource reads (which also run inside the connection's context) resolve the
+ * same per-connection store, so a screenshot URI only resolves for the client
+ * that produced it.
+ */
+interface ImageStoreState {
+  images: Map<string, StoredImage>;
+  counter: number;
+}
+
+function imageStoreState(): ImageStoreState {
+  const adapter = core.getAdapterState();
+  let state = adapter.imageStore as ImageStoreState | undefined;
+  if (!state) {
+    state = { images: new Map(), counter: 0 };
+    adapter.imageStore = state;
+  }
+  return state;
+}
+
 /**
  * Store an image and return its unique resource URI
  */
 function storeImage(data: string, type: "screenshot" | "cropped"): string {
-  const id = `${type}-${++imageIdCounter}`;
-  
+  const state = imageStoreState();
+  const imageStore = state.images;
+  const id = `${type}-${++state.counter}`;
+
   // Clean up old images if we exceed the limit
   if (imageStore.size >= MAX_STORED_IMAGES) {
     // Remove oldest images (first entries in the map)
@@ -250,15 +271,15 @@ function storeImage(data: string, type: "screenshot" | "cropped"): string {
     }
     logger.debug("storeImage: Cleaned up old images", { removed: entriesToRemove, remaining: imageStore.size });
   }
-  
+
   imageStore.set(id, {
     data,
     type,
     timestamp: Date.now(),
   });
-  
+
   logger.debug("storeImage: Stored image", { id, type, dataLength: data.length });
-  
+
   const base = type === "screenshot" ? SCREENSHOT_RESOURCE_BASE : CROPPED_IMAGE_RESOURCE_BASE;
   return `${base}/${id}`;
 }
@@ -267,7 +288,13 @@ function storeImage(data: string, type: "screenshot" | "cropped"): string {
  * Get an image by its ID
  */
 function getStoredImage(id: string): StoredImage | undefined {
-  return imageStore.get(id);
+  return imageStoreState().images.get(id);
+}
+
+/** Diagnostics for the active connection's image store (used in debug logs). */
+function imageStoreDiagnostics(): { imageStoreSize: number; availableKeys: string[] } {
+  const images = imageStoreState().images;
+  return { imageStoreSize: images.size, availableKeys: Array.from(images.keys()) };
 }
 
 /**
@@ -277,7 +304,7 @@ function getSessionData(session: SessionState | null) {
   if (!session) return { id: null, expiresIn: 0 };
   return {
     id: session.sessionId,
-    expiresIn: sessionManager.getTimeRemaining(session.sessionId),
+    expiresIn: core.getSessionTimeRemaining(session.sessionId),
   };
 }
 
@@ -478,7 +505,7 @@ function createToolResult(
   let fullText = textContent;
   if (generatedCode && success) {
     // Get the test file from the current session
-    const session = sessionManager.getCurrentSession();
+    const session = core.getCurrentSession();
     const testFile = session?.testFile;
     
     if (testFile) {
@@ -501,7 +528,7 @@ function createToolResult(
   // structuredContent goes to UI (includes imageUrl for display)
   // Always include success flag so UI can display correct status indicator
   // Include generatedCode and testFile in structured data so agents can programmatically handle it
-  const session = sessionManager.getCurrentSession();
+  const session = core.getCurrentSession();
   return {
     content,
     structuredContent: { 
@@ -513,24 +540,40 @@ function createToolResult(
   };
 }
 
-// Create MCP server wrapped with Sentry for automatic tracing
-const server = isSentryEnabled()
-  ? Sentry.wrapMcpServerWithSentry(
-      new McpServer({
+/**
+ * Build a fully-registered MCP server instance (all resources + tools).
+ *
+ * This used to run once at module load against a single shared `server`. It is
+ * now a factory so the HTTP transport can mint one server PER connection — each
+ * bound (via the surrounding {@link core.runInContext}) to its own isolated
+ * {@link core.CoreContext}. That's what stops one Streamable-HTTP client's tool
+ * call from reading another client's sandbox. The stdio host calls this once and
+ * runs it over the global context, so its behavior is unchanged.
+ *
+ * The tool/resource handlers below don't reference any per-connection state
+ * directly — they go through the context-aware helpers (`storeImage`,
+ * `getStoredImage`, `core.*`), which resolve the active context. So a single
+ * copy of the registration code serves every connection correctly.
+ */
+function buildServer(): McpServer {
+  // Create MCP server wrapped with Sentry for automatic tracing
+  const server = isSentryEnabled()
+    ? Sentry.wrapMcpServerWithSentry(
+        new McpServer({
+          name: "testdriver",
+          version: version,
+        })
+      )
+    : new McpServer({
         name: "testdriver",
         version: version,
-      })
-    )
-  : new McpServer({
-      name: "testdriver",
-      version: version,
-    });
+      });
 
-// =============================================================================
-// Register UI Resource
-// =============================================================================
+  // =============================================================================
+  // Register UI Resource
+  // =============================================================================
 
-registerAppResource(
+  registerAppResource(
   server,
   RESOURCE_URI,
   RESOURCE_URI,
@@ -1294,8 +1337,7 @@ You can optionally provide a reference image URI to compare against a previous s
         logger.info("check: Looking up reference image", {
           referenceImageUri: params.referenceImageUri,
           extractedImageId: imageId,
-          imageStoreSize: imageStore.size,
-          availableKeys: Array.from(imageStore.keys()),
+          ...imageStoreDiagnostics(),
         });
 
         const storedImage = getStoredImage(imageId);
@@ -1312,8 +1354,7 @@ You can optionally provide a reference image URI to compare against a previous s
           logger.warn("check: Reference image NOT found in store, falling back to last screenshot", {
             referenceImageUri: params.referenceImageUri,
             imageId,
-            imageStoreSize: imageStore.size,
-            availableKeys: Array.from(imageStore.keys()),
+            ...imageStoreDiagnostics(),
           });
         }
       }
@@ -1870,20 +1911,56 @@ Learn more at https://docs.testdriver.ai/v7/getting-started/
   }
 );
 
+  return server;
+}
+
 
 // =============================================================================
 // HTTP transport (Streamable HTTP, no auth)
 // =============================================================================
 
+/** A live MCP connection: its transport, the server bound to it, and the
+ *  isolated action context every one of its tool calls runs inside. */
+interface HttpConnection {
+  transport: StreamableHTTPServerTransport;
+  server: McpServer;
+  ctx: core.CoreContext;
+}
+
+/** Read and JSON-parse a request body. Returns undefined for an empty/invalid
+ *  body (GET/DELETE carry none) so callers can treat "no body" uniformly. */
+function readJsonBody(req: http.IncomingMessage): Promise<unknown> {
+  return new Promise((resolve) => {
+    const chunks: Buffer[] = [];
+    req.on("data", (chunk) => chunks.push(chunk as Buffer));
+    req.on("end", () => {
+      if (chunks.length === 0) return resolve(undefined);
+      try {
+        resolve(JSON.parse(Buffer.concat(chunks).toString("utf-8")));
+      } catch {
+        resolve(undefined);
+      }
+    });
+    req.on("error", () => resolve(undefined));
+  });
+}
+
 /**
  * Serve the MCP server over Streamable HTTP with no authentication.
  *
  * This is what the eve agent (and any other Streamable-HTTP MCP client) connects
- * to. The server holds single-session global state (the active `sdk`, element
- * refs, image store), so we mirror that with a single long-lived MCP transport:
- * the first `initialize` request mints a session id and every subsequent request
- * reuses the same `server` + transport. Concurrent sessions are intentionally not
- * supported — this server provisions one sandbox at a time.
+ * to. Each MCP session gets its OWN transport + server + isolated action context,
+ * keyed by the `mcp-session-id` header:
+ *
+ *  - An `initialize` request (no session id yet) mints a fresh connection: a new
+ *    isolated {@link core.CoreContext}, a server built over it, and a transport
+ *    whose generated session id we store. Every later request carrying that id
+ *    routes back to the same connection and runs inside its context via
+ *    {@link core.runInContext} — so one client's tool call can only ever touch
+ *    its own sandbox / element refs / image store. This is the fix for the
+ *    cross-session bleed: state is per-connection, not process-global.
+ *  - `DELETE` (or transport close) tears the connection down and drops it from
+ *    the map, so a disconnecting client frees its slot and its context is GC'd.
  *
  * No auth is applied here by design (the connection is expected to be local-only
  * or otherwise protected outside the MCP layer). Do not expose this on a public
@@ -1894,33 +1971,92 @@ async function startHttpServer() {
   const port = Number(process.env.TD_MCP_PORT || process.env.PORT || 8788);
   const mcpPath = process.env.TD_MCP_PATH || "/mcp";
 
-  // One transport for the lifetime of the process (single-session design).
-  const transport = new StreamableHTTPServerTransport({
-    sessionIdGenerator: () => randomUUID(),
-  });
-  await server.connect(transport);
+  // Live connections keyed by MCP session id. One entry per connected client.
+  const connections = new Map<string, HttpConnection>();
 
   const httpServer = http.createServer((req, res) => {
-    const url = new URL(req.url || "/", `http://${req.headers.host}`);
+    void (async () => {
+      const url = new URL(req.url || "/", `http://${req.headers.host}`);
+
+      // Lightweight health check for readiness probes / `eve dev`.
+      if (req.method === "GET" && url.pathname === "/health") {
+        res.writeHead(200, { "content-type": "application/json" });
+        res.end(JSON.stringify({ status: "ok", server: "testdriver", version, sessions: connections.size }));
+        return;
+      }
 
-    // Lightweight health check for readiness probes / `eve dev`.
-    if (req.method === "GET" && url.pathname === "/health") {
-      res.writeHead(200, { "content-type": "application/json" });
-      res.end(JSON.stringify({ status: "ok", server: "testdriver", version }));
-      return;
-    }
+      if (url.pathname !== mcpPath) {
+        res.writeHead(404, { "content-type": "application/json" });
+        res.end(JSON.stringify({ error: "Not found", hint: `MCP endpoint is ${mcpPath}` }));
+        return;
+      }
 
-    if (url.pathname !== mcpPath) {
-      res.writeHead(404, { "content-type": "application/json" });
-      res.end(JSON.stringify({ error: "Not found", hint: `MCP endpoint is ${mcpPath}` }));
-      return;
-    }
+      const sessionId = req.headers["mcp-session-id"];
+      const sid = Array.isArray(sessionId) ? sessionId[0] : sessionId;
 
-    // The Streamable HTTP transport handles POST (requests), GET (SSE stream),
-    // and DELETE (session teardown) on the same path. Body parsing is left to
-    // the transport so it can read the raw stream.
-    transport.handleRequest(req, res).catch((error) => {
-      logger.error("http: handleRequest failed", { error: String(error) });
+      // POST bodies must be parsed here (to detect `initialize` and to hand the
+      // transport a pre-parsed body). GET (SSE) and DELETE carry no JSON body.
+      const body = req.method === "POST" ? await readJsonBody(req) : undefined;
+
+      let connection: HttpConnection | undefined = sid ? connections.get(sid) : undefined;
+
+      // A brand-new session: only an `initialize` request may create one.
+      if (!connection) {
+        if (sid) {
+          // Client presented a session id we don't know — it was torn down or is
+          // stale. 404 so the client re-initializes (matches SDK stateful mode).
+          res.writeHead(404, { "content-type": "application/json" });
+          res.end(JSON.stringify({ error: "Unknown or expired MCP session", sessionId: sid }));
+          return;
+        }
+        if (!isInitializeRequest(body)) {
+          res.writeHead(400, { "content-type": "application/json" });
+          res.end(JSON.stringify({ error: "Bad Request: no MCP session id and not an initialize request" }));
+          return;
+        }
+
+        // Mint an isolated context + a server bound to it + a transport.
+        const ctx = core.createIsolatedContext();
+        const mcpServer = buildServer();
+        const transport = new StreamableHTTPServerTransport({
+          sessionIdGenerator: () => randomUUID(),
+          onsessioninitialized: (newId: string) => {
+            connections.set(newId, { transport, server: mcpServer, ctx });
+            logger.info("http: MCP session initialized", { sessionId: newId, activeSessions: connections.size });
+          },
+        });
+        // When the transport closes (client DELETE, network drop, or shutdown),
+        // drop the connection and best-effort tear down its sandbox so a
+        // disconnecting client doesn't leak a live VM.
+        transport.onclose = () => {
+          const closedId = transport.sessionId;
+          if (closedId && connections.delete(closedId)) {
+            logger.info("http: MCP session closed", { sessionId: closedId, activeSessions: connections.size });
+          }
+          // Disconnect the SDK inside this connection's context (best-effort).
+          core.runInContext(ctx, () => { void core.disconnect(); });
+        };
+
+        await mcpServer.connect(transport);
+        connection = { transport, server: mcpServer, ctx };
+      }
+
+      // Route the request through THIS connection's transport, inside THIS
+      // connection's action context so every core.* call the handlers make
+      // resolves the right sandbox / refs / image store.
+      const conn = connection;
+      await core.runInContext(conn.ctx, () =>
+        conn.transport.handleRequest(req, res, body),
+      ).catch((error: unknown) => {
+        logger.error("http: handleRequest failed", { error: String(error) });
+        captureException(error as Error, { tags: { phase: "http" } });
+        if (!res.headersSent) {
+          res.writeHead(500, { "content-type": "application/json" });
+          res.end(JSON.stringify({ error: "Internal server error" }));
+        }
+      });
+    })().catch((error: unknown) => {
+      logger.error("http: request handling failed", { error: String(error) });
       captureException(error as Error, { tags: { phase: "http" } });
       if (!res.headersSent) {
         res.writeHead(500, { "content-type": "application/json" });
@@ -1935,9 +2071,12 @@ async function startHttpServer() {
   });
 
   const shutdown = async () => {
-    logger.info("Shutting down MCP Server");
+    logger.info("Shutting down MCP Server", { activeSessions: connections.size });
     httpServer.close();
-    await transport.close().catch(() => {});
+    for (const conn of connections.values()) {
+      await conn.transport.close().catch(() => {});
+    }
+    connections.clear();
     await flushSentry();
     process.exit(0);
   };
@@ -1964,6 +2103,10 @@ async function main() {
     return;
   }
 
+  // stdio: one long-lived server for the process. It runs over the global
+  // action context (no runInContext wrap), so behavior is unchanged — a single
+  // local CLI client, one sandbox at a time, exactly as before.
+  const server = buildServer();
   const transport = new StdioServerTransport();
   await server.connect(transport);
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@testdriverai/mcp",
-  "version": "7.11.6-test",
+  "version": "7.11.7-canary",
   "description": "Next generation autonomous AI agent for end-to-end testing of web & desktop",
   "main": "sdk.js",
   "types": "sdk.d.ts",