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