Skip to content

Commit 62ca0f3

Browse files
betegonbetegon
committed
refactor(mcp-server): refactor span handling and utility functions for MCP server instrumentation
1 parent edc4e3c commit 62ca0f3

File tree

5 files changed

+625
-308
lines changed

5 files changed

+625
-308
lines changed

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

Lines changed: 60 additions & 244 deletions
Original file line numberDiff line numberDiff line change
@@ -1,48 +1,17 @@
1-
import {
2-
SEMANTIC_ATTRIBUTE_SENTRY_OP,
3-
SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
4-
SEMANTIC_ATTRIBUTE_SENTRY_SOURCE,
5-
} from './semanticAttributes';
6-
import { startSpan, withActiveSpan } from './tracing';
71
import type { Span } from './types-hoist/span';
8-
import { getActiveSpan } from './utils/spanUtils';
9-
import {
10-
MCP_METHOD_NAME_ATTRIBUTE,
11-
MCP_REQUEST_ID_ATTRIBUTE,
12-
MCP_SESSION_ID_ATTRIBUTE,
13-
MCP_TRANSPORT_ATTRIBUTE,
14-
NETWORK_TRANSPORT_ATTRIBUTE,
15-
NETWORK_PROTOCOL_VERSION_ATTRIBUTE,
16-
CLIENT_ADDRESS_ATTRIBUTE,
17-
CLIENT_PORT_ATTRIBUTE,
18-
MCP_NOTIFICATION_DIRECTION_ATTRIBUTE,
19-
MCP_SERVER_OP_VALUE,
20-
MCP_FUNCTION_ORIGIN_VALUE,
21-
MCP_NOTIFICATION_ORIGIN_VALUE,
22-
MCP_ROUTE_SOURCE_VALUE,
23-
} from './utils/mcp-server/attributes';
242
import type {
25-
JsonRpcRequest,
26-
JsonRpcNotification,
27-
MCPTransport,
3+
ExtraHandlerData,
284
MCPServerInstance,
5+
MCPTransport,
296
SessionId,
30-
RequestId,
31-
ExtraHandlerData,
327
} from './utils/mcp-server/types';
338
import {
34-
isJsonRpcRequest,
9+
createMcpNotificationSpan,
10+
createMcpOutgoingNotificationSpan,
11+
createMcpServerSpan,
3512
isJsonRpcNotification,
36-
extractTarget,
37-
getTargetAttributes,
38-
getRequestArguments,
39-
getTransportTypes,
40-
getNotificationDescription,
41-
getNotificationAttributes,
42-
extractClientAddress,
43-
extractClientPort,
13+
isJsonRpcRequest,
4414
validateMcpServerInstance,
45-
createSpanName,
4615
} from './utils/mcp-server/utils';
4716

4817
const wrappedMcpServerInstances = new WeakSet();
@@ -65,55 +34,56 @@ export function wrapMcpServerWithSentry<S extends object>(mcpServerInstance: S):
6534
const serverInstance = mcpServerInstance as MCPServerInstance;
6635

6736
// Wrap connect() to intercept AFTER Protocol sets up transport handlers
68-
serverInstance.connect = new Proxy(serverInstance.connect, {
37+
const originalConnect = serverInstance.connect;
38+
serverInstance.connect = new Proxy(originalConnect, {
6939
async apply(target, thisArg, argArray) {
7040
const [transport, ...restArgs] = argArray as [MCPTransport, ...unknown[]];
7141

7242
// Call the original connect first to let Protocol set up its handlers
7343
const result = await Reflect.apply(target, thisArg, [transport, ...restArgs]);
74-
44+
7545
// Intercept incoming messages via onmessage
7646
if (transport.onmessage) {
7747
const protocolOnMessage = transport.onmessage;
78-
48+
7949
transport.onmessage = new Proxy(protocolOnMessage, {
8050
apply(onMessageTarget, onMessageThisArg, onMessageArgs) {
8151
const [jsonRpcMessage, extra] = onMessageArgs;
82-
52+
8353
// Instrument both requests and notifications
8454
if (isJsonRpcRequest(jsonRpcMessage)) {
8555
return createMcpServerSpan(jsonRpcMessage, transport, extra as ExtraHandlerData, () => {
86-
return onMessageTarget.apply(onMessageThisArg, onMessageArgs);
56+
return Reflect.apply(onMessageTarget, onMessageThisArg, onMessageArgs);
8757
});
8858
}
8959
if (isJsonRpcNotification(jsonRpcMessage)) {
9060
return createMcpNotificationSpan(jsonRpcMessage, transport, extra as ExtraHandlerData, () => {
91-
return onMessageTarget.apply(onMessageThisArg, onMessageArgs);
61+
return Reflect.apply(onMessageTarget, onMessageThisArg, onMessageArgs);
9262
});
9363
}
94-
95-
return onMessageTarget.apply(onMessageThisArg, onMessageArgs);
96-
}
64+
65+
return Reflect.apply(onMessageTarget, onMessageThisArg, onMessageArgs);
66+
},
9767
});
9868
}
9969

10070
// Intercept outgoing messages via send
10171
if (transport.send) {
10272
const originalSend = transport.send;
103-
73+
10474
transport.send = new Proxy(originalSend, {
10575
async apply(sendTarget, sendThisArg, sendArgs) {
10676
const [message, options] = sendArgs;
107-
77+
10878
// Instrument outgoing notifications (but not requests/responses)
10979
if (isJsonRpcNotification(message)) {
11080
return createMcpOutgoingNotificationSpan(message, transport, options as Record<string, unknown>, () => {
111-
return sendTarget.apply(sendThisArg, sendArgs);
81+
return Reflect.apply(sendTarget, sendThisArg, sendArgs);
11282
});
11383
}
114-
115-
return sendTarget.apply(sendThisArg, sendArgs);
116-
}
84+
85+
return Reflect.apply(sendTarget, sendThisArg, sendArgs);
86+
},
11787
});
11888
}
11989

@@ -125,8 +95,8 @@ export function wrapMcpServerWithSentry<S extends object>(mcpServerInstance: S):
12595
if (transport.sessionId) {
12696
handleTransportOnClose(transport.sessionId);
12797
}
128-
return onCloseTarget.apply(onCloseThisArg, onCloseArgs);
129-
}
98+
return Reflect.apply(onCloseTarget, onCloseThisArg, onCloseArgs);
99+
},
130100
});
131101
}
132102
return result;
@@ -137,203 +107,49 @@ export function wrapMcpServerWithSentry<S extends object>(mcpServerInstance: S):
137107
return mcpServerInstance as S;
138108
}
139109

140-
function createMcpServerSpan(
141-
jsonRpcMessage: JsonRpcRequest,
142-
transport: MCPTransport,
143-
extra: ExtraHandlerData,
144-
callback: () => unknown
145-
) {
146-
const { method, id: requestId, params } = jsonRpcMessage;
147-
148-
// Extract target from method and params for proper description
149-
const target = extractTarget(method, params as Record<string, unknown>);
150-
const description = createSpanName(method, target);
151-
152-
// Session ID should come from the transport itself, not the RPC message
153-
const sessionId = transport.sessionId;
154-
155-
// Extract client information from extra/request data
156-
const clientAddress = extractClientAddress(extra);
157-
const clientPort = extractClientPort(extra);
158-
159-
// Determine transport types
160-
const { mcpTransport, networkTransport } = getTransportTypes(transport);
161-
162-
const attributes: Record<string, string | number> = {
163-
[MCP_METHOD_NAME_ATTRIBUTE]: method,
164-
165-
...(requestId !== undefined && { [MCP_REQUEST_ID_ATTRIBUTE]: String(requestId) }),
166-
...(target && getTargetAttributes(method, target)),
167-
...(sessionId && { [MCP_SESSION_ID_ATTRIBUTE]: sessionId }),
168-
...(clientAddress && { [CLIENT_ADDRESS_ATTRIBUTE]: clientAddress }),
169-
...(clientPort && { [CLIENT_PORT_ATTRIBUTE]: clientPort }),
170-
[MCP_TRANSPORT_ATTRIBUTE]: mcpTransport, // Application level: "http", "sse", "stdio", "websocket"
171-
[NETWORK_TRANSPORT_ATTRIBUTE]: networkTransport, // Network level: "tcp", "pipe", "udp", "quic"
172-
[NETWORK_PROTOCOL_VERSION_ATTRIBUTE]: '2.0', // JSON-RPC version
173-
174-
// Opt-in: Tool arguments (if enabled)
175-
...getRequestArguments(method, params as Record<string, unknown>),
176-
177-
// Sentry-specific attributes
178-
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: MCP_SERVER_OP_VALUE,
179-
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: MCP_FUNCTION_ORIGIN_VALUE,
180-
[SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: MCP_ROUTE_SOURCE_VALUE
181-
};
182-
183-
return startSpan({
184-
name: description,
185-
forceTransaction: true,
186-
attributes
187-
}, () => {
188-
// TODO(bete): add proper error handling. Handle JSON RPC errors in the result
189-
return callback();
190-
});
191-
}
192-
193-
function createMcpNotificationSpan(
194-
jsonRpcMessage: JsonRpcNotification,
195-
transport: MCPTransport,
196-
extra: ExtraHandlerData,
197-
callback: () => unknown
198-
) {
199-
const { method, params } = jsonRpcMessage;
200-
201-
const description = getNotificationDescription(method, params as Record<string, unknown>);
202-
203-
const sessionId = transport.sessionId;
204-
205-
// Extract client information
206-
const clientAddress = extractClientAddress(extra);
207-
const clientPort = extractClientPort(extra);
208-
209-
// Determine transport types
210-
const { mcpTransport, networkTransport } = getTransportTypes(transport);
211-
212-
const notificationAttribs = getNotificationAttributes(method, params as Record<string, unknown>);
213-
214-
const attributes: Record<string, string | number> = {
215-
[MCP_METHOD_NAME_ATTRIBUTE]: method,
216-
[MCP_NOTIFICATION_DIRECTION_ATTRIBUTE]: 'client_to_server', // Incoming notification
217-
218-
...(sessionId && { [MCP_SESSION_ID_ATTRIBUTE]: sessionId }),
219-
...(clientAddress && { [CLIENT_ADDRESS_ATTRIBUTE]: clientAddress }),
220-
...(clientPort && { [CLIENT_PORT_ATTRIBUTE]: clientPort }),
221-
[MCP_TRANSPORT_ATTRIBUTE]: mcpTransport,
222-
[NETWORK_TRANSPORT_ATTRIBUTE]: networkTransport,
223-
[NETWORK_PROTOCOL_VERSION_ATTRIBUTE]: '2.0',
224-
225-
// Notification-specific attributes
226-
...notificationAttribs,
227-
228-
// Sentry-specific attributes
229-
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: MCP_SERVER_OP_VALUE,
230-
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: MCP_NOTIFICATION_ORIGIN_VALUE,
231-
[SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: MCP_ROUTE_SOURCE_VALUE
232-
};
233-
234-
return startSpan({
235-
name: description,
236-
forceTransaction: true,
237-
attributes
238-
}, () => {
239-
const result = callback();
240-
return result;
241-
});
242-
}
243-
244-
function createMcpOutgoingNotificationSpan(
245-
jsonRpcMessage: JsonRpcNotification,
246-
transport: MCPTransport,
247-
options: Record<string, unknown>,
248-
callback: () => unknown
249-
) {
250-
const { method, params } = jsonRpcMessage;
251-
252-
const description = getNotificationDescription(method, params as Record<string, unknown>);
253-
254-
const sessionId = transport.sessionId;
255-
256-
// Determine transport types
257-
const { mcpTransport, networkTransport } = getTransportTypes(transport);
258-
259-
const notificationAttribs = getNotificationAttributes(method, params as Record<string, unknown>);
260-
261-
const attributes: Record<string, string | number> = {
262-
[MCP_METHOD_NAME_ATTRIBUTE]: method,
263-
[MCP_NOTIFICATION_DIRECTION_ATTRIBUTE]: 'server_to_client', // Outgoing notification
264-
265-
...(sessionId && { [MCP_SESSION_ID_ATTRIBUTE]: sessionId }),
266-
[MCP_TRANSPORT_ATTRIBUTE]: mcpTransport,
267-
[NETWORK_TRANSPORT_ATTRIBUTE]: networkTransport,
268-
[NETWORK_PROTOCOL_VERSION_ATTRIBUTE]: '2.0',
269-
270-
// Notification-specific attributes
271-
...notificationAttribs,
272-
273-
// Sentry-specific attributes
274-
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: MCP_SERVER_OP_VALUE,
275-
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: MCP_NOTIFICATION_ORIGIN_VALUE,
276-
[SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: MCP_ROUTE_SOURCE_VALUE
277-
};
278-
279-
return startSpan({
280-
name: description,
281-
forceTransaction: true,
282-
attributes
283-
}, () => {
284-
const result = callback();
285-
return result;
286-
});
287-
}
288-
289110
// =============================================================================
290111
// SESSION AND REQUEST CORRELATION (Legacy support)
291112
// =============================================================================
292113

293-
interface ExtraHandlerDataWithRequestId {
294-
sessionId: SessionId;
295-
requestId: RequestId;
296-
}
297-
298-
const sessionAndRequestToRequestParentSpanMap = new Map<SessionId, Map<RequestId, Span>>();
114+
const sessionAndRequestToRequestParentSpanMap = new Map<SessionId, Map<string, Span>>();
299115

300116
function handleTransportOnClose(sessionId: SessionId): void {
301117
sessionAndRequestToRequestParentSpanMap.delete(sessionId);
302118
}
303119

304120
// TODO(bete): refactor this and associateContextWithRequestSpan to use the new span API.
305-
function handleTransportOnMessage(sessionId: SessionId, requestId: RequestId): void {
306-
const activeSpan = getActiveSpan();
307-
if (activeSpan) {
308-
const requestIdToSpanMap = sessionAndRequestToRequestParentSpanMap.get(sessionId) ?? new Map();
309-
requestIdToSpanMap.set(requestId, activeSpan);
310-
sessionAndRequestToRequestParentSpanMap.set(sessionId, requestIdToSpanMap);
311-
}
312-
}
313-
314-
function associateContextWithRequestSpan<T>(
315-
extraHandlerData: ExtraHandlerDataWithRequestId | undefined,
316-
cb: () => T,
317-
): T {
318-
if (extraHandlerData) {
319-
const { sessionId, requestId } = extraHandlerData;
320-
const requestIdSpanMap = sessionAndRequestToRequestParentSpanMap.get(sessionId);
321-
322-
if (!requestIdSpanMap) {
323-
return cb();
324-
}
325-
326-
const span = requestIdSpanMap.get(requestId);
327-
if (!span) {
328-
return cb();
329-
}
330-
331-
// remove the span from the map so it can be garbage collected
332-
requestIdSpanMap.delete(requestId);
333-
return withActiveSpan(span, () => {
334-
return cb();
335-
});
336-
}
337-
338-
return cb();
339-
}
121+
// function handleTransportOnMessage(sessionId: SessionId, requestId: string): void {
122+
// const activeSpan = getActiveSpan();
123+
// if (activeSpan) {
124+
// const requestIdToSpanMap = sessionAndRequestToRequestParentSpanMap.get(sessionId) ?? new Map();
125+
// requestIdToSpanMap.set(requestId, activeSpan);
126+
// sessionAndRequestToRequestParentSpanMap.set(sessionId, requestIdToSpanMap);
127+
// }
128+
// }
129+
130+
// function associateContextWithRequestSpan<T>(
131+
// extraHandlerData: { sessionId: SessionId; requestId: string } | undefined,
132+
// cb: () => T,
133+
// ): T {
134+
// if (extraHandlerData) {
135+
// const { sessionId, requestId } = extraHandlerData;
136+
// const requestIdSpanMap = sessionAndRequestToRequestParentSpanMap.get(sessionId);
137+
138+
// if (!requestIdSpanMap) {
139+
// return cb();
140+
// }
141+
142+
// const span = requestIdSpanMap.get(requestId);
143+
// if (!span) {
144+
// return cb();
145+
// }
146+
147+
// // remove the span from the map so it can be garbage collected
148+
// requestIdSpanMap.delete(requestId);
149+
// return withActiveSpan(span, () => {
150+
// return cb();
151+
// });
152+
// }
153+
154+
// return cb();
155+
// }

‎packages/core/src/utils/mcp-server/attributes.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
/**
22
* Essential MCP attribute constants for Sentry instrumentation
3-
*
3+
*
44
* Based on OpenTelemetry MCP semantic conventions
55
* @see https://github.com/open-telemetry/semantic-conventions/blob/3097fb0af5b9492b0e3f55dc5f6c21a3dc2be8df/docs/gen-ai/mcp.md
66
*/
@@ -115,4 +115,4 @@ export const MCP_NOTIFICATION_ORIGIN_VALUE = 'auto.mcp.notification';
115115
/**
116116
* Sentry source value for MCP route spans
117117
*/
118-
export const MCP_ROUTE_SOURCE_VALUE = 'route';
118+
export const MCP_ROUTE_SOURCE_VALUE = 'route';

0 commit comments

Comments
 (0)