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.
This commit is contained in:
Shayan12456 2026-10-29 17:20:45 +00:00
parent c98641fd5a
commit 01c1426abf
11 changed files with 111 additions and 6 deletions

View file

@ -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"],

View file

@ -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),

View file

@ -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,

View file

@ -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,
},
)

View file

@ -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":

View file

@ -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.

View file

@ -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

View file

@ -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"

View file

@ -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",
[

View file

@ -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():

View file

@ -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():