Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 13 additions & 7 deletions litellm/integrations/otel/logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
is_mcp_list_tools,
is_mcp_tool_call,
)
from litellm.integrations.otel.model.semconv import Error
from litellm.integrations.otel.model.spans import SpanRole, span_role_for_service
from litellm.integrations.otel.model.utils import to_ns
from litellm.integrations.otel.plumbing.context import (
Expand Down Expand Up @@ -634,18 +635,23 @@ def record_error_attributes_on_span(
"""Stamp the v2 error.* attributes on the FastAPI-owned SERVER span for a
failure that dies before any LLM-call span exists (malformed body, auth /
validation rejection). Called from the proxy's global exception handler via
``_close_dangling_otel_server_span``. The instrumentor still owns the span's
status and lifecycle, so this only decorates it — never sets status, never
ends it — and emits no exception event, matching v1's SERVER-span behavior
and avoiding a duplicate of the event ``async_post_call_failure_hook`` or
the ``auth`` phase span already records."""
``_close_dangling_otel_server_span``, which swallows the exception into a
``JSONResponse`` so the instrumentor never sees it and leaves the span
``UNSET``; the status is set here instead (v1 did the same from the handler)
so a failed request reads as failed and not merely as a span carrying an
error message. The instrumentor still owns the span's lifecycle, so this
never ends it. The exception event is recorded only when nothing stamped
this span already — ``async_post_call_failure_hook`` and the ``auth`` phase
span record their own, and a second event would duplicate it — while the
attributes are always restamped so ``error.code`` stays pinned to the real
response status."""
if span is None or not is_recordable_span(span):
return
already_stamped: Final = Error.TYPE in (getattr(span, "attributes", None) or ())
stamp_error(
span,
_span_error_from_exception(exception, status_code=status_code),
record_event=False,
set_status=False,
record_event=not already_stamped,
)

async def async_post_call_failure_hook(
Expand Down
168 changes: 112 additions & 56 deletions litellm/proxy/auth/user_api_key_auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -1044,6 +1044,22 @@ def _ensure_parent_otel_span_on_request_state(request: Request) -> None:
request.state.parent_otel_span = parent_otel_span


async def _read_request_body_deferring_parse_failure(
request: Request,
) -> tuple[dict, ProxyException | None]:
"""Parse the body, returning a parse failure instead of raising it.

A body that fails to parse is still a request from a known caller, so auth
must run (resolving identity onto the request's trace) before the 400 goes
out; the caller re-raises the returned exception once identity is seeded.
"""
try:
parsed_body: Final = await _read_request_body(request=request)
except ProxyException as parse_exception:
return {}, parse_exception # mutable-ok: request_data is a plain dict across the whole auth path
return populate_request_with_path_params(request_data=parsed_body, request=request), None


async def _user_api_key_auth_builder(
request: Request,
api_key: str,
Expand Down Expand Up @@ -2516,6 +2532,72 @@ def _resolve_request_principal(request: Request, valid_token: UserAPIKeyAuth) ->
)


async def _authorize_authenticated_request(
user_api_key_auth_obj: UserAPIKeyAuth,
request: Request,
request_data: dict,
route: str,
api_key: str,
) -> UserAPIKeyAuth | None:
"""Authorize an already-authenticated request: disabled-route check, the single
``common_checks`` gate (which also reserves budget), and end-user fallback
resolution. Returns the auth object the exception handler recovered when a check
failed but the request may proceed anyway, else ``None``.
"""
## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ##
RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request)

# Single authorization point. Builder paths MUST NOT call common_checks.
# Route through the same exception handler the builder uses so
# authorization failures (ProxyException, or plain Exception from
# admin-only-route / model-access / budget checks) surface as
# ProxyException consistently with pre-refactor behavior.
try:
await _run_centralized_common_checks(
user_api_key_auth_obj=user_api_key_auth_obj,
request=request,
request_data=request_data,
route=route,
)
except Exception as e:
return await UserAPIKeyAuthExceptionHandler._handle_authentication_error(
e=e,
request=request,
request_data=request_data,
route=route,
parent_otel_span=user_api_key_auth_obj.parent_otel_span,
api_key=api_key,
resolved_identity=user_api_key_auth_obj,
)

# Defense-in-depth: ``_user_api_key_auth_builder`` has multiple early-return
# paths (no master key, /user/auth route, JWT short-circuits) that bypass
# the end-user resolution block. If those paths produced an auth obj
# without an ``end_user_id`` set, fall back to extracting from the request
# body so spend logs are still attributed correctly. Validation honours
# ``litellm.validate_end_user_id_in_db``.
if user_api_key_auth_obj.end_user_id is None:
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)

raw_end_user_id: Final = get_end_user_id_from_request_body(request_data, _safe_get_request_headers(request))
if raw_end_user_id is not None:
resolved_end_user_id: Final = await resolve_and_validate_end_user_id(
raw_end_user_id=raw_end_user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_auth_obj.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route=route,
)
if resolved_end_user_id is not None:
user_api_key_auth_obj.end_user_id = resolved_end_user_id
return None


@tracer.wrap()
async def user_api_key_auth(
request: Request,
Expand All @@ -2536,78 +2618,49 @@ async def user_api_key_auth(
# close, and the trace never reaches the backend.
_ensure_parent_otel_span_on_request_state(request)

request_data = await _read_request_body(request=request)
request_data = populate_request_with_path_params(request_data=request_data, request=request)
request_data, body_parse_exception = await _read_request_body_deferring_parse_failure(request=request)
route: Final[str] = get_request_route(request=request)
## CHECK IF ROUTE IS ALLOWED

# Run the whole auth phase inside a live ``auth`` span so the DB lookups it
# triggers (key/user/team object reads) nest under it instead of flattening
# onto the server span. No-op when OTel V2 isn't active.
with phase_span(f"auth {route}"):
user_api_key_auth_obj: Final = await _user_api_key_auth_builder(
request=request,
api_key=api_key,
azure_api_key_header=azure_api_key_header,
anthropic_api_key_header=anthropic_api_key_header,
google_ai_studio_api_key_header=google_ai_studio_api_key_header,
azure_apim_header=azure_apim_header,
request_data=request_data,
custom_litellm_key_header=custom_litellm_key_header,
)
user_api_key_auth_obj.budget_reservation = None

## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ##
RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request)

# Single authorization point. Builder paths MUST NOT call common_checks.
# Route through the same exception handler the builder uses so
# authorization failures (ProxyException, or plain Exception from
# admin-only-route / model-access / budget checks) surface as
# ProxyException consistently with pre-refactor behavior.
try:
await _run_centralized_common_checks(
user_api_key_auth_obj=user_api_key_auth_obj,
user_api_key_auth_obj: Final = await _user_api_key_auth_builder(
request=request,
api_key=api_key,
azure_api_key_header=azure_api_key_header,
anthropic_api_key_header=anthropic_api_key_header,
google_ai_studio_api_key_header=google_ai_studio_api_key_header,
azure_apim_header=azure_apim_header,
request_data=request_data,
route=route,
custom_litellm_key_header=custom_litellm_key_header,
)
except Exception as e:
return await UserAPIKeyAuthExceptionHandler._handle_authentication_error(
e=e,
except Exception:
# The body was read first, so a caller who sent both a malformed body and
# a rejected key used to get the 400; the response is unchanged, and the
# auth failure is still recorded on the trace by the handler that ran.
if body_parse_exception is not None:
raise body_parse_exception
raise
user_api_key_auth_obj.budget_reservation = None

# A body that never parsed is authenticated (so the trace carries identity
# and this ``auth`` span) but not authorized: there is no model to check it
# against, and budget reservation would increment live spend counters that
# only the endpoint's post-call path releases; the endpoint never runs, since
# the parse failure is raised below.
if body_parse_exception is None:
recovered_auth_obj: Final = await _authorize_authenticated_request(
user_api_key_auth_obj=user_api_key_auth_obj,
request=request,
request_data=request_data,
route=route,
parent_otel_span=user_api_key_auth_obj.parent_otel_span,
api_key=api_key,
resolved_identity=user_api_key_auth_obj,
)

# Defense-in-depth: ``_user_api_key_auth_builder`` has multiple early-return
# paths (no master key, /user/auth route, JWT short-circuits) that bypass
# the end-user resolution block. If those paths produced an auth obj
# without an ``end_user_id`` set, fall back to extracting from the request
# body so spend logs are still attributed correctly. Validation honours
# ``litellm.validate_end_user_id_in_db``.
if user_api_key_auth_obj.end_user_id is None:
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)

raw_end_user_id: Final = get_end_user_id_from_request_body(request_data, _safe_get_request_headers(request))
if raw_end_user_id is not None:
resolved_end_user_id: Final = await resolve_and_validate_end_user_id(
raw_end_user_id=raw_end_user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_auth_obj.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
route=route,
)
if resolved_end_user_id is not None:
user_api_key_auth_obj.end_user_id = resolved_end_user_id
if recovered_auth_obj is not None:
return recovered_auth_obj

# Identity is now resolved. Seed it AFTER the auth span closes so the Baggage
# persists on the request task (detaching the span's context token inside the
Expand All @@ -2619,6 +2672,9 @@ async def user_api_key_auth(
)
user_api_key_auth_obj.request_route = normalize_request_route(route)

if body_parse_exception is not None:
raise body_parse_exception

# Resolve caller identity once, here at the seam, into a single per-request
# Principal projected off the key object the builder already fetched (no
# second lookup). Downstream consumers read identity off this instead of
Expand Down
35 changes: 32 additions & 3 deletions tests/test_litellm/integrations/otel/test_otel_v2_logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -1195,8 +1195,13 @@ def test_async_post_call_failure_hook_skips_a_transport_that_already_answered():
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
not end the span or emit a duplicate exception event, and must pin error.code
to the real response status (not the exception's own code)."""
not end the span, and must pin error.code to the real response status (not the
exception's own code).

LIT-4780: the instrumentor never sees the exception (the proxy handler turns it
into a JSONResponse), so nothing else marks the span as failed; the status and
the exception event have to come from here or the trace shows the error message
on an otherwise successful-looking request."""
logger, exporter = _logger()
server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME)
logger.record_error_attributes_on_span(server, _proxy_exc("Invalid JSON body", 400), 422)
Expand All @@ -1206,7 +1211,31 @@ def test_record_error_attributes_on_span_decorates_without_ending():
assert span.attributes["error.type"] == "ProxyException"
assert span.attributes["error.message"] == "Invalid JSON body"
assert span.attributes["litellm.provider.error.code"] == "422"
assert all(e.name != "exception" for e in span.events)
assert span.status.status_code is StatusCode.ERROR
assert [e.name for e in span.events] == ["exception"]


def test_record_error_attributes_on_span_does_not_duplicate_an_already_stamped_error():
"""A failure that already went through ``async_post_call_failure_hook`` reaches
the exception handler too; the second stamp must keep one exception event while
still repinning error.code to the real response status."""
from litellm.proxy._types import UserAPIKeyAuth

logger, exporter = _logger()
server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME)
set_request_root_span(server)
exc = _proxy_exc("Authentication Error, invalid key", 401)
asyncio.run(
logger.async_post_call_failure_hook(
request_data={}, original_exception=exc, user_api_key_dict=UserAPIKeyAuth()
)
)
logger.record_error_attributes_on_span(server, exc, 400)
server.end()
(span,) = exporter.get_finished_spans()
assert [e.name for e in span.events] == ["exception"]
assert span.attributes["litellm.provider.error.code"] == "400"
assert span.status.status_code is StatusCode.ERROR


def test_record_error_attributes_on_span_ignores_below_400_and_missing_span():
Expand Down
Loading
Loading