From 1b5dacc717ee3d26eeac603b74cf93a75544a540 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 16 Sep 2026 18:05:38 +0000 Subject: [PATCH] fix(proxy): tolerate malformed auth spans and read the OTel span from request state for custom auth Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/litellm_pre_call_utils.py | 14 +++++-- .../proxy/test_litellm_pre_call_utils.py | 38 ++++++++++++++++--- 2 files changed, 43 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 19f7d075922..18f4280c01c 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -111,9 +111,13 @@ def _trace_id_from_traceparent(traceparent: str) -> str | None: def _trace_id_from_otel_span(span: "OtelSpan | None") -> str | None: if span is None: return None - span_context: Final = span.get_span_context() - trace_id: Final = span_context.trace_id - if not span_context.is_valid or not isinstance(trace_id, int): + try: + span_context: Final = span.get_span_context() + is_valid: Final = span_context.is_valid + trace_id: Final = span_context.trace_id + except AttributeError: + return None + if not is_valid or not isinstance(trace_id, int): return None return format(trace_id, "032x") @@ -2074,7 +2078,9 @@ async def add_litellm_data_to_request( add_otel_trace_id_to_request( data=data, _metadata_variable_name=_metadata_variable_name, - parent_otel_span=user_api_key_dict.parent_otel_span, + parent_otel_span=user_api_key_dict.parent_otel_span + if user_api_key_dict.parent_otel_span is not None + else getattr(request.state, "parent_otel_span", None), ) apply_missing_session_id_policy( data=data, diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 07365500016..72668dd3528 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -560,6 +560,7 @@ def _batches_request_mock() -> MagicMock: request_mock.headers = {"Content-Type": "application/json"} request_mock.client = MagicMock() request_mock.client.host = "127.0.0.1" + request_mock.state.parent_otel_span = None return request_mock @@ -3578,6 +3579,28 @@ async def test_add_litellm_data_to_request_defaults_trace_id_to_otel_server_span assert "litellm_session_id" not in data +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_falls_back_to_request_state_otel_span(): + """Custom auth hooks return a UserAPIKeyAuth without parent_otel_span even + though user_api_key_auth already opened the server span on request.state, + so the fallback must read the span from there or custom-auth requests would + keep getting an unrelated session id.""" + otel_trace_id: Final = 0x4BF92F3577B34DA6A3CE929D0E0E4736 + request_mock: Final = _request_mock_without_trace_headers() + request_mock.state.parent_otel_span = _otel_span_with_trace_id(otel_trace_id) + + data: Final = await add_litellm_data_to_request( + data={"model": "gpt-5.6"}, + request=request_mock, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key", parent_otel_span=None), + proxy_config=MagicMock(), + general_settings={}, + ) + + assert data["litellm_trace_id"] == format(otel_trace_id, "032x") + assert data["metadata"]["trace_id"] == format(otel_trace_id, "032x") + + @pytest.mark.asyncio async def test_add_litellm_data_to_request_otel_span_does_not_override_caller_trace_id(): """A caller's own trace identity (x-litellm-trace-id header or body @@ -3650,12 +3673,17 @@ async def test_add_litellm_data_to_request_otel_span_fills_empty_body_trace_id(e @pytest.mark.asyncio -@pytest.mark.parametrize("parent_otel_span", [None, "invalid_span", "not_a_span"]) +@pytest.mark.parametrize("parent_otel_span", [None, "invalid_span", "not_a_span", "plain_string"]) async def test_add_litellm_data_to_request_no_trace_id_without_valid_otel_span(parent_otel_span): - """No OTel span (OTel off), a span with an invalid context, or an object - that only quacks like a span (auth is typed loosely and often stubbed) must - leave litellm_trace_id unset so downstream keeps generating its own id.""" - span: Final = {"invalid_span": INVALID_SPAN, "not_a_span": MagicMock()}.get(parent_otel_span) + """No OTel span (OTel off), a span with an invalid context, an object that + only quacks like a span, or a value that is not a span at all (custom auth + is typed loosely and can hand back anything) must leave litellm_trace_id + unset, and never fail the request, so downstream keeps generating its own id.""" + span: Final = { + "invalid_span": INVALID_SPAN, + "not_a_span": MagicMock(), + "plain_string": "not-a-span", + }.get(parent_otel_span) data: Final = await add_litellm_data_to_request( data={"model": "gpt-5.6"}, request=_request_mock_without_trace_headers(),