diff --git a/packages/core/src/integrations/mcp-server/attributeExtraction.ts b/packages/core/src/integrations/mcp-server/attributeExtraction.ts index 68eade987a08..8f1e5a77d94d 100644 --- a/packages/core/src/integrations/mcp-server/attributeExtraction.ts +++ b/packages/core/src/integrations/mcp-server/attributeExtraction.ts @@ -1,66 +1,18 @@ /** - * Attribute extraction and building functions for MCP server instrumentation + * Core attribute extraction and building functions for MCP server instrumentation */ import { isURLObjectRelative, parseStringToURLObject } from '../../utils/url'; import { - CLIENT_ADDRESS_ATTRIBUTE, - CLIENT_PORT_ATTRIBUTE, MCP_LOGGING_DATA_TYPE_ATTRIBUTE, MCP_LOGGING_LEVEL_ATTRIBUTE, MCP_LOGGING_LOGGER_ATTRIBUTE, MCP_LOGGING_MESSAGE_ATTRIBUTE, - MCP_PROTOCOL_VERSION_ATTRIBUTE, MCP_REQUEST_ID_ATTRIBUTE, MCP_RESOURCE_URI_ATTRIBUTE, - MCP_SERVER_NAME_ATTRIBUTE, - MCP_SERVER_TITLE_ATTRIBUTE, - MCP_SERVER_VERSION_ATTRIBUTE, - MCP_SESSION_ID_ATTRIBUTE, - MCP_TOOL_RESULT_CONTENT_COUNT_ATTRIBUTE, - MCP_TOOL_RESULT_IS_ERROR_ATTRIBUTE, - MCP_TRANSPORT_ATTRIBUTE, - NETWORK_PROTOCOL_VERSION_ATTRIBUTE, - NETWORK_TRANSPORT_ATTRIBUTE, } from './attributes'; import { extractTargetInfo, getRequestArguments } from './methodConfig'; -import { - getClientInfoForTransport, - getProtocolVersionForTransport, - getSessionDataForTransport, -} from './sessionManagement'; -import type { - ExtraHandlerData, - JsonRpcNotification, - JsonRpcRequest, - McpSpanType, - MCPTransport, - PartyInfo, - SessionData, -} from './types'; - -/** - * Extracts transport types based on transport constructor name - * @param transport - MCP transport instance - * @returns Transport type mapping for span attributes - */ -export function getTransportTypes(transport: MCPTransport): { mcpTransport: string; networkTransport: string } { - const transportName = transport.constructor?.name?.toLowerCase() || ''; - - if (transportName.includes('stdio')) { - return { mcpTransport: 'stdio', networkTransport: 'pipe' }; - } - - if (transportName.includes('streamablehttp') || transportName.includes('streamable')) { - return { mcpTransport: 'http', networkTransport: 'tcp' }; - } - - if (transportName.includes('sse')) { - return { mcpTransport: 'sse', networkTransport: 'tcp' }; - } - - return { mcpTransport: 'unknown', networkTransport: 'unknown' }; -} +import type { JsonRpcNotification, JsonRpcRequest, McpSpanType } from './types'; /** * Extracts additional attributes for specific notification types @@ -138,155 +90,6 @@ export function getNotificationAttributes( return attributes; } -/** - * Extracts and validates PartyInfo from an unknown object - * @param obj - Unknown object that might contain party info - * @returns Validated PartyInfo object with only string properties - */ -function extractPartyInfo(obj: unknown): PartyInfo { - const partyInfo: PartyInfo = {}; - - if (obj && typeof obj === 'object' && obj !== null) { - const source = obj as Record; - if (typeof source.name === 'string') partyInfo.name = source.name; - if (typeof source.title === 'string') partyInfo.title = source.title; - if (typeof source.version === 'string') partyInfo.version = source.version; - } - - return partyInfo; -} - -/** - * Extracts session data from "initialize" requests - * @param request - JSON-RPC "initialize" request containing client info and protocol version - * @returns Session data extracted from request parameters including protocol version and client info - */ -export function extractSessionDataFromInitializeRequest(request: JsonRpcRequest): SessionData { - const sessionData: SessionData = {}; - if (request.params && typeof request.params === 'object' && request.params !== null) { - const params = request.params as Record; - if (typeof params.protocolVersion === 'string') { - sessionData.protocolVersion = params.protocolVersion; - } - if (params.clientInfo) { - sessionData.clientInfo = extractPartyInfo(params.clientInfo); - } - } - return sessionData; -} - -/** - * Extracts session data from "initialize" response - * @param result - "initialize" response result containing server info and protocol version - * @returns Partial session data extracted from response including protocol version and server info - */ -export function extractSessionDataFromInitializeResponse(result: unknown): Partial { - const sessionData: Partial = {}; - if (result && typeof result === 'object') { - const resultObj = result as Record; - if (typeof resultObj.protocolVersion === 'string') sessionData.protocolVersion = resultObj.protocolVersion; - if (resultObj.serverInfo) { - sessionData.serverInfo = extractPartyInfo(resultObj.serverInfo); - } - } - return sessionData; -} - -/** - * Build client attributes from stored client info - * @param transport - MCP transport instance - * @returns Client attributes for span instrumentation - */ -export function getClientAttributes(transport: MCPTransport): Record { - const clientInfo = getClientInfoForTransport(transport); - const attributes: Record = {}; - - if (clientInfo?.name) { - attributes['mcp.client.name'] = clientInfo.name; - } - if (clientInfo?.title) { - attributes['mcp.client.title'] = clientInfo.title; - } - if (clientInfo?.version) { - attributes['mcp.client.version'] = clientInfo.version; - } - - return attributes; -} - -/** - * Build server attributes from stored server info - * @param transport - MCP transport instance - * @returns Server attributes for span instrumentation - */ -export function getServerAttributes(transport: MCPTransport): Record { - const serverInfo = getSessionDataForTransport(transport)?.serverInfo; - const attributes: Record = {}; - - if (serverInfo?.name) { - attributes[MCP_SERVER_NAME_ATTRIBUTE] = serverInfo.name; - } - if (serverInfo?.title) { - attributes[MCP_SERVER_TITLE_ATTRIBUTE] = serverInfo.title; - } - if (serverInfo?.version) { - attributes[MCP_SERVER_VERSION_ATTRIBUTE] = serverInfo.version; - } - - return attributes; -} - -/** - * Extracts client connection info from extra handler data - * @param extra - Extra handler data containing connection info - * @returns Client address and port information - */ -export 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, - }; -} - -/** - * Build transport and network attributes - * @param transport - MCP transport instance - * @param extra - Optional extra handler data - * @returns Transport attributes for span instrumentation - */ -export function buildTransportAttributes( - transport: MCPTransport, - extra?: ExtraHandlerData, -): Record { - const sessionId = transport.sessionId; - const clientInfo = extra ? extractClientInfo(extra) : {}; - const { mcpTransport, networkTransport } = getTransportTypes(transport); - const clientAttributes = getClientAttributes(transport); - const serverAttributes = getServerAttributes(transport); - const protocolVersion = getProtocolVersionForTransport(transport); - - const attributes = { - ...(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', - ...(protocolVersion && { [MCP_PROTOCOL_VERSION_ATTRIBUTE]: protocolVersion }), - ...clientAttributes, - ...serverAttributes, - }; - - return attributes; -} - /** * Build type-specific attributes based on message type * @param type - Span type (request or notification) @@ -313,67 +116,5 @@ export function buildTypeSpecificAttributes( return getNotificationAttributes(message.method, params || {}); } -/** - * Build attributes for tool result content items - * @param content - Array of content items from tool result - * @returns Attributes extracted from each content item including type, text, mime type, URI, and resource info - */ -function buildAllContentItemAttributes(content: unknown[]): Record { - const attributes: Record = { - [MCP_TOOL_RESULT_CONTENT_COUNT_ATTRIBUTE]: content.length, - }; - - for (const [i, item] of content.entries()) { - if (typeof item !== 'object' || item === null) continue; - - const contentItem = item as Record; - const prefix = content.length === 1 ? 'mcp.tool.result' : `mcp.tool.result.${i}`; - - const safeSet = (key: string, value: unknown): void => { - if (typeof value === 'string') attributes[`${prefix}.${key}`] = value; - }; - - safeSet('content_type', contentItem.type); - safeSet('mime_type', contentItem.mimeType); - safeSet('uri', contentItem.uri); - safeSet('name', contentItem.name); - - if (typeof contentItem.text === 'string') { - const text = contentItem.text; - const maxLength = 500; - attributes[`${prefix}.content`] = text.length > maxLength ? `${text.slice(0, maxLength - 3)}...` : text; - } - - if (typeof contentItem.data === 'string') { - attributes[`${prefix}.data_size`] = contentItem.data.length; - } - - const resource = contentItem.resource; - if (typeof resource === 'object' && resource !== null) { - const res = resource as Record; - safeSet('resource_uri', res.uri); - safeSet('resource_mime_type', res.mimeType); - } - } - - return attributes; -} - -/** - * Extract tool result attributes for span instrumentation - * @param result - Tool execution result - * @returns Attributes extracted from tool result content - */ -export function extractToolResultAttributes(result: unknown): Record { - let attributes: Record = {}; - if (typeof result !== 'object' || result === null) return attributes; - - const resultObj = result as Record; - if (typeof resultObj.isError === 'boolean') { - attributes[MCP_TOOL_RESULT_IS_ERROR_ATTRIBUTE] = resultObj.isError; - } - if (Array.isArray(resultObj.content)) { - attributes = { ...attributes, ...buildAllContentItemAttributes(resultObj.content) }; - } - return attributes; -} +// Re-export buildTransportAttributes for spans.ts +export { buildTransportAttributes } from './sessionExtraction'; diff --git a/packages/core/src/integrations/mcp-server/attributes.ts b/packages/core/src/integrations/mcp-server/attributes.ts index 074bd09b7bdf..273bdbdb9560 100644 --- a/packages/core/src/integrations/mcp-server/attributes.ts +++ b/packages/core/src/integrations/mcp-server/attributes.ts @@ -76,6 +76,28 @@ export const MCP_TOOL_RESULT_CONTENT_COUNT_ATTRIBUTE = 'mcp.tool.result.content_ /** Serialized content of the tool result */ export const MCP_TOOL_RESULT_CONTENT_ATTRIBUTE = 'mcp.tool.result.content'; +/** Prefix for tool result attributes that contain sensitive content */ +export const MCP_TOOL_RESULT_PREFIX = 'mcp.tool.result'; + +// ============================================================================= +// PROMPT RESULT ATTRIBUTES +// ============================================================================= + +/** Description of the prompt result */ +export const MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE = 'mcp.prompt.result.description'; + +/** Number of messages in the prompt result */ +export const MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE = 'mcp.prompt.result.message_count'; + +/** Role of the message in the prompt result (for single message results) */ +export const MCP_PROMPT_RESULT_MESSAGE_ROLE_ATTRIBUTE = 'mcp.prompt.result.message_role'; + +/** Content of the message in the prompt result (for single message results) */ +export const MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE = 'mcp.prompt.result.message_content'; + +/** Prefix for prompt result attributes that contain sensitive content */ +export const MCP_PROMPT_RESULT_PREFIX = 'mcp.prompt.result'; + // ============================================================================= // REQUEST ARGUMENT ATTRIBUTES // ============================================================================= diff --git a/packages/core/src/integrations/mcp-server/correlation.ts b/packages/core/src/integrations/mcp-server/correlation.ts index 7f00341bdd5a..7a73f63f64e3 100644 --- a/packages/core/src/integrations/mcp-server/correlation.ts +++ b/packages/core/src/integrations/mcp-server/correlation.ts @@ -9,8 +9,8 @@ import { getClient } from '../../currentScopes'; import { SPAN_STATUS_ERROR } from '../../tracing'; import type { Span } from '../../types-hoist/span'; -import { extractToolResultAttributes } from './attributeExtraction'; import { filterMcpPiiFromSpanData } from './piiFiltering'; +import { extractPromptResultAttributes, extractToolResultAttributes } from './resultExtraction'; import type { MCPTransport, RequestId, RequestSpanMapValue } from './types'; /** @@ -69,6 +69,13 @@ export function completeSpanWithResults(transport: MCPTransport, requestId: Requ const toolAttributes = filterMcpPiiFromSpanData(rawToolAttributes, sendDefaultPii); span.setAttributes(toolAttributes); + } else if (method === 'prompts/get') { + const rawPromptAttributes = extractPromptResultAttributes(result); + const client = getClient(); + const sendDefaultPii = Boolean(client?.getOptions().sendDefaultPii); + const promptAttributes = filterMcpPiiFromSpanData(rawPromptAttributes, sendDefaultPii); + + span.setAttributes(promptAttributes); } span.end(); @@ -83,7 +90,9 @@ export function completeSpanWithResults(transport: MCPTransport, requestId: Requ */ export function cleanupPendingSpansForTransport(transport: MCPTransport): number { const spanMap = transportToSpanMap.get(transport); - if (!spanMap) return 0; + if (!spanMap) { + return 0; + } const pendingCount = spanMap.size; diff --git a/packages/core/src/integrations/mcp-server/piiFiltering.ts b/packages/core/src/integrations/mcp-server/piiFiltering.ts index 654427ca2d6d..ff801cbf2a1e 100644 --- a/packages/core/src/integrations/mcp-server/piiFiltering.ts +++ b/packages/core/src/integrations/mcp-server/piiFiltering.ts @@ -9,9 +9,13 @@ import { CLIENT_ADDRESS_ATTRIBUTE, CLIENT_PORT_ATTRIBUTE, MCP_LOGGING_MESSAGE_ATTRIBUTE, + MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE, + MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE, + MCP_PROMPT_RESULT_PREFIX, MCP_REQUEST_ARGUMENT, MCP_RESOURCE_URI_ATTRIBUTE, MCP_TOOL_RESULT_CONTENT_ATTRIBUTE, + MCP_TOOL_RESULT_PREFIX, } from './attributes'; /** @@ -22,16 +26,42 @@ const PII_ATTRIBUTES = new Set([ CLIENT_ADDRESS_ATTRIBUTE, CLIENT_PORT_ATTRIBUTE, MCP_LOGGING_MESSAGE_ATTRIBUTE, + MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE, + MCP_PROMPT_RESULT_MESSAGE_CONTENT_ATTRIBUTE, MCP_RESOURCE_URI_ATTRIBUTE, MCP_TOOL_RESULT_CONTENT_ATTRIBUTE, ]); /** - * Checks if an attribute key should be considered PII + * Checks if an attribute key should be considered PII. + * + * Returns true for: + * - Explicit PII attributes (client.address, client.port, mcp.logging.message, etc.) + * - All request arguments (mcp.request.argument.*) + * - Tool and prompt result content (mcp.tool.result.*, mcp.prompt.result.*) except metadata + * + * Preserves metadata attributes ending with _count, _error, or .is_error as they don't contain sensitive data. + * + * @param key - Attribute key to evaluate + * @returns true if the attribute should be filtered out (is PII), false if it should be preserved * @internal */ function isPiiAttribute(key: string): boolean { - return PII_ATTRIBUTES.has(key) || key.startsWith(`${MCP_REQUEST_ARGUMENT}.`); + if (PII_ATTRIBUTES.has(key)) { + return true; + } + + if (key.startsWith(`${MCP_REQUEST_ARGUMENT}.`)) { + return true; + } + + if (key.startsWith(`${MCP_TOOL_RESULT_PREFIX}.`) || key.startsWith(`${MCP_PROMPT_RESULT_PREFIX}.`)) { + if (!key.endsWith('_count') && !key.endsWith('_error') && !key.endsWith('.is_error')) { + return true; + } + } + + return false; } /** diff --git a/packages/core/src/integrations/mcp-server/resultExtraction.ts b/packages/core/src/integrations/mcp-server/resultExtraction.ts new file mode 100644 index 000000000000..34dc2be9d09c --- /dev/null +++ b/packages/core/src/integrations/mcp-server/resultExtraction.ts @@ -0,0 +1,126 @@ +/** + * Result extraction functions for MCP server instrumentation + * + * Handles extraction of attributes from tool and prompt execution results. + */ + +import { + MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE, + MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE, + MCP_TOOL_RESULT_CONTENT_COUNT_ATTRIBUTE, + MCP_TOOL_RESULT_IS_ERROR_ATTRIBUTE, +} from './attributes'; +import { isValidContentItem } from './validation'; + +/** + * Build attributes for tool result content items + * @param content - Array of content items from tool result + * @returns Attributes extracted from each content item including type, text, mime type, URI, and resource info + */ +function buildAllContentItemAttributes(content: unknown[]): Record { + const attributes: Record = { + [MCP_TOOL_RESULT_CONTENT_COUNT_ATTRIBUTE]: content.length, + }; + + for (const [i, item] of content.entries()) { + if (!isValidContentItem(item)) { + continue; + } + + const prefix = content.length === 1 ? 'mcp.tool.result' : `mcp.tool.result.${i}`; + + const safeSet = (key: string, value: unknown): void => { + if (typeof value === 'string') { + attributes[`${prefix}.${key}`] = value; + } + }; + + safeSet('content_type', item.type); + safeSet('mime_type', item.mimeType); + safeSet('uri', item.uri); + safeSet('name', item.name); + + if (typeof item.text === 'string') { + attributes[`${prefix}.content`] = item.text; + } + + if (typeof item.data === 'string') { + attributes[`${prefix}.data_size`] = item.data.length; + } + + const resource = item.resource; + if (isValidContentItem(resource)) { + safeSet('resource_uri', resource.uri); + safeSet('resource_mime_type', resource.mimeType); + } + } + + return attributes; +} + +/** + * Extract tool result attributes for span instrumentation + * @param result - Tool execution result + * @returns Attributes extracted from tool result content + */ +export function extractToolResultAttributes(result: unknown): Record { + if (!isValidContentItem(result)) { + return {}; + } + + const attributes = Array.isArray(result.content) ? buildAllContentItemAttributes(result.content) : {}; + + if (typeof result.isError === 'boolean') { + attributes[MCP_TOOL_RESULT_IS_ERROR_ATTRIBUTE] = result.isError; + } + + return attributes; +} + +/** + * Extract prompt result attributes for span instrumentation + * @param result - Prompt execution result + * @returns Attributes extracted from prompt result + */ +export function extractPromptResultAttributes(result: unknown): Record { + const attributes: Record = {}; + if (!isValidContentItem(result)) { + return attributes; + } + + if (typeof result.description === 'string') { + attributes[MCP_PROMPT_RESULT_DESCRIPTION_ATTRIBUTE] = result.description; + } + + if (Array.isArray(result.messages)) { + attributes[MCP_PROMPT_RESULT_MESSAGE_COUNT_ATTRIBUTE] = result.messages.length; + + const messages = result.messages; + for (const [i, message] of messages.entries()) { + if (!isValidContentItem(message)) { + continue; + } + + const prefix = messages.length === 1 ? 'mcp.prompt.result' : `mcp.prompt.result.${i}`; + + const safeSet = (key: string, value: unknown): void => { + if (typeof value === 'string') { + const attrName = messages.length === 1 ? `${prefix}.message_${key}` : `${prefix}.${key}`; + attributes[attrName] = value; + } + }; + + safeSet('role', message.role); + + if (isValidContentItem(message.content)) { + const content = message.content; + if (typeof content.text === 'string') { + const attrName = messages.length === 1 ? `${prefix}.message_content` : `${prefix}.content`; + attributes[attrName] = content.text; + } + } + } + } + + return attributes; +} diff --git a/packages/core/src/integrations/mcp-server/sessionExtraction.ts b/packages/core/src/integrations/mcp-server/sessionExtraction.ts new file mode 100644 index 000000000000..90e235d4e544 --- /dev/null +++ b/packages/core/src/integrations/mcp-server/sessionExtraction.ts @@ -0,0 +1,202 @@ +/** + * Session and party info extraction functions for MCP server instrumentation + * + * Handles extraction of client/server info and session data from MCP messages. + */ + +import { + CLIENT_ADDRESS_ATTRIBUTE, + CLIENT_PORT_ATTRIBUTE, + MCP_PROTOCOL_VERSION_ATTRIBUTE, + MCP_SERVER_NAME_ATTRIBUTE, + MCP_SERVER_TITLE_ATTRIBUTE, + MCP_SERVER_VERSION_ATTRIBUTE, + MCP_SESSION_ID_ATTRIBUTE, + MCP_TRANSPORT_ATTRIBUTE, + NETWORK_PROTOCOL_VERSION_ATTRIBUTE, + NETWORK_TRANSPORT_ATTRIBUTE, +} from './attributes'; +import { + getClientInfoForTransport, + getProtocolVersionForTransport, + getSessionDataForTransport, +} from './sessionManagement'; +import type { ExtraHandlerData, JsonRpcRequest, MCPTransport, PartyInfo, SessionData } from './types'; +import { isValidContentItem } from './validation'; + +/** + * Extracts and validates PartyInfo from an unknown object + * @param obj - Unknown object that might contain party info + * @returns Validated PartyInfo object with only string properties + */ +function extractPartyInfo(obj: unknown): PartyInfo { + const partyInfo: PartyInfo = {}; + + if (isValidContentItem(obj)) { + if (typeof obj.name === 'string') { + partyInfo.name = obj.name; + } + if (typeof obj.title === 'string') { + partyInfo.title = obj.title; + } + if (typeof obj.version === 'string') { + partyInfo.version = obj.version; + } + } + + return partyInfo; +} + +/** + * Extracts session data from "initialize" requests + * @param request - JSON-RPC "initialize" request containing client info and protocol version + * @returns Session data extracted from request parameters including protocol version and client info + */ +export function extractSessionDataFromInitializeRequest(request: JsonRpcRequest): SessionData { + const sessionData: SessionData = {}; + if (isValidContentItem(request.params)) { + if (typeof request.params.protocolVersion === 'string') { + sessionData.protocolVersion = request.params.protocolVersion; + } + if (request.params.clientInfo) { + sessionData.clientInfo = extractPartyInfo(request.params.clientInfo); + } + } + return sessionData; +} + +/** + * Extracts session data from "initialize" response + * @param result - "initialize" response result containing server info and protocol version + * @returns Partial session data extracted from response including protocol version and server info + */ +export function extractSessionDataFromInitializeResponse(result: unknown): Partial { + const sessionData: Partial = {}; + if (isValidContentItem(result)) { + if (typeof result.protocolVersion === 'string') { + sessionData.protocolVersion = result.protocolVersion; + } + if (result.serverInfo) { + sessionData.serverInfo = extractPartyInfo(result.serverInfo); + } + } + return sessionData; +} + +/** + * Build client attributes from stored client info + * @param transport - MCP transport instance + * @returns Client attributes for span instrumentation + */ +export function getClientAttributes(transport: MCPTransport): Record { + const clientInfo = getClientInfoForTransport(transport); + const attributes: Record = {}; + + if (clientInfo?.name) { + attributes['mcp.client.name'] = clientInfo.name; + } + if (clientInfo?.title) { + attributes['mcp.client.title'] = clientInfo.title; + } + if (clientInfo?.version) { + attributes['mcp.client.version'] = clientInfo.version; + } + + return attributes; +} + +/** + * Build server attributes from stored server info + * @param transport - MCP transport instance + * @returns Server attributes for span instrumentation + */ +export function getServerAttributes(transport: MCPTransport): Record { + const serverInfo = getSessionDataForTransport(transport)?.serverInfo; + const attributes: Record = {}; + + if (serverInfo?.name) { + attributes[MCP_SERVER_NAME_ATTRIBUTE] = serverInfo.name; + } + if (serverInfo?.title) { + attributes[MCP_SERVER_TITLE_ATTRIBUTE] = serverInfo.title; + } + if (serverInfo?.version) { + attributes[MCP_SERVER_VERSION_ATTRIBUTE] = serverInfo.version; + } + + return attributes; +} + +/** + * Extracts client connection info from extra handler data + * @param extra - Extra handler data containing connection info + * @returns Client address and port information + */ +export 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, + }; +} + +/** + * Extracts transport types based on transport constructor name + * @param transport - MCP transport instance + * @returns Transport type mapping for span attributes + */ +export function getTransportTypes(transport: MCPTransport): { mcpTransport: string; networkTransport: string } { + const transportName = transport.constructor?.name?.toLowerCase() || ''; + + if (transportName.includes('stdio')) { + return { mcpTransport: 'stdio', networkTransport: 'pipe' }; + } + + if (transportName.includes('streamablehttp') || transportName.includes('streamable')) { + return { mcpTransport: 'http', networkTransport: 'tcp' }; + } + + if (transportName.includes('sse')) { + return { mcpTransport: 'sse', networkTransport: 'tcp' }; + } + + return { mcpTransport: 'unknown', networkTransport: 'unknown' }; +} + +/** + * Build transport and network attributes + * @param transport - MCP transport instance + * @param extra - Optional extra handler data + * @returns Transport attributes for span instrumentation + */ +export function buildTransportAttributes( + transport: MCPTransport, + extra?: ExtraHandlerData, +): Record { + const sessionId = transport.sessionId; + const clientInfo = extra ? extractClientInfo(extra) : {}; + const { mcpTransport, networkTransport } = getTransportTypes(transport); + const clientAttributes = getClientAttributes(transport); + const serverAttributes = getServerAttributes(transport); + const protocolVersion = getProtocolVersionForTransport(transport); + + const attributes = { + ...(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', + ...(protocolVersion && { [MCP_PROTOCOL_VERSION_ATTRIBUTE]: protocolVersion }), + ...clientAttributes, + ...serverAttributes, + }; + + return attributes; +} diff --git a/packages/core/src/integrations/mcp-server/sessionManagement.ts b/packages/core/src/integrations/mcp-server/sessionManagement.ts index 99ba2e0d8806..9d9c8b48f27d 100644 --- a/packages/core/src/integrations/mcp-server/sessionManagement.ts +++ b/packages/core/src/integrations/mcp-server/sessionManagement.ts @@ -16,7 +16,9 @@ const transportToSessionData = new WeakMap(); * @param sessionData - Session data to store */ export function storeSessionDataForTransport(transport: MCPTransport, sessionData: SessionData): void { - if (transport.sessionId) transportToSessionData.set(transport, sessionData); + if (transport.sessionId) { + transportToSessionData.set(transport, sessionData); + } } /** diff --git a/packages/core/src/integrations/mcp-server/transport.ts b/packages/core/src/integrations/mcp-server/transport.ts index 3244ce73e49a..6943ac3e8850 100644 --- a/packages/core/src/integrations/mcp-server/transport.ts +++ b/packages/core/src/integrations/mcp-server/transport.ts @@ -8,12 +8,9 @@ import { getIsolationScope, withIsolationScope } from '../../currentScopes'; import { startInactiveSpan, withActiveSpan } from '../../tracing'; import { fill } from '../../utils/object'; -import { - extractSessionDataFromInitializeRequest, - extractSessionDataFromInitializeResponse, -} from './attributeExtraction'; import { cleanupPendingSpansForTransport, completeSpanWithResults, storeSpanForRequest } from './correlation'; import { captureError } from './errorCapture'; +import { extractSessionDataFromInitializeRequest, extractSessionDataFromInitializeResponse } from './sessionExtraction'; import { cleanupSessionDataForTransport, storeSessionDataForTransport, @@ -21,7 +18,7 @@ import { } from './sessionManagement'; import { buildMcpServerSpanConfig, createMcpNotificationSpan, createMcpOutgoingNotificationSpan } from './spans'; import type { ExtraHandlerData, MCPTransport } from './types'; -import { isJsonRpcNotification, isJsonRpcRequest, isJsonRpcResponse } from './validation'; +import { isJsonRpcNotification, isJsonRpcRequest, isJsonRpcResponse, isValidContentItem } from './validation'; /** * Wraps transport.onmessage to create spans for incoming messages. @@ -93,9 +90,8 @@ export function wrapTransportSend(transport: MCPTransport): void { captureJsonRpcErrorResponse(message.error); } - if (message.result && typeof message.result === 'object') { - const result = message.result as Record; - if (result.protocolVersion || result.serverInfo) { + if (isValidContentItem(message.result)) { + if (message.result.protocolVersion || message.result.serverInfo) { try { const serverData = extractSessionDataFromInitializeResponse(message.result); updateSessionDataForTransport(this, serverData); diff --git a/packages/core/src/integrations/mcp-server/validation.ts b/packages/core/src/integrations/mcp-server/validation.ts index 21d257c01aeb..9ed21b290728 100644 --- a/packages/core/src/integrations/mcp-server/validation.ts +++ b/packages/core/src/integrations/mcp-server/validation.ts @@ -75,3 +75,12 @@ export function validateMcpServerInstance(instance: unknown): boolean { DEBUG_BUILD && debug.warn('Did not patch MCP server. Interface is incompatible.'); return false; } + +/** + * Check if the item is a valid content item + * @param item - The item to check + * @returns True if the item is a valid content item, false otherwise + */ +export function isValidContentItem(item: unknown): item is Record { + return item != null && typeof item === 'object'; +} diff --git a/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts b/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts index 14f803b28ccc..a86ccbd534d0 100644 --- a/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts +++ b/packages/core/test/lib/integrations/mcp-server/piiFiltering.test.ts @@ -124,7 +124,7 @@ describe('MCP Server PII Filtering', () => { setAttributes: vi.fn(), setStatus: vi.fn(), end: vi.fn(), - } as any; + } as unknown as ReturnType; startInactiveSpanSpy.mockReturnValueOnce(mockSpan); const toolCallRequest = { @@ -147,7 +147,7 @@ describe('MCP Server PII Filtering', () => { mockTransport.send?.(toolResponse); - // Tool result content should be filtered out + // Tool result content should be filtered out, but metadata should remain const setAttributesCall = mockSpan.setAttributes.mock.calls[0]?.[0]; expect(setAttributesCall).toBeDefined(); expect(setAttributesCall).not.toHaveProperty('mcp.tool.result.content'); @@ -163,6 +163,11 @@ describe('MCP Server PII Filtering', () => { 'client.port': 54321, 'mcp.request.argument.location': '"San Francisco"', 'mcp.tool.result.content': 'Weather data: 18°C', + 'mcp.tool.result.content_count': 1, + 'mcp.prompt.result.description': 'Code review prompt for sensitive analysis', + 'mcp.prompt.result.message_content': 'Please review this confidential code.', + 'mcp.prompt.result.message_count': 1, + 'mcp.resource.result.content': 'Sensitive resource content', 'mcp.logging.message': 'User requested weather', 'mcp.resource.uri': 'file:///private/docs/secret.txt', 'mcp.method.name': 'tools/call', // Non-PII should remain @@ -180,6 +185,16 @@ describe('MCP Server PII Filtering', () => { 'mcp.request.argument.location': '"San Francisco"', 'mcp.request.argument.units': '"celsius"', 'mcp.tool.result.content': 'Weather data: 18°C', + 'mcp.tool.result.content_count': 1, + 'mcp.prompt.result.description': 'Code review prompt for sensitive analysis', + 'mcp.prompt.result.message_count': 2, + 'mcp.prompt.result.0.role': 'user', + 'mcp.prompt.result.0.content': 'Sensitive prompt content', + 'mcp.prompt.result.1.role': 'assistant', + 'mcp.prompt.result.1.content': 'Another sensitive response', + 'mcp.resource.result.content_count': 1, + 'mcp.resource.result.uri': 'file:///private/file.txt', + 'mcp.resource.result.content': 'Sensitive resource content', 'mcp.logging.message': 'User requested weather', 'mcp.resource.uri': 'file:///private/docs/secret.txt', 'mcp.method.name': 'tools/call', // Non-PII should remain @@ -188,14 +203,37 @@ describe('MCP Server PII Filtering', () => { const result = filterMcpPiiFromSpanData(spanData, false); + // Client info should be filtered expect(result).not.toHaveProperty('client.address'); expect(result).not.toHaveProperty('client.port'); + + // Request arguments should be filtered expect(result).not.toHaveProperty('mcp.request.argument.location'); expect(result).not.toHaveProperty('mcp.request.argument.units'); + + // Specific PII content attributes should be filtered expect(result).not.toHaveProperty('mcp.tool.result.content'); + expect(result).not.toHaveProperty('mcp.prompt.result.description'); + + // Count attributes should remain as they don't contain sensitive content + expect(result).toHaveProperty('mcp.tool.result.content_count', 1); + expect(result).toHaveProperty('mcp.prompt.result.message_count', 2); + + // All tool and prompt result content should be filtered (including indexed attributes) + expect(result).not.toHaveProperty('mcp.prompt.result.0.role'); + expect(result).not.toHaveProperty('mcp.prompt.result.0.content'); + expect(result).not.toHaveProperty('mcp.prompt.result.1.role'); + expect(result).not.toHaveProperty('mcp.prompt.result.1.content'); + + expect(result).toHaveProperty('mcp.resource.result.content_count', 1); + expect(result).toHaveProperty('mcp.resource.result.uri', 'file:///private/file.txt'); + expect(result).toHaveProperty('mcp.resource.result.content', 'Sensitive resource content'); + + // Other PII attributes should be filtered expect(result).not.toHaveProperty('mcp.logging.message'); expect(result).not.toHaveProperty('mcp.resource.uri'); + // Non-PII attributes should remain expect(result).toHaveProperty('mcp.method.name', 'tools/call'); expect(result).toHaveProperty('mcp.session.id', 'test-session-123'); }); diff --git a/packages/core/test/lib/integrations/mcp-server/semanticConventions.test.ts b/packages/core/test/lib/integrations/mcp-server/semanticConventions.test.ts index 7b110a0b2756..5437d4a5a13a 100644 --- a/packages/core/test/lib/integrations/mcp-server/semanticConventions.test.ts +++ b/packages/core/test/lib/integrations/mcp-server/semanticConventions.test.ts @@ -434,5 +434,77 @@ describe('MCP Server Semantic Conventions', () => { expect(setStatusSpy).not.toHaveBeenCalled(); expect(endSpy).toHaveBeenCalled(); }); + + it('should instrument prompt call results and complete span with enriched attributes', async () => { + await wrappedMcpServer.connect(mockTransport); + + const setAttributesSpy = vi.fn(); + const setStatusSpy = vi.fn(); + const endSpy = vi.fn(); + const mockSpan = { + setAttributes: setAttributesSpy, + setStatus: setStatusSpy, + end: endSpy, + }; + startInactiveSpanSpy.mockReturnValueOnce( + mockSpan as unknown as ReturnType, + ); + + const promptCallRequest = { + jsonrpc: '2.0', + method: 'prompts/get', + id: 'req-prompt-result', + params: { + name: 'code-review', + arguments: { language: 'typescript', complexity: 'high' }, + }, + }; + + mockTransport.onmessage?.(promptCallRequest, {}); + + expect(startInactiveSpanSpy).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'prompts/get code-review', + op: 'mcp.server', + forceTransaction: true, + attributes: expect.objectContaining({ + 'mcp.method.name': 'prompts/get', + 'mcp.prompt.name': 'code-review', + 'mcp.request.id': 'req-prompt-result', + }), + }), + ); + + const promptResponse = { + jsonrpc: '2.0', + id: 'req-prompt-result', + result: { + description: 'Code review prompt for TypeScript with high complexity analysis', + messages: [ + { + role: 'user', + content: { + type: 'text', + text: 'Please review this TypeScript code for complexity and best practices.', + }, + }, + ], + }, + }; + + mockTransport.send?.(promptResponse); + + expect(setAttributesSpy).toHaveBeenCalledWith( + expect.objectContaining({ + 'mcp.prompt.result.description': 'Code review prompt for TypeScript with high complexity analysis', + 'mcp.prompt.result.message_count': 1, + 'mcp.prompt.result.message_role': 'user', + 'mcp.prompt.result.message_content': 'Please review this TypeScript code for complexity and best practices.', + }), + ); + + expect(setStatusSpy).not.toHaveBeenCalled(); + expect(endSpy).toHaveBeenCalled(); + }); }); }); diff --git a/packages/core/test/lib/integrations/mcp-server/transportInstrumentation.test.ts b/packages/core/test/lib/integrations/mcp-server/transportInstrumentation.test.ts index 7f06eb886cdb..008942ac4099 100644 --- a/packages/core/test/lib/integrations/mcp-server/transportInstrumentation.test.ts +++ b/packages/core/test/lib/integrations/mcp-server/transportInstrumentation.test.ts @@ -4,7 +4,7 @@ import { wrapMcpServerWithSentry } from '../../../../src/integrations/mcp-server import { extractSessionDataFromInitializeRequest, extractSessionDataFromInitializeResponse, -} from '../../../../src/integrations/mcp-server/attributeExtraction'; +} from '../../../../src/integrations/mcp-server/sessionExtraction'; import { cleanupSessionDataForTransport, getClientInfoForTransport,