fix(core): handle stateless MCP wrapper transport correlation (#20293)

This PR handles stateless MCP wrappers by using the connected transport
instance as the stable MCP correlation key across wrapper layers.

I verified the issue and the fix by adding a regression test for
stateless wrapper transports where the request path runs on the inner
transport and the response path runs on the wrapper.

Closes #20290

---------

Co-authored-by: GPT-5 <noreply@anthropic.com>

Author: logaretm

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

@@ -43,7 +43,7 @@ export function wrapTransportOnMessage(transport: MCPTransport, options: Resolve
           if (isInitialize) {
             try {
               initSessionData = extractSessionDataFromInitializeRequest(message);
-              storeSessionDataForTransport(this, initSessionData);
+              storeSessionDataForTransport(transport, initSessionData);
             } catch {
               // noop
             }
@@ -52,7 +52,7 @@ export function wrapTransportOnMessage(transport: MCPTransport, options: Resolve
           const isolationScope = getIsolationScope().clone();
 
           return withIsolationScope(isolationScope, () => {
-            const spanConfig = buildMcpServerSpanConfig(message, this, extra as ExtraHandlerData, options);
+            const spanConfig = buildMcpServerSpanConfig(message, transport, extra as ExtraHandlerData, options);
             const span = startInactiveSpan(spanConfig);
 
             // For initialize requests, add client info directly to span (works even for stateless transports)
@@ -65,7 +65,7 @@ export function wrapTransportOnMessage(transport: MCPTransport, options: Resolve
               });
             }
 
-            storeSpanForRequest(this, message.id, span, message.method);
+            storeSpanForRequest(transport, message.id, span, message.method);
 
             return withActiveSpan(span, () => {
               return (originalOnMessage as (...args: unknown[]) => unknown).call(this, message, extra);
@@ -74,7 +74,7 @@ export function wrapTransportOnMessage(transport: MCPTransport, options: Resolve
         }
 
         if (isJsonRpcNotification(message)) {
-          return createMcpNotificationSpan(message, this, extra as ExtraHandlerData, options, () => {
+          return createMcpNotificationSpan(message, transport, extra as ExtraHandlerData, options, () => {
             return (originalOnMessage as (...args: unknown[]) => unknown).call(this, message, extra);
           });
         }
@@ -99,7 +99,7 @@ export function wrapTransportSend(transport: MCPTransport, options: ResolvedMcpO
         const [message] = args;
 
         if (isJsonRpcNotification(message)) {
-          return createMcpOutgoingNotificationSpan(message, this, options, () => {
+          return createMcpOutgoingNotificationSpan(message, transport, options, () => {
             return (originalSend as (...args: unknown[]) => unknown).call(this, ...args);
           });
         }
@@ -114,14 +114,14 @@ export function wrapTransportSend(transport: MCPTransport, options: ResolvedMcpO
               if (message.result.protocolVersion || message.result.serverInfo) {
                 try {
                   const serverData = extractSessionDataFromInitializeResponse(message.result);
-                  updateSessionDataForTransport(this, serverData);
+                  updateSessionDataForTransport(transport, serverData);
                 } catch {
                   // noop
                 }
               }
             }
 
-            completeSpanWithResults(this, message.id, message.result, options, !!message.error);
+            completeSpanWithResults(transport, message.id, message.result, options, !!message.error);
           }
         }
 
@@ -139,8 +139,8 @@ export function wrapTransportOnClose(transport: MCPTransport): void {
   if (transport.onclose) {
     fill(transport, 'onclose', originalOnClose => {
       return function (this: MCPTransport, ...args: unknown[]) {
-        cleanupPendingSpansForTransport(this);
-        cleanupSessionDataForTransport(this);
+        cleanupPendingSpansForTransport(transport);
+        cleanupSessionDataForTransport(transport);
         return (originalOnClose as (...args: unknown[]) => unknown).call(this, ...args);
       };
     });

packages/core/test/lib/integrations/mcp-server/transportInstrumentation.test.ts

@@ -880,6 +880,95 @@ describe('MCP Server Transport Instrumentation', () => {
       expect(mockSpan.end).toHaveBeenCalled();
     });
 
+    it('should correlate spans correctly for stateless wrapper transports', async () => {
+      const { wrapper, inner } = createMockWrapperTransport('stateless-wrapper-session');
+      inner.sessionId = undefined;
+
+      const mockMcpServer = createMockMcpServer();
+      const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
+
+      await wrappedMcpServer.connect(wrapper);
+
+      const mockSpan = { setAttributes: vi.fn(), end: vi.fn() };
+      startInactiveSpanSpy.mockReturnValue(mockSpan as any);
+
+      inner.onmessage?.call(
+        inner,
+        {
+          jsonrpc: '2.0',
+          method: 'tools/call',
+          id: 'stateless-wrapper-req-1',
+          params: { name: 'test-tool' },
+        },
+        {},
+      );
+
+      await wrapper.send({
+        jsonrpc: '2.0',
+        id: 'stateless-wrapper-req-1',
+        result: { content: [{ type: 'text', text: 'success' }] },
+      });
+
+      expect(mockSpan.end).toHaveBeenCalled();
+    });
+
+    it('should preserve session metadata for later stateless wrapper spans', async () => {
+      const { wrapper, inner } = createMockWrapperTransport('stateless-wrapper-session');
+      inner.sessionId = undefined;
+
+      const mockMcpServer = createMockMcpServer();
+      const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
+
+      await wrappedMcpServer.connect(wrapper);
+
+      inner.onmessage?.call(
+        inner,
+        {
+          jsonrpc: '2.0',
+          method: 'initialize',
+          id: 'init-stateless',
+          params: {
+            protocolVersion: '2025-06-18',
+            clientInfo: { name: 'test-client', version: '1.0.0' },
+          },
+        },
+        {},
+      );
+
+      await wrapper.send({
+        jsonrpc: '2.0',
+        id: 'init-stateless',
+        result: {
+          protocolVersion: '2025-06-18',
+          serverInfo: { name: 'test-server', version: '2.0.0' },
+          capabilities: {},
+        },
+      });
+
+      inner.onmessage?.call(
+        inner,
+        {
+          jsonrpc: '2.0',
+          method: 'tools/call',
+          id: 'stateless-wrapper-req-2',
+          params: { name: 'test-tool' },
+        },
+        {},
+      );
+
+      expect(startInactiveSpanSpy).toHaveBeenCalledWith(
+        expect.objectContaining({
+          attributes: expect.objectContaining({
+            'mcp.client.name': 'test-client',
+            'mcp.client.version': '1.0.0',
+            'mcp.protocol.version': '2025-06-18',
+            'mcp.server.name': 'test-server',
+            'mcp.server.version': '2.0.0',
+          }),
+        }),
+      );
+    });
+
     it('should handle initialize request/response with wrapper transport', async () => {
       const { wrapper } = createMockWrapperTransport('init-wrapper-session');
       const mockMcpServer = createMockMcpServer();