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

@@ -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

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,

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`

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",