mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
fix(otel): keep an MCP tool call in one trace, anchored to its own request (#34537)
Under otel_v2 a single MCP tool call surfaced in APM as two disconnected
traces joined only by a span link: the HTTP transport transaction
POST /{mcp_server_name}/mcp and the tools/call span carrying
error.type=MCPToolResultError. resolve_mcp_span_context parented the MCP
span to the W3C trace context the client propagates in params._meta
(SEP-414) and recorded the transport as a link, so with no traceparent
propagated (the common case today, including MCP Inspector) the span
started its own root trace.
Nest the MCP span under the transport span when nothing is propagated, so
the call stays in one trace; the propagated-context path is unchanged and
still parents to the remote context and links the transport per the OTel
GenAI MCP semconv.
The transport has to be resolved per message rather than read from the
request-root ContextVar. A stateful streamable-HTTP session runs every
message on the single task the session's initialize POST spawned, so that
ContextVar is frozen at initialize inside the handler: live capture on
staging showed the tools/call span linking the initialize POST rather than
the POST that carried it, and nesting on that anchor would hang every tool
call of a session off the first request's already-ended span. The gateway
now resolves the current request's span on the ASGI task and carries it to
the handler on the authenticated-user object, the same way per-request auth
already crosses that boundary.
This commit is contained in:
Yassin Kortam
GitHub
parent
61d32c9aac
commit
7263aa0028
5 changed files with 323 additions and 51 deletions
|
|
@ -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
|
||||
),
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue