diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index fa331b37950..c7b698b7f18 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1081,6 +1081,15 @@ BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: Final = frozenset( } ) +VIDEO_RETRIEVAL_ROUTES: Final = frozenset( + { + "/videos/{video_id}", + "/v1/videos/{video_id}", + "/videos/{video_id}/content", + "/v1/videos/{video_id}/content", + } +) + MCP_DISCOVERY_ROUTES: Final = frozenset({"/mcp-rest/tools/list", "/v1/mcp/tools"}) MCP_ZERO_SPEND_JSONRPC_METHODS: Final = frozenset({"initialize", "notifications/initialized", "ping", "tools/list"}) @@ -1097,6 +1106,10 @@ def is_mcp_discovery_request(route: str, request_body: Mapping[str, object]) -> return isinstance(method, str) and method in MCP_ZERO_SPEND_JSONRPC_METHODS +def is_video_retrieval_request(route: str, http_method: str | None) -> bool: + return http_method == "GET" and RouteChecks.check_route_access(route=route, allowed_routes=VIDEO_RETRIEVAL_ROUTES) + + class ModelMaxBudgetHolder(Protocol): @property def model_max_budget(self) -> Mapping[str, object] | None: ... @@ -1126,8 +1139,9 @@ def request_skips_budget_checks( llm_router: Router | None, *, request_body: Mapping[str, object] | None = None, + http_method: str | None = None, ) -> bool: - if route_skips_budget_checks(route=route): + if route_skips_budget_checks(route=route) or is_video_retrieval_request(route=route, http_method=http_method): return True if request_body is not None and is_mcp_discovery_request(route=route, request_body=request_body): return True @@ -1189,6 +1203,7 @@ async def common_checks( model=_model, llm_router=llm_router, request_body=_typed_request_body(request_body), + http_method=RouteChecks.get_request_method(request), ) ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 1b2e4e791e3..ff6e7f84bed 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -85,6 +85,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo is_model_cost_zero, is_user_proxy_admin, is_valid_fallback_model, + is_video_retrieval_request, jwt_key_mapping_cache_key, key_access_group_grants_model, key_model_aliases_for_auth_check, @@ -2680,7 +2681,11 @@ async def user_api_key_auth_builder( team_id=valid_token.team_id, ) skip_budget_checks = request_skips_budget_checks( - route=route, model=model, llm_router=llm_router, request_body=request_data + route=route, + model=model, + llm_router=llm_router, + request_body=request_data, + http_method=RouteChecks.get_request_method(request), ) # Fetch project object for JWT path if project_id is set @@ -3198,6 +3203,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e model=await budget_served_model(model, request_data, request, llm_router, valid_token), llm_router=llm_router, request_body=request_data, + http_method=RouteChecks.get_request_method(request), ) # Check 3. Check if user is in their team budget @@ -4299,6 +4305,7 @@ async def _should_skip_budget_checks( model=(model if sub_call else await budget_served_model(model, request_data, request, llm_router, valid_token)), llm_router=llm_router, request_body=request_data, + http_method=RouteChecks.get_request_method(request), ) @@ -4976,7 +4983,9 @@ async def _run_post_custom_auth_checks( # spend some other model accrued. The JWT and virtual-key paths already skip # every budget check for these; this path did not, so the same request could # be refused under custom auth and served under the other two. - skip_budget_checks: Final = ( + skip_budget_checks: Final = is_video_retrieval_request( + route=route, http_method=RouteChecks.get_request_method(request) + ) or ( is_model_cost_zero( model=await budget_served_model(current_model, request_data, request, llm_router, valid_token), llm_router=llm_router, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index e585cd7b1ff..3efdbfd30e2 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -111,6 +111,7 @@ from litellm.proxy.auth.auth_checks import ( tag_max_budget_check_for_tags, ) from litellm.proxy.auth.auth_utils import check_response_size_is_safe, get_request_route +from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.bug_report_config import build_proxy_bug_report from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, @@ -1048,11 +1049,14 @@ async def _enforce_guardrail_added_tag_budgets( llm_router: Router | None, user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, + http_method: str | None = None, ) -> None: added_tags: Final = tuple( tag for tag in get_tags_from_request_body(request_body=data) if tag not in tags_before_guardrails ) - if not added_tags or request_skips_budget_checks(route=route, model=_request_model(data), llm_router=llm_router): + if not added_tags or request_skips_budget_checks( + route=route, model=_request_model(data), llm_router=llm_router, http_method=http_method + ): return from litellm.proxy.proxy_server import prisma_client, user_api_key_cache @@ -2864,6 +2868,7 @@ class ProxyBaseLLMRequestProcessing: llm_router=llm_router, user_api_key_dict=user_api_key_dict, proxy_logging_obj=proxy_logging_obj, + http_method=RouteChecks.get_request_method(request), ) if route_type == "aget_responses": attach_post_call_pipelines_to_retrieval( 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 e7bcc696c0d..4cbe5320aa1 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 @@ -7286,6 +7286,45 @@ async def test_model_discovery_route_bypasses_user_budget(): assert result is True +@pytest.mark.parametrize( + "method, route, expect_blocked", + [ + ("GET", "/v1/videos/video_abc", False), + ("GET", "/videos/video_abc/content", False), + ("POST", "/v1/videos", True), + ("POST", "/v1/videos/video_abc/remix", True), + ("GET", "/v1/videos", True), + ], +) +@pytest.mark.asyncio +async def test_exhausted_team_budget_spares_only_video_retrieval(method: str, route: str, expect_blocked: bool) -> None: + from litellm.proxy.auth.auth_checks import common_checks + + request: Final = MagicMock() + request.method = method + + async def run() -> bool: + return await common_checks( + request_body={}, + team_object=LiteLLM_TeamTable(team_id="test-team", spend=0.4, max_budget=0.3), + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route=route, + llm_router=None, + proxy_logging_obj=AsyncMock(), + valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"), + request=request, + ) + + if expect_blocked: + with pytest.raises(litellm.BudgetExceededError): + await run() + else: + assert await run() is True + + @pytest.mark.parametrize("route", ["/health/services", "/auto_router/test_routing"]) @pytest.mark.asyncio async def test_side_effectful_info_route_still_enforces_budget(route: str) -> None: @@ -11202,6 +11241,34 @@ def test_request_skips_budget_checks_for_non_spending_requests( ) +@pytest.mark.parametrize( + "route, http_method, expected", + [ + ("/v1/videos/video_abc", "GET", True), + ("/videos/video_abc", "GET", True), + ("/v1/videos/video_abc/content", "GET", True), + ("/videos/video_abc/content", "GET", True), + ("/v1/videos/video_abc", "POST", False), + ("/v1/videos/video_abc/content", None, False), + ("/v1/videos", "GET", False), + ("/v1/videos", "POST", False), + ("/v1/videos/video_abc/remix", "POST", False), + ("/v1/videos/edits", "POST", False), + ("/v1/videos/video_abc/remix", "GET", False), + ("/v1/chat/completions", "GET", False), + ], +) +def test_request_skips_budget_checks_only_for_paid_video_retrieval( + route: str, http_method: str | None, expected: bool, free_and_paid_router: litellm.Router +) -> None: + assert ( + request_skips_budget_checks( + route=route, model="paid-model", llm_router=free_and_paid_router, http_method=http_method + ) + is expected + ) + + def _agent_model_ceiling_resolver( models: frozenset[str] | None, ) -> tuple[CeilingResolver, list[str]]: diff --git a/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py b/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py index e75a77325c6..2a5ee88dbf9 100644 --- a/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py @@ -1,3 +1,4 @@ +from datetime import datetime, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -6,15 +7,24 @@ from starlette.requests import Request from starlette.routing import Route import litellm +from litellm.caching.dual_cache import DualCache from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_EndUserTable, + Litellm_EntityType, + ProxyErrorTypes, + ProxyException, UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import ( _run_post_custom_auth_checks, update_valid_token_with_end_user_params, ) +from litellm.proxy.hooks.model_max_budget_limiter import ( + PROXY_VirtualKeyModelMaxBudgetLimiter, + model_budget_spend_cache_key, +) +from litellm.types.videos.utils import encode_video_id_with_provider @pytest.mark.asyncio @@ -513,3 +523,86 @@ async def test_custom_auth_judges_the_zero_cost_waiver_by_the_key_alias_target( result: Final = await run() assert result.token == "test_token" limiter.is_end_user_within_model_budget.assert_not_awaited() + + +_VIDEO_ID: Final = encode_video_id_with_provider(video_id="video_raw", provider="openai", model_id="sora-2") + + +async def _post_custom_auth_video_request( + valid_token: UserAPIKeyAuth, method: str, template: str, request_data: dict[str, object] +) -> UserAPIKeyAuth: + limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + await limiter.dual_cache.async_set_cache( + key=model_budget_spend_cache_key( + entity_type=Litellm_EntityType.KEY, entity_id=valid_token.token, budget_model="sora-2", budget_duration="1d" + ), + value=0.4, + ) + path: Final = template.format(video_id=_VIDEO_ID) + http_request: Final = Request( + { + "type": "http", + "method": method, + "path": path, + "headers": [], + "query_string": b"", + "route": Route(template, endpoint=lambda request: None), + } + ) + with ( + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.model_max_budget_limiter", limiter), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + return await _run_post_custom_auth_checks( + valid_token=valid_token, + request=http_request, + request_data=request_data, + route=path, + parent_otel_span=None, + ) + + +def _exhausted_model_budget_token() -> UserAPIKeyAuth: + return UserAPIKeyAuth(token="test_token", model_max_budget={"sora-2": {"max_budget": 0.3, "budget_duration": "1d"}}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "template", + ["/v1/videos/{video_id}", "/videos/{video_id}", "/v1/videos/{video_id}/content", "/videos/{video_id}/content"], +) +async def test_custom_auth_exhausted_model_budget_still_retrieves_an_already_paid_video(template: str) -> None: + result: Final = await _post_custom_auth_video_request( + _exhausted_model_budget_token(), "GET", template, {"video_id": _VIDEO_ID} + ) + + assert result.token == "test_token" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "method, template, request_data", + [ + ("POST", "/v1/videos", {"model": "sora-2"}), + ("POST", "/v1/videos/{video_id}/remix", {"video_id": _VIDEO_ID}), + ("GET", "/v1/videos", {"model": "sora-2"}), + ], +) +async def test_custom_auth_exhausted_model_budget_still_blocks_video_requests_other_than_retrieval( + method: str, template: str, request_data: dict[str, object] +) -> None: + with pytest.raises(litellm.BudgetExceededError): + await _post_custom_auth_video_request(_exhausted_model_budget_token(), method, template, request_data) + + +@pytest.mark.asyncio +async def test_custom_auth_video_retrieval_still_rejects_an_expired_token() -> None: + expired_token: Final = _exhausted_model_budget_token().model_copy( + update={"expires": datetime(2000, 1, 1, tzinfo=timezone.utc)} + ) + + with pytest.raises(ProxyException) as exc_info: + await _post_custom_auth_video_request(expired_token, "GET", "/v1/videos/{video_id}", {"video_id": _VIDEO_ID}) + + assert exc_info.value.type == ProxyErrorTypes.expired_key diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index 7746fc96673..d39c8cb8c85 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -2360,6 +2360,115 @@ async def test_jwt_path_skips_user_model_budget_only_for_zero_spend_mcp_methods( assert user_model_budget_check.await_count == (1 if expect_user_model_budget_check else 0) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "method, path, expect_rejected", + [ + ("GET", "/v1/videos/{video_id}", False), + ("GET", "/videos/{video_id}/content", False), + ("POST", "/v1/videos", True), + ("POST", "/v1/videos/{video_id}/remix", True), + ], +) +async def test_jwt_path_spares_only_video_retrieval_from_an_exhausted_user_model_budget( + monkeypatch: pytest.MonkeyPatch, method: str, path: str, expect_rejected: bool +) -> None: + from litellm.caching.dual_cache import DualCache + from litellm.proxy._types import JWTAuthBuilderResult, LiteLLM_JWTAuth, Litellm_EntityType, ProxyErrorTypes + from litellm.proxy.auth import user_api_key_auth as user_api_key_auth_module + from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_builder + from litellm.proxy.hooks.model_max_budget_limiter import ( + PROXY_VirtualKeyModelMaxBudgetLimiter, + model_budget_spend_cache_key, + ) + from litellm.types.videos.utils import encode_video_id_with_provider + + video_id: Final = encode_video_id_with_provider(video_id="video_raw", provider="openai", model_id="sora-2") + jwt_result: Final = JWTAuthBuilderResult( + is_proxy_admin=False, + team_object=None, + user_object=LiteLLM_UserTable( + user_id="jwt-user", + model_max_budget={"sora-2": {"budget_limit": 0.3, "time_period": "1d"}}, + ), + end_user_object=None, + org_object=None, + token="fake.jwt.token", + team_id=None, + user_id="jwt-user", + user_email=None, + end_user_id=None, + org_id=None, + team_membership=None, + jwt_claims={}, + agent_id=None, + ) + + class _VerifiedJWTAuthManager(JWTAuthManager): + @staticmethod + async def auth_builder( + api_key: object, + jwt_handler: object, + request_data: object, + general_settings: object, + route: object, + prisma_client: object, + user_api_key_cache: object, + parent_otel_span: object, + proxy_logging_obj: object, + request_headers: object = None, + request_method: object = None, + ) -> JWTAuthBuilderResult: + return jwt_result + + limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + await limiter.dual_cache.async_set_cache( + key=model_budget_spend_cache_key( + entity_type=Litellm_EntityType.USER, entity_id="jwt-user", budget_model="sora-2", budget_duration="1d" + ), + value=0.4, + ) + jwt_handler: Final = JWTHandler() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth(), + ) + mock_request: Final = MagicMock() + mock_request.url.path = path.format(video_id=video_id) + mock_request.method = method + mock_request.headers = {"authorization": "Bearer fake.jwt.token"} + mock_request.query_params = {} + request_data: Final[dict[str, object]] = {"model": "sora-2"} if path == "/v1/videos" else {"video_id": video_id} + + monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"enable_jwt_auth": True}) + monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) + monkeypatch.setattr(litellm.proxy.proxy_server, "jwt_handler", jwt_handler) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", None) + monkeypatch.setattr(litellm.proxy.proxy_server, "model_max_budget_limiter", limiter) + monkeypatch.setattr(user_api_key_auth_module, "JWTAuthManager", _VerifiedJWTAuthManager) + + async def authenticate() -> str | None: + result: Final = await user_api_key_auth_builder( + request=mock_request, + api_key="Bearer fake.jwt.token", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data=request_data, + ) + return result.user_id + + if not expect_rejected: + assert await authenticate() == "jwt-user" + return + with pytest.raises(ProxyException) as exc_info: + await authenticate() + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + + def _qa_websocket(messages): async def receive(): return {"type": "websocket.connect"} diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 5824b5fccab..67a76906dc8 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -13,8 +13,8 @@ from functools import partial from io import BytesIO from pathlib import Path from textwrap import dedent -from types import SimpleNamespace -from typing import Final +from types import MappingProxyType, SimpleNamespace +from typing import Final, Literal from unittest.mock import ANY, AsyncMock, MagicMock, patch @@ -98,6 +98,7 @@ from litellm.proxy.auth.user_api_key_auth import ( ) from litellm.proxy.pass_through_endpoints import pass_through_endpoints as pass_through from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata +from litellm.types.videos.utils import encode_video_id_with_provider from tests.unit.proxy.db.fake_prisma_engine import engine_call from tests._master_key import MASTER_KEY @@ -13956,6 +13957,143 @@ async def test_exhausted_key_budget_spares_only_zero_spend_mcp_discovery( assert result.spend == 5.0 +_VIDEO_ID: Final = encode_video_id_with_provider(video_id="video_raw", provider="openai", model_id="sora-2") +_VIDEO_RETRIEVAL_ROUTES: Final = ( + f"/v1/videos/{_VIDEO_ID}", + f"/videos/{_VIDEO_ID}", + f"/v1/videos/{_VIDEO_ID}/content", + f"/videos/{_VIDEO_ID}/content", +) +_EXPIRED_BEFORE_ANY_TEST_RUN: Final = datetime(2000, 1, 1, tzinfo=timezone.utc) +_VIDEO_SPENDING_REQUESTS: Final = ( + ("POST", "/v1/videos", MappingProxyType({"model": "sora-2"})), + ("POST", f"/v1/videos/{_VIDEO_ID}/remix", MappingProxyType({"video_id": _VIDEO_ID})), + ("POST", "/v1/videos/edits", MappingProxyType({"model": "sora-2"})), + ("POST", f"/v1/videos/{_VIDEO_ID}", MappingProxyType({"video_id": _VIDEO_ID})), + ("GET", "/v1/videos", MappingProxyType({"model": "sora-2"})), +) + + +async def _authenticate_video_request( + key: UserAPIKeyAuth, api_key: str, method: str, route: str, request_data: dict[str, object] +) -> UserAPIKeyAuth: + user_api_key_cache: Final = DualCache() + await cache_key_object( + hashed_token=key.token or "", + user_api_key_obj=key, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=None, + ) + model_budget_limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + await model_budget_limiter.dual_cache.async_set_cache( + key=model_budget_spend_cache_key( + entity_type=Litellm_EntityType.KEY, entity_id=key.token, budget_model="sora-2", budget_duration="1d" + ), + value=0.4, + ) + request: Final = Request( + { + "type": "http", + "method": method, + "path": route, + "headers": [(b"authorization", f"Bearer {api_key}".encode())], + "query_string": b"", + } + ) + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + with ( + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.llm_model_list", []), + patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), + patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj), + patch("litellm.proxy.proxy_server.model_max_budget_limiter", model_budget_limiter), + patch( + "litellm.proxy.auth.user_api_key_auth.virtual_key_max_budget_alert_check", + new_callable=AsyncMock, + ), + ): + resolved: Final = await user_api_key_auth_builder( + request=request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data=request_data, + ) + await _authorize_authenticated_request(resolved, request, request_data, route, api_key) + return resolved + + +def _exhausted_video_key(budget: Literal["max_budget", "model_max_budget"]) -> tuple[UserAPIKeyAuth, str]: + from litellm.proxy.utils import hash_token + + api_key: Final = f"sk-exhausted-video-{budget}" + if budget == "max_budget": + return UserAPIKeyAuth(token=hash_token(api_key), max_budget=0.3, spend=0.4), api_key + return ( + UserAPIKeyAuth( + token=hash_token(api_key), model_max_budget={"sora-2": {"max_budget": 0.3, "budget_duration": "1d"}} + ), + api_key, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("budget", ["max_budget", "model_max_budget"]) +@pytest.mark.parametrize("route", _VIDEO_RETRIEVAL_ROUTES) +async def test_exhausted_key_budget_still_retrieves_an_already_paid_video( + budget: Literal["max_budget", "model_max_budget"], route: str +) -> None: + key, api_key = _exhausted_video_key(budget) + + result: Final = await _authenticate_video_request(key, api_key, "GET", route, {"video_id": _VIDEO_ID}) + + assert result.token == key.token + + +@pytest.mark.asyncio +@pytest.mark.parametrize("budget", ["max_budget", "model_max_budget"]) +@pytest.mark.parametrize("method, route, request_data", _VIDEO_SPENDING_REQUESTS) +async def test_exhausted_key_budget_still_blocks_video_requests_other_than_retrieval( + budget: Literal["max_budget", "model_max_budget"], method: str, route: str, request_data: Mapping[str, object] +) -> None: + key, api_key = _exhausted_video_key(budget) + + with pytest.raises(ProxyException) as exc_info: + await _authenticate_video_request(key, api_key, method, route, dict(request_data)) + + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "overrides, expected_type", + [ + ({"expires": _EXPIRED_BEFORE_ANY_TEST_RUN}, ProxyErrorTypes.expired_key), + ({"blocked": True}, ProxyErrorTypes.auth_error), + ({"models": ["other-model"]}, ProxyErrorTypes.key_model_access_denied), + ], +) +@pytest.mark.parametrize("route", (_VIDEO_RETRIEVAL_ROUTES[0], _VIDEO_RETRIEVAL_ROUTES[2])) +async def test_video_retrieval_on_exhausted_key_keeps_key_access_controls( + overrides: dict[str, object], expected_type: ProxyErrorTypes, route: str +) -> None: + exhausted_key, api_key = _exhausted_video_key("max_budget") + + with pytest.raises(ProxyException) as exc_info: + await _authenticate_video_request( + exhausted_key.model_copy(update=overrides), api_key, "GET", route, {"video_id": _VIDEO_ID} + ) + + assert exc_info.value.type == expected_type + + @pytest.mark.asyncio @pytest.mark.parametrize( "requested_model, key_aliases, team_model_aliases, expect_blocked", diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index ccfac3d4e30..bbc77b8645b 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -1020,6 +1020,49 @@ class TestProxyBaseLLMRequestProcessing: ) tag_budget_check.assert_not_awaited() + @pytest.mark.asyncio + @pytest.mark.parametrize( + "http_method, route, route_type, expect_rejected", + [ + ("GET", "/v1/videos/video_abc", "avideo_status", False), + ("GET", "/videos/video_abc/content", "avideo_content", False), + ("POST", "/v1/videos", "avideo_generation", True), + ("POST", "/v1/videos/video_abc/remix", "avideo_remix", True), + ("GET", "/v1/videos", "avideo_list", True), + ], + ) + async def test_common_processing_pre_call_logic_spares_only_video_retrieval_from_exhausted_guardrail_tags( + self, monkeypatch: pytest.MonkeyPatch, http_method: str, route: str, route_type: str, expect_rejected: bool + ) -> None: + processing_obj, mock_request, mock_proxy_logging_obj, mock_proxy_config, tag_budget_check = ( + self._guardrail_tag_budget_harness( + monkeypatch, + request_body={"model": "paid-model", "metadata": {"tags": ["existing-tag"]}}, + guardrail_tags=["guardrail-tag"], + route=route, + ) + ) + tag_budget_check.side_effect = litellm.BudgetExceededError(current_cost=0.4, max_budget=0.3) + mock_request.method = http_method + + async def pre_call() -> None: + await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=mock_proxy_config, + route_type=route_type, + llm_router=self._router_with_free_and_paid_models(), + ) + + if not expect_rejected: + await pre_call() + return + with pytest.raises(ProxyException) as exc_info: + await pre_call() + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + @pytest.mark.asyncio async def test_enforce_guardrail_added_tag_budgets_raises_budget_exceeded_proxy_exception(self, monkeypatch): from litellm.proxy.common_request_processing import _enforce_guardrail_added_tag_budgets