From b60246c35267b5d41e78953ba1f6b9c1528c9554 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 13 Nov 2026 03:39:51 +0000 Subject: [PATCH] fix(proxy): keep call id, model and key logging on keys refused before the body is read (#51525) * Fix failure logging context for cached key preflight rejections Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Centralize preflight failure route classification and restore precedence Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * Bound preflight logging reads and tolerate model lookup failures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_exception_handler.py | 2 + litellm/proxy/auth/user_api_key_auth.py | 86 ++++++++++- litellm/proxy/utils.py | 25 +++- .../test_user_api_key_auth_request_flow.py | 141 ++++++++++++++++++ ..._utils_model_creation_and_error_logging.py | 20 ++- 5 files changed, 263 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 2ec2274db2b..a004f6df392 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -146,6 +146,7 @@ class UserAPIKeyAuthExceptionHandler: parent_otel_span: Span | None, api_key: str, resolved_identity: UserAPIKeyAuth | None = None, + dispatched_call_type: str | None = None, ) -> UserAPIKeyAuth: """ Handles Connection Errors when reading a Virtual Key from LiteLLM DB @@ -258,6 +259,7 @@ class UserAPIKeyAuthExceptionHandler: user_api_key_dict=user_api_key_dict, error_type=ProxyErrorTypes.auth_error, route=route, + dispatched_call_type=dispatched_call_type, ) # Use transformed exception if callback returned one, otherwise use original if transformed_exception is not None: diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index ada6d90c128..cad35b669c2 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -18,11 +18,13 @@ from functools import partial from types import MappingProxyType from typing import Final, NamedTuple, Protocol, Union, cast +import anyio import fastapi from fastapi import HTTPException, Request, WebSocket, status from fastapi.security.api_key import APIKeyHeader -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError from starlette.exceptions import WebSocketException +from starlette.requests import ClientDisconnect import litellm from litellm._internal_context import service_target, set_request_parent_otel_span @@ -2233,6 +2235,80 @@ async def _read_request_body_deferring_parse_failure( return populate_request_with_path_params(request_data=parsed_body, request=request), None +PREFLIGHT_FAILURE_JSON_BODY_MAX_BYTES: Final = 64 * 1024 +PREFLIGHT_FAILURE_JSON_BODY_TIMEOUT_SECONDS: Final = 1.0 + + +def _preflight_failure_model( + body: Mapping[str, object], route: str, request: Request, team_id: str | None +) -> str | Sequence[str] | None: + from litellm.proxy.proxy_server import llm_router + + try: + return _get_model_from_request_context(dict(body), route, request, llm_router, team_id) + except Exception: # noqa: BLE001 # logging enrichment must never replace the existing preflight rejection + return None + + +async def _preflight_failure_request_data( + request: Request, route: str, cached_key: UserAPIKeyAuth +) -> Mapping[str, object]: + from litellm.proxy.common_request_processing import resolve_litellm_call_id + from litellm.proxy.litellm_pre_call_utils import key_and_team_logging_fields + from litellm.proxy.proxy_server import proxy_config + from litellm.proxy.utils import ( + _call_type_for_route, # pyright: ignore[reportPrivateUsage] # share proxy failure route classification + ) + + body: Final = await _small_preflight_failure_json(request) + model: Final = _preflight_failure_model(body, route, request, cached_key.team_id) + call_type: Final = _call_type_for_route(route, method=request.method) + _seed_request_destinations(cached_key, request) + seed_request_identity(cached_key, model=model, request_data=body) + return { + **{key: value for key, value in body.items() if key != "model"}, + **({"model": model} if model is not None else {}), + "litellm_call_id": resolve_litellm_call_id(request.headers.get("x-litellm-call-id")), + "call_type": call_type, + **key_and_team_logging_fields(user_api_key_dict=cached_key, proxy_config=proxy_config), + } + + +async def _small_preflight_failure_json(request: Request) -> Mapping[str, object]: + declared: Final = request.headers.get("content-length", "") + content_type: Final = request.headers.get("content-type", "").split(";", 1)[0].strip().lower() + if ( + not declared.isdecimal() + or len(declared) > 10 + or not 0 < int(declared) <= PREFLIGHT_FAILURE_JSON_BODY_MAX_BYTES + or request.headers.get("transfer-encoding") is not None + or content_type != "application/json" + ): + return {} + buffered: Final = bytearray() + try: + with anyio.fail_after(PREFLIGHT_FAILURE_JSON_BODY_TIMEOUT_SECONDS): + async for chunk in request.stream(): + if len(buffered) + len(chunk) > PREFLIGHT_FAILURE_JSON_BODY_MAX_BYTES: + return {} + buffered.extend(chunk) + parsed: Final = _GENERAL_SETTINGS_ADAPTER.validate_json(buffered) + except (ValidationError, RuntimeError, ClientDisconnect, TimeoutError): + return {} + return { + **({"model": parsed["model"]} if isinstance(parsed.get("model"), str) else {}), + **{ + slot: { + name: value + for name, value in _GENERAL_SETTINGS_ADAPTER.validate_python(metadata).items() + if name in ("tags", "session_id", "trace_user_id", "trace_id", "trace_name") + } + for slot in ("metadata", "litellm_metadata") + if isinstance(metadata := parsed.get(slot), dict) + }, + } + + async def _reject_from_headers_before_reading_body( request: Request, api_key: str | None, custom_litellm_key_header: str | None, route: str ) -> None: @@ -2271,6 +2347,7 @@ async def _reject_from_headers_before_reading_body( return parent_otel_span: Final[Span | None] = getattr(request.state, "parent_otel_span", None) if cached_key.blocked is True: + request_data: Final = dict(await _preflight_failure_request_data(request, route, cached_key)) await UserAPIKeyAuthExceptionHandler.handle_authentication_error( e=ProxyException( message="Authentication Error, Key is blocked. Update via `/key/unblock` if you're an admin.", @@ -2279,11 +2356,12 @@ async def _reject_from_headers_before_reading_body( code=status.HTTP_401_UNAUTHORIZED, ), request=request, - request_data={}, + request_data=request_data, route=route, parent_otel_span=parent_otel_span, api_key=bearer_token, resolved_identity=cached_key, + dispatched_call_type=str(request_data["call_type"]) if request_data["call_type"] is not None else None, ) return limiter: Final = proxy_logging_obj.get_proxy_hook("parallel_request_limiter") @@ -2299,11 +2377,13 @@ async def _reject_from_headers_before_reading_body( try: await limiter.raise_if_over_limit_before_body(rate_limited) except HTTPException as rejection: + failure_data: Final = dict(await _preflight_failure_request_data(request, route, rate_limited)) transformed: Final = await proxy_logging_obj.post_call_failure_hook( # pyright: ignore[reportUnknownMemberType] # bare dict in sig - request_data={}, + request_data=failure_data, original_exception=rejection, user_api_key_dict=rate_limited, route=route, + dispatched_call_type=str(failure_data["call_type"]) if failure_data["call_type"] is not None else None, ) failure: Final = transformed if transformed is not None else rejection if isinstance(failure, ProxyException): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ee78db34d62..7ff2b9e236e 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1689,11 +1689,20 @@ def _deployment_attribution_for_model_group(model_group: object, team_id: str | ) -def _call_type_for_route(route: str | None) -> str | None: +def _call_type_for_route(route: str | None, *, method: str | None = None) -> str | None: """The route's call type when it maps to a single operation (its async and sync variants); - None for routes shared by several operations, since the method is not known here.""" + None for routes shared by several operations unless the method disambiguates them.""" if route is None: return None + normalized_route: Final = re.sub(r"^/v1(?:beta)?/", "/", route) + endpoint_call_type: Final = { + "/responses/compact": "acompact_responses", + "/interactions": "acreate_interaction", + }.get(normalized_route) + if endpoint_call_type is not None: + return endpoint_call_type + if normalized_route == "/videos" and method == "POST": + return "avideo_generation" call_types: Final = get_call_types_for_route(route) if not call_types: return None @@ -1769,7 +1778,14 @@ def _as_str_mapping(value: object) -> Mapping[str, object]: def _proxy_only_error_logged_metadata_keys(*, authenticated: bool, past_pre_call: bool) -> frozenset[str]: return frozenset( key - for key, logged in (("tags", authenticated), ("standard_logging_guardrail_information", past_pre_call)) + for key, logged in ( + ("tags", authenticated), + ("session_id", authenticated), + ("trace_user_id", authenticated), + ("trace_id", authenticated), + ("trace_name", authenticated), + ("standard_logging_guardrail_information", past_pre_call), + ) if logged ) @@ -4270,6 +4286,8 @@ class ProxyLogging: ) if litellm_logging_obj is None: from litellm._uuid import uuid + from litellm.proxy.litellm_pre_call_utils import key_and_team_logging_fields + from litellm.proxy.proxy_server import proxy_config request_data.setdefault("litellm_call_id", str(uuid.uuid4())) @@ -4280,6 +4298,7 @@ class ProxyLogging: **_proxy_only_error_log_fields( # pyright: ignore[reportArgumentType] # the allowlist never carries is_async_call request_fields, logged_metadata_keys ), + **key_and_team_logging_fields(user_api_key_dict=user_api_key_dict, proxy_config=proxy_config), ) request_data["litellm_logging_obj"] = litellm_logging_obj # rebind-ok: lifted then popped by the caller request_data.update({REJECTED_BEFORE_REQUEST_SETUP_KEY: REJECTED_BEFORE_REQUEST_SETUP}) diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index d0acb13230b..e76b2e85afe 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -14349,3 +14349,144 @@ async def test_blocked_cached_key_is_refused_before_the_body_is_read() -> None: assert refusal.code == "401", refusal.message assert "Key is blocked" in refusal.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize("body_mode", ("small-json", "large-unread", "stalled", "laya", "bespoke")) +@pytest.mark.parametrize("blocked", (True, False), ids=("blocked", "rate-limited")) +async def test_preflight_rejection_preserves_failure_logging_context(body_mode: str, blocked: bool) -> None: + from litellm.proxy.utils import hash_token + + api_key: Final = "sk-preflight-logging" + cache: Final = DualCache() + identity: Final = UserAPIKeyAuth( + token=hash_token(api_key), + blocked=blocked, + rpm_limit=1, + metadata={ + "logging": [{"callback_name": "langfuse", "callback_type": "success_and_failure", "callback_vars": {}}] + }, + ) + await cache_key_object( + hashed_token=hash_token(api_key), + user_api_key_obj=identity, + user_api_key_cache=cache, + proxy_logging_obj=None, + ) + limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache())) + if not blocked: + await limiter.async_pre_call_hook( + UserAPIKeyAuth(api_key=hash_token(api_key), rpm_limit=1), + DualCache(), + {"model": "requested-model", "messages": []}, + "acompletion", + ) + body: Final = json.dumps( + { + "model": "requested-model", + "metadata": { + "tags": ["caller-tag"], + "session_id": "caller-session", + "trace_user_id": "caller-user", + }, + "failure_callback": ["untrusted-callback"], + } + ).encode() + + upload: Final[asyncio.Queue[dict[str, object]]] = asyncio.Queue() + upload.put_nowait({"type": "http.request", "body": b"{", "more_body": True}) + + async def receive() -> dict[str, object]: + assert body_mode in ("small-json", "stalled"), "large refused bodies must not be read" + if body_mode == "stalled": + return await upload.get() + return {"type": "http.request", "body": body, "more_body": False} + + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": f"/{body_mode}/v1/systemone" if body_mode in ("laya", "bespoke") else "/v1/chat/completions", + "headers": [ + (b"content-type", b"application/json"), + ( + b"content-length", + str(len(body) if body_mode == "small-json" else 2 if body_mode == "stalled" else 8388608).encode(), + ), + (b"x-litellm-call-id", b"caller-call-id"), + ], + }, + receive=receive, + ) + logging_obj: Final = MagicMock(post_call_failure_hook=AsyncMock(return_value=None)) + logging_obj.get_proxy_hook.return_value = limiter + with ( + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.open_telemetry_logger", None), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), + patch("litellm.proxy.proxy_server.proxy_logging_obj", logging_obj), + pytest.raises(ProxyException) as failure, + ): + await asyncio.wait_for(user_api_key_auth(request=request, api_key=f"Bearer {api_key}"), timeout=3) + assert failure.value.code == ("401" if blocked else "429") + logging_obj.post_call_failure_hook.assert_awaited_once() + data: Final = logging_obj.post_call_failure_hook.call_args.kwargs["request_data"] + assert data["litellm_call_id"] == "caller-call-id" + assert data["failure_callback"] == ["langfuse"] + if body_mode not in ("laya", "bespoke"): + assert data["call_type"] == "acompletion" + if body_mode == "small-json": + assert data["model"] == "requested-model" + assert data["metadata"]["tags"] == ["caller-tag"] + assert data["metadata"]["session_id"] == "caller-session" + assert data["metadata"]["trace_user_id"] == "caller-user" + else: + assert "model" not in data + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "headers,body,reads", + ( + ((), b"{}", 0), + (((b"content-length", b"2"), (b"transfer-encoding", b"chunked")), b"{}", 0), + (((b"content-length", b"2"), (b"content-type", b"text/plain")), b"{}", 0), + (((b"content-length", b"bad"),), b"{}", 0), + (((b"content-length", b"10"),), b"{invalid", 1), + (((b"content-length", b"2"),), b"[]", 1), + (((b"content-length", b"2"),), b"x" * (64 * 1024 + 1), 1), + ), + ids=("no-length", "chunked", "not-json", "invalid-length", "invalid-json", "not-object", "actual-over-cap"), +) +async def test_preflight_logging_drops_unsafe_or_unparsable_bodies( + headers: tuple[tuple[bytes, bytes], ...], body: bytes, reads: int +) -> None: + from litellm.proxy.auth.user_api_key_auth import _small_preflight_failure_json + + receive: Final = AsyncMock( + return_value={ + "type": "http.request", + "body": body, + "more_body": len(body) > 64 * 1024, + } + ) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/chat/completions", + "headers": [ + *headers, + *( + () + if any(name == b"content-type" for name, _ in headers) + else ((b"content-type", b"application/json"),) + ), + ], + }, + receive=receive, + ) + assert await _small_preflight_failure_json(request) == {} + assert receive.await_count == reads diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index 5c399f96e83..e17d830c664 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -936,11 +936,19 @@ class TestPostCallFailureHookLiftsCallTypeAndStartTime: @pytest.mark.asyncio @pytest.mark.parametrize( - ("request_data", "route", "endpoint_call_type"), + ("request_data", "route", "endpoint_call_type", "dispatched_call_type"), [ - ({"model": "gemini-tts", "input": "Hello there", "voice": "Kore"}, "/v1/audio/speech", "aspeech"), - ({"model": "omni-moderation", "input": ["one", "two"]}, "/v1/moderations", "amoderation"), - ({"model": "gpt-transcribe"}, "/v1/audio/transcriptions", "atranscription"), + ( + {"model": "gemini-tts", "input": "Hello there", "voice": "Kore"}, + "/v1/audio/speech", + "aspeech", + "aspeech", + ), + ({"model": "omni-moderation", "input": ["one", "two"]}, "/v1/moderations", "amoderation", "aspeech"), + ({"model": "gpt-transcribe"}, "/v1/audio/transcriptions", "atranscription", "atranscription"), + ({"model": "test"}, "/v1/responses/compact", "acompact_responses", None), + ({"model": "test"}, "/v1beta/interactions", "acreate_interaction", None), + ({"model": "test"}, "/v1/videos", "avideo_generation", "avideo_generation"), ], ) async def test_request_rejected_by_the_proxy_records_the_endpoint_call_type( @@ -949,6 +957,7 @@ class TestPostCallFailureHookLiftsCallTypeAndStartTime: request_data: dict[str, object], route: str, endpoint_call_type: str, + dispatched_call_type: str | None, ): from types import SimpleNamespace @@ -967,7 +976,8 @@ class TestPostCallFailureHookLiftsCallTypeAndStartTime: request_data=dict(request_data), original_exception=HTTPException(status_code=429, detail="Rate limit exceeded"), user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", request_route=route), - dispatched_call_type=endpoint_call_type, + route=route, + dispatched_call_type=dispatched_call_type, ) finally: ProxyLogging._callback_capabilities_cache.clear()