fix(budget): match prefixed key/end-user entries on bare requests and reject non-enforceable team caps

The key/end-user matcher normalized only the request side, so an entry keyed
provider/model (e.g. openai/gpt-4) never matched a bare model-group request:
the key cap was silently unenforced, and on team keys the request was charged
to, and blocked by, the shared team counter instead of the key's own cap. The
matcher now delegates to the same both-sides normalization the team path uses.

Team model_max_budget validation now rejects zero, negative, and non-finite
budget_limit values: enforcement skips non-positive caps, so a team admin
could null an imposed cap by lowering it to zero, which the update authority
gate counted as a permitted decrease.

Also brings the branch inside the tightened Final-annotation and lazy-logging
budgets.

litellm/_redis.py

@@ -395,9 +395,7 @@ def _get_redis_client_logic(**env_overrides):
     if _sentinel_password is not None:
         redis_kwargs["sentinel_password"] = _sentinel_password
 
-    _service_name: Final[str | None] = redis_kwargs.get("service_name", None) or get_secret(
-        "REDIS_SERVICE_NAME"
-    )
+    _service_name: Final[str | None] = redis_kwargs.get("service_name", None) or get_secret("REDIS_SERVICE_NAME")
 
     if _service_name is not None:
         redis_kwargs["service_name"] = _service_name

litellm/proxy/auth/auth_checks.py

@@ -4081,14 +4081,14 @@ async def _team_model_max_budget_check(
     """
     if team_object is None or team_object.team_id is None or model is None:
         return
-    team_model_max_budget = team_object.model_max_budget
+    team_model_max_budget: Final = team_object.model_max_budget
     if not isinstance(team_model_max_budget, dict) or len(team_model_max_budget) == 0:
         return
 
     from litellm.proxy.proxy_server import model_max_budget_limiter
 
-    models = (model,) if isinstance(model, str) else tuple(model)
-    key_model_max_budget = valid_token.model_max_budget if valid_token is not None else None
+    models: Final = (model,) if isinstance(model, str) else tuple(model)
+    key_model_max_budget: Final = valid_token.model_max_budget if valid_token is not None else None
     for model_name in models:
         await model_max_budget_limiter.is_team_within_model_budget(
             team_id=team_object.team_id,

litellm/proxy/auth/user_api_key_auth.py

@@ -2337,7 +2337,9 @@ async def _run_centralized_common_checks(
     if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None:
         user_api_key_auth_obj.org_id = team_object.organization_id
 
-    user_api_key_auth_obj.team_model_max_budget = team_object.model_max_budget if team_object is not None else None
+    user_api_key_auth_obj.team_model_max_budget = (  # rebind-ok: this hydration point mutates the caller's auth object by contract, like org_id above
+        team_object.model_max_budget if team_object is not None else None
+    )
 
     # common_checks identifies admin via user_object, not the token
     # (non_proxy_admin_allowed_routes_check). JWT admin shortcut and

litellm/proxy/hooks/model_max_budget_limiter.py

@@ -71,11 +71,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
                 model=model,
                 key_budget_config=_current_model_budget_info,
             )
-            if (
-                _current_spend is not None
-                and _current_model_budget_info.max_budget is not None
-                and _current_spend > _current_model_budget_info.max_budget
-            ):
+            if _current_spend is not None and _current_spend > _current_model_budget_info.max_budget:
                 raise litellm.BudgetExceededError(
                     message=f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, exceeded budget for model={model}",
                     current_cost=_current_spend,
@@ -137,11 +133,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
                 model=model,
                 key_budget_config=_current_model_budget_info,
             )
-            if (
-                _current_spend is not None
-                and _current_model_budget_info.max_budget is not None
-                and _current_spend > _current_model_budget_info.max_budget
-            ):
+            if _current_spend is not None and _current_spend > _current_model_budget_info.max_budget:
                 raise litellm.BudgetExceededError(
                     message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}",
                     current_cost=_current_spend,
@@ -170,7 +162,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
         """
         if model in internal_model_max_budget:
             return model
-        stripped_model = self._get_model_without_custom_llm_provider(model)
+        stripped_model: Final = self._get_model_without_custom_llm_provider(model)
         if stripped_model in internal_model_max_budget:
             return stripped_model
         return next(
@@ -195,7 +187,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
         """
         if not key_model_max_budget:
             return False
-        budget_config = self._get_request_model_budget_config(
+        budget_config: Final = self._get_request_model_budget_config(
             model=model, internal_model_max_budget=self._coerce_budget_configs(key_model_max_budget)
         )
         return (
@@ -226,22 +218,22 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
             verbose_proxy_logger.debug("Team model budget check skipped for model=%s: key has its own entry", model)
             return True
 
-        internal_model_max_budget = self._coerce_budget_configs(team_model_max_budget)
+        internal_model_max_budget: Final = self._coerce_budget_configs(team_model_max_budget)
 
         verbose_proxy_logger.debug("team internal_model_max_budget %s", internal_model_max_budget)
 
-        matched_model = self._get_matched_budget_model_name(
+        matched_model: Final = self._get_matched_budget_model_name(
             model=model, internal_model_max_budget=internal_model_max_budget
         )
         if matched_model is None:
             verbose_proxy_logger.debug("Model %s not found in team_model_max_budget", model)
             return True
-        budget_config = internal_model_max_budget[matched_model]
+        budget_config: Final = internal_model_max_budget[matched_model]
 
         if not budget_config.max_budget or budget_config.max_budget <= 0:
             return True
 
-        current_spend = await self._get_team_spend_for_model(
+        current_spend: Final = await self._get_team_spend_for_model(
             team_id=team_id,
             matched_model=matched_model,
             budget_config=budget_config,
@@ -276,7 +268,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
         write always address one counter regardless of how the request spelled the
         model.
         """
-        team_model_spend_cache_key = (
+        team_model_spend_cache_key: Final = (
             f"{TEAM_MODEL_SPEND_CACHE_KEY_PREFIX}:{team_id}:{matched_model}:{budget_config.budget_duration}"
         )
         return await self.dual_cache.async_get_cache(
@@ -304,20 +296,22 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
         if self._key_already_covers_model(key_model_max_budget, model):
             return
 
-        internal_model_max_budget = self._coerce_budget_configs(team_model_max_budget)
-        matched_model = self._get_matched_budget_model_name(
+        internal_model_max_budget: Final = self._coerce_budget_configs(team_model_max_budget)
+        matched_model: Final = self._get_matched_budget_model_name(
             model=model, internal_model_max_budget=internal_model_max_budget
         )
         if matched_model is None:
             return
-        budget_config = internal_model_max_budget[matched_model]
+        budget_config: Final = internal_model_max_budget[matched_model]
         if not budget_config.budget_duration:
             return
 
-        team_spend_key = (
+        team_spend_key: Final = (
             f"{TEAM_MODEL_SPEND_CACHE_KEY_PREFIX}:{team_id}:{matched_model}:{budget_config.budget_duration}"
         )
-        team_start_time_key = f"team_model_budget_start_time:{team_id}:{matched_model}:{budget_config.budget_duration}"
+        team_start_time_key: Final = (
+            f"team_model_budget_start_time:{team_id}:{matched_model}:{budget_config.budget_duration}"
+        )
         await self._increment_spend_for_key(
             budget_config=budget_config,
             spend_key=team_spend_key,
@@ -382,14 +376,14 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
         self, model: str, internal_model_max_budget: Mapping[str, BudgetConfig]
     ) -> BudgetConfig | None:
         """
-        Get the budget config for the request model
-
-        1. Check if `model` is in `internal_model_max_budget`
-        2. If not, check if `model` without custom llm provider is in `internal_model_max_budget`
+        Get the budget config for the request model via _get_matched_budget_model_name,
+        so request and config-entry names are both `{provider}/`-normalized and a
+        prefixed entry like `openai/gpt-4` binds a bare `gpt-4` request.
         """
-        return internal_model_max_budget.get(model, None) or internal_model_max_budget.get(
-            self._get_model_without_custom_llm_provider(model), None
+        matched_model: Final = self._get_matched_budget_model_name(
+            model=model, internal_model_max_budget=internal_model_max_budget
         )
+        return internal_model_max_budget[matched_model] if matched_model is not None else None
 
     @staticmethod
     def _get_model_without_custom_llm_provider(model: str) -> str:
@@ -427,7 +421,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
         user_api_key_end_user_model_max_budget: Final[dict | None] = _metadata.get(
             "user_api_key_end_user_model_max_budget", None
         )
-        user_api_key_team_model_max_budget: Mapping[str, Mapping[str, str | float]] | None = _metadata.get(
+        user_api_key_team_model_max_budget: Final[Mapping[str, Mapping[str, str | float]] | None] = _metadata.get(
             "user_api_key_team_model_max_budget", None
         )
         if (
@@ -503,7 +497,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
                     response_cost=response_cost,
                 )
 
-        user_api_key_team_id: str | None = _metadata.get("user_api_key_team_id", None)
+        user_api_key_team_id: Final[str | None] = _metadata.get("user_api_key_team_id", None)
         if (
             user_api_key_team_id is not None
             and user_api_key_team_model_max_budget is not None

litellm/proxy/litellm_pre_call_utils.py

@@ -1744,7 +1744,9 @@ async def add_litellm_data_to_request(
     data[_metadata_variable_name]["user_api_key_end_user_model_max_budget"] = (
         user_api_key_dict.end_user_model_max_budget
     )
-    data[_metadata_variable_name]["user_api_key_team_model_max_budget"] = user_api_key_dict.team_model_max_budget
+    data[_metadata_variable_name]["user_api_key_team_model_max_budget"] = (  # rebind-ok: pre-call metadata injection
+        user_api_key_dict.team_model_max_budget
+    )
 
     # User spend, budget - used by prometheus.py
     # Follow same pattern as team and API key budgets

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -425,7 +425,7 @@ def _team_key_model_max_budget_check(
         return
     if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
         return
-    team_member_object = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id)
+    team_member_object: Final = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id)
     if team_member_object is not None and team_member_object.role == "admin":
         return
     raise HTTPException(

litellm/proxy/management_endpoints/team_endpoints.py

@@ -1056,11 +1056,11 @@ def _check_team_model_max_budget_update_authority(
 
     from litellm.types.utils import BudgetConfig
 
-    existing_configs = MappingProxyType(
+    existing_configs: Final = MappingProxyType(
         {_model: BudgetConfig(**_info) for _model, _info in existing_model_max_budget.items()}
     )
-    requested_items = data.model_max_budget.items() if data.model_max_budget else ()
-    requested_configs = MappingProxyType({_model: BudgetConfig(**_info) for _model, _info in requested_items})
+    requested_items: Final = data.model_max_budget.items() if data.model_max_budget else ()
+    requested_configs: Final = MappingProxyType({_model: BudgetConfig(**_info) for _model, _info in requested_items})
     for model_name, existing_config in existing_configs.items():
         requested_config = requested_configs.get(model_name)
         if requested_config is None:
@@ -1111,11 +1111,11 @@ def _validate_team_model_max_budget(
             code="400",
         )
 
-    normalized_names = tuple(
+    normalized_names: Final = tuple(
         _PROXY_VirtualKeyModelMaxBudgetLimiter._get_model_without_custom_llm_provider(entry_name)
         for entry_name in model_max_budget
     )
-    colliding_entries = tuple(
+    colliding_entries: Final = tuple(
         sorted(
             entry_name
             for entry_name, normalized_name in zip(model_max_budget, normalized_names)
@@ -1134,6 +1134,32 @@ def _validate_team_model_max_budget(
             code="400",
         )
 
+    from types import MappingProxyType
+
+    from litellm.types.utils import BudgetConfig
+
+    entry_configs: Final = MappingProxyType(
+        {entry_name: BudgetConfig(**entry_info) for entry_name, entry_info in model_max_budget.items()}
+    )
+    non_enforceable: Final = tuple(
+        sorted(
+            entry_name
+            for entry_name, config in entry_configs.items()
+            if config.max_budget is None or not math.isfinite(config.max_budget) or config.max_budget <= 0
+        )
+    )
+    if non_enforceable:
+        raise ProxyException(
+            message=(
+                f"model_max_budget entries {non_enforceable} have no enforceable budget_limit; "
+                "enforcement skips non-positive caps, so a zero or negative value would leave the "
+                "model uncapped. Use a positive finite budget_limit, or remove the entry"
+            ),
+            type=ProxyErrorTypes.bad_request_error,
+            param="model_max_budget",
+            code="400",
+        )
+
 
 def _should_auto_add_team_creator(
     user_api_key_dict: UserAPIKeyAuth,

litellm/router.py

@@ -3410,9 +3410,7 @@ class Router:
 
         # Await the first task to complete successfully
         while pending_tasks:
-            done, pending_tasks = await asyncio.wait(
-                pending_tasks, return_when=asyncio.FIRST_COMPLETED
-            )
+            done, pending_tasks = await asyncio.wait(pending_tasks, return_when=asyncio.FIRST_COMPLETED)
             for completed_task in done:
                 result = await check_response(completed_task)
 
@@ -5240,9 +5238,7 @@ class Router:
                     # Update kwargs with the current model name or any other model-specific adjustments
                     ## SET CUSTOM PROVIDER TO SELECTED DEPLOYMENT ##
                     if not custom_llm_provider:
-                        _, custom_llm_provider, _, _ = get_llm_provider(
-                            model=model
-                        )
+                        _, custom_llm_provider, _, _ = get_llm_provider(model=model)
                     new_kwargs: Final = safe_deep_copy(kwargs)
                     self._update_kwargs_with_deployment(
                         deployment=cast(dict, model_name),
@@ -6029,9 +6025,7 @@ class Router:
                 raise Exception(
                     "'custom_llm_provider' must be set. Either via:\n `Router(assistants_config={'custom_llm_provider': ..})` \nor\n `router.arun_thread(custom_llm_provider=..)`"
                 )
-        return await original_function(
-            custom_llm_provider=custom_llm_provider, client=client, **kwargs
-        )
+        return await original_function(custom_llm_provider=custom_llm_provider, client=client, **kwargs)
 
     #### [END] ASSISTANTS API ####
 
@@ -6364,9 +6358,7 @@ class Router:
                 mask_sensitive_structure(fallback_model_group),
             )
             if len(fallback_failure_exception_str) > 0:
-                original_exception.message += (
-                    f"\nError doing the fallback: {fallback_failure_exception_str}"
-                )
+                original_exception.message += f"\nError doing the fallback: {fallback_failure_exception_str}"
 
         raise original_exception
 
@@ -9124,9 +9116,7 @@ class Router:
                     and model_info["supports_parallel_function_calling"] is True
                 ):
                     model_group_info.supports_parallel_function_calling = True
-                if (
-                    model_info.get("supports_vision", None) is not None and model_info["supports_vision"] is True
-                ):
+                if model_info.get("supports_vision", None) is not None and model_info["supports_vision"] is True:
                     model_group_info.supports_vision = True
                 if (
                     model_info.get("supports_function_calling", None) is not None
@@ -9144,9 +9134,7 @@ class Router:
                 ):
                     model_group_info.supports_url_context = True
 
-                if (
-                    model_info.get("supports_reasoning", None) is not None and model_info["supports_reasoning"] is True
-                ):
+                if model_info.get("supports_reasoning", None) is not None and model_info["supports_reasoning"] is True:
                     model_group_info.supports_reasoning = True
                 if (
                     model_info.get("supported_openai_params", None) is not None
@@ -10876,9 +10864,7 @@ class Router:
                         args=(e, traceback_exception),
                     ).start()  # log response
                     # Handle any exceptions that might occur during streaming
-                    asyncio.create_task(
-                        logging_obj.async_failure_handler(e, traceback_exception)
-                    )
+                    asyncio.create_task(logging_obj.async_failure_handler(e, traceback_exception))
             raise e
 
     async def async_get_available_deployment_for_pass_through(
@@ -11003,9 +10989,7 @@ class Router:
                         target=logging_obj.failure_handler,
                         args=(e, traceback_exception),
                     ).start()
-                    asyncio.create_task(
-                        logging_obj.async_failure_handler(e, traceback_exception)
-                    )
+                    asyncio.create_task(logging_obj.async_failure_handler(e, traceback_exception))
             raise e
 
     async def _run_routing_plugins(

litellm/utils.py

@@ -766,16 +766,8 @@ def function_setup(
 
         if (
             len(litellm.input_callback) > 0 or len(litellm.success_callback) > 0 or len(litellm.failure_callback) > 0
-        ) and len(
-            callback_list
-        ) == 0:
-            callback_list = list(
-                set(
-                    litellm.input_callback
-                    + litellm.success_callback
-                    + litellm.failure_callback
-                )
-            )
+        ) and len(callback_list) == 0:
+            callback_list = list(set(litellm.input_callback + litellm.success_callback + litellm.failure_callback))
             get_set_callbacks: Final = getattr(sys.modules[__name__], "get_set_callbacks")
             get_set_callbacks()(callback_list=callback_list, function_id=function_id)
         ## ASYNC CALLBACKS - safety net for callbacks added via direct append

tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py

@@ -61,6 +61,61 @@ def test_get_request_model_budget_config(budget_limiter):
     assert config is None
 
 
+def test_get_request_model_budget_config_strips_entry_prefix(budget_limiter):
+    """Key/end-user entries keyed provider/model must bind bare (and cross-prefixed)
+    requests; only the request side was normalized before, so an openai/gpt-4
+    entry never matched a gpt-4 request and the cap was inert (Greptile finding,
+    key-side sibling of the team matcher fix)."""
+    internal_budget = {"openai/gpt-4": GenericBudgetInfo(budget_limit=100.0, time_period="1d")}
+    matched = budget_limiter._get_request_model_budget_config(
+        model="gpt-4", internal_model_max_budget=internal_budget
+    )
+    assert matched is not None and matched.max_budget == 100.0
+    cross = budget_limiter._get_request_model_budget_config(
+        model="azure/gpt-4", internal_model_max_budget=internal_budget
+    )
+    assert cross is not None and cross.max_budget == 100.0
+    assert (
+        budget_limiter._get_request_model_budget_config(
+            model="claude-3", internal_model_max_budget=internal_budget
+        )
+        is None
+    )
+
+
+@pytest.mark.asyncio
+async def test_is_key_within_model_budget_prefixed_entry_bare_request(budget_limiter):
+    """A key cap keyed openai/gpt-4 must enforce on a bare gpt-4 request."""
+    user_api_key = UserAPIKeyAuth(
+        token="test-key",
+        key_alias="test-alias",
+        model_max_budget={"openai/gpt-4": {"budget_limit": 100.0, "time_period": "1d"}},
+    )
+    with patch.object(budget_limiter, "_get_virtual_key_spend_for_model", return_value=150.0):
+        with pytest.raises(litellm.BudgetExceededError):
+            await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4")
+
+
+@pytest.mark.asyncio
+async def test_team_check_skipped_for_prefixed_key_entry_on_bare_request(budget_limiter):
+    """A key whose own entry is keyed openai/gpt-4 is exempt from the team cap for
+    a bare gpt-4 request; before the fix it charged, and was blocked by, the
+    shared team counter."""
+    with patch.object(
+        budget_limiter, "_get_team_spend_for_model", AsyncMock(return_value=10_000.0)
+    ) as mock_team_spend:
+        assert (
+            await budget_limiter.is_team_within_model_budget(
+                team_id="team-1",
+                team_model_max_budget={"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}},
+                model="gpt-4",
+                key_model_max_budget={"openai/gpt-4": {"budget_limit": 50.0, "time_period": "1d"}},
+            )
+            is True
+        )
+    mock_team_spend.assert_not_called()
+
+
 # Test is_key_within_model_budget
 @pytest.mark.asyncio
 async def test_is_key_within_model_budget(budget_limiter):

tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py

@@ -11108,3 +11108,16 @@ class TestValidateTeamModelMaxBudget:
         with pytest.raises(ProxyException) as exc_info:
             self._validate({"gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}})
         assert exc_info.value.code == "400"
+
+    @pytest.mark.parametrize("bad_limit", [0.0, -5.0, float("inf")])
+    def test_non_positive_cap_rejected(self, bad_limit):
+        """Enforcement skips non-positive caps, so accepting one silently uncaps the
+        model; a team admin could exploit that by lowering an imposed cap to zero
+        (Veria finding). Validation now rejects the value outright."""
+        from litellm.proxy._types import ProxyException
+
+        with pytest.raises(ProxyException) as exc_info:
+            self._validate({"gpt-4": {"budget_limit": bad_limit, "time_period": "1d"}})
+        assert exc_info.value.code == "400"
+        assert "gpt-4" in exc_info.value.message
+        assert "budget_limit" in exc_info.value.message