fix(v9/mcp): capture prompt results and add defensive edge cases for transport instrumentation (#17401)

backport of #17284 and #17291

Author: betegon

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<string, unknown>;
-    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<string, unknown>;
-    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<SessionData> {
-  const sessionData: Partial<SessionData> = {};
-  if (result && typeof result === 'object') {
-    const resultObj = result as Record<string, unknown>;
-    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<string, string> {
-  const clientInfo = getClientInfoForTransport(transport);
-  const attributes: Record<string, string> = {};
-
-  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<string, string> {
-  const serverInfo = getSessionDataForTransport(transport)?.serverInfo;
-  const attributes: Record<string, string> = {};
-
-  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<string, string | number> {
-  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<string, string | number> {
-  const attributes: Record<string, string | number> = {
-    [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<string, unknown>;
-    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<string, unknown>;
-      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<string, string | number | boolean> {
-  let attributes: Record<string, string | number | boolean> = {};
-  if (typeof result !== 'object' || result === null) return attributes;
-
-  const resultObj = result as Record<string, unknown>;
-  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';

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
 // =============================================================================

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;