fix(proxy): enforce key budgets on provider-prefixed batch and file routes

litellm/proxy/auth/auth_checks.py

@@ -1080,10 +1080,7 @@ def model_max_budget_of(holder: ModelMaxBudgetHolder) -> Mapping[str, object] |
 
 def route_skips_budget_checks(route: str) -> bool:
     return route not in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES and (
-        route in MODEL_DISCOVERY_ROUTES
-        or not (
-            RouteChecks.is_llm_api_route(route=route) or RouteChecks.is_provider_prefixed_batch_or_file_route(route)
-        )
+        route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_handler_route(route)
     )
 
 

litellm/proxy/auth/route_checks.py

@@ -113,8 +113,10 @@ _TOKEN_COUNTING_ACTION: Final = "countTokens"
 
 class RouteChecks:
     @staticmethod
-    def is_provider_prefixed_batch_or_file_route(route: str) -> bool:
-        return _PROVIDER_PREFIXED_BATCH_OR_FILE_ROUTE.fullmatch(route) is not None
+    def is_llm_handler_route(route: str) -> bool:
+        return RouteChecks.is_llm_api_route(route=route) or (
+            _PROVIDER_PREFIXED_BATCH_OR_FILE_ROUTE.fullmatch(route) is not None
+        )
 
     @staticmethod
     def check_route_enabled(route: str) -> None:

litellm/proxy/auth/user_api_key_auth.py

@@ -3144,7 +3144,7 @@ async def validate_resolved_virtual_key(  # noqa: C901  # Preserve ordering of e
                 )
 
                 # Check 5. Token Spend is under budget
-                if RouteChecks.is_llm_api_route(route=route):
+                if RouteChecks.is_llm_handler_route(route):
                     await virtual_key_max_budget_check(
                         valid_token=valid_token,
                         proxy_logging_obj=proxy_logging_obj,

litellm/proxy/common_utils/request_tag_mappings.py

@@ -20,9 +20,8 @@ def is_baggage_mapped_route(route: str) -> bool:
 
     return (
         (
-            RouteChecks.is_llm_api_route(route)
+            RouteChecks.is_llm_handler_route(route)
             or InitPassThroughEndpointHelpers.is_registered_pass_through_route(route)
-            or RouteChecks.is_provider_prefixed_batch_or_file_route(route)
         )
         and not RouteChecks.is_unbilled_route(route)
         and not RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.mcp_inference_routes.value)

tests/integration/spend/test_baggage_tag_mappings.py

@@ -625,6 +625,26 @@ def test_spent_mapped_tag_budget_rejects_a_provider_prefixed_batch_before_the_pr
         assert mapped.provider.drain() == ()
 
 
+@pytest.mark.parametrize("route", ["/litellm_proxy/v1/batches", "/hosted_vllm/v1/batches"])
+def test_spent_key_budget_rejects_a_provider_prefixed_batch_before_the_provider(mapped: Rig, route: str) -> None:
+    entry: Final = uuid.uuid4().hex
+    baggage: Final = f"entrypoint={entry},tenant={entry}"
+    with mapped.proxy.scenario() as scenario:
+        key: Final = scenario.key(allowed_routes=["/v1/chat/completions", route], max_budget=0.0001)
+
+        def batch() -> httpx.Response:
+            return mapped.proxy.request("POST", route, _batch_body(mapped), key=key, headers={"baggage": baggage})
+
+        first: Final = _chat(mapped, key, baggage)
+        assert first.status_code == 200, first.text
+        eventually(lambda: batch().status_code, lambda status: status == 422, seconds=70)
+        mapped.provider.drain()
+        blocked: Final = batch()
+        assert blocked.status_code == 422, blocked.text
+        assert "budget" in blocked.text.lower(), blocked.text
+        assert mapped.provider.drain() == ()
+
+
 def test_require_budgeted_tag_rejects_an_unbudgeted_mapped_tag_until_the_tag_has_a_budget(mapped: Rig) -> None:
     value: Final = uuid.uuid4().hex
     tag: Final = mapped.budgeted_tag(value)