fix(proxy): run budget checks on provider-prefixed batch and file routes

litellm/proxy/auth/auth_checks.py

@@ -1080,7 +1080,10 @@ 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)
+        route in MODEL_DISCOVERY_ROUTES
+        or not (
+            RouteChecks.is_llm_api_route(route=route) or RouteChecks.is_provider_prefixed_batch_or_file_route(route)
+        )
     )
 
 

litellm/proxy/auth/route_checks.py

@@ -24,6 +24,8 @@ from .auth_checks_organization import (  # noqa: F401  # legacy module exports
     user_is_org_admin,
 )
 
+_PROVIDER_PREFIXED_BATCH_OR_FILE_ROUTE: Final = re.compile(r"/[^/]+/v1/(?:batches|files)(?:/.*)?")
+
 # Management write routes denied to PROXY_ADMIN_VIEW_ONLY. Adding a new write
 # endpoint to a management router REQUIRES adding it here too — the surrounding
 # check falls through to "allow" if the route is not matched, which previously
@@ -110,6 +112,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
+
     @staticmethod
     def check_route_enabled(route: str) -> None:
         try:

litellm/proxy/common_utils/request_tag_mappings.py

@@ -1,4 +1,3 @@
-import re
 from collections.abc import Mapping, Sequence
 from typing import Final
 from urllib.parse import unquote_plus
@@ -14,7 +13,6 @@ from litellm.proxy.auth.route_checks import RouteChecks
 from litellm.types.proxy.request_tag_mappings import RequestTagMapping
 
 REQUEST_TAG_MAPPINGS_ADAPTER: Final = TypeAdapter(tuple[RequestTagMapping, ...])
-_PROVIDER_PREFIXED_BATCH_OR_FILE_ROUTE: Final = re.compile(r"/[^/]+/v1/(?:batches|files)(?:/.*)?")
 
 
 def is_baggage_mapped_route(route: str) -> bool:
@@ -24,7 +22,7 @@ def is_baggage_mapped_route(route: str) -> bool:
         (
             RouteChecks.is_llm_api_route(route)
             or InitPassThroughEndpointHelpers.is_registered_pass_through_route(route)
-            or _PROVIDER_PREFIXED_BATCH_OR_FILE_ROUTE.fullmatch(route) is not None
+            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

@@ -602,6 +602,29 @@ def test_provider_prefixed_batch_spend_row_carries_the_mapped_tags(mapped: Rig)
         assert _spend_tags(mapped.proxy, key) == {mapped.exact_tag, mapped.tenant_tag(tenant)}
 
 
+@pytest.mark.parametrize("route", ["/litellm_proxy/v1/batches", "/hosted_vllm/v1/batches"])
+def test_spent_mapped_tag_budget_rejects_a_provider_prefixed_batch_before_the_provider(mapped: Rig, route: str) -> None:
+    entry: Final = uuid.uuid4().hex
+    tag: Final = mapped.usecase_tag(entry)
+    baggage: Final = f"entrypoint={entry},tenant={entry}"
+    with mapped.proxy.scenario() as scenario:
+        key: Final = scenario.key(allowed_routes=["/v1/chat/completions", route])
+        mapped.proxy.post("/tag/new", {"name": tag, "max_budget": 0.0001})
+        scenario.cleanups.callback(lambda: mapped.proxy.request("POST", "/tag/delete", {"name": tag}))
+
+        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)

tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py

@@ -11201,6 +11201,18 @@ def test_route_skips_budget_checks_marks_only_spend_free_routes() -> None:
     assert route_skips_budget_checks(route="/v1/chat/completions") is False
 
 
+@pytest.mark.parametrize(
+    "route", ["/litellm_proxy/v1/batches", "/hosted_vllm/v1/batches/batch_qa/cancel", "/xai/v1/files", "/manus/v1/files/f"]
+)
+def test_provider_prefixed_batch_and_file_routes_keep_budget_checks(route: str) -> None:
+    assert route_skips_budget_checks(route=route) is False
+
+
+@pytest.mark.parametrize("route", ["/hosted_vllm/nested/v1/batches", "/hosted_vllm/v1/batches_extra"])
+def test_routes_outside_the_batch_and_file_handlers_still_skip_budget_checks(route: str) -> None:
+    assert route_skips_budget_checks(route=route) is True
+
+
 def test_request_skips_budget_checks_extends_route_rule_with_zero_cost_models() -> None:
     assert request_skips_budget_checks(route="/v1/models", model=None, llm_router=None) is True
     assert request_skips_budget_checks(route="/v1/chat/completions", model=None, llm_router=None) is False