Skip to content

Commit e0e9a2b

Browse files
lforstlforst
committed
feat(core): Add wrapMcpServerWithSentry to instrument MCP servers from @modelcontextprotocol/sdk 
1 parent 8046e14 commit e0e9a2b

File tree

3 files changed

+372
-0
lines changed

3 files changed

+372
-0
lines changed

Diff for: ‎packages/core/src/index.ts

+1
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,7 @@ export { thirdPartyErrorFilterIntegration } from './integrations/third-party-err
111111
export { profiler } from './profiling';
112112
export { instrumentFetchRequest } from './fetch';
113113
export { trpcMiddleware } from './trpc';
114+
export { wrapMcpServerWithSentry } from './mcp-server';
114115
export { captureFeedback } from './feedback';
115116
export type { ReportDialogOptions } from './report-dialog';
116117
export { _INTERNAL_captureLog, _INTERNAL_flushLogsBuffer } from './logs/exports';

Diff for: ‎packages/core/src/mcp-server.ts

+129
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
import { DEBUG_BUILD } from './debug-build';
2+
import {
3+
SEMANTIC_ATTRIBUTE_SENTRY_OP,
4+
SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
5+
SEMANTIC_ATTRIBUTE_SENTRY_SOURCE,
6+
} from './semanticAttributes';
7+
import { startSpan } from './tracing';
8+
import { logger } from './utils-hoist';
9+
10+
interface MCPServerInstance {
11+
// The first arg is always a name, the last arg should always be a callback function (ie a handler).
12+
// TODO: We could also make use of the resource uri argument somehow.
13+
resource: (name: string, ...args: unknown[]) => void;
14+
// The first arg is always a name, the last arg should always be a callback function (ie a handler).
15+
tool: (name: string, ...args: unknown[]) => void;
16+
// The first arg is always a name, the last arg should always be a callback function (ie a handler).
17+
prompt: (name: string, ...args: unknown[]) => void;
18+
}
19+
20+
const wrappedMcpServerInstances = new WeakSet();
21+
22+
/**
23+
* Wraps a MCP Server instance from the `@modelcontextprotocol/sdk` package with Sentry instrumentation.
24+
*
25+
* Compatible with versions `^1.9.0` of the `@modelcontextprotocol/sdk` package.
26+
*/
27+
// We are exposing this API for non-node runtimes that cannot rely on auto-instrumentation.
28+
export function wrapMcpServerWithSentry<S extends object>(mcpServerInstance: S): S {
29+
if (wrappedMcpServerInstances.has(mcpServerInstance)) {
30+
return mcpServerInstance;
31+
}
32+
33+
if (!isMcpServerInstance(mcpServerInstance)) {
34+
DEBUG_BUILD && logger.warn('Did not patch MCP server. Interface is incompatible.');
35+
return mcpServerInstance;
36+
}
37+
38+
mcpServerInstance.resource = new Proxy(mcpServerInstance.resource, {
39+
apply(target, thisArg, argArray) {
40+
const resourceName: unknown = argArray[0];
41+
const resourceHandler: unknown = argArray[argArray.length - 1];
42+
43+
if (typeof resourceName !== 'string' || typeof resourceHandler !== 'function') {
44+
return target.apply(thisArg, argArray);
45+
}
46+
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+
},
57+
},
58+
() => target.apply(thisArg, argArray),
59+
);
60+
},
61+
});
62+
63+
mcpServerInstance.tool = new Proxy(mcpServerInstance.tool, {
64+
apply(target, thisArg, argArray) {
65+
const toolName: unknown = argArray[0];
66+
const toolHandler: unknown = argArray[argArray.length - 1];
67+
68+
if (typeof toolName !== 'string' || typeof toolHandler !== 'function') {
69+
return target.apply(thisArg, argArray);
70+
}
71+
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+
},
82+
},
83+
() => target.apply(thisArg, argArray),
84+
);
85+
},
86+
});
87+
88+
mcpServerInstance.prompt = new Proxy(mcpServerInstance.prompt, {
89+
apply(target, thisArg, argArray) {
90+
const promptName: unknown = argArray[0];
91+
const promptHandler: unknown = argArray[argArray.length - 1];
92+
93+
if (typeof promptName !== 'string' || typeof promptHandler !== 'function') {
94+
return target.apply(thisArg, argArray);
95+
}
96+
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+
},
107+
},
108+
() => target.apply(thisArg, argArray),
109+
);
110+
},
111+
});
112+
113+
wrappedMcpServerInstances.add(mcpServerInstance);
114+
115+
return mcpServerInstance as S;
116+
}
117+
118+
function isMcpServerInstance(mcpServerInstance: unknown): mcpServerInstance is MCPServerInstance {
119+
return (
120+
typeof mcpServerInstance === 'object' &&
121+
mcpServerInstance !== null &&
122+
'resource' in mcpServerInstance &&
123+
typeof mcpServerInstance.resource === 'function' &&
124+
'tool' in mcpServerInstance &&
125+
typeof mcpServerInstance.tool === 'function' &&
126+
'prompt' in mcpServerInstance &&
127+
typeof mcpServerInstance.prompt === 'function'
128+
);
129+
}

Diff for: ‎packages/core/test/lib/mcp-server.test.ts

+242
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,242 @@
1+
import { beforeEach, describe, expect, it, vi } from 'vitest';
2+
import { wrapMcpServerWithSentry } from '../../src/mcp-server';
3+
import {
4+
SEMANTIC_ATTRIBUTE_SENTRY_OP,
5+
SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
6+
SEMANTIC_ATTRIBUTE_SENTRY_SOURCE,
7+
} from '../../src/semanticAttributes';
8+
import * as tracingModule from '../../src/tracing';
9+
10+
vi.mock('../../src/tracing');
11+
12+
describe('wrapMcpServerWithSentry', () => {
13+
beforeEach(() => {
14+
vi.clearAllMocks();
15+
// @ts-expect-error mocking span is annoying
16+
vi.mocked(tracingModule.startSpan).mockImplementation((_, cb) => cb());
17+
});
18+
19+
it('should wrap valid MCP server instance methods with Sentry spans', () => {
20+
// Create a mock MCP server instance
21+
const mockResource = vi.fn();
22+
const mockTool = vi.fn();
23+
const mockPrompt = vi.fn();
24+
25+
const mockMcpServer = {
26+
resource: mockResource,
27+
tool: mockTool,
28+
prompt: mockPrompt,
29+
};
30+
31+
// Wrap the MCP server
32+
const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
33+
34+
// Verify it returns the same instance (modified)
35+
expect(wrappedMcpServer).toBe(mockMcpServer);
36+
37+
// Original methods should be wrapped
38+
expect(wrappedMcpServer.resource).not.toBe(mockResource);
39+
expect(wrappedMcpServer.tool).not.toBe(mockTool);
40+
expect(wrappedMcpServer.prompt).not.toBe(mockPrompt);
41+
});
42+
43+
it('should return the input unchanged if it is not a valid MCP server instance', () => {
44+
const invalidMcpServer = {
45+
// Missing required methods
46+
resource: () => {},
47+
tool: () => {},
48+
// No prompt method
49+
};
50+
51+
const result = wrapMcpServerWithSentry(invalidMcpServer);
52+
expect(result).toBe(invalidMcpServer);
53+
54+
// Methods should not be wrapped
55+
expect(result.resource).toBe(invalidMcpServer.resource);
56+
expect(result.tool).toBe(invalidMcpServer.tool);
57+
58+
// No calls to startSpan
59+
expect(tracingModule.startSpan).not.toHaveBeenCalled();
60+
});
61+
62+
it('should not wrap the same instance twice', () => {
63+
const mockMcpServer = {
64+
resource: vi.fn(),
65+
tool: vi.fn(),
66+
prompt: vi.fn(),
67+
};
68+
69+
// First wrap
70+
const wrappedOnce = wrapMcpServerWithSentry(mockMcpServer);
71+
72+
// Store references to wrapped methods
73+
const wrappedResource = wrappedOnce.resource;
74+
const wrappedTool = wrappedOnce.tool;
75+
const wrappedPrompt = wrappedOnce.prompt;
76+
77+
// Second wrap
78+
const wrappedTwice = wrapMcpServerWithSentry(wrappedOnce);
79+
80+
// Should be the same instance with the same wrapped methods
81+
expect(wrappedTwice).toBe(wrappedOnce);
82+
expect(wrappedTwice.resource).toBe(wrappedResource);
83+
expect(wrappedTwice.tool).toBe(wrappedTool);
84+
expect(wrappedTwice.prompt).toBe(wrappedPrompt);
85+
});
86+
87+
describe('resource method wrapping', () => {
88+
it('should create a span with proper attributes when resource is called', () => {
89+
const mockResourceHandler = vi.fn();
90+
const resourceName = 'test-resource';
91+
92+
const mockMcpServer = {
93+
resource: vi.fn(),
94+
tool: vi.fn(),
95+
prompt: vi.fn(),
96+
};
97+
98+
const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
99+
wrappedMcpServer.resource(resourceName, {}, mockResourceHandler);
100+
101+
expect(tracingModule.startSpan).toHaveBeenCalledTimes(1);
102+
expect(tracingModule.startSpan).toHaveBeenCalledWith(
103+
{
104+
name: `mcp-server/resource:${resourceName}`,
105+
forceTransaction: true,
106+
attributes: {
107+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server',
108+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server',
109+
[SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
110+
'mcp_server.resource': resourceName,
111+
},
112+
},
113+
expect.any(Function),
114+
);
115+
116+
// Verify the original method was called with all arguments
117+
expect(mockMcpServer.resource).toHaveBeenCalledWith(resourceName, {}, mockResourceHandler);
118+
});
119+
120+
it('should call the original resource method directly if name or handler is not valid', () => {
121+
const mockMcpServer = {
122+
resource: vi.fn(),
123+
tool: vi.fn(),
124+
prompt: vi.fn(),
125+
};
126+
127+
const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
128+
129+
// Call without string name
130+
wrappedMcpServer.resource({} as any, 'handler');
131+
132+
// Call without function handler
133+
wrappedMcpServer.resource('name', 'not-a-function');
134+
135+
// Original method should be called directly without creating spans
136+
expect(mockMcpServer.resource).toHaveBeenCalledTimes(2);
137+
expect(tracingModule.startSpan).not.toHaveBeenCalled();
138+
});
139+
});
140+
141+
describe('tool method wrapping', () => {
142+
it('should create a span with proper attributes when tool is called', () => {
143+
const mockToolHandler = vi.fn();
144+
const toolName = 'test-tool';
145+
146+
const mockMcpServer = {
147+
resource: vi.fn(),
148+
tool: vi.fn(),
149+
prompt: vi.fn(),
150+
};
151+
152+
const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
153+
wrappedMcpServer.tool(toolName, {}, mockToolHandler);
154+
155+
expect(tracingModule.startSpan).toHaveBeenCalledTimes(1);
156+
expect(tracingModule.startSpan).toHaveBeenCalledWith(
157+
{
158+
name: `mcp-server/tool:${toolName}`,
159+
forceTransaction: true,
160+
attributes: {
161+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server',
162+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server',
163+
[SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
164+
'mcp_server.tool': toolName,
165+
},
166+
},
167+
expect.any(Function),
168+
);
169+
170+
// Verify the original method was called with all arguments
171+
expect(mockMcpServer.tool).toHaveBeenCalledWith(toolName, {}, mockToolHandler);
172+
});
173+
174+
it('should call the original tool method directly if name or handler is not valid', () => {
175+
const mockMcpServer = {
176+
resource: vi.fn(),
177+
tool: vi.fn(),
178+
prompt: vi.fn(),
179+
};
180+
181+
const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
182+
183+
// Call without string name
184+
wrappedMcpServer.tool({} as any, 'handler');
185+
186+
// Original method should be called directly without creating spans
187+
expect(mockMcpServer.tool).toHaveBeenCalledTimes(1);
188+
expect(tracingModule.startSpan).not.toHaveBeenCalled();
189+
});
190+
});
191+
192+
describe('prompt method wrapping', () => {
193+
it('should create a span with proper attributes when prompt is called', () => {
194+
const mockPromptHandler = vi.fn();
195+
const promptName = 'test-prompt';
196+
197+
const mockMcpServer = {
198+
resource: vi.fn(),
199+
tool: vi.fn(),
200+
prompt: vi.fn(),
201+
};
202+
203+
const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
204+
wrappedMcpServer.prompt(promptName, {}, mockPromptHandler);
205+
206+
expect(tracingModule.startSpan).toHaveBeenCalledTimes(1);
207+
expect(tracingModule.startSpan).toHaveBeenCalledWith(
208+
{
209+
name: `mcp-server/resource:${promptName}`,
210+
forceTransaction: true,
211+
attributes: {
212+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server',
213+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server',
214+
[SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
215+
'mcp_server.prompt': promptName,
216+
},
217+
},
218+
expect.any(Function),
219+
);
220+
221+
// Verify the original method was called with all arguments
222+
expect(mockMcpServer.prompt).toHaveBeenCalledWith(promptName, {}, mockPromptHandler);
223+
});
224+
225+
it('should call the original prompt method directly if name or handler is not valid', () => {
226+
const mockMcpServer = {
227+
resource: vi.fn(),
228+
tool: vi.fn(),
229+
prompt: vi.fn(),
230+
};
231+
232+
const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
233+
234+
// Call without function handler
235+
wrappedMcpServer.prompt('name', 'not-a-function');
236+
237+
// Original method should be called directly without creating spans
238+
expect(mockMcpServer.prompt).toHaveBeenCalledTimes(1);
239+
expect(tracingModule.startSpan).not.toHaveBeenCalled();
240+
});
241+
});
242+
});

0 commit comments

Comments
 (0)