Draft
Show file tree
Hide file tree
Changes from all commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Failed to load files.
Original file line numberDiff line numberDiff line change
Expand Up@@ -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';
Expand Down
Original file line numberDiff line numberDiff line change
@@ -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();
});
});
});
Loading
Loading