scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix(otel): stamp an MCP tool failure on the request that carried it (#34551)

Browse source 
A failed MCP tool call aimed its error.* attributes at request_root_span(),
a ContextVar written on the ASGI request task. A stateful streamable-HTTP
session runs every message on the single task the session's initialize POST
spawned, so inside the message handler that ContextVar still holds the
initialize request's SERVER span. That span ended long ago, so the SDK
dropped every write (five 'Setting attribute on ended span' warnings plus
set_status and _add_event per failed call) and the POST that actually
failed carried no error at all. The identity attributes seeded onto the
server span went the same way.

Publish the live transport span on the ASGI scope of the request being
handled and read it back in the message handler through req_ctx.request,
the Request the streamable-HTTP transport attaches to each message. That
replaces the session-scoped field with a per-message one: a JSON-RPC
response POST deliberately skips the per-session lock, since it can arrive
while the tool call awaiting it is still in flight, so a field on the
shared auth object could be overwritten mid-call and send the tool call's
telemetry to the response's request. A scope also dies with its request
rather than holding a finished span on idle session state.

Publishing re-anchors the request root for the message so guardrail spans
and identity seeding follow, and only a transport still open for writes is
anchored or stamped: a notification POST can answer before the session task
is done, and moving dropped writes from one finished span to another is no
fix. Live capture goes from seven ended-span warnings and an unmarked
transaction to zero warnings and ERROR on the POST that carried the call.
This commit is contained in:
Yassin Kortam 2026-07-25 10:32:53 -07:00 • committed by GitHub
parent fa9e0f180c
commit 502d3609af
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 218 additions and 99 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/integrations/otel/logger.py
View file

@ -17,6 +17,7 @@ from litellm.integrations.otel.model.baggage import promoted_baggage
from litellm.integrations.otel.model.config import OpenTelemetryV2Config
from litellm.integrations.otel.plumbing.context import (
is_recordable_span,
mcp_message_transport_span,
request_root_span,
resolve_mcp_span_context,
resolve_parent_context,
@ -641,8 +642,14 @@ class OpenTelemetryV2(CustomLogger):
endpoint, auth failure), so the failed request carries the same error keys
a failed LLM call does. v1's ``OpenTelemetry`` implemented this same hook;
v2 lost it when it stopped subclassing ``OpenTelemetry``, which is the
LIT-4179 regression for pre-call failures."""
span = request_root_span() or user_api_key_dict.parent_otel_span
LIT-4179 regression for pre-call failures.
An MCP message is handled on the session's task, where the request-root
anchor is whatever request opened the session, so prefer the transport the
gateway published for this specific message. Without that, a failed tool
call aimed its error at the ``initialize`` request's finished span and the
SDK dropped it, leaving the POST that actually failed unmarked."""
span = mcp_message_transport_span() or request_root_span() or user_api_key_dict.parent_otel_span
if span is None or not is_recordable_span(span):
return None
stamp_error(span, _span_error_from_exception(original_exception, traceback_str=traceback_str))

78
litellm/integrations/otel/plumbing/context.py
View file

@ -79,47 +79,67 @@ 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.
# The transport span of the HTTP request carrying the CURRENT MCP message.
#
# ``_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
# Reading it from the message handler parents every tool call in the session to the
# first request's server span and aims that call's ``error.*`` at it — a span that
# ended long ago, so the SDK drops the write and the failure reaches no request at
# all. 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 and the failure hook.
_mcp_message_transport_span: "ContextVar[Span | None]" = ContextVar(
"litellm_otel_mcp_message_transport_span", default=None
)
def set_mcp_message_transport_span_context(
span_context: "SpanContext | None",
) -> "Token[SpanContext | None]":
def set_mcp_message_transport_span(span: object) -> "Token[Span | None]":
"""Publish the transport span of the request carrying the current MCP message.
Also re-anchors the request root, so everything else the message emits or stamps
— the identity attributes seeded onto the server span, a guardrail span, a
proxy-level failure — lands on this request instead of on the one that opened
the session. The MCP SDK dispatches each message on its own task, so the anchor
is scoped to this message; the handler re-publishes it for the next one either
way. Only a transport still open for writes is anchored: replacing the anchor
with a request that already answered would just move the dropped writes from one
finished span to another.
Takes ``object`` because the gateway reads it back out of the ASGI scope, whose
values are untyped; anything that is not a usable span is stored as ``None``
rather than trusted.
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)
transport = span if isinstance(span, Span) and is_recordable_span(span) else None
if transport is not None and transport.is_recording():
set_request_root_span(transport)
return _mcp_message_transport_span.set(transport)
def reset_mcp_message_transport_span_context(token: "Token[SpanContext | None]") -> None:
_mcp_message_transport_span_context.reset(token)
def reset_mcp_message_transport_span(token: "Token[Span | None]") -> None:
_mcp_message_transport_span.reset(token)
def request_root_span_context() -> "SpanContext | None":
"""The anchored request root span's context, safe to hand to another task.
def mcp_message_transport_span() -> "Span | None":
"""The published transport span, only while it is still open for writes.
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.
Recording — not merely valid — is the bar here because this span is the target
of ``error.*`` stamping from another task, and the publisher's validity check
cannot speak for a span that has since ended. A finished span keeps a valid
context forever, so it would otherwise be handed back for a write the SDK then
refuses. The POST carrying a ``tools/call`` stays open until the result is
written, so it is recording for the life of the call; a notification POST can
answer first, and this returns ``None`` for it rather than writing into the void.
"""
span = request_root_span()
return span.get_span_context() if span is not None else None
span = _mcp_message_transport_span.get()
if span is None or not span.is_recording():
return None
return span
def _mcp_transport_span_context() -> "SpanContext | None":
@ -127,12 +147,16 @@ def _mcp_transport_span_context() -> "SpanContext | None":
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).
request task itself (the REST MCP endpoints, the SDK). Parenting and linking
only need the immutable context, and unlike ``mcp_message_transport_span`` they
stay correct against a transport that has already finished, so this does not
require the span to still be recording.
"""
published = _mcp_message_transport_span_context.get()
if published is not None and published.is_valid:
return published
return request_root_span_context()
published = _mcp_message_transport_span.get()
if published is not None:
return published.get_span_context()
span = request_root_span()
return span.get_span_context() if span is not None else None
def set_request_baggage(values: Mapping[str, str], context: Context | None = None) -> Context:

9
litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py
View file

@ -1,12 +1,9 @@
from typing import TYPE_CHECKING, Dict, List, Optional
from typing import 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):
"""
@ -19,8 +16,6 @@ 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__(
@ -33,7 +28,6 @@ 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
@ -43,4 +37,3 @@ 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

104
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -15,7 +15,6 @@ import types
import uuid
from datetime import datetime
from typing import (
TYPE_CHECKING,
Any,
AsyncIterator,
Callable,
@ -107,9 +106,9 @@ _MAX_STATEFUL_SESSIONS_PER_OWNER = 100
# prevents an authenticated client from forcing the proxy to buffer an
# arbitrarily large body just to make a routing decision.
_MCP_ROUTING_PEEK_MAX_BYTES = 4096
if TYPE_CHECKING:
from opentelemetry.trace import SpanContext
# ASGI scope key holding the tracing span of the request carrying an MCP
# message, written on the request task and read back by the message handler.
_MCP_TRANSPORT_SPAN_SCOPE_KEY = "litellm_otel_transport_span"
def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None:
@ -294,52 +293,78 @@ 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.
def _otel_publish_transport_span_on_scope(scope: Scope) -> None:
"""Record this request's tracing span on its own ASGI scope.
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
and read back by the MCP message handler through ``req_ctx.request`` — the
``Request`` the transport attaches to each message. 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.
The scope, not the shared session auth context: a JSON-RPC *response* POST
deliberately skips the per-session lock (it can arrive while the tool call that
awaits it is still in flight), so a field on that shared object would be
overwritten mid-call and the tool call would attribute itself to the response's
request. A scope belongs to exactly one request and dies with it, which also
keeps a finished span from being retained by an idle session.
The live span, not just its context: a failed tool call stamps ``error.*`` on it,
which needs a span still open for writes. Lazily imported so opentelemetry stays
an optional dependency; a no-op when otel_v2 is unavailable or no request span is
anchored."""
try:
from litellm.integrations.otel.plumbing.context import (
request_root_span_context,
request_root_span,
)
return request_root_span_context()
span = request_root_span()
except ImportError:
return
if span is not None:
scope[_MCP_TRANSPORT_SPAN_SCOPE_KEY] = span
def _otel_transport_span_from_message(req_ctx: object) -> object:
"""The tracing span of the HTTP request that carried this MCP message.
Read off that request's ASGI scope, reached through the ``Request`` the
streamable-HTTP transport attaches to each message, so it is this message's
transport and not whichever request happens to have touched the session last.
Returns whatever the scope holds; the otel plumbing validates it."""
request = getattr(req_ctx, "request", None)
scope = getattr(request, "scope", None)
if not isinstance(scope, Mapping):
return None
return scope.get(_MCP_TRANSPORT_SPAN_SCOPE_KEY)
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:
def _otel_set_mcp_transport_span(span: object) -> object:
"""Publish the current message's transport span, which the otel_v2 MCP span
attaches to and a failed tool call stamps its error on. Returns a reset token,
or ``None`` when otel_v2 is unavailable."""
if span is None:
return None
try:
from litellm.integrations.otel.plumbing.context import (
set_mcp_message_transport_span_context,
set_mcp_message_transport_span,
)
return set_mcp_message_transport_span_context(span_context)
return set_mcp_message_transport_span(span)
except ImportError:
return None
def _otel_reset_mcp_transport_span_context(token: object) -> None:
"""Paired with ``_otel_set_mcp_transport_span_context``."""
def _otel_reset_mcp_transport_span(token: object) -> None:
"""Paired with ``_otel_set_mcp_transport_span``."""
if token is None:
return
try:
from litellm.integrations.otel.plumbing.context import (
reset_mcp_message_transport_span_context,
reset_mcp_message_transport_span,
)
reset_mcp_message_transport_span_context(token)
reset_mcp_message_transport_span(token)
except ImportError:
return
@ -710,18 +735,6 @@ 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]":
"""
@ -742,7 +755,7 @@ if MCP_AVAILABLE:
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())
_transport_token = _otel_set_mcp_transport_span(_otel_transport_span_from_message(req_ctx))
# Get user authentication from context variable
(
user_api_key_auth,
@ -798,7 +811,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_transport_span(_transport_token)
_otel_reset_mcp_trace_carrier(_trace_token)
if _session_reset_token is not None:
active_mcp_session_var.reset(_session_reset_token)
@ -976,7 +989,7 @@ if MCP_AVAILABLE:
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())
_transport_token = _otel_set_mcp_transport_span(_otel_transport_span_from_message(req_ctx))
# Validate arguments
(
user_api_key_auth,
@ -1115,7 +1128,7 @@ if MCP_AVAILABLE:
return response
finally:
_otel_reset_mcp_transport_span_context(_transport_token)
_otel_reset_mcp_transport_span(_transport_token)
_otel_reset_mcp_trace_carrier(_trace_token)
if _session_reset_token is not None:
active_mcp_session_var.reset(_session_reset_token)
@ -4260,6 +4273,7 @@ if MCP_AVAILABLE:
_increment_active_request_session(initialized_session_id)
async def _dispatch() -> None:
_otel_publish_transport_span_on_scope(scope)
auth_user = _set_or_update_auth_context(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
@ -4271,7 +4285,6 @@ 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:
@ -4496,7 +4509,6 @@ 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
@ -4505,7 +4517,6 @@ 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],
@ -4515,7 +4526,6 @@ 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.
@ -4526,7 +4536,6 @@ 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,
@ -4536,7 +4545,6 @@ 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
@ -4552,7 +4560,6 @@ 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:
@ -4567,7 +4574,6 @@ 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,
@ -4578,7 +4584,6 @@ 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
@ -4590,7 +4595,6 @@ 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(

115
tests/test_litellm/integrations/otel/test_otel_v2_logger.py
View file

@ -29,9 +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,
reset_mcp_message_transport_span,
set_mcp_message_trace_carrier,
set_mcp_message_transport_span_context,
set_mcp_message_transport_span,
set_request_root_span,
)
from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402
@ -58,11 +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)
_otel_context._mcp_message_transport_span.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)
_otel_context._mcp_message_transport_span.set(None)
def _payload(**overrides):
@ -580,15 +580,13 @@ def test_mcp_span_nests_under_this_messages_transport_not_the_session_opener(
)
async def session_task():
token = set_mcp_message_transport_span_context(
this_message.get_span_context()
)
token = set_mcp_message_transport_span(this_message)
try:
await logger.async_log_success_event(
{"standard_logging_object": make_payload()}, None, None, None
)
finally:
reset_mcp_message_transport_span_context(token)
reset_mcp_message_transport_span(token)
async def initialize_request():
# The anchor the session task inherits is the one ``initialize`` left behind;
@ -750,9 +748,7 @@ def test_mcp_span_links_this_messages_transport_when_context_is_propagated():
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()
)
transport_token = set_mcp_message_transport_span(this_message)
try:
asyncio.run(
logger.async_log_success_event(
@ -760,7 +756,7 @@ def test_mcp_span_links_this_messages_transport_when_context_is_propagated():
)
)
finally:
reset_mcp_message_transport_span_context(transport_token)
reset_mcp_message_transport_span(transport_token)
reset_mcp_message_trace_carrier(trace_token)
session_opener.end()
this_message.end()
@ -1022,6 +1018,101 @@ def test_async_post_call_failure_hook_falls_back_to_user_api_key_parent_span():
assert span.attributes["litellm.provider.error.code"] == "401"
def test_async_post_call_failure_hook_stamps_the_mcp_messages_own_transport():
"""A failed MCP tool call is handled on the session's task, where the request
root anchor is still the request that opened the session — an ended span, so the
SDK dropped the write and the POST that actually failed carried no error at all.
The hook must stamp the transport the gateway published for this message."""
from litellm.proxy._types import UserAPIKeyAuth
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(this_message)
try:
await logger.async_post_call_failure_hook(
request_data={},
original_exception=_proxy_exc("Authorization failed for tool 'x'", 403),
user_api_key_dict=UserAPIKeyAuth(),
)
finally:
reset_mcp_message_transport_span(token)
async def initialize_request():
set_request_root_span(session_opener)
await asyncio.create_task(session_task())
asyncio.run(initialize_request())
session_opener.end()
this_message.end()
by_id = {s.context.span_id: s for s in exporter.get_finished_spans()}
failed = by_id[this_message.get_span_context().span_id]
opener = by_id[session_opener.get_span_context().span_id]
assert failed.attributes["error.type"] == "ProxyException"
assert failed.status.status_code is StatusCode.ERROR
assert "error.type" not in opener.attributes
assert opener.status.status_code is not StatusCode.ERROR
def test_mcp_message_transport_reanchors_request_level_spans():
"""Publishing the message's transport also re-anchors the request root, so
everything else the message emits lands on the request that carried it. Without
that, a guardrail run during a tool call — like the identity attributes seeded
onto the server span — attaches to the request that 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)
async def session_task():
token = set_mcp_message_transport_span(this_message)
try:
logger.emit_guardrail_span({"guardrail_name": "my_guard", "guardrail_status": "success"})
finally:
reset_mcp_message_transport_span(token)
async def initialize_request():
set_request_root_span(session_opener)
await asyncio.create_task(session_task())
asyncio.run(initialize_request())
session_opener.end()
this_message.end()
guard = next(s for s in exporter.get_finished_spans() if s.name == "execute_guardrail my_guard")
assert guard.parent.span_id == this_message.get_span_context().span_id
assert guard.parent.span_id != session_opener.get_span_context().span_id
def test_async_post_call_failure_hook_skips_a_transport_that_already_answered():
"""The published transport is only writable while its request is open. A
notification POST answers before the session task is done with the message, and
writing to the finished span is a no-op the SDK logs and discards, so the hook
must fall through to the anchor instead of aiming at it."""
from litellm.proxy._types import UserAPIKeyAuth
logger, exporter = _logger()
anchor = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME)
answered = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME)
answered.end()
set_request_root_span(anchor)
token = set_mcp_message_transport_span(answered)
try:
asyncio.run(
logger.async_post_call_failure_hook(
request_data={},
original_exception=_proxy_exc("boom", 403),
user_api_key_dict=UserAPIKeyAuth(),
)
)
finally:
reset_mcp_message_transport_span(token)
anchor.end()
by_id = {s.context.span_id: s for s in exporter.get_finished_spans()}
assert by_id[anchor.get_span_context().span_id].attributes["error.type"] == "ProxyException"
assert "error.type" not in by_id[answered.get_span_context().span_id].attributes
def test_record_error_attributes_on_span_decorates_without_ending():
"""PATH A: a failure that dies before any LLM-call span (malformed body,
validation) is stamped onto the instrumentor-owned SERVER span. The method must