From c9598532da353fab799b77f8a005af315b8233e5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Miguel=20Beteg=C3=B3n?= Date: Wed, 6 Aug 2025 17:19:56 +0200 Subject: [PATCH] feat(core): MCP Server - Capture prompt results from prompt function calls (#17284) closes #17283 includes these attributes for `mcp.server` spans: - `mcp.prompt.result.description` - `mcp.prompt.result.message_content` - `mcp.prompt.result.message_role` - `mcp.prompt.result.message_count` Example: Screenshot 2025-08-01 at 12 40 46 Needed to make `attributeExtraction.ts` <300 lines of code (requirement) so it's now split between `sessionExtraction.ts`, `sessionExtraction.ts` and `resultExtraction.ts`. So changes explained so it's easier to review: - The only function this PR adds is `extractPromptResultAttributes` inside `resultExtraction.ts`. - It adds the prompt results as PII in `piiFiltering.ts`. Just add them to the `set`. - adds a `else if (method === 'prompts/get')` to execute the `extractPromptResultAttributes` function. - adds a test that checks we're capturing the results and updates the PII test to check PII result attributes are being removed if sending PII is not enabled. (cherry picked from commit 0e05a40b133e5fa2b61734c0f183e4c44fc990e1) --- .../mcp-server/attributeExtraction.ts | 267 +----------------- .../src/integrations/mcp-server/attributes.ts | 22 ++ .../integrations/mcp-server/correlation.ts | 13 +- .../integrations/mcp-server/piiFiltering.ts | 34 ++- .../mcp-server/resultExtraction.ts | 126 +++++++++ .../mcp-server/sessionExtraction.ts | 202 +++++++++++++ .../mcp-server/sessionManagement.ts | 4 +- .../src/integrations/mcp-server/transport.ts | 12 +- .../src/integrations/mcp-server/validation.ts | 9 + .../mcp-server/piiFiltering.test.ts | 42 ++- .../mcp-server/semanticConventions.test.ts | 72 +++++ .../transportInstrumentation.test.ts | 2 +- 12 files changed, 526 insertions(+), 279 deletions(-) create mode 100644 packages/core/src/integrations/mcp-server/resultExtraction.ts create mode 100644 packages/core/src/integrations/mcp-server/sessionExtraction.ts 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,