From 9d4bab3b700c1170b7451bedfbda67cf830f58f3 Mon Sep 17 00:00:00 2001 From: zoroyihan7 Date: Tue, 8 Sep 2026 10:46:16 +0000 Subject: [PATCH 01/28] fix(responses): emit typed streaming failure events --- litellm/exceptions.py | 3 +- .../common_utils/responses_stream_errors.py | 119 ++++++++++++++++++ litellm/proxy/proxy_server.py | 21 +++- .../proxy/response_api_endpoints/endpoints.py | 6 +- litellm/responses/streaming_iterator.py | 5 + .../proxy_server/test_streaming_helpers.py | 85 ++++++++++++- .../response_api_endpoints/test_endpoints.py | 95 +++++++++++++- 7 files changed, 328 insertions(+), 6 deletions(-) create mode 100644 litellm/proxy/common_utils/responses_stream_errors.py diff --git a/litellm/exceptions.py b/litellm/exceptions.py index f9215267bf3..318b227ffa8 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -789,6 +789,7 @@ class APIError(openai.APIError): litellm_debug_info: str | None = None, max_retries: int | None = None, num_retries: int | None = None, + body: object | None = None, ): self.status_code = status_code self.message = f"litellm.APIError: {message}" @@ -799,7 +800,7 @@ class APIError(openai.APIError): self.num_retries = num_retries if request is None: request = httpx.Request(method="POST", url="https://api.openai.com/v1") - super().__init__(self.message, request=request, body=None) + super().__init__(self.message, request=request, body=body) def __str__(self): _message = self.message diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py new file mode 100644 index 00000000000..a3bee06e912 --- /dev/null +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -0,0 +1,119 @@ +import time +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from pydantic import BaseModel, ConfigDict + +from litellm._logging import redact_internal_details_from_client_message +from litellm._uuid import uuid +from litellm.exceptions import MidStreamFallbackError +from litellm.types.llms.openai import ResponseFailedEvent, ResponsesAPIResponse, ResponsesAPIStreamEvents + + +class _ResponseIdentity(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + id: str | None = None + model: str | None = None + created_at: int | None = None + + +class _StreamEvent(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + type: str | None = None + sequence_number: int | None = None + response: _ResponseIdentity | None = None + + +class _FailureDetails(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + message: str | None = None + code: str | int | None = None + type: str | None = None + status_code: int | None = None + + +def _original_failure(exception: Exception) -> Exception: + if isinstance(exception, MidStreamFallbackError) and exception.original_exception is not None: + return _original_failure(exception.original_exception) + return exception + + +def _response_error_code(details: _FailureDetails) -> str: + for value in (details.code, details.type): + if value == "insufficient_quota": + return "insufficient_quota" + if value in (429, "429") or isinstance(value, str) and value.startswith("rate_limit"): + return "rate_limit_exceeded" + if isinstance(details.code, str) and details.code and not details.code.isdecimal(): + return details.code + if details.status_code == 429: + return "rate_limit_exceeded" + return "server_error" + + +class ResponsesStreamErrorState: + def __init__(self) -> None: + self.response_id: str | None = None + self.model: str | None = None + self.created_at: int | None = None + self.sequence_number = -1 + self.terminal_emitted = False + + @staticmethod + def observe_chunk(chunk: object) -> _StreamEvent | None: + if not isinstance(chunk, (BaseModel, Mapping)): + return None + return _StreamEvent.model_validate(chunk) + + def mark_emitted(self, event: _StreamEvent | None) -> None: + if event is None: + return + if event.sequence_number is not None: + self.sequence_number = max(self.sequence_number, event.sequence_number) + if event.response is not None: + self.response_id = event.response.id or self.response_id + self.model = event.response.model or self.model + if event.response.created_at is not None: + self.created_at = event.response.created_at + if event.type in ("response.completed", "response.failed", "response.incomplete"): + self.terminal_emitted = True + + def format_failure(self, exception: Exception) -> str | None: + if self.terminal_emitted: + return None + original: Final = _original_failure(exception) + details: Final = _FailureDetails.model_validate(original) + response: Final = ResponsesAPIResponse.model_validate( + MappingProxyType( + { + "id": self.response_id or f"resp_{uuid.uuid4().hex}", + "object": "response", + "created_at": self.created_at if self.created_at is not None else int(time.time()), + "model": self.model, + "status": "failed", + "output": (), + "error": MappingProxyType( + { + "code": _response_error_code(details), + "message": redact_internal_details_from_client_message(details.message or str(original)), + } + ), + } + ) + ) + event: Final = ResponseFailedEvent.model_validate( + MappingProxyType( + { + "type": ResponsesAPIStreamEvents.RESPONSE_FAILED, + "response": response, + "sequence_number": self.sequence_number + 1, + } + ) + ) + payload: Final = event.model_dump_json(exclude_none=True) + self.terminal_emitted = True + return f"event: response.failed\ndata: {payload}\n\n" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 32b6b841af7..9f1d69acefa 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -385,6 +385,7 @@ from litellm.proxy.common_utils.periodic_reload_schedule import ( ) from litellm.proxy.common_utils.proxy_state import ProxyState from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob +from litellm.proxy.common_utils.responses_stream_errors import ResponsesStreamErrorState from litellm.proxy.common_utils.scheduled_job_stagger import ( apply_scheduled_job_stagger, attach_job_timing_logger, @@ -8723,10 +8724,13 @@ async def async_data_generator( user_api_key_dict: UserAPIKeyAuth, request_data: dict, request: Request | None = None, + *, + responses_stream_errors: bool = False, ): verbose_proxy_logger.debug("inside generator") stream_completed = False client_disconnected = False + error_state: Final = ResponsesStreamErrorState() if responses_stream_errors else None try: error_message: str | None = None requested_model_from_client: Final = _get_client_requested_model_for_streaming(request_data=request_data) @@ -8837,6 +8841,7 @@ async def async_data_generator( fallback_metadata_event_sent = True continue + responses_event: Final = error_state.observe_chunk(chunk) if error_state is not None else None raw_passthrough = False if isinstance(chunk, BaseModel): chunk = _serialize_streaming_chunk(chunk) @@ -8871,8 +8876,13 @@ async def async_data_generator( if not raw_passthrough: try: - yield _format_streaming_sse_chunk(chunk=chunk) + formatted_chunk: Final = _format_streaming_sse_chunk(chunk=chunk) + if error_state is not None: + error_state.mark_emitted(responses_event) + yield formatted_chunk except Exception as e: + if error_state is not None: + raise yield f"data: {e}\n\n" if pending_fallback_event: @@ -8922,6 +8932,12 @@ async def async_data_generator( e, ) + if error_state is not None: + stream_completed = True + error_frame: Final = error_state.format_failure(e) + if error_frame is not None: + yield error_frame + return if isinstance(e, HTTPException): raise e elif isinstance(e, StreamingCallbackError): @@ -8958,12 +8974,15 @@ def select_data_generator( user_api_key_dict: UserAPIKeyAuth, request_data: dict, request: Request | None = None, + *, + responses_stream_errors: bool = False, ): return async_data_generator( response=response, user_api_key_dict=user_api_key_dict, request_data=request_data, request=request, + responses_stream_errors=responses_stream_errors, ) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 5907ffc64eb..3202fabe74e 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -3,6 +3,7 @@ import json import time from collections.abc import AsyncIterator, Awaitable, Mapping from enum import Enum +from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, cast, get_args from uuid import uuid4 @@ -243,6 +244,7 @@ async def responses_api( version, ) + native_data_generator: Final = partial(select_data_generator, responses_stream_errors=True) data = await _read_request_body(request=request) # Check if polling via cache should be used for this request @@ -329,7 +331,7 @@ async def responses_api( llm_router=llm_router, proxy_config=proxy_config, proxy_logging_obj=proxy_logging_obj, - select_data_generator=select_data_generator, + select_data_generator=native_data_generator, user_model=user_model, user_temperature=user_temperature, user_request_timeout=user_request_timeout, @@ -355,7 +357,7 @@ async def responses_api( llm_router=llm_router, general_settings=general_settings, proxy_config=proxy_config, - select_data_generator=select_data_generator, + select_data_generator=native_data_generator, model=None, user_model=user_model, user_temperature=user_temperature, diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 9f9016c5a7f..1134f6b07e3 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -579,6 +579,11 @@ class BaseResponsesAPIStreamingIterator: message=error_message, llm_provider=self.custom_llm_provider or "", model=self.model or "", + body={ # mutable-ok: OpenAI APIError reads code/type only from a dict body + "code": error_code, + "type": error_type, + "message": error_message, + }, ) if 400 <= status_code < 500 and status_code != 429: raise mapped_exception diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index 87e10ce7e8d..6f3a47a4b79 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -17,15 +17,18 @@ from __future__ import annotations import asyncio import json +from collections.abc import AsyncIterator +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import Response from fastapi.responses import StreamingResponse +from pydantic import BaseModel import litellm -from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY import litellm.proxy.proxy_server as ps +from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.proxy_server import ( _apply_streaming_chunk_hooks, @@ -42,6 +45,12 @@ from litellm.proxy.proxy_server import ( data_generator, select_data_generator, ) +from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponseCreatedEvent, + ResponseFailedEvent, + ResponsesAPIResponse, +) from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage from .conftest import normalize @@ -872,6 +881,80 @@ async def test_async_data_generator_mid_stream_exception_yields_error_payload( assert any(isinstance(item, str) and item.startswith('data: {"error":') for item in out) +@pytest.mark.asyncio +@pytest.mark.parametrize("terminal", ["completed", "serialization_failure", "failure_after_completed"]) +async def test_responses_stream_keeps_tool_deltas_and_only_emits_a_valid_terminal( + terminal: Literal["completed", "serialization_failure", "failure_after_completed"], +) -> None: + class ToolDelta(BaseModel): + type: Literal["response.function_call_arguments.delta"] + sequence_number: int + item_id: str + output_index: int + delta: str + + class UnserializableTerminal(BaseModel): + type: Literal["response.completed"] + sequence_number: int + response: ResponsesAPIResponse + invalid: object + + response: Final = ResponsesAPIResponse(id="resp_visible", created_at=1, model="gpt-6-astra", output=[]) + created: Final = ResponseCreatedEvent.model_validate( + {"type": "response.created", "sequence_number": 0, "response": response} + ) + completed: Final = ResponseCompletedEvent.model_validate( + {"type": "response.completed", "sequence_number": 2, "response": response} + ) + tool_delta: Final = ToolDelta( + type="response.function_call_arguments.delta", sequence_number=1, item_id="fc_stream_error", + output_index=0, delta='{"path":"partial', + ) + + async def upstream() -> AsyncIterator[BaseModel]: + yield created + yield tool_delta + yield ( + UnserializableTerminal(type="response.completed", sequence_number=2, response=response, invalid=object()) + if terminal == "serialization_failure" else completed + ) + if terminal == "failure_after_completed": + raise litellm.APIError(status_code=500, message="Stream close failed", llm_provider="openai", model="gpt-6-astra") + + frames: Final = [ + frame + async for frame in select_data_generator( + response=upstream(), + user_api_key_dict=_user_auth(), + request_data={}, + responses_stream_errors=True, + ) + ] + decoded: Final = tuple(frame.decode() if isinstance(frame, bytes) else frame for frame in frames) + event_frames: Final = tuple(frame for frame in decoded if frame != "data: [DONE]\n\n") + payloads: Final = tuple( + json.loads(next(line[6:] for line in frame.splitlines() if line.startswith("data: "))) + for frame in event_frames + ) + + assert payloads[0]["response"]["id"] == "resp_visible" + assert payloads[1] == tool_delta.model_dump() + assert len(payloads) == 3 + if terminal == "serialization_failure": + failure: Final = ResponseFailedEvent.model_validate(payloads[-1]) + assert event_frames[-1].startswith("event: response.failed\n") + assert failure.response.id == "resp_visible" + assert failure.response.status == "failed" + assert failure.response.error is not None + assert failure.response.error["code"] == "server_error" + assert "serialize" in failure.response.error["message"].lower() + assert payloads[-1]["sequence_number"] > payloads[1]["sequence_number"] + else: + assert payloads[-1]["type"] == "response.completed" + assert payloads[-1]["sequence_number"] == 2 + assert "error" not in payloads[-1] + + # --------------------------------------------------------------------------- # select_data_generator # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index d7010de6405..9e2f70a3d64 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -3,10 +3,12 @@ Test for response_api_endpoints/endpoints.py """ import unittest -from typing import Any +from typing import Any, Final, Literal from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx from fastapi.testclient import TestClient from httpx import Response @@ -14,6 +16,97 @@ import litellm from litellm.proxy.proxy_server import app +@pytest.mark.asyncio +@pytest.mark.parametrize( + "path,error_kind", + [ + ("/v1/responses", "rate_limit"), + ("/v1/responses", "numeric_rate_limit"), + ("/v1/responses", "server_error"), + ("/v1/responses", "response_failed"), + ("/cursor/chat/completions", "server_error"), + ("/v1/chat/completions", "server_error"), + ], +) +async def test_streaming_upstream_errors_keep_the_client_protocol( + monkeypatch: pytest.MonkeyPatch, + path: str, + error_kind: Literal["rate_limit", "numeric_rate_limit", "server_error", "response_failed"], +) -> None: + import litellm.proxy.proxy_server as ps + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + model: Final = "gpt-6-astra" + message: Final = "Upstream cannot complete this response" + code: Final = { + "rate_limit": "rate_limit_exceeded", "numeric_rate_limit": "429", + "server_error": "server_error", "response_failed": "server_error", + }[error_kind] + error: Final = {"message": message, "code": code, "type": None, "param": "input"} + response: Final = {"id": "resp_upstream", "object": "response", "created_at": 1, + "status": "in_progress", "model": model, "output": [], + "parallel_tool_calls": True, "tool_choice": "auto", "tools": []} + created: Final = {"type": "response.created", "sequence_number": 0, "response": response} + tool_added: Final = {"type": "response.output_item.added", "sequence_number": 1, "output_index": 0, + "item": {"type": "function_call", "id": "fc_partial", "call_id": "call_partial", + "name": "read_file", "arguments": "", "status": "in_progress"}} + tool_delta: Final = {"type": "response.function_call_arguments.delta", "sequence_number": 2, + "item_id": "fc_partial", "output_index": 0, "delta": '{"path":"partial'} + failed: Final = ( + {"type": "response.failed", "sequence_number": 9, + "response": {**response, "status": "failed", "error": error}} + if error_kind == "response_failed" else {"type": "error", "error": error} + ) + chat: Final = {"id": "chatcmpl_partial", "object": "chat.completion.chunk", "created": 1, + "model": model, "choices": [{"index": 0, "delta": {"content": "partial"}, + "finish_reason": None}]} + is_chat: Final = path == "/v1/chat/completions" + upstream_events: Final = (chat, {"error": error}) if is_chat else (created, tool_added, tool_delta, failed) + wire: Final = "".join("data: " + json.dumps(event) + "\n\n" for event in upstream_events) + upstream_url: Final = "https://streaming.example/v1" + router: Final = litellm.Router( + model_list=[{"model_name": model, "litellm_params": { + "model": "openai/" + model, "api_base": upstream_url, "api_key": "fixture-key"}}], + num_retries=0, + ) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override) + with respx.mock as transport: + transport.post(upstream_url + ("/chat/completions" if is_chat else "/responses")).respond( + 200, content=wire, headers={"Content-Type": "text/event-stream"} + ) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: + result: Final = await client.post( + path, json={"model": model, "stream": True, + **({"messages": [{"role": "user", "content": "hello"}]} if is_chat else {"input": "hello"})}, + ) + frames: Final = tuple(frame for frame in result.text.split("\n\n") if "data: " in frame) + events: Final = tuple( + json.loads(next(line[6:] for line in frame.splitlines() if line.startswith("data: "))) + for frame in frames if "data: [DONE]" not in frame + ) + + assert result.status_code == 200, result.text + assert message in result.text + if path == "/v1/responses": + assert frames[-1].startswith("event: response.failed\n"), result.text + assert [event["type"] for event in events] == [ + "response.created", "response.output_item.added", "response.function_call_arguments.delta", "response.failed" + ] + assert events[2]["delta"] == tool_delta["delta"] + assert events[-1]["sequence_number"] == events[-2]["sequence_number"] + 1 + assert events[-1]["response"]["id"] == events[0]["response"]["id"] + assert events[-1]["response"]["status"] == "failed" + assert events[-1]["response"]["error"]["code"] == ( + "rate_limit_exceeded" if error_kind in ("rate_limit", "numeric_rate_limit") else "server_error" + ) + else: + assert events[0]["object"] == "chat.completion.chunk", result.text + assert "response.failed" not in result.text + assert "error" in events[-1] + + class TestResponsesAPIEndpoints(unittest.TestCase): @pytest.mark.asyncio @patch("litellm.proxy.proxy_server.llm_router") From 089703d10b2eebbec1a5b780494686451ab72161 Mon Sep 17 00:00:00 2001 From: zoroyihan7 Date: Tue, 8 Sep 2026 11:02:38 +0000 Subject: [PATCH 02/28] fix(responses): preserve failure metadata at streaming boundaries --- .../common_utils/responses_stream_errors.py | 45 +++++++++---- litellm/proxy/proxy_server.py | 9 +-- .../proxy_server/test_streaming_helpers.py | 63 ++++++++++++++++--- .../response_api_endpoints/test_endpoints.py | 28 ++++++--- 4 files changed, 114 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index a3bee06e912..356e948a2df 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -1,9 +1,10 @@ import time from collections.abc import Mapping +from http import HTTPStatus from types import MappingProxyType from typing import Final -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, field_validator from litellm._logging import redact_internal_details_from_client_message from litellm._uuid import uuid @@ -35,6 +36,16 @@ class _FailureDetails(BaseModel): type: str | None = None status_code: int | None = None + @field_validator("code", mode="before") + @classmethod + def normalize_code(cls, value: object) -> str | int | None: + return value if isinstance(value, (str, int)) and not isinstance(value, bool) else None + + @field_validator("type", mode="before") + @classmethod + def normalize_type(cls, value: object) -> str | None: + return value if isinstance(value, str) else None + def _original_failure(exception: Exception) -> Exception: if isinstance(exception, MidStreamFallbackError) and exception.original_exception is not None: @@ -50,9 +61,21 @@ def _response_error_code(details: _FailureDetails) -> str: return "rate_limit_exceeded" if isinstance(details.code, str) and details.code and not details.code.isdecimal(): return details.code - if details.status_code == 429: - return "rate_limit_exceeded" - return "server_error" + match details.status_code: + case HTTPStatus.UNAUTHORIZED: + return "authentication_error" + case HTTPStatus.FORBIDDEN: + return "permission_error" + case HTTPStatus.NOT_FOUND: + return "not_found_error" + case HTTPStatus.REQUEST_TIMEOUT: + return "request_timeout" + case HTTPStatus.TOO_MANY_REQUESTS: + return "rate_limit_exceeded" + case int(status) if HTTPStatus.BAD_REQUEST <= status < HTTPStatus.INTERNAL_SERVER_ERROR: + return "invalid_request_error" + case _: + return "server_error" class ResponsesStreamErrorState: @@ -62,16 +85,15 @@ class ResponsesStreamErrorState: self.created_at: int | None = None self.sequence_number = -1 self.terminal_emitted = False + self._pending_event: _StreamEvent | None = None - @staticmethod - def observe_chunk(chunk: object) -> _StreamEvent | None: - if not isinstance(chunk, (BaseModel, Mapping)): - return None - return _StreamEvent.model_validate(chunk) + def observe_chunk(self, chunk: object) -> None: + self._pending_event = _StreamEvent.model_validate(chunk) if isinstance(chunk, (BaseModel, Mapping)) else None - def mark_emitted(self, event: _StreamEvent | None) -> None: + def mark_emitted(self, frame: str | bytes) -> str | bytes: + event: Final = self._pending_event if event is None: - return + return frame if event.sequence_number is not None: self.sequence_number = max(self.sequence_number, event.sequence_number) if event.response is not None: @@ -81,6 +103,7 @@ class ResponsesStreamErrorState: self.created_at = event.response.created_at if event.type in ("response.completed", "response.failed", "response.incomplete"): self.terminal_emitted = True + return frame def format_failure(self, exception: Exception) -> str | None: if self.terminal_emitted: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9f1d69acefa..870cd78aa85 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8841,7 +8841,8 @@ async def async_data_generator( fallback_metadata_event_sent = True continue - responses_event: Final = error_state.observe_chunk(chunk) if error_state is not None else None + if error_state is not None: + error_state.observe_chunk(cast(object, chunk)) # cast-ok: the helper validates legacy untyped chunks raw_passthrough = False if isinstance(chunk, BaseModel): chunk = _serialize_streaming_chunk(chunk) @@ -8876,10 +8877,10 @@ async def async_data_generator( if not raw_passthrough: try: - formatted_chunk: Final = _format_streaming_sse_chunk(chunk=chunk) if error_state is not None: - error_state.mark_emitted(responses_event) - yield formatted_chunk + yield error_state.mark_emitted(_format_streaming_sse_chunk(chunk=chunk)) + else: + yield _format_streaming_sse_chunk(chunk=chunk) except Exception as e: if error_state is not None: raise diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index 6f3a47a4b79..53e055882e8 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -21,9 +21,11 @@ from collections.abc import AsyncIterator from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock +import httpx import pytest -from fastapi import Response +from fastapi import HTTPException, Response from fastapi.responses import StreamingResponse +from openai import APIError as OpenAIAPIError from pydantic import BaseModel import litellm @@ -882,9 +884,44 @@ async def test_async_data_generator_mid_stream_exception_yields_error_payload( @pytest.mark.asyncio -@pytest.mark.parametrize("terminal", ["completed", "serialization_failure", "failure_after_completed"]) +@pytest.mark.parametrize( + "terminal,upstream_error,expected_code", + [ + ("completed", None, None), + ("serialization_failure", None, "server_error"), + ("failure_after_completed", None, None), + pytest.param( + "upstream_failure", + litellm.AuthenticationError( + message="Upstream rejected request", llm_provider="openai", model="gpt-6-astra" + ), + "authentication_error", id="authentication_error", + ), + pytest.param( + "upstream_failure", + OpenAIAPIError( + message="Upstream rejected request", + request=httpx.Request("POST", "https://streaming.example/v1/responses"), + body={"code": {"reason": "overloaded"}, "type": {"unexpected": "object"}}, + ), + "server_error", id="structured_provider_error_fields", + ), + *( + pytest.param( + "upstream_failure", HTTPException(status_code=status, detail="Upstream rejected request"), + code, id=f"http_{status}", + ) + for status, code in ( + (400, "invalid_request_error"), (403, "permission_error"), (404, "not_found_error"), + (408, "request_timeout"), (422, "invalid_request_error"), (500, "server_error"), (503, "server_error"), + ) + ), + ], +) async def test_responses_stream_keeps_tool_deltas_and_only_emits_a_valid_terminal( - terminal: Literal["completed", "serialization_failure", "failure_after_completed"], + terminal: Literal["completed", "serialization_failure", "failure_after_completed", "upstream_failure"], + upstream_error: HTTPException | OpenAIAPIError | None, + expected_code: str | None, ) -> None: class ToolDelta(BaseModel): type: Literal["response.function_call_arguments.delta"] @@ -910,16 +947,23 @@ async def test_responses_stream_keeps_tool_deltas_and_only_emits_a_valid_termina type="response.function_call_arguments.delta", sequence_number=1, item_id="fc_stream_error", output_index=0, delta='{"path":"partial', ) + original_status: Final = ( + upstream_error.status_code if isinstance(upstream_error, (HTTPException, litellm.AuthenticationError)) else None + ) async def upstream() -> AsyncIterator[BaseModel]: yield created yield tool_delta + if upstream_error is not None: + raise upstream_error yield ( UnserializableTerminal(type="response.completed", sequence_number=2, response=response, invalid=object()) if terminal == "serialization_failure" else completed ) if terminal == "failure_after_completed": - raise litellm.APIError(status_code=500, message="Stream close failed", llm_provider="openai", model="gpt-6-astra") + raise litellm.APIError( + status_code=500, message="Stream close failed", llm_provider="openai", model="gpt-6-astra" + ) frames: Final = [ frame @@ -940,14 +984,19 @@ async def test_responses_stream_keeps_tool_deltas_and_only_emits_a_valid_termina assert payloads[0]["response"]["id"] == "resp_visible" assert payloads[1] == tool_delta.model_dump() assert len(payloads) == 3 - if terminal == "serialization_failure": + if terminal in ("serialization_failure", "upstream_failure"): failure: Final = ResponseFailedEvent.model_validate(payloads[-1]) assert event_frames[-1].startswith("event: response.failed\n") assert failure.response.id == "resp_visible" assert failure.response.status == "failed" assert failure.response.error is not None - assert failure.response.error["code"] == "server_error" - assert "serialize" in failure.response.error["message"].lower() + assert failure.response.error["code"] == expected_code + if upstream_error is None: + assert "serialize" in failure.response.error["message"].lower() + else: + assert "Upstream rejected request" in failure.response.error["message"] + if isinstance(upstream_error, (HTTPException, litellm.AuthenticationError)): + assert upstream_error.status_code == original_status assert payloads[-1]["sequence_number"] > payloads[1]["sequence_number"] else: assert payloads[-1]["type"] == "response.completed" diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index 9e2f70a3d64..7b79e1b8613 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -61,7 +61,9 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( "model": model, "choices": [{"index": 0, "delta": {"content": "partial"}, "finish_reason": None}]} is_chat: Final = path == "/v1/chat/completions" - upstream_events: Final = (chat, {"error": error}) if is_chat else (created, tool_added, tool_delta, failed) + partial: Final = path != "/v1/responses" or error_kind in ("numeric_rate_limit", "response_failed") + response_events: Final = (created, tool_added, tool_delta, failed) if partial else (failed,) + upstream_events: Final = (chat, {"error": error}) if is_chat else response_events wire: Final = "".join("data: " + json.dumps(event) + "\n\n" for event in upstream_events) upstream_url: Final = "https://streaming.example/v1" router: Final = litellm.Router( @@ -78,8 +80,10 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( ) async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: result: Final = await client.post( - path, json={"model": model, "stream": True, - **({"messages": [{"role": "user", "content": "hello"}]} if is_chat else {"input": "hello"})}, + path, json={ + "model": model, "stream": True, + **({"messages": [{"role": "user", "content": "hello"}]} if is_chat else {"input": "hello"}), + }, ) frames: Final = tuple(frame for frame in result.text.split("\n\n") if "data: " in frame) events: Final = tuple( @@ -91,12 +95,18 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( assert message in result.text if path == "/v1/responses": assert frames[-1].startswith("event: response.failed\n"), result.text - assert [event["type"] for event in events] == [ - "response.created", "response.output_item.added", "response.function_call_arguments.delta", "response.failed" - ] - assert events[2]["delta"] == tool_delta["delta"] - assert events[-1]["sequence_number"] == events[-2]["sequence_number"] + 1 - assert events[-1]["response"]["id"] == events[0]["response"]["id"] + if partial: + assert [event["type"] for event in events] == [ + "response.created", "response.output_item.added", + "response.function_call_arguments.delta", "response.failed", + ] + assert events[2]["delta"] == tool_delta["delta"] + assert events[-1]["sequence_number"] == events[-2]["sequence_number"] + 1 + assert events[-1]["response"]["id"] == events[0]["response"]["id"] + else: + assert [event["type"] for event in events] == ["response.failed"] + assert events[0]["sequence_number"] == 0 + assert events[0]["response"]["id"].startswith("resp_") assert events[-1]["response"]["status"] == "failed" assert events[-1]["response"]["error"]["code"] == ( "rate_limit_exceeded" if error_kind in ("rate_limit", "numeric_rate_limit") else "server_error" From dc7895c1eacf55733953c3f802f46ea9103a76b4 Mon Sep 17 00:00:00 2001 From: zoroyihan7 Date: Tue, 8 Sep 2026 11:44:41 +0000 Subject: [PATCH 03/28] fix(responses): satisfy streaming regression checks --- litellm/proxy/common_utils/responses_stream_errors.py | 7 ++++--- .../test_response_polling_pre_call_checks.py | 1 - 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index 356e948a2df..706b4c298d7 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -48,9 +48,10 @@ class _FailureDetails(BaseModel): def _original_failure(exception: Exception) -> Exception: - if isinstance(exception, MidStreamFallbackError) and exception.original_exception is not None: - return _original_failure(exception.original_exception) - return exception + current = exception # rebind-ok: the recursion gate requires iterative wrapper traversal + while isinstance(current, MidStreamFallbackError) and current.original_exception is not None: + current = current.original_exception + return current def _response_error_code(details: _FailureDetails) -> str: diff --git a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py index 459834d0fd2..3fabdcefe5a 100644 --- a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py +++ b/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py @@ -130,7 +130,6 @@ class TestPollingEndpointPreCallGuard: "litellm.proxy.proxy_server.proxy_config": MagicMock(), "litellm.proxy.proxy_server.proxy_logging_obj": AsyncMock(), "litellm.proxy.proxy_server.redis_usage_cache": AsyncMock(), - "litellm.proxy.proxy_server.select_data_generator": None, "litellm.proxy.proxy_server.user_api_base": None, "litellm.proxy.proxy_server.user_max_tokens": None, "litellm.proxy.proxy_server.user_model": None, From 1a749d84bdd66706bb41cafdba28e7a8b6a20fa9 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 01:56:00 +0000 Subject: [PATCH 04/28] fix(proxy): track project spend and enforce project budgets additively Project-scoped keys never wrote spend to LiteLLM_ProjectTable, so /project/info stayed at 0 and project budgets could not block. Wire the PROJECT entity through the spend queue, redis buffer, and db writer, reserve and increment a spend:project counter, reseed it from the project row, reset project spend in the budget cascade, and read the live counter in the project max budget check. Team member budgets keep gating project-scoped keys alongside the project budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 1 + litellm/proxy/auth/auth_checks.py | 34 ++-- .../proxy/common_utils/reset_budget_job.py | 24 +++ .../proxy/common_utils/user_api_key_cache.py | 10 ++ litellm/proxy/db/db_spend_update_writer.py | 79 ++++++++- .../redis_update_buffer.py | 7 + .../spend_update_queue.py | 4 + litellm/proxy/db/spend_counter_reseed.py | 5 + .../proxy/hooks/proxy_track_cost_callback.py | 6 + litellm/proxy/proxy_server.py | 30 ++++ .../spend_tracking/budget_reservation.py | 41 +++++ .../spend_tracking/spend_counter_batch.py | 11 +- litellm/repositories/prisma_protocols.py | 3 + litellm/repositories/unit_of_work.py | 2 + .../proxy/auth/test_auth_checks.py | 61 +++++++ .../common_utils/test_reset_budget_job.py | 30 +++- .../proxy/db/test_db_spend_update_writer.py | 76 +++++++++ .../proxy/db/test_spend_counter_reseed.py | 19 +++ .../hooks/test_proxy_track_cost_callback.py | 1 + .../test_spend_tracking_utils.py | 1 + .../proxy/test_budget_reservation.py | 156 ++++++++++++++++++ .../repositories/test_unit_of_work.py | 3 + 22 files changed, 583 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 321f8190f13..228a91ad446 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5261,6 +5261,7 @@ class DBSpendUpdateTransactions(TypedDict): team_member_list_transactions: dict[str, float] | None org_list_transactions: dict[str, float] | None org_member_list_transactions: ReadOnly[dict[str, float] | None] + project_list_transactions: ReadOnly[dict[str, float] | None] tag_list_transactions: dict[str, float] | None agent_list_transactions: dict[str, float] | None model_access_group_list_transactions: ReadOnly[dict[str, float] | None] diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e3783c94dc7..ef17913d9ec 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -92,6 +92,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( model_access_group_registry_cache_key, model_access_group_spend_counter_key, object_permission_cache_key, + project_cache_key, + project_spend_counter_key, tag_cache_key, tag_registry_cache_key, team_membership_auth_cache_key, @@ -5586,16 +5588,22 @@ async def _project_max_budget_check( if project_object.litellm_budget_table is not None: max_budget = project_object.litellm_budget_table.max_budget - if ( - max_budget is not None - and project_object.spend is not None - and math.isfinite(max_budget) - and project_object.spend > max_budget - ): + if max_budget is None or not math.isfinite(max_budget): + return + + from litellm.proxy.proxy_server import get_current_spend + + project_spend: Final = await get_current_spend( + counter_key=project_spend_counter_key(project_object.project_id), + fallback_spend=project_object.spend or 0.0, + max_budget=max_budget, + ) + + if project_spend >= max_budget: if valid_token: call_info: Final = CallInfo( token=valid_token.token, - spend=project_object.spend, + spend=project_spend, max_budget=max_budget, user_id=valid_token.user_id, team_id=valid_token.team_id, @@ -5611,9 +5619,9 @@ async def _project_max_budget_check( ) raise litellm.BudgetExceededError( - current_cost=project_object.spend, + current_cost=project_spend, max_budget=max_budget, - message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_object.spend}, Max budget: {max_budget}", + message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_spend}, Max budget: {max_budget}", entity_type=Litellm_EntityType.PROJECT.value, entity_id=project_object.project_id, ) @@ -5663,10 +5671,6 @@ async def _project_soft_budget_check( ) -def _project_cache_key(project_id: str) -> str: - return f"project_id:{project_id}" - - async def get_project_object( project_id: str, prisma_client: PrismaClient | None, @@ -5684,7 +5688,7 @@ async def get_project_object( return None # Check cache first - cache_key: Final = _project_cache_key(project_id) + cache_key: Final = project_cache_key(project_id) deserialized_project: Final = await user_api_key_cache.async_get_cache( key=cache_key, model_type=LiteLLM_ProjectTableCachedObj, @@ -5726,7 +5730,7 @@ async def delete_cached_project_object( from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast await evict_and_broadcast( - cache_keys=(_project_cache_key(project_id),), + cache_keys=(project_cache_key(project_id),), user_api_key_cache=user_api_key_cache, ) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index acb51e73daf..2baefa89943 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -41,6 +41,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( end_user_cache_key, model_access_group_cache_key, model_access_group_spend_counter_key, + project_cache_key, + project_spend_counter_key, tag_cache_key, ) from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row @@ -49,6 +51,7 @@ from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.prisma_protocols import SpendLinkedTable +from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ( EndUserRepository, ModelAccessGroupBudgetRepository, @@ -115,6 +118,11 @@ class _ModelAccessGroupRow(_BudgetLinkedRow, Protocol): def access_group_name(self) -> str: ... +class _ProjectRow(_BudgetLinkedRow, Protocol): + @property + def project_id(self) -> str: ... + + class _EndUserRow(_BudgetLinkedRow, Protocol): @property def user_id(self) -> str: ... @@ -185,6 +193,14 @@ def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...] return (model_access_group_cache_key(row.access_group_name),) +def _project_counter_key(row: _ProjectRow) -> str: + return project_spend_counter_key(row.project_id) + + +def _project_cache_keys(row: _ProjectRow) -> tuple[str, ...]: + return (project_cache_key(row.project_id),) + + def _enduser_counter_key(row: _EndUserRow) -> str: return f"spend:end_user:{row.user_id}" @@ -661,6 +677,11 @@ class ResetBudgetJob: where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), log_subject="model access groups", ) + projects: Final[tuple[_ProjectRow, ...]] = await self._fetch_linked_rows( + table=ProjectRepository(self.prisma_client).table, + where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), + log_subject="projects", + ) rollover_caps: Final[Mapping[str, float]] = MappingProxyType( { # mutable-ok: MappingProxyType wraps a one-shot dict comprehension b.budget_id: cap @@ -695,6 +716,7 @@ class ResetBudgetJob: (_model_access_group_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in model_access_groups ), + *((_project_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in projects), *((_enduser_counter_key(row), _enduser_carried_spend(row, rollover_caps)) for row in endusers), ), rollover_caps=rollover_caps, @@ -704,6 +726,7 @@ class ResetBudgetJob: *(key for row in orgs for key in _org_cache_keys(row)), *(key for row in tags for key in _tag_cache_keys(row)), *(key for row in model_access_groups for key in _model_access_group_cache_keys(row)), + *(key for row in projects for key in _project_cache_keys(row)), *(key for row in endusers for key in _enduser_cache_keys(row)), ), ) @@ -731,6 +754,7 @@ class ResetBudgetJob: _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE) + _queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE) _queue_enduser_resets(uow.endusers, cascade) for budget_id, budget_reset_at in cascade.budget_resets: uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 1c7a379897f..2187ed63ea5 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -306,6 +306,16 @@ def model_access_group_spend_counter_key(access_group_name: str) -> str: return f"spend:model_access_group:{access_group_name}" +def project_cache_key(project_id: str) -> str: + """Cache key one project row is stored under; shared by auth, spend tracking and the spend writer.""" + return f"project_id:{project_id}" + + +def project_spend_counter_key(project_id: str) -> str: + """Spend counter key for one project; the reservation, cost callback, auth and reseed paths all read it.""" + return f"spend:project:{project_id}" + + #: Cached under ``end_user_restricted_registry_cache_key`` when the restricted set exceeds #: ``END_USER_RESTRICTED_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-id fetch. END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL: Final = "__end_user_restricted_registry_overflow__" diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index a90d1351fd7..51d00b9789c 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -44,6 +44,7 @@ from litellm.proxy._types import ( SpendUpdateQueueItem, ToolDiscoveryQueueItem, ) +from litellm.proxy.common_utils.user_api_key_cache import project_cache_key from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, build_bulk_upsert, @@ -116,6 +117,7 @@ class _SpendBatch(Protocol): litellm_teammembership: BatchTable litellm_organizationtable: BatchTable litellm_organizationmembership: BatchTable + litellm_projecttable: BatchTable litellm_tagtable: BatchTable litellm_agentstable: BatchTable litellm_modelaccessgroupbudgettable: BatchTable @@ -254,6 +256,7 @@ class DBSpendUpdateWriter: start_time: datetime | None, end_time: datetime | None, response_cost: float | None, + project_id: str | None = None, ) -> bool: """Record the request's spend, answering whether its cost still needs charging. @@ -335,6 +338,7 @@ class DBSpendUpdateWriter: hashed_token=hashed_token, team_id=team_id, org_id=org_id, + project_id=project_id, end_user_id=end_user_id, prisma_client=prisma_client, litellm_proxy_budget_name=litellm_proxy_budget_name, @@ -631,6 +635,7 @@ class DBSpendUpdateWriter: litellm_proxy_budget_name: str | None, payload: SpendLogsPayload, request_model_access_groups: Sequence[str] = (), + project_id: str | None = None, ): """ Runs all 13 spend-update helpers sequentially inside a single asyncio task. @@ -694,6 +699,18 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) + try: + await self._update_project_db( + response_cost=response_cost, + project_id=project_id, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: _update_project_db failed: %s", + traceback.format_exc(), + ) + try: await self._update_tag_db( response_cost=response_cost, @@ -956,6 +973,33 @@ class DBSpendUpdateWriter: ) raise e + async def _update_project_db( + self, + response_cost: float | None, + project_id: str | None, + prisma_client: PrismaClient | None, + ): + try: + if project_id is None or prisma_client is None: + return + + await self.spend_update_queue.add_update( + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.PROJECT, + entity_id=project_id, + response_cost=response_cost, + ) + ) + except Exception as e: + spend_log_error( + "Spend tracking - failed to enqueue project spend update. project_id=%s, response_cost=%s - %s", + project_id, + response_cost, + str(e), + exc=e, + ) + raise e + async def _update_agent_db( self, response_cost: float | None, @@ -1193,8 +1237,8 @@ class DBSpendUpdateWriter: if db_spend_update_transactions is not None: verbose_proxy_logger.info( "Spend tracking - committing spend updates from Redis to DB: " - "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, org_members=%d, tags=%d, " - "agents=%d, model_access_groups=%d", + "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, org_members=%d, " + "projects=%d, tags=%d, agents=%d, model_access_groups=%d", len(db_spend_update_transactions.get("key_list_transactions") or {}), len(db_spend_update_transactions.get("user_list_transactions") or {}), len(db_spend_update_transactions.get("team_list_transactions") or {}), @@ -1202,6 +1246,7 @@ class DBSpendUpdateWriter: len(db_spend_update_transactions.get("end_user_list_transactions") or {}), len(db_spend_update_transactions.get("team_member_list_transactions") or {}), len(db_spend_update_transactions.get("org_member_list_transactions") or {}), + len(db_spend_update_transactions.get("project_list_transactions") or {}), len(db_spend_update_transactions.get("tag_list_transactions") or {}), len(db_spend_update_transactions.get("agent_list_transactions") or {}), len(db_spend_update_transactions.get("model_access_group_list_transactions") or {}), @@ -1762,6 +1807,22 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) + ### UPDATE PROJECT TABLE ### + project_list_transactions: Final = db_spend_update_transactions.get("project_list_transactions") + await DBSpendUpdateWriter._update_entity_spend_in_db( + entity_name="Project", + transactions=project_list_transactions, + table_accessor="litellm_projecttable", + where_field="project_id", + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ) + await DBSpendUpdateWriter._invalidate_project_caches( + project_ids=tuple(project_list_transactions or ()), + proxy_logging_obj=proxy_logging_obj, + ) + ### UPDATE TAG TABLE ### tag_list_transactions: Final = db_spend_update_transactions["tag_list_transactions"] await DBSpendUpdateWriter._update_entity_spend_in_db( @@ -1800,11 +1861,23 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) + @staticmethod + async def _invalidate_project_caches(project_ids: Sequence[str], proxy_logging_obj: ProxyLogging | None) -> None: + if not project_ids or proxy_logging_obj is None: + return + user_api_key_cache: Final = proxy_logging_obj.call_details.get("user_api_key_cache") + if user_api_key_cache is None: + return + for project_id in project_ids: + await user_api_key_cache.async_delete_cache(key=project_cache_key(project_id)) + @staticmethod async def _update_entity_spend_in_db( entity_name: str, transactions: dict[str, float] | None, - table_accessor: Literal["litellm_tagtable", "litellm_agentstable", "litellm_modelaccessgroupbudgettable"], + table_accessor: Literal[ + "litellm_tagtable", "litellm_agentstable", "litellm_modelaccessgroupbudgettable", "litellm_projecttable" + ], where_field: str, n_retry_times: int, prisma_client: PrismaClient, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 6f49a00b763..cead63795a2 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -70,6 +70,7 @@ _SpendTransactionField: TypeAlias = Literal[ "team_member_list_transactions", "org_list_transactions", "org_member_list_transactions", + "project_list_transactions", "tag_list_transactions", "agent_list_transactions", "model_access_group_list_transactions", @@ -83,6 +84,7 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = ( "team_member_list_transactions", "org_list_transactions", "org_member_list_transactions", + "project_list_transactions", "tag_list_transactions", "agent_list_transactions", "model_access_group_list_transactions", @@ -418,6 +420,10 @@ class RedisUpdateBuffer: Litellm_EntityType.ORGANIZATION_MEMBER, db_spend_update_transactions.get("org_member_list_transactions"), ), + ( + Litellm_EntityType.PROJECT, + db_spend_update_transactions.get("project_list_transactions"), + ), ( Litellm_EntityType.TAG, db_spend_update_transactions.get("tag_list_transactions"), @@ -885,6 +891,7 @@ class RedisUpdateBuffer: org_member_list_transactions=_merged_entity_transactions( list_of_transactions, "org_member_list_transactions" ), + project_list_transactions=_merged_entity_transactions(list_of_transactions, "project_list_transactions"), tag_list_transactions=_merged_entity_transactions(list_of_transactions, "tag_list_transactions"), agent_list_transactions=_merged_entity_transactions(list_of_transactions, "agent_list_transactions"), model_access_group_list_transactions=_merged_entity_transactions( diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index bc068d10daf..2b8535cb113 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -138,6 +138,7 @@ class SpendUpdateQueue(BaseUpdateQueue): team_member_list_transactions={}, org_list_transactions={}, org_member_list_transactions={}, + project_list_transactions={}, tag_list_transactions={}, agent_list_transactions={}, model_access_group_list_transactions={}, @@ -152,6 +153,7 @@ class SpendUpdateQueue(BaseUpdateQueue): Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions", Litellm_EntityType.ORGANIZATION: "org_list_transactions", Litellm_EntityType.ORGANIZATION_MEMBER: "org_member_list_transactions", + Litellm_EntityType.PROJECT: "project_list_transactions", Litellm_EntityType.TAG: "tag_list_transactions", Litellm_EntityType.AGENT: "agent_list_transactions", Litellm_EntityType.MODEL_ACCESS_GROUP: "model_access_group_list_transactions", @@ -192,6 +194,8 @@ class SpendUpdateQueue(BaseUpdateQueue): transactions_dict = db_spend_update_transactions["org_list_transactions"] elif dict_key == "org_member_list_transactions": transactions_dict = db_spend_update_transactions["org_member_list_transactions"] + elif dict_key == "project_list_transactions": + transactions_dict = db_spend_update_transactions["project_list_transactions"] elif dict_key == "tag_list_transactions": transactions_dict = db_spend_update_transactions["tag_list_transactions"] elif dict_key == "agent_list_transactions": diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 89a07234c6c..2dd028454d6 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -26,6 +26,7 @@ from litellm.proxy._types import Litellm_EntityType from litellm.proxy.db.db_lookup_gate import db_lookup_gate from litellm.proxy.spend_tracking.spend_counter_batch import read_batched_spend_counter, record_spend_counter_value from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ( BudgetWindowSpendRepository, EndUserRepository, @@ -77,6 +78,7 @@ class SpendCounterReseed: spend:team_member:{uid}:{tid} -> LiteLLM_TeamMembership.spend spend:user:{user_id} -> LiteLLM_UserTable.spend spend:org:{org_id} -> LiteLLM_OrganizationTable.spend + spend:project:{project_id} -> LiteLLM_ProjectTable.spend End-user and tag spend counters intentionally do not reseed here. Their auth paths already load the corresponding objects via get_end_user_object() @@ -157,6 +159,9 @@ class SpendCounterReseed: row = await OrganizationRepository(prisma_client).table.find_unique( where={"organization_id": org_id} ) + elif counter_key.startswith("spend:project:"): + project_id: Final = counter_key[len("spend:project:") :] + row = await ProjectRepository(prisma_client).table.find_unique(where={"project_id": project_id}) else: return None except Exception: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 1ae106be390..0c562cf37ef 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -267,6 +267,7 @@ class _ProxyDBLogger(CustomLogger): start_time=actual_start_time, end_time=datetime.now(), org_id=user_api_key_dict.org_id, + project_id=user_api_key_dict.project_id, ) @log_db_metrics @@ -318,6 +319,7 @@ class _ProxyDBLogger(CustomLogger): user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None)) team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None)) org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None)) + project_id: Final = cast(str | None, metadata.get("user_api_key_project_id", None)) key_alias: Final = cast(str | None, metadata.get("user_api_key_alias", None)) end_user_max_budget: Final = metadata.get("user_api_end_user_max_budget", None) sl_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) @@ -368,6 +370,7 @@ class _ProxyDBLogger(CustomLogger): budget_reservation=budget_reservation, request_tags=tags, model_access_groups=model_access_groups, + project_id=project_id, ) if not charged: return @@ -651,6 +654,7 @@ async def _update_database_and_spend_counters( budget_reservation: dict | None, request_tags: list[str] | None = None, model_access_groups: Sequence[str] | None = None, + project_id: str | None = None, ) -> bool: if budget_reservation is not None: await _reconcile_budget_reservation_before_db_update( @@ -668,6 +672,7 @@ async def _update_database_and_spend_counters( start_time=start_time, end_time=end_time, org_id=org_id, + project_id=project_id, ) except Exception: if budget_reservation is not None: @@ -698,6 +703,7 @@ async def _update_database_and_spend_counters( tags=request_tags, request_started_at=start_time, model_access_groups=model_access_groups, + project_id=project_id, ) except Exception: if budget_reservation is not None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d7964556531..1e84e5f56e2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -417,6 +417,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( get_management_object_ttl, model_access_group_cache_key, model_access_group_spend_counter_key, + project_cache_key, + project_spend_counter_key, tag_cache_key, ) from litellm.proxy.config_resolvers import resolve_fields @@ -2780,6 +2782,7 @@ async def increment_spend_counters( tags: list[str] | None = None, request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, + project_id: str | None = None, ): """ Atomically increment spend counters for budget enforcement. @@ -2801,6 +2804,7 @@ async def increment_spend_counters( end_user_id=end_user_id, tags=tags, model_access_groups=model_access_groups, + project_id=project_id, ), ): await _increment_spend_counters_batched( @@ -2814,6 +2818,7 @@ async def increment_spend_counters( tags=tags, request_started_at=request_started_at, model_access_groups=model_access_groups, + project_id=project_id, ) @@ -2828,6 +2833,7 @@ async def _increment_spend_counters_batched( tags: list[str] | None, request_started_at: datetime | None, model_access_groups: Sequence[str] | None, + project_id: str | None = None, ): """Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET.""" reserved_counter_keys: Final = await _reconcile_budget_reservation_for_counter_update( @@ -3028,6 +3034,13 @@ async def _increment_spend_counters_batched( ) if org_id is not None else None, + _prepare_project_spend_increment( + project_id=project_id, + response_cost=cost, + reserved_counter_keys=reserved_counter_keys, + ) + if project_id is not None + else None, ) if coro is not None ) @@ -3180,6 +3193,23 @@ async def _prepare_org_spend_increment( return (pending,) if pending is not None else () +async def _prepare_project_spend_increment( + project_id: str | None, + response_cost: float, + reserved_counter_keys: set[str], +) -> tuple[PendingSpendIncrement, ...]: + if project_id is None: + return () + + pending: Final = await _prepare_unreserved_spend_counter_increment( + counter_key=project_spend_counter_key(project_id), + source_cache_key=project_cache_key(project_id), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + return (pending,) if pending is not None else () + + async def _prepare_unreserved_spend_counter_increment( counter_key: str, source_cache_key: str | list[str], diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 373f2d0fe36..24b8470145c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import json +import math from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -29,6 +30,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( end_user_cache_key, model_access_group_cache_key, model_access_group_spend_counter_key, + project_cache_key, + project_spend_counter_key, tag_cache_key, team_membership_reservation_cache_key, ) @@ -62,6 +65,7 @@ _COUNTER_ENTITY_TYPES: Final[Mapping[str, str]] = { "Tag": Litellm_EntityType.TAG.value, "Model access group": Litellm_EntityType.MODEL_ACCESS_GROUP.value, "Organization": Litellm_EntityType.ORGANIZATION.value, + "Project": Litellm_EntityType.PROJECT.value, } @@ -542,6 +546,13 @@ async def _get_budget_counters( if org_counter is not None: counters.append(org_counter) + project_counter: Final = await _get_project_budget_counter( + valid_token=valid_token, + user_api_key_cache=user_api_key_cache, + ) + if project_counter is not None: + counters.append(project_counter) + return counters @@ -751,6 +762,36 @@ async def _get_org_budget_counter( ) +async def _get_project_budget_counter( + valid_token: UserAPIKeyAuth, + user_api_key_cache: UserApiKeyCache, +) -> _BudgetCounter | None: + if valid_token.project_id is None: + return None + + source_cache_key: Final = project_cache_key(valid_token.project_id) + project_object: Final = await user_api_key_cache.async_get_cache(key=source_cache_key) + if project_object is None: + return None + + project_budget_table: Final = _get_value(project_object, "litellm_budget_table") + if project_budget_table is None: + return None + + project_max_budget: Final = _to_float(_get_value(project_budget_table, "max_budget")) + if project_max_budget is None or project_max_budget <= 0 or not math.isfinite(project_max_budget): + return None + + return _BudgetCounter( + counter_key=project_spend_counter_key(valid_token.project_id), + source_cache_key=source_cache_key, + max_budget=project_max_budget, + fallback_spend=_to_float(_get_value(project_object, "spend")) or 0.0, + entity_type="Project", + entity_id=valid_token.project_id, + ) + + def _get_budget_limit_counters( entity_prefix: str, entity_type: str, diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py index 7106d88c655..ddb074ae023 100644 --- a/litellm/proxy/spend_tracking/spend_counter_batch.py +++ b/litellm/proxy/spend_tracking/spend_counter_batch.py @@ -12,7 +12,10 @@ from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.caching.redis_cache import RedisCache from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.common_utils.user_api_key_cache import model_access_group_spend_counter_key +from litellm.proxy.common_utils.user_api_key_cache import ( + model_access_group_spend_counter_key, + project_spend_counter_key, +) _CounterValues: Final = TypeAdapter(dict[str, float | None]) _NO_VALUES: Final[Mapping[str, float | None]] = MappingProxyType({}) @@ -154,6 +157,8 @@ def _iter_admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) yield f"spend:end_user:{end_user_id}" if token.org_id is not None: yield f"spend:org:{token.org_id}" + if token.project_id is not None: + yield project_spend_counter_key(token.project_id) def admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) -> frozenset[str]: @@ -168,10 +173,12 @@ def post_call_counter_keys( end_user_id: str | None, tags: Sequence[object] | None, model_access_groups: Sequence[object] | None, + project_id: str | None = None, ) -> frozenset[str]: """Every counter ``increment_spend_counters`` warm-checks, except budget windows which bind on read.""" entity_keys: Final = admission_counter_keys( - UserAPIKeyAuth(token=token, team_id=team_id, user_id=user_id, org_id=org_id), end_user_id + UserAPIKeyAuth(token=token, team_id=team_id, user_id=user_id, org_id=org_id, project_id=project_id), + end_user_id, ) tag_keys: Final = frozenset(f"spend:tag:{tag}" for tag in tags or () if tag and isinstance(tag, str)) group_keys: Final = frozenset( diff --git a/litellm/repositories/prisma_protocols.py b/litellm/repositories/prisma_protocols.py index 93b8c5c7cd7..60c16fbd746 100644 --- a/litellm/repositories/prisma_protocols.py +++ b/litellm/repositories/prisma_protocols.py @@ -152,4 +152,7 @@ class PrismaBatch(Protocol): @property def litellm_modelaccessgroupbudgettable(self) -> BatchTable: ... + @property + def litellm_projecttable(self) -> BatchTable: ... + async def commit(self) -> None: ... diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index 0cdce307f9b..c09e5eb75d4 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -109,6 +109,7 @@ class BudgetCascadeUnitOfWork: organizations: LinkedSpendResetWrites tags: LinkedSpendResetWrites model_access_groups: LinkedSpendResetWrites + projects: LinkedSpendResetWrites endusers: LinkedSpendResetWrites budgets: BudgetWindowWrites @@ -135,6 +136,7 @@ async def budget_cascade_unit_of_work( organizations=LinkedSpendResetWrites(table=batch.litellm_organizationtable), tags=LinkedSpendResetWrites(table=batch.litellm_tagtable), model_access_groups=LinkedSpendResetWrites(table=batch.litellm_modelaccessgroupbudgettable), + projects=LinkedSpendResetWrites(table=batch.litellm_projecttable), endusers=LinkedSpendResetWrites(table=batch.litellm_endusertable), budgets=BudgetWindowWrites(table=batch.litellm_budgettable), ) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 8c8b755195f..0d5c0dd5d72 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7138,6 +7138,67 @@ async def test_project_allowlist_enforced_when_key_models_empty(): assert exc_info.value.code == "403" +def _project_with_budget(spend: float, max_budget: float): + from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_ProjectTableCachedObj + + return LiteLLM_ProjectTableCachedObj( + project_id="p-budget", + team_id="t-1", + budget_id="b-1", + spend=spend, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b-1", max_budget=max_budget), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "counter_spend, db_spend, blocks", + [ + pytest.param(5.0, 0.0, True, id="counter-at-budget-blocks-despite-stale-db-row"), + pytest.param(4.99, 0.0, False, id="counter-under-budget-admits"), + pytest.param(None, 5.0, True, id="no-counter-falls-back-to-persisted-spend"), + pytest.param(None, 0.0, False, id="no-counter-and-no-persisted-spend-admits"), + ], +) +async def test_project_max_budget_check_reads_live_spend_counter(counter_spend, db_spend, blocks): + """LIT-3269: project budget enforcement must read the cross-pod + ``spend:project:{id}`` counter first and only fall back to the cached row's + spend, matching key/team/org checks. The boundary is inclusive (>=).""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.auth.auth_checks import _project_max_budget_check + + real_spend_counter_cache = DualCache() + if counter_spend is not None: + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:project:p-budget", value=counter_spend) + valid_token = UserAPIKeyAuth(api_key="hashed-key", project_id="p-budget", team_id="t-1", user_id="u-1") + proxy_logging_obj = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + + with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ): + if not blocks: + await _project_max_budget_check( + project_object=_project_with_budget(spend=db_spend, max_budget=5.0), + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + return + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _project_max_budget_check( + project_object=_project_with_budget(spend=db_spend, max_budget=5.0), + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) + await asyncio.sleep(0) + + assert exc_info.value.entity_type == Litellm_EntityType.PROJECT.value + assert exc_info.value.entity_id == "p-budget" + assert exc_info.value.current_cost == 5.0 + proxy_logging_obj.budget_alerts.assert_awaited_once() + assert proxy_logging_obj.budget_alerts.await_args.kwargs["type"] == "project_budget" + + def test_is_user_proxy_admin_rejects_view_only_admin(): """This predicate skips `non_proxy_admin_allowed_routes_check` entirely, so an Admin Viewer answering True here would gain every write route. Read parity for diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 943a6c905c0..8c254686385 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -78,6 +78,7 @@ class MockBatcher: self.litellm_organizationtable = _Table("org", self) self.litellm_tagtable = _Table("tag", self) self.litellm_modelaccessgroupbudgettable = _Table("model_access_group", self) + self.litellm_projecttable = _Table("project", self) self.litellm_endusertable = _Table("enduser", self) async def commit(self): @@ -93,6 +94,7 @@ class MockDB: self.litellm_organizationtable = MockTable() self.litellm_tagtable = MockTable() self.litellm_modelaccessgroupbudgettable = MockTable() + self.litellm_projecttable = MockTable() self.batch_calls: List[Dict[str, Any]] = [] self.batchers: List[MockBatcher] = [] @@ -1507,13 +1509,19 @@ _INVALIDATION_CASES = [ "spend:model_access_group:gpt-4-group", {"model_access_group:gpt-4-group"}, ), + ( + "litellm_projecttable", + type("Project", (), {"project_id": "proj-1"}), + "spend:project:proj-1", + {"project_id:proj-1"}, + ), ] @pytest.mark.parametrize( "table_attr, linked_row, counter_key, cache_keys", _INVALIDATION_CASES, - ids=["team_membership", "key", "org", "tag", "model_access_group"], + ids=["team_membership", "key", "org", "tag", "model_access_group", "project"], ) def test_budget_table_reset_invalidates_counters_and_management_cache( reset_budget_job, mock_prisma_client, monkeypatch, table_attr, linked_row, counter_key, cache_keys @@ -1657,6 +1665,25 @@ def test_budget_table_reset_invalidates_every_access_group_not_just_the_first( counter_cache.in_memory_cache.delete_cache.assert_any_call(key=f"spend:model_access_group:{name}") +def test_project_reset_zeroes_spend_on_due_tiers(reset_budget_job, mock_prisma_client, monkeypatch): + """A project linked to an expiring budget tier has its spend zeroed in the same cascade transaction.""" + _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-due", budget_duration="7d")] + mock_prisma_client.db.litellm_projecttable.set_find_many_results( + [type("Project", (), {"project_id": "proj-1", "spend": 12.0, "budget_id": "budget-due"})] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + expected_where = {"budget_id": {"in": ["budget-due"]}, "spend": {"gt": 0}} + assert mock_prisma_client.db.litellm_projecttable.find_many_calls == [{"where": expected_where}] + writes = _batch_writes(mock_prisma_client, "project", op="update_many") + assert len(writes) == 1 + assert writes[0]["where"] == expected_where + assert writes[0]["data"] == {"spend": 0} + assert mock_prisma_client.db.batchers[0].committed is True + + def test_budget_cascade_carries_access_group_overage_when_rollover_enabled( rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch ): @@ -1802,6 +1829,7 @@ def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mo ("org", "update_many"), ("tag", "update_many"), ("model_access_group", "update_many"), + ("project", "update_many"), ("enduser", "update_many"), ("budget", "update_many"), } diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index c547d06904b..da5879a375a 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1060,6 +1060,82 @@ async def test_batch_database_updates_queues_org_member_spend_for_the_request_us assert transactions["org_member_list_transactions"] == {"organization_id::org1::user_id::u1": 0.1} +@pytest.mark.asyncio +async def test_project_spend_is_persisted_to_project_table_and_project_cache_is_evicted(): + """Regression for LIT-3269: a request made with a project-scoped key must + increment LiteLLM_ProjectTable.spend, otherwise /project/info stays at 0 + and the project budget never blocks. The cached project row is evicted so + the next auth check reads the fresh spend.""" + db_writer: Final = DBSpendUpdateWriter() + await db_writer._batch_database_updates( + response_cost=0.25, + user_id="u1", + hashed_token="t1", + team_id="team-1", + org_id=None, + end_user_id=None, + prisma_client=MagicMock(), + litellm_proxy_budget_name=None, + payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.25}, + project_id="proj-1", + ) + await db_writer._batch_database_updates( + response_cost=0.5, + user_id="u1", + hashed_token="t1", + team_id="team-1", + org_id=None, + end_user_id=None, + prisma_client=MagicMock(), + litellm_proxy_budget_name=None, + payload={"request_id": "req-2", "model": "gpt-4o-mini", "spend": 0.5}, + project_id="proj-1", + ) + transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + assert transactions["project_list_transactions"] == {"proj-1": 0.75} + assert transactions["team_member_list_transactions"] == {"team_id::team-1::user_id::u1": 0.75} + + mock_batcher: Final = MagicMock() + mock_prisma_client: Final = MagicMock() + mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(mock_batcher)) + user_api_key_cache: Final = MagicMock() + user_api_key_cache.async_delete_cache = AsyncMock() + proxy_logging: Final = MagicMock() + proxy_logging.call_details = {"user_api_key_cache": user_api_key_cache} + + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=proxy_logging, + db_spend_update_transactions=transactions, + ) + + mock_batcher.litellm_projecttable.update_many.assert_called_once_with( + where={"project_id": "proj-1"}, + data={"spend": {"increment": 0.75}}, + ) + user_api_key_cache.async_delete_cache.assert_any_await(key="project_id:proj-1") + + +@pytest.mark.asyncio +async def test_batch_database_updates_without_project_id_touches_no_project_row(): + db_writer: Final = DBSpendUpdateWriter() + await db_writer._batch_database_updates( + response_cost=0.1, + user_id="u1", + hashed_token="t1", + team_id=None, + org_id=None, + end_user_id=None, + prisma_client=MagicMock(), + litellm_proxy_budget_name=None, + payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.1}, + ) + transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + + assert transactions["project_list_transactions"] == {} + + @pytest.mark.asyncio async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id(): """ diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index bca6344b3f7..53e91b8792e 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -67,12 +67,14 @@ class _FakePrismaClient: error: Exception | None = None, end_user_row: SimpleNamespace | None = None, end_user_error: Exception | None = None, + project_row: SimpleNamespace | None = None, ) -> None: self.db = SimpleNamespace( litellm_budgetwindowspend=_FakeFindUniqueTable(row=row, error=error), litellm_spendlogs=_FakeSpendLogsTable(total=spend_logs_total), litellm_endusertable=_FakeFindUniqueTable(row=end_user_row, error=end_user_error), litellm_verificationtoken=_InFlightCountingTable(), + litellm_projecttable=_FakeFindUniqueTable(row=project_row), ) @@ -428,6 +430,23 @@ async def test_from_db_bounds_in_flight_prisma_requests_across_counter_keys(): assert prisma.db.litellm_verificationtoken.max_in_flight == PROXY_DB_LOOKUP_MAX_CONCURRENCY +@pytest.mark.asyncio +async def test_from_db_reseeds_project_counter_from_the_project_row(): + """LIT-3269: a cold ``spend:project:{id}`` counter seeds from LiteLLM_ProjectTable.spend, + so a fresh pod enforces the project budget against persisted spend rather than 0.""" + prisma: Final = _FakePrismaClient(project_row=SimpleNamespace(project_id="proj-1", spend=7.25)) + + assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:project:proj-1") == 7.25 + assert prisma.db.litellm_projecttable.where_clauses == [{"project_id": "proj-1"}] + + +@pytest.mark.asyncio +async def test_from_db_returns_none_for_a_missing_project_row(): + prisma: Final = _FakePrismaClient(project_row=None) + + assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:project:proj-1") is None + + @pytest.mark.asyncio async def test_from_db_still_never_reads_the_end_user_row(): """A cold end-user counter keeps seeding from the cached end-user object the auth diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index dfc95db3e14..965e134772d 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -586,6 +586,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda tags=["tag-a"], request_started_at=start_time, model_access_groups=("premium",), + project_id=None, ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 8b105e94d19..834cf8d100d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -1187,6 +1187,7 @@ async def test_api_key_preserved_through_failure_hook_to_database(): start_time, end_time, org_id, + project_id=None, ): """Mock update_database and capture the payload it creates""" from litellm.proxy.spend_tracking.spend_tracking_utils import ( diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 032722d3259..014f240d9cc 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -22,6 +22,7 @@ from litellm.proxy._types import ( LiteLLM_EndUserTable, Litellm_EntityType, LiteLLM_OrganizationTable, + LiteLLM_ProjectTableCachedObj, LiteLLM_TagTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, @@ -631,6 +632,161 @@ async def test_should_reserve_team_member_and_org_budget_counters(spend_counter_ await release_budget_reservation(reservation) +def _project_scoped_token() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + token="key-project-scoped", + spend=0.0, + user_id="user-proj", + team_id="team-proj", + project_id="proj-1", + ) + + +async def _seed_project_scoped_budgets( + key_cache: DualCache, + team_member_spend: float, + team_member_max_budget: float, + project_spend: float, + project_max_budget: float, +) -> None: + await key_cache.async_set_cache( + key="team_membership:user-proj:team-proj", + value=LiteLLM_TeamMembership( + user_id="user-proj", + team_id="team-proj", + spend=team_member_spend, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=team_member_max_budget), + ).model_dump(), + ) + await key_cache.async_set_cache( + key="project_id:proj-1", + value=LiteLLM_ProjectTableCachedObj( + project_id="proj-1", + team_id="team-proj", + budget_id="project-budget-id", + spend=project_spend, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=project_max_budget), + ).model_dump(), + ) + + +@pytest.mark.asyncio +async def test_should_reserve_project_and_team_member_counters_for_project_scoped_key(spend_counter_state): + """LIT-3269: a key carrying user_id, team_id and project_id reserves against + both the team member counter and the project counter; neither replaces the + other. After the call the project counter reflects the real cost once, not + the reservation plus the post-call increment.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + await _seed_project_scoped_budgets( + key_cache, + team_member_spend=0.1, + team_member_max_budget=1.0, + project_spend=0.2, + project_max_budget=1.0, + ) + + estimated = estimate_request_max_cost(request_body=_request_body(), route="/chat/completions", llm_router=None) + assert estimated is not None and estimated > 0 + + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=_project_scoped_token(), + team_object=LiteLLM_TeamTable(team_id="team-proj", spend=0.0, max_budget=None), + user_object=LiteLLM_UserTable(user_id="user-proj", spend=0.0), + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is not None + assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-proj:team-proj") == pytest.approx( + 0.1 + estimated + ) + assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") == pytest.approx(0.2 + estimated) + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token="key-project-scoped", + team_id="team-proj", + user_id="user-proj", + response_cost=0.05, + budget_reservation=reservation, + project_id="proj-1", + ) + + assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") == pytest.approx(0.25) + assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-proj:team-proj") == pytest.approx(0.15) + + +@pytest.mark.asyncio +async def test_exhausted_team_member_budget_still_blocks_project_scoped_key(spend_counter_state): + """LIT-3269: the project budget is additive. A project with plenty of + headroom must not let a key through once its team member budget is spent.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + await _seed_project_scoped_budgets( + key_cache, + team_member_spend=1.0, + team_member_max_budget=1.0, + project_spend=0.0, + project_max_budget=100.0, + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=_project_scoped_token(), + team_object=LiteLLM_TeamTable(team_id="team-proj", spend=0.0, max_budget=None), + user_object=LiteLLM_UserTable(user_id="user-proj", spend=0.0), + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert "TeamMember=user-proj:team-proj" in str(exc_info.value) + assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") in (None, pytest.approx(0.0)) + + +@pytest.mark.asyncio +async def test_exhausted_project_budget_blocks_project_scoped_key(spend_counter_state): + """LIT-3269: with team member headroom left, the project budget alone blocks the key.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + await _seed_project_scoped_budgets( + key_cache, + team_member_spend=0.0, + team_member_max_budget=100.0, + project_spend=5.0, + project_max_budget=5.0, + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=_project_scoped_token(), + team_object=LiteLLM_TeamTable(team_id="team-proj", spend=0.0, max_budget=None), + user_object=LiteLLM_UserTable(user_id="user-proj", spend=0.0), + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert "Project=proj-1" in str(exc_info.value) + assert exc_info.value.entity_type == Litellm_EntityType.PROJECT.value + assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-proj:team-proj") in ( + None, + pytest.approx(0.0), + ) + + @pytest.mark.asyncio async def test_should_not_reserve_user_budget_counter_for_team_key(spend_counter_state): """The reservation path mirrors the read path: no personal user counter for a team key. diff --git a/tests/test_litellm/repositories/test_unit_of_work.py b/tests/test_litellm/repositories/test_unit_of_work.py index b52b8ced31e..9eff248b917 100644 --- a/tests/test_litellm/repositories/test_unit_of_work.py +++ b/tests/test_litellm/repositories/test_unit_of_work.py @@ -35,6 +35,7 @@ class FakeBatch: self.litellm_organizationtable = FakeBatchTable("litellm_organizationtable", self.calls) self.litellm_tagtable = FakeBatchTable("litellm_tagtable", self.calls) self.litellm_modelaccessgroupbudgettable = FakeBatchTable("litellm_modelaccessgroupbudgettable", self.calls) + self.litellm_projecttable = FakeBatchTable("litellm_projecttable", self.calls) self.litellm_endusertable = FakeBatchTable("litellm_endusertable", self.calls) async def commit(self) -> None: @@ -94,6 +95,7 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): uow.organizations.queue_spend_zero(where=linked) uow.tags.queue_spend_zero(where=linked) uow.model_access_groups.queue_spend_zero(where=linked) + uow.projects.queue_spend_zero(where=linked) uow.endusers.queue_spend_zero(where={"user_id": {"in": ["enduser-1"]}}) uow.budgets.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at) assert batch.commit_count == 0 @@ -105,6 +107,7 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): ("litellm_organizationtable.update_many", linked, {"spend": 0}), ("litellm_tagtable.update_many", linked, {"spend": 0}), ("litellm_modelaccessgroupbudgettable.update_many", linked, {"spend": 0}), + ("litellm_projecttable.update_many", linked, {"spend": 0}), ("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-1"]}}, {"spend": 0}), ("litellm_budgettable.update_many", {"budget_id": "budget-1"}, {"budget_reset_at": reset_at}), ] From b00cd15bd73a380c77aad0d04672d27ee4bc13cf Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 02:07:11 +0000 Subject: [PATCH 05/28] test(proxy): import project_cache_key from user_api_key_cache Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/test_key_management_endpoints.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 4e70063015d..b57f8b7857d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -38,11 +38,10 @@ from litellm.proxy._types import ( from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.proxy.auth.auth_checks import ( _delete_cache_key_object, - _project_cache_key, jwt_key_mapping_cache_key, ) from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, project_cache_key from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_org_key_limits, @@ -18951,7 +18950,7 @@ async def test_regenerate_key_repoints_live_membership_not_the_key_row_it_read( async def _cache_with_project(project_id: str, project_models: list[str]) -> UserApiKeyCache: user_api_key_cache = UserApiKeyCache() await user_api_key_cache.async_set_cache( - key=_project_cache_key(project_id), + key=project_cache_key(project_id), value=LiteLLM_ProjectTableCachedObj(project_id=project_id, team_id="team-lit-5823", models=project_models), model_type=LiteLLM_ProjectTableCachedObj, ) From b20f1422eb204c2cb2fba26912a548454cd18b02 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 02:28:38 +0000 Subject: [PATCH 06/28] fix(proxy): carry project_id through key metadata enrichment and drop docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/user_api_key_cache.py | 2 -- litellm/proxy/hooks/proxy_track_cost_callback.py | 2 ++ tests/test_litellm/proxy/auth/test_auth_checks.py | 3 --- .../proxy/common_utils/test_reset_budget_job.py | 1 - tests/test_litellm/proxy/db/test_db_spend_update_writer.py | 4 ---- tests/test_litellm/proxy/db/test_spend_counter_reseed.py | 2 -- .../proxy/hooks/test_proxy_track_cost_callback.py | 3 +++ tests/test_litellm/proxy/test_budget_reservation.py | 7 ------- 8 files changed, 5 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 2187ed63ea5..0386b58070d 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -307,12 +307,10 @@ def model_access_group_spend_counter_key(access_group_name: str) -> str: def project_cache_key(project_id: str) -> str: - """Cache key one project row is stored under; shared by auth, spend tracking and the spend writer.""" return f"project_id:{project_id}" def project_spend_counter_key(project_id: str) -> str: - """Spend counter key for one project; the reservation, cost callback, auth and reseed paths all read it.""" return f"spend:project:{project_id}" diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 0c562cf37ef..5e525108ade 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -504,6 +504,8 @@ class _ProxyDBLogger(CustomLogger): metadata["user_api_key_team_id"] = key_obj.team_id if metadata.get("user_api_key_org_id") is None: metadata["user_api_key_org_id"] = key_obj.org_id + if metadata.get("user_api_key_project_id") is None: + metadata["user_api_key_project_id"] = key_obj.project_id except Exception: verbose_proxy_logger.debug( "Failed to enrich failure metadata with key info for api_key=%s", diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 0d5c0dd5d72..fe469fc574e 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7161,9 +7161,6 @@ def _project_with_budget(spend: float, max_budget: float): ], ) async def test_project_max_budget_check_reads_live_spend_counter(counter_spend, db_spend, blocks): - """LIT-3269: project budget enforcement must read the cross-pod - ``spend:project:{id}`` counter first and only fall back to the cached row's - spend, matching key/team/org checks. The boundary is inclusive (>=).""" from litellm.caching.dual_cache import DualCache from litellm.proxy.auth.auth_checks import _project_max_budget_check diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 8c254686385..48b06649237 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -1666,7 +1666,6 @@ def test_budget_table_reset_invalidates_every_access_group_not_just_the_first( def test_project_reset_zeroes_spend_on_due_tiers(reset_budget_job, mock_prisma_client, monkeypatch): - """A project linked to an expiring budget tier has its spend zeroed in the same cascade transaction.""" _make_counter_invalidation_job(monkeypatch) mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-due", budget_duration="7d")] mock_prisma_client.db.litellm_projecttable.set_find_many_results( diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index da5879a375a..60cc742577b 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1062,10 +1062,6 @@ async def test_batch_database_updates_queues_org_member_spend_for_the_request_us @pytest.mark.asyncio async def test_project_spend_is_persisted_to_project_table_and_project_cache_is_evicted(): - """Regression for LIT-3269: a request made with a project-scoped key must - increment LiteLLM_ProjectTable.spend, otherwise /project/info stays at 0 - and the project budget never blocks. The cached project row is evicted so - the next auth check reads the fresh spend.""" db_writer: Final = DBSpendUpdateWriter() await db_writer._batch_database_updates( response_cost=0.25, diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index 53e91b8792e..ff0b67d426b 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -432,8 +432,6 @@ async def test_from_db_bounds_in_flight_prisma_requests_across_counter_keys(): @pytest.mark.asyncio async def test_from_db_reseeds_project_counter_from_the_project_row(): - """LIT-3269: a cold ``spend:project:{id}`` counter seeds from LiteLLM_ProjectTable.spend, - so a fresh pod enforces the project budget against persisted spend rather than 0.""" prisma: Final = _FakePrismaClient(project_row=SimpleNamespace(project_id="proj-1", spend=7.25)) assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:project:proj-1") == 7.25 diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 965e134772d..202495517ad 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1372,6 +1372,7 @@ async def test_enrich_failure_metadata_with_full_key_lookup(): mock_key_obj.user_id = "fetched-user-id" mock_key_obj.team_id = "fetched-team-id" mock_key_obj.org_id = "fetched-org-id" + mock_key_obj.project_id = "fetched-project-id" mock_team_obj = MagicMock() mock_team_obj.team_alias = "fetched-team-alias" @@ -1395,12 +1396,14 @@ async def test_enrich_failure_metadata_with_full_key_lookup(): "user_api_key_team_id": None, "user_api_key_team_alias": None, "user_api_key_org_id": None, + "user_api_key_project_id": None, } result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) assert result["user_api_key_alias"] == "fetched-key-alias" assert result["user_api_key_user_id"] == "fetched-user-id" assert result["user_api_key_team_id"] == "fetched-team-id" assert result["user_api_key_org_id"] == "fetched-org-id" + assert result["user_api_key_project_id"] == "fetched-project-id" assert result["user_api_key_team_alias"] == "fetched-team-alias" diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 014f240d9cc..c834ac05f0a 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -672,10 +672,6 @@ async def _seed_project_scoped_budgets( @pytest.mark.asyncio async def test_should_reserve_project_and_team_member_counters_for_project_scoped_key(spend_counter_state): - """LIT-3269: a key carrying user_id, team_id and project_id reserves against - both the team member counter and the project counter; neither replaces the - other. After the call the project counter reflects the real cost once, not - the reservation plus the post-call increment.""" counter_cache, key_cache = spend_counter_state proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) await _seed_project_scoped_budgets( @@ -724,8 +720,6 @@ async def test_should_reserve_project_and_team_member_counters_for_project_scope @pytest.mark.asyncio async def test_exhausted_team_member_budget_still_blocks_project_scoped_key(spend_counter_state): - """LIT-3269: the project budget is additive. A project with plenty of - headroom must not let a key through once its team member budget is spent.""" counter_cache, key_cache = spend_counter_state proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) await _seed_project_scoped_budgets( @@ -755,7 +749,6 @@ async def test_exhausted_team_member_budget_still_blocks_project_scoped_key(spen @pytest.mark.asyncio async def test_exhausted_project_budget_blocks_project_scoped_key(spend_counter_state): - """LIT-3269: with team member headroom left, the project budget alone blocks the key.""" counter_cache, key_cache = spend_counter_state proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) await _seed_project_scoped_budgets( From 7e5b3b49d46ae1773e3f0450adbda334c198f108 Mon Sep 17 00:00:00 2001 From: ryan Date: Thu, 17 Sep 2026 22:33:54 +0000 Subject: [PATCH 07/28] fix(proxy): treat non-positive project max_budget as unbudgeted Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 2 +- .../proxy/auth/test_auth_checks.py | 23 +++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e80a4a677ce..e6313790380 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5612,7 +5612,7 @@ async def _project_max_budget_check( if project_object.litellm_budget_table is not None: max_budget = project_object.litellm_budget_table.max_budget - if max_budget is None or not math.isfinite(max_budget): + if max_budget is None or max_budget <= 0 or not math.isfinite(max_budget): return from litellm.proxy.proxy_server import get_current_spend diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 56ee0aac249..08ed1575fdd 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7378,6 +7378,29 @@ async def test_project_max_budget_check_reads_live_spend_counter(counter_spend, assert proxy_logging_obj.budget_alerts.await_args.kwargs["type"] == "project_budget" +@pytest.mark.asyncio +@pytest.mark.parametrize("max_budget", [0.0, -1.0]) +async def test_project_max_budget_check_treats_non_positive_budget_as_unbudgeted(max_budget): + from litellm.caching.dual_cache import DualCache + from litellm.proxy.auth.auth_checks import _project_max_budget_check + + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:project:p-budget", value=12.5) + proxy_logging_obj = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + + with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ): + await _project_max_budget_check( + project_object=_project_with_budget(spend=12.5, max_budget=max_budget), + valid_token=UserAPIKeyAuth(api_key="hashed-key", project_id="p-budget"), + proxy_logging_obj=proxy_logging_obj, + ) + + proxy_logging_obj.budget_alerts.assert_not_awaited() + + def test_is_user_proxy_admin_rejects_view_only_admin(): """This predicate skips `non_proxy_admin_allowed_routes_check` entirely, so an Admin Viewer answering True here would gain every write route. Read parity for From 1c0342332beb96ba3c13ede34563c7e1ccd3c8a2 Mon Sep 17 00:00:00 2001 From: ryan Date: Thu, 17 Sep 2026 22:40:07 +0000 Subject: [PATCH 08/28] refactor(proxy): keep project spend enqueue within lint ceilings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/db_spend_update_writer.py | 67 ++++++++-------------- 1 file changed, 25 insertions(+), 42 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 4f04a1b3a1a..e0985d5b5ae 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -706,17 +706,11 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) - try: - await self._update_project_db( - response_cost=response_cost, - project_id=project_id, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_project_db failed: %s", - traceback.format_exc(), - ) + await self._update_project_db( + response_cost=response_cost, + project_id=project_id, + prisma_client=prisma_client, + ) try: await self._update_tag_db( @@ -985,27 +979,16 @@ class DBSpendUpdateWriter: response_cost: float | None, project_id: str | None, prisma_client: PrismaClient | None, - ): - try: - if project_id is None or prisma_client is None: - return - - await self.spend_update_queue.add_update( - update=SpendUpdateQueueItem( - entity_type=Litellm_EntityType.PROJECT, - entity_id=project_id, - response_cost=response_cost, - ) + ) -> None: + if project_id is None or prisma_client is None: + return + await self.spend_update_queue.add_update( + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.PROJECT, + entity_id=project_id, + response_cost=response_cost, ) - except Exception as e: - spend_log_error( - "Spend tracking - failed to enqueue project spend update. project_id=%s, response_cost=%s - %s", - project_id, - response_cost, - str(e), - exc=e, - ) - raise e + ) async def _update_agent_db( self, @@ -1246,17 +1229,17 @@ class DBSpendUpdateWriter: "Spend tracking - committing spend updates from Redis to DB: " "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, org_members=%d, " "projects=%d, tags=%d, agents=%d, model_access_groups=%d", - len(db_spend_update_transactions.get("key_list_transactions") or {}), - len(db_spend_update_transactions.get("user_list_transactions") or {}), - len(db_spend_update_transactions.get("team_list_transactions") or {}), - len(db_spend_update_transactions.get("org_list_transactions") or {}), - len(db_spend_update_transactions.get("end_user_list_transactions") or {}), - len(db_spend_update_transactions.get("team_member_list_transactions") or {}), - len(db_spend_update_transactions.get("org_member_list_transactions") or {}), - len(db_spend_update_transactions.get("project_list_transactions") or {}), - len(db_spend_update_transactions.get("tag_list_transactions") or {}), - len(db_spend_update_transactions.get("agent_list_transactions") or {}), - len(db_spend_update_transactions.get("model_access_group_list_transactions") or {}), + len(db_spend_update_transactions.get("key_list_transactions") or ()), + len(db_spend_update_transactions.get("user_list_transactions") or ()), + len(db_spend_update_transactions.get("team_list_transactions") or ()), + len(db_spend_update_transactions.get("org_list_transactions") or ()), + len(db_spend_update_transactions.get("end_user_list_transactions") or ()), + len(db_spend_update_transactions.get("team_member_list_transactions") or ()), + len(db_spend_update_transactions.get("org_member_list_transactions") or ()), + len(db_spend_update_transactions.get("project_list_transactions") or ()), + len(db_spend_update_transactions.get("tag_list_transactions") or ()), + len(db_spend_update_transactions.get("agent_list_transactions") or ()), + len(db_spend_update_transactions.get("model_access_group_list_transactions") or ()), ) await self._commit_spend_updates_to_db( prisma_client=prisma_client, From 70164dabd2795613a814f3e6d76c344e5f5ce618 Mon Sep 17 00:00:00 2001 From: ryan Date: Thu, 17 Sep 2026 23:20:10 +0000 Subject: [PATCH 09/28] fix(proxy): isolate project spend enqueue failures from sibling spend writes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/db_spend_update_writer.py | 16 +++++++---- .../proxy/db/test_db_spend_update_writer.py | 27 +++++++++++++++++++ 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e0985d5b5ae..2d6c7a0a74c 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -706,11 +706,17 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) - await self._update_project_db( - response_cost=response_cost, - project_id=project_id, - prisma_client=prisma_client, - ) + try: + await self._update_project_db( + response_cost=response_cost, + project_id=project_id, + prisma_client=prisma_client, + ) + except Exception: # noqa: BLE001 # a project enqueue failure must not skip the sibling spend writes + verbose_proxy_logger.debug( + "_batch_database_updates: _update_project_db failed: %s", + traceback.format_exc(), + ) try: await self._update_tag_db( diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index f9494cd6332..9a2b15d8a30 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1175,6 +1175,33 @@ async def test_batch_database_updates_without_project_id_touches_no_project_row( assert transactions["project_list_transactions"] == {} +@pytest.mark.asyncio +async def test_project_enqueue_failure_does_not_stop_sibling_spend_updates(): + db_writer: Final = DBSpendUpdateWriter() + db_writer._update_project_db = AsyncMock(side_effect=RuntimeError("project queue boom")) + db_writer._update_tag_db = AsyncMock() + db_writer._update_agent_db = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock() + + await db_writer._batch_database_updates( + response_cost=0.1, + user_id="u1", + hashed_token="t1", + team_id="team-1", + org_id=None, + end_user_id=None, + prisma_client=MagicMock(), + litellm_proxy_budget_name=None, + payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.1, "request_tags": ["t"]}, + project_id="proj-1", + ) + + db_writer._update_project_db.assert_awaited_once() + db_writer._update_tag_db.assert_awaited_once() + db_writer._update_agent_db.assert_awaited_once() + db_writer.add_spend_log_transaction_to_daily_user_transaction.assert_awaited_once() + + @pytest.mark.asyncio async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id(): """ From e4778cd2d1d14a5efed40546031d474fe53eab63 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 17 Sep 2026 16:24:51 -0700 Subject: [PATCH 10/28] test(proxy): drive the real spend queue in the project enqueue isolation test Substitutes a queue that rejects project items instead of replacing writer methods with mocks, so the test asserts the sibling key, team and tag spend actually landed in the queue rather than that a mock was awaited. --- .../proxy/db/test_db_spend_update_writer.py | 31 +++++++++++-------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 9a2b15d8a30..3d61b1dba12 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -15,8 +15,9 @@ import pytest from redis.exceptions import DataError import litellm -from litellm.proxy._types import Litellm_EntityType +from litellm.proxy._types import Litellm_EntityType, SpendUpdateQueueItem from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( build_window_spend_transaction, ) @@ -1176,30 +1177,34 @@ async def test_batch_database_updates_without_project_id_touches_no_project_row( @pytest.mark.asyncio -async def test_project_enqueue_failure_does_not_stop_sibling_spend_updates(): +async def test_failed_project_enqueue_does_not_drop_the_rest_of_the_batch(): + class _ProjectRejectingQueue(SpendUpdateQueue): + async def add_update(self, update: SpendUpdateQueueItem): + if update.get("entity_type") is Litellm_EntityType.PROJECT: + raise RuntimeError("project enqueue boom") + await super().add_update(update) + db_writer: Final = DBSpendUpdateWriter() - db_writer._update_project_db = AsyncMock(side_effect=RuntimeError("project queue boom")) - db_writer._update_tag_db = AsyncMock() - db_writer._update_agent_db = AsyncMock() - db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock() + db_writer.spend_update_queue = _ProjectRejectingQueue() await db_writer._batch_database_updates( - response_cost=0.1, + response_cost=0.25, user_id="u1", hashed_token="t1", team_id="team-1", - org_id=None, + org_id="org-1", end_user_id=None, prisma_client=MagicMock(), litellm_proxy_budget_name=None, - payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.1, "request_tags": ["t"]}, + payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.25, "request_tags": ["tag-1"]}, project_id="proj-1", ) - db_writer._update_project_db.assert_awaited_once() - db_writer._update_tag_db.assert_awaited_once() - db_writer._update_agent_db.assert_awaited_once() - db_writer.add_spend_log_transaction_to_daily_user_transaction.assert_awaited_once() + transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + assert transactions["project_list_transactions"] == {} + assert transactions["tag_list_transactions"] == {"tag-1": 0.25} + assert transactions["key_list_transactions"] == {"t1": 0.25} + assert transactions["team_list_transactions"] == {"team-1": 0.25} @pytest.mark.asyncio From 89289d4f7753d3d780ce0cf3b679c409220c2b88 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 17 Sep 2026 16:26:26 -0700 Subject: [PATCH 11/28] fix(proxy): surface a failed project spend enqueue at error level The call-site guard kept a project enqueue failure from skipping the sibling spend writes, but logged it at debug only, so a dropped project charge was invisible on a default log level. Log it through spend_log_error inside the helper and re-raise, the way the org and agent helpers already do. --- litellm/proxy/db/db_spend_update_writer.py | 22 +++++++++---- .../proxy/db/test_db_spend_update_writer.py | 33 +++++++++++-------- 2 files changed, 36 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 2d6c7a0a74c..f715897b314 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -988,13 +988,23 @@ class DBSpendUpdateWriter: ) -> None: if project_id is None or prisma_client is None: return - await self.spend_update_queue.add_update( - update=SpendUpdateQueueItem( - entity_type=Litellm_EntityType.PROJECT, - entity_id=project_id, - response_cost=response_cost, + try: + await self.spend_update_queue.add_update( + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.PROJECT, + entity_id=project_id, + response_cost=response_cost, + ) ) - ) + except Exception as e: + spend_log_error( + "Spend tracking - failed to enqueue project spend update. project_id=%s, response_cost=%s - %s", + project_id, + response_cost, + str(e), + exc=e, + ) + raise e async def _update_agent_db( self, diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 3d61b1dba12..2efc0d99b8a 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1,6 +1,7 @@ import asyncio import copy import json +import logging import re @@ -15,6 +16,7 @@ import pytest from redis.exceptions import DataError import litellm +from litellm._logging import verbose_proxy_logger from litellm.proxy._types import Litellm_EntityType, SpendUpdateQueueItem from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue @@ -1177,7 +1179,9 @@ async def test_batch_database_updates_without_project_id_touches_no_project_row( @pytest.mark.asyncio -async def test_failed_project_enqueue_does_not_drop_the_rest_of_the_batch(): +async def test_failed_project_enqueue_is_reported_and_does_not_drop_the_rest_of_the_batch( + caplog: pytest.LogCaptureFixture, +): class _ProjectRejectingQueue(SpendUpdateQueue): async def add_update(self, update: SpendUpdateQueueItem): if update.get("entity_type") is Litellm_EntityType.PROJECT: @@ -1187,18 +1191,21 @@ async def test_failed_project_enqueue_does_not_drop_the_rest_of_the_batch(): db_writer: Final = DBSpendUpdateWriter() db_writer.spend_update_queue = _ProjectRejectingQueue() - await db_writer._batch_database_updates( - response_cost=0.25, - user_id="u1", - hashed_token="t1", - team_id="team-1", - org_id="org-1", - end_user_id=None, - prisma_client=MagicMock(), - litellm_proxy_budget_name=None, - payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.25, "request_tags": ["tag-1"]}, - project_id="proj-1", - ) + with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name): + await db_writer._batch_database_updates( + response_cost=0.25, + user_id="u1", + hashed_token="t1", + team_id="team-1", + org_id="org-1", + end_user_id=None, + prisma_client=MagicMock(), + litellm_proxy_budget_name=None, + payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.25, "request_tags": ["tag-1"]}, + project_id="proj-1", + ) + + assert any("proj-1" in record.getMessage() for record in caplog.records if record.levelno >= logging.ERROR) transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() assert transactions["project_list_transactions"] == {} From 7cfec730d5df0b0994922608984b1660207ba735 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 12:42:08 -0700 Subject: [PATCH 12/28] fix(proxy): refuse config-owned keys on POST /config/update POST /config/update stored keys the config file owns and answered 200 while the file silently kept winning. Run the config-owned check for general_settings, litellm_settings, and router_settings before the first database read, the way /config/field/update already does, so a refused request stores nothing. The config reload re-loaded the merged settings as yaml settings, which turned every saved router setting read-only after one tick. Read the saved router settings row instead so database-owned values stay writable. The integration harness seeds num_retries through /config/update instead of the config file, which is what the effective-settings and observed routing tests need to keep exercising a database-owned value. --- .circleci/scripts/run_integration.sh | 3 + litellm/proxy/proxy_server.py | 63 ++++++++++++------- tests/integration/proxy_config.yaml | 1 - .../proxy/proxy_server/test_routes_config.py | 37 +++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 53 ++++++++++++---- 5 files changed, 123 insertions(+), 34 deletions(-) diff --git a/.circleci/scripts/run_integration.sh b/.circleci/scripts/run_integration.sh index 6fab6dd57db..80f9cb0a4e0 100644 --- a/.circleci/scripts/run_integration.sh +++ b/.circleci/scripts/run_integration.sh @@ -125,6 +125,9 @@ start_proxy() { start_proxy 4000 proxy.log proxy_pid="$launched_pid" .venv/bin/python .circleci/scripts/wait_integration_services.py +curl --noproxy '*' -sSf -X POST "$INTEGRATION_PROXY_URL/config/update" \ + -H "Authorization: Bearer $LITELLM_MASTER_KEY" -H 'Content-Type: application/json' \ + -d '{"router_settings": {"num_retries": 0}}' > "$results/seed-router-settings.json" if [ "$suite" = management ]; then export INTEGRATION_PEER_URL=http://127.0.0.1:4001 start_proxy 4001 peer.log diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bb0e432af51..4654dffcb11 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6869,9 +6869,7 @@ class ProxyConfig: self._add_callbacks_from_db_config(config_data) # router settings - await self._add_router_settings_from_db_config( - config_data=config_data, llm_router=llm_router, prisma_client=prisma_client - ) + await self._add_router_settings_from_db_config(llm_router=llm_router, prisma_client=prisma_client) return still_desired_ids @@ -7079,13 +7077,11 @@ class ProxyConfig: async def _add_router_settings_from_db_config( self, - config_data: Mapping[str, object], llm_router: Router | None, prisma_client: PrismaClient | None, ) -> None: if llm_router is None or prisma_client is None: return - self.router_settings.load_yaml(_as_settings_mapping(config_data.get("router_settings"))) db_router_settings: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first( where={"param_name": "router_settings"} ) @@ -16944,6 +16940,46 @@ async def update_config( if prisma_client is None: raise Exception("No DB Connected") + requested_general_settings: Final[Mapping[str, JsonValue]] = ( + config_info.general_settings.model_dump(exclude_none=True, exclude_unset=True) + if config_info.general_settings is not None + else {} + ) + raw_litellm_settings: Final[Mapping[str, JsonValue]] = _CONFIG_SECTION_VALUES.validate_python( + config_info.litellm_settings if config_info.litellm_settings is not None else {} + ) + incoming_success_callback: Final = raw_litellm_settings.get("success_callback") + updated_litellm_settings: Final[Mapping[str, JsonValue]] = _CONFIG_SECTION_VALUES.validate_python( + { + **raw_litellm_settings, + **( + {"success_callback": normalize_callback_names(incoming_success_callback)} + if isinstance(incoming_success_callback, list) + else {} + ), + } + ) + typed_router_settings: Final[Mapping[str, JsonValue]] = ( + config_info.router_settings.model_dump(exclude_none=True) if config_info.router_settings is not None else {} + ) + router_settings_updates: Final[Mapping[str, JsonValue]] = { + **typed_router_settings, + **( + { + key: value + for key, value in raw_router_settings.items() + if key not in typed_router_settings and value is not None + } + if isinstance(raw_router_settings, dict) + else {} + ), + } + proxy_config.reject_config_owned_writes( + section_name="general_settings", changed_keys=requested_general_settings + ) + proxy_config.reject_config_owned_writes(section_name="litellm_settings", changed_keys=updated_litellm_settings) + proxy_config.reject_config_owned_writes(section_name="router_settings", changed_keys=router_settings_updates) + async def _read_section(param_name: str) -> dict: row: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first( where={"param_name": param_name} @@ -17014,15 +17050,9 @@ async def update_config( if config_info.litellm_settings is not None: existing = await _read_section("litellm_settings") before_litellm_settings: Final = copy.deepcopy(existing) - updated_litellm_settings: Final = dict(config_info.litellm_settings) - - incoming_cb = updated_litellm_settings.get("success_callback") - if isinstance(incoming_cb, list): - updated_litellm_settings["success_callback"] = normalize_callback_names(incoming_cb) - merged: Final = {**existing, **updated_litellm_settings} - incoming_cb = updated_litellm_settings.get("success_callback") + incoming_cb: Final = updated_litellm_settings.get("success_callback") existing_cb: Final = existing.get("success_callback") if isinstance(incoming_cb, list): if isinstance(existing_cb, list): @@ -17045,15 +17075,6 @@ async def update_config( if isinstance(raw_router_settings, dict): existing = await _read_section("router_settings") before_router_settings: Final = copy.deepcopy(existing) - typed_router_settings: Final = ( - config_info.router_settings.dict(exclude_none=True) if config_info.router_settings is not None else {} - ) - raw_router_settings_without_none: Final = { - key: value - for key, value in raw_router_settings.items() - if key not in typed_router_settings and value is not None - } - router_settings_updates: Final = {**typed_router_settings, **raw_router_settings_without_none} new_router_settings: Final = {**existing, **router_settings_updates} await _upsert_section("router_settings", new_router_settings) asyncio.create_task( diff --git a/tests/integration/proxy_config.yaml b/tests/integration/proxy_config.yaml index b6f9767c210..a3b07f76d2f 100644 --- a/tests/integration/proxy_config.yaml +++ b/tests/integration/proxy_config.yaml @@ -13,5 +13,4 @@ litellm_settings: host: os.environ/REDIS_HOST port: os.environ/REDIS_PORT router_settings: - num_retries: 0 disable_cooldowns: true diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py index 19b36b026f9..0a2daf64463 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py @@ -139,6 +139,43 @@ def test_config_update_persists_disable_cooldowns(client, auth_as, mock_prisma, assert persisted["disable_cooldowns"] is True +@pytest.mark.parametrize( + ("section", "store_attr", "yaml_values", "changed_values"), + [ + ("general_settings", "settings", {"alerting": ["slack"]}, {"alerting": ["email"]}), + ("litellm_settings", "litellm_settings", {"success_callback": ["langfuse"]}, {"success_callback": ["otel"]}), + ("router_settings", "router_settings", {"num_retries": 0}, {"num_retries": 2}), + ], +) +def test_config_update_rejects_config_owned_keys_and_accepts_the_same_value( + client, auth_as, mock_prisma, monkeypatch, section, store_attr, yaml_values, changed_values +): + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps.proxy_config, "add_deployment", AsyncMock()) + store = getattr(ps.proxy_config, store_attr) + store.load_yaml(yaml_values) + try: + with auth_as(LitellmUserRoles.PROXY_ADMIN): + rejected = client.post("/config/update", json={section: changed_values}) + rejected_message = rejected.json()["error"]["message"] + table.upsert.assert_not_called() + accepted = client.post("/config/update", json={section: yaml_values}) + finally: + store.load_yaml({}) + + assert rejected.status_code == 400 + assert f"{section} key '{next(iter(yaml_values))}' is set in the config file and cannot be changed here" in ( + rejected_message + ) + assert accepted.status_code == 200 + persisted = json.loads(table.upsert.call_args.kwargs["data"]["create"]["param_value"]) + assert persisted[next(iter(yaml_values))] == yaml_values[next(iter(yaml_values))] + + def test_config_update_rejects_assistants_config(client, auth_as, mock_prisma, monkeypatch): from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index bc7ae556e7b..b6f10ca8f76 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -4919,8 +4919,8 @@ async def test_add_router_settings_from_db_config_merge_logic(): mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=mock_db_config) # Call the method under test + proxy_config.router_settings.load_yaml(config_data["router_settings"]) await proxy_config._add_router_settings_from_db_config( - config_data=config_data, llm_router=mock_router, prisma_client=mock_prisma_client, ) @@ -4980,8 +4980,8 @@ async def test_add_router_settings_from_db_config_empty_db_lists_do_not_clobber_ mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=mock_db_config) + proxy_config.router_settings.load_yaml(config_data["router_settings"]) await proxy_config._add_router_settings_from_db_config( - config_data=config_data, llm_router=mock_router, prisma_client=mock_prisma_client, ) @@ -5017,8 +5017,8 @@ async def test_add_router_settings_from_db_config_empty_db_list_still_clears_unc mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=mock_db_config) + proxy_config.router_settings.load_yaml(config_data["router_settings"]) await proxy_config._add_router_settings_from_db_config( - config_data=config_data, llm_router=mock_router, prisma_client=mock_prisma_client, ) @@ -5042,8 +5042,8 @@ async def test_add_router_settings_from_db_config_edge_cases(): mock_router.update_settings = MagicMock() # Test Case 1: No router provided + proxy_config.router_settings.load_yaml({"test": "value"}) await proxy_config._add_router_settings_from_db_config( - config_data={"router_settings": {"test": "value"}}, llm_router=None, prisma_client=MagicMock(), ) @@ -5051,8 +5051,8 @@ async def test_add_router_settings_from_db_config_edge_cases(): mock_router.update_settings.assert_not_called() # Test Case 2: No prisma client provided + proxy_config.router_settings.load_yaml({"test": "value"}) await proxy_config._add_router_settings_from_db_config( - config_data={"router_settings": {"test": "value"}}, llm_router=mock_router, prisma_client=None, ) @@ -5065,8 +5065,8 @@ async def test_add_router_settings_from_db_config_edge_cases(): config_data = {"router_settings": {"routing_strategy": "usage-based"}} + proxy_config.router_settings.load_yaml(config_data["router_settings"]) await proxy_config._add_router_settings_from_db_config( - config_data=config_data, llm_router=mock_router, prisma_client=mock_prisma_client, ) @@ -5080,8 +5080,8 @@ async def test_add_router_settings_from_db_config_edge_cases(): mock_db_config.param_value = {"db_setting": "db_value"} mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=mock_db_config) + proxy_config.router_settings.load_yaml({}) await proxy_config._add_router_settings_from_db_config( - config_data={}, # No router_settings in config llm_router=mock_router, prisma_client=mock_prisma_client, ) @@ -5093,9 +5093,8 @@ async def test_add_router_settings_from_db_config_edge_cases(): # Test Case 5: Both config and DB router_settings are None/empty mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=None) - await proxy_config._add_router_settings_from_db_config( - config_data={}, llm_router=mock_router, prisma_client=mock_prisma_client - ) + proxy_config.router_settings.load_yaml({}) + await proxy_config._add_router_settings_from_db_config(llm_router=mock_router, prisma_client=mock_prisma_client) # Should not call update_settings when no settings exist mock_router.update_settings.assert_not_called() @@ -5107,8 +5106,8 @@ async def test_add_router_settings_from_db_config_edge_cases(): config_data = {"router_settings": {"config_setting": "config_value"}} + proxy_config.router_settings.load_yaml(config_data["router_settings"]) await proxy_config._add_router_settings_from_db_config( - config_data=config_data, llm_router=mock_router, prisma_client=mock_prisma_client, ) @@ -5157,8 +5156,8 @@ async def test_add_router_settings_shallow_merge_behavior(): mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=mock_db_config) + proxy_config.router_settings.load_yaml(config_data["router_settings"]) await proxy_config._add_router_settings_from_db_config( - config_data=config_data, llm_router=mock_router, prisma_client=mock_prisma_client, ) @@ -5180,6 +5179,36 @@ async def test_add_router_settings_shallow_merge_behavior(): assert merged_settings["top_level"] == "config_top" +@pytest.mark.asyncio +async def test_router_settings_reload_keeps_db_values_writable(tmp_path, monkeypatch): + from litellm.proxy.proxy_server import ProxyConfig + + config_path: Final = tmp_path / "config.yaml" + config_path.write_text(yaml.safe_dump({"model_list": [], "router_settings": {"disable_cooldowns": True}})) + db_row: Final = types.SimpleNamespace(param_value={"num_retries": 0}) + + async def read_config_row(_prisma_client, param_name): + return db_row if param_name == "router_settings" else None + + mock_prisma_client: Final = MagicMock() + mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=db_row) + mock_router: Final = MagicMock() + monkeypatch.setattr(proxy_server_module, "get_config_param", read_config_row) + monkeypatch.setattr(proxy_server_module, "prisma_client", mock_prisma_client) + monkeypatch.setattr(proxy_server_module, "store_model_in_db", True) + monkeypatch.setattr(proxy_server_module, "user_config_file_path", None) + proxy_config: Final = ProxyConfig() + + for _ in range(2): + await proxy_config.get_config(config_file_path=str(config_path)) + await proxy_config._add_router_settings_from_db_config(llm_router=mock_router, prisma_client=mock_prisma_client) + + assert mock_router.update_settings.call_args.kwargs == {"disable_cooldowns": True, "num_retries": 0} + assert proxy_config.router_settings.source("num_retries") == "db" + assert proxy_config.router_settings.rejected_writes({"num_retries": 3}) == () + assert proxy_config.router_settings.rejected_writes({"disable_cooldowns": False}) == ("disable_cooldowns",) + + @pytest.mark.asyncio async def test_model_info_v1_oci_secrets_not_leaked(): """ From 8533cb9d5915de44bb8bcd6ca1955d5e0ec51989 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 12:44:42 -0700 Subject: [PATCH 13/28] test(proxy): drop the pass-through routes the reload tests add to the shared app The two pass-through reload tests register routes on the shared app and never take them off, so a later allowlist test in the same xdist worker finds routes that no component exposes. Restore the app's route list when each of those tests finishes. --- tests/test_litellm/proxy/test_proxy_server.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index b6f10ca8f76..87114e183a9 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -7496,7 +7496,15 @@ async def test_update_general_settings_db_pass_through_endpoint_cannot_override_ assert still_open.api_key is None +@pytest.fixture +def app_routes_restored(): + routes_before: Final = tuple(app.router.routes) + yield + app.router.routes[:] = routes_before + + @pytest.mark.asyncio +@pytest.mark.usefixtures("app_routes_restored") async def test_deleting_the_stored_pass_through_row_takes_the_route_out_of_service(): """A pass-through route the database declared has to stop serving when that row is deleted. The proxy's own registry of live pass-through routes is what decides whether @@ -7524,6 +7532,7 @@ async def test_deleting_the_stored_pass_through_row_takes_the_route_out_of_servi @pytest.mark.asyncio +@pytest.mark.usefixtures("app_routes_restored") async def test_a_stored_pass_through_row_never_disturbs_the_config_declared_routes(): """``pass_through_endpoints`` is config-owned once the file declares it, so writing and then deleting a stored row resolves to the same list both times and the config file's routes keep From 1581c45615ccd841560b5123b3843f7b864c6c6f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 13:14:22 -0700 Subject: [PATCH 14/28] fix(proxy): check and persist only the settings the request sent POST /config/update compared litellm_settings after lowercasing the callback list, so a config file spelling a callback in mixed case refused the same list sent back, and it stored every general_settings model default next to the keys the request set. Both now use the request as sent; only the stored callback list is lowercased. Also drops config_data from the router settings reload callers the previous commit left behind and teaches the legacy MockProxyConfig the ownership check. --- litellm/proxy/proxy_server.py | 5 +-- tests/proxy_unit_tests/test_proxy_server.py | 3 ++ .../proxy/proxy_server/test_proxy_config.py | 5 +-- .../proxy/proxy_server/test_routes_config.py | 42 +++++++++++++++++++ 4 files changed, 49 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4654dffcb11..ddee8337609 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -16977,7 +16977,7 @@ async def update_config( proxy_config.reject_config_owned_writes( section_name="general_settings", changed_keys=requested_general_settings ) - proxy_config.reject_config_owned_writes(section_name="litellm_settings", changed_keys=updated_litellm_settings) + proxy_config.reject_config_owned_writes(section_name="litellm_settings", changed_keys=raw_litellm_settings) proxy_config.reject_config_owned_writes(section_name="router_settings", changed_keys=router_settings_updates) async def _read_section(param_name: str) -> dict: @@ -17006,8 +17006,7 @@ async def update_config( if config_info.general_settings is not None: existing = await _read_section("general_settings") before_general_settings: Final = copy.deepcopy(existing) - updates: Mapping[str, JsonValue] = config_info.general_settings.dict(exclude_none=True) - for k, v in updates.items(): + for k, v in requested_general_settings.items(): if k == "alert_to_webhook_url": if "alerting" not in existing: existing["alerting"] = ["slack"] diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 1fcdaa67143..8d1daa8a3d5 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -3069,6 +3069,9 @@ async def test_update_config_success_callback_normalization(): async def add_deployment(self, prisma_client=None, proxy_logging_obj=None): # noqa: F811 # pytest fixture, not a redefinition return None + def reject_config_owned_writes(self, *, section_name, changed_keys): + return None + setattr(proxy_server, "proxy_config", MockProxyConfig()) config_update = ConfigYAML(litellm_settings={"success_callback": ["SQS", "sQs"]}) diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 42cd6e4ed78..1e9a93bd019 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -3399,9 +3399,8 @@ async def test_ProxyConfig__add_router_settings_from_db_config_updates_router(): fake_prisma.db.litellm_config.find_first = AsyncMock( return_value=SimpleNamespace(param_value={"timeout": 30, "retries": 2, "fallbacks": []}) ) - config_data = {"router_settings": {"timeout": 10}} + pc.router_settings.load_yaml({"timeout": 10}) await pc._add_router_settings_from_db_config( - config_data=config_data, llm_router=fake_router, prisma_client=fake_prisma, ) @@ -3421,7 +3420,7 @@ async def test_ProxyConfig__add_router_settings_from_db_config_updates_router(): async def test_ProxyConfig__add_router_settings_from_db_config_none_router_noop(): pc = ProxyConfig() # No router and no prisma — should silently return. - await pc._add_router_settings_from_db_config(config_data={}, llm_router=None, prisma_client=None) + await pc._add_router_settings_from_db_config(llm_router=None, prisma_client=None) # Error-style: bad call signature raises. with pytest.raises(TypeError): await pc._add_router_settings_from_db_config() # type: ignore[call-arg] diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py index 0a2daf64463..e4d80c86f57 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py @@ -176,6 +176,48 @@ def test_config_update_rejects_config_owned_keys_and_accepts_the_same_value( assert persisted[next(iter(yaml_values))] == yaml_values[next(iter(yaml_values))] +def test_config_update_persists_only_the_general_settings_keys_the_request_set( + client, auth_as, mock_prisma, monkeypatch +): + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps.proxy_config, "add_deployment", AsyncMock()) + ps.proxy_config.settings.load_yaml({"health_check_interval": 60}) + try: + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.post("/config/update", json={"general_settings": {"alerting_threshold": 600}}) + finally: + ps.proxy_config.settings.load_yaml({}) + + assert response.status_code == 200 + persisted = json.loads(table.upsert.call_args.kwargs["data"]["create"]["param_value"]) + assert persisted == {"alerting_threshold": 600} + + +def test_config_update_accepts_a_config_owned_success_callback_the_file_spells_in_mixed_case( + client, auth_as, mock_prisma, monkeypatch +): + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps.proxy_config, "add_deployment", AsyncMock()) + ps.proxy_config.litellm_settings.load_yaml({"success_callback": ["Langfuse"]}) + try: + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.post("/config/update", json={"litellm_settings": {"success_callback": ["Langfuse"]}}) + finally: + ps.proxy_config.litellm_settings.load_yaml({}) + + assert response.status_code == 200 + persisted = json.loads(table.upsert.call_args.kwargs["data"]["create"]["param_value"]) + assert persisted["success_callback"] == ["langfuse"] + + def test_config_update_rejects_assistants_config(client, auth_as, mock_prisma, monkeypatch): from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles From 72a14246a6ffef7724492b8c0da2798c30a32148 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 13:50:56 -0700 Subject: [PATCH 15/28] fix(proxy): requeue daily spend rows when the commit fails without the Redis buffer With the Redis transaction buffer off, each daily spend queue (user, team, org, end user, agent) was drained into a dict and handed to the bulk upsert. When the upsert raised after its retries, the drained dict was discarded and the exception escaped update_spend, so those rows never reached the daily rollup tables and /user/daily/activity stayed short forever while /spend/logs had every request. Each daily queue now flushes through one helper that puts the uncommitted remainder back on the queue for the next tick and moves on to the next table, the same shape the window-spend step already used. --- litellm/proxy/db/db_spend_update_writer.py | 99 ++++++++++++------- .../proxy/db/test_db_spend_update_writer.py | 55 +++++++++++ 2 files changed, 118 insertions(+), 36 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index c13b852484e..a4f1713d61d 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -15,7 +15,7 @@ import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, overload from urllib.parse import quote, unquote from typing_extensions import ReadOnly, TypedDict @@ -143,6 +143,20 @@ class _SpendTransactionManager(Protocol): async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ... +_DailySpendTransactionT = TypeVar("_DailySpendTransactionT", bound=BaseDailySpendTransaction) + + +class _DailySpendCommit(Protocol[_DailySpendTransactionT]): + async def __call__( + self, + *, + n_retry_times: int, + prisma_client: PrismaClient, + proxy_logging_obj: ProxyLogging, + daily_spend_transactions: dict[str, _DailySpendTransactionT], + ) -> None: ... + + def _timed_request_duration_ms( payload: dict | SpendLogsPayload, request_status: Literal["success", "failure"], @@ -1288,6 +1302,34 @@ class DBSpendUpdateWriter: cronjob_id=DB_SPEND_UPDATE_JOB_NAME, ) + async def _flush_daily_spend_queue( + self, + queue: DailySpendUpdateQueue, + entity_type: Literal["user", "team", "org", "end_user", "agent"], + commit: _DailySpendCommit[_DailySpendTransactionT], + n_retry_times: int, + prisma_client: PrismaClient, + proxy_logging_obj: ProxyLogging, + ) -> None: + transactions: Final = await queue.flush_and_get_aggregated_daily_spend_update_transactions() + try: + await commit( + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + daily_spend_transactions=cast(dict[str, _DailySpendTransactionT], transactions), + ) + except Exception as e: # noqa: BLE001 # the uncommitted rows go back on the queue; the other tables must still flush + spend_log_error( + "Spend tracking - failed to commit daily %s spend updates. " + "Re-queued %d rows for retry on next tick. Error: %s", + entity_type, + len(transactions), + str(e), + exc=e, + ) + await queue.add_update(transactions) + async def _commit_spend_updates_to_db_without_redis_buffer( self, prisma_client: PrismaClient, @@ -1316,74 +1358,59 @@ class DBSpendUpdateWriter: ################## Daily Spend Update Transactions ################## # Aggregate all in memory daily spend transactions and commit to db - daily_spend_update_transactions: Final = cast( - dict[str, DailyUserSpendTransaction], - await self.daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(), - ) - - await DBSpendUpdateWriter.update_daily_user_spend( + await self._flush_daily_spend_queue( + queue=self.daily_spend_update_queue, + entity_type="user", + commit=DBSpendUpdateWriter.update_daily_user_spend, n_retry_times=n_retry_times, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, - daily_spend_transactions=daily_spend_update_transactions, ) ################## Daily Team Spend Update Transactions ################## # Aggregate all in memory daily team spend transactions and commit to db - daily_team_spend_update_transactions: Final = cast( - dict[str, DailyTeamSpendTransaction], - await self.daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(), - ) - - await DBSpendUpdateWriter.update_daily_team_spend( + await self._flush_daily_spend_queue( + queue=self.daily_team_spend_update_queue, + entity_type="team", + commit=DBSpendUpdateWriter.update_daily_team_spend, n_retry_times=n_retry_times, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, - daily_spend_transactions=daily_team_spend_update_transactions, ) ################## Daily Organization Spend Update Transactions ################## # Aggregate all in memory daily org spend transactions and commit to db - daily_org_spend_update_transactions: Final = cast( - dict[str, DailyOrganizationSpendTransaction], - await self.daily_org_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(), - ) - - await DBSpendUpdateWriter.update_daily_org_spend( + await self._flush_daily_spend_queue( + queue=self.daily_org_spend_update_queue, + entity_type="org", + commit=DBSpendUpdateWriter.update_daily_org_spend, n_retry_times=n_retry_times, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, - daily_spend_transactions=daily_org_spend_update_transactions, ) # NOTE: Daily tag spend is committed by a separate scheduler job. ################## Daily End-User Spend Update Transactions ################## # Aggregate all in memory daily end-user spend transactions and commit to db - daily_end_user_spend_update_transactions: Final = cast( - dict[str, DailyEndUserSpendTransaction], - await self.daily_end_user_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(), - ) - - await DBSpendUpdateWriter.update_daily_end_user_spend( + await self._flush_daily_spend_queue( + queue=self.daily_end_user_spend_update_queue, + entity_type="end_user", + commit=DBSpendUpdateWriter.update_daily_end_user_spend, n_retry_times=n_retry_times, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, - daily_spend_transactions=daily_end_user_spend_update_transactions, ) ################## Daily Agent Spend Update Transactions ################## # Aggregate all in memory daily agent spend transactions and commit to db - daily_agent_spend_update_transactions: Final = cast( - dict[str, DailyAgentSpendTransaction], - await self.daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(), - ) - - await DBSpendUpdateWriter.update_daily_agent_spend( + await self._flush_daily_spend_queue( + queue=self.daily_agent_spend_update_queue, + entity_type="agent", + commit=DBSpendUpdateWriter.update_daily_agent_spend, n_retry_times=n_retry_times, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, - daily_spend_transactions=daily_agent_spend_update_transactions, ) ################## Budget Window Spend Update Transactions ################## diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index b3f5a60877d..3feb117edb7 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2809,6 +2809,61 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_ assert requeued == (transaction,) +class _DailySpendFakeDB(_WindowSpendFakeDB): + """Records the daily rollup upserts it is handed and fails the ones aimed at one table.""" + + def __init__(self, failing_table: str | None) -> None: + super().__init__() + self.failing_table = failing_table + self.execute_raw_calls: list[Statement] = [] + + async def execute_raw(self, query: str, *args: object) -> int: + if self.failing_table is not None and self.failing_table in query: + raise Exception("connection reset") + self.execute_raw_calls.append((query, args)) + return len(args) + + +def _daily_upserts(db: _DailySpendFakeDB, table: str) -> list[Statement]: + return [statement for statement in db.execute_raw_calls if table in statement[0]] + + +@pytest.mark.asyncio +async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other_tables(): + """With the Redis buffer off, a daily batch that failed to commit was discarded along + with the tick's exception, so the Usage page stayed short of LiteLLM_SpendLogs for good. + The uncommitted rows must go back on their queue and land on the next tick, and the + other daily tables must still be flushed on the failing tick.""" + db_writer = DBSpendUpdateWriter() + await db_writer.daily_spend_update_queue.add_update({"user-key": _daily_txn(user_id="user-1")}) + team_txn = {key: value for key, value in _daily_txn().items() if key != "user_id"} | {"team_id": "team-1"} + await db_writer.daily_team_spend_update_queue.add_update({"team-key": team_txn}) + db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend") + db_writer._flush_tool_discovery_queue = AsyncMock() + proxy_logging_obj = MagicMock() + proxy_logging_obj.failure_handler = AsyncMock() + + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj + ) + + assert _daily_upserts(db, "LiteLLM_DailyUserSpend") == [] + (team_upsert,) = _daily_upserts(db, "LiteLLM_DailyTeamSpend") + assert _row_values(team_upsert, "team_id") == ["team-1"] + db_writer._flush_tool_discovery_queue.assert_called_once() + + db.failing_table = None + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj + ) + + (user_upsert,) = _daily_upserts(db, "LiteLLM_DailyUserSpend") + assert _row_values(user_upsert, "user_id") == ["user-1"] + assert _row_values(user_upsert, "spend") == [0.1] + assert len(_daily_upserts(db, "LiteLLM_DailyTeamSpend")) == 1 + assert db_writer.daily_spend_update_queue.update_queue.empty() + + @pytest.mark.asyncio async def test_failed_window_spend_commit_from_redis_is_restored_to_redis(): """The Redis drain is destructive, so a failed window commit has to push From df1b3c849b76b4dff04441a5838427b5b6be9827 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 13:58:53 -0700 Subject: [PATCH 16/28] fix(responses): keep upstream error details in response.failed RateLimitError and InternalServerError now carry the provider body, so the OpenAI exception mapper keeps upstream codes like cyber_policy and the upstream message instead of a generic mapped one The proxy's response.failed event prefers the upstream body's code, message, and type over the mapped exception's, and numeric error codes in an error event map to their own HTTP status --- litellm/exceptions.py | 6 ++-- .../exception_mapping_utils.py | 5 +++ litellm/proxy/_lazy_openapi_snapshot.json | 2 +- .../common_utils/responses_stream_errors.py | 21 +++++++++++- litellm/responses/streaming_iterator.py | 13 +++++--- .../test_exception_mapping_utils.py | 23 +++++++++++++ .../proxy_server/test_streaming_helpers.py | 16 +++++++++ .../response_api_endpoints/test_endpoints.py | 17 ++++++---- .../test_streaming_iterator_error_events.py | 33 +++++++++++++++++++ 9 files changed, 120 insertions(+), 16 deletions(-) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index c638e8a0f86..de9f5c692a1 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -464,6 +464,7 @@ class RateLimitError(openai.RateLimitError): rate_limit_type: str | RateLimitType | None = None, headers: dict[str, str] | None = None, detail: Any = None, + body: object | None = None, ): self.status_code = 429 self.message = f"litellm.RateLimitError: {message}" @@ -507,7 +508,7 @@ class RateLimitError(openai.RateLimitError): ), ) super().__init__( - self.message, response=self.response, body=None + self.message, response=self.response, body=body ) # Call the base class constructor with the parameters it needs self.code = "429" self.type = "throttling_error" @@ -765,6 +766,7 @@ class InternalServerError(openai.InternalServerError): litellm_debug_info: str | None = None, max_retries: int | None = None, num_retries: int | None = None, + body: object | None = None, ): self.status_code = 500 self.message = f"litellm.InternalServerError: {message}" @@ -783,7 +785,7 @@ class InternalServerError(openai.InternalServerError): ), ) super().__init__( - self.message, response=self.response, body=None + self.message, response=self.response, body=body ) # Call the base class constructor with the parameters it needs def __str__(self): diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 70675966dfc..61e2698dd6f 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -307,6 +307,7 @@ def _map_openai_exception( model=model, llm_provider=custom_llm_provider, response=response, + body=getattr(original_exception, "body", None), ) elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str): raise ContextWindowExceededError( @@ -381,6 +382,7 @@ def _map_openai_exception( message=f"{exception_provider} - {message}", model=model, llm_provider=custom_llm_provider, + body=getattr(original_exception, "body", None), ) elif "Request too large" in error_str: raise RateLimitError( @@ -389,6 +391,7 @@ def _map_openai_exception( llm_provider=custom_llm_provider, response=response, litellm_debug_info=extra_information, + body=getattr(original_exception, "body", None), ) elif ( "The api_key client option must be set either by passing api_key to the client or by setting the OPENAI_API_KEY environment variable" @@ -460,6 +463,7 @@ def _map_openai_exception( llm_provider=custom_llm_provider, response=response, litellm_debug_info=extra_information, + body=getattr(original_exception, "body", None), ) elif original_exception.status_code == 500: raise InternalServerError( @@ -468,6 +472,7 @@ def _map_openai_exception( llm_provider=custom_llm_provider, response=response, litellm_debug_info=extra_information, + body=getattr(original_exception, "body", None), ) elif original_exception.status_code == 502: raise BadGatewayError( diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 8a8d08c6887..2aa2cf15ac1 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19394,7 +19394,7 @@ } } }, - "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " + "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" }, "500": { "content": { diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index 706b4c298d7..63f881c126e 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -36,6 +36,11 @@ class _FailureDetails(BaseModel): type: str | None = None status_code: int | None = None + @field_validator("message", mode="before") + @classmethod + def normalize_message(cls, value: object) -> str | None: + return value if isinstance(value, str) else None + @field_validator("code", mode="before") @classmethod def normalize_code(cls, value: object) -> str | int | None: @@ -54,6 +59,20 @@ def _original_failure(exception: Exception) -> Exception: return current +def _failure_details(original: Exception) -> _FailureDetails: + mapped: Final = _FailureDetails.model_validate(original) + body: Final = getattr(original, "body", None) + if not isinstance(body, Mapping): + return mapped + upstream: Final = _FailureDetails.model_validate(body) + return _FailureDetails( + message=upstream.message or mapped.message, + code=upstream.code if upstream.code is not None else mapped.code, + type=upstream.type or mapped.type, + status_code=mapped.status_code, + ) + + def _response_error_code(details: _FailureDetails) -> str: for value in (details.code, details.type): if value == "insufficient_quota": @@ -110,7 +129,7 @@ class ResponsesStreamErrorState: if self.terminal_emitted: return None original: Final = _original_failure(exception) - details: Final = _FailureDetails.model_validate(original) + details: Final = _failure_details(original) response: Final = ResponsesAPIResponse.model_validate( MappingProxyType( { diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 8d766cf1cd0..9e7cfc0fbe7 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -211,18 +211,21 @@ def _error_event_fields(error_obj: object) -> tuple[str, str | None, str | None] raw_code = None message: Final = str(raw_message) if raw_message is not None else "Response API in-stream error" error_type: Final = raw_type if isinstance(raw_type, str) else None - code: Final = raw_code if isinstance(raw_code, str) else None + code: Final = str(raw_code) if isinstance(raw_code, (str, int)) and not isinstance(raw_code, bool) else None return message, error_type, code +def _status_code_for_error_field(field: str) -> int | None: + if field.isdecimal() and 400 <= int(field) <= 599: + return int(field) + return _ERROR_CODE_HTTP_STATUS.get(field) + + def _status_code_for_error_fields(error_type: str | None, error_code: str | None) -> int: fields: Final = tuple(field for field in (error_code, error_type) if field is not None) if any(field.startswith("rate_limit") or field == "insufficient_quota" for field in fields): return 429 - return next( - (_ERROR_CODE_HTTP_STATUS[field] for field in fields if field in _ERROR_CODE_HTTP_STATUS), - 500, - ) + return next((status for status in map(_status_code_for_error_field, fields) if status is not None), 500) def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool: diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index acc6248bf3e..0970526956e 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -1437,6 +1437,29 @@ def test_openai_compatible_vendor_400_keeps_body_but_not_headers(): assert not exc_info.value.response.headers +@pytest.mark.parametrize( + ("status_code", "mapped_class"), [(429, litellm.RateLimitError), (500, litellm.InternalServerError)] +) +def test_openai_429_and_500_keep_body(status_code: int, mapped_class: type[openai.APIError]): + with pytest.raises(mapped_class) as exc_info: + exception_type( + model="gpt-5.4-mini", + original_exception=_openai_handler_error( + "server_error", {}, status_code=status_code, message="upstream cannot complete this response" + ), + custom_llm_provider="openai", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.body == { + **_GUARDRAIL_BLOCK_ERROR, + "type": "server_error", + "code": str(status_code), + "message": "upstream cannot complete this response", + } + + def test_litellm_proxy_repeated_response_header_keeps_each_value(): repeated = [("x-litellm-call-id", "call-guardrail"), ("set-cookie", "a=1"), ("set-cookie", "b=2")] diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index 53e055882e8..9adc2b80741 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -883,6 +883,13 @@ async def test_async_data_generator_mid_stream_exception_yields_error_payload( assert any(isinstance(item, str) and item.startswith('data: {"error":') for item in out) +_UPSTREAM_BODY: Final = { + "code": "cyber_policy", + "message": "Upstream rejected request: flagged for possible cybersecurity risk", + "type": None, +} + + @pytest.mark.asyncio @pytest.mark.parametrize( "terminal,upstream_error,expected_code", @@ -906,6 +913,13 @@ async def test_async_data_generator_mid_stream_exception_yields_error_payload( ), "server_error", id="structured_provider_error_fields", ), + pytest.param( + "upstream_failure", + litellm.InternalServerError( + message="Upstream rejected request", llm_provider="openai", model="gpt-6-astra", body=_UPSTREAM_BODY + ), + "cyber_policy", id="upstream_body_code_and_message", + ), *( pytest.param( "upstream_failure", HTTPException(status_code=status, detail="Upstream rejected request"), @@ -995,6 +1009,8 @@ async def test_responses_stream_keeps_tool_deltas_and_only_emits_a_valid_termina assert "serialize" in failure.response.error["message"].lower() else: assert "Upstream rejected request" in failure.response.error["message"] + if isinstance(upstream_error, litellm.InternalServerError): + assert failure.response.error["message"] == _UPSTREAM_BODY["message"] if isinstance(upstream_error, (HTTPException, litellm.AuthenticationError)): assert upstream_error.status_code == original_status assert payloads[-1]["sequence_number"] > payloads[1]["sequence_number"] diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index 7b79e1b8613..1fb3fef8fb3 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -24,6 +24,7 @@ from litellm.proxy.proxy_server import app ("/v1/responses", "numeric_rate_limit"), ("/v1/responses", "server_error"), ("/v1/responses", "response_failed"), + ("/v1/responses", "cyber_policy"), ("/cursor/chat/completions", "server_error"), ("/v1/chat/completions", "server_error"), ], @@ -31,7 +32,7 @@ from litellm.proxy.proxy_server import app async def test_streaming_upstream_errors_keep_the_client_protocol( monkeypatch: pytest.MonkeyPatch, path: str, - error_kind: Literal["rate_limit", "numeric_rate_limit", "server_error", "response_failed"], + error_kind: Literal["rate_limit", "numeric_rate_limit", "server_error", "response_failed", "cyber_policy"], ) -> None: import litellm.proxy.proxy_server as ps from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -40,7 +41,7 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( message: Final = "Upstream cannot complete this response" code: Final = { "rate_limit": "rate_limit_exceeded", "numeric_rate_limit": "429", - "server_error": "server_error", "response_failed": "server_error", + "server_error": "server_error", "response_failed": "server_error", "cyber_policy": "cyber_policy", }[error_kind] error: Final = {"message": message, "code": code, "type": None, "param": "input"} response: Final = {"id": "resp_upstream", "object": "response", "created_at": 1, @@ -55,13 +56,13 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( failed: Final = ( {"type": "response.failed", "sequence_number": 9, "response": {**response, "status": "failed", "error": error}} - if error_kind == "response_failed" else {"type": "error", "error": error} + if error_kind in ("response_failed", "cyber_policy") else {"type": "error", "error": error} ) chat: Final = {"id": "chatcmpl_partial", "object": "chat.completion.chunk", "created": 1, "model": model, "choices": [{"index": 0, "delta": {"content": "partial"}, "finish_reason": None}]} is_chat: Final = path == "/v1/chat/completions" - partial: Final = path != "/v1/responses" or error_kind in ("numeric_rate_limit", "response_failed") + partial: Final = path != "/v1/responses" or error_kind in ("numeric_rate_limit", "response_failed", "cyber_policy") response_events: Final = (created, tool_added, tool_delta, failed) if partial else (failed,) upstream_events: Final = (chat, {"error": error}) if is_chat else response_events wire: Final = "".join("data: " + json.dumps(event) + "\n\n" for event in upstream_events) @@ -108,9 +109,11 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( assert events[0]["sequence_number"] == 0 assert events[0]["response"]["id"].startswith("resp_") assert events[-1]["response"]["status"] == "failed" - assert events[-1]["response"]["error"]["code"] == ( - "rate_limit_exceeded" if error_kind in ("rate_limit", "numeric_rate_limit") else "server_error" - ) + assert events[-1]["response"]["error"]["code"] == { + "rate_limit": "rate_limit_exceeded", "numeric_rate_limit": "rate_limit_exceeded", + "server_error": "server_error", "response_failed": "server_error", "cyber_policy": "cyber_policy", + }[error_kind] + assert events[-1]["response"]["error"]["message"] == message else: assert events[0]["object"] == "chat.completion.chunk", result.text assert "response.failed" not in result.text diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/test_litellm/responses/test_streaming_iterator_error_events.py index 3d7c220804a..e7cf09909fe 100644 --- a/tests/test_litellm/responses/test_streaming_iterator_error_events.py +++ b/tests/test_litellm/responses/test_streaming_iterator_error_events.py @@ -340,6 +340,36 @@ def test_maybe_raise_for_response_failed_event_with_dict_error(): assert exc_info.value.status_code == 429 +@pytest.mark.parametrize("code", [429, "429"]) +def test_response_failed_numeric_code_maps_to_its_http_status(code: int | str): + iterator = _make_iterator() + mock_response_obj = Mock() + mock_response_obj.error = {"code": code, "message": "throttled"} + chunk = Mock() + chunk.type = "response.failed" + chunk.response = mock_response_obj + with pytest.raises(MidStreamFallbackError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert exc_info.value.status_code == 429 + assert isinstance(exc_info.value.original_exception, litellm.RateLimitError) + + +def test_response_failed_unknown_code_keeps_upstream_code_and_message_on_mapped_exception(): + iterator = _make_iterator() + upstream_message = "This content was flagged for possible cybersecurity risk." + mock_response_obj = Mock() + mock_response_obj.error = {"code": "cyber_policy", "message": upstream_message} + chunk = Mock() + chunk.type = "response.failed" + chunk.response = mock_response_obj + with pytest.raises(MidStreamFallbackError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + mapped = exc_info.value.original_exception + assert isinstance(mapped, litellm.InternalServerError) + assert mapped.code == "cyber_policy" + assert mapped.body == {"message": upstream_message, "type": None, "code": "cyber_policy"} + + def test_maybe_raise_for_error_event_null_error_obj(): """error chunk with no error field: message and code default; wrapped as 500.""" iterator = _make_iterator() @@ -523,6 +553,9 @@ def test_every_openai_sdk_response_error_code_has_explicit_status_mapping(): ("failed_to_download_image", 400), ("image_file_not_found", 400), ("totally_unknown_future_code", 500), + ("429", 429), + ("503", 503), + ("200", 500), ], ) def test_status_code_for_documented_response_error_codes(code: str, expected_status: int): From 6449632c7bc65015431abc08920737248b64b544 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:01:05 -0700 Subject: [PATCH 17/28] fix(proxy): persist only the router settings keys the request set --- litellm/proxy/proxy_server.py | 4 +++- .../proxy/proxy_server/test_routes_config.py | 23 +++++++++++++++++++ .../test_router_retry_policy_update.py | 2 +- 3 files changed, 27 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f43042942de..c7f91ad004e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -16982,7 +16982,9 @@ async def update_config( } ) typed_router_settings: Final[Mapping[str, JsonValue]] = ( - config_info.router_settings.model_dump(exclude_none=True) if config_info.router_settings is not None else {} + config_info.router_settings.model_dump(exclude_none=True, exclude_unset=True) + if config_info.router_settings is not None + else {} ) router_settings_updates: Final[Mapping[str, JsonValue]] = { **typed_router_settings, diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_config.py b/tests/test_litellm/proxy/proxy_server/test_routes_config.py index e4d80c86f57..4234cdad23d 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_config.py @@ -197,6 +197,29 @@ def test_config_update_persists_only_the_general_settings_keys_the_request_set( assert persisted == {"alerting_threshold": 600} +def test_config_update_persists_only_the_router_settings_keys_the_request_set( + client, auth_as, mock_prisma, monkeypatch +): + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + table = _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps.proxy_config, "add_deployment", AsyncMock()) + ps.proxy_config.router_settings.load_yaml({"model_group_alias": {"opus": "claude-opus-5"}}) + try: + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.post( + "/config/update", json={"router_settings": {"retry_policy": {"TimeoutErrorRetries": 3}}} + ) + finally: + ps.proxy_config.router_settings.load_yaml({}) + + assert response.status_code == 200, response.text + persisted = json.loads(table.upsert.call_args.kwargs["data"]["create"]["param_value"]) + assert persisted == {"retry_policy": {"TimeoutErrorRetries": 3}} + + def test_config_update_accepts_a_config_owned_success_callback_the_file_spells_in_mixed_case( client, auth_as, mock_prisma, monkeypatch ): diff --git a/tests/test_litellm/test_router_retry_policy_update.py b/tests/test_litellm/test_router_retry_policy_update.py index be568134763..0a3dcba325a 100644 --- a/tests/test_litellm/test_router_retry_policy_update.py +++ b/tests/test_litellm/test_router_retry_policy_update.py @@ -360,7 +360,7 @@ async def test_config_update_persists_and_reads_back_retry_policy(monkeypatch): async def _apply_router_settings(*args, **kwargs): await proxy_server.proxy_config._add_router_settings_from_db_config( - config_data={}, llm_router=router, prisma_client=prisma_client + llm_router=router, prisma_client=prisma_client ) monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) From eb96d885ceff1b8daa6c04330e819e39653ba33b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:16:47 -0700 Subject: [PATCH 18/28] fix(proxy): end failed responses streams with [DONE] Emit data: [DONE] after event: response.failed, and after a late failure when a terminal event already went out, so OpenAI SDK clients see the same stream end as a completed response. Restore the lazy OpenAPI snapshot to its Python 3.12 rendering, which is what CI regenerates. --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- litellm/proxy/proxy_server.py | 6 ++++-- .../proxy/proxy_server/test_streaming_helpers.py | 2 ++ .../proxy/response_api_endpoints/test_endpoints.py | 3 ++- 4 files changed, 9 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 2aa2cf15ac1..8a8d08c6887 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19394,7 +19394,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9cd233e5ad5..38a883a3b1b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8865,6 +8865,7 @@ def _format_streaming_sse_chunk(chunk: str | bytes) -> str | bytes: _SSE_FRAME_DELIMITERS: Final = ("\r\n\r\n", "\n\n", "\r\r") +_OPENAI_STREAM_DONE_FRAME: Final = "data: [DONE]\n\n" _MAX_RAW_SSE_BUFFER_CHARS: Final = 8 * 1024 * 1024 @@ -9272,8 +9273,7 @@ async def async_data_generator( yield error_message # OpenAI-compatible streams terminate with data: [DONE]; Google GenAI (?alt=sse) does not. if not request_data.get("_litellm_skip_openai_stream_done"): - done_message: Final = "[DONE]" - yield f"data: {done_message}\n\n" + yield _OPENAI_STREAM_DONE_FRAME except (asyncio.CancelledError, GeneratorExit): # Client disconnected mid-stream. CancelledError / GeneratorExit are # BaseException, so they bypass the success/failure logging callbacks @@ -9303,6 +9303,8 @@ async def async_data_generator( error_frame: Final = error_state.format_failure(e) if error_frame is not None: yield error_frame + if not request_data.get("_litellm_skip_openai_stream_done"): + yield _OPENAI_STREAM_DONE_FRAME return if isinstance(e, HTTPException): raise e diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index 9adc2b80741..86dd356e5f5 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -995,6 +995,8 @@ async def test_responses_stream_keeps_tool_deltas_and_only_emits_a_valid_termina for frame in event_frames ) + assert decoded[-1] == "data: [DONE]\n\n" + assert len(decoded) == len(event_frames) + 1 assert payloads[0]["response"]["id"] == "resp_visible" assert payloads[1] == tool_delta.model_dump() assert len(payloads) == 3 diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index 1fb3fef8fb3..63faef2366f 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -95,7 +95,8 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( assert result.status_code == 200, result.text assert message in result.text if path == "/v1/responses": - assert frames[-1].startswith("event: response.failed\n"), result.text + assert frames[-1] == "data: [DONE]", result.text + assert frames[-2].startswith("event: response.failed\n"), result.text if partial: assert [event["type"] for event in events] == [ "response.created", "response.output_item.added", From 595768b54b7ccfefdee9234c0bb70fbe219ee255 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 18 Sep 2026 14:33:41 -0700 Subject: [PATCH 19/28] fix(proxy): narrow project_id without a cast and fold the unbudgeted cases into the budget matrix test The lint job failed on one new typing.cast (LIT006) in the cost callback, and the test-quality gate behind it would have failed next on a test whose only assertion inspected a mock (TQ002). project_id is now narrowed with isinstance, and the zero and negative max_budget cases run through the existing parametrized budget test, which asserts the raised error or a clean admit with no alert --- .../proxy/hooks/proxy_track_cost_callback.py | 6 ++- .../proxy/auth/test_auth_checks.py | 45 ++++++------------- 2 files changed, 19 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 5e525108ade..903255c7b6c 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -319,7 +319,11 @@ class _ProxyDBLogger(CustomLogger): user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None)) team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None)) org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None)) - project_id: Final = cast(str | None, metadata.get("user_api_key_project_id", None)) + project_id: Final = ( + project_id_value + if isinstance(project_id_value := metadata.get("user_api_key_project_id"), str) + else None + ) key_alias: Final = cast(str | None, metadata.get("user_api_key_alias", None)) end_user_max_budget: Final = metadata.get("user_api_end_user_max_budget", None) sl_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index beadaa27831..85673df57ba 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7559,15 +7559,19 @@ def _project_with_budget(spend: float, max_budget: float): @pytest.mark.asyncio @pytest.mark.parametrize( - "counter_spend, db_spend, blocks", + "counter_spend, db_spend, max_budget, blocks", [ - pytest.param(5.0, 0.0, True, id="counter-at-budget-blocks-despite-stale-db-row"), - pytest.param(4.99, 0.0, False, id="counter-under-budget-admits"), - pytest.param(None, 5.0, True, id="no-counter-falls-back-to-persisted-spend"), - pytest.param(None, 0.0, False, id="no-counter-and-no-persisted-spend-admits"), + pytest.param(5.0, 0.0, 5.0, True, id="counter-at-budget-blocks-despite-stale-db-row"), + pytest.param(4.99, 0.0, 5.0, False, id="counter-under-budget-admits"), + pytest.param(None, 5.0, 5.0, True, id="no-counter-falls-back-to-persisted-spend"), + pytest.param(None, 0.0, 5.0, False, id="no-counter-and-no-persisted-spend-admits"), + pytest.param(12.5, 12.5, 0.0, False, id="zero-budget-is-unbudgeted"), + pytest.param(12.5, 12.5, -1.0, False, id="negative-budget-is-unbudgeted"), ], ) -async def test_project_max_budget_check_reads_live_spend_counter(counter_spend, db_spend, blocks): +async def test_project_max_budget_check_blocks_only_when_live_spend_reaches_a_positive_budget( + counter_spend, db_spend, max_budget, blocks +): from litellm.caching.dual_cache import DualCache from litellm.proxy.auth.auth_checks import _project_max_budget_check @@ -7583,14 +7587,16 @@ async def test_project_max_budget_check_reads_live_spend_counter(counter_spend, ): if not blocks: await _project_max_budget_check( - project_object=_project_with_budget(spend=db_spend, max_budget=5.0), + project_object=_project_with_budget(spend=db_spend, max_budget=max_budget), valid_token=valid_token, proxy_logging_obj=proxy_logging_obj, ) + await asyncio.sleep(0) + proxy_logging_obj.budget_alerts.assert_not_awaited() return with pytest.raises(litellm.BudgetExceededError) as exc_info: await _project_max_budget_check( - project_object=_project_with_budget(spend=db_spend, max_budget=5.0), + project_object=_project_with_budget(spend=db_spend, max_budget=max_budget), valid_token=valid_token, proxy_logging_obj=proxy_logging_obj, ) @@ -7603,29 +7609,6 @@ async def test_project_max_budget_check_reads_live_spend_counter(counter_spend, assert proxy_logging_obj.budget_alerts.await_args.kwargs["type"] == "project_budget" -@pytest.mark.asyncio -@pytest.mark.parametrize("max_budget", [0.0, -1.0]) -async def test_project_max_budget_check_treats_non_positive_budget_as_unbudgeted(max_budget): - from litellm.caching.dual_cache import DualCache - from litellm.proxy.auth.auth_checks import _project_max_budget_check - - real_spend_counter_cache = DualCache() - real_spend_counter_cache.in_memory_cache.set_cache(key="spend:project:p-budget", value=12.5) - proxy_logging_obj = MagicMock() - proxy_logging_obj.budget_alerts = AsyncMock() - - with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock - "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache - ): - await _project_max_budget_check( - project_object=_project_with_budget(spend=12.5, max_budget=max_budget), - valid_token=UserAPIKeyAuth(api_key="hashed-key", project_id="p-budget"), - proxy_logging_obj=proxy_logging_obj, - ) - - proxy_logging_obj.budget_alerts.assert_not_awaited() - - def test_is_user_proxy_admin_rejects_view_only_admin(): """This predicate skips `non_proxy_admin_allowed_routes_check` entirely, so an Admin Viewer answering True here would gain every write route. Read parity for From 76d1abba723116da947f81368fe70939ad976b3c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:57:15 -0700 Subject: [PATCH 20/28] refactor(responses): map status codes to error codes with a lookup Ends _response_error_code in an unconditional return so CodeQL stops flagging mixed explicit and implicit returns. No behavior change: every status maps as before. --- .../common_utils/responses_stream_errors.py | 33 ++++++++++--------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index 63f881c126e..25b83a5ba34 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -73,6 +73,23 @@ def _failure_details(original: Exception) -> _FailureDetails: ) +_CLIENT_ERROR_CODES: Final = MappingProxyType( + { + int(HTTPStatus.UNAUTHORIZED): "authentication_error", + int(HTTPStatus.FORBIDDEN): "permission_error", + int(HTTPStatus.NOT_FOUND): "not_found_error", + int(HTTPStatus.REQUEST_TIMEOUT): "request_timeout", + int(HTTPStatus.TOO_MANY_REQUESTS): "rate_limit_exceeded", + } +) + + +def _status_error_code(status_code: int | None) -> str: + if status_code is None or not HTTPStatus.BAD_REQUEST <= status_code < HTTPStatus.INTERNAL_SERVER_ERROR: + return "server_error" + return _CLIENT_ERROR_CODES.get(status_code, "invalid_request_error") + + def _response_error_code(details: _FailureDetails) -> str: for value in (details.code, details.type): if value == "insufficient_quota": @@ -81,21 +98,7 @@ def _response_error_code(details: _FailureDetails) -> str: return "rate_limit_exceeded" if isinstance(details.code, str) and details.code and not details.code.isdecimal(): return details.code - match details.status_code: - case HTTPStatus.UNAUTHORIZED: - return "authentication_error" - case HTTPStatus.FORBIDDEN: - return "permission_error" - case HTTPStatus.NOT_FOUND: - return "not_found_error" - case HTTPStatus.REQUEST_TIMEOUT: - return "request_timeout" - case HTTPStatus.TOO_MANY_REQUESTS: - return "rate_limit_exceeded" - case int(status) if HTTPStatus.BAD_REQUEST <= status < HTTPStatus.INTERNAL_SERVER_ERROR: - return "invalid_request_error" - case _: - return "server_error" + return _status_error_code(details.status_code) class ResponsesStreamErrorState: From 3edbf60e9c1d3070ffe6e1e70bea39c688f39a21 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:58:34 -0700 Subject: [PATCH 21/28] fix(proxy): requeue the daily tag rollup on commit failure without the Redis buffer --- litellm/proxy/db/db_spend_update_writer.py | 20 +++++-------- .../proxy/db/test_db_spend_update_writer.py | 29 +++++++++++++++++++ 2 files changed, 37 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index a4f1713d61d..b2e6f9dc54d 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1305,7 +1305,7 @@ class DBSpendUpdateWriter: async def _flush_daily_spend_queue( self, queue: DailySpendUpdateQueue, - entity_type: Literal["user", "team", "org", "end_user", "agent"], + entity_type: Literal["user", "team", "org", "tag", "end_user", "agent"], commit: _DailySpendCommit[_DailySpendTransactionT], n_retry_times: int, prisma_client: PrismaClient, @@ -1447,19 +1447,15 @@ class DBSpendUpdateWriter: Commit only tag spend updates to database. This is called by a separate scheduler job at a longer interval. """ - daily_tag_spend_update_transactions: Final = cast( - dict[str, DailyTagSpendTransaction], - await self.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(), + await self._flush_daily_spend_queue( + queue=self.daily_tag_spend_update_queue, + entity_type="tag", + commit=DBSpendUpdateWriter.update_daily_tag_spend, + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, ) - if daily_tag_spend_update_transactions: - await DBSpendUpdateWriter.update_daily_tag_spend( - n_retry_times=n_retry_times, - prisma_client=prisma_client, - proxy_logging_obj=proxy_logging_obj, - daily_spend_transactions=daily_tag_spend_update_transactions, - ) - async def _commit_daily_tag_spend_to_db_with_redis( self, prisma_client: PrismaClient, diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 3feb117edb7..b27b838133b 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2864,6 +2864,35 @@ async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other assert db_writer.daily_spend_update_queue.update_queue.empty() +@pytest.mark.asyncio +async def test_failed_daily_tag_spend_commit_requeues_the_rows(): + """The tag rollup drains on its own scheduler job with the same no-Redis drop: + a failed LiteLLM_DailyTagSpend commit has to put the rows back for the next tick.""" + db_writer = DBSpendUpdateWriter() + tag_txn = {key: value for key, value in _daily_txn().items() if key != "user_id"} | {"tag": "tag-1"} + await db_writer.daily_tag_spend_update_queue.add_update({"tag-key": tag_txn}) + db = _DailySpendFakeDB(failing_table="LiteLLM_DailyTagSpend") + proxy_logging_obj = MagicMock() + proxy_logging_obj.failure_handler = AsyncMock() + + await db_writer._commit_daily_tag_spend_to_db( + prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj + ) + + assert _daily_upserts(db, "LiteLLM_DailyTagSpend") == [] + assert not db_writer.daily_tag_spend_update_queue.update_queue.empty() + + db.failing_table = None + await db_writer._commit_daily_tag_spend_to_db( + prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj + ) + + (tag_upsert,) = _daily_upserts(db, "LiteLLM_DailyTagSpend") + assert _row_values(tag_upsert, "tag") == ["tag-1"] + assert _row_values(tag_upsert, "spend") == [0.1] + assert db_writer.daily_tag_spend_update_queue.update_queue.empty() + + @pytest.mark.asyncio async def test_failed_window_spend_commit_from_redis_is_restored_to_redis(): """The Redis drain is destructive, so a failed window commit has to push From d057e82e6482ec5a75f7952db45cb5a7fe7aa973 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:13:51 -0700 Subject: [PATCH 22/28] test(proxy): assert stored login throttle limits never outrank the config file --- tests/test_litellm/proxy/test_proxy_server.py | 43 +++++++++---------- 1 file changed, 20 insertions(+), 23 deletions(-) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 9309c318573..1b3a460c730 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -14060,32 +14060,29 @@ async def test_authoritative_floor_spend_keeps_a_reset_marker_written_during_the @pytest.mark.asyncio -async def test_login_throttle_settings_are_not_hot_applied_from_the_database(): - """LIT-5285: a stored sign-in limit does not take effect on a live worker. - - _update_general_settings copies an allowlist of keys out of the DB row on every config - poll. Adding these to it would let a stored value outrank config.yaml without a restart, - so an operator locked out by a bad value could not fix it by editing YAML and restarting. - """ +async def test_login_throttle_limits_from_the_config_file_outrank_the_database(monkeypatch): import litellm.proxy.proxy_server as ps from litellm.proxy.proxy_server import ProxyConfig - original = dict(ps.general_settings) - try: - ps.general_settings.clear() - await ProxyConfig()._update_general_settings( - db_general_settings={ - "max_failed_login_attempts_per_source": 999, - "failed_login_window_seconds": 1, - "failed_login_block_seconds": 1, - } - ) - assert "max_failed_login_attempts_per_source" not in ps.general_settings - assert "failed_login_window_seconds" not in ps.general_settings - assert "failed_login_block_seconds" not in ps.general_settings - finally: - ps.general_settings.clear() - ps.general_settings.update(original) + monkeypatch.setattr( + ps, + "general_settings", + { + "max_failed_login_attempts_per_source": 10, + "failed_login_window_seconds": 60, + "failed_login_block_seconds": 300, + }, + ) + await ProxyConfig()._update_general_settings( + db_general_settings={ + "max_failed_login_attempts_per_source": 999, + "failed_login_window_seconds": 1, + "failed_login_block_seconds": 1, + } + ) + assert ps.general_settings.get("max_failed_login_attempts_per_source") == 10 + assert ps.general_settings.get("failed_login_window_seconds") == 60 + assert ps.general_settings.get("failed_login_block_seconds") == 300 @pytest.mark.asyncio From c181c927d0b7a4ad214a4dd160c2f16f1373385e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:58:15 -0700 Subject: [PATCH 23/28] fix(proxy): record response.failed frames in background polling --- .../response_polling/background_streaming.py | 7 +++- .../test_response_polling_handler.py | 38 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index fac45d4391c..b13042dfb6c 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -75,6 +75,10 @@ class _StreamEventParser: parse: Callable[[str], _StreamEvent] = staticmethod(json.loads) +def _sse_frame_data(frame: str) -> str | None: + return next((line[6:].strip() for line in frame.splitlines() if line.startswith("data: ")), None) + + async def _never_receive() -> Message: await asyncio.Event().wait() raise AssertionError("unreachable") @@ -224,8 +228,7 @@ async def background_streaming_task( if isinstance(chunk, bytes): chunk = chunk.decode("utf-8") - if isinstance(chunk, str) and chunk.startswith("data: "): - chunk_data = chunk[6:].strip() + if isinstance(chunk, str) and (chunk_data := _sse_frame_data(chunk)) is not None: if chunk_data == "[DONE]": break diff --git a/tests/proxy_unit_tests/test_response_polling_handler.py b/tests/proxy_unit_tests/test_response_polling_handler.py index 467c1332325..81ca0114a8d 100644 --- a/tests/proxy_unit_tests/test_response_polling_handler.py +++ b/tests/proxy_unit_tests/test_response_polling_handler.py @@ -1482,6 +1482,44 @@ class TestBackgroundStreamingTerminalEvents: assert final_call.kwargs["status"] == "failed" assert final_call.kwargs["error"] == error_payload + @pytest.mark.asyncio + async def test_named_event_failed_frame_sets_failed_status_and_error(self): + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) + + error_payload = { + "code": "cyber_policy", + "message": "Your request was flagged for possible cybersecurity risk and was not completed", + } + failed_event = { + "type": "response.failed", + "sequence_number": 5, + "response": {"id": "resp_123", "status": "failed", "error": error_payload, "output": []}, + } + + async def _body_iterator(): + yield b'data: {"type": "response.in_progress"}\n\n' + yield f"event: response.failed\ndata: {json.dumps(failed_event)}\n\n".encode() + yield b"data: [DONE]\n\n" + + mock_response = Mock() + mock_response.body_iterator = _body_iterator() + handler = AsyncMock(spec=ResponsePollingHandler) + kwargs = _make_background_streaming_kwargs("poll_named_event", handler) + + with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = AsyncMock( + return_value=mock_response + ) + await background_streaming_task(**kwargs) + + final_call = handler.update_state.call_args_list[-1] + assert final_call.kwargs["status"] == "failed" + assert final_call.kwargs["error"] == error_payload + @pytest.mark.asyncio async def test_response_incomplete_sets_incomplete_status_and_details(self): """Test that a response.incomplete stream event results in incomplete status""" From b3cf45e9f232c094e2f2b2a6bf59f609464bf742 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:58:24 -0700 Subject: [PATCH 24/28] fix(proxy): drop daily spend batches that cannot be re-sent safely instead of requeueing them --- litellm/proxy/db/db_spend_update_writer.py | 24 ++++++++- litellm/proxy/db/exception_handler.py | 18 +++++++ .../proxy/db/test_db_spend_update_writer.py | 51 ++++++++++++++++++- .../proxy/db/test_exception_handler.py | 22 ++++++++ 4 files changed, 112 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index b2e6f9dc54d..e9967fb0d67 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -31,6 +31,7 @@ from litellm.constants import ( from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( + DB_CONNECTION_ERROR_TYPES, DB_RETRY_SAFE_ERROR_TYPES, BaseDailySpendTransaction, DailyAgentSpendTransaction, @@ -64,6 +65,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( WindowSpendTransaction, WindowSpendUpdateQueue, ) +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING from litellm.proxy.spend_tracking.compression_savings import ( extract_compression_saved_tokens, @@ -157,6 +159,16 @@ class _DailySpendCommit(Protocol[_DailySpendTransactionT]): ) -> None: ... +_DATA_REJECTED_SQLSTATE_CLASSES: Final = frozenset({"22", "23"}) + + +def _daily_spend_commit_failure_is_requeue_safe(e: Exception) -> bool: + if isinstance(e, DB_CONNECTION_ERROR_TYPES): + return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) + sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) + return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES + + def _timed_request_duration_ms( payload: dict | SpendLogsPayload, request_status: Literal["success", "failure"], @@ -1319,7 +1331,17 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=cast(dict[str, _DailySpendTransactionT], transactions), ) - except Exception as e: # noqa: BLE001 # the uncommitted rows go back on the queue; the other tables must still flush + except Exception as e: # noqa: BLE001 # whatever failed here, the other tables must still flush + if not _daily_spend_commit_failure_is_requeue_safe(e): + spend_log_error( + "Spend tracking - dropped %d daily %s spend rows: the failed commit may have applied " + "or the database refused the data, so re-sending it is not safe. Error: %s", + len(transactions), + entity_type, + str(e), + exc=e, + ) + return spend_log_error( "Spend tracking - failed to commit daily %s spend updates. " "Re-queued %d rows for retry on next tick. Error: %s", diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 2cee5128c66..460bf5db3b1 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -1,6 +1,8 @@ from collections.abc import Awaitable, Callable, Iterator from typing import Any, Final, TypeVar +from pydantic import TypeAdapter, ValidationError + from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, @@ -17,6 +19,8 @@ _TRANSIENT_DB_UNAVAILABLE_MESSAGE: Final = ( "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly." ) +_DATABASE_ERROR_META: Final = TypeAdapter(dict[str, object]) + def _exception_chain(e: BaseException) -> Iterator[BaseException]: current = e # rebind-ok: advances one link per iteration of the bounded walk @@ -221,6 +225,20 @@ class PrismaDBExceptionHandler: or "write conflict or a deadlock" in error_message ) + @staticmethod + def postgres_sqlstate(e: Exception) -> str | None: + """The SQLSTATE Postgres attached to a failed statement, as prisma surfaces it, or None.""" + import prisma + + if not isinstance(e, _exception_types(prisma.errors.DataError)): + return None + try: + meta: Final = _DATABASE_ERROR_META.validate_python(getattr(e, "meta", None)) + except ValidationError: + return None + code: Final = meta.get("code") + return code if isinstance(code, str) else None + @staticmethod def is_read_only_transaction_error(e: Exception) -> bool: """True iff ``e`` is Postgres SQLSTATE 25006 surfaced through prisma: the diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index b27b838133b..155bca656d5 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -11,7 +11,9 @@ from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, call, patch +import httpx import pytest +from prisma.errors import RawQueryError from redis.exceptions import DataError import litellm @@ -2812,14 +2814,15 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_ class _DailySpendFakeDB(_WindowSpendFakeDB): """Records the daily rollup upserts it is handed and fails the ones aimed at one table.""" - def __init__(self, failing_table: str | None) -> None: + def __init__(self, failing_table: str | None, failure: Exception | None = None) -> None: super().__init__() self.failing_table = failing_table + self.failure = failure self.execute_raw_calls: list[Statement] = [] async def execute_raw(self, query: str, *args: object) -> int: if self.failing_table is not None and self.failing_table in query: - raise Exception("connection reset") + raise self.failure if self.failure is not None else Exception("connection reset") self.execute_raw_calls.append((query, args)) return len(args) @@ -2828,6 +2831,50 @@ def _daily_upserts(db: _DailySpendFakeDB, table: str) -> list[Statement]: return [statement for statement in db.execute_raw_calls if table in statement[0]] +def _postgres_rejection(sqlstate: str) -> RawQueryError: + return RawQueryError( + data={"user_facing_error": {"error_code": "P2010", "meta": {"code": sqlstate, "message": "db error"}}} + ) + + +@pytest.mark.parametrize( + ("failure", "lands_on_the_next_tick"), + [ + pytest.param(httpx.ReadTimeout("no reply"), False, id="reply lost after the statement was sent"), + pytest.param(httpx.ConnectError("refused"), True, id="statement never reached the database"), + pytest.param(_postgres_rejection("22021"), False, id="postgres refused the data itself"), + pytest.param(_postgres_rejection("23502"), False, id="postgres refused a constraint violation"), + pytest.param(_postgres_rejection("42P01"), True, id="table missing"), + pytest.param(_postgres_rejection("57014"), True, id="statement cancelled"), + ], +) +@pytest.mark.asyncio +async def test_failed_daily_spend_commit_is_requeued_only_when_the_rows_are_provably_uncommitted( + failure: Exception, lands_on_the_next_tick: bool +): + """A lost reply means the statement may already have applied, and re-sending it stacks a + second increment into the same transaction (LIT-4823); a row Postgres refuses would fail + every tick forever. Both are dropped loudly. Every other failure left nothing committed, + so its rows go back on the queue and land on the next tick.""" + db_writer = DBSpendUpdateWriter() + await db_writer.daily_spend_update_queue.add_update({"user-key": _daily_txn(user_id="user-1")}) + db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=failure) + db_writer._flush_tool_discovery_queue = AsyncMock() + proxy_logging_obj = MagicMock() + proxy_logging_obj.failure_handler = AsyncMock() + + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj + ) + db.failing_table = None + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj + ) + + assert len(_daily_upserts(db, "LiteLLM_DailyUserSpend")) == (1 if lands_on_the_next_tick else 0) + assert db_writer.daily_spend_update_queue.update_queue.empty() + + @pytest.mark.asyncio async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other_tables(): """With the Redis buffer off, a daily batch that failed to commit was discarded along diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index 3f009137a1c..26ac1ea65ad 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -665,6 +665,28 @@ def test_is_deadlock_error_excludes_non_deadlocks(error): assert PrismaDBExceptionHandler.is_deadlock_error(error) is False +@pytest.mark.parametrize( + ("error", "sqlstate"), + [ + ( + RawQueryError( + data={"user_facing_error": {"error_code": "P2010", "meta": {"code": "22021", "message": "m"}}} + ), + "22021", + ), + (RawQueryError(data={"user_facing_error": {"error_code": "P2010", "meta": {"message": "m"}}}), None), + (RawQueryError(data={"user_facing_error": {"error_code": "P2010", "meta": {"code": 42, "message": "m"}}}), None), + (prisma_errors.DataError(data={"user_facing_error": {"meta": None}}), None), + (PrismaError("db error"), None), + (httpx.ReadTimeout("no reply"), None), + ], +) +def test_postgres_sqlstate_reads_the_code_prisma_attached_to_the_failed_statement(error, sqlstate): + """Only a prisma data error carrying Postgres's own error code yields a SQLSTATE; a + codeless or malformed payload, an engine-level error, and a transport error yield None.""" + assert PrismaDBExceptionHandler.postgres_sqlstate(error) == sqlstate + + READ_ONLY_CONNECTOR_ERROR: Final = ( "Error occurred during query execution:\nConnectorError(ConnectorError { user_facing_error: None, " 'kind: QueryError(PostgresError { code: "25006", message: "cannot execute UPDATE in a read-only transaction", ' From 8005856411ae43091c86ba2632d389916b8fa0ec Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:03:27 -0700 Subject: [PATCH 25/28] ci(build_and_test): seed the routing strategy through /config/update --- .circleci/config.yml | 6 ++++++ proxy_server_config.yaml | 1 - 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index df17a9e4402..87f1ee604cf 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1785,6 +1785,12 @@ jobs: - wait_for_service: url: http://localhost:4000 timeout: "300" + - run: + name: Seed the routing strategy through /config/update + command: | + curl --noproxy '*' -sSf -X POST http://localhost:4000/config/update \ + -H 'Authorization: Bearer sk-1234' -H 'Content-Type: application/json' \ + -d '{"router_settings": {"routing_strategy": "usage-based-routing-v2"}}' - run: name: Run tests command: | diff --git a/proxy_server_config.yaml b/proxy_server_config.yaml index 73990153227..703d56bc0cd 100644 --- a/proxy_server_config.yaml +++ b/proxy_server_config.yaml @@ -213,7 +213,6 @@ files_settings: api_key: os.environ/OPENAI_API_KEY router_settings: - routing_strategy: usage-based-routing-v2 redis_host: os.environ/REDIS_HOST redis_password: os.environ/REDIS_PASSWORD redis_port: os.environ/REDIS_PORT From 5db2a0c8850385b9e56b17bcd8053bab4fac0b59 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:05:42 -0700 Subject: [PATCH 26/28] test(proxy): type the sqlstate test parameters --- tests/test_litellm/proxy/db/test_exception_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index 26ac1ea65ad..f7cc5e3ed83 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -681,7 +681,7 @@ def test_is_deadlock_error_excludes_non_deadlocks(error): (httpx.ReadTimeout("no reply"), None), ], ) -def test_postgres_sqlstate_reads_the_code_prisma_attached_to_the_failed_statement(error, sqlstate): +def test_postgres_sqlstate_reads_the_code_prisma_attached_to_the_failed_statement(error: Exception, sqlstate: str | None): """Only a prisma data error carrying Postgres's own error code yields a SQLSTATE; a codeless or malformed payload, an engine-level error, and a transport error yield None.""" assert PrismaDBExceptionHandler.postgres_sqlstate(error) == sqlstate From c6c8aed3f8594f89a86b91df553b84f6aab2fb20 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:22:04 -0700 Subject: [PATCH 27/28] fix(proxy): drop only the daily spend batch whose failure cannot be re-sent, requeue the unsent ones --- litellm/proxy/db/db_spend_update_writer.py | 30 ++++++----- .../proxy/db/test_db_spend_update_writer.py | 54 +++++++++++++++++++ 2 files changed, 71 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e9967fb0d67..37eac8604bd 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1332,15 +1332,7 @@ class DBSpendUpdateWriter: daily_spend_transactions=cast(dict[str, _DailySpendTransactionT], transactions), ) except Exception as e: # noqa: BLE001 # whatever failed here, the other tables must still flush - if not _daily_spend_commit_failure_is_requeue_safe(e): - spend_log_error( - "Spend tracking - dropped %d daily %s spend rows: the failed commit may have applied " - "or the database refused the data, so re-sending it is not safe. Error: %s", - len(transactions), - entity_type, - str(e), - exc=e, - ) + if not transactions: return spend_log_error( "Spend tracking - failed to commit daily %s spend updates. " @@ -2050,13 +2042,25 @@ class DBSpendUpdateWriter: sql, params = build_bulk_upsert(table=table, batch=merged_batch) await prisma_client.db.execute_raw(sql, *params) except Exception as batch_error: - # Log detailed error information for debugging batch upsert failures - # This helps diagnose issues like unique constraint violations + if _daily_spend_commit_failure_is_requeue_safe(batch_error): + spend_log_error( + "Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s", + entity_type, + table.name, + len(transactions_to_process), + str(batch_error), + exc=batch_error, + ) + raise + for key in transactions_to_process: + daily_spend_transactions.pop(key, None) spend_log_error( - "Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s", + "Spend tracking - dropped %d daily %s spend rows: the failed statement may have " + "applied or the database refused the data, so re-sending it is not safe. " + "Table: %s, Error: %s", + len(transactions_to_process), entity_type, table.name, - len(transactions_to_process), str(batch_error), exc=batch_error, ) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 155bca656d5..20f1fa9d363 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1624,6 +1624,33 @@ async def test_update_daily_spend_keeps_failed_transactions_for_retry(): assert daily_spend_transactions == expected +@pytest.mark.asyncio +async def test_update_daily_spend_drops_the_batch_whose_failure_cannot_be_resent(): + """A reply lost after the statement was sent may already have applied, so the batch is + taken out of the caller's dict before the error propagates: whichever requeue the caller + runs afterwards, the Redis restore included, cannot send it a second time.""" + + def lose_the_reply() -> int: + raise httpx.ReadTimeout("no reply") + + prisma_client = _RecordingPrisma(execute_raw=lose_the_reply) + daily_spend_transactions = {"user-key": _daily_txn(user_id="user-1")} + proxy_logging_obj = MagicMock() + proxy_logging_obj.failure_handler = AsyncMock() + + with pytest.raises(httpx.ReadTimeout): + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=0, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + daily_spend_transactions=daily_spend_transactions, + entity_type="user", + entity_id_field="user_id", + ) + + assert daily_spend_transactions == {} + + @pytest.mark.asyncio async def test_commit_key_spend_updates_includes_last_active(): """ @@ -2875,6 +2902,33 @@ async def test_failed_daily_spend_commit_is_requeued_only_when_the_rows_are_prov assert db_writer.daily_spend_update_queue.update_queue.empty() +@pytest.mark.asyncio +async def test_failed_daily_spend_commit_drops_only_the_batch_that_was_sent(): + """A tick holding more than one batch of 100 rows sends them one statement at a time, and + a reply lost on one statement says nothing about the batches after it: only the batch that + was on the wire is dropped, the ones never sent go back on the queue and land next tick.""" + db_writer = DBSpendUpdateWriter() + await db_writer.daily_spend_update_queue.add_update( + {f"user-{i:03d}": _daily_txn(user_id=f"user-{i:03d}") for i in range(150)} + ) + db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=httpx.ReadTimeout("no reply")) + db_writer._flush_tool_discovery_queue = AsyncMock() + proxy_logging_obj = MagicMock() + proxy_logging_obj.failure_handler = AsyncMock() + + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj + ) + db.failing_table = None + await db_writer._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj + ) + + (upsert,) = _daily_upserts(db, "LiteLLM_DailyUserSpend") + assert _row_values(upsert, "user_id") == [f"user-{i:03d}" for i in range(100, 150)] + assert db_writer.daily_spend_update_queue.update_queue.empty() + + @pytest.mark.asyncio async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other_tables(): """With the Redis buffer off, a daily batch that failed to commit was discarded along From 83d89aa134bbeac391ce4a49dd63223f52b97019 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 19 Sep 2026 00:08:54 +0000 Subject: [PATCH 28/28] test(integration): cover off-peak pricing on a live proxy Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/_support/client.py | 4 +- tests/integration/contracts.json | 6 ++ .../pricing/test_off_peak_pricing.py | 82 +++++++++++++++++++ 3 files changed, 90 insertions(+), 2 deletions(-) create mode 100644 tests/integration/pricing/test_off_peak_pricing.py diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index 97522e5728c..9f1118ab1e3 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -161,7 +161,7 @@ class Scenario: assert all(object_value(object_value(entry)["model_info"])["id"] != identity for entry in entries) assert read_rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == [] - def model(self, **parameters: JsonValue) -> str: + def model(self, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue) -> str: name: Final = f"integration-{uuid.uuid4().hex}" created: Final = self.gateway.post( "/model/new", @@ -173,7 +173,7 @@ class Scenario: "api_base": f"{self.gateway.upstream_url}/v1", **parameters, }, - "model_info": {}, + "model_info": dict(model_info) if model_info is not None else {}, }, ) identity: Final = string_value(object_value(created["model_info"])["id"]) diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index 91b1bd86954..6958ade50f7 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -92,6 +92,12 @@ "tests/integration/pricing/test_price_precedence.py::test_same_upstream_aliases_keep_distinct_prices_after_reload": [ "quota_management.spend_tracking.alias_prices.remain_independent_on_reload" ], + "tests/integration/pricing/test_off_peak_pricing.py::test_open_off_peak_window_bills_off_peak_rates": [ + "quota_management.spend_tracking.off_peak_pricing.open_window_bills_off_peak_rates" + ], + "tests/integration/pricing/test_off_peak_pricing.py::test_closed_off_peak_window_bills_standard_rates": [ + "quota_management.spend_tracking.off_peak_pricing.closed_window_bills_standard_rates" + ], "tests/integration/spend/test_cache_and_quota.py::test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost": [ "quota_management.response_cache.generated_sequences_preserve_content_and_accounting" ], diff --git a/tests/integration/pricing/test_off_peak_pricing.py b/tests/integration/pricing/test_off_peak_pricing.py new file mode 100644 index 00000000000..5623356c078 --- /dev/null +++ b/tests/integration/pricing/test_off_peak_pricing.py @@ -0,0 +1,82 @@ +import json +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +from typing import Final + +import pytest +from pydantic import JsonValue + +from tests.integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from tests.integration._support.database import read_rows + +STANDARD_INPUT_RATE: Final = 0.001 +STANDARD_OUTPUT_RATE: Final = 0.002 +OFF_PEAK_INPUT_RATE: Final = 0.0001 +OFF_PEAK_OUTPUT_RATE: Final = 0.0002 + + +def off_peak_window(start_offset_hours: int, end_offset_hours: int) -> Mapping[str, JsonValue]: + now: Final = datetime.now(timezone.utc) + start: Final = now + timedelta(hours=start_offset_hours) + end: Final = now + timedelta(hours=end_offset_hours) + return { + "hours_utc": f"{start:%H:%M}-{end:%H:%M}", + "input_cost_per_token": OFF_PEAK_INPUT_RATE, + "output_cost_per_token": OFF_PEAK_OUTPUT_RATE, + } + + +def billed_model(scenario: Scenario, off_peak: Mapping[str, JsonValue]) -> str: + return scenario.model( + input_cost_per_token=STANDARD_INPUT_RATE, + output_cost_per_token=STANDARD_OUTPUT_RATE, + model_info={"off_peak_pricing": dict(off_peak)}, + ) + + +def assert_chat_bills_rates(gateway: Gateway, model: str, input_rate: float, output_rate: float) -> None: + response: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "off peak control"}]} + ) + assert response.status_code == 200, response.text + expected: Final = 20 * input_rate + 20 * output_rate + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected, rel=1e-6) + request_id: Final = string_value(object_value(response.json())["id"]) + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id = %s', + (request_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert rows[0]["prompt_tokens"] == 20 + assert rows[0]["completion_tokens"] == 20 + assert float(rows[0]["spend"]) == pytest.approx(expected, rel=1e-6) + metadata: Final = rows[0]["metadata"] + parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) + breakdown: Final = object_value(parsed["cost_breakdown"]) + assert float(breakdown["input_cost"]) == pytest.approx(20 * input_rate, rel=1e-6) + assert float(breakdown["output_cost"]) == pytest.approx(20 * output_rate, rel=1e-6) + + +@pytest.mark.covers("quota_management.spend_tracking.off_peak_pricing.open_window_bills_off_peak_rates") +def test_open_off_peak_window_bills_off_peak_rates(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = billed_model(scenario, off_peak_window(-1, 1)) + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list) + matching: Final = tuple(object_value(entry) for entry in entries if object_value(entry)["model_name"] == model) + assert len(matching) == 1 + info: Final = object_value(matching[0]["model_info"]) + off_peak: Final = object_value(info["off_peak_pricing"]) + assert off_peak["input_cost_per_token"] == OFF_PEAK_INPUT_RATE + assert off_peak["output_cost_per_token"] == OFF_PEAK_OUTPUT_RATE + assert_chat_bills_rates(gateway, model, OFF_PEAK_INPUT_RATE, OFF_PEAK_OUTPUT_RATE) + + +@pytest.mark.covers("quota_management.spend_tracking.off_peak_pricing.closed_window_bills_standard_rates") +def test_closed_off_peak_window_bills_standard_rates(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = billed_model(scenario, off_peak_window(2, 3)) + assert_chat_bills_rates(gateway, model, STANDARD_INPUT_RATE, STANDARD_OUTPUT_RATE)