Visitar URL original
feat(core): Inject trace context into outgoing MCP server messages · getsentry/sentry-javascript@740f240 · GitHub
Skip to content

Commit 740f240

Browse files
sentry-junior[bot]mydea
sentry-junior[bot]
and
mydea
committed
feat(core): Inject trace context into outgoing MCP server messages
Co-Authored-By: Francesco Gringl-Novy <francesco.novy@sentry.io>
1 parent 074bcf2 commit 740f240

7 files changed

Lines changed: 255 additions & 6 deletions

File tree

‎dev-packages/node-integration-tests/suites/tracing/mcp-server/trace-injection/instrument.mjs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
import * as Sentry from '@sentry/node';
2+
import { loggingTransport } from '@sentry-internal/node-integration-tests';
3+
4+
Sentry.init({
5+
dsn: 'https://public@dsn.ingest.sentry.io/1337',
6+
release: '1.0',
7+
tracesSampleRate: 1.0,
8+
transport: loggingTransport,
9+
});

‎dev-packages/node-integration-tests/suites/tracing/mcp-server/trace-injection/scenario.mjs‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
import * as Sentry from '@sentry/node';
2+
import { Client } from '@modelcontextprotocol/client';
3+
import { InMemoryTransport, McpServer } from '@modelcontextprotocol/server';
4+
5+
const server = new McpServer({ name: 'Echo', version: '1.0.0' }, { capabilities: { logging: {} } });
6+
7+
server.registerTool('echo', {}, async ctx => {
8+
await ctx.mcpReq.notify({ method: 'notifications/message', params: { level: 'info', data: 'echoing' } });
9+
return { content: [{ type: 'text', text: 'Tool echo' }] };
10+
});
11+
12+
async function run() {
13+
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
14+
const client = new Client({ name: 'test-client', version: '1.0.0' });
15+
16+
// The client is not instrumented, so it continues the trace by hand to show what the server put on the wire.
17+
client.setNotificationHandler('notifications/message', notification => {
18+
const meta = notification.params._meta ?? {};
19+
Sentry.continueTrace({ sentryTrace: meta['sentry-trace'], baggage: meta.baggage }, () => {
20+
Sentry.startSpan({ name: 'handle log message', op: 'test.client', forceTransaction: true }, () => undefined);
21+
});
22+
});
23+
24+
await Promise.all([server.connect(serverTransport), client.connect(clientTransport)]);
25+
26+
await client.callTool({ name: 'echo', arguments: {} });
27+
28+
await client.close();
29+
await server.close();
30+
}
31+
32+
run();

‎dev-packages/node-integration-tests/suites/tracing/mcp-server/trace-injection/test.ts‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
import { afterAll, describe, expect } from 'vitest';
2+
import { cleanupChildProcesses, createEsmAndCjsTests } from '../../../../utils/runner';
3+
4+
describe('MCP server trace injection into outgoing params._meta', () => {
5+
afterAll(() => {
6+
cleanupChildProcesses();
7+
});
8+
9+
createEsmAndCjsTests(__dirname, 'scenario.mjs', 'instrument.mjs', (createTestRunner, test) => {
10+
test('lets the client continue the trace from an outgoing notification', async () => {
11+
await createTestRunner()
12+
.unordered()
13+
.expect({
14+
span: container => {
15+
const notification = container.items.find(
16+
item => item.attributes['sentry.op']?.value === 'mcp.notification.server_to_client',
17+
);
18+
const clientSpan = container.items.find(item => item.attributes['sentry.op']?.value === 'test.client');
19+
expect(notification, 'expected the outgoing notification span').toBeDefined();
20+
expect(clientSpan, 'expected the client span').toBeDefined();
21+
expect(clientSpan!.trace_id).toBe(notification!.trace_id);
22+
expect(clientSpan!.parent_span_id).toBe(notification!.span_id);
23+
},
24+
})
25+
.start()
26+
.completed();
27+
});
28+
});
29+
});

‎packages/core/src/integrations/mcp-server/tracePropagation.ts‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,9 @@ import { continueTrace } from '../../tracing/trace';
77
import type { SpanLink } from '../../types/link';
88
import { isObjectLike } from '../../utils/is';
99
import { getActiveSpan } from '../../utils/spanUtils';
10+
import { getTraceData } from '../../utils/traceData';
1011
import { extractTraceparentData, sentryTraceFromTraceparent } from '../../utils/tracing';
12+
import type { JsonRpcNotification, JsonRpcRequest } from './types';
1113

1214
function getSentryTraceFromMeta(meta: Record<PropertyKey, unknown>): string | undefined {
1315
const sentryTrace = meta['sentry-trace'];
@@ -42,3 +44,38 @@ export function continueMcpTrace<T>(params: unknown, callback: (links: SpanLink[
4244

4345
return continueTrace({ sentryTrace, baggage }, () => callback(links));
4446
}
47+
48+
/**
49+
* Returns a copy of an outgoing request or notification with the active trace context added to `params._meta`, so the
50+
* receiving side can continue the trace. Messages that already carry trace context are returned unchanged.
51+
*
52+
* `traceparent` is always added, independent of the `propagateTraceparent` option: that option guards HTTP requests
53+
* against CORS failures, which do not apply here, and W3C Trace Context is what non-Sentry MCP peers read.
54+
*/
55+
export function injectMcpTraceContext<M extends JsonRpcRequest | JsonRpcNotification>(message: M): M {
56+
const meta = message.params?._meta;
57+
if (
58+
meta !== undefined &&
59+
(!isObjectLike(meta) || meta['sentry-trace'] !== undefined || meta.traceparent !== undefined)
60+
) {
61+
return message;
62+
}
63+
64+
const { 'sentry-trace': sentryTrace, baggage, traceparent } = getTraceData({ propagateTraceparent: true });
65+
if (!sentryTrace) {
66+
return message;
67+
}
68+
69+
return {
70+
...message,
71+
params: {
72+
...message.params,
73+
_meta: {
74+
...meta,
75+
'sentry-trace': sentryTrace,
76+
...(traceparent && { traceparent }),
77+
...(baggage && { baggage }),
78+
},
79+
},
80+
};
81+
}

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

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ import { captureError, isJsonRpcServerError } from './errorCapture';
1818
import { buildClientAttributesFromInfo, extractSessionDataFromInitializeRequest } from './sessionExtraction';
1919
import { cleanupSessionDataForTransport, updateSessionDataForTransport } from './sessionManagement';
2020
import { buildMcpServerSpanConfig, createMcpNotificationSpan, createMcpOutgoingNotificationSpan } from './spans';
21-
import { continueMcpTrace } from './tracePropagation';
21+
import { continueMcpTrace, injectMcpTraceContext } from './tracePropagation';
2222
import type { ExtraHandlerData, McpServerWrapperOptions, MCPTransport, ResolvedMcpOptions, SessionData } from './types';
2323
import { isJsonRpcNotification, isJsonRpcRequest, isJsonRpcResponse } from './validation';
2424

@@ -153,15 +153,27 @@ export function wrapTransportSend(transport: MCPTransport, options: McpServerWra
153153
if (transport.send) {
154154
fill(transport, 'send', originalSend => {
155155
return async function (this: MCPTransport, ...args: unknown[]) {
156-
const [message] = args;
156+
const [message, ...restArgs] = args;
157157

158158
if (isJsonRpcNotification(message)) {
159159
const resolvedOptions = resolveMcpOptions(options);
160160
return createMcpOutgoingNotificationSpan(message, transport, resolvedOptions, () => {
161-
return (originalSend as (...args: unknown[]) => unknown).call(this, ...args);
161+
return (originalSend as (...args: unknown[]) => unknown).call(
162+
this,
163+
injectMcpTraceContext(message),
164+
...restArgs,
165+
);
162166
});
163167
}
164168

169+
if (isJsonRpcRequest(message)) {
170+
return (originalSend as (...args: unknown[]) => unknown).call(
171+
this,
172+
injectMcpTraceContext(message),
173+
...restArgs,
174+
);
175+
}
176+
165177
if (isJsonRpcResponse(message)) {
166178
if (message.id !== null && message.id !== undefined) {
167179
if (message.error) {

‎packages/core/test/lib/integrations/mcp-server/testUtils.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ export function createMockClient(
3131
frameContextLines: 5,
3232
}),
3333
getDsn: () => ({ publicKey: 'test-key', host: 'test-host' }),
34+
getTransport: () => undefined,
3435
emit: vi.fn(),
3536
} as unknown as Client;
3637
}
@@ -43,6 +44,7 @@ export function createMockClientWithDataCollection(dataCollection: DataCollectio
4344
getOptions: () => ({ traceLifecycle: 'static', dataCollection }),
4445
getDataCollectionOptions: () => resolveDataCollectionOptions({ dataCollection }),
4546
getDsn: () => ({ publicKey: 'test-key', host: 'test-host' }),
47+
getTransport: () => undefined,
4648
emit: vi.fn(),
4749
} as unknown as Client;
4850
}

‎packages/core/test/lib/integrations/mcp-server/tracePropagation.test.ts‎

Lines changed: 131 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,12 @@
1-
import { afterEach, beforeEach, describe, expect, it } from 'vitest';
2-
import { getCurrentScope, getDynamicSamplingContextFromSpan, setCurrentClient, startSpan } from '../../../../src';
3-
import { wrapTransportOnMessage } from '../../../../src/integrations/mcp-server/transport';
1+
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
2+
import {
3+
getClient,
4+
getCurrentScope,
5+
getDynamicSamplingContextFromSpan,
6+
setCurrentClient,
7+
startSpan,
8+
} from '../../../../src';
9+
import { wrapTransportOnMessage, wrapTransportSend } from '../../../../src/integrations/mcp-server/transport';
410
import type { Span } from '../../../../src/types/span';
511
import { getActiveSpan, spanToJSON } from '../../../../src/utils/spanUtils';
612
import { getDefaultTestClientOptions, TestClient } from '../../../mocks/client';
@@ -112,3 +118,125 @@ describe('MCP server trace propagation via params._meta', () => {
112118
expect(json.links).toBeUndefined();
113119
});
114120
});
121+
122+
describe('MCP server trace injection into outgoing params._meta', () => {
123+
beforeEach(() => {
124+
const client = new TestClient(
125+
getDefaultTestClientOptions({ dsn: 'https://public@dsn.ingest.sentry.io/1337', tracesSampleRate: 1 }),
126+
);
127+
setCurrentClient(client);
128+
client.init();
129+
});
130+
131+
afterEach(() => {
132+
getCurrentScope().setClient(undefined);
133+
});
134+
135+
function wrappedSend() {
136+
const transport = createMockStdioTransport();
137+
const originalSend = vi.fn().mockResolvedValue(undefined);
138+
transport.send = originalSend;
139+
wrapTransportSend(transport, {});
140+
return { send: transport.send, originalSend };
141+
}
142+
143+
it('injects the active span into outgoing requests', async () => {
144+
const { send, originalSend } = wrappedSend();
145+
const request = { jsonrpc: '2.0', method: 'sampling/createMessage', id: 7, params: { maxTokens: 10 } };
146+
147+
let toolSpan: Span | undefined;
148+
await startSpan({ name: 'tools/call get-weather' }, async span => {
149+
toolSpan = span;
150+
await send(request, { relatedRequestId: 1 });
151+
});
152+
153+
const { traceId, spanId } = toolSpan!.spanContext();
154+
expect(originalSend).toHaveBeenCalledWith(
155+
{
156+
...request,
157+
params: {
158+
maxTokens: 10,
159+
_meta: {
160+
'sentry-trace': `${traceId}-${spanId}-1`,
161+
traceparent: `00-${traceId}-${spanId}-01`,
162+
baggage: expect.stringContaining(`sentry-trace_id=${traceId}`),
163+
},
164+
},
165+
},
166+
{ relatedRequestId: 1 },
167+
);
168+
expect(request.params).toEqual({ maxTokens: 10 });
169+
});
170+
171+
it('parents outgoing notifications to their own notification span', async () => {
172+
const { send, originalSend } = wrappedSend();
173+
let notificationSpan: Span | undefined;
174+
getClient()!.on('spanStart', span => {
175+
if (spanToJSON(span).name === 'notifications/progress') {
176+
notificationSpan = span;
177+
}
178+
});
179+
180+
await startSpan({ name: 'tools/call get-weather' }, () =>
181+
send({ jsonrpc: '2.0', method: 'notifications/progress', params: { progressToken: 1, progress: 50 } }),
182+
);
183+
184+
const { traceId, spanId } = notificationSpan!.spanContext();
185+
expect(originalSend).toHaveBeenCalledWith({
186+
jsonrpc: '2.0',
187+
method: 'notifications/progress',
188+
params: {
189+
progressToken: 1,
190+
progress: 50,
191+
_meta: expect.objectContaining({
192+
'sentry-trace': `${traceId}-${spanId}-1`,
193+
traceparent: `00-${traceId}-${spanId}-01`,
194+
}),
195+
},
196+
});
197+
});
198+
199+
it('keeps existing _meta entries and adds params when there are none', async () => {
200+
const { send, originalSend } = wrappedSend();
201+
202+
await startSpan({ name: 'tools/call get-weather' }, async () => {
203+
await send({ jsonrpc: '2.0', method: 'ping', id: 1 });
204+
await send({ jsonrpc: '2.0', method: 'roots/list', id: 2, params: { _meta: { 'com.example/key': 'value' } } });
205+
});
206+
207+
expect(originalSend).toHaveBeenCalledWith({
208+
jsonrpc: '2.0',
209+
method: 'ping',
210+
id: 1,
211+
params: { _meta: expect.objectContaining({ 'sentry-trace': expect.any(String) }) },
212+
});
213+
expect(originalSend).toHaveBeenCalledWith({
214+
jsonrpc: '2.0',
215+
method: 'roots/list',
216+
id: 2,
217+
params: { _meta: expect.objectContaining({ 'com.example/key': 'value', 'sentry-trace': expect.any(String) }) },
218+
});
219+
});
220+
221+
it.each([
222+
['a traceparent', { traceparent: `00-${TRACE_ID}-${PARENT_SPAN_ID}-01` }],
223+
['a sentry-trace', { 'sentry-trace': `${TRACE_ID}-${PARENT_SPAN_ID}-1` }],
224+
['a non-object _meta', 'not-an-object'],
225+
])('leaves requests that already carry %s unchanged', async (_, meta) => {
226+
const { send, originalSend } = wrappedSend();
227+
const request = { jsonrpc: '2.0', method: 'elicitation/create', id: 3, params: { _meta: meta } };
228+
229+
await startSpan({ name: 'tools/call get-weather' }, () => send(request));
230+
231+
expect(originalSend).toHaveBeenCalledWith(request);
232+
});
233+
234+
it('does not touch responses', async () => {
235+
const { send, originalSend } = wrappedSend();
236+
const response = { jsonrpc: '2.0', id: 1, result: { content: [] } };
237+
238+
await startSpan({ name: 'tools/call get-weather' }, () => send(response));
239+
240+
expect(originalSend).toHaveBeenCalledWith(response);
241+
});
242+
});

0 commit comments

Comments
 (0)