mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-20 00:11:50 +00:00
chore(otel/v2): merge latest litellm_internal_staging
Resolve the import-list conflict in test_otel_v2_logger.py by keeping both staging's set_mcp_message_transport_span_context and the PR's set_request_destinations. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
shivam
commit
7e34a24e1d
7 changed files with 337 additions and 52 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 TYPE_CHECKING, 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,
|
||||
)
|
||||
|
|
@ -100,6 +107,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
|
||||
|
|
@ -160,33 +223,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:
|
||||
|
|
|
|||
|
|
@ -123,7 +123,8 @@ class VertexAIBatchTransformation:
|
|||
Gets the output file id from the Vertex AI Batch response
|
||||
"""
|
||||
|
||||
output_file_id: str = response.get("outputInfo", OutputInfo()).get("gcsOutputDirectory", "")
|
||||
output_info = response.get("outputInfo") or OutputInfo()
|
||||
output_file_id: str = output_info.get("gcsOutputDirectory", "")
|
||||
if output_file_id:
|
||||
output_file_id = output_file_id.rstrip("/") + "/predictions.jsonl"
|
||||
if output_file_id and output_file_id != "/predictions.jsonl":
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -30,7 +30,9 @@ from litellm.integrations.otel.plumbing import providers # noqa: E402
|
|||
from litellm.integrations.otel.plumbing.context import ( # noqa: E402
|
||||
_request_destinations,
|
||||
reset_mcp_message_trace_carrier,
|
||||
reset_mcp_message_transport_span_context,
|
||||
set_mcp_message_trace_carrier,
|
||||
set_mcp_message_transport_span_context,
|
||||
set_request_destinations,
|
||||
set_request_root_span,
|
||||
)
|
||||
|
|
@ -77,9 +79,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):
|
||||
|
|
@ -549,15 +553,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
|
||||
|
|
@ -570,11 +574,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)
|
||||
|
|
@ -662,10 +730,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
|
||||
|
|
@ -682,9 +751,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
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -226,6 +226,18 @@ def test_get_output_file_id_empty_output_info_falls_through_to_output_config():
|
|||
assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl"
|
||||
|
||||
|
||||
def test_get_output_file_id_output_info_explicit_none_falls_through_to_output_config():
|
||||
resp = {
|
||||
"outputInfo": None,
|
||||
"outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://b/cfg"}},
|
||||
}
|
||||
assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl"
|
||||
|
||||
|
||||
def test_get_output_file_id_output_info_explicit_none_and_no_output_config():
|
||||
assert T._get_output_file_id_from_vertex_ai_batch_response({"outputInfo": None}) == ""
|
||||
|
||||
|
||||
def test_get_output_file_id_no_output_info_and_no_output_config():
|
||||
assert T._get_output_file_id_from_vertex_ai_batch_response({}) == ""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue