fix: create a map for each transport to identify different sessions to assign context through request cycle (transport to handlers)

Author: betegon

packages/core/src/integrations/mcp-server/correlation.ts

@@ -4,56 +4,46 @@
  */
 
 import { getClient } from '../../currentScopes';
-import { SPAN_STATUS_ERROR, withActiveSpan } from '../../tracing';
+import { SPAN_STATUS_ERROR } from '../../tracing';
 import type { Span } from '../../types-hoist/span';
 import { extractToolResultAttributes } from './attributeExtraction';
 import { filterMcpPiiFromSpanData } from './piiFiltering';
-import type { RequestId, RequestSpanMapValue } from './types';
+import type { MCPTransport, RequestId, RequestSpanMapValue } from './types';
 
-// Simplified correlation system that works with or without sessionId
-// Maps requestId directly to span data for stateless operation
-const requestIdToSpanMap = new Map<RequestId, RequestSpanMapValue>();
+// Transport-scoped correlation system that prevents collisions between different MCP sessions
+// Each transport instance gets its own correlation map, eliminating request ID conflicts
+const transportToSpanMap = new WeakMap<MCPTransport, Map<RequestId, RequestSpanMapValue>>();
+
+/**
+ * Gets or creates the span map for a specific transport instance
+ */
+function getOrCreateSpanMap(transport: MCPTransport): Map<RequestId, RequestSpanMapValue> {
+  let spanMap = transportToSpanMap.get(transport);
+  if (!spanMap) {
+    spanMap = new Map();
+    transportToSpanMap.set(transport, spanMap);
+  }
+  return spanMap;
+}
 
 /**
  * Stores span context for later correlation with handler execution
  */
-export function storeSpanForRequest(requestId: RequestId, span: Span, method: string): void {
-  requestIdToSpanMap.set(requestId, {
+export function storeSpanForRequest(transport: MCPTransport, requestId: RequestId, span: Span, method: string): void {
+  const spanMap = getOrCreateSpanMap(transport);
+  spanMap.set(requestId, {
     span,
     method,
     startTime: Date.now(),
   });
 }
 
-/**
- * Associates handler execution with the corresponding request span
- */
-export function associateContextWithRequestSpan<T>(
-  extraHandlerData: { requestId: RequestId } | undefined,
-  cb: () => T,
-): T {
-  if (extraHandlerData) {
-    const { requestId } = extraHandlerData;
-
-    const spanData = requestIdToSpanMap.get(requestId);
-    if (!spanData) {
-      return cb();
-    }
-
-    // Keep span in map for response enrichment (don't delete yet)
-    return withActiveSpan(spanData.span, () => {
-      return cb();
-    });
-  }
-
-  return cb();
-}
-
 /**
  * Completes span with tool results and cleans up correlation
  */
-export function completeSpanWithResults(requestId: RequestId, result: unknown): void {
-  const spanData = requestIdToSpanMap.get(requestId);
+export function completeSpanWithResults(transport: MCPTransport, requestId: RequestId, result: unknown): void {
+  const spanMap = getOrCreateSpanMap(transport);
+  const spanData = spanMap.get(requestId);
   if (spanData) {
     const { span, method } = spanData;
 
@@ -68,24 +58,27 @@ export function completeSpanWithResults(requestId: RequestId, result: unknown):
     }
 
     span.end();
-    requestIdToSpanMap.delete(requestId);
+    spanMap.delete(requestId);
   }
 }
 
 /**
- * Cleans up all pending spans (for transport close)
+ * Cleans up pending spans for a specific transport (when that transport closes)
  */
-export function cleanupAllPendingSpans(): number {
-  const pendingCount = requestIdToSpanMap.size;
+export function cleanupPendingSpansForTransport(transport: MCPTransport): number {
+  const spanMap = transportToSpanMap.get(transport);
+  if (!spanMap) return 0;
+
+  const pendingCount = spanMap.size;
 
-  for (const [, spanData] of requestIdToSpanMap) {
+  for (const [, spanData] of spanMap) {
     spanData.span.setStatus({
       code: SPAN_STATUS_ERROR,
       message: 'cancelled',
     });
     spanData.span.end();
   }
 
-  requestIdToSpanMap.clear();
+  spanMap.clear();
   return pendingCount;
 }

packages/core/src/integrations/mcp-server/handlers.ts

@@ -1,12 +1,11 @@
 /**
- * Handler wrapping functions for MCP server methods
- * Provides span correlation for tool, resource, and prompt handlers
+ * Handler method wrapping for MCP server instrumentation
+ * Provides automatic error capture and span correlation for tool, resource, and prompt handlers
  */
 
 import { DEBUG_BUILD } from '../../debug-build';
 import { debug } from '../../utils/debug-logger';
 import { fill } from '../../utils/object';
-import { associateContextWithRequestSpan } from './correlation';
 import { captureError } from './errorCapture';
 import type { HandlerExtraData, MCPHandler, MCPServerInstance } from './types';
 
@@ -36,25 +35,23 @@ function createWrappedHandler(originalHandler: MCPHandler, methodName: keyof MCP
     try {
       const extraHandlerData = findExtraHandlerData(handlerArgs);
 
-      return associateContextWithRequestSpan(extraHandlerData, () => {
-        return createErrorCapturingHandler.call(
-          this,
-          originalHandler,
-          methodName,
-          handlerName,
-          handlerArgs,
-          extraHandlerData,
-        );
-      });
-    } catch (error) {
-      DEBUG_BUILD && debug.warn('MCP handler wrapping failed:', error);
-      return originalHandler.apply(this, handlerArgs);
-    }
+      return createErrorCapturingHandler.call(
+        this,
+        originalHandler,
+        methodName,
+        handlerName,
+        handlerArgs,
+        extraHandlerData,
+      );
+          } catch (error) {
+        DEBUG_BUILD && debug.warn('MCP handler wrapping failed:', error);
+        return originalHandler.apply(this, handlerArgs);
+      }
   };
 }
 
 /**
- * Creates a handler that captures execution errors for Sentry
+ * Creates an error-capturing wrapper for handler execution
  */
 function createErrorCapturingHandler(
   this: MCPServerInstance,
@@ -67,19 +64,15 @@ function createErrorCapturingHandler(
   try {
     const result = originalHandler.apply(this, handlerArgs);
 
-    // Handle both sync and async handlers
-    if (result && typeof result === 'object' && 'then' in result) {
-      // Async handler - wrap with error capture
-      return (result as Promise<unknown>).catch((error: Error) => {
+    if (result && typeof result === 'object' && typeof (result as { then?: unknown }).then === 'function') {
+      return Promise.resolve(result).catch(error => {
         captureHandlerError(error, methodName, handlerName, handlerArgs, extraHandlerData);
         throw error; // Re-throw to maintain MCP error handling behavior
       });
     }
 
-    // Sync handler - return result as-is
     return result;
   } catch (error) {
-    // Sync handler threw an error - capture it
     captureHandlerError(error as Error, methodName, handlerName, handlerArgs, extraHandlerData);
     throw error; // Re-throw to maintain MCP error handling behavior
   }

packages/core/src/integrations/mcp-server/transport.ts

@@ -6,7 +6,7 @@
 import { getIsolationScope, withIsolationScope } from '../../currentScopes';
 import { startInactiveSpan, withActiveSpan } from '../../tracing';
 import { fill } from '../../utils/object';
-import { cleanupAllPendingSpans, completeSpanWithResults, storeSpanForRequest } from './correlation';
+import { cleanupPendingSpansForTransport, completeSpanWithResults, storeSpanForRequest } from './correlation';
 import { captureError } from './errorCapture';
 import { buildMcpServerSpanConfig, createMcpNotificationSpan, createMcpOutgoingNotificationSpan } from './spans';
 import type { ExtraHandlerData, MCPTransport } from './types';
@@ -30,8 +30,7 @@ export function wrapTransportOnMessage(transport: MCPTransport): void {
             const spanConfig = buildMcpServerSpanConfig(jsonRpcMessage, this, extra as ExtraHandlerData);
             const span = startInactiveSpan(spanConfig);
 
-            // Store span context for handler correlation using requestId
-            storeSpanForRequest(messageTyped.id, span, messageTyped.method);
+            storeSpanForRequest(this, messageTyped.id, span, messageTyped.method);
 
             // Execute handler within span context
             return withActiveSpan(span, () => {
@@ -74,7 +73,7 @@ export function wrapTransportSend(transport: MCPTransport): void {
               captureJsonRpcErrorResponse(messageTyped.error, messageTyped.id, this);
             }
 
-            completeSpanWithResults(messageTyped.id, messageTyped.result);
+            completeSpanWithResults(this, messageTyped.id, messageTyped.result);
           }
         }
 
@@ -85,13 +84,14 @@ export function wrapTransportSend(transport: MCPTransport): void {
 }
 
 /**
- * Wraps transport.onclose to clean up pending spans
+ * Wraps transport.onclose to clean up pending spans for this transport only
  */
 export function wrapTransportOnClose(transport: MCPTransport): void {
   if (transport.onclose) {
     fill(transport, 'onclose', originalOnClose => {
       return function (this: MCPTransport, ...args: unknown[]) {
-        cleanupAllPendingSpans();
+        // Clean up only spans associated with this specific transport
+        cleanupPendingSpansForTransport(this);
 
         return (originalOnClose as (...args: unknown[]) => unknown).call(this, ...args);
       };