diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index c93f95ec97d..fa41070a8de 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -18,12 +18,14 @@ before the LLM call even starts), so a guardrail is a sibling of the LLM call, not a child of it. The emitter parents every span to the ambient OTel context (the active server span), which matches this. -MCP spans (``MCP_TOOL_CALL``, ``MCP_LIST_TOOLS``) are intentionally NOT in this -tree. Per the OTel GenAI MCP semconv, MCP and the HTTP transport are independent -contexts, so an MCP span parents to the trace context the client propagated in -``params._meta`` (or starts its own root when none is propagated) and records the -``PROXY_REQUEST`` transport span as a span *link*, never a parent. The registry -encodes this as ``parent=None, links=PROXY_REQUEST``. +MCP spans (``MCP_TOOL_CALL``, ``MCP_LIST_TOOLS``) have two shapes, chosen at emit +time by :func:`resolve_mcp_span_context`. When the client propagates trace context +in ``params._meta`` MCP and the HTTP transport are independent contexts per the +OTel GenAI MCP semconv, so the span parents to that propagated context and records +the ``PROXY_REQUEST`` transport span as a span *link*, never a parent — the shape +this registry's ``parent=None, links=PROXY_REQUEST`` entry encodes. When nothing is +propagated (the common case) the span nests under the transport span of the request +carrying that message, so the tool call stays in one trace. Not every service call becomes a span — :func:`span_role_for_service` decides: @@ -89,12 +91,13 @@ class SpanSpec: SPAN_REGISTRY: dict[SpanRole, SpanSpec] = { SpanRole.PROXY_REQUEST: SpanSpec(SpanRole.PROXY_REQUEST, LiteLLMSpanKind.SERVER, parent=None), SpanRole.LLM_CALL: SpanSpec(SpanRole.LLM_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST), - # MCP and the HTTP transport are independent contexts (OTel GenAI MCP semconv), - # so an MCP span does not nest under the transport span. The proxy is an MCP - # client to the upstream server, so it's a CLIENT span; it parents to the trace - # context the client propagated in ``params._meta`` (or starts its own root when - # none is propagated) and records the PROXY_REQUEST transport span as a span - # *link*, never a parent — hence ``parent=None, links=PROXY_REQUEST``. + # The proxy is an MCP client to the upstream server, so MCP spans are CLIENT + # spans. With trace context propagated in ``params._meta``, MCP and the HTTP + # transport are independent contexts (OTel GenAI MCP semconv): the span parents + # to the propagated context and records the PROXY_REQUEST transport span as a + # span *link*, never a parent — the shape ``parent=None, links=PROXY_REQUEST`` + # encodes. With nothing propagated, ``resolve_mcp_span_context`` nests the span + # under that message's transport span instead, keeping the call in one trace. SpanRole.MCP_TOOL_CALL: SpanSpec( SpanRole.MCP_TOOL_CALL, LiteLLMSpanKind.CLIENT, parent=None, links=SpanRole.PROXY_REQUEST ), diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 8acac112c3d..939559347b1 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -5,7 +5,14 @@ from typing import Mapping from opentelemetry import baggage from opentelemetry.context import Context, get_current -from opentelemetry.trace import Link, Span, get_current_span, set_span_in_context +from opentelemetry.trace import ( + Link, + NonRecordingSpan, + Span, + SpanContext, + get_current_span, + set_span_in_context, +) from opentelemetry.trace.propagation.tracecontext import ( TraceContextTextMapPropagator, ) @@ -72,6 +79,62 @@ def reset_mcp_message_trace_carrier(token: "Token[Mapping[str, str] | None]") -> _mcp_message_trace_carrier.reset(token) +# The transport span of the HTTP request carrying the CURRENT MCP message, as a +# plain ``SpanContext`` so it can cross a task boundary. +# +# ``_request_root_span`` above cannot be used for MCP: a *stateful* streamable-HTTP +# session runs every message on the single task spawned by that session's +# ``initialize`` POST, so the ContextVar the ASGI request task writes at auth time +# is frozen at ``initialize`` there and never sees the later ``tools/call`` POSTs. +# Reading it from the message handler would parent every tool call in the session +# to the first request's (already ended) server span. The gateway instead resolves +# the current message's transport span on the request task and hands it over the +# same way it hands over per-request auth, and the handler publishes it here for +# the span emitter to pick up. +_mcp_message_transport_span_context: "ContextVar[SpanContext | None]" = ContextVar( + "litellm_otel_mcp_message_transport_span_context", default=None +) + + +def set_mcp_message_transport_span_context( + span_context: "SpanContext | None", +) -> "Token[SpanContext | None]": + """Publish the transport span of the request carrying the current MCP message. + + Returns the reset token; the caller must reset it once the message is handled + so the transport never leaks to the next message on the same session task. + """ + return _mcp_message_transport_span_context.set(span_context) + + +def reset_mcp_message_transport_span_context(token: "Token[SpanContext | None]") -> None: + _mcp_message_transport_span_context.reset(token) + + +def request_root_span_context() -> "SpanContext | None": + """The anchored request root span's context, safe to hand to another task. + + A ``SpanContext`` is an immutable value, unlike the live ``Span``, so passing it + across the MCP session-task boundary cannot keep a finished span alive or invite + writes to it from the wrong request. + """ + span = request_root_span() + return span.get_span_context() if span is not None else None + + +def _mcp_transport_span_context() -> "SpanContext | None": + """The transport span an MCP message span should attach to. + + Prefers the transport the gateway published for this specific message; falls + back to the ambient request anchor for paths that emit an MCP span on the + request task itself (the REST MCP endpoints, the SDK). + """ + published = _mcp_message_transport_span_context.get() + if published is not None and published.is_valid: + return published + return request_root_span_context() + + def set_request_baggage(values: Mapping[str, str], context: Context | None = None) -> Context: """Return a context with ``values`` written into Baggage.""" ctx = context @@ -132,33 +195,44 @@ def resolve_request_span_context() -> Context: def resolve_mcp_span_context( carrier: "Mapping[str, str] | None" = None, ) -> "tuple[Context, tuple[Link, ...]]": - """Parent context + links for an MCP message span, per the OTel GenAI MCP semconv. + """Parent context + links for an MCP message span. - MCP and the underlying transport (HTTP) are independent lifecycles — one - streamable-HTTP session multiplexes many messages, so nesting the message span - under the HTTP/session span is wrong (it renders the message at the session's - start, skewed by however long the session has been open). Instead: + When the client propagates W3C trace context in the request's ``params._meta`` + (SEP-414), MCP and the underlying transport are independent lifecycles — one + streamable-HTTP session multiplexes many messages, and the client's own span is + the truthful parent. So, per the OTel GenAI MCP semconv: - * parent to the trace context the client propagated in the request's - ``params._meta`` (a *remote* parent), and - * record the transport/session span as a *link*, never the parent. + * parent to the trace context the client propagated (a *remote* parent), and + * record the transport span as a *link*, never the parent. + + Almost no client implements SEP-414 yet, so in practice nothing is propagated. + Rooting the span there splits a single tool call into two disconnected traces + joined only by a link, which is how it surfaces in APM: the ``POST`` transaction + and the ``tools/call`` span share no trace. With no remote parent to honor, + parent to the transport span of the request carrying this message instead, so + the call stays in one trace; no link is added since the transport is now the + real parent. The transport comes from :func:`_mcp_transport_span_context`, which + is the *current message's* POST rather than whatever request happened to open + the session, so a long-lived session does not glue every message under its + first request. With neither a remote parent nor a transport the returned context + carries no span and the span legitimately starts its own root trace. Only trace context (``traceparent``/``tracestate``) is extracted, never the client's W3C Baggage: ``params._meta`` is caller-controlled, and the otel baggage processor stamps allowlisted baggage keys (``litellm.team.id``, ``litellm.metadata.*``, ...) onto the span as attributes, so honoring remote - baggage would let a client spoof a span's identity attribution. - - With no propagated context the returned context carries no span, so the span - starts its own root trace (still linked to the transport). The base context is - explicitly empty so an absent ``traceparent`` can never fall through to the - ambient (stale session) span. + baggage would let a client spoof a span's identity attribution. The base context + for extraction is explicitly empty so an absent or malformed ``traceparent`` can + never fall through to the ambient (stale session) span. """ source = carrier if carrier is not None else _mcp_message_trace_carrier.get() parent = _PROPAGATOR.extract(dict(source or {}), context=Context()) - transport = request_root_span() - links = (Link(transport.get_span_context()),) if transport is not None else () - return parent, links + transport = _mcp_transport_span_context() + if is_recordable_span(get_current_span(parent)): + return parent, (Link(transport),) if transport is not None else () + if transport is not None: + return context_from_span(NonRecordingSpan(transport)), () + return parent, () def is_recordable_span(obj: object) -> bool: diff --git a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py index 7122c64ec64..f7bc14575c7 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py +++ b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py @@ -1,9 +1,12 @@ -from typing import Dict, List, Optional +from typing import TYPE_CHECKING, Dict, List, Optional from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser from litellm.proxy._types import UserAPIKeyAuth +if TYPE_CHECKING: + from opentelemetry.trace import SpanContext + class MCPAuthenticatedUser(AuthenticatedUser): """ @@ -16,6 +19,8 @@ class MCPAuthenticatedUser(AuthenticatedUser): 4. Server-specific authentication headers 5. OAuth2 headers 6. Raw headers - allows forwarding specific headers to the MCP server, specified by the admin. + 7. Transport span context - the tracing span of the HTTP request carrying the current + message, which a stateful session's message handler cannot read from its own task. """ def __init__( @@ -28,6 +33,7 @@ class MCPAuthenticatedUser(AuthenticatedUser): mcp_protocol_version: Optional[str] = None, raw_headers: Optional[Dict[str, str]] = None, client_ip: Optional[str] = None, + transport_span_context: Optional["SpanContext"] = None, ): self.user_api_key_auth = user_api_key_auth self.mcp_auth_header = mcp_auth_header @@ -37,3 +43,4 @@ class MCPAuthenticatedUser(AuthenticatedUser): self.oauth2_headers = oauth2_headers self.raw_headers = raw_headers self.client_ip = client_ip + self.transport_span_context = transport_span_context diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 4fca4406a6f..483f57f9139 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -15,6 +15,7 @@ import types import uuid from datetime import datetime from typing import ( + TYPE_CHECKING, Any, AsyncIterator, Callable, @@ -107,6 +108,9 @@ _MAX_STATEFUL_SESSIONS_PER_OWNER = 100 # arbitrarily large body just to make a routing decision. _MCP_ROUTING_PEEK_MAX_BYTES = 4096 +if TYPE_CHECKING: + from opentelemetry.trace import SpanContext + def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None: """Remove a (user_id, server_id) entry from the BYOK credential cache. @@ -242,10 +246,12 @@ def _mcp_meta_trace_carrier(req_ctx: object) -> Optional[dict[str, str]]: """The W3C trace context (``traceparent``/``tracestate``) the MCP client propagated in the request's ``params._meta`` (SEP-414), or ``None``. - Per the OTel MCP semconv the MCP span parents to this propagated context rather - than to the HTTP/session transport (which is recorded as a link instead), so a - streamable-HTTP session that multiplexes many messages does not glue every - message under the session's first request. The client's W3C Baggage is + When present, per the OTel MCP semconv the MCP span parents to this propagated + context rather than to the HTTP transport (which is recorded as a link instead). + When absent, the span nests under the transport span of the request carrying + this specific message, so a streamable-HTTP session that multiplexes many + messages still does not glue every message under the session's first request; + see ``resolve_mcp_span_context``. The client's W3C Baggage is deliberately excluded: it is caller-controlled, and the otel baggage processor stamps allowlisted baggage keys (``litellm.team.id``, ``litellm.metadata.*``, ...) onto the span, so honoring remote baggage would let a client spoof a @@ -288,6 +294,56 @@ def _otel_reset_mcp_trace_carrier(token: object) -> None: return +def _otel_request_transport_span_context() -> Optional["SpanContext"]: + """The tracing span of the HTTP request being handled, as a portable value. + + Resolved on the ASGI request task, where the proxy's server span is anchored, + and carried to the MCP message handler on the authenticated-user object. A + stateful streamable-HTTP session handles every message on the task spawned by + its ``initialize`` POST, so the handler's own task cannot see later requests' + spans; this is the same reason per-request auth is carried across rather than + read from a ContextVar. Lazily imported so opentelemetry stays an optional + dependency; returns ``None`` when otel_v2 is unavailable or no request span is + anchored.""" + try: + from litellm.integrations.otel.plumbing.context import ( + request_root_span_context, + ) + + return request_root_span_context() + except ImportError: + return None + + +def _otel_set_mcp_transport_span_context(span_context: Optional["SpanContext"]) -> object: + """Publish the current message's transport span for the otel_v2 MCP span and + return a reset token, or ``None`` when otel_v2 is unavailable.""" + if span_context is None: + return None + try: + from litellm.integrations.otel.plumbing.context import ( + set_mcp_message_transport_span_context, + ) + + return set_mcp_message_transport_span_context(span_context) + except ImportError: + return None + + +def _otel_reset_mcp_transport_span_context(token: object) -> None: + """Paired with ``_otel_set_mcp_transport_span_context``.""" + if token is None: + return + try: + from litellm.integrations.otel.plumbing.context import ( + reset_mcp_message_transport_span_context, + ) + + reset_mcp_message_transport_span_context(token) + except ImportError: + return + + def _proxy_exception_to_http_exception(exc: ProxyException) -> HTTPException: """Map a ``ProxyException`` to an ``HTTPException`` that preserves its real status code and headers. @@ -654,6 +710,18 @@ if MCP_AVAILABLE: ############### MCP Server Routes ####################### ######################################################## + def _current_transport_span_context() -> Optional["SpanContext"]: + """The transport span of the HTTP request carrying the message being handled. + + Published by the ASGI request task onto the authenticated-user object, because + a stateful session's message handler runs on the task spawned by that session's + ``initialize`` POST and so cannot read later requests' spans from its own task. + """ + auth_user = auth_context_var.get() + if not isinstance(auth_user, MCPAuthenticatedUser): + auth_user = _recover_auth_from_session() + return auth_user.transport_span_context if auth_user is not None else None + @server.list_tools() async def handle_list_tools() -> "ListToolsResult | List[Tool]": """ @@ -670,9 +738,11 @@ if MCP_AVAILABLE: if req_ctx: _session_reset_token = active_mcp_session_var.set(req_ctx.session) _trace_token = None + _transport_token = None try: _trace_token = _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(req_ctx)) + _transport_token = _otel_set_mcp_transport_span_context(_current_transport_span_context()) # Get user authentication from context variable ( user_api_key_auth, @@ -728,6 +798,7 @@ if MCP_AVAILABLE: # This prevents the HTTP stream from failing and allows the client to get a response return [] finally: + _otel_reset_mcp_transport_span_context(_transport_token) _otel_reset_mcp_trace_carrier(_trace_token) if _session_reset_token is not None: active_mcp_session_var.reset(_session_reset_token) @@ -901,9 +972,11 @@ if MCP_AVAILABLE: if req_ctx: _session_reset_token = active_mcp_session_var.set(req_ctx.session) _trace_token = None + _transport_token = None try: _trace_token = _otel_set_mcp_trace_carrier(_mcp_meta_trace_carrier(req_ctx)) + _transport_token = _otel_set_mcp_transport_span_context(_current_transport_span_context()) # Validate arguments ( user_api_key_auth, @@ -1042,6 +1115,7 @@ if MCP_AVAILABLE: return response finally: + _otel_reset_mcp_transport_span_context(_transport_token) _otel_reset_mcp_trace_carrier(_trace_token) if _session_reset_token is not None: active_mcp_session_var.reset(_session_reset_token) @@ -4197,6 +4271,7 @@ if MCP_AVAILABLE: session_id=session_id if use_stateful else None, touch_last_seen=(scope.get("method") or "").upper() != "DELETE", copy_existing_session_auth_context=is_initialize, + transport_span_context=_otel_request_transport_span_context(), ) local_send = send if use_stateful and is_initialize: @@ -4421,6 +4496,7 @@ if MCP_AVAILABLE: oauth2_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, client_ip: Optional[str] = None, + transport_span_context: Optional["SpanContext"] = None, ) -> None: auth_user.user_api_key_auth = user_api_key_auth auth_user.mcp_auth_header = mcp_auth_header @@ -4429,6 +4505,7 @@ if MCP_AVAILABLE: auth_user.oauth2_headers = oauth2_headers auth_user.raw_headers = raw_headers auth_user.client_ip = client_ip + auth_user.transport_span_context = transport_span_context def set_auth_context( user_api_key_auth: Optional[UserAPIKeyAuth], @@ -4438,6 +4515,7 @@ if MCP_AVAILABLE: oauth2_headers: Optional[Dict[str, str]] = None, raw_headers: Optional[Dict[str, str]] = None, client_ip: Optional[str] = None, + transport_span_context: Optional["SpanContext"] = None, ) -> MCPAuthenticatedUser: """ Set the UserAPIKeyAuth in the auth context variable. @@ -4448,6 +4526,7 @@ if MCP_AVAILABLE: mcp_servers: Optional list of server names and access groups to filter by mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value} client_ip: Client IP address for MCP access control + transport_span_context: Tracing span of the HTTP request carrying this message """ auth_user = MCPAuthenticatedUser( user_api_key_auth=user_api_key_auth, @@ -4457,6 +4536,7 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, client_ip=client_ip, + transport_span_context=transport_span_context, ) auth_context_var.set(auth_user) return auth_user @@ -4472,6 +4552,7 @@ if MCP_AVAILABLE: session_id: Optional[str] = None, touch_last_seen: bool = True, copy_existing_session_auth_context: bool = False, + transport_span_context: Optional["SpanContext"] = None, ) -> MCPAuthenticatedUser: auth_user = _stateful_session_auth_contexts.get(session_id) if session_id else None if auth_user is not None and session_id is not None: @@ -4486,6 +4567,7 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, client_ip=client_ip, + transport_span_context=transport_span_context, ) _update_auth_context( auth_user=auth_user, @@ -4496,6 +4578,7 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, client_ip=client_ip, + transport_span_context=transport_span_context, ) auth_context_var.set(auth_user) return auth_user @@ -4507,6 +4590,7 @@ if MCP_AVAILABLE: oauth2_headers=oauth2_headers, raw_headers=raw_headers, client_ip=client_ip, + transport_span_context=transport_span_context, ) def _wrap_send_with_stateful_session_auth_context( diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index 5f6002f4cdf..d3593f2c06b 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -29,7 +29,9 @@ from litellm.integrations.otel import ( # noqa: E402 from litellm.integrations.otel.plumbing import providers # noqa: E402 from litellm.integrations.otel.plumbing.context import ( # noqa: E402 reset_mcp_message_trace_carrier, + reset_mcp_message_transport_span_context, set_mcp_message_trace_carrier, + set_mcp_message_transport_span_context, set_request_root_span, ) from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402 @@ -56,9 +58,11 @@ def _reset_request_root_span(): _otel_context._request_root_span.set(None) _otel_context._mcp_message_trace_carrier.set(None) + _otel_context._mcp_message_transport_span_context.set(None) yield _otel_context._request_root_span.set(None) _otel_context._mcp_message_trace_carrier.set(None) + _otel_context._mcp_message_transport_span_context.set(None) def _payload(**overrides): @@ -528,15 +532,15 @@ _MCP_SPAN_CASES = [ @pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES) -def test_mcp_span_roots_and_links_transport_without_propagated_context( +def test_mcp_span_nests_under_transport_without_propagated_context( make_payload, span_name ): - """MCP and the HTTP transport are independent lifecycles (one streamable-HTTP - session multiplexes many messages), so per the MCP semconv the message span - must NOT nest under the session/transport span — that is what made it render - skewed at the session's start. With no propagated ``params._meta`` context it - starts its own root trace and records the transport span as a *link*, never - the parent.""" + """Almost no MCP client implements SEP-414, so ``params._meta`` normally carries + no trace context. Rooting the span there split one tool call into two traces + joined only by a link, which is how it surfaced in APM: the ``POST`` transaction + and the ``tools/call`` span shared no ``trace_id``. With no remote parent to + honor the span nests under the transport span instead, and records no link since + the transport is now the real parent.""" logger, exporter = _logger() transport = logger._emitter.start_span( SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME @@ -549,11 +553,75 @@ def test_mcp_span_roots_and_links_transport_without_propagated_context( ) transport.end() span = next(s for s in exporter.get_finished_spans() if s.name == span_name) + assert span.parent is not None + assert span.parent.span_id == transport.get_span_context().span_id + assert span.context.trace_id == transport.get_span_context().trace_id + assert span.links == () + + +@pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES) +def test_mcp_span_nests_under_this_messages_transport_not_the_session_opener( + make_payload, span_name +): + """A *stateful* streamable-HTTP session runs every message on the single task + spawned by that session's ``initialize`` POST, so the ``_request_root_span`` + ContextVar the ASGI request task writes is frozen at ``initialize`` inside the + handler and never sees the later ``tools/call`` POST. Nesting on that anchor + would hang every tool call of the session off the first request's (already + ended) span, rendering skewed at the session's start. The gateway resolves the + current message's transport on the request task and publishes it, so the span + parents to the POST that actually carried this message.""" + logger, exporter = _logger() + session_opener = logger._emitter.start_span( + SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME + ) + this_message = logger._emitter.start_span( + SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME + ) + + async def session_task(): + token = set_mcp_message_transport_span_context( + this_message.get_span_context() + ) + try: + await logger.async_log_success_event( + {"standard_logging_object": make_payload()}, None, None, None + ) + finally: + reset_mcp_message_transport_span_context(token) + + async def initialize_request(): + # The anchor the session task inherits is the one ``initialize`` left behind; + # spawning here reproduces the SDK's session task, which outlives this request. + set_request_root_span(session_opener) + await asyncio.create_task(session_task()) + + asyncio.run(initialize_request()) + session_opener.end() + this_message.end() + span = next(s for s in exporter.get_finished_spans() if s.name == span_name) + assert span.parent is not None + assert span.parent.span_id == this_message.get_span_context().span_id + assert span.context.trace_id == this_message.get_span_context().trace_id + assert span.parent.span_id != session_opener.get_span_context().span_id + assert span.context.trace_id != session_opener.get_span_context().trace_id + + +@pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES) +def test_mcp_span_roots_without_transport_or_propagated_context( + make_payload, span_name +): + """With neither a remote parent nor a transport span there is nothing to nest + under, so the span legitimately starts its own root trace with no links.""" + logger, exporter = _logger() + asyncio.run( + logger.async_log_success_event( + {"standard_logging_object": make_payload()}, None, None, None + ) + ) + span = next(s for s in exporter.get_finished_spans() if s.name == span_name) assert span.parent is None - assert span.context.trace_id != transport.get_span_context().trace_id - assert [link.context.span_id for link in span.links] == [ - transport.get_span_context().span_id - ] + assert span.links == () @pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES) @@ -641,10 +709,11 @@ def test_mcp_span_carries_authenticated_identity(make_payload, span_name): assert span.attributes[LiteLLM.TEAM_ID] == "t1" -def test_mcp_span_malformed_traceparent_starts_root(): +def test_mcp_span_malformed_traceparent_nests_under_transport(): """A malformed traceparent in ``params._meta`` must not crash or parent to a - bogus span: the propagator ignores it, so the span starts its own root trace and - still links the transport span.""" + bogus span: the propagator ignores it, leaving no remote parent, so the span + falls back to nesting under the transport span rather than starting a + disconnected root trace.""" logger, exporter = _logger() transport = logger._emitter.start_span( SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME @@ -661,9 +730,44 @@ def test_mcp_span_malformed_traceparent_starts_root(): reset_mcp_message_trace_carrier(token) transport.end() span = next(s for s in exporter.get_finished_spans() if s.name == "tools/list") - assert span.parent is None + assert span.parent is not None + assert span.parent.span_id == transport.get_span_context().span_id + assert span.links == () + + +def test_mcp_span_links_this_messages_transport_when_context_is_propagated(): + """On the semconv path the transport is recorded as a link, and that link must + point at the POST carrying this message too. Reading the stale session anchor + would attribute the tool call to whichever request opened the session.""" + logger, exporter = _logger() + session_opener = logger._emitter.start_span( + SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME + ) + this_message = logger._emitter.start_span( + SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME + ) + set_request_root_span(session_opener) + trace_token = set_mcp_message_trace_carrier( + {"traceparent": "00-11111111111111111111111111111111-2222222222222222-01"} + ) + transport_token = set_mcp_message_transport_span_context( + this_message.get_span_context() + ) + try: + asyncio.run( + logger.async_log_success_event( + {"standard_logging_object": _mcp_list_payload()}, None, None, None + ) + ) + finally: + reset_mcp_message_transport_span_context(transport_token) + reset_mcp_message_trace_carrier(trace_token) + session_opener.end() + this_message.end() + span = next(s for s in exporter.get_finished_spans() if s.name == "tools/list") + assert span.parent is not None and span.parent.span_id == 0x2222222222222222 assert [link.context.span_id for link in span.links] == [ - transport.get_span_context().span_id + this_message.get_span_context().span_id ]