Add tests

Author: lforst

packages/core/src/mcp-server.ts

@@ -25,8 +25,7 @@ const wrappedMcpServerInstances = new WeakSet();
  * 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.
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-export function wrapMcpServerWithSentry<S>(mcpServerInstance: any): S {
+export function wrapMcpServerWithSentry<S extends object>(mcpServerInstance: S): S {
   if (wrappedMcpServerInstances.has(mcpServerInstance)) {
     return mcpServerInstance;
   }

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();
+    });
+  });
+});