diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 9d800a1fd63..d94f045de25 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -3056,6 +3056,7 @@ class JWTAuthManager: agent_id=result.get("agent_id"), user_tpm_limit=user.tpm_limit if user is not None and not admin else None, user_rpm_limit=user.rpm_limit if user is not None and not admin else None, + user_max_parallel_requests=user.max_parallel_requests if user is not None and not admin else None, user_model_max_budget=user.model_max_budget if user is not None and not admin else None, **team_grants( team_object=result["team_object"], diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 292a147d200..8774895ef31 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -4464,6 +4464,7 @@ async def _return_user_api_key_auth_obj( user_api_key_kwargs.update( user_tpm_limit=user_obj.tpm_limit, user_rpm_limit=user_obj.rpm_limit, + user_max_parallel_requests=user_obj.max_parallel_requests, user_email=user_obj.user_email, user_spend=getattr(user_obj, "spend", None), user_max_budget=getattr(user_obj, "max_budget", None), diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index 05c64317a44..87e16965fb9 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -422,10 +422,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): if user_id is not None: user_tpm_limit = user_api_key_dict.user_tpm_limit user_rpm_limit = user_api_key_dict.user_rpm_limit + user_max_parallel_requests = user_api_key_dict.user_max_parallel_requests if user_tpm_limit is None: user_tpm_limit = sys.maxsize if user_rpm_limit is None: user_rpm_limit = sys.maxsize + if user_max_parallel_requests is None: + user_max_parallel_requests = sys.maxsize request_count_api_key = f"{user_id}::{precise_minute}::request_count" # print(f"Checking if {request_count_api_key} is allowed to make request for minute {precise_minute}") @@ -434,7 +437,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): cache=cache, data=data, call_type=call_type, - max_parallel_requests=sys.maxsize, # TODO: Support max parallel requests for a user + max_parallel_requests=user_max_parallel_requests, current=cache_objects["request_count_user_id"], request_count_api_key=request_count_api_key, tpm_limit=user_tpm_limit, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index a3a6f582dce..1c176f9de24 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -3435,7 +3435,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # User rate limits if user_api_key_dict.user_id and ( - user_api_key_dict.user_rpm_limit is not None or user_api_key_dict.user_tpm_limit is not None + user_api_key_dict.user_rpm_limit is not None + or user_api_key_dict.user_tpm_limit is not None + or user_api_key_dict.user_max_parallel_requests is not None ): descriptors.append( RateLimitDescriptor( @@ -3444,6 +3446,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): rate_limit={ "requests_per_unit": user_api_key_dict.user_rpm_limit, "tokens_per_unit": user_api_key_dict.user_tpm_limit, + "max_parallel_requests": user_api_key_dict.user_max_parallel_requests, "window_size": self.window_size, }, ) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index bb901f22649..02621a5f2cd 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1286,7 +1286,7 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda fields_set: Final = data.fields_set() if hasattr(data, "fields_set") else set() for k, v in data_json.items(): - if k in ("max_budget", "budget_duration", "tpm_limit", "rpm_limit"): + if k in ("max_budget", "budget_duration", "tpm_limit", "rpm_limit", "max_parallel_requests"): if k in fields_set: non_default_values[k] = v elif k == "model_max_budget": diff --git a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py index cc7080738d0..fb450f6356d 100644 --- a/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/mcp_server/auth/user_api_key_auth_mcp.py @@ -1308,6 +1308,7 @@ class MCPRequestHandler: # per-KEY limits do not apply, there being no key.) user_tpm_limit=user_object.tpm_limit, user_rpm_limit=user_object.rpm_limit, + user_max_parallel_requests=user_object.max_parallel_requests, ) # Server-only marker, set AFTER construction: the before-validator strips it from any validated # input, so caller-supplied data (key metadata, JWT claims) can never forge it. diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index 2c1c75fbe04..9dde52636cd 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -139,6 +139,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob tpm_limit_per_model: dict[str, int] | None = None user_tpm_limit: int | None = None user_rpm_limit: int | None = None + user_max_parallel_requests: int | None = None user_email: str | None = None user_spend: float | None = None user_max_budget: float | None = None 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 b7f3ded9291..c74fd76bd1f 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 @@ -1924,6 +1924,7 @@ async def test_return_user_api_key_auth_obj_user_spend_and_budget(): { "tpm_limit": 1000, "rpm_limit": 100, + "max_parallel_requests": 5, "user_email": "test@example.com", "spend": 250.0, "max_budget": 1000.0, @@ -1961,6 +1962,7 @@ async def test_return_user_api_key_auth_obj_user_spend_and_budget(): assert result.user_max_budget == 1000.0 assert result.user_tpm_limit == 1000 assert result.user_rpm_limit == 100 + assert result.user_max_parallel_requests == 5 assert result.user_email == "test@example.com" diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter.py b/tests/unit/proxy/hooks/test_parallel_request_limiter.py index 08840978169..5a34e8a59ad 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter.py @@ -97,6 +97,35 @@ async def test_pre_call_hook_counts_a_cli_session_under_the_per_user_alias_not_t assert counted == {"current_requests": 1, "current_tpm": 0, "current_rpm": 1}, counted +@pytest.mark.asyncio +async def test_user_max_parallel_requests_is_enforced(): + """ + A user's own `max_parallel_requests` must be enforced on the legacy + (LEGACY_MULTI_INSTANCE_RATE_LIMITING) limiter too, the same as it already + is for an API key: a second request from the same user, under a + different API key, in the same minute bucket, is refused. + """ + handler: Final = PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), clock=_frozen_clock + ) + user_id: Final = "pia" + first_key: Final = UserAPIKeyAuth( + api_key="sk-pia-key-one", user_id=user_id, user_max_parallel_requests=1 + ) + second_key: Final = UserAPIKeyAuth( + api_key="sk-pia-key-two", user_id=user_id, user_max_parallel_requests=1 + ) + + await handler.async_pre_call_hook( + user_api_key_dict=first_key, cache=DualCache(), data={"model": "gpt-4o-mini"}, call_type="completion" + ) + + with pytest.raises(ProxyRateLimitError): + await handler.async_pre_call_hook( + user_api_key_dict=second_key, cache=DualCache(), data={"model": "gpt-4o-mini"}, call_type="completion" + ) + + @pytest.mark.parametrize( "response_obj", [ diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index a38f1a54944..fc797adf049 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -1029,6 +1029,47 @@ async def test_rejected_request_does_not_consume_parallel_slot_v3(): ) +@pytest.mark.asyncio +async def test_user_max_parallel_requests_is_enforced_v3(): + """ + A user's own `max_parallel_requests` (set via /user/new or /user/update, + independent of any key-level limit) must be enforced: a second concurrent + request from the same user, admitted under a different API key, is + refused with 429 once the user's slot is full. + """ + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_id = "pia" + first_key = hash_token("sk-pia-key-one") + second_key = hash_token("sk-pia-key-two") + user_api_key_dict = UserAPIKeyAuth( + api_key=first_key, user_id=user_id, user_max_parallel_requests=1 + ) + other_key_same_user = UserAPIKeyAuth( + api_key=second_key, user_id=user_id, user_max_parallel_requests=1 + ) + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + acquisition = get_request_stash().parallel_slot + assert acquisition is not None + assert acquisition["counter_keys"] == [f"{{user:{user_id}}}:max_parallel_requests"] + + with pytest.raises(HTTPException): + await handler.async_pre_call_hook( + user_api_key_dict=other_key_same_user, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + + @pytest.mark.asyncio async def test_parallel_gauge_uses_atomic_redis_script_v3(): """ @@ -1242,6 +1283,25 @@ async def test_should_rate_limit_only_called_when_limits_exist_v3(): should_rate_limit_called ), "should_rate_limit should be called when max parallel requests are configured" + # Test 7: User-level max parallel requests configured - should_rate_limit SHOULD be called + should_rate_limit_called = False + user_api_key_dict_with_user_parallel_limits = UserAPIKeyAuth( + api_key=_api_key, + user_id="test_user", + user_max_parallel_requests=5, # User-level max parallel requests configured + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict_with_user_parallel_limits, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + + assert ( + should_rate_limit_called + ), "should_rate_limit should be called when user-level max parallel requests are configured" + @pytest.mark.asyncio async def test_model_specific_rate_limits_only_called_when_configured_v3(): diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index d1bbf444d2c..961c798032d 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2255,16 +2255,20 @@ def test_update_internal_user_params_ignores_other_nones(): assert non_default_values["max_budget"] == 100.0 -@pytest.mark.parametrize("field", ["tpm_limit", "rpm_limit"], ids=["tpm_limit", "rpm_limit"]) +@pytest.mark.parametrize( + "field", + ["tpm_limit", "rpm_limit", "max_parallel_requests"], + ids=["tpm_limit", "rpm_limit", "max_parallel_requests"], +) def test_update_internal_user_params_explicit_null_clears_rate_limit_but_omitted_is_untouched( field: str, ) -> None: data: Final = UpdateUserRequest(user_id="limit-clear", **{field: None}) result: Final = _update_internal_user_params(data_json=data.model_dump(exclude_unset=True), data=data) - other_field: Final = "rpm_limit" if field == "tpm_limit" else "tpm_limit" + other_fields: Final = {"tpm_limit", "rpm_limit", "max_parallel_requests"} - {field} assert result[field] is None - assert other_field not in result + assert other_fields.isdisjoint(result) def test_update_internal_user_params_keeps_original_max_budget_when_not_provided():