adding client and server attribute builders, updating span completion logic for initialize requests, and improving session data handling

Author: betegon

packages/core/src/integrations/mcp-server/correlation.ts

@@ -9,8 +9,10 @@
 import { getClient } from '../../currentScopes';
 import { SPAN_STATUS_ERROR } from '../../tracing';
 import type { Span } from '../../types-hoist/span';
+import { MCP_PROTOCOL_VERSION_ATTRIBUTE } from './attributes';
 import { filterMcpPiiFromSpanData } from './piiFiltering';
 import { extractPromptResultAttributes, extractToolResultAttributes } from './resultExtraction';
+import { buildServerAttributesFromInfo, extractSessionDataFromInitializeResponse } from './sessionExtraction';
 import type { MCPTransport, RequestId, RequestSpanMapValue } from './types';
 
 /**
@@ -51,18 +53,30 @@ export function storeSpanForRequest(transport: MCPTransport, requestId: RequestI
 }
 
 /**
- * Completes span with tool results and cleans up correlation
+ * Completes span with results and cleans up correlation
  * @param transport - MCP transport instance
  * @param requestId - Request identifier
- * @param result - Tool execution result for attribute extraction
+ * @param result - Execution result for attribute extraction
  */
 export function completeSpanWithResults(transport: MCPTransport, requestId: RequestId, result: unknown): void {
   const spanMap = getOrCreateSpanMap(transport);
   const spanData = spanMap.get(requestId);
   if (spanData) {
     const { span, method } = spanData;
 
-    if (method === 'tools/call') {
+    if (method === 'initialize') {
+      const sessionData = extractSessionDataFromInitializeResponse(result);
+      const serverAttributes = buildServerAttributesFromInfo(sessionData.serverInfo);
+
+      const initAttributes: Record<string, string | number> = {
+        ...serverAttributes,
+      };
+      if (sessionData.protocolVersion) {
+        initAttributes[MCP_PROTOCOL_VERSION_ATTRIBUTE] = sessionData.protocolVersion;
+      }
+
+      span.setAttributes(initAttributes);
+    } else if (method === 'tools/call') {
       const rawToolAttributes = extractToolResultAttributes(result);
       const client = getClient();
       const sendDefaultPii = Boolean(client?.getOptions().sendDefaultPii);

packages/core/src/integrations/mcp-server/sessionExtraction.ts

@@ -105,6 +105,27 @@ export function getClientAttributes(transport: MCPTransport): Record<string, str
   return attributes;
 }
 
+/**
+ * Build client attributes from PartyInfo directly
+ * @param clientInfo - Client party info
+ * @returns Client attributes for span instrumentation
+ */
+export function buildClientAttributesFromInfo(clientInfo?: PartyInfo): Record<string, string> {
+  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
@@ -127,6 +148,27 @@ export function getServerAttributes(transport: MCPTransport): Record<string, str
   return attributes;
 }
 
+/**
+ * Build server attributes from PartyInfo directly
+ * @param serverInfo - Server party info
+ * @returns Server attributes for span instrumentation
+ */
+export function buildServerAttributesFromInfo(serverInfo?: PartyInfo): Record<string, string> {
+  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

packages/core/src/integrations/mcp-server/transport.ts

@@ -8,16 +8,21 @@
 import { getIsolationScope, withIsolationScope } from '../../currentScopes';
 import { startInactiveSpan, withActiveSpan } from '../../tracing';
 import { fill } from '../../utils/object';
+import { MCP_PROTOCOL_VERSION_ATTRIBUTE } from './attributes';
 import { cleanupPendingSpansForTransport, completeSpanWithResults, storeSpanForRequest } from './correlation';
 import { captureError } from './errorCapture';
-import { extractSessionDataFromInitializeRequest, extractSessionDataFromInitializeResponse } from './sessionExtraction';
+import {
+  buildClientAttributesFromInfo,
+  extractSessionDataFromInitializeRequest,
+  extractSessionDataFromInitializeResponse,
+} from './sessionExtraction';
 import {
   cleanupSessionDataForTransport,
   storeSessionDataForTransport,
   updateSessionDataForTransport,
 } from './sessionManagement';
 import { buildMcpServerSpanConfig, createMcpNotificationSpan, createMcpOutgoingNotificationSpan } from './spans';
-import type { ExtraHandlerData, MCPTransport } from './types';
+import type { ExtraHandlerData, MCPTransport, SessionData } from './types';
 import { isJsonRpcNotification, isJsonRpcRequest, isJsonRpcResponse, isValidContentItem } from './validation';
 
 /**
@@ -31,10 +36,13 @@ export function wrapTransportOnMessage(transport: MCPTransport): void {
     fill(transport, 'onmessage', originalOnMessage => {
       return function (this: MCPTransport, message: unknown, extra?: unknown) {
         if (isJsonRpcRequest(message)) {
-          if (message.method === 'initialize') {
+          const isInitialize = message.method === 'initialize';
+          let initSessionData: SessionData | undefined;
+
+          if (isInitialize) {
             try {
-              const sessionData = extractSessionDataFromInitializeRequest(message);
-              storeSessionDataForTransport(this, sessionData);
+              initSessionData = extractSessionDataFromInitializeRequest(message);
+              storeSessionDataForTransport(this, initSessionData);
             } catch {
               // noop
             }
@@ -46,6 +54,16 @@ export function wrapTransportOnMessage(transport: MCPTransport): void {
             const spanConfig = buildMcpServerSpanConfig(message, this, extra as ExtraHandlerData);
             const span = startInactiveSpan(spanConfig);
 
+            // For initialize requests, add client info directly to span (works even for stateless transports)
+            if (isInitialize && initSessionData) {
+              span.setAttributes({
+                ...buildClientAttributesFromInfo(initSessionData.clientInfo),
+                ...(initSessionData.protocolVersion && {
+                  [MCP_PROTOCOL_VERSION_ATTRIBUTE]: initSessionData.protocolVersion,
+                }),
+              });
+            }
+
             storeSpanForRequest(this, message.id, span, message.method);
 
             return withActiveSpan(span, () => {