diff --git a/packages/core/src/mcp-server.ts b/packages/core/src/mcp-server.ts index 1a6f626a83f2..85b1e03ec621 100644 --- a/packages/core/src/mcp-server.ts +++ b/packages/core/src/mcp-server.ts @@ -1,31 +1,17 @@ -import { DEBUG_BUILD } from './debug-build'; +import type { + ExtraHandlerData, + MCPServerInstance, + MCPTransport, +} from './utils/mcp-server/types'; import { - SEMANTIC_ATTRIBUTE_SENTRY_OP, - SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, - SEMANTIC_ATTRIBUTE_SENTRY_SOURCE, -} from './semanticAttributes'; -import { startSpan, withActiveSpan } from './tracing'; -import type { Span } from './types-hoist/span'; -import { logger } from './utils/logger'; -import { getActiveSpan } from './utils/spanUtils'; - -interface MCPTransport { - // The first argument is a JSON RPC message - onmessage?: (...args: unknown[]) => void; - onclose?: (...args: unknown[]) => void; - sessionId?: string; -} - -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; - connect(transport: MCPTransport): Promise; -} + createMcpNotificationSpan, + createMcpOutgoingNotificationSpan, + createMcpServerSpan, + isJsonRpcNotification, + isJsonRpcRequest, + validateMcpServerInstance, +} from './utils/mcp-server/utils'; +import { fill } from './utils/object'; const wrappedMcpServerInstances = new WeakSet(); @@ -40,253 +26,105 @@ export function wrapMcpServerWithSentry(mcpServerInstance: S): return mcpServerInstance; } - if (!isMcpServerInstance(mcpServerInstance)) { - DEBUG_BUILD && logger.warn('Did not patch MCP server. Interface is incompatible.'); + if (!validateMcpServerInstance(mcpServerInstance)) { return mcpServerInstance; } - // eslint-disable-next-line @typescript-eslint/unbound-method - mcpServerInstance.connect = new Proxy(mcpServerInstance.connect, { - apply(target, thisArg, argArray) { - const [transport, ...restArgs] = argArray as [MCPTransport, ...unknown[]]; - - if (!transport.onclose) { - transport.onclose = () => { - if (transport.sessionId) { - handleTransportOnClose(transport.sessionId); - } - }; + const serverInstance = mcpServerInstance as MCPServerInstance; + + fill(serverInstance, 'connect', (originalConnect) => { + return async function(this: MCPServerInstance, transport: MCPTransport, ...restArgs: unknown[]) { + const result = await originalConnect.call(this, transport, ...restArgs); + + if (transport.onmessage) { + fill(transport, 'onmessage', (originalOnMessage) => { + return function(this: MCPTransport, jsonRpcMessage: unknown, extra?: unknown) { + if (isJsonRpcRequest(jsonRpcMessage)) { + return createMcpServerSpan(jsonRpcMessage, this, extra as ExtraHandlerData, () => { + return originalOnMessage.call(this, jsonRpcMessage, extra); + }); + } + if (isJsonRpcNotification(jsonRpcMessage)) { + return createMcpNotificationSpan(jsonRpcMessage, this, extra as ExtraHandlerData, () => { + return originalOnMessage.call(this, jsonRpcMessage, extra); + }); + } + return originalOnMessage.call(this, jsonRpcMessage, extra); + }; + }); } - if (!transport.onmessage) { - transport.onmessage = jsonRpcMessage => { - if (transport.sessionId && isJsonRPCMessageWithRequestId(jsonRpcMessage)) { - handleTransportOnMessage(transport.sessionId, jsonRpcMessage.id); - } - }; + if (transport.send) { + fill(transport, 'send', (originalSend) => { + return async function(this: MCPTransport, message: unknown) { + if (isJsonRpcNotification(message)) { + return createMcpOutgoingNotificationSpan(message, this, () => { + return originalSend.call(this, message); + }); + } + return originalSend.call(this, message); + }; + }); } - const patchedTransport = new Proxy(transport, { - set(target, key, value) { - if (key === 'onmessage') { - target[key] = new Proxy(value, { - apply(onMessageTarget, onMessageThisArg, onMessageArgArray) { - const [jsonRpcMessage] = onMessageArgArray; - if (transport.sessionId && isJsonRPCMessageWithRequestId(jsonRpcMessage)) { - handleTransportOnMessage(transport.sessionId, jsonRpcMessage.id); - } - return Reflect.apply(onMessageTarget, onMessageThisArg, onMessageArgArray); - }, - }); - } else if (key === 'onclose') { - target[key] = new Proxy(value, { - apply(onCloseTarget, onCloseThisArg, onCloseArgArray) { - if (transport.sessionId) { - handleTransportOnClose(transport.sessionId); - } - return Reflect.apply(onCloseTarget, onCloseThisArg, onCloseArgArray); - }, - }); - } else { - target[key as keyof MCPTransport] = value; - } - return true; - }, - }); - - return Reflect.apply(target, thisArg, [patchedTransport, ...restArgs]); - }, - }); - - 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); + if (transport.onclose) { + fill(transport, 'onclose', (originalOnClose) => { + return function(this: MCPTransport, ...args: unknown[]) { + return originalOnClose.call(this, ...args); + }; + }); } - - const wrappedResourceHandler = new Proxy(resourceHandler, { - apply(resourceHandlerTarget, resourceHandlerThisArg, resourceHandlerArgArray) { - const extraHandlerDataWithRequestId = resourceHandlerArgArray.find(isExtraHandlerDataWithRequestId); - return associateContextWithRequestSpan(extraHandlerDataWithRequestId, () => { - 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, - }, - }, - () => resourceHandlerTarget.apply(resourceHandlerThisArg, resourceHandlerArgArray), - ); - }); - }, - }); - - return Reflect.apply(target, thisArg, [...argArray.slice(0, -1), wrappedResourceHandler]); - }, - }); - - 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); - } - - const wrappedToolHandler = new Proxy(toolHandler, { - apply(toolHandlerTarget, toolHandlerThisArg, toolHandlerArgArray) { - const extraHandlerDataWithRequestId = toolHandlerArgArray.find(isExtraHandlerDataWithRequestId); - return associateContextWithRequestSpan(extraHandlerDataWithRequestId, () => { - 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, - }, - }, - () => toolHandlerTarget.apply(toolHandlerThisArg, toolHandlerArgArray), - ); - }); - }, - }); - - return Reflect.apply(target, thisArg, [...argArray.slice(0, -1), wrappedToolHandler]); - }, - }); - - 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); - } - - const wrappedPromptHandler = new Proxy(promptHandler, { - apply(promptHandlerTarget, promptHandlerThisArg, promptHandlerArgArray) { - const extraHandlerDataWithRequestId = promptHandlerArgArray.find(isExtraHandlerDataWithRequestId); - return associateContextWithRequestSpan(extraHandlerDataWithRequestId, () => { - return startSpan( - { - name: `mcp-server/prompt:${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, - }, - }, - () => promptHandlerTarget.apply(promptHandlerThisArg, promptHandlerArgArray), - ); - }); - }, - }); - - return Reflect.apply(target, thisArg, [...argArray.slice(0, -1), wrappedPromptHandler]); - }, + return result; + }; }); 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' && - 'connect' in mcpServerInstance && - typeof mcpServerInstance.connect === 'function' - ); -} - -function isJsonRPCMessageWithRequestId(target: unknown): target is { id: RequestId } { - return ( - typeof target === 'object' && - target !== null && - 'id' in target && - (typeof target.id === 'number' || typeof target.id === 'string') - ); -} - -interface ExtraHandlerDataWithRequestId { - sessionId: SessionId; - requestId: RequestId; -} - -// Note that not all versions of the MCP library have `requestId` as a field on the extra data. -function isExtraHandlerDataWithRequestId(target: unknown): target is ExtraHandlerDataWithRequestId { - return ( - typeof target === 'object' && - target !== null && - 'sessionId' in target && - typeof target.sessionId === 'string' && - 'requestId' in target && - (typeof target.requestId === 'number' || typeof target.requestId === 'string') - ); -} - -type SessionId = string; -type RequestId = string | number; - -const sessionAndRequestToRequestParentSpanMap = new Map>(); - -function handleTransportOnClose(sessionId: SessionId): void { - sessionAndRequestToRequestParentSpanMap.delete(sessionId); -} - -function handleTransportOnMessage(sessionId: SessionId, requestId: RequestId): void { - const activeSpan = getActiveSpan(); - if (activeSpan) { - const requestIdToSpanMap = sessionAndRequestToRequestParentSpanMap.get(sessionId) ?? new Map(); - requestIdToSpanMap.set(requestId, activeSpan); - sessionAndRequestToRequestParentSpanMap.set(sessionId, requestIdToSpanMap); - } -} - -function associateContextWithRequestSpan( - extraHandlerData: ExtraHandlerDataWithRequestId | undefined, - cb: () => T, -): T { - if (extraHandlerData) { - const { sessionId, requestId } = extraHandlerData; - const requestIdSpanMap = sessionAndRequestToRequestParentSpanMap.get(sessionId); - - if (!requestIdSpanMap) { - return cb(); - } - - const span = requestIdSpanMap.get(requestId); - if (!span) { - return cb(); - } - - // remove the span from the map so it can be garbage collected - requestIdSpanMap.delete(requestId); - return withActiveSpan(span, () => { - return cb(); - }); - } - - return cb(); -} +// ============================================================================= +// SESSION AND REQUEST CORRELATION (Legacy support) +// ============================================================================= + +// const sessionAndRequestToRequestParentSpanMap = new Map>(); + +// function handleTransportOnClose(sessionId: SessionId): void { +// sessionAndRequestToRequestParentSpanMap.delete(sessionId); +// } + +// TODO(bete): refactor this and associateContextWithRequestSpan to use the new span API. +// function handleTransportOnMessage(sessionId: SessionId, requestId: string): void { +// const activeSpan = getActiveSpan(); +// if (activeSpan) { +// const requestIdToSpanMap = sessionAndRequestToRequestParentSpanMap.get(sessionId) ?? new Map(); +// requestIdToSpanMap.set(requestId, activeSpan); +// sessionAndRequestToRequestParentSpanMap.set(sessionId, requestIdToSpanMap); +// } +// } + +// function associateContextWithRequestSpan( +// extraHandlerData: { sessionId: SessionId; requestId: string } | undefined, +// cb: () => T, +// ): T { +// if (extraHandlerData) { +// const { sessionId, requestId } = extraHandlerData; +// const requestIdSpanMap = sessionAndRequestToRequestParentSpanMap.get(sessionId); + +// if (!requestIdSpanMap) { +// return cb(); +// } + +// const span = requestIdSpanMap.get(requestId); +// if (!span) { +// return cb(); +// } + +// // remove the span from the map so it can be garbage collected +// requestIdSpanMap.delete(requestId); +// return withActiveSpan(span, () => { +// return cb(); +// }); +// } + +// return cb(); +// } diff --git a/packages/core/src/utils/mcp-server/attributes.ts b/packages/core/src/utils/mcp-server/attributes.ts new file mode 100644 index 000000000000..1df3952401d3 --- /dev/null +++ b/packages/core/src/utils/mcp-server/attributes.ts @@ -0,0 +1,120 @@ +/** + * Essential MCP attribute constants for Sentry instrumentation + * + * Based on OpenTelemetry MCP semantic conventions + * @see https://github.com/open-telemetry/semantic-conventions/blob/3097fb0af5b9492b0e3f55dc5f6c21a3dc2be8df/docs/gen-ai/mcp.md + */ + +// ============================================================================= +// CORE MCP ATTRIBUTES +// ============================================================================= + +/** + * The name of the request or notification method + * @see https://github.com/open-telemetry/semantic-conventions/blob/3097fb0af5b9492b0e3f55dc5f6c21a3dc2be8df/docs/gen-ai/mcp.md#required-attributes + */ +export const MCP_METHOD_NAME_ATTRIBUTE = 'mcp.method.name'; + +/** + * Unique identifier for the request + * @see https://github.com/open-telemetry/semantic-conventions/blob/3097fb0af5b9492b0e3f55dc5f6c21a3dc2be8df/docs/gen-ai/mcp.md#recommended-attributes + */ +export const MCP_REQUEST_ID_ATTRIBUTE = 'mcp.request.id'; + +/** + * Identifies the MCP session + * @see https://github.com/open-telemetry/semantic-conventions/blob/3097fb0af5b9492b0e3f55dc5f6c21a3dc2be8df/docs/gen-ai/mcp.md#recommended-attributes + */ +export const MCP_SESSION_ID_ATTRIBUTE = 'mcp.session.id'; + +/** + * Transport method used for MCP communication + * @see https://github.com/open-telemetry/semantic-conventions/blob/3097fb0af5b9492b0e3f55dc5f6c21a3dc2be8df/docs/gen-ai/mcp.md#recommended-attributes + */ +export const MCP_TRANSPORT_ATTRIBUTE = 'mcp.transport'; + +// ============================================================================= +// METHOD-SPECIFIC ATTRIBUTES +// ============================================================================= + +/** + * Name of the tool being called + * @see https://github.com/open-telemetry/semantic-conventions/blob/3097fb0af5b9492b0e3f55dc5f6c21a3dc2be8df/docs/gen-ai/mcp.md#method-specific-attributes + */ +export const MCP_TOOL_NAME_ATTRIBUTE = 'mcp.tool.name'; + +/** + * The resource URI being accessed + * @see https://github.com/open-telemetry/semantic-conventions/blob/3097fb0af5b9492b0e3f55dc5f6c21a3dc2be8df/docs/gen-ai/mcp.md#method-specific-attributes + */ +export const MCP_RESOURCE_URI_ATTRIBUTE = 'mcp.resource.uri'; + +/** + * Name of the prompt template + * @see https://github.com/open-telemetry/semantic-conventions/blob/3097fb0af5b9492b0e3f55dc5f6c21a3dc2be8df/docs/gen-ai/mcp.md#method-specific-attributes + */ +export const MCP_PROMPT_NAME_ATTRIBUTE = 'mcp.prompt.name'; + +// ============================================================================= +// NETWORK ATTRIBUTES (OpenTelemetry Standard) +// ============================================================================= + +/** + * OSI transport layer protocol + * @see https://github.com/open-telemetry/semantic-conventions/blob/3097fb0af5b9492b0e3f55dc5f6c21a3dc2be8df/docs/gen-ai/mcp.md#network-attributes + */ +export const NETWORK_TRANSPORT_ATTRIBUTE = 'network.transport'; + +/** + * The version of JSON RPC protocol used + * @see https://github.com/open-telemetry/semantic-conventions/blob/3097fb0af5b9492b0e3f55dc5f6c21a3dc2be8df/docs/gen-ai/mcp.md#network-attributes + */ +export const NETWORK_PROTOCOL_VERSION_ATTRIBUTE = 'network.protocol.version'; + +/** + * Client address - domain name if available without reverse DNS lookup; otherwise, IP address or Unix domain socket name + * @see https://github.com/open-telemetry/semantic-conventions/blob/3097fb0af5b9492b0e3f55dc5f6c21a3dc2be8df/docs/gen-ai/mcp.md#network-attributes + */ +export const CLIENT_ADDRESS_ATTRIBUTE = 'client.address'; + +/** + * Client port number + * @see https://github.com/open-telemetry/semantic-conventions/blob/3097fb0af5b9492b0e3f55dc5f6c21a3dc2be8df/docs/gen-ai/mcp.md#network-attributes + */ +export const CLIENT_PORT_ATTRIBUTE = 'client.port'; + +// ============================================================================= +// SENTRY-SPECIFIC MCP ATTRIBUTE VALUES +// ============================================================================= + +/** + * Sentry operation value for MCP server spans + */ +export const MCP_SERVER_OP_VALUE = 'mcp.server'; + +/** + * Sentry operation value for client-to-server notifications + * Following OpenTelemetry MCP semantic conventions + */ +export const MCP_NOTIFICATION_CLIENT_TO_SERVER_OP_VALUE = 'mcp.notification.client_to_server'; + +/** + * Sentry operation value for server-to-client notifications + * Following OpenTelemetry MCP semantic conventions + */ +export const MCP_NOTIFICATION_SERVER_TO_CLIENT_OP_VALUE = 'mcp.notification.server_to_client'; + +/** + * Sentry origin value for MCP function spans + */ +export const MCP_FUNCTION_ORIGIN_VALUE = 'auto.function.mcp_server'; + +/** + * Sentry origin value for MCP notification spans + */ +export const MCP_NOTIFICATION_ORIGIN_VALUE = 'auto.mcp.notification'; + +/** + * Sentry source value for MCP route spans + */ +export const MCP_ROUTE_SOURCE_VALUE = 'route'; diff --git a/packages/core/src/utils/mcp-server/types.ts b/packages/core/src/utils/mcp-server/types.ts new file mode 100644 index 000000000000..a3d592b78257 --- /dev/null +++ b/packages/core/src/utils/mcp-server/types.ts @@ -0,0 +1,137 @@ +/** + * types for MCP server instrumentation + */ + + +/** Method configuration type */ +export type MethodConfig = { + targetField: string; + targetAttribute: string; + captureArguments?: boolean; + argumentsField?: string; + captureUri?: boolean; + captureName?: boolean; +}; + +/** + * JSON-RPC 2.0 request object + */ +export interface JsonRpcRequest { + jsonrpc: '2.0'; + method: string; + id: string | number; + params?: Record; +} + +/** + * JSON-RPC 2.0 response object + */ +export interface JsonRpcResponse { + jsonrpc: '2.0'; + id: string | number | null; + result?: unknown; + error?: JsonRpcError; +} + +/** + * JSON-RPC 2.0 error object + */ +export interface JsonRpcError { + code: number; + message: string; + data?: unknown; +} + +/** + * JSON-RPC 2.0 notification object + * Note: Notifications do NOT have an 'id' field - this is what distinguishes them from requests + */ +export interface JsonRpcNotification { + jsonrpc: '2.0'; + method: string; + params?: Record; +} + +/** + * MCP transport interface + */ +export interface MCPTransport { + /** + * Message handler for incoming JSON-RPC messages + * The first argument is a JSON RPC message + */ + onmessage?: (...args: unknown[]) => void; + + /** + * Close handler for transport lifecycle + */ + onclose?: (...args: unknown[]) => void; + + /** + * Send method for outgoing messages + */ + send?: (message: JsonRpcMessage, options?: Record) => Promise; + + /** + * Optional session identifier + */ + sessionId?: string; +} + +/** + * Union type for all JSON-RPC message types + */ +export type JsonRpcMessage = JsonRpcRequest | JsonRpcNotification | JsonRpcResponse; + +/** + * MCP server instance interface + */ +export interface MCPServerInstance { + /** + * Register a resource handler + * The first arg is always a name, the last arg should always be a callback function (ie a handler). + */ + resource: (name: string, ...args: unknown[]) => void; + + /** + * Register a tool handler + * 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; + + /** + * Register a prompt handler + * 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; + + /** + * Connect the server to a transport + */ + connect(transport: MCPTransport): Promise; +} + +export interface ExtraHandlerData { + requestInfo?: { remoteAddress?: string; remotePort?: number }; + clientAddress?: string; + clientPort?: number; + request?: { + ip?: string; + connection?: { remoteAddress?: string; remotePort?: number }; + }; +} + +/** + * Configuration for creating MCP spans + */ +export interface McpSpanConfig { + type: 'request' | 'notification-incoming' | 'notification-outgoing'; + message: JsonRpcRequest | JsonRpcNotification; + transport: MCPTransport; + extra?: ExtraHandlerData; + callback: () => unknown; +} + + +export type SessionId = string; +export type RequestId = string | number; diff --git a/packages/core/src/utils/mcp-server/utils.ts b/packages/core/src/utils/mcp-server/utils.ts new file mode 100644 index 000000000000..33383542101f --- /dev/null +++ b/packages/core/src/utils/mcp-server/utils.ts @@ -0,0 +1,448 @@ +/** + * Essential utility functions for MCP server instrumentation + */ + +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 '../logger'; +import { + CLIENT_ADDRESS_ATTRIBUTE, + CLIENT_PORT_ATTRIBUTE, + MCP_FUNCTION_ORIGIN_VALUE, + MCP_METHOD_NAME_ATTRIBUTE, + MCP_NOTIFICATION_CLIENT_TO_SERVER_OP_VALUE, + MCP_NOTIFICATION_ORIGIN_VALUE, + MCP_NOTIFICATION_SERVER_TO_CLIENT_OP_VALUE, + MCP_PROMPT_NAME_ATTRIBUTE, + MCP_REQUEST_ID_ATTRIBUTE, + MCP_RESOURCE_URI_ATTRIBUTE, + MCP_ROUTE_SOURCE_VALUE, + MCP_SERVER_OP_VALUE, + MCP_SESSION_ID_ATTRIBUTE, + MCP_TOOL_NAME_ATTRIBUTE, + MCP_TRANSPORT_ATTRIBUTE, + NETWORK_PROTOCOL_VERSION_ATTRIBUTE, + NETWORK_TRANSPORT_ATTRIBUTE, +} from './attributes'; +import type { ExtraHandlerData, JsonRpcNotification, JsonRpcRequest, McpSpanConfig, MCPTransport, MethodConfig } from './types'; + +/** Validates if a message is a JSON-RPC request */ +export function isJsonRpcRequest(message: unknown): message is JsonRpcRequest { + return ( + typeof message === 'object' && + message !== null && + 'jsonrpc' in message && + (message as JsonRpcRequest).jsonrpc === '2.0' && + 'method' in message && + 'id' in message + ); +} + +/** Validates if a message is a JSON-RPC notification */ +export function isJsonRpcNotification(message: unknown): message is JsonRpcNotification { + return ( + typeof message === 'object' && + message !== null && + 'jsonrpc' in message && + (message as JsonRpcNotification).jsonrpc === '2.0' && + 'method' in message && + !('id' in message) + ); +} + +/** Validates MCP server instance with comprehensive type checking */ +export function validateMcpServerInstance(instance: unknown): boolean { + if ( + typeof instance === 'object' && + instance !== null && + 'resource' in instance && + 'tool' in instance && + 'prompt' in instance && + 'connect' in instance + ) { + return true; + } + DEBUG_BUILD && logger.warn('Did not patch MCP server. Interface is incompatible.'); + return false; +} + +/** Configuration for MCP methods to extract targets and arguments */ +const METHOD_CONFIGS: Record = { + 'tools/call': { + targetField: 'name', + targetAttribute: MCP_TOOL_NAME_ATTRIBUTE, + captureArguments: true, + argumentsField: 'arguments', + }, + 'resources/read': { + targetField: 'uri', + targetAttribute: MCP_RESOURCE_URI_ATTRIBUTE, + captureUri: true, + }, + 'resources/subscribe': { + targetField: 'uri', + targetAttribute: MCP_RESOURCE_URI_ATTRIBUTE, + }, + 'resources/unsubscribe': { + targetField: 'uri', + targetAttribute: MCP_RESOURCE_URI_ATTRIBUTE, + }, + 'prompts/get': { + targetField: 'name', + targetAttribute: MCP_PROMPT_NAME_ATTRIBUTE, + captureName: true, + captureArguments: true, + argumentsField: 'arguments', + }, +}; + +/** Extracts target info from method and params based on method type */ +function extractTargetInfo(method: string, params: Record): { + target?: string; + attributes: Record +} { + const config = METHOD_CONFIGS[method as keyof typeof METHOD_CONFIGS]; + if (!config) { + return { attributes: {} }; + } + + const target = config.targetField && typeof params?.[config.targetField] === 'string' + ? params[config.targetField] as string + : undefined; + + return { + target, + attributes: target && config.targetAttribute ? { [config.targetAttribute]: target } : {} + }; +} + +/** Extracts request arguments based on method type */ +function getRequestArguments(method: string, params: Record): Record { + const args: Record = {}; + const config = METHOD_CONFIGS[method as keyof typeof METHOD_CONFIGS]; + + if (!config) { + return args; + } + + // Capture arguments from the configured field + if (config.captureArguments && config.argumentsField && params?.[config.argumentsField]) { + const argumentsObj = params[config.argumentsField]; + if (typeof argumentsObj === 'object' && argumentsObj !== null) { + for (const [key, value] of Object.entries(argumentsObj as Record)) { + args[`mcp.request.argument.${key.toLowerCase()}`] = JSON.stringify(value); + } + } + } + + // Capture specific fields as arguments + if (config.captureUri && params?.uri) { + args['mcp.request.argument.uri'] = JSON.stringify(params.uri); + } + + if (config.captureName && params?.name) { + args['mcp.request.argument.name'] = JSON.stringify(params.name); + } + + return args; +} + +/** Extracts transport types based on transport constructor name */ +function getTransportTypes(transport: MCPTransport): { mcpTransport: string; networkTransport: string } { + const transportName = transport.constructor?.name?.toLowerCase() || ''; + + // Standard MCP transports per specification + if (transportName.includes('stdio')) { + return { mcpTransport: 'stdio', networkTransport: 'pipe' }; + } + + // Streamable HTTP is the standard HTTP-based transport + // The official SDK uses 'StreamableHTTPServerTransport' / 'StreamableHTTPClientTransport' + if (transportName.includes('streamablehttp') || transportName.includes('streamable')) { + return { mcpTransport: 'http', networkTransport: 'tcp' }; + } + + // SSE is the deprecated HTTP+SSE transport (backwards compatibility) + // Note: Modern Streamable HTTP can use SSE internally, but SSE transport is deprecated + if (transportName.includes('sse')) { + return { mcpTransport: 'sse', networkTransport: 'tcp' }; + } + + // For custom transports, mark as unknown + // TODO(bete): Add support for custom transports + return { mcpTransport: 'unknown', networkTransport: 'unknown' }; +} + +/** Extracts additional attributes for specific notification types */ +function getNotificationAttributes( + method: string, + params: Record, +): Record { + const attributes: Record = {}; + + switch (method) { + case 'notifications/cancelled': + if (params?.requestId) { + attributes['mcp.cancelled.request_id'] = String(params.requestId); + } + if (params?.reason) { + attributes['mcp.cancelled.reason'] = String(params.reason); + } + break; + + case 'notifications/message': + if (params?.level) { + attributes['mcp.logging.level'] = String(params.level); + } + if (params?.logger) { + attributes['mcp.logging.logger'] = String(params.logger); + } + if (params?.data !== undefined) { + attributes['mcp.logging.data_type'] = typeof params.data; + // Store the actual message content + if (typeof params.data === 'string') { + attributes['mcp.logging.message'] = params.data; + } else { + attributes['mcp.logging.message'] = JSON.stringify(params.data); + } + } + break; + + case 'notifications/progress': + if (params?.progressToken) { + attributes['mcp.progress.token'] = String(params.progressToken); + } + if (typeof params?.progress === 'number') { + attributes['mcp.progress.current'] = params.progress; + } + if (typeof params?.total === 'number') { + attributes['mcp.progress.total'] = params.total; + if (typeof params?.progress === 'number') { + attributes['mcp.progress.percentage'] = (params.progress / params.total) * 100; + } + } + if (params?.message) { + attributes['mcp.progress.message'] = String(params.message); + } + break; + + case 'notifications/resources/updated': + if (params?.uri) { + attributes['mcp.resource.uri'] = String(params.uri); + // Extract protocol from URI + try { + const url = new URL(String(params.uri)); + attributes['mcp.resource.protocol'] = url.protocol; + } catch { + // Ignore invalid URIs + } + } + break; + + case 'notifications/initialized': + attributes['mcp.lifecycle.phase'] = 'initialization_complete'; + attributes['mcp.protocol.ready'] = 1; + break; + } + + return attributes; +} + + +/** + * Creates a span name based on the method and target + */ +function createSpanName(method: string, target?: string): string { + return target ? `${method} ${target}` : method; +} + + +/** + * Unified builder for creating MCP spans + * Follows OpenTelemetry semantic conventions for span naming + */ +function createMcpSpan(config: McpSpanConfig): unknown { + const { type, message, transport, extra, callback } = config; + const { method } = message; + const params = message.params as Record | undefined; + + // Determine span name based on type and OTEL conventions + let spanName: string; + if (type === 'request') { + const targetInfo = extractTargetInfo(method, params || {}); + spanName = createSpanName(method, targetInfo.target); + } else { + // For notifications, use method name directly per OpenTelemetry conventions + spanName = method; + } + + // Build attributes + const attributes: Record = { + // Base attributes + ...buildTransportAttributes(transport, extra), + // Method name (required for all spans) + [MCP_METHOD_NAME_ATTRIBUTE]: method, + // Type-specific attributes + ...buildTypeSpecificAttributes(type, message, params), + // Sentry attributes + ...buildSentryAttributes(type), + }; + + return startSpan( + { + name: spanName, + forceTransaction: true, + attributes, + }, + callback, + ); +} + +/** + * Build transport and network attributes + */ +function buildTransportAttributes( + transport: MCPTransport, + extra?: ExtraHandlerData, +): Record { + const sessionId = transport.sessionId; + const clientInfo = extra ? extractClientInfo(extra) : {}; + const { mcpTransport, networkTransport } = getTransportTypes(transport); + + return { + ...(sessionId && { [MCP_SESSION_ID_ATTRIBUTE]: sessionId }), + ...(clientInfo.address && { [CLIENT_ADDRESS_ATTRIBUTE]: clientInfo.address }), + ...(clientInfo.port && { [CLIENT_PORT_ATTRIBUTE]: clientInfo.port }), + [MCP_TRANSPORT_ATTRIBUTE]: mcpTransport, + [NETWORK_TRANSPORT_ATTRIBUTE]: networkTransport, + [NETWORK_PROTOCOL_VERSION_ATTRIBUTE]: '2.0', + }; +} + +/** + * Build type-specific attributes based on message type + */ +function buildTypeSpecificAttributes( + type: McpSpanConfig['type'], + message: JsonRpcRequest | JsonRpcNotification, + params?: Record, +): Record { + if (type === 'request') { + const request = message as JsonRpcRequest; + const targetInfo = extractTargetInfo(request.method, params || {}); + + return { + ...(request.id !== undefined && { [MCP_REQUEST_ID_ATTRIBUTE]: String(request.id) }), + ...targetInfo.attributes, + ...getRequestArguments(request.method, params || {}), + }; + } + + // For notifications, only include notification-specific attributes + return getNotificationAttributes(message.method, params || {}); +} + +/** + * Build Sentry-specific attributes based on span type + * Uses specific operations for notification direction + */ +function buildSentryAttributes(type: McpSpanConfig['type']): Record { + let op: string; + let origin: string; + + switch (type) { + case 'request': + op = MCP_SERVER_OP_VALUE; + origin = MCP_FUNCTION_ORIGIN_VALUE; + break; + case 'notification-incoming': + op = MCP_NOTIFICATION_CLIENT_TO_SERVER_OP_VALUE; + origin = MCP_NOTIFICATION_ORIGIN_VALUE; + break; + case 'notification-outgoing': + op = MCP_NOTIFICATION_SERVER_TO_CLIENT_OP_VALUE; + origin = MCP_NOTIFICATION_ORIGIN_VALUE; + break; + } + + return { + [SEMANTIC_ATTRIBUTE_SENTRY_OP]: op, + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: origin, + [SEMANTIC_ATTRIBUTE_SENTRY_SOURCE]: MCP_ROUTE_SOURCE_VALUE, + }; +} + +/** + * Creates a span for MCP server request handling + */ +export function createMcpServerSpan( + jsonRpcMessage: JsonRpcRequest, + transport: MCPTransport, + extra: ExtraHandlerData, + callback: () => unknown, +): unknown { + return createMcpSpan({ + type: 'request', + message: jsonRpcMessage, + transport, + extra, + callback, + }); +} + +/** + * Creates a span for incoming MCP notifications + */ +export function createMcpNotificationSpan( + jsonRpcMessage: JsonRpcNotification, + transport: MCPTransport, + extra: ExtraHandlerData, + callback: () => unknown, +): unknown { + return createMcpSpan({ + type: 'notification-incoming', + message: jsonRpcMessage, + transport, + extra, + callback, + }); +} + +/** + * Creates a span for outgoing MCP notifications + */ +export function createMcpOutgoingNotificationSpan( + jsonRpcMessage: JsonRpcNotification, + transport: MCPTransport, + callback: () => unknown, +): unknown { + return createMcpSpan({ + type: 'notification-outgoing', + message: jsonRpcMessage, + transport, + callback, + }); +} + +/** + * Combine the two extraction functions into one + */ +function extractClientInfo(extra: ExtraHandlerData): { + address?: string; + port?: number +} { + return { + address: extra?.requestInfo?.remoteAddress || + extra?.clientAddress || + extra?.request?.ip || + extra?.request?.connection?.remoteAddress, + port: extra?.requestInfo?.remotePort || + extra?.clientPort || + extra?.request?.connection?.remotePort + }; +} diff --git a/packages/core/test/lib/mcp-server.test.ts b/packages/core/test/lib/mcp-server.test.ts index 12e85f9f370e..91b1410b8209 100644 --- a/packages/core/test/lib/mcp-server.test.ts +++ b/packages/core/test/lib/mcp-server.test.ts @@ -1,52 +1,26 @@ 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', () => { + const startSpanSpy = vi.spyOn(tracingModule, 'startSpan'); + 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, - connect: vi.fn(), - }; - - // Wrap the MCP server + it('should return the same instance (modified) if it is a valid MCP server instance', () => { + const mockMcpServer = createMockMcpServer(); 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 + // Missing required methods }; const result = wrapMcpServerWithSentry(invalidMcpServer); @@ -57,214 +31,672 @@ describe('wrapMcpServerWithSentry', () => { expect(result.tool).toBe(invalidMcpServer.tool); // No calls to startSpan - expect(tracingModule.startSpan).not.toHaveBeenCalled(); + expect(startSpanSpy).not.toHaveBeenCalled(); }); it('should not wrap the same instance twice', () => { - const mockMcpServer = { - resource: vi.fn(), - tool: vi.fn(), - prompt: vi.fn(), - }; + const mockMcpServer = createMockMcpServer(); - // 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'; + describe('Transport-level instrumentation', () => { + let mockMcpServer: ReturnType; + let wrappedMcpServer: ReturnType; + let mockTransport: ReturnType; + let originalConnect: any; + + beforeEach(() => { + mockMcpServer = createMockMcpServer(); + originalConnect = mockMcpServer.connect; + wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer); + mockTransport = createMockTransport(); + }); + + it('should proxy the connect method', () => { + // We need to test this before connection, so create fresh instances + const freshMockMcpServer = createMockMcpServer(); + const originalConnect = freshMockMcpServer.connect; + + const freshWrappedMcpServer = wrapMcpServerWithSentry(freshMockMcpServer); + + expect(freshWrappedMcpServer.connect).not.toBe(originalConnect); + }); + + it('should intercept transport onmessage handler', async () => { + const originalOnMessage = mockTransport.onmessage; + + await wrappedMcpServer.connect(mockTransport); + + // onmessage should be wrapped after connection + expect(mockTransport.onmessage).not.toBe(originalOnMessage); + }); + + it('should intercept transport send handler', async () => { + const originalSend = mockTransport.send; + + await wrappedMcpServer.connect(mockTransport); + + // send should be wrapped after connection + expect(mockTransport.send).not.toBe(originalSend); + }); + + it('should intercept transport onclose handler', async () => { + const originalOnClose = mockTransport.onclose; + + await wrappedMcpServer.connect(mockTransport); + + // onclose should be wrapped after connection + expect(mockTransport.onclose).not.toBe(originalOnClose); + }); + + it('should call original connect and preserve functionality', async () => { + await wrappedMcpServer.connect(mockTransport); + + // Check the original spy was called + expect(originalConnect).toHaveBeenCalledWith(mockTransport); + }); + + it('should create spans for incoming JSON-RPC requests', async () => { + await wrappedMcpServer.connect(mockTransport); + + const jsonRpcRequest = { + jsonrpc: '2.0', + method: 'tools/call', + id: 'req-1', + params: { name: 'get-weather' }, + }; + + // Simulate incoming message + mockTransport.onmessage?.(jsonRpcRequest, {}); + + expect(startSpanSpy).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'tools/call get-weather', + forceTransaction: true, + }), + expect.any(Function), + ); + }); + + it('should create spans for incoming JSON-RPC notifications', async () => { + await wrappedMcpServer.connect(mockTransport); - const mockMcpServer = { - resource: vi.fn(), - tool: vi.fn(), - prompt: vi.fn(), - connect: vi.fn(), + const jsonRpcNotification = { + jsonrpc: '2.0', + method: 'notifications/initialized', + // No 'id' field - this makes it a notification }; - const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer); - wrappedMcpServer.resource(resourceName, {}, mockResourceHandler); + // Simulate incoming notification + mockTransport.onmessage?.(jsonRpcNotification, {}); + + expect(startSpanSpy).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'notifications/initialized', + forceTransaction: true, + }), + expect.any(Function), + ); + }); + + it('should create spans for outgoing notifications', async () => { + await wrappedMcpServer.connect(mockTransport); + + const outgoingNotification = { + jsonrpc: '2.0', + method: 'notifications/tools/list_changed', + // No 'id' field + }; + + // Simulate outgoing notification + await mockTransport.send?.(outgoingNotification); + + expect(startSpanSpy).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'notifications/tools/list_changed', + forceTransaction: true, + }), + expect.any(Function), + ); + }); + + it('should not create spans for non-JSON-RPC messages', async () => { + await wrappedMcpServer.connect(mockTransport); + + // Simulate non-JSON-RPC message + mockTransport.onmessage?.({ some: 'data' }, {}); + + expect(startSpanSpy).not.toHaveBeenCalled(); + }); + + it('should handle transport onclose events', async () => { + await wrappedMcpServer.connect(mockTransport); + mockTransport.sessionId = 'test-session-123'; + + // Trigger onclose - should not throw + expect(() => mockTransport.onclose?.()).not.toThrow(); + }); + }); - // The original registration should use a wrapped handler - expect(mockMcpServer.resource).toHaveBeenCalledWith(resourceName, {}, expect.any(Function)); + describe('Span Creation & Semantic Conventions', () => { + let mockMcpServer: ReturnType; + let wrappedMcpServer: ReturnType; + let mockTransport: ReturnType; + + beforeEach(() => { + mockMcpServer = createMockMcpServer(); + wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer); + mockTransport = createMockTransport(); + mockTransport.sessionId = 'test-session-123'; + // Don't connect here - let individual tests control when connection happens + }); + + it('should create spans with correct MCP server semantic attributes for tool operations', async () => { + await wrappedMcpServer.connect(mockTransport); - // Invoke the wrapped handler to trigger Sentry span - const wrappedResourceHandler = (mockMcpServer.resource as any).mock.calls[0][2]; - wrappedResourceHandler('test-uri', { foo: 'bar' }); + const jsonRpcRequest = { + jsonrpc: '2.0', + method: 'tools/call', + id: 'req-1', + params: { name: 'get-weather', arguments: { location: 'Seattle, WA' } }, + }; - expect(tracingModule.startSpan).toHaveBeenCalledTimes(1); - expect(tracingModule.startSpan).toHaveBeenCalledWith( + const extraWithClientInfo = { + requestInfo: { + remoteAddress: '192.168.1.100', + remotePort: 54321, + }, + }; + + mockTransport.onmessage?.(jsonRpcRequest, extraWithClientInfo); + + expect(startSpanSpy).toHaveBeenCalledWith( { - name: `mcp-server/resource:${resourceName}`, + name: 'tools/call get-weather', 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, + 'mcp.method.name': 'tools/call', + 'mcp.tool.name': 'get-weather', + 'mcp.request.id': 'req-1', + 'mcp.session.id': 'test-session-123', + 'client.address': '192.168.1.100', + 'client.port': 54321, + 'mcp.transport': 'http', + 'network.transport': 'tcp', + 'network.protocol.version': '2.0', + 'mcp.request.argument.location': '"Seattle, WA"', + 'sentry.op': 'mcp.server', + 'sentry.origin': 'auto.function.mcp_server', + 'sentry.source': 'route', }, }, expect.any(Function), ); + }); - // Verify the original handler was called within the span - expect(mockResourceHandler).toHaveBeenCalledWith('test-uri', { foo: 'bar' }); + it('should create spans with correct attributes for resource operations', async () => { + await wrappedMcpServer.connect(mockTransport); + + const jsonRpcRequest = { + jsonrpc: '2.0', + method: 'resources/read', + id: 'req-2', + params: { uri: 'file:///docs/api.md' }, + }; + + mockTransport.onmessage?.(jsonRpcRequest, {}); + + expect(startSpanSpy).toHaveBeenCalledWith( + { + name: 'resources/read file:///docs/api.md', + forceTransaction: true, + attributes: { + 'mcp.method.name': 'resources/read', + 'mcp.resource.uri': 'file:///docs/api.md', + 'mcp.request.id': 'req-2', + 'mcp.session.id': 'test-session-123', + 'mcp.transport': 'http', + 'network.transport': 'tcp', + 'network.protocol.version': '2.0', + 'mcp.request.argument.uri': '"file:///docs/api.md"', + 'sentry.op': 'mcp.server', + 'sentry.origin': 'auto.function.mcp_server', + 'sentry.source': 'route', + }, + }, + expect.any(Function), + ); }); - 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(), - connect: vi.fn(), + it('should create spans with correct attributes for prompt operations', async () => { + await wrappedMcpServer.connect(mockTransport); + + const jsonRpcRequest = { + jsonrpc: '2.0', + method: 'prompts/get', + id: 'req-3', + params: { name: 'analyze-code' }, }; - const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer); + mockTransport.onmessage?.(jsonRpcRequest, {}); - // Call without string name - wrappedMcpServer.resource({} as any, 'handler'); + expect(startSpanSpy).toHaveBeenCalledWith( + { + name: 'prompts/get analyze-code', + forceTransaction: true, + attributes: { + 'mcp.method.name': 'prompts/get', + 'mcp.prompt.name': 'analyze-code', + 'mcp.request.id': 'req-3', + 'mcp.session.id': 'test-session-123', + 'mcp.transport': 'http', + 'network.transport': 'tcp', + 'network.protocol.version': '2.0', + 'mcp.request.argument.name': '"analyze-code"', + 'sentry.op': 'mcp.server', + 'sentry.origin': 'auto.function.mcp_server', + 'sentry.source': 'route', + }, + }, + expect.any(Function), + ); + }); + + it('should create spans with correct attributes for notifications (no request id)', async () => { + await wrappedMcpServer.connect(mockTransport); - // Call without function handler - wrappedMcpServer.resource('name', 'not-a-function'); + const jsonRpcNotification = { + jsonrpc: '2.0', + method: 'notifications/tools/list_changed', + params: {}, + }; + + mockTransport.onmessage?.(jsonRpcNotification, {}); + + expect(startSpanSpy).toHaveBeenCalledWith( + { + name: 'notifications/tools/list_changed', + forceTransaction: true, + attributes: { + 'mcp.method.name': 'notifications/tools/list_changed', + 'mcp.session.id': 'test-session-123', + 'mcp.transport': 'http', + 'network.transport': 'tcp', + 'network.protocol.version': '2.0', + 'sentry.op': 'mcp.notification.client_to_server', + 'sentry.origin': 'auto.mcp.notification', + 'sentry.source': 'route', + }, + }, + expect.any(Function), + ); - // Original method should be called directly without creating spans - expect(mockMcpServer.resource).toHaveBeenCalledTimes(2); - expect(tracingModule.startSpan).not.toHaveBeenCalled(); + // Should not include mcp.request.id for notifications + const callArgs = vi.mocked(tracingModule.startSpan).mock.calls[0]; + expect(callArgs).toBeDefined(); + const attributes = callArgs?.[0]?.attributes; + expect(attributes).not.toHaveProperty('mcp.request.id'); }); - }); - describe('tool method wrapping', () => { - it('should create a span with proper attributes when tool is called', () => { - const mockToolHandler = vi.fn(); - const toolName = 'test-tool'; + it('should create spans for list operations without target in name', async () => { + await wrappedMcpServer.connect(mockTransport); - const mockMcpServer = { - resource: vi.fn(), - tool: vi.fn(), - prompt: vi.fn(), - connect: vi.fn(), + const jsonRpcRequest = { + jsonrpc: '2.0', + method: 'tools/list', + id: 'req-4', + params: {}, }; - const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer); - wrappedMcpServer.tool(toolName, {}, mockToolHandler); + mockTransport.onmessage?.(jsonRpcRequest, {}); - // The original registration should use a wrapped handler - expect(mockMcpServer.tool).toHaveBeenCalledWith(toolName, {}, expect.any(Function)); + expect(startSpanSpy).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'tools/list', + forceTransaction: true, + attributes: expect.objectContaining({ + 'mcp.method.name': 'tools/list', + 'mcp.request.id': 'req-4', + 'mcp.session.id': 'test-session-123', + // Transport attributes + 'mcp.transport': 'http', + 'network.transport': 'tcp', + 'network.protocol.version': '2.0', + // Sentry-specific + 'sentry.op': 'mcp.server', + 'sentry.origin': 'auto.function.mcp_server', + 'sentry.source': 'route', + }), + }), + expect.any(Function), + ); + }); + + it('should create spans with logging attributes for notifications/message', async () => { + await wrappedMcpServer.connect(mockTransport); + + const loggingNotification = { + jsonrpc: '2.0', + method: 'notifications/message', + params: { + level: 'info', + logger: 'math-service', + data: 'Addition completed: 2 + 5 = 7', + }, + }; - // Invoke the wrapped handler to trigger Sentry span - const wrappedToolHandler = (mockMcpServer.tool as any).mock.calls[0][2]; - wrappedToolHandler({ arg: 'value' }, { foo: 'baz' }); + mockTransport.onmessage?.(loggingNotification, {}); - expect(tracingModule.startSpan).toHaveBeenCalledTimes(1); - expect(tracingModule.startSpan).toHaveBeenCalledWith( + expect(startSpanSpy).toHaveBeenCalledWith( { - name: `mcp-server/tool:${toolName}`, + name: 'notifications/message', 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, + 'mcp.method.name': 'notifications/message', + 'mcp.session.id': 'test-session-123', + 'mcp.transport': 'http', + 'network.transport': 'tcp', + 'network.protocol.version': '2.0', + 'mcp.logging.level': 'info', + 'mcp.logging.logger': 'math-service', + 'mcp.logging.data_type': 'string', + 'mcp.logging.message': 'Addition completed: 2 + 5 = 7', + 'sentry.op': 'mcp.notification.client_to_server', + 'sentry.origin': 'auto.mcp.notification', + 'sentry.source': 'route', }, }, expect.any(Function), ); - - // Verify the original handler was called within the span - expect(mockToolHandler).toHaveBeenCalledWith({ arg: 'value' }, { foo: 'baz' }); }); - 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(), - connect: vi.fn(), + it('should create spans with attributes for other notification types', async () => { + await wrappedMcpServer.connect(mockTransport); + + // Test notifications/cancelled + const cancelledNotification = { + jsonrpc: '2.0', + method: 'notifications/cancelled', + params: { + requestId: 'req-123', + reason: 'user_requested', + }, }; - const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer); + mockTransport.onmessage?.(cancelledNotification, {}); + + expect(startSpanSpy).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'notifications/cancelled', + attributes: expect.objectContaining({ + 'mcp.method.name': 'notifications/cancelled', + 'mcp.cancelled.request_id': 'req-123', + 'mcp.cancelled.reason': 'user_requested', + 'sentry.op': 'mcp.notification.client_to_server', + 'sentry.origin': 'auto.mcp.notification', + 'sentry.source': 'route', + }), + }), + expect.any(Function), + ); + + vi.clearAllMocks(); + + // Test notifications/progress + const progressNotification = { + jsonrpc: '2.0', + method: 'notifications/progress', + params: { + progressToken: 'token-456', + progress: 75, + total: 100, + message: 'Processing files...', + }, + }; + + mockTransport.onmessage?.(progressNotification, {}); + + expect(startSpanSpy).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'notifications/progress', + attributes: expect.objectContaining({ + 'mcp.method.name': 'notifications/progress', + 'mcp.progress.token': 'token-456', + 'mcp.progress.current': 75, + 'mcp.progress.total': 100, + 'mcp.progress.percentage': 75, + 'mcp.progress.message': 'Processing files...', + 'sentry.op': 'mcp.notification.client_to_server', + 'sentry.origin': 'auto.mcp.notification', + 'sentry.source': 'route', + }), + }), + expect.any(Function), + ); + + vi.clearAllMocks(); - // Call without string name - wrappedMcpServer.tool({} as any, 'handler'); + // Test notifications/resources/updated + const resourceUpdatedNotification = { + jsonrpc: '2.0', + method: 'notifications/resources/updated', + params: { + uri: 'file:///tmp/data.json', + }, + }; - // Original method should be called directly without creating spans - expect(mockMcpServer.tool).toHaveBeenCalledTimes(1); - expect(tracingModule.startSpan).not.toHaveBeenCalled(); + mockTransport.onmessage?.(resourceUpdatedNotification, {}); + + expect(startSpanSpy).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'notifications/resources/updated', + attributes: expect.objectContaining({ + 'mcp.method.name': 'notifications/resources/updated', + 'mcp.resource.uri': 'file:///tmp/data.json', + 'mcp.resource.protocol': 'file:', + 'sentry.op': 'mcp.notification.client_to_server', + 'sentry.origin': 'auto.mcp.notification', + 'sentry.source': 'route', + }), + }), + expect.any(Function), + ); }); - }); - describe('prompt method wrapping', () => { - it('should create a span with proper attributes when prompt is called', () => { - const mockPromptHandler = vi.fn(); - const promptName = 'test-prompt'; + it('should create spans with correct operation for outgoing notifications', async () => { + await wrappedMcpServer.connect(mockTransport); - const mockMcpServer = { - resource: vi.fn(), - tool: vi.fn(), - prompt: vi.fn(), - connect: vi.fn(), + const outgoingNotification = { + jsonrpc: '2.0', + method: 'notifications/tools/list_changed', }; - const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer); - wrappedMcpServer.prompt(promptName, {}, mockPromptHandler); + await mockTransport.send?.(outgoingNotification); + + expect(startSpanSpy).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'notifications/tools/list_changed', + attributes: expect.objectContaining({ + 'mcp.method.name': 'notifications/tools/list_changed', + 'sentry.op': 'mcp.notification.server_to_client', + 'sentry.origin': 'auto.mcp.notification', + 'sentry.source': 'route', + }), + }), + expect.any(Function), + ); + }); + }); - // The original registration should use a wrapped handler - expect(mockMcpServer.prompt).toHaveBeenCalledWith(promptName, {}, expect.any(Function)); + describe('Stdio Transport Tests', () => { + let mockMcpServer: ReturnType; + let wrappedMcpServer: ReturnType; + let mockStdioTransport: ReturnType; + + beforeEach(() => { + mockMcpServer = createMockMcpServer(); + wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer); + mockStdioTransport = createMockStdioTransport(); + mockStdioTransport.sessionId = 'stdio-session-456'; + }); - // Invoke the wrapped handler to trigger Sentry span - const wrappedPromptHandler = (mockMcpServer.prompt as any).mock.calls[0][2]; - wrappedPromptHandler({ msg: 'hello' }, { data: 123 }); + it('should detect stdio transport and set correct attributes', async () => { + await wrappedMcpServer.connect(mockStdioTransport); - expect(tracingModule.startSpan).toHaveBeenCalledTimes(1); - expect(tracingModule.startSpan).toHaveBeenCalledWith( + const jsonRpcRequest = { + jsonrpc: '2.0', + method: 'tools/call', + id: 'req-stdio-1', + params: { name: 'process-file', arguments: { path: '/tmp/data.txt' } }, + }; + + mockStdioTransport.onmessage?.(jsonRpcRequest, {}); + + expect(startSpanSpy).toHaveBeenCalledWith( { - name: `mcp-server/prompt:${promptName}`, + name: 'tools/call process-file', 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, + 'mcp.method.name': 'tools/call', + 'mcp.tool.name': 'process-file', + 'mcp.request.id': 'req-stdio-1', + 'mcp.session.id': 'stdio-session-456', + 'mcp.transport': 'stdio', // Should be stdio, not http + 'network.transport': 'pipe', // Should be pipe, not tcp + 'network.protocol.version': '2.0', + 'mcp.request.argument.path': '"/tmp/data.txt"', + 'sentry.op': 'mcp.server', + 'sentry.origin': 'auto.function.mcp_server', + 'sentry.source': 'route', }, }, expect.any(Function), ); - - // Verify the original handler was called within the span - expect(mockPromptHandler).toHaveBeenCalledWith({ msg: 'hello' }, { data: 123 }); }); - 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(), - connect: vi.fn(), + it('should handle stdio transport notifications correctly', async () => { + await wrappedMcpServer.connect(mockStdioTransport); + + const notification = { + jsonrpc: '2.0', + method: 'notifications/message', + params: { + level: 'debug', + data: 'Processing stdin input', + }, }; - const wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer); + mockStdioTransport.onmessage?.(notification, {}); + + expect(startSpanSpy).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'notifications/message', + attributes: expect.objectContaining({ + 'mcp.method.name': 'notifications/message', + 'mcp.session.id': 'stdio-session-456', + 'mcp.transport': 'stdio', + 'network.transport': 'pipe', + 'mcp.logging.level': 'debug', + 'mcp.logging.message': 'Processing stdin input', + }), + }), + expect.any(Function), + ); + }); + }); + + describe('SSE Transport Tests (Backwards Compatibility)', () => { + let mockMcpServer: ReturnType; + let wrappedMcpServer: ReturnType; + let mockSseTransport: ReturnType; - // Call without function handler - wrappedMcpServer.prompt('name', 'not-a-function'); + beforeEach(() => { + mockMcpServer = createMockMcpServer(); + wrappedMcpServer = wrapMcpServerWithSentry(mockMcpServer); + mockSseTransport = createMockSseTransport(); + mockSseTransport.sessionId = 'sse-session-789'; + }); - // Original method should be called directly without creating spans - expect(mockMcpServer.prompt).toHaveBeenCalledTimes(1); - expect(tracingModule.startSpan).not.toHaveBeenCalled(); + it('should detect SSE transport for backwards compatibility', async () => { + await wrappedMcpServer.connect(mockSseTransport); + + const jsonRpcRequest = { + jsonrpc: '2.0', + method: 'resources/read', + id: 'req-sse-1', + params: { uri: 'https://api.example.com/data' }, + }; + + mockSseTransport.onmessage?.(jsonRpcRequest, {}); + + expect(startSpanSpy).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'resources/read https://api.example.com/data', + attributes: expect.objectContaining({ + 'mcp.method.name': 'resources/read', + 'mcp.resource.uri': 'https://api.example.com/data', + 'mcp.transport': 'sse', // Deprecated but supported + 'network.transport': 'tcp', + 'mcp.session.id': 'sse-session-789', + }), + }), + expect.any(Function), + ); }); }); }); + +// Test helpers +function createMockMcpServer() { + return { + resource: vi.fn(), + tool: vi.fn(), + prompt: vi.fn(), + connect: vi.fn().mockResolvedValue(undefined), + server: { + setRequestHandler: vi.fn(), + }, + }; +} + +function createMockTransport() { + // exact naming pattern from the official SDK + class StreamableHTTPServerTransport { + onmessage = vi.fn(); + onclose = vi.fn(); + send = vi.fn().mockResolvedValue(undefined); + sessionId = 'test-session-123'; + } + + return new StreamableHTTPServerTransport(); +} + +function createMockStdioTransport() { + // Create a mock that mimics StdioServerTransport + // Using the exact naming pattern from the official SDK + class StdioServerTransport { + onmessage = vi.fn(); + onclose = vi.fn(); + send = vi.fn().mockResolvedValue(undefined); + sessionId = 'stdio-session-456'; + } + + return new StdioServerTransport(); +} + +function createMockSseTransport() { + // Create a mock that mimics the deprecated SSEServerTransport + // For backwards compatibility testing + class SSEServerTransport { + onmessage = vi.fn(); + onclose = vi.fn(); + send = vi.fn().mockResolvedValue(undefined); + sessionId = 'sse-session-789'; + } + + return new SSEServerTransport(); +}