klz-cables.com/.pnpm-store/v10/files/8c/97437fc1d5c397c67450a6d2360a51dcc45006e1c792c34dfeb24b258e11126936b2023581b322d842d09bc3a50ccf23e385e918f622b6f1f75a6e84bfcf1b

{"version":3,"file":"correlation.js","sources":["../../../../src/integrations/mcp-server/correlation.ts"],"sourcesContent":["/**\n * Request-span correlation system for MCP server instrumentation\n *\n * Handles mapping requestId to span data for correlation with handler execution.\n *\n * Uses sessionId as the primary key for stateful transports. This handles the wrapper\n * transport pattern (e.g., NodeStreamableHTTPServerTransport wrapping WebStandardStreamableHTTPServerTransport)\n * where onmessage and send may receive different `this` values but share the same sessionId.\n *\n * Falls back to WeakMap by transport instance for stateless transports (no sessionId).\n */\n\nimport { SPAN_STATUS_ERROR } from '../../tracing';\nimport type { Span } from '../../types-hoist/span';\nimport { MCP_PROTOCOL_VERSION_ATTRIBUTE } from './attributes';\nimport { extractPromptResultAttributes, extractToolResultAttributes } from './resultExtraction';\nimport { buildServerAttributesFromInfo, extractSessionDataFromInitializeResponse } from './sessionExtraction';\nimport type { MCPTransport, RequestId, RequestSpanMapValue, ResolvedMcpOptions } from './types';\n\n/**\n * Session-scoped correlation for stateful transports (with sessionId)\n * @internal Using sessionId as key handles wrapper transport patterns where\n * different transport objects share the same logical session\n */\nconst sessionToSpanMap = new Map<string, Map<RequestId, RequestSpanMapValue>>();\n\n/**\n * Transport-scoped correlation fallback for stateless transports (no sessionId)\n * @internal WeakMap allows automatic cleanup when transport is garbage collected\n */\nconst statelessSpanMap = new WeakMap<MCPTransport, Map<RequestId, RequestSpanMapValue>>();\n\n/**\n * Gets or creates the span map for a transport, using sessionId when available\n * @internal\n * @param transport - MCP transport instance\n * @returns Span map for the transport/session\n */\nfunction getOrCreateSpanMap(transport: MCPTransport): Map<RequestId, RequestSpanMapValue> {\n  const sessionId = transport.sessionId;\n\n  if (sessionId) {\n    // Stateful transport - use sessionId as key (handles wrapper pattern)\n    let spanMap = sessionToSpanMap.get(sessionId);\n    if (!spanMap) {\n      spanMap = new Map();\n      sessionToSpanMap.set(sessionId, spanMap);\n    }\n    return spanMap;\n  }\n\n  // Stateless fallback - use transport instance as key\n  let spanMap = statelessSpanMap.get(transport);\n  if (!spanMap) {\n    spanMap = new Map();\n    statelessSpanMap.set(transport, spanMap);\n  }\n  return spanMap;\n}\n\n/**\n * Stores span context for later correlation with handler execution\n * @param transport - MCP transport instance\n * @param requestId - Request identifier\n * @param span - Active span to correlate\n * @param method - MCP method name\n */\nexport function storeSpanForRequest(transport: MCPTransport, requestId: RequestId, span: Span, method: string): void {\n  const spanMap = getOrCreateSpanMap(transport);\n  spanMap.set(requestId, {\n    span,\n    method,\n    // eslint-disable-next-line @sentry-internal/sdk/no-unsafe-random-apis\n    startTime: Date.now(),\n  });\n}\n\n/**\n * Completes span with results and cleans up correlation\n * @param transport - MCP transport instance\n * @param requestId - Request identifier\n * @param result - Execution result for attribute extraction\n * @param options - Resolved MCP options\n */\nexport function completeSpanWithResults(\n  transport: MCPTransport,\n  requestId: RequestId,\n  result: unknown,\n  options: ResolvedMcpOptions,\n): void {\n  const spanMap = getOrCreateSpanMap(transport);\n  const spanData = spanMap.get(requestId);\n  if (spanData) {\n    const { span, method } = spanData;\n\n    if (method === 'initialize') {\n      const sessionData = extractSessionDataFromInitializeResponse(result);\n      const serverAttributes = buildServerAttributesFromInfo(sessionData.serverInfo);\n\n      const initAttributes: Record<string, string | number> = {\n        ...serverAttributes,\n      };\n      if (sessionData.protocolVersion) {\n        initAttributes[MCP_PROTOCOL_VERSION_ATTRIBUTE] = sessionData.protocolVersion;\n      }\n\n      span.setAttributes(initAttributes);\n    } else if (method === 'tools/call') {\n      const toolAttributes = extractToolResultAttributes(result, options.recordOutputs);\n      span.setAttributes(toolAttributes);\n    } else if (method === 'prompts/get') {\n      const promptAttributes = extractPromptResultAttributes(result, options.recordOutputs);\n      span.setAttributes(promptAttributes);\n    }\n\n    span.end();\n    spanMap.delete(requestId);\n  }\n}\n\n/**\n * Cleans up pending spans for a specific transport (when that transport closes)\n * @param transport - MCP transport instance\n */\nexport function cleanupPendingSpansForTransport(transport: MCPTransport): void {\n  const sessionId = transport.sessionId;\n\n  // Try sessionId-based cleanup first (for stateful transports)\n  if (sessionId) {\n    const spanMap = sessionToSpanMap.get(sessionId);\n    if (spanMap) {\n      for (const [, spanData] of spanMap) {\n        spanData.span.setStatus({\n          code: SPAN_STATUS_ERROR,\n          message: 'cancelled',\n        });\n        spanData.span.end();\n      }\n      sessionToSpanMap.delete(sessionId);\n    }\n    return;\n  }\n\n  // Fallback to transport-based cleanup (for stateless transports)\n  const spanMap = statelessSpanMap.get(transport);\n  if (spanMap) {\n    for (const [, spanData] of spanMap) {\n      spanData.span.setStatus({\n        code: SPAN_STATUS_ERROR,\n        message: 'cancelled',\n      });\n      spanData.span.end();\n    }\n    spanMap.clear();\n    // Note: WeakMap entries are automatically cleaned up when transport is GC'd\n  }\n}\n"],"names":[],"mappings":";;;;;AAAA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;AACA;;;AASA;AACA;AACA;AACA;AACA;AACA,MAAM,gBAAA,GAAmB,IAAI,GAAG,EAA+C;;AAE/E;AACA;AACA;AACA;AACA,MAAM,gBAAA,GAAmB,IAAI,OAAO,EAAqD;;AAEzF;AACA;AACA;AACA;AACA;AACA;AACA,SAAS,kBAAkB,CAAC,SAAS,EAAqD;AAC1F,EAAE,MAAM,SAAA,GAAY,SAAS,CAAC,SAAS;;AAEvC,EAAE,IAAI,SAAS,EAAE;AACjB;AACA,IAAI,IAAI,UAAU,gBAAgB,CAAC,GAAG,CAAC,SAAS,CAAC;AACjD,IAAI,IAAI,CAAC,OAAO,EAAE;AAClB,MAAM,OAAA,GAAU,IAAI,GAAG,EAAE;AACzB,MAAM,gBAAgB,CAAC,GAAG,CAAC,SAAS,EAAE,OAAO,CAAC;AAC9C,IAAI;AACJ,IAAI,OAAO,OAAO;AAClB,EAAE;;AAEF;AACA,EAAE,IAAI,UAAU,gBAAgB,CAAC,GAAG,CAAC,SAAS,CAAC;AAC/C,EAAE,IAAI,CAAC,OAAO,EAAE;AAChB,IAAI,OAAA,GAAU,IAAI,GAAG,EAAE;AACvB,IAAI,gBAAgB,CAAC,GAAG,CAAC,SAAS,EAAE,OAAO,CAAC;AAC5C,EAAE;AACF,EAAE,OAAO,OAAO;AAChB;;AAEA;AACA;AACA;AACA;AACA;AACA;AACA;AACO,SAAS,mBAAmB,CAAC,SAAS,EAAgB,SAAS,EAAa,IAAI,EAAQ,MAAM,EAAgB;AACrH,EAAE,MAAM,OAAA,GAAU,kBAAkB,CAAC,SAAS,CAAC;AAC/C,EAAE,OAAO,CAAC,GAAG,CAAC,SAAS,EAAE;AACzB,IAAI,IAAI;AACR,IAAI,MAAM;AACV;AACA,IAAI,SAAS,EAAE,IAAI,CAAC,GAAG,EAAE;AACzB,GAAG,CAAC;AACJ;;AAEA;AACA;AACA;AACA;AACA;AACA;AACA;AACO,SAAS,uBAAuB;AACvC,EAAE,SAAS;AACX,EAAE,SAAS;AACX,EAAE,MAAM;AACR,EAAE,OAAO;AACT,EAAQ;AACR,EAAE,MAAM,OAAA,GAAU,kBAAkB,CAAC,SAAS,CAAC;AAC/C,EAAE,MAAM,WAAW,OAAO,CAAC,GAAG,CAAC,SAAS,CAAC;AACzC,EAAE,IAAI,QAAQ,EAAE;AAChB,IAAI,MAAM,EAAE,IAAI,EAAE,MAAA,EAAO,GAAI,QAAQ;;AAErC,IAAI,IAAI,MAAA,KAAW,YAAY,EAAE;AACjC,MAAM,MAAM,WAAA,GAAc,wCAAwC,CAAC,MAAM,CAAC;AAC1E,MAAM,MAAM,mBAAmB,6BAA6B,CAAC,WAAW,CAAC,UAAU,CAAC;;AAEpF,MAAM,MAAM,cAAc,GAAoC;AAC9D,QAAQ,GAAG,gBAAgB;AAC3B,OAAO;AACP,MAAM,IAAI,WAAW,CAAC,eAAe,EAAE;AACvC,QAAQ,cAAc,CAAC,8BAA8B,IAAI,WAAW,CAAC,eAAe;AACpF,MAAM;;AAEN,MAAM,IAAI,CAAC,aAAa,CAAC,cAAc,CAAC;AACxC,IAAI,OAAO,IAAI,MAAA,KAAW,YAAY,EAAE;AACxC,MAAM,MAAM,cAAA,GAAiB,2BAA2B,CAAC,MAAM,EAAE,OAAO,CAAC,aAAa,CAAC;AACvF,MAAM,IAAI,CAAC,aAAa,CAAC,cAAc,CAAC;AACxC,IAAI,OAAO,IAAI,MAAA,KAAW,aAAa,EAAE;AACzC,MAAM,MAAM,gBAAA,GAAmB,6BAA6B,CAAC,MAAM,EAAE,OAAO,CAAC,aAAa,CAAC;AAC3F,MAAM,IAAI,CAAC,aAAa,CAAC,gBAAgB,CAAC;AAC1C,IAAI;;AAEJ,IAAI,IAAI,CAAC,GAAG,EAAE;AACd,IAAI,OAAO,CAAC,MAAM,CAAC,SAAS,CAAC;AAC7B,EAAE;AACF;;AAEA;AACA;AACA;AACA;AACO,SAAS,+BAA+B,CAAC,SAAS,EAAsB;AAC/E,EAAE,MAAM,SAAA,GAAY,SAAS,CAAC,SAAS;;AAEvC;AACA,EAAE,IAAI,SAAS,EAAE;AACjB,IAAI,MAAM,UAAU,gBAAgB,CAAC,GAAG,CAAC,SAAS,CAAC;AACnD,IAAI,IAAI,OAAO,EAAE;AACjB,MAAM,KAAK,MAAM,GAAG,QAAQ,CAAA,IAAK,OAAO,EAAE;AAC1C,QAAQ,QAAQ,CAAC,IAAI,CAAC,SAAS,CAAC;AAChC,UAAU,IAAI,EAAE,iBAAiB;AACjC,UAAU,OAAO,EAAE,WAAW;AAC9B,SAAS,CAAC;AACV,QAAQ,QAAQ,CAAC,IAAI,CAAC,GAAG,EAAE;AAC3B,MAAM;AACN,MAAM,gBAAgB,CAAC,MAAM,CAAC,SAAS,CAAC;AACxC,IAAI;AACJ,IAAI;AACJ,EAAE;;AAEF;AACA,EAAE,MAAM,UAAU,gBAAgB,CAAC,GAAG,CAAC,SAAS,CAAC;AACjD,EAAE,IAAI,OAAO,EAAE;AACf,IAAI,KAAK,MAAM,GAAG,QAAQ,CAAA,IAAK,OAAO,EAAE;AACxC,MAAM,QAAQ,CAAC,IAAI,CAAC,SAAS,CAAC;AAC9B,QAAQ,IAAI,EAAE,iBAAiB;AAC/B,QAAQ,OAAO,EAAE,WAAW;AAC5B,OAAO,CAAC;AACR,MAAM,QAAQ,CAAC,IAAI,CAAC,GAAG,EAAE;AACzB,IAAI;AACJ,IAAI,OAAO,CAAC,KAAK,EAAE;AACnB;AACA,EAAE;AACF;;;;"}