From 832eb193ab5a3c43254ebc43dfbd3ef693b0dfa3 Mon Sep 17 00:00:00 2001 From: jacksonriding Date: Sat, 15 Aug 2026 12:01:33 +1000 Subject: [PATCH 1/2] feat(proxy): allow custom client-facing budget exceeded message --- litellm/__init__.py | 1 + litellm/constants.py | 1 + litellm/exceptions.py | 6 ++++ .../mcp_server/auth/user_api_key_auth_mcp.py | 2 +- litellm/proxy/auth/auth_exception_handler.py | 2 +- litellm/proxy/common_request_processing.py | 11 +++++- .../auto_router_endpoints.py | 2 +- .../proxy/auth/test_auth_exception_handler.py | 36 +++++++++++++++++++ .../proxy/test_common_request_processing.py | 22 ++++++++++++ 9 files changed, 79 insertions(+), 4 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 8961de940a0..802109eae07 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -394,6 +394,7 @@ budget_duration: Optional[str] = ( ) default_soft_budget: float = DEFAULT_SOFT_BUDGET # by default all litellm proxy keys have a soft budget of 50.0 budget_exceeded_throttle_percentage: Optional[float] = None +budget_exceeded_error_message: Optional[str] = None forward_traceparent_to_llm_provider: bool = False diff --git a/litellm/constants.py b/litellm/constants.py index 6449834d6a4..458b4cf119a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1570,6 +1570,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [ "cost_discount_config", "cost_margin_config", "budget_exceeded_throttle_percentage", + "budget_exceeded_error_message", # Every field editable from the Admin UI (proxy_server._GENERAL_SETTINGS_UI_LITELLM_FIELDS) # must be listed here so a DB write from one worker overrides the live litellm attribute on # the others when config reloads; otherwise peer workers stay on their startup value. diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 2eb4232fef9..d15a683290a 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -985,6 +985,12 @@ class BudgetExceededError(Exception): self.message = message super().__init__(message) + @property + def client_facing_message(self) -> str: + import litellm + + return litellm.budget_exceeded_error_message or self.message + ## DEPRECATED ## class InvalidRequestError(openai.BadRequestError): diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 95a3806e8ad..a58f48cab02 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1135,7 +1135,7 @@ class MCPRequestHandler: except (HTTPException, ProxyException): raise except litellm.BudgetExceededError as e: - raise HTTPException(status_code=getattr(e, "status_code", 429), detail=str(e)) from None + raise HTTPException(status_code=getattr(e, "status_code", 429), detail=e.client_facing_message) from None except Exception as e: # noqa: BLE001 # untyped gate failure: retryable 503 for a DB outage, else fail closed 401 MCPRequestHandler._raise_503_if_db_unavailable(e) raise HTTPException(status_code=401, detail="Invalid or expired credential") from None diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 603e72463bc..4031c17289e 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -145,7 +145,7 @@ class UserAPIKeyAuthExceptionHandler: if isinstance(e, litellm.BudgetExceededError): raise ProxyException( - message=e.message, + message=e.client_facing_message, type=ProxyErrorTypes.budget_exceeded, param=None, code=getattr(e, "status_code", status.HTTP_429_TOO_MANY_REQUESTS), diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 891915eb357..f6630c6a79f 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -40,7 +40,7 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import can_key_call_resolved_model from litellm.proxy.auth.auth_utils import check_response_size_is_safe from litellm.proxy.common_utils.callback_utils import ( @@ -2881,6 +2881,15 @@ class ProxyBaseLLMRequestProcessing: code=status.HTTP_400_BAD_REQUEST, headers=safe_headers, ) + if isinstance(e, litellm.BudgetExceededError): + raise ProxyException( + message=e.client_facing_message, + type=ProxyErrorTypes.budget_exceeded, + param=None, + code=e.status_code, + headers=safe_headers, + ) + # Extract status_code from the exception if it carries one. # Provider exceptions (NotFoundError, BadRequestError, GeminiError, # VertexAIError, etc.) all have a status_code attribute reflecting diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 4b2569fa9fa..8c5ba4c79f6 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -162,7 +162,7 @@ async def _authorize_models_this_test_can_call( ) except BudgetExceededError as e: raise ProxyException( - message=e.message, + message=e.client_facing_message, type=ProxyErrorTypes.budget_exceeded, param=None, code=status.HTTP_400_BAD_REQUEST, diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index 27798ec0bff..1d3a869a96f 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -30,6 +30,7 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path +import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler @@ -356,6 +357,41 @@ async def test_handle_authentication_error_budget_exceeded(): assert int(exc_info.value.code) == status.HTTP_429_TOO_MANY_REQUESTS +@pytest.mark.asyncio +async def test_handle_authentication_error_budget_exceeded_custom_message(monkeypatch): + from litellm.exceptions import BudgetExceededError + + monkeypatch.setattr( + litellm, + "budget_exceeded_error_message", + "Your AI usage allowance has been reached. Please contact the AI team.", + ) + + budget_error = BudgetExceededError( + message="Budget has been exceeded! Current cost: 10.51, Max budget: 10.00", + current_cost=10.51, + max_budget=10.00, + ) + + with pytest.raises(ProxyException) as exc_info: + await UserAPIKeyAuthExceptionHandler()._handle_authentication_error( + budget_error, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-test", + ) + + assert ( + exc_info.value.message + == "Your AI usage allowance has been reached. Please contact the AI team." + ) + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + assert int(exc_info.value.code) == status.HTTP_429_TOO_MANY_REQUESTS + assert "10.51" in budget_error.message + + @pytest.mark.asyncio async def test_route_passed_to_post_call_failure_hook(): """ diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 9ddd74a46a8..4c560623818 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -3190,6 +3190,28 @@ class TestHandleLLMApiExceptionRetryAfter: assert proxy_exc.headers["retry-after"] == "43" assert proxy_exc.headers["x-custom"] == "1" + async def test_handle_llm_api_exception_budget_exceeded_uses_custom_message( + self, monkeypatch + ): + monkeypatch.setattr( + litellm, "budget_exceeded_error_message", "Allowance reached, contact the AI team." + ) + exc = litellm.BudgetExceededError(current_cost=10.51, max_budget=10.0) + + proxy_exc = await self._invoke(exc) + + assert proxy_exc.message == "Allowance reached, contact the AI team." + assert proxy_exc.type == "budget_exceeded" + assert proxy_exc.code == "429" + + async def test_handle_llm_api_exception_budget_exceeded_defaults_to_detail(self): + exc = litellm.BudgetExceededError(current_cost=10.51, max_budget=10.0) + + proxy_exc = await self._invoke(exc) + + assert "10.51" in proxy_exc.message + assert proxy_exc.type == "budget_exceeded" + class TestHandleLLMApiExceptionFramingHeaders: """HTTP-framing headers on the provider exception must be stripped before the From 72aa5b0b1a92459584503910e75dec1c738fce83 Mon Sep 17 00:00:00 2001 From: jacksonriding Date: Sat, 15 Aug 2026 12:19:11 +1000 Subject: [PATCH 2/2] test(proxy): cover custom MCP budget message --- .../mcp_server/auth/test_user_api_key_auth_mcp.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 52dc91ce24d..9f5ff944b77 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -5671,15 +5671,17 @@ class TestMCPDcrBridgeDelegateAdmission: ) return exc_info.value - async def test_over_budget_admission_surfaces_429_not_401(self): + async def test_over_budget_admission_surfaces_429_not_401(self, monkeypatch): """A validly-authenticated but over-budget identity surfaces the standard pipeline's 429, not a misleading 401. Flattening budget to 401 told the caller their credential was invalid, which on a DCR client reads as broken auth and triggers a re-authorize that cannot fix a budget problem. Regression for the status-flattening finding on the live-policy gate.""" import litellm + monkeypatch.setattr(litellm, "budget_exceeded_error_message", "Allowance reached") mapped = await self._enforce_with_gate_error(litellm.BudgetExceededError(current_cost=10.0, max_budget=1.0)) assert mapped.status_code == 429 + assert mapped.detail == "Allowance reached" async def test_db_outage_during_policy_surfaces_503_not_401(self): """A transient database outage during the live-policy gate surfaces a retryable 503, not a 401