feat(core): Add wrapMcpServerWithSentry to instrument MCP servers from @modelcontextprotocol/sdk

packages/core/src/index.ts

@@ -111,6 +111,7 @@ export { thirdPartyErrorFilterIntegration } from './integrations/third-party-err
 export { profiler } from './profiling';
 export { instrumentFetchRequest } from './fetch';
 export { trpcMiddleware } from './trpc';
+export { wrapMcpServerWithSentry } from './mcp-server';
 export { captureFeedback } from './feedback';
 export type { ReportDialogOptions } from './report-dialog';
 export { _INTERNAL_captureLog, _INTERNAL_flushLogsBuffer } from './logs/exports';

packages/core/src/mcp-server.ts

@@ -0,0 +1,129 @@
+import { DEBUG_BUILD } from './debug-build';
+import {
+  SEMANTIC_ATTRIBUTE_SENTRY_OP,
+  SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
+  SEMANTIC_ATTRIBUTE_SENTRY_SOURCE,
+} from './semanticAttributes';
+import { startSpan } from './tracing';
+import { logger } from './utils-hoist';
+
+interface MCPServerInstance {
+  // The first arg is always a name, the last arg should always be a callback function (ie a handler).
+  // TODO: We could also make use of the resource uri argument somehow.
+  resource: (name: string, ...args: unknown[]) => void;
+  // The first arg is always a name, the last arg should always be a callback function (ie a handler).
+  tool: (name: string, ...args: unknown[]) => void;
+  // The first arg is always a name, the last arg should always be a callback function (ie a handler).
+  prompt: (name: string, ...args: unknown[]) => void;
+}
+
+const wrappedMcpServerInstances = new WeakSet();
+
+/**
+ * Wraps a MCP Server instance from the `@modelcontextprotocol/sdk` package with Sentry instrumentation.
+ *
+ * Compatible with versions `^1.9.0` of the `@modelcontextprotocol/sdk` package.
+ */
+// We are exposing this API for non-node runtimes that cannot rely on auto-instrumentation.
+export function wrapMcpServerWithSentry<S extends object>(mcpServerInstance: S): S {
+  if (wrappedMcpServerInstances.has(mcpServerInstance)) {
+    return mcpServerInstance;
+  }
+
+  if (!isMcpServerInstance(mcpServerInstance)) {
+    DEBUG_BUILD && logger.warn('Did not patch MCP server. Interface is incompatible.');
+    return mcpServerInstance;
+  }
+
+  mcpServerInstance.resource = new Proxy(mcpServerInstance.resource, {
+    apply(target, thisArg, argArray) {
+      const resourceName: unknown = argArray[0];
+      const resourceHandler: unknown = argArray[argArray.length - 1];
+
+      if (typeof resourceName !== 'string' || typeof resourceHandler !== 'function') {
+        return target.apply(thisArg, argArray);
+      }
+
+      return startSpan(
+        {
+          name: `mcp-server/resource:${resourceName}`,
+          forceTransaction: true,
+          attributes: {
+            [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server',
+            [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server',
+            [SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
+            'mcp_server.resource': resourceName,
+          },
+        },
+        () => target.apply(thisArg, argArray),
+      );
+    },
+  });
+
+  mcpServerInstance.tool = new Proxy(mcpServerInstance.tool, {
+    apply(target, thisArg, argArray) {
+      const toolName: unknown = argArray[0];
+      const toolHandler: unknown = argArray[argArray.length - 1];
+
+      if (typeof toolName !== 'string' || typeof toolHandler !== 'function') {
+        return target.apply(thisArg, argArray);
+      }
+
+      return startSpan(
+        {
+          name: `mcp-server/tool:${toolName}`,
+          forceTransaction: true,
+          attributes: {
+            [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server',
+            [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server',
+            [SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
+            'mcp_server.tool': toolName,
+          },
+        },
+        () => target.apply(thisArg, argArray),
+      );
+    },
+  });
+
+  mcpServerInstance.prompt = new Proxy(mcpServerInstance.prompt, {
+    apply(target, thisArg, argArray) {
+      const promptName: unknown = argArray[0];
+      const promptHandler: unknown = argArray[argArray.length - 1];
+
+      if (typeof promptName !== 'string' || typeof promptHandler !== 'function') {
+        return target.apply(thisArg, argArray);
+      }
+
+      return startSpan(
+        {
+          name: `mcp-server/resource:${promptName}`,
+          forceTransaction: true,
+          attributes: {
+            [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server',
+            [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server',
+            [SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
+            'mcp_server.prompt': promptName,
+          },
+        },
+        () => target.apply(thisArg, argArray),
+      );
+    },
+  });
+
+  wrappedMcpServerInstances.add(mcpServerInstance);
+
+  return mcpServerInstance as S;
+}
+
+function isMcpServerInstance(mcpServerInstance: unknown): mcpServerInstance is MCPServerInstance {
+  return (
+    typeof mcpServerInstance === 'object' &&
+    mcpServerInstance !== null &&
+    'resource' in mcpServerInstance &&
+    typeof mcpServerInstance.resource === 'function' &&
+    'tool' in mcpServerInstance &&
+    typeof mcpServerInstance.tool === 'function' &&
+    'prompt' in mcpServerInstance &&
+    typeof mcpServerInstance.prompt === 'function'
+  );
+}

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

@@ -0,0 +1,242 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+import { wrapMcpServerWithSentry } from '../../src/mcp-server';
+import {
+  SEMANTIC_ATTRIBUTE_SENTRY_OP,
+  SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
+  SEMANTIC_ATTRIBUTE_SENTRY_SOURCE,
+} from '../../src/semanticAttributes';
+import * as tracingModule from '../../src/tracing';
+
+vi.mock('../../src/tracing');
+
+describe('wrapMcpServerWithSentry', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    // @ts-expect-error mocking span is annoying
+    vi.mocked(tracingModule.startSpan).mockImplementation((_, cb) => cb());
+  });
+
+  it('should wrap valid MCP server instance methods with Sentry spans', () => {
+    // Create a mock MCP server instance
+    const mockResource = vi.fn();
+    const mockTool = vi.fn();
+    const mockPrompt = vi.fn();
+
+    const mockMcpServer = {
+      resource: mockResource,
+      tool: mockTool,
+      prompt: mockPrompt,
+    };
+
+    // Wrap the MCP server
+    const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
+
+    // Verify it returns the same instance (modified)
+    expect(wrappedMcpServer).toBe(mockMcpServer);
+
+    // Original methods should be wrapped
+    expect(wrappedMcpServer.resource).not.toBe(mockResource);
+    expect(wrappedMcpServer.tool).not.toBe(mockTool);
+    expect(wrappedMcpServer.prompt).not.toBe(mockPrompt);
+  });
+
+  it('should return the input unchanged if it is not a valid MCP server instance', () => {
+    const invalidMcpServer = {
+      // Missing required methods
+      resource: () => {},
+      tool: () => {},
+      // No prompt method
+    };
+
+    const result = wrapMcpServerWithSentry(invalidMcpServer);
+    expect(result).toBe(invalidMcpServer);
+
+    // Methods should not be wrapped
+    expect(result.resource).toBe(invalidMcpServer.resource);
+    expect(result.tool).toBe(invalidMcpServer.tool);
+
+    // No calls to startSpan
+    expect(tracingModule.startSpan).not.toHaveBeenCalled();
+  });
+
+  it('should not wrap the same instance twice', () => {
+    const mockMcpServer = {
+      resource: vi.fn(),
+      tool: vi.fn(),
+      prompt: vi.fn(),
+    };
+
+    // First wrap
+    const wrappedOnce = wrapMcpServerWithSentry(mockMcpServer);
+
+    // Store references to wrapped methods
+    const wrappedResource = wrappedOnce.resource;
+    const wrappedTool = wrappedOnce.tool;
+    const wrappedPrompt = wrappedOnce.prompt;
+
+    // Second wrap
+    const wrappedTwice = wrapMcpServerWithSentry(wrappedOnce);
+
+    // Should be the same instance with the same wrapped methods
+    expect(wrappedTwice).toBe(wrappedOnce);
+    expect(wrappedTwice.resource).toBe(wrappedResource);
+    expect(wrappedTwice.tool).toBe(wrappedTool);
+    expect(wrappedTwice.prompt).toBe(wrappedPrompt);
+  });
+
+  describe('resource method wrapping', () => {
+    it('should create a span with proper attributes when resource is called', () => {
+      const mockResourceHandler = vi.fn();
+      const resourceName = 'test-resource';
+
+      const mockMcpServer = {
+        resource: vi.fn(),
+        tool: vi.fn(),
+        prompt: vi.fn(),
+      };
+
+      const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
+      wrappedMcpServer.resource(resourceName, {}, mockResourceHandler);
+
+      expect(tracingModule.startSpan).toHaveBeenCalledTimes(1);
+      expect(tracingModule.startSpan).toHaveBeenCalledWith(
+        {
+          name: `mcp-server/resource:${resourceName}`,
+          forceTransaction: true,
+          attributes: {
+            [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server',
+            [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server',
+            [SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
+            'mcp_server.resource': resourceName,
+          },
+        },
+        expect.any(Function),
+      );
+
+      // Verify the original method was called with all arguments
+      expect(mockMcpServer.resource).toHaveBeenCalledWith(resourceName, {}, mockResourceHandler);
+    });
+
+    it('should call the original resource method directly if name or handler is not valid', () => {
+      const mockMcpServer = {
+        resource: vi.fn(),
+        tool: vi.fn(),
+        prompt: vi.fn(),
+      };
+
+      const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
+
+      // Call without string name
+      wrappedMcpServer.resource({} as any, 'handler');
+
+      // Call without function handler
+      wrappedMcpServer.resource('name', 'not-a-function');
+
+      // Original method should be called directly without creating spans
+      expect(mockMcpServer.resource).toHaveBeenCalledTimes(2);
+      expect(tracingModule.startSpan).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('tool method wrapping', () => {
+    it('should create a span with proper attributes when tool is called', () => {
+      const mockToolHandler = vi.fn();
+      const toolName = 'test-tool';
+
+      const mockMcpServer = {
+        resource: vi.fn(),
+        tool: vi.fn(),
+        prompt: vi.fn(),
+      };
+
+      const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
+      wrappedMcpServer.tool(toolName, {}, mockToolHandler);
+
+      expect(tracingModule.startSpan).toHaveBeenCalledTimes(1);
+      expect(tracingModule.startSpan).toHaveBeenCalledWith(
+        {
+          name: `mcp-server/tool:${toolName}`,
+          forceTransaction: true,
+          attributes: {
+            [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server',
+            [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server',
+            [SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
+            'mcp_server.tool': toolName,
+          },
+        },
+        expect.any(Function),
+      );
+
+      // Verify the original method was called with all arguments
+      expect(mockMcpServer.tool).toHaveBeenCalledWith(toolName, {}, mockToolHandler);
+    });
+
+    it('should call the original tool method directly if name or handler is not valid', () => {
+      const mockMcpServer = {
+        resource: vi.fn(),
+        tool: vi.fn(),
+        prompt: vi.fn(),
+      };
+
+      const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
+
+      // Call without string name
+      wrappedMcpServer.tool({} as any, 'handler');
+
+      // Original method should be called directly without creating spans
+      expect(mockMcpServer.tool).toHaveBeenCalledTimes(1);
+      expect(tracingModule.startSpan).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('prompt method wrapping', () => {
+    it('should create a span with proper attributes when prompt is called', () => {
+      const mockPromptHandler = vi.fn();
+      const promptName = 'test-prompt';
+
+      const mockMcpServer = {
+        resource: vi.fn(),
+        tool: vi.fn(),
+        prompt: vi.fn(),
+      };
+
+      const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
+      wrappedMcpServer.prompt(promptName, {}, mockPromptHandler);
+
+      expect(tracingModule.startSpan).toHaveBeenCalledTimes(1);
+      expect(tracingModule.startSpan).toHaveBeenCalledWith(
+        {
+          name: `mcp-server/resource:${promptName}`,
+          forceTransaction: true,
+          attributes: {
+            [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'auto.function.mcp-server',
+            [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.function.mcp-server',
+            [SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: 'route',
+            'mcp_server.prompt': promptName,
+          },
+        },
+        expect.any(Function),
+      );
+
+      // Verify the original method was called with all arguments
+      expect(mockMcpServer.prompt).toHaveBeenCalledWith(promptName, {}, mockPromptHandler);
+    });
+
+    it('should call the original prompt method directly if name or handler is not valid', () => {
+      const mockMcpServer = {
+        resource: vi.fn(),
+        tool: vi.fn(),
+        prompt: vi.fn(),
+      };
+
+      const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer);
+
+      // Call without function handler
+      wrappedMcpServer.prompt('name', 'not-a-function');
+
+      // Original method should be called directly without creating spans
+      expect(mockMcpServer.prompt).toHaveBeenCalledTimes(1);
+      expect(tracingModule.startSpan).not.toHaveBeenCalled();
+    });
+  });
+});