Skip to content

Commit 8b9fd2c

Browse files
betegonclaude
andcommitted
fix(mcp): retroactively wrap handlers registered before wrapMcpServerWithSentry
wrapMcpServerWithSentry patches the registration methods (registerTool etc.) so only handlers registered after wrapping get instrumented. When wrapping happens after registration, the already-stored executors/callbacks are never touched, silently skipping error capture for those handlers. Fix by scanning the McpServer's internal registries (_registeredTools, _registeredResources, _registeredResourceTemplates, _registeredPrompts) after patching the registration methods, and wrapping the stored callables in-place: - tools: wrap `executor` (what the SDK calls on tools/call) - resources/templates: wrap `readCallback` - prompts: wrap `handler` Both call orderings are now fully supported. JSDoc updated to reflect this and reference the upstream SDK source line where executor is invoked. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
1 parent cea58ff commit 8b9fd2c

4 files changed

Lines changed: 157 additions & 3 deletions

File tree

packages/core/src/integrations/mcp-server/handlers.ts

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -167,3 +167,74 @@ export function wrapAllMCPHandlers(serverInstance: MCPServerInstance): void {
167167
wrapResourceHandlers(serverInstance);
168168
wrapPromptHandlers(serverInstance);
169169
}
170+
171+
/**
172+
* Retroactively wraps handlers on tools, resources, and prompts that were registered
173+
* before `wrapMcpServerWithSentry` was called.
174+
*
175+
* The MCP SDK stores registered entries in private maps and invokes them via the entry's
176+
* own property at call time — `executor` for tools, `readCallback` for resources, and
177+
* `handler` for prompts. Replacing those properties
178+
* in-place is therefore equivalent to having wrapped the original registration call.
179+
*
180+
* NOTE: This intentionally accesses private MCP SDK internals (`_registeredTools` etc.).
181+
* The properties and their shapes are verified against @modelcontextprotocol/sdk source:
182+
* https://github.com/modelcontextprotocol/typescript-sdk/blob/2c0c481cb9dbfd15c8613f765c940a5f5bace94d/packages/server/src/server/mcp.ts#L304
183+
* When upgrading the MCP SDK, re-verify that these internal maps and their callable
184+
* properties still exist and are invoked directly (not captured by closure at registration).
185+
* All access is defensive — if a property is absent or not a function we skip silently.
186+
* @internal
187+
*/
188+
export function wrapExistingHandlers(serverInstance: MCPServerInstance): void {
189+
const server = serverInstance as unknown as Record<string, unknown>;
190+
191+
// Tools: MCP SDK calls registeredTool.executor (generated from handler at registration time)
192+
const registeredTools = server['_registeredTools'];
193+
if (registeredTools && typeof registeredTools === 'object') {
194+
for (const [name, tool] of Object.entries(registeredTools as Record<string, Record<string, unknown>>)) {
195+
if (typeof tool['executor'] === 'function') {
196+
tool['executor'] = createWrappedHandler(tool['executor'] as MCPHandler, 'registerTool', name);
197+
}
198+
}
199+
}
200+
201+
// Resources: MCP SDK calls registeredResource.readCallback
202+
const registeredResources = server['_registeredResources'];
203+
if (registeredResources && typeof registeredResources === 'object') {
204+
for (const [name, resource] of Object.entries(registeredResources as Record<string, Record<string, unknown>>)) {
205+
if (typeof resource['readCallback'] === 'function') {
206+
resource['readCallback'] = createWrappedHandler(
207+
resource['readCallback'] as MCPHandler,
208+
'registerResource',
209+
name,
210+
);
211+
}
212+
}
213+
}
214+
215+
// Resource templates: MCP SDK calls registeredResourceTemplate.readCallback
216+
const registeredResourceTemplates = server['_registeredResourceTemplates'];
217+
if (registeredResourceTemplates && typeof registeredResourceTemplates === 'object') {
218+
for (const [name, template] of Object.entries(
219+
registeredResourceTemplates as Record<string, Record<string, unknown>>,
220+
)) {
221+
if (typeof template['readCallback'] === 'function') {
222+
template['readCallback'] = createWrappedHandler(
223+
template['readCallback'] as MCPHandler,
224+
'registerResource',
225+
name,
226+
);
227+
}
228+
}
229+
}
230+
231+
// Prompts: MCP SDK calls registeredPrompt.handler
232+
const registeredPrompts = server['_registeredPrompts'];
233+
if (registeredPrompts && typeof registeredPrompts === 'object') {
234+
for (const [name, prompt] of Object.entries(registeredPrompts as Record<string, Record<string, unknown>>)) {
235+
if (typeof prompt['handler'] === 'function') {
236+
prompt['handler'] = createWrappedHandler(prompt['handler'] as MCPHandler, 'registerPrompt', name);
237+
}
238+
}
239+
}
240+
}

packages/core/src/integrations/mcp-server/index.ts

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { getClient } from '../../currentScopes';
22
import { fill } from '../../utils/object';
3-
import { wrapAllMCPHandlers } from './handlers';
3+
import { wrapAllMCPHandlers, wrapExistingHandlers } from './handlers';
44
import { wrapTransportError, wrapTransportOnClose, wrapTransportOnMessage, wrapTransportSend } from './transport';
55
import type { MCPServerInstance, McpServerWrapperOptions, MCPTransport, ResolvedMcpOptions } from './types';
66
import { validateMcpServerInstance } from './validation';
@@ -18,17 +18,24 @@ const wrappedMcpServerInstances = new WeakSet();
1818
* and versions that expose the newer `registerTool`/`registerResource`/`registerPrompt` API (introduced in 1.x, sole API in 2.x).
1919
* Automatically instruments transport methods and handler functions for comprehensive monitoring.
2020
*
21+
* Both call orderings are supported: wrapping before or after registering tools, resources,
22+
* and prompts. Sentry patches the registration methods for future handlers and retroactively
23+
* wraps any already-registered ones. Wrapping at construction time is recommended by
24+
* convention (consistent with other SDK integrations), but is not required.
25+
*
2126
* @example
2227
* ```typescript
2328
* import * as Sentry from '@sentry/core';
2429
* import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
2530
* import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
2631
*
27-
* // Default: inputs/outputs captured based on sendDefaultPii option
32+
* // Wrap first, then register tools — this is the correct order
2833
* const server = Sentry.wrapMcpServerWithSentry(
2934
* new McpServer({ name: "my-server", version: "1.0.0" })
3035
* );
3136
*
37+
* server.registerTool('my-tool', schema, handler);
38+
*
3239
* // Explicitly control input/output capture
3340
* const server = Sentry.wrapMcpServerWithSentry(
3441
* new McpServer({ name: "my-server", version: "1.0.0" }),
@@ -80,6 +87,8 @@ export function wrapMcpServerWithSentry<S extends object>(mcpServerInstance: S,
8087

8188
wrapAllMCPHandlers(serverInstance);
8289

90+
wrapExistingHandlers(serverInstance);
91+
8392
wrappedMcpServerInstances.add(mcpServerInstance);
8493
return mcpServerInstance;
8594
}

packages/core/test/lib/integrations/mcp-server/mcpServerWrapper.test.ts

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,11 @@ import { beforeEach, describe, expect, it, vi } from 'vitest';
22
import * as currentScopes from '../../../../src/currentScopes';
33
import { wrapMcpServerWithSentry } from '../../../../src/integrations/mcp-server';
44
import * as tracingModule from '../../../../src/tracing';
5-
import { createMockMcpServer, createMockMcpServerWithRegisterApi } from './testUtils';
5+
import {
6+
createMockMcpServer,
7+
createMockMcpServerWithPreregisteredHandlers,
8+
createMockMcpServerWithRegisterApi,
9+
} from './testUtils';
610

711
describe('wrapMcpServerWithSentry', () => {
812
const startSpanSpy = vi.spyOn(tracingModule, 'startSpan');
@@ -145,6 +149,41 @@ describe('wrapMcpServerWithSentry', () => {
145149
});
146150
});
147151

152+
describe('Retroactive handler wrapping (handlers registered before wrapMcpServerWithSentry)', () => {
153+
it('should replace executor/readCallback/handler on pre-registered entries with wrapped versions', () => {
154+
const server = createMockMcpServerWithPreregisteredHandlers();
155+
const { toolExecutor, resourceReadCallback, resourceTemplateReadCallback, promptHandler } = server._originals;
156+
157+
wrapMcpServerWithSentry(server);
158+
159+
expect(server._registeredTools['my-tool']!.executor).not.toBe(toolExecutor);
160+
expect(server._registeredResources['res://my-resource']!.readCallback).not.toBe(resourceReadCallback);
161+
expect(server._registeredResourceTemplates['my-template']!.readCallback).not.toBe(resourceTemplateReadCallback);
162+
expect(server._registeredPrompts['my-prompt']!.handler).not.toBe(promptHandler);
163+
});
164+
165+
it('should still wrap the registration methods for future handlers', () => {
166+
const server = createMockMcpServerWithPreregisteredHandlers();
167+
const originalRegisterTool = server.registerTool;
168+
169+
wrapMcpServerWithSentry(server);
170+
171+
expect(server.registerTool).not.toBe(originalRegisterTool);
172+
});
173+
174+
it('should not double-wrap if called twice on the same instance with pre-registered handlers', () => {
175+
const server = createMockMcpServerWithPreregisteredHandlers();
176+
177+
wrapMcpServerWithSentry(server);
178+
const executorAfterFirstWrap = server._registeredTools['my-tool']!.executor;
179+
180+
wrapMcpServerWithSentry(server);
181+
const executorAfterSecondWrap = server._registeredTools['my-tool']!.executor;
182+
183+
expect(executorAfterFirstWrap).toBe(executorAfterSecondWrap);
184+
});
185+
});
186+
148187
describe('Handler Wrapping (register* API)', () => {
149188
let mockServer: ReturnType<typeof createMockMcpServerWithRegisterApi>;
150189
let wrappedServer: ReturnType<typeof createMockMcpServerWithRegisterApi>;

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,41 @@ export function createMockMcpServer() {
1515
};
1616
}
1717

18+
/**
19+
* Create a mock MCP server that simulates already having tools/resources/prompts registered
20+
* (i.e. wrapMcpServerWithSentry is called after registration). Mirrors the internal shape
21+
* used by McpServer v2: tools have an `executor`, resources/prompts have `readCallback`/`handler`.
22+
*/
23+
export function createMockMcpServerWithPreregisteredHandlers() {
24+
const toolExecutor = vi.fn().mockResolvedValue({ content: [] });
25+
const resourceReadCallback = vi.fn().mockResolvedValue({ contents: [] });
26+
const resourceTemplateReadCallback = vi.fn().mockResolvedValue({ contents: [] });
27+
const promptHandler = vi.fn().mockResolvedValue({ messages: [] });
28+
29+
return {
30+
registerTool: vi.fn(),
31+
registerResource: vi.fn(),
32+
registerPrompt: vi.fn(),
33+
connect: vi.fn().mockResolvedValue(undefined),
34+
server: { setRequestHandler: vi.fn() },
35+
// Simulated internal registries (mirrors McpServer v2 private fields)
36+
_registeredTools: {
37+
'my-tool': { executor: toolExecutor },
38+
},
39+
_registeredResources: {
40+
'res://my-resource': { readCallback: resourceReadCallback },
41+
},
42+
_registeredResourceTemplates: {
43+
'my-template': { readCallback: resourceTemplateReadCallback },
44+
},
45+
_registeredPrompts: {
46+
'my-prompt': { handler: promptHandler },
47+
},
48+
// Expose the original fns so tests can assert wrapping happened
49+
_originals: { toolExecutor, resourceReadCallback, resourceTemplateReadCallback, promptHandler },
50+
};
51+
}
52+
1853
/**
1954
* Create a mock MCP server instance using the new register* API (SDK >=1.x / 2.x)
2055
*/

0 commit comments

Comments
 (0)