From 01c1426abf43057395dd92fcfbe783bfbdb2176f Mon Sep 17 00:00:00 2001 From: Shayan12456 <136985581+Shayan12456@users.noreply.github.com> Date: Thu, 29 Oct 2026 17:20:45 +0000 Subject: [PATCH] fix(proxy): enforce a user's max_parallel_requests and allow clearing it with null /user/new and /user/update store max_parallel_requests on a user, but nothing ever read it back out. UserAPIKeyAuth carried the user's tpm/rpm limits but had no slot for max_parallel_requests, so none of the three auth paths that build it (key auth, JWT auth, keyless MCP auth) could pass it to the rate limiter. The v3 limiter's user-scope descriptor builder never looked for it either, so the concurrency gauge was never created for that scope even though the same gauge mechanism already works for API keys. The legacy limiter (used behind LEGACY_MULTI_INSTANCE_RATE_LIMITING) had the identical gap, hardcoding sys.maxsize with a TODO instead of reading the field. Separately, /user/update with max_parallel_requests: null silently kept the old value, unlike tpm_limit/rpm_limit/max_budget/budget_duration, which null already clears. Adds the missing field, wires it through every path that already carries user_tpm_limit/user_rpm_limit, lets both limiters enforce it the same way they already enforce the API key's own limit, and adds max_parallel_requests to the set of fields a null can clear. --- litellm/proxy/auth/handle_jwt.py | 1 + litellm/proxy/auth/user_api_key_auth.py | 1 + .../proxy/hooks/parallel_request_limiter.py | 5 +- .../hooks/parallel_request_limiter_v3.py | 5 +- .../internal_user_endpoints.py | 2 +- .../mcp_server/auth/user_api_key_auth_mcp.py | 1 + litellm/types/proxy/auth/user_api_key_auth.py | 1 + .../test_user_api_key_auth_request_flow.py | 2 + .../hooks/test_parallel_request_limiter.py | 29 +++++++++ .../hooks/test_parallel_request_limiter_v3.py | 60 +++++++++++++++++++ .../test_internal_user_endpoints.py | 10 +++- 11 files changed, 111 insertions(+), 6 deletions(-) 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():