61 lines
1.7 KiB
Python
61 lines
1.7 KiB
Python
"""OpenTelemetry helpers for MCP."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Generator, Mapping
|
|
from contextlib import contextmanager
|
|
from typing import Any
|
|
|
|
from opentelemetry.context import Context
|
|
from opentelemetry.propagate import extract, inject
|
|
from opentelemetry.trace import SpanKind, get_current_span, get_tracer
|
|
from opentelemetry.trace.span import Span
|
|
|
|
_tracer = get_tracer("mcp-python-sdk")
|
|
|
|
|
|
@contextmanager
|
|
def otel_span(
|
|
name: str,
|
|
*,
|
|
kind: SpanKind,
|
|
attributes: dict[str, Any] | None = None,
|
|
context: Context | None = None,
|
|
record_exception: bool = True,
|
|
set_status_on_exception: bool = True,
|
|
) -> Generator[Span]:
|
|
"""Create an OTel span."""
|
|
with _tracer.start_as_current_span(
|
|
name,
|
|
kind=kind,
|
|
attributes=attributes,
|
|
context=context,
|
|
record_exception=record_exception,
|
|
set_status_on_exception=set_status_on_exception,
|
|
) as span:
|
|
yield span
|
|
|
|
|
|
def inject_trace_context(meta: dict[str, Any]) -> None:
|
|
"""Inject W3C trace context (traceparent/tracestate) into a `_meta` dict."""
|
|
inject(meta)
|
|
|
|
|
|
def extract_trace_context(meta: Mapping[str, Any] | None) -> Context | None:
|
|
"""Extract W3C trace context from a `_meta` dict.
|
|
|
|
Returns `None` when the carrier is absent, malformed, or carries no
|
|
valid `traceparent`, so callers fall through to ambient parenting; an
|
|
explicit empty `Context` would orphan the span instead of nesting under
|
|
the current one.
|
|
"""
|
|
if not meta:
|
|
return None
|
|
try:
|
|
ctx = extract(meta)
|
|
except (ValueError, TypeError):
|
|
return None
|
|
if not get_current_span(ctx).get_span_context().is_valid:
|
|
return None
|
|
return ctx
|