From d11917e89cd66ce43fb0bfa36d6c0b6909346998 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 13 Nov 2026 16:13:50 -0800 Subject: [PATCH] fix(proxy): preserve caller metadata on auth failures (#51745) * fix(proxy): keep x-litellm-spend-logs-metadata on auth-rejected failure spend logs Budget, invalid-key and other auth-stage rejections are raised before the request headers are merged into metadata, so the persisted failure row lost the caller's spend logs metadata header. The auth failure handler now parses it with the same pre-call helper and stamps it into the metadata bucket it already uses for the requester IP and user agent. * fix(proxy): type the spend logs metadata header parser honestly The parser takes the request headers as a str mapping and returns the parsed JSON as object, which is what successful requests already store. The auth failure path keeps only JSON objects. The precedence regression now also asserts the caller's request body is left untouched. * fix(proxy): stamp the auth failure spend logs metadata into the route's bucket A chat request whose body also carried litellm_metadata got the header in litellm_metadata, which the failure spend row never reads on chat routes. The header now goes into the bucket picked by metadata_variable_name_for_request for the route, and the failure row lifts spend_logs_metadata from that bucket the same way it already lifts the requester IP and user agent. The regression tests now assert on the built failure spend row. * fix(proxy): keep the auth failure spend logs metadata fix inside the auth handler Restores the failure spend row's lifted fields to main, so failures after auth on litellm_metadata routes keep their current logging. The auth handler now writes the parsed header into metadata, which the failure spend row reads on every route, and into the route's own bucket for callbacks on litellm_metadata routes. * fix(proxy): make the spend logs metadata header field read-only and test the real spend writer The header collector now sets spend_logs_metadata when it is built, so the field can be ReadOnly with the same stored value. The auth failure tests run the real DBSpendUpdateWriter and read the queued spend log row from a fake database client instead of an autospec writer. * test(proxy): read auth failure spend rows from a real unconnected PrismaClient Replaces the mocked database client and the class-level queue byte reset. The test now enqueues into a real PrismaClient's own queue and dequeues its row in cleanup, which keeps the shared byte counter balanced. --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- litellm/proxy/_types.py | 2 +- litellm/proxy/auth/auth_exception_handler.py | 35 +++- litellm/proxy/litellm_pre_call_utils.py | 12 +- .../proxy/auth/test_auth_exception_handler.py | 160 ++++++++++++++++++ 4 files changed, 202 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3d8a142636f..4d832febe92 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4676,7 +4676,7 @@ class LitellmMetadataFromRequestHeaders(TypedDict, total=False): Headers a user can pass that will get added to litellm metadata for the request """ - spend_logs_metadata: dict | None + spend_logs_metadata: ReadOnly[object] agent_id: str | None trace_id: str | None session_id: str | None diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index a004f6df392..cadaba7bc44 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -90,6 +90,14 @@ def _get_user_agent(request: Request) -> str | None: return request.headers.get("user-agent") +def _get_spend_logs_metadata(request: Request) -> Mapping[str, object] | None: + if "headers" not in request.scope: + return None + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + return as_str_mapping(LiteLLMProxyRequestSetup.get_spend_logs_metadata_from_request_headers(request.headers)) + + def _with_client_context( request_data: dict[str, object], requester_ip: str | None, user_agent: str | None ) -> dict[str, object]: @@ -108,6 +116,27 @@ def _with_client_context( return {**request_data, key: {**base, **stamped}} +def _metadata_bucket(request_data: Mapping[str, object], key: str) -> Mapping[str, object]: + metadata: Final = request_data.get(key) + return metadata if isinstance(metadata, Mapping) else EMPTY_MAPPING + + +def _with_spend_logs_metadata( + request_data: dict[str, object], route: str, spend_logs_metadata: Mapping[str, object] | None +) -> dict[str, object]: + """The failure spend row reads `metadata` on every route, while callbacks on + `litellm_metadata` routes read that bucket, so the header goes into both.""" + if spend_logs_metadata is None: + return request_data + from litellm.proxy.litellm_pre_call_utils import metadata_variable_name_for_request + + keys: Final = frozenset(("metadata", metadata_variable_name_for_request(route, request_data))) + return { + **request_data, + **{key: {**_metadata_bucket(request_data, key), "spend_logs_metadata": spend_logs_metadata} for key in keys}, + } + + def _escape_control_chars(value: str) -> str: return "".join(ch if ch.isprintable() else ch.encode("unicode_escape").decode("ascii") for ch in value) @@ -254,7 +283,11 @@ class UserAPIKeyAuthExceptionHandler: # Allow callbacks to transform the error response transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook( - request_data=_with_client_context(request_context, requester_ip, user_agent), + request_data=_with_spend_logs_metadata( + _with_client_context(request_context, requester_ip, user_agent), + normalize_request_route(route), + _get_spend_logs_metadata(request), + ), original_exception=e, user_api_key_dict=user_api_key_dict, error_type=ProxyErrorTypes.auth_error, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index f2f06df3ca4..63e0e6a4cc1 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1820,7 +1820,7 @@ class LiteLLMProxyRequestSetup: return None @staticmethod - def _get_spend_logs_metadata_from_request_headers(headers: dict) -> dict | None: + def get_spend_logs_metadata_from_request_headers(headers: Mapping[str, str]) -> object: """ Get the `spend_logs_metadata` from the request headers. """ @@ -2062,12 +2062,14 @@ class LiteLLMProxyRequestSetup: """ from litellm.proxy._types import LitellmMetadataFromRequestHeaders - metadata_from_headers: Final = LitellmMetadataFromRequestHeaders() request_body: Final = cast(Mapping[str, object], data) # cast-ok: request body is a str-keyed JSON object _drop_client_session_metadata(request_body) - spend_logs_metadata: Final = LiteLLMProxyRequestSetup._get_spend_logs_metadata_from_request_headers(headers) - if spend_logs_metadata is not None: - metadata_from_headers["spend_logs_metadata"] = spend_logs_metadata + spend_logs_metadata: Final = LiteLLMProxyRequestSetup.get_spend_logs_metadata_from_request_headers(headers) + metadata_from_headers: Final = ( + LitellmMetadataFromRequestHeaders(spend_logs_metadata=spend_logs_metadata) + if spend_logs_metadata is not None + else LitellmMetadataFromRequestHeaders() + ) ######################################################################################### # Finally update the requests metadata with the `metadata_from_headers` diff --git a/tests/unit/proxy/auth/test_auth_exception_handler.py b/tests/unit/proxy/auth/test_auth_exception_handler.py index 2523284ec2c..804e0d78967 100644 --- a/tests/unit/proxy/auth/test_auth_exception_handler.py +++ b/tests/unit/proxy/auth/test_auth_exception_handler.py @@ -1,4 +1,5 @@ import asyncio +import copy import json from collections.abc import Mapping from typing import Final @@ -39,7 +40,13 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler, _as_proxy_exception +from litellm.proxy.auth.auth_utils import is_invalid_virtual_key_error from litellm.proxy.auth.model_access_denied import ModelAccessDeniedHTTPException +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.db.request_error_tracking import RequestErrorAccumulator +from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger +from litellm.proxy.utils import PrismaClient, dequeue_spend_logs +from litellm.responses.streaming_iterator import drain_logging_tasks class _EngineHttp500: @@ -1140,6 +1147,159 @@ def _marked_malformed_key_error() -> HTTPException: return error +async def _reject_and_read_failure_row( + monkeypatch: pytest.MonkeyPatch, + error: Exception, + route: str, + headers: dict[str, str], + request_data: dict[str, object], + recorder: _AuthFailureRecorder | None = None, +) -> tuple[ProxyException, Mapping[str, object]]: + from litellm.proxy import proxy_server + from litellm.proxy.db import db_spend_update_writer as writer_module + + database: Final = PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_server.proxy_logging_obj) + monkeypatch.setattr(database, "spend_log_transactions", []) + monkeypatch.setattr(database, "_spend_log_transactions_lock", asyncio.Lock()) + spend_writer: Final = DBSpendUpdateWriter() + monkeypatch.setattr(writer_module, "request_error_accumulator", RequestErrorAccumulator()) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + monkeypatch.setattr( + litellm, + "callbacks", + [*([recorder] if recorder is not None else []), ProxyDBLogger(spend_writer=lambda: spend_writer)], + ) + monkeypatch.setattr(proxy_server, "general_settings", {"allow_requests_on_db_unavailable": False}) + try: + with pytest.raises(ProxyException) as exc_info: + await UserAPIKeyAuthExceptionHandler.handle_authentication_error( + error, _http_request(headers=headers), request_data, route, None, "sk-test" + ) + finally: + await drain_logging_tasks() + rows: Final = await dequeue_spend_logs(database, len(database.spend_log_transactions)) + (row,) = rows + row_metadata: Final = json.loads(row["metadata"]) + assert isinstance(row_metadata, dict) + return exc_info.value, row_metadata + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "error_kind, expect_invalid_key", + [("budget", False), ("unauthorized", False), ("malformed_key", True)], +) +async def test_auth_failure_row_keeps_spend_logs_metadata_header( + monkeypatch: pytest.MonkeyPatch, error_kind: str, expect_invalid_key: bool +) -> None: + monkeypatch.setattr(litellm, "budget_exceeded_status_code", 429) + errors: Final = { + "budget": BudgetExceededError(current_cost=2, max_budget=1), + "unauthorized": HTTPException(status_code=401, detail="Invalid API key"), + "malformed_key": _marked_malformed_key_error(), + } + error: Final = errors[error_kind] + expected: Final = _as_proxy_exception(error) + + raised, row_metadata = await _reject_and_read_failure_row( + monkeypatch, + error, + "/v1/chat/completions", + {"x-litellm-spend-logs-metadata": '{"marker": "rejected-call", "attempt": 2}'}, + {"model": "gpt-4o"}, + ) + + assert (raised.code, raised.type, raised.message) == (expected.code, expected.type, expected.message) + assert is_invalid_virtual_key_error(raised) is expect_invalid_key + assert row_metadata["status"] == "failure" + assert row_metadata["spend_logs_metadata"] == {"marker": "rejected-call", "attempt": 2} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "header, body, expected", + [ + pytest.param( + '{"marker": "header"}', + {"metadata": {"spend_logs_metadata": {"marker": "body"}}}, + {"marker": "header"}, + id="header_overrides_body_like_successful_requests", + ), + pytest.param( + '{"marker": "header"}', + {"metadata": {}, "litellm_metadata": {"spend_logs_metadata": {"marker": "body"}}}, + {"marker": "header"}, + id="chat_body_with_both_metadata_buckets", + ), + pytest.param("{}", {"metadata": {"spend_logs_metadata": {"marker": "body"}}}, {}, id="empty_object_header"), + pytest.param( + "{not json", {"metadata": {"spend_logs_metadata": {"marker": "body"}}}, {"marker": "body"}, id="malformed" + ), + pytest.param( + '["marker"]', {"metadata": {"spend_logs_metadata": {"marker": "body"}}}, {"marker": "body"}, id="non_object" + ), + pytest.param(None, {"metadata": {"spend_logs_metadata": {"marker": "body"}}}, {"marker": "body"}, id="missing"), + pytest.param("{not json", {"metadata": {}}, None, id="malformed_without_body_value"), + pytest.param(None, {"metadata": {}}, None, id="missing_without_body_value"), + ], +) +async def test_auth_failure_row_spend_logs_metadata_precedence( + monkeypatch: pytest.MonkeyPatch, + header: str | None, + body: dict[str, object], + expected: dict[str, object] | None, +) -> None: + request_data: Final[dict[str, object]] = {"model": "gpt-4o", **body} + original_request_data: Final = copy.deepcopy(request_data) + + _, row_metadata = await _reject_and_read_failure_row( + monkeypatch, + HTTPException(status_code=401, detail="Invalid API key"), + "/v1/chat/completions", + {"x-litellm-spend-logs-metadata": header} if header is not None else {}, + request_data, + ) + + assert row_metadata["spend_logs_metadata"] == expected + assert request_data == original_request_data + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route, body", + [ + pytest.param("/v1/responses", {"input": "hi", "metadata": {"native": "keep-me"}}, id="responses"), + pytest.param("/v1/messages", {"messages": [], "metadata": {"native": "keep-me"}}, id="messages"), + ], +) +async def test_auth_failure_row_keeps_spend_logs_metadata_header_on_litellm_metadata_routes( + monkeypatch: pytest.MonkeyPatch, route: str, body: dict[str, object] +) -> None: + request_data: Final[dict[str, object]] = {"model": "gpt-4o", **body} + original_request_data: Final = copy.deepcopy(request_data) + recorder: Final = _AuthFailureRecorder() + + _, row_metadata = await _reject_and_read_failure_row( + monkeypatch, + BudgetExceededError(current_cost=2, max_budget=1), + route, + {"x-litellm-spend-logs-metadata": '{"marker": "routed"}'}, + request_data, + recorder, + ) + + assert row_metadata["spend_logs_metadata"] == {"marker": "routed"} + assert recorder.data is not None + litellm_metadata: Final = recorder.data["litellm_metadata"] + metadata: Final = recorder.data["metadata"] + assert isinstance(litellm_metadata, Mapping) + assert isinstance(metadata, Mapping) + assert litellm_metadata["spend_logs_metadata"] == {"marker": "routed"} + assert metadata["native"] == "keep-me" + assert request_data == original_request_data + + @pytest.mark.asyncio @pytest.mark.parametrize( "auth_error,expect_traceback,expect_level",