From 2c46be688ac05fd17fd2bce8abaad1e5f2389a84 Mon Sep 17 00:00:00 2001 From: betegon Date: Wed, 2 Sep 2026 20:03:29 +0200 Subject: [PATCH] fix(core): instrument MCP transports before start Co-Authored-By: OpenAI Codex --- .../mcp-server-streamed/instrument-otel.mjs | 19 +++ .../mcp-server-streamed/instrument.mjs | 9 ++ .../mcp-server-streamed/scenario-start-v2.mjs | 30 ++++ .../mcp-server-streamed/scenario-v1.mjs | 31 ++++ .../tracing/mcp-server-streamed/test.ts | 38 +++++ .../core/src/integrations/mcp-server/index.ts | 109 ++++++++++++-- .../core/src/integrations/mcp-server/types.ts | 3 + .../transportInstrumentation.test.ts | 135 ++++++++++++++++++ 8 files changed, 362 insertions(+), 12 deletions(-) create mode 100644 dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/instrument-otel.mjs create mode 100644 dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/scenario-start-v2.mjs create mode 100644 dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/scenario-v1.mjs diff --git a/dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/instrument-otel.mjs b/dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/instrument-otel.mjs new file mode 100644 index 000000000000..d13348de08b0 --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/instrument-otel.mjs @@ -0,0 +1,19 @@ +import * as Sentry from '@sentry/node'; +import { loggingTransport } from '@sentry-internal/node-integration-tests'; + +Sentry.init({ + dsn: 'https://public@dsn.ingest.sentry.io/1337', + release: '1.0', + tracesSampleRate: 1.0, + transport: loggingTransport, + enableOpenTelemetrySetup: true, +}); + +let initializeSpansStarted = 0; +Sentry.getClient()?.on('spanStart', span => { + const attributes = Sentry.spanToJSON(span).attributes; + if (attributes['sentry.op'] === 'mcp.server' && attributes['mcp.method.name'] === 'initialize') { + initializeSpansStarted += 1; + span.setAttribute('test.mcp.initialize_spans_started', initializeSpansStarted); + } +}); diff --git a/dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/instrument.mjs b/dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/instrument.mjs index 46a27dd03b74..4666b1a70da3 100644 --- a/dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/instrument.mjs +++ b/dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/instrument.mjs @@ -7,3 +7,12 @@ Sentry.init({ tracesSampleRate: 1.0, transport: loggingTransport, }); + +let initializeSpansStarted = 0; +Sentry.getClient()?.on('spanStart', span => { + const attributes = Sentry.spanToJSON(span).attributes; + if (attributes['sentry.op'] === 'mcp.server' && attributes['mcp.method.name'] === 'initialize') { + initializeSpansStarted += 1; + span.setAttribute('test.mcp.initialize_spans_started', initializeSpansStarted); + } +}); diff --git a/dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/scenario-start-v2.mjs b/dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/scenario-start-v2.mjs new file mode 100644 index 000000000000..1f51dbd77342 --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/scenario-start-v2.mjs @@ -0,0 +1,30 @@ +import { Client } from '@modelcontextprotocol/client'; +import { InMemoryTransport, McpServer } from '@modelcontextprotocol/server'; +import { wrapMcpServerWithSentry } from '@sentry/node'; + +const server = wrapMcpServerWithSentry(new McpServer({ name: 'Echo', version: '1.0.0' })); + +async function run() { + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + const client = new Client({ name: 'test-client', version: '1.0.0' }, { versionNegotiation: { mode: 'legacy' } }); + const originalSend = clientTransport.send.bind(clientTransport); + const requestQueued = new Promise(resolve => { + clientTransport.send = async (...args) => { + const result = await originalSend(...args); + if (args[0]?.method === 'initialize') { + resolve(); + } + return result; + }; + }); + + const clientConnection = client.connect(clientTransport); + await requestQueued; + await server.connect(serverTransport); + await clientConnection; + + await client.close(); + await server.close(); +} + +run(); diff --git a/dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/scenario-v1.mjs b/dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/scenario-v1.mjs new file mode 100644 index 000000000000..1bc775aef562 --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/scenario-v1.mjs @@ -0,0 +1,31 @@ +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { InMemoryTransport } from '@modelcontextprotocol/sdk/inMemory.js'; +import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; +import { wrapMcpServerWithSentry } from '@sentry/node'; + +const server = wrapMcpServerWithSentry(new McpServer({ name: 'Echo', version: '1.0.0' })); + +async function run() { + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + const client = new Client({ name: 'test-client', version: '1.0.0' }); + const originalSend = clientTransport.send.bind(clientTransport); + const requestQueued = new Promise(resolve => { + clientTransport.send = async (...args) => { + const result = await originalSend(...args); + if (args[0]?.method === 'initialize') { + resolve(); + } + return result; + }; + }); + + const clientConnection = client.connect(clientTransport); + await requestQueued; + await server.connect(serverTransport); + await clientConnection; + + await client.close(); + await server.close(); +} + +run(); diff --git a/dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/test.ts b/dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/test.ts index eb146ca33fbd..c6ccc8fbfe72 100644 --- a/dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/test.ts +++ b/dev-packages/node-integration-tests/suites/tracing/mcp-server-streamed/test.ts @@ -6,6 +6,20 @@ function mcpSpans(container: SerializedStreamedSpanContainer): SerializedStreame return container.items.filter(item => item.attributes['sentry.op']?.value === 'mcp.server'); } +function assertInitializeSpan(container: SerializedStreamedSpanContainer): void { + const initializeSpans = mcpSpans(container).filter( + span => span.attributes['mcp.method.name']?.value === 'initialize', + ); + + expect(initializeSpans).toHaveLength(1); + const initializeSpan = initializeSpans[0]!; + expect(initializeSpan.name).toBe('initialize'); + expect(initializeSpan.status).toBe('ok'); + expect(initializeSpan.attributes['sentry.op']).toEqual({ type: 'string', value: 'mcp.server' }); + expect(initializeSpan.attributes['sentry.origin']).toEqual({ type: 'string', value: 'auto.function.mcp_server' }); + expect(initializeSpan.attributes['test.mcp.initialize_spans_started']).toEqual({ type: 'integer', value: 1 }); +} + describe('MCP server spans (streamed)', () => { afterAll(() => { cleanupChildProcesses(); @@ -43,4 +57,28 @@ describe('MCP server spans (streamed)', () => { .completed(); }); }); + + createEsmAndCjsTests(__dirname, 'scenario-start-v2.mjs', 'instrument.mjs', (createTestRunner, test) => { + test('captures an MCP v2 initialize request queued before transport start once', async () => { + await createTestRunner().expect({ span: assertInitializeSpan }).start().completed(); + }); + }); + + createEsmAndCjsTests( + __dirname, + 'scenario-v1.mjs', + 'instrument.mjs', + (createTestRunner, test) => { + test('captures an MCP v1 initialize request queued before transport start once', async () => { + await createTestRunner().expect({ span: assertInitializeSpan }).start().completed(); + }); + }, + { additionalDependencies: { '@modelcontextprotocol/sdk': '1.30.0' } }, + ); + + createEsmAndCjsTests(__dirname, 'scenario-start-v2.mjs', 'instrument-otel.mjs', (createTestRunner, test) => { + test('captures the queued request with Sentry OpenTelemetry setup enabled', async () => { + await createTestRunner().expect({ span: assertInitializeSpan }).start().completed(); + }); + }); }); diff --git a/packages/core/src/integrations/mcp-server/index.ts b/packages/core/src/integrations/mcp-server/index.ts index 62521def4d01..192372462a9e 100644 --- a/packages/core/src/integrations/mcp-server/index.ts +++ b/packages/core/src/integrations/mcp-server/index.ts @@ -10,6 +10,79 @@ import { validateMcpServerInstance } from './validation'; */ const wrappedMcpServerInstances = new WeakSet(); +function instrumentTransport(transport: MCPTransport, options: McpServerWrapperOptions): void { + wrapTransportOnMessage(transport, options); + wrapTransportSend(transport, options); + wrapTransportOnClose(transport); + wrapTransportError(transport); +} + +function interceptTransportStart(transport: MCPTransport, beforeStart: () => void): () => void { + let transportStart: MCPTransport['start']; + let originalDescriptor: PropertyDescriptor | undefined; + + try { + transportStart = transport.start; + originalDescriptor = Object.getOwnPropertyDescriptor(transport, 'start'); + } catch { + return () => undefined; + } + + if (typeof transportStart !== 'function') { + return () => undefined; + } + + const originalStart = transportStart; + let isInstalled = false; + + const restoreStart = (): void => { + if (!isInstalled) { + return; + } + + try { + const currentDescriptor = Object.getOwnPropertyDescriptor(transport, 'start'); + if (currentDescriptor?.value !== interceptedStart) { + isInstalled = false; + return; + } + + if (originalDescriptor) { + Object.defineProperty(transport, 'start', originalDescriptor); + isInstalled = false; + } else if (Reflect.deleteProperty(transport, 'start')) { + isInstalled = false; + } + } catch {} + }; + + function interceptedStart(this: MCPTransport): Promise { + // Restoring first keeps recursive calls and user-observed method identity identical to the original transport. + restoreStart(); + beforeStart(); + return originalStart.call(this); + } + + const replacementDescriptor: PropertyDescriptor = + originalDescriptor && 'value' in originalDescriptor + ? { ...originalDescriptor, value: interceptedStart } + : { + configurable: originalDescriptor?.configurable ?? true, + enumerable: originalDescriptor?.enumerable ?? false, + writable: true, + value: interceptedStart, + }; + + try { + Object.defineProperty(transport, 'start', replacementDescriptor); + isInstalled = true; + } catch { + // The post-connect fallback preserves the previous behavior for transports which cannot be patched. + } + + return restoreStart; +} + /** * Wraps an MCP Server instance with Sentry instrumentation. * @@ -63,18 +136,30 @@ export function wrapMcpServerWithSentry(mcpServerInstance: S, fill(serverInstance, 'connect', originalConnect => { return async function (this: MCPServerInstance, transport: MCPTransport, ...restArgs: unknown[]) { - const result = await (originalConnect as (...args: unknown[]) => Promise).call( - this, - transport, - ...restArgs, - ); - - wrapTransportOnMessage(transport, captureOptions); - wrapTransportSend(transport, captureOptions); - wrapTransportOnClose(transport); - wrapTransportError(transport); - - return result; + let isTransportInstrumented = false; + const instrumentTransportOnce = (): void => { + if (isTransportInstrumented) { + return; + } + + isTransportInstrumented = true; + instrumentTransport(transport, captureOptions); + }; + const restoreStart = interceptTransportStart(transport, instrumentTransportOnce); + + try { + const result = await (originalConnect as (...args: unknown[]) => Promise).call( + this, + transport, + ...restArgs, + ); + + instrumentTransportOnce(); + + return result; + } finally { + restoreStart(); + } }; }); diff --git a/packages/core/src/integrations/mcp-server/types.ts b/packages/core/src/integrations/mcp-server/types.ts index 7cca7051b785..dc8cceb7d896 100644 --- a/packages/core/src/integrations/mcp-server/types.ts +++ b/packages/core/src/integrations/mcp-server/types.ts @@ -65,6 +65,9 @@ export interface JsonRpcNotification { * @description Abstraction for MCP communication transport layer */ export interface MCPTransport { + /** Starts the transport lifecycle. */ + start?: () => Promise; + /** * Message handler for incoming JSON-RPC messages * The first argument is a JSON RPC message 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 e97745c57253..1c4b0d3234bf 100644 --- a/packages/core/test/lib/integrations/mcp-server/transportInstrumentation.test.ts +++ b/packages/core/test/lib/integrations/mcp-server/transportInstrumentation.test.ts @@ -34,6 +34,33 @@ import { createMockWrapperTransport, } from './testUtils'; +type StartImplementation = (transport: InMemoryTransport) => Promise; + +class InMemoryTransport { + public onmessage?: (...args: unknown[]) => void; + public onclose?: (...args: unknown[]) => void; + public onerror?: (error: Error) => void; + public send = vi.fn().mockResolvedValue(undefined); + + public constructor(private readonly startImplementation: StartImplementation = () => Promise.resolve()) {} + + public start(): Promise { + return this.startImplementation(this); + } +} + +function createStartingMcpServer() { + return { + ...createMockMcpServer(), + connect: vi.fn(async (transport: InMemoryTransport) => { + transport.onmessage = vi.fn(); + transport.onclose = vi.fn(); + transport.onerror = vi.fn(); + await transport.start(); + }), + }; +} + describe('MCP Server Transport Instrumentation', () => { const startSpanSpy = vi.spyOn(tracingModule, 'startSpan'); const startInactiveSpanSpy = vi.spyOn(tracingModule, 'startInactiveSpan'); @@ -101,6 +128,114 @@ describe('MCP Server Transport Instrumentation', () => { expect(originalConnect).toHaveBeenCalledWith(mockTransport); }); + it('instruments a request delivered while the transport starts', async () => { + const transport = new InMemoryTransport(connectedTransport => { + connectedTransport.onmessage?.({ + jsonrpc: '2.0', + method: 'tools/call', + id: 'queued-request', + params: { name: 'get-weather' }, + }); + return Promise.resolve(); + }); + + await wrapMcpServerWithSentry(createStartingMcpServer()).connect(transport); + + expect(startInactiveSpanSpy).toHaveBeenCalledOnce(); + expect(startInactiveSpanSpy).toHaveBeenCalledWith({ + name: 'tools/call get-weather', + forceTransaction: true, + attributes: { + 'mcp.method.name': 'tools/call', + 'mcp.tool.name': 'get-weather', + 'mcp.request.id': 'queued-request', + 'mcp.transport': 'InMemoryTransport', + 'network.transport': 'unknown', + 'network.protocol.version': '2.0', + 'sentry.op': 'mcp.server', + 'sentry.origin': 'auto.function.mcp_server', + 'sentry.segment.name.source': 'route', + }, + }); + }); + + it('preserves the start receiver and restores an inherited method before calling it', async () => { + let receivedExpectedThis = false; + let wasRestoredBeforeStart = false; + let originalStart: InMemoryTransport['start']; + const transport = new InMemoryTransport(connectedTransport => { + receivedExpectedThis = connectedTransport === transport; + wasRestoredBeforeStart = connectedTransport.start === originalStart; + return Promise.resolve(); + }); + originalStart = transport.start; + + await wrapMcpServerWithSentry(createStartingMcpServer()).connect(transport); + + expect(receivedExpectedThis).toBe(true); + expect(wasRestoredBeforeStart).toBe(true); + expect(transport.start).toBe(originalStart); + expect(Object.prototype.hasOwnProperty.call(transport, 'start')).toBe(false); + }); + + it('restores start when connect rejects before starting the transport', async () => { + const connectionError = new Error('connection failed'); + const transport = new InMemoryTransport(); + const server = { + ...createMockMcpServer(), + connect: vi.fn().mockRejectedValue(connectionError), + }; + + const connection = wrapMcpServerWithSentry(server).connect(transport); + + await expect(connection).rejects.toBe(connectionError); + expect(Object.prototype.hasOwnProperty.call(transport, 'start')).toBe(false); + }); + + it('restores start and preserves a synchronous start error', async () => { + const startError = new Error('start failed'); + const originalStart = vi.fn(() => { + throw startError; + }); + const transport = new InMemoryTransport(); + Object.defineProperty(transport, 'start', { + configurable: true, + enumerable: false, + value: originalStart, + writable: false, + }); + const originalDescriptor = Object.getOwnPropertyDescriptor(transport, 'start'); + + const connection = wrapMcpServerWithSentry(createStartingMcpServer()).connect(transport); + + await expect(connection).rejects.toBe(startError); + expect(Object.getOwnPropertyDescriptor(transport, 'start')).toEqual(originalDescriptor); + expect(originalStart).toHaveBeenCalledOnce(); + }); + + it('falls back to post-connect instrumentation when start cannot be replaced', async () => { + const originalStart = vi.fn().mockResolvedValue(undefined); + const transport = new InMemoryTransport(); + Object.defineProperty(transport, 'start', { + configurable: false, + enumerable: false, + value: originalStart, + writable: false, + }); + + await wrapMcpServerWithSentry(createStartingMcpServer()).connect(transport); + + transport.onmessage?.({ + jsonrpc: '2.0', + method: 'tools/call', + id: 'post-start-request', + params: { name: 'get-weather' }, + }); + + expect(originalStart).toHaveBeenCalledOnce(); + expect(startInactiveSpanSpy).toHaveBeenCalledOnce(); + }); + it('should create spans for incoming JSON-RPC requests', async () => { await wrappedMcpServer.connect(mockTransport);