diff --git a/dev-packages/node-integration-tests/suites/tracing/mcp-server/trace-injection/instrument.mjs b/dev-packages/node-integration-tests/suites/tracing/mcp-server/trace-injection/instrument.mjs new file mode 100644 index 000000000000..46a27dd03b74 --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/mcp-server/trace-injection/instrument.mjs @@ -0,0 +1,9 @@ +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, +}); diff --git a/dev-packages/node-integration-tests/suites/tracing/mcp-server/trace-injection/scenario.mjs b/dev-packages/node-integration-tests/suites/tracing/mcp-server/trace-injection/scenario.mjs new file mode 100644 index 000000000000..d095f50ae58e --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/mcp-server/trace-injection/scenario.mjs @@ -0,0 +1,44 @@ +import * as Sentry from '@sentry/node'; +import { Client } from '@modelcontextprotocol/client'; +import { InMemoryTransport, McpServer } from '@modelcontextprotocol/server'; + +const server = new McpServer({ name: 'Echo', version: '1.0.0' }, { capabilities: { logging: {} } }); + +server.registerTool('echo', {}, async ctx => { + await ctx.mcpReq.notify({ method: 'notifications/message', params: { level: 'info', data: 'echoing' } }); + await ctx.mcpReq.elicitInput({ + message: 'Confirm echo', + requestedSchema: { type: 'object', properties: { confirm: { type: 'boolean' } } }, + }); + return { content: [{ type: 'text', text: 'Tool echo' }] }; +}); + +async function run() { + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + const client = new Client({ name: 'test-client', version: '1.0.0' }, { capabilities: { elicitation: { form: {} } } }); + + // The client is not instrumented, so it continues the trace by hand to show what the server put on the wire. + const continueFromMeta = (params, name) => { + const meta = params._meta ?? {}; + Sentry.continueTrace({ sentryTrace: meta['sentry-trace'], baggage: meta.baggage }, () => { + Sentry.startSpan({ name, op: 'test.client', forceTransaction: true }, () => undefined); + }); + }; + + client.setNotificationHandler('notifications/message', notification => { + continueFromMeta(notification.params, 'handle log message'); + }); + client.setRequestHandler('elicitation/create', request => { + continueFromMeta(request.params, 'handle elicitation'); + return { action: 'accept', content: { confirm: true } }; + }); + + await Promise.all([server.connect(serverTransport), client.connect(clientTransport)]); + + await client.callTool({ name: 'echo', arguments: {} }); + + await client.close(); + await server.close(); +} + +run(); diff --git a/dev-packages/node-integration-tests/suites/tracing/mcp-server/trace-injection/test.ts b/dev-packages/node-integration-tests/suites/tracing/mcp-server/trace-injection/test.ts new file mode 100644 index 000000000000..1b38c962efd7 --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/mcp-server/trace-injection/test.ts @@ -0,0 +1,46 @@ +import { MCP_METHOD_NAME, SENTRY_OP } from '@sentry/conventions/attributes'; +import { MCP_NOTIFICATION_SERVER_TO_CLIENT, MCP_SERVER } from '@sentry/conventions/op'; +import type { SerializedStreamedSpan } from '@sentry/core'; +import { afterAll, describe, expect } from 'vitest'; +import { cleanupChildProcesses, createEsmAndCjsTests } from '../../../../utils/runner'; + +describe('MCP server trace injection into outgoing params._meta', () => { + afterAll(() => { + cleanupChildProcesses(); + }); + + createEsmAndCjsTests(__dirname, 'scenario.mjs', 'instrument.mjs', (createTestRunner, test) => { + test('lets the client continue the trace from outgoing notifications and requests', async () => { + const isToolCall = (span: SerializedStreamedSpan): boolean => + span.attributes[SENTRY_OP]?.value === MCP_SERVER && span.attributes[MCP_METHOD_NAME]?.value === 'tools/call'; + const isNotification = (span: SerializedStreamedSpan): boolean => + span.attributes[SENTRY_OP]?.value === MCP_NOTIFICATION_SERVER_TO_CLIENT; + const isLogHandler = (span: SerializedStreamedSpan): boolean => span.name === 'handle log message'; + const isElicitationHandler = (span: SerializedStreamedSpan): boolean => span.name === 'handle elicitation'; + + const runner = createTestRunner(); + // The spans are separate segments and can arrive in different envelopes, so collect them across envelopes by trace. + const spansPromise = runner.collectStreamedSpans( + spans => + spans.some(isToolCall) && + spans.some(isNotification) && + spans.some(isLogHandler) && + spans.some(isElicitationHandler), + ); + + await runner.start().completed(); + + const spans = await spansPromise; + const toolCall = spans.find(isToolCall)!; + const notification = spans.find(isNotification)!; + const logHandler = spans.find(isLogHandler)!; + const elicitationHandler = spans.find(isElicitationHandler)!; + + expect(notification.parent_span_id).toBe(toolCall.span_id); + expect(logHandler.trace_id).toBe(toolCall.trace_id); + expect(logHandler.parent_span_id).toBe(notification.span_id); + expect(elicitationHandler.trace_id).toBe(toolCall.trace_id); + expect(elicitationHandler.parent_span_id).toBe(toolCall.span_id); + }); + }); +}); diff --git a/packages/core/src/integrations/mcp-server/tracePropagation.ts b/packages/core/src/integrations/mcp-server/tracePropagation.ts index 2ea59139a5f3..5a40f7f65422 100644 --- a/packages/core/src/integrations/mcp-server/tracePropagation.ts +++ b/packages/core/src/integrations/mcp-server/tracePropagation.ts @@ -6,9 +6,12 @@ import { getClient } from '../../currentScopes'; import { continueTrace } from '../../tracing/trace'; import type { SpanLink } from '../../types/link'; +import { mergeBaggageHeaders } from '../../utils/baggage'; import { isObjectLike } from '../../utils/is'; import { getActiveSpan } from '../../utils/spanUtils'; +import { getTraceData } from '../../utils/traceData'; import { extractTraceparentData, sentryTraceFromTraceparent } from '../../utils/tracing'; +import type { JsonRpcNotification, JsonRpcRequest } from './types'; function getSentryTraceFromMeta(meta: Record): string | undefined { const sentryTrace = meta['sentry-trace']; @@ -46,3 +49,43 @@ export function continueMcpTrace(params: unknown, callback: (links: SpanLink[ return continueTrace({ sentryTrace, baggage }, () => callback(links)); } + +/** + * Returns a copy of an outgoing request or notification with the active trace context added to `params._meta`, so the + * receiving side can continue the trace. Messages that already carry trace context are returned unchanged. + */ +export function injectMcpTraceContext(message: M): M { + const meta = message.params?._meta; + if ( + meta !== undefined && + (!isObjectLike(meta) || meta['sentry-trace'] !== undefined || meta.traceparent !== undefined) + ) { + return message; + } + + const { + 'sentry-trace': sentryTrace, + baggage, + traceparent, + } = getTraceData({ + propagateTraceparent: getClient()?.getOptions().propagateTraceparent, + }); + if (!sentryTrace) { + return message; + } + + const existingBaggage = isObjectLike(meta) && typeof meta.baggage === 'string' ? meta.baggage : undefined; + + return { + ...message, + params: { + ...message.params, + _meta: { + ...meta, + 'sentry-trace': sentryTrace, + ...(traceparent && { traceparent }), + ...(baggage && { baggage: mergeBaggageHeaders(existingBaggage, baggage) }), + }, + }, + }; +} diff --git a/packages/core/src/integrations/mcp-server/transport.ts b/packages/core/src/integrations/mcp-server/transport.ts index aed5745c521d..168cefaef2a4 100644 --- a/packages/core/src/integrations/mcp-server/transport.ts +++ b/packages/core/src/integrations/mcp-server/transport.ts @@ -18,7 +18,7 @@ import { captureError, isJsonRpcServerError } from './errorCapture'; import { buildClientAttributesFromInfo, extractSessionDataFromInitializeRequest } from './sessionExtraction'; import { cleanupSessionDataForTransport, updateSessionDataForTransport } from './sessionManagement'; import { buildMcpServerSpanConfig, createMcpNotificationSpan, createMcpOutgoingNotificationSpan } from './spans'; -import { continueMcpTrace } from './tracePropagation'; +import { continueMcpTrace, injectMcpTraceContext } from './tracePropagation'; import type { ExtraHandlerData, McpServerWrapperOptions, MCPTransport, ResolvedMcpOptions, SessionData } from './types'; import { isJsonRpcNotification, isJsonRpcRequest, isJsonRpcResponse } from './validation'; @@ -153,15 +153,27 @@ export function wrapTransportSend(transport: MCPTransport, options: McpServerWra if (transport.send) { fill(transport, 'send', originalSend => { return async function (this: MCPTransport, ...args: unknown[]) { - const [message] = args; + const [message, ...restArgs] = args; if (isJsonRpcNotification(message)) { const resolvedOptions = resolveMcpOptions(options); return createMcpOutgoingNotificationSpan(message, transport, resolvedOptions, () => { - return (originalSend as (...args: unknown[]) => unknown).call(this, ...args); + return (originalSend as (...args: unknown[]) => unknown).call( + this, + injectMcpTraceContext(message), + ...restArgs, + ); }); } + if (isJsonRpcRequest(message)) { + return (originalSend as (...args: unknown[]) => unknown).call( + this, + injectMcpTraceContext(message), + ...restArgs, + ); + } + if (isJsonRpcResponse(message)) { if (message.id !== null && message.id !== undefined) { if (message.error) { diff --git a/packages/core/test/lib/integrations/mcp-server/testUtils.ts b/packages/core/test/lib/integrations/mcp-server/testUtils.ts index 27d81a65e5cb..d00d222056e6 100644 --- a/packages/core/test/lib/integrations/mcp-server/testUtils.ts +++ b/packages/core/test/lib/integrations/mcp-server/testUtils.ts @@ -31,6 +31,7 @@ export function createMockClient( frameContextLines: 5, }), getDsn: () => ({ publicKey: 'test-key', host: 'test-host' }), + getTransport: () => undefined, emit: vi.fn(), } as unknown as Client; } @@ -43,6 +44,7 @@ export function createMockClientWithDataCollection(dataCollection: DataCollectio getOptions: () => ({ traceLifecycle: 'static', dataCollection }), getDataCollectionOptions: () => resolveDataCollectionOptions({ dataCollection }), getDsn: () => ({ publicKey: 'test-key', host: 'test-host' }), + getTransport: () => undefined, emit: vi.fn(), } as unknown as Client; } diff --git a/packages/core/test/lib/integrations/mcp-server/tracePropagation.test.ts b/packages/core/test/lib/integrations/mcp-server/tracePropagation.test.ts index 5ca20b4eaed5..9fb47454fd52 100644 --- a/packages/core/test/lib/integrations/mcp-server/tracePropagation.test.ts +++ b/packages/core/test/lib/integrations/mcp-server/tracePropagation.test.ts @@ -1,6 +1,12 @@ -import { afterEach, beforeEach, describe, expect, it } from 'vitest'; -import { getCurrentScope, getDynamicSamplingContextFromSpan, setCurrentClient, startSpan } from '../../../../src'; -import { wrapTransportOnMessage } from '../../../../src/integrations/mcp-server/transport'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { + getClient, + getCurrentScope, + getDynamicSamplingContextFromSpan, + setCurrentClient, + startSpan, +} from '../../../../src'; +import { wrapTransportOnMessage, wrapTransportSend } from '../../../../src/integrations/mcp-server/transport'; import type { Span } from '../../../../src/types/span'; import { getActiveSpan, spanToJSON } from '../../../../src/utils/spanUtils'; import { getDefaultTestClientOptions, TestClient } from '../../../mocks/client'; @@ -138,3 +144,160 @@ describe('MCP server trace propagation via params._meta', () => { expect(json.links).toBeUndefined(); }); }); + +describe('MCP server trace injection into outgoing params._meta', () => { + beforeEach(() => { + initClient({ dsn: 'https://public@dsn.ingest.sentry.io/1337', propagateTraceparent: true }); + }); + + afterEach(() => { + getCurrentScope().setClient(undefined); + }); + + function wrappedSend() { + const transport = createMockStdioTransport(); + const originalSend = vi.fn().mockResolvedValue(undefined); + transport.send = originalSend; + wrapTransportSend(transport, {}); + return { send: transport.send, originalSend }; + } + + it('injects the active span into outgoing requests', async () => { + const { send, originalSend } = wrappedSend(); + const request = { jsonrpc: '2.0', method: 'sampling/createMessage', id: 7, params: { maxTokens: 10 } }; + + let toolSpan: Span | undefined; + await startSpan({ name: 'tools/call get-weather' }, async span => { + toolSpan = span; + await send(request, { relatedRequestId: 1 }); + }); + + const { traceId, spanId } = toolSpan!.spanContext(); + expect(originalSend).toHaveBeenCalledWith( + { + ...request, + params: { + maxTokens: 10, + _meta: { + 'sentry-trace': `${traceId}-${spanId}-1`, + traceparent: `00-${traceId}-${spanId}-01`, + baggage: expect.stringContaining(`sentry-trace_id=${traceId}`), + }, + }, + }, + { relatedRequestId: 1 }, + ); + expect(request.params).toEqual({ maxTokens: 10 }); + }); + + it('parents outgoing notifications to their own notification span', async () => { + const { send, originalSend } = wrappedSend(); + let notificationSpan: Span | undefined; + getClient()!.on('spanStart', span => { + if (spanToJSON(span).name === 'notifications/progress') { + notificationSpan = span; + } + }); + + await startSpan({ name: 'tools/call get-weather' }, () => + send({ jsonrpc: '2.0', method: 'notifications/progress', params: { progressToken: 1, progress: 50 } }), + ); + + const { traceId, spanId } = notificationSpan!.spanContext(); + expect(originalSend).toHaveBeenCalledWith({ + jsonrpc: '2.0', + method: 'notifications/progress', + params: { + progressToken: 1, + progress: 50, + _meta: expect.objectContaining({ + 'sentry-trace': `${traceId}-${spanId}-1`, + traceparent: `00-${traceId}-${spanId}-01`, + }), + }, + }); + }); + + it('only injects traceparent when propagateTraceparent is enabled', async () => { + initClient({ dsn: 'https://public@dsn.ingest.sentry.io/1337' }); + const { send, originalSend } = wrappedSend(); + + await startSpan({ name: 'tools/call get-weather' }, () => send({ jsonrpc: '2.0', method: 'ping', id: 1 })); + + expect(originalSend).toHaveBeenCalledWith({ + jsonrpc: '2.0', + method: 'ping', + id: 1, + params: { _meta: { 'sentry-trace': expect.any(String), baggage: expect.any(String) } }, + }); + }); + + it('keeps existing _meta entries and adds params when there are none', async () => { + const { send, originalSend } = wrappedSend(); + + await startSpan({ name: 'tools/call get-weather' }, async () => { + await send({ jsonrpc: '2.0', method: 'ping', id: 1 }); + await send({ jsonrpc: '2.0', method: 'roots/list', id: 2, params: { _meta: { 'com.example/key': 'value' } } }); + }); + + expect(originalSend).toHaveBeenCalledWith({ + jsonrpc: '2.0', + method: 'ping', + id: 1, + params: { _meta: expect.objectContaining({ 'sentry-trace': expect.any(String) }) }, + }); + expect(originalSend).toHaveBeenCalledWith({ + jsonrpc: '2.0', + method: 'roots/list', + id: 2, + params: { _meta: expect.objectContaining({ 'com.example/key': 'value', 'sentry-trace': expect.any(String) }) }, + }); + }); + + it('merges existing third-party baggage entries with the Sentry entries', async () => { + const { send, originalSend } = wrappedSend(); + + await startSpan({ name: 'tools/call get-weather' }, () => + send({ + jsonrpc: '2.0', + method: 'ping', + id: 1, + params: { _meta: { baggage: 'vendor=abc,sentry-release=stale' } }, + }), + ); + + expect(originalSend).toHaveBeenCalledWith({ + jsonrpc: '2.0', + method: 'ping', + id: 1, + params: { + _meta: expect.objectContaining({ + // The stale Sentry entry is replaced, third-party entries are kept. + baggage: expect.stringMatching(/^vendor=abc,(?!.*stale).*sentry-trace_id=/), + }), + }, + }); + }); + + it.each([ + ['a traceparent', { traceparent: `00-${TRACE_ID}-${PARENT_SPAN_ID}-01` }], + ['a sentry-trace', { 'sentry-trace': `${TRACE_ID}-${PARENT_SPAN_ID}-1` }], + ['a non-object _meta', 'not-an-object'], + ])('leaves requests that already carry %s unchanged', async (_, meta) => { + const { send, originalSend } = wrappedSend(); + const request = { jsonrpc: '2.0', method: 'elicitation/create', id: 3, params: { _meta: meta } }; + + await startSpan({ name: 'tools/call get-weather' }, () => send(request)); + + expect(originalSend).toHaveBeenCalledWith(request); + }); + + it('does not touch responses', async () => { + const { send, originalSend } = wrappedSend(); + const response = { jsonrpc: '2.0', id: 1, result: { content: [] } }; + + await startSpan({ name: 'tools/call get-weather' }, () => send(response)); + + expect(originalSend).toHaveBeenCalledWith(response); + }); +});