diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b8d9530d2eb..e202c163aff 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/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) + ) ) diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 125bfab4edf..5c60c31f03e 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/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: diff --git a/litellm/proxy/common_utils/request_tag_mappings.py b/litellm/proxy/common_utils/request_tag_mappings.py index 3673b8916d9..0e7b0401f4e 100644 --- a/litellm/proxy/common_utils/request_tag_mappings.py +++ b/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) diff --git a/tests/integration/spend/test_baggage_tag_mappings.py b/tests/integration/spend/test_baggage_tag_mappings.py index 4511df58842..c497b28a02b 100644 --- a/tests/integration/spend/test_baggage_tag_mappings.py +++ b/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) diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 0bdc5f1910a..93f42a195cb 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/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