Skip to content

Commit 32fa70a

Browse files
author
Luca Forstner
committed
feat(core): Associate resource/tool/prompt invocations with request span instead of response span
1 parent d1ddb53 commit 32fa70a

1 file changed

Lines changed: 201 additions & 38 deletions

File tree

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

Lines changed: 201 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,17 @@ import {
44
SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
55
SEMANTIC_ATTRIBUTE_SENTRY_SOURCE,
66
} from './semanticAttributes';
7-
import { startSpan } from './tracing';
7+
import { startSpan, withActiveSpan } from './tracing';
8+
import type { Span } from './types-hoist/span';
89
import { logger } from './utils-hoist/logger';
10+
import { getActiveSpan } from './utils/spanUtils';
11+
12+
interface MCPTransport {
13+
// The first argument is a JSON RPC message
14+
onmessage?: (...args: unknown[]) => void;
15+
onclose?: (...args: unknown[]) => void;
16+
sessionId?: string;
17+
}
918

1019
interface MCPServerInstance {
1120
// The first arg is always a name, the last arg should always be a callback function (ie a handler).
@@ -15,6 +24,7 @@ interface MCPServerInstance {
1524
tool: (name: string, ...args: unknown[]) => void;
1625
// The first arg is always a name, the last arg should always be a callback function (ie a handler).
1726
prompt: (name: string, ...args: unknown[]) => void;
27+
connect(transport: MCPTransport): Promise<void>;
1828
}
1929

2030
const wrappedMcpServerInstances = new WeakSet();
@@ -35,6 +45,59 @@ export function wrapMcpServerWithSentry<S extends object>(mcpServerInstance: S):
3545
return mcpServerInstance;
3646
}
3747

48+
// eslint-disable-next-line @typescript-eslint/unbound-method
49+
mcpServerInstance.connect = new Proxy(mcpServerInstance.connect, {
50+
apply(target, thisArg, argArray) {
51+
const [transport, ...restArgs] = argArray as [MCPTransport, ...unknown[]];
52+
53+
if (!transport.onclose) {
54+
transport.onclose = () => {
55+
if (transport.sessionId) {
56+
handleTransportOnClose(transport.sessionId);
57+
}
58+
};
59+
}
60+
61+
if (!transport.onmessage) {
62+
transport.onmessage = jsonRpcMessage => {
63+
if (transport.sessionId && isJsonRPCMessageWithRequestId(jsonRpcMessage)) {
64+
handleTransportOnMessage(transport.sessionId, jsonRpcMessage.id);
65+
}
66+
};
67+
}
68+
69+
const patchedTransport = new Proxy(transport, {
70+
set(target, key, value) {
71+
if (key === 'onmessage') {
72+
target[key] = new Proxy(value, {
73+
apply(onMessageTarget, onMessageThisArg, onMessageArgArray) {
74+
const [jsonRpcMessage] = onMessageArgArray;
75+
if (transport.sessionId && isJsonRPCMessageWithRequestId(jsonRpcMessage)) {
76+
handleTransportOnMessage(transport.sessionId, jsonRpcMessage.id);
77+
}
78+
return Reflect.apply(onMessageTarget, onMessageThisArg, onMessageArgArray);
79+
},
80+
});
81+
} else if (key === 'onclose') {
82+
target[key] = new Proxy(value, {
83+
apply(onCloseTarget, onCloseThisArg, onCloseArgArray) {
84+
if (transport.sessionId) {
85+
handleTransportOnClose(transport.sessionId);
86+
}
87+
return Reflect.apply(onCloseTarget, onCloseThisArg, onCloseArgArray);
88+
},
89+
});
90+
} else {
91+
target[key as keyof MCPTransport] = value;
92+
}
93+
return true;
94+
},
95+
});
96+
97+
return Reflect.apply(target, thisArg, [patchedTransport, ...restArgs]);
98+
},
99+
});
100+
38101
mcpServerInstance.resource = new Proxy(mcpServerInstance.resource, {
39102
apply(target, thisArg, argArray) {
40103
const resourceName: unknown = argArray[0];
@@ -44,19 +107,28 @@ export function wrapMcpServerWithSentry<S extends object>(mcpServerInstance: S):
44107
return target.apply(thisArg, argArray);
45108
}
46109

47-
return startSpan(
48-
{
49-
name: `mcp-server/resource:${resourceName}`,
50-
forceTransaction: true,
51-
attributes: {
52-
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server',
53-
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server',
54-
[SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
55-
'mcp_server.resource': resourceName,
56-
},
110+
const wrappedResourceHandler = new Proxy(resourceHandler, {
111+
apply(resourceHandlerTarget, resourceHandlerThisArg, resourceHandlerArgArray) {
112+
const extraHandlerDataWithRequestId = resourceHandlerArgArray.find(isExtraHandlerDataWithRequestId);
113+
return associateContextWithRequestSpan(extraHandlerDataWithRequestId, () => {
114+
return startSpan(
115+
{
116+
name: `mcp-server/resource:${resourceName}`,
117+
forceTransaction: true,
118+
attributes: {
119+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server',
120+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server',
121+
[SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
122+
'mcp_server.resource': resourceName,
123+
},
124+
},
125+
() => resourceHandlerTarget.apply(resourceHandlerThisArg, resourceHandlerArgArray),
126+
);
127+
});
57128
},
58-
() => target.apply(thisArg, argArray),
59-
);
129+
});
130+
131+
return Reflect.apply(target, thisArg, [...argArray.slice(0, -1), wrappedResourceHandler]);
60132
},
61133
});
62134

@@ -69,19 +141,28 @@ export function wrapMcpServerWithSentry<S extends object>(mcpServerInstance: S):
69141
return target.apply(thisArg, argArray);
70142
}
71143

72-
return startSpan(
73-
{
74-
name: `mcp-server/tool:${toolName}`,
75-
forceTransaction: true,
76-
attributes: {
77-
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server',
78-
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server',
79-
[SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
80-
'mcp_server.tool': toolName,
81-
},
144+
const wrappedToolHandler = new Proxy(toolHandler, {
145+
apply(toolHandlerTarget, toolHandlerThisArg, toolHandlerArgArray) {
146+
const extraHandlerDataWithRequestId = toolHandlerArgArray.find(isExtraHandlerDataWithRequestId);
147+
return associateContextWithRequestSpan(extraHandlerDataWithRequestId, () => {
148+
return startSpan(
149+
{
150+
name: `mcp-server/tool:${toolName}`,
151+
forceTransaction: true,
152+
attributes: {
153+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server',
154+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server',
155+
[SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
156+
'mcp_server.tool': toolName,
157+
},
158+
},
159+
() => toolHandlerTarget.apply(toolHandlerThisArg, toolHandlerArgArray),
160+
);
161+
});
82162
},
83-
() => target.apply(thisArg, argArray),
84-
);
163+
});
164+
165+
return Reflect.apply(target, thisArg, [...argArray.slice(0, -1), wrappedToolHandler]);
85166
},
86167
});
87168

@@ -94,19 +175,28 @@ export function wrapMcpServerWithSentry<S extends object>(mcpServerInstance: S):
94175
return target.apply(thisArg, argArray);
95176
}
96177

97-
return startSpan(
98-
{
99-
name: `mcp-server/resource:${promptName}`,
100-
forceTransaction: true,
101-
attributes: {
102-
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server',
103-
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server',
104-
[SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
105-
'mcp_server.prompt': promptName,
106-
},
178+
const wrappedPromptHandler = new Proxy(promptHandler, {
179+
apply(promptHandlerTarget, promptHandlerThisArg, promptHandlerArgArray) {
180+
const extraHandlerDataWithRequestId = promptHandlerArgArray.find(isExtraHandlerDataWithRequestId);
181+
return associateContextWithRequestSpan(extraHandlerDataWithRequestId, () => {
182+
return startSpan(
183+
{
184+
name: `mcp-server/prompt:${promptName}`,
185+
forceTransaction: true,
186+
attributes: {
187+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server',
188+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server',
189+
[SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
190+
'mcp_server.prompt': promptName,
191+
},
192+
},
193+
() => promptHandlerTarget.apply(promptHandlerThisArg, promptHandlerArgArray),
194+
);
195+
});
107196
},
108-
() => target.apply(thisArg, argArray),
109-
);
197+
});
198+
199+
return Reflect.apply(target, thisArg, [...argArray.slice(0, -1), wrappedPromptHandler]);
110200
},
111201
});
112202

@@ -124,6 +214,79 @@ function isMcpServerInstance(mcpServerInstance: unknown): mcpServerInstance is M
124214
'tool' in mcpServerInstance &&
125215
typeof mcpServerInstance.tool === 'function' &&
126216
'prompt' in mcpServerInstance &&
127-
typeof mcpServerInstance.prompt === 'function'
217+
typeof mcpServerInstance.prompt === 'function' &&
218+
'connect' in mcpServerInstance &&
219+
typeof mcpServerInstance.connect === 'function'
220+
);
221+
}
222+
223+
function isJsonRPCMessageWithRequestId(target: unknown): target is { id: RequestId } {
224+
return (
225+
typeof target === 'object' &&
226+
target !== null &&
227+
'id' in target &&
228+
(typeof target.id === 'number' || typeof target.id === 'string')
128229
);
129230
}
231+
232+
interface ExtraHandlerDataWithRequestId {
233+
sessionId: SessionId;
234+
requestId: RequestId;
235+
}
236+
237+
// Note that not all versions of the MCP library have `requestId` as a field on the extra data.
238+
function isExtraHandlerDataWithRequestId(target: unknown): target is ExtraHandlerDataWithRequestId {
239+
return (
240+
typeof target === 'object' &&
241+
target !== null &&
242+
'sessionId' in target &&
243+
typeof target.sessionId === 'string' &&
244+
'requestId' in target &&
245+
(typeof target.requestId === 'number' || typeof target.requestId === 'string')
246+
);
247+
}
248+
249+
type SessionId = string;
250+
type RequestId = string | number;
251+
252+
const sessionAndRequestToRequestParentSpanMap = new Map<SessionId, Map<RequestId, Span>>();
253+
254+
function handleTransportOnClose(sessionId: SessionId): void {
255+
sessionAndRequestToRequestParentSpanMap.delete(sessionId);
256+
}
257+
258+
function handleTransportOnMessage(sessionId: SessionId, requestId: RequestId): void {
259+
const activeSpan = getActiveSpan();
260+
if (activeSpan) {
261+
const requestIdToSpanMap = sessionAndRequestToRequestParentSpanMap.get(sessionId) ?? new Map();
262+
requestIdToSpanMap.set(requestId, activeSpan);
263+
sessionAndRequestToRequestParentSpanMap.set(sessionId, requestIdToSpanMap);
264+
}
265+
}
266+
267+
function associateContextWithRequestSpan<T>(
268+
extraHandlerData: ExtraHandlerDataWithRequestId | undefined,
269+
cb: () => T,
270+
): T {
271+
if (extraHandlerData) {
272+
const { sessionId, requestId } = extraHandlerData;
273+
const requestIdSpanMap = sessionAndRequestToRequestParentSpanMap.get(sessionId);
274+
275+
if (!requestIdSpanMap) {
276+
return cb();
277+
}
278+
279+
const span = requestIdSpanMap.get(requestId);
280+
if (!span) {
281+
return cb();
282+
}
283+
284+
// remove the span from the map so it can be garbage collected
285+
requestIdSpanMap.delete(requestId);
286+
return withActiveSpan(span, () => {
287+
return cb();
288+
});
289+
}
290+
291+
return cb();
292+
}

0 commit comments

Comments
 (0)