diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index e439f6a5998..b805310bcb4 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2435,6 +2435,7 @@ async def _run_centralized_common_checks( await _reserve_budget_after_common_checks( user_api_key_auth_obj=user_api_key_auth_obj, + request=request, request_data=request_data, route=route, llm_router=llm_router, @@ -2458,6 +2459,7 @@ async def _noop_none() -> None: async def _reserve_budget_after_common_checks( user_api_key_auth_obj: UserAPIKeyAuth, + request: Optional[Request], request_data: dict, route: str, llm_router: Optional[Any], @@ -2474,6 +2476,8 @@ async def _reserve_budget_after_common_checks( user_api_key_auth_obj.budget_reservation = None if skip_budget_checks: return + if request is not None and getattr(request.state, "skip_budget_reservation", False): + return if general_settings.get("disable_budget_reservation") is True: verbose_proxy_logger.warning( "disable_budget_reservation is enabled: skipping optimistic budget " diff --git a/litellm/proxy/rust_control_plane/auth_endpoints.py b/litellm/proxy/rust_control_plane/auth_endpoints.py index 629e67060b1..b8aad7e7ecb 100644 --- a/litellm/proxy/rust_control_plane/auth_endpoints.py +++ b/litellm/proxy/rust_control_plane/auth_endpoints.py @@ -96,7 +96,12 @@ def _synthetic_request( "client": ("127.0.0.1", 0), "server": ("127.0.0.1", 4000), } - return Request(scope, receive) + request = Request(scope, receive) + # This endpoint verifies admission only. The realtime gateway records actual + # session spend through callback logs, so a pre-call optimistic reservation + # here would have no matching request lifecycle to reconcile. + request.state.skip_budget_reservation = True + return request router = APIRouter(prefix="/v1/rust_control_plane", tags=["rust control plane"]) diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 7219ab58799..26a9eae8bd1 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -103,6 +103,7 @@ async def test_should_clear_stale_budget_reservation_when_budget_checks_skip(): await _reserve_budget_after_common_checks( user_api_key_auth_obj=user_api_key_auth_obj, + request=None, request_data={"model": "free-model"}, route="/v1/chat/completions", llm_router=None, @@ -130,6 +131,7 @@ async def test_disable_budget_reservation_skips_reservation(): ) as mock_reserve: await _reserve_budget_after_common_checks( user_api_key_auth_obj=user_api_key_auth_obj, + request=None, request_data={"model": "gpt-4o"}, route="/v1/chat/completions", llm_router=None, @@ -146,6 +148,34 @@ async def test_disable_budget_reservation_skips_reservation(): assert user_api_key_auth_obj.budget_reservation is None +@pytest.mark.asyncio +async def test_request_scoped_skip_budget_reservation_skips_only_reservation(): + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + request = SimpleNamespace(state=SimpleNamespace(skip_budget_reservation=True)) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request", + new=AsyncMock(return_value={"reserved_cost": 0.5, "entries": []}), + ) as mock_reserve: + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request=request, + request_data={"model": "gpt-4o"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=False, + general_settings={}, + ) + + mock_reserve.assert_not_called() + assert user_api_key_auth_obj.budget_reservation is None + + @pytest.mark.asyncio async def test_budget_reservation_runs_when_not_disabled(): """Control for #27639: with the flag absent, the reservation still runs and is stored.""" @@ -161,6 +191,7 @@ async def test_budget_reservation_runs_when_not_disabled(): ) as mock_reserve: await _reserve_budget_after_common_checks( user_api_key_auth_obj=user_api_key_auth_obj, + request=None, request_data={"model": "gpt-4o"}, route="/v1/chat/completions", llm_router=None, diff --git a/tests/test_litellm/proxy/rust_control_plane/test_auth_endpoints.py b/tests/test_litellm/proxy/rust_control_plane/test_auth_endpoints.py index 597ea405306..1e91e0332a1 100644 --- a/tests/test_litellm/proxy/rust_control_plane/test_auth_endpoints.py +++ b/tests/test_litellm/proxy/rust_control_plane/test_auth_endpoints.py @@ -8,6 +8,7 @@ from litellm.proxy.rust_control_plane.auth_endpoints import ( DATA_PLANE_KEY_ENV_VAR, DATA_PLANE_KEY_HEADER, VerifyKeyRequest, + _synthetic_request, require_data_plane_key, router, verify_key, @@ -85,6 +86,19 @@ def test_router_mounts_auth_verify_under_rust_control_plane(): ) +@pytest.mark.asyncio +async def test_synthetic_request_skips_budget_reservation(): + request = _synthetic_request( + route="/v1/realtime", + authorization_header="Bearer sk-test-key", + model="gpt-realtime", + ) + + assert request.url.path == "/v1/realtime" + assert request.state.skip_budget_reservation is True + assert (await request.json()) == {"model": "gpt-realtime"} + + @pytest.mark.asyncio async def test_verify_key_returns_model_dump(monkeypatch): expected_auth = UserAPIKeyAuth(