diff --git a/packages/core/src/integrations/mcp-server/transport.ts b/packages/core/src/integrations/mcp-server/transport.ts index 4b04a78f43b0..49b141964fa9 100644 --- a/packages/core/src/integrations/mcp-server/transport.ts +++ b/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); }; }); diff --git a/packages/core/test/lib/integrations/mcp-server/transportInstrumentation.test.ts b/packages/core/test/lib/integrations/mcp-server/transportInstrumentation.test.ts index 71590fd17712..d8cc546393ec 100644 --- a/packages/core/test/lib/integrations/mcp-server/transportInstrumentation.test.ts +++ b/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();