diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 56703d8a971..22bb93a1c49 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -18,7 +18,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast from fastapi import HTTPException, Request, status -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm @@ -133,7 +133,7 @@ from .auth_checks_organization import ( add_team_org_context_to_request_body, organization_role_based_access_check, ) -from .auth_utils import get_model_from_request, get_request_route_template +from .auth_utils import get_model_from_request, get_request_route_template, route_in_additonal_public_routes if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -863,22 +863,52 @@ def route_skips_budget_checks(route: str) -> bool: _AUTHN_FLAGS: Final = ("enable_jwt_auth", "enable_oauth2_auth", "enable_oauth2_proxy_auth") +class _PassThroughEndpointAuth(BaseModel): + """The two fields of a ``pass_through_endpoints`` entry that decide whether auth runs on it.""" + + path: str = "" + auth: bool | str | None = None + + +_PASS_THROUGH_ENDPOINTS_ADAPTER: Final = TypeAdapter(tuple[_PassThroughEndpointAuth, ...]) + + +def _is_unauthenticated_pass_through(route: str, general_settings: Mapping[str, object]) -> bool: + configured: Final = general_settings.get("pass_through_endpoints") + if configured is None: + return False + try: + endpoints: Final = _PASS_THROUGH_ENDPOINTS_ADAPTER.validate_python(configured) + except ValidationError: + return False + return any(endpoint.path == route and endpoint.auth is not True for endpoint in endpoints) + + def auth_skips_common_checks( - general_settings: Mapping[str, object], master_key: str | None, custom_auth_configured: bool + route: str, general_settings: Mapping[str, object], master_key: str | None, custom_auth_configured: bool ) -> bool: """ - Whether ``user_api_key_auth`` runs no ``common_checks`` at all for this deployment. + Whether ``user_api_key_auth`` runs no ``common_checks`` at all for this request. - That is the case in no-auth dev mode (no master key and no JWT or OAuth2 - auth configured, so the proxy is unauthenticated by configuration) and behind - a custom auth hook that did not opt in with ``custom_auth_run_common_checks``. + That is the case on a public route, on a user-configured pass-through endpoint + that did not ask for auth, in no-auth dev mode (no master key and no JWT or + OAuth2 auth configured, so the proxy is unauthenticated by configuration) and + behind a custom auth hook that did not opt in with ``custom_auth_run_common_checks``. Post-auth checks that mirror ``common_checks`` skip themselves on the same terms. """ + public_route: Final = route in LiteLLMRoutes.public_routes.value or route_in_additonal_public_routes( + current_route=route + ) no_auth_mode: Final = master_key is None and not any(general_settings.get(flag, False) for flag in _AUTHN_FLAGS) custom_auth_opted_out: Final = custom_auth_configured and not general_settings.get( "custom_auth_run_common_checks", False ) - return no_auth_mode or custom_auth_opted_out + return ( + public_route + or _is_unauthenticated_pass_through(route=route, general_settings=general_settings) + or no_auth_mode + or custom_auth_opted_out + ) async def common_checks( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 095f61c666b..27994042c3d 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2521,27 +2521,11 @@ async def _run_centralized_common_checks( user_custom_auth, ) - # Public routes (e.g. /health/liveness) are exempt from - # auth in the builder — the wrapper must not retroactively apply - # authz on top, or k8s readiness probes and other unauthenticated - # callers get 401. - if route in LiteLLMRoutes.public_routes.value or route_in_additonal_public_routes(current_route=route): - return - - # User-configured pass-through endpoints with ``auth: false`` are - # explicitly unauthenticated — the builder returns an empty - # UserAPIKeyAuth() and the request is forwarded as-is. Running - # common_checks on the empty token would reject the request as - # admin-only. The "auth" flag on the endpoint config is the - # contract; honor it. - pass_through_endpoints: Final = general_settings.get("pass_through_endpoints", None) - if pass_through_endpoints is not None: - for endpoint in pass_through_endpoints: - if isinstance(endpoint, dict) and endpoint.get("path", "") == route and endpoint.get("auth") is not True: - return - if auth_skips_common_checks( - general_settings=general_settings, master_key=master_key, custom_auth_configured=user_custom_auth is not None + route=route, + general_settings=general_settings, + master_key=master_key, + custom_auth_configured=user_custom_auth is not None, ): return diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 481b7d87382..44163ee9e53 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -695,7 +695,10 @@ async def _enforce_tag_budgets_for_added_tags( from litellm.proxy.proxy_server import master_key, prisma_client, user_api_key_cache, user_custom_auth if auth_skips_common_checks( - general_settings=general_settings, master_key=master_key, custom_auth_configured=user_custom_auth is not None + route=route, + general_settings=general_settings, + master_key=master_key, + custom_auth_configured=user_custom_auth is not None, ): return () await tag_max_budget_check_for_tags( diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 40e934c197f..11bd770842c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -360,10 +360,9 @@ async def _reserve_counters( fail_closed_budget_enforcement=fail_closed_budget_enforcement, ) continue - except Exception: - await _release_applied_entries_best_effort( - entries=applied_entries, - default_reserved_cost=reservation_cost, + except BaseException: + await asyncio.shield( + _release_applied_entries_best_effort(entries=applied_entries, default_reserved_cost=reservation_cost) ) raise diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 09809ee6316..d24b83a69e1 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2482,27 +2482,60 @@ def test_route_skips_budget_checks_matches_auth_scope(route, expected): @pytest.mark.parametrize( - ("general_settings", "master_key", "custom_auth_configured", "expected"), + ("route", "general_settings", "master_key", "custom_auth_configured", "expected"), [ - ({}, None, False, True), - ({"enable_jwt_auth": True}, None, False, False), - ({"enable_oauth2_auth": True}, None, False, False), - ({"enable_oauth2_proxy_auth": True}, None, False, False), - ({}, "sk-master", False, False), - ({}, "sk-master", True, True), - ({"custom_auth_run_common_checks": True}, "sk-master", True, False), - ({"custom_auth_run_common_checks": False}, "sk-master", True, True), + ("/v1/chat/completions", {}, None, False, True), + ("/v1/chat/completions", {"enable_jwt_auth": True}, None, False, False), + ("/v1/chat/completions", {"enable_oauth2_auth": True}, None, False, False), + ("/v1/chat/completions", {"enable_oauth2_proxy_auth": True}, None, False, False), + ("/v1/chat/completions", {}, "sk-master", False, False), + ("/v1/chat/completions", {}, "sk-master", True, True), + ("/v1/chat/completions", {"custom_auth_run_common_checks": True}, "sk-master", True, False), + ("/v1/chat/completions", {"custom_auth_run_common_checks": False}, "sk-master", True, True), + ("/health/liveliness", {}, "sk-master", False, True), + ("/v1/chat/completions", {"public_routes": ["/v1/chat/completions"]}, "sk-master", False, True), + ("/v1/chat/completions", {"public_routes": ["/v1/embeddings"]}, "sk-master", False, False), + ("/bria", {"pass_through_endpoints": [{"path": "/bria", "target": "https://x"}]}, "sk-master", False, True), + ( + "/bria", + {"pass_through_endpoints": [{"path": "/bria", "target": "https://x", "auth": False}]}, + "sk-master", + False, + True, + ), + ( + "/bria", + {"pass_through_endpoints": [{"path": "/bria", "target": "https://x", "auth": True}]}, + "sk-master", + False, + False, + ), + ( + "/other", + {"pass_through_endpoints": [{"path": "/bria", "target": "https://x", "auth": False}]}, + "sk-master", + False, + False, + ), ], ) -def test_auth_skips_common_checks_names_the_deployments_that_never_run_them( - general_settings, master_key, custom_auth_configured, expected +def test_auth_skips_common_checks_names_the_requests_that_never_run_them( + monkeypatch, route, general_settings, master_key, custom_auth_configured, expected ): - """No-auth dev mode and a custom auth hook without the opt-in run no common_checks, so no budget checks.""" + """Public routes, pass-through endpoints without auth, no-auth dev mode and a custom auth hook without + the opt-in run no common_checks, so no budget checks.""" + from litellm.proxy import proxy_server from litellm.proxy.auth.auth_checks import auth_skips_common_checks + monkeypatch.setattr(proxy_server, "general_settings", general_settings) + monkeypatch.setattr(proxy_server, "premium_user", True) + assert ( auth_skips_common_checks( - general_settings=general_settings, master_key=master_key, custom_auth_configured=custom_auth_configured + route=route, + general_settings=general_settings, + master_key=master_key, + custom_auth_configured=custom_auth_configured, ) is expected ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index ecea188e864..e0070318a96 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -1,7 +1,9 @@ from __future__ import annotations +import asyncio import json import math +from collections.abc import Mapping from types import MappingProxyType, SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -139,6 +141,7 @@ async def test_repeated_token_counting_never_touches_a_tiny_budget( HOOK_TAG: Final = "hook-added-tag" +SECOND_HOOK_TAG: Final = "second-hook-added-tag" BODY_TAG: Final = "body-tag" CHAT_BODY: Final[dict[str, object]] = { "model": "gpt-4o", @@ -176,11 +179,14 @@ def _budgeted_tag_prisma(tag_names: tuple[str, ...], max_budget: float) -> Magic async def _reserve_added_tags( - route: str, prisma: MagicMock, tags: tuple[str, ...] = (HOOK_TAG,) + route: str, + prisma: MagicMock, + tags: tuple[str, ...] = (HOOK_TAG,), + request_body: Mapping[str, object] = CHAT_BODY, ) -> dict[str, object] | None: return await reserve_budget_for_added_tags( tags=tags, - request_body=dict(CHAT_BODY), + request_body=dict(request_body), route=route, llm_router=None, valid_token=UserAPIKeyAuth(token="hashed-hook-tag-key", max_budget=100.0, spend=0.0), @@ -231,11 +237,63 @@ async def test_reserve_budget_for_added_tags_skips_routes_auth_never_reserves(sp assert spend_counter_cache.in_memory_cache.get_cache(key=f"spend:tag:{HOOK_TAG}") is None +class _ParkingIncrementCache(DualCache): + """A spend-counter cache whose increment of ``parked_key`` never returns, so a test can cancel mid-reservation.""" + + def __init__(self, parked_key: str) -> None: + super().__init__() + self.parked_key: Final = parked_key + self.parked: Final = asyncio.Event() + + async def async_increment_cache(self, key: str, value: float, **kwargs: object) -> float | None: + if key == self.parked_key: + self.parked.set() + await asyncio.Event().wait() + return await super().async_increment_cache(key=key, value=value, **kwargs) + + +@pytest.mark.asyncio +async def test_reserve_budget_for_added_tags_releases_the_reserved_tag_when_cancelled_mid_acquisition( + monkeypatch: pytest.MonkeyPatch, +): + """A client disconnect under SSE keepalives cancels the request while the second tag is being reserved; + the first tag's counter must not stay charged for a request that never reached the provider.""" + cache: Final = _ParkingIncrementCache(parked_key=f"spend:tag:{SECOND_HOOK_TAG}") + monkeypatch.setattr(proxy_server, "spend_counter_cache", cache) + monkeypatch.setattr(proxy_server, "prisma_client", None) + prisma: Final = _budgeted_tag_prisma((HOOK_TAG, SECOND_HOOK_TAG), max_budget=1.0) + + reserving: Final = asyncio.create_task( + _reserve_added_tags("/v1/chat/completions", prisma, tags=(HOOK_TAG, SECOND_HOOK_TAG)) + ) + await cache.parked.wait() + assert cache.in_memory_cache.get_cache(key=f"spend:tag:{HOOK_TAG}") > 0 + reserving.cancel() + with pytest.raises(asyncio.CancelledError): + await reserving + + assert cache.in_memory_cache.get_cache(key=f"spend:tag:{HOOK_TAG}") == pytest.approx(0.0) + + @pytest.mark.asyncio async def test_reserve_budget_for_added_tags_ignores_tags_without_a_budget(spend_counter_cache: DualCache): assert await _reserve_added_tags("/v1/chat/completions", _budgeted_tag_prisma((), max_budget=1.0)) is None +@pytest.mark.asyncio +async def test_reserve_budget_for_added_tags_skips_a_request_with_no_model_to_price(spend_counter_cache: DualCache): + """Without a model there is no estimate to reserve, same as the auth-time reservation.""" + body: Final = MappingProxyType({key: value for key, value in CHAT_BODY.items() if key != "model"}) + + assert ( + await _reserve_added_tags( + "/v1/chat/completions", _budgeted_tag_prisma((HOOK_TAG,), max_budget=1.0), request_body=body + ) + is None + ) + assert spend_counter_cache.in_memory_cache.get_cache(key=f"spend:tag:{HOOK_TAG}") is None + + BEDROCK_SONNET: Final = "us.anthropic.claude-sonnet-4-6" CONVERSE_BODY: Final = { "messages": [{"role": "user", "content": [{"text": "Reply with one word: pong"}]}], diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 193d5fa6a4d..77a03087f97 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -621,13 +621,15 @@ class TestProxyBaseLLMRequestProcessing: ("sk-master", object(), {}, False), ("sk-master", object(), {"custom_auth_run_common_checks": True}, True), ("sk-master", None, {}, True), + ("sk-master", None, {"public_routes": ["/v1/chat/completions"]}, False), + ("sk-master", None, {"public_routes": ["/v1/embeddings"]}, True), ], ) async def test_common_processing_pre_call_logic_enforces_hook_added_tags_only_where_auth_runs_common_checks( self, monkeypatch, master_key, user_custom_auth, general_settings, checked ): - """A deployment whose auth wrapper skips common_checks (no-auth dev mode, custom auth without opt-in) - never budget-checked tags before, so a hook-added tag must not start 429ing it.""" + """A request whose auth wrapper skips common_checks (a public route, no-auth dev mode, custom auth + without opt-in) never budget-checked tags before, so a hook-added tag must not start 429ing it.""" processing_obj = ProxyBaseLLMRequestProcessing(data={}) async def mock_pre_call_hook(user_api_key_dict, data, call_type): @@ -639,6 +641,8 @@ class TestProxyBaseLLMRequestProcessing: ) monkeypatch.setattr("litellm.proxy.proxy_server.master_key", master_key) monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_auth", user_custom_auth) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) await processing_obj.common_processing_pre_call_logic( request=mock_request,