From ec0af8d5f8e6a94d82f4e3831acb0746c1b208af Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:02:32 -0700 Subject: [PATCH] fix(proxy): evict cached user on every proxy for tpm/rpm updates and edit limits in the users UI (#44130) * test(proxy): cover cross-worker cache eviction for user tpm/rpm limit updates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): broadcast user cache eviction when tpm_limit or rpm_limit changes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(proxy): return tpm_limit and rpm_limit from /v2/user/info Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(ui): edit user tpm and rpm limits from the user edit form Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): clear user tpm_limit and rpm_limit when sent as null Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover user tpm/rpm limit updates across proxies Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover user rate-limit routes and bulk updates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): reject unsafe rate limit integers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(ui): cover user rate limit seed and saved state Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): type user endpoint test doubles Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * revert: drop rate limit upper bound Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(ui): validate user rate limits with zod and clear user edit lint warnings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(ui): rename user edit schema and name its input and output types Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): drop redundant event loop yields in rate limit eviction test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mrinal Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 2 + .../internal_user_endpoints.py | 10 +- .../test_user_rate_limit_updates.py | 381 ++++++++++++++++++ .../test_internal_user_endpoints.py | 262 ++++++++++-- .../users/_components/user_edit_view.test.tsx | 162 +++++++- .../users/_components/user_edit_view.tsx | 99 ++++- .../view_users/user_info_view.test.tsx | 38 +- .../_components/view_users/user_info_view.tsx | 4 + .../src/components/networking.tsx | 2 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 + 10 files changed, 915 insertions(+), 49 deletions(-) create mode 100644 tests/integration/management/test_user_rate_limit_updates.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d02d1f29ded..a314fdbf060 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3645,6 +3645,8 @@ class UserInfoV2Response(LiteLLMPydanticObjectBase): user_role: str | None = None spend: float = 0.0 max_budget: float | None = None + tpm_limit: int | None = None + rpm_limit: int | None = None models: list[str] = [] budget_duration: str | None = None budget_reset_at: datetime | None = None diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index e04d92a5398..1953370be39 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -119,7 +119,7 @@ if TYPE_CHECKING: router: Final = APIRouter() _USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig]) _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE: Final = 50 -_USER_BUDGET_CACHE_FIELDS: Final = frozenset({"max_budget", "model_max_budget"}) +_USER_LIMIT_CACHE_FIELDS: Final = frozenset({"max_budget", "model_max_budget", "tpm_limit", "rpm_limit"}) def _user_table( @@ -1158,6 +1158,8 @@ async def user_info_v2( user_role=user_data.get("user_role"), spend=user_data.get("spend", 0.0), max_budget=user_data.get("max_budget"), + tpm_limit=user_data.get("tpm_limit"), + rpm_limit=user_data.get("rpm_limit"), models=user_data.get("models") or [], budget_duration=user_data.get("budget_duration"), budget_reset_at=user_data.get("budget_reset_at"), @@ -1298,7 +1300,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"): + if k in ("max_budget", "budget_duration", "tpm_limit", "rpm_limit"): if k in fields_set: non_default_values[k] = v elif k == "model_max_budget": @@ -1627,7 +1629,7 @@ async def _update_single_user_helper( await _invalidate_user_spend_counter_if_changed(non_default_values) - if not _USER_BUDGET_CACHE_FIELDS.isdisjoint(non_default_values) or "metadata" in data_json: + if not _USER_LIMIT_CACHE_FIELDS.isdisjoint(non_default_values) or "metadata" in data_json: await evict_and_broadcast( cache_keys=(non_default_values["user_id"],), user_api_key_cache=user_api_key_cache, @@ -1985,7 +1987,7 @@ async def bulk_user_update( ), ) - if not _USER_BUDGET_CACHE_FIELDS.isdisjoint(non_default_values): + if not _USER_LIMIT_CACHE_FIELDS.isdisjoint(non_default_values): for start in range(0, len(all_users_in_db), _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE): await asyncio.gather( *( diff --git a/tests/integration/management/test_user_rate_limit_updates.py b/tests/integration/management/test_user_rate_limit_updates.py new file mode 100644 index 00000000000..ce66da78e7f --- /dev/null +++ b/tests/integration/management/test_user_rate_limit_updates.py @@ -0,0 +1,381 @@ +import json +from collections.abc import Mapping +from contextlib import ExitStack +from typing import Final +from uuid import uuid4 + +import httpx +import pytest +from pydantic import JsonValue, TypeAdapter + +from tests.integration._support.client import ( + JSON_OBJECT, + Gateway, + eventually, + object_value, + string_value, +) +from tests.integration._support.database import read_rows +from tests.integration._support.wire import Reply, Request, wire_server + +_HEADER_VALUE_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) +_UPSTREAM_REPLIES: Final[Mapping[str, Mapping[str, JsonValue]]] = { + "/v1/chat/completions": { + "id": "chatcmpl_hook_isolation", + "object": "chat.completion", + "created": 1, + "model": "gpt-5.6", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + "/v1/responses": { + "id": "resp_hook_isolation", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.6", + "output": [ + { + "type": "message", + "id": "msg_hook_isolation", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "parallel_tool_calls": False, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, +} + + +def _chat(proxy: Gateway, model: str, key: str) -> httpx.Response: + return proxy.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"user rate limit probe {uuid4().hex}"}]}, + key=key, + ) + + +def _assert_user_rate_limit_error(response: httpx.Response, user: str, limit_type: str) -> None: + request_limit: Final[str | None] = _HEADER_VALUE_ADAPTER.validate_python( + response.headers.get("x-ratelimit-user-limit-requests") + ) + token_limit: Final[str | None] = _HEADER_VALUE_ADAPTER.validate_python( + response.headers.get("x-ratelimit-user-limit-tokens") + ) + context: Final = ( + f"Expected a user {limit_type} limit error for {user}, received HTTP {response.status_code} with " + f"user limits requests={request_limit}, tokens={token_limit}: {response.text}" + ) + assert response.status_code == 429, context + body: Final = JSON_OBJECT.validate_json(response.content) + error: Final = object_value(body["error"]) + message: Final = string_value(error["message"]) + assert error.get("type") == "throttling_error", context + assert message.startswith(f"Rate limit exceeded for user: {user}. Limit type: {limit_type}. Current limit: 1,"), ( + context + ) + + +def _route_request(proxy: Gateway, route: str, model: str, key: str, stream: bool) -> httpx.Response: + marker: Final = f"user rpm route probe {uuid4().hex}" + if route == "/v1/messages": + return proxy.request( + "POST", + route, + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": marker}]}, + key=key, + headers={"anthropic-version": "2023-06-01"}, + ) + if route == "/v1/responses": + return proxy.request( + "POST", + route, + {"model": model, "input": marker, "max_output_tokens": 16, "store": False}, + key=key, + ) + return proxy.request( + "POST", + route, + {"model": model, "messages": [{"role": "user", "content": marker}], "stream": stream}, + key=key, + ) + + +def _assert_route_user_requests_limit_error(response: httpx.Response, user: str, route: str) -> None: + request_limit: Final[str | None] = _HEADER_VALUE_ADAPTER.validate_python( + response.headers.get("x-ratelimit-user-limit-requests") + ) + token_limit: Final[str | None] = _HEADER_VALUE_ADAPTER.validate_python( + response.headers.get("x-ratelimit-user-limit-tokens") + ) + context: Final = ( + f"Expected a user requests limit error for {user} on {route}, received HTTP {response.status_code} with " + f"user limits requests={request_limit}, tokens={token_limit}: {response.text}" + ) + assert response.status_code == 429, context + body: Final = JSON_OBJECT.validate_json(response.content) + error: Final = object_value(body["error"]) + assert route != "/v1/messages" or body.get("type") == "error", context + expected_error_type: Final = "rate_limit_error" if route == "/v1/messages" else "throttling_error" + assert error.get("type") == expected_error_type, context + message: Final = string_value(error["message"]) + assert message.startswith(f"Rate limit exceeded for user: {user}. Limit type: requests. Current limit: 1,"), context + + +def _assert_user_rate_limit_on_every_proxy( + gateway: Gateway, + peer: Gateway, + model: str, + user: str, + key: str, +) -> None: + responses: Final = eventually( + lambda: (_chat(gateway, model, key), _chat(peer, model, key)), + lambda observed: all(response.status_code == 429 for response in observed), + seconds=10, + return_last_on_timeout=True, + ) + context: Final = tuple( + ( + response.status_code, + response.headers.get("x-ratelimit-user-limit-requests"), + response.headers.get("x-ratelimit-user-limit-tokens"), + response.text, + ) + for response in responses + ) + assert tuple(response.status_code for response in responses) == (429, 429), ( + f"Expected the user RPM limit on gateway and peer for {user}, received {context!r}" + ) + _assert_user_rate_limit_error(responses[0], user, "requests") + _assert_user_rate_limit_error(responses[1], user, "requests") + + +@pytest.mark.parametrize("field", ("tpm_limit", "rpm_limit")) +def test_user_rate_limit_lowered_on_gateway_is_enforced_by_peer(gateway: Gateway, peer: Gateway, field: str) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model() + user: Final = scenario.user(tpm_limit=100000, rpm_limit=1000) + key: Final = scenario.key(user_id=user, models=[model]) + + gateway_warm: Final = _chat(gateway, model, key) + peer_warm: Final = _chat(peer, model, key) + assert gateway_warm.status_code == 200, ( + f"Gateway rejected the initial user-limited request: {gateway_warm.text}" + ) + assert peer_warm.status_code == 200, f"Peer rejected the initial user-limited request: {peer_warm.text}" + assert peer_warm.headers.get("x-ratelimit-user-limit-requests") == "1000", peer_warm.headers + assert peer_warm.headers.get("x-ratelimit-user-limit-tokens") == "100000", peer_warm.headers + + gateway.post("/user/update", {"user_id": user, field: 1}) + + expected_tpm: Final = 1 if field == "tpm_limit" else 100000 + expected_rpm: Final = 1 if field == "rpm_limit" else 1000 + rows: Final = read_rows( + 'SELECT tpm_limit, rpm_limit FROM "LiteLLM_UserTable" WHERE user_id = %s', + (user,), + ) + assert rows == [{"tpm_limit": expected_tpm, "rpm_limit": expected_rpm}], ( + f"User {field} update did not persist without changing the other limit: {rows!r}" + ) + + limit_type: Final = "tokens" if field == "tpm_limit" else "requests" + peer_limited: Final = eventually( + lambda: _chat(peer, model, key), + lambda response: response.status_code == 429, + seconds=10, + return_last_on_timeout=True, + ) + _assert_user_rate_limit_error(peer_limited, user, limit_type) + + +def test_user_rate_limit_explicit_null_clears_and_omitted_limit_is_untouched(gateway: Gateway, peer: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model() + user: Final = scenario.user(tpm_limit=100000, rpm_limit=1) + key: Final = scenario.key(user_id=user, models=[model]) + + gateway_warm: Final = _chat(gateway, model, key) + assert gateway_warm.status_code == 200, f"Gateway rejected the initial request under RPM 1: {gateway_warm.text}" + peer_limited: Final = _chat(peer, model, key) + _assert_user_rate_limit_error(peer_limited, user, "requests") + + gateway.post("/user/update", {"user_id": user, "rpm_limit": None}) + + cleared_rows: Final = read_rows( + 'SELECT tpm_limit, rpm_limit FROM "LiteLLM_UserTable" WHERE user_id = %s', + (user,), + ) + assert cleared_rows == [{"tpm_limit": 100000, "rpm_limit": None}], ( + f"Clearing RPM changed the wrong user limits: {cleared_rows!r}" + ) + info: Final = gateway.get("/v2/user/info", {"user_id": user}) + assert info["tpm_limit"] == 100000, f"User info omitted or changed TPM after RPM clear: {info!r}" + assert info["rpm_limit"] is None, f"User info did not report the cleared RPM limit: {info!r}" + + gateway_after_clear: Final = _chat(gateway, model, key) + assert gateway_after_clear.status_code == 200, ( + f"Gateway still enforced RPM after it was cleared: {gateway_after_clear.text}" + ) + peer_after_clear: Final = eventually( + lambda: _chat(peer, model, key), + lambda response: response.status_code == 200, + seconds=10, + ) + assert peer_after_clear.status_code == 200, ( + f"Peer did not stop enforcing RPM after it was cleared: {peer_after_clear.text}" + ) + + gateway.post("/user/update", {"user_id": user, "tpm_limit": 50000}) + omitted_rows: Final = read_rows( + 'SELECT tpm_limit, rpm_limit FROM "LiteLLM_UserTable" WHERE user_id = %s', + (user,), + ) + assert omitted_rows == [{"tpm_limit": 50000, "rpm_limit": None}], ( + f"Omitting RPM during the TPM update changed it: {omitted_rows!r}" + ) + + +@pytest.mark.parametrize( + ("route", "stream", "upstream_target", "expected_rpm_header"), + ( + pytest.param("/v1/messages", False, "/v1/responses", "1000", id="messages"), + pytest.param("/v1/responses", False, "/v1/responses", "1000", id="responses"), + pytest.param("/v1/chat/completions", True, None, None, id="streaming-chat-completions"), + ), +) +def test_user_rpm_lowered_on_gateway_is_enforced_by_peer_on_llm_route( + gateway: Gateway, + peer: Gateway, + route: str, + stream: bool, + upstream_target: str | None, + expected_rpm_header: str | None, +) -> None: + with gateway.scenario() as scenario, ExitStack() as resources: + + def upstream(request: Request) -> Reply: + assert request.target == upstream_target, request.target + reply: Final = _UPSTREAM_REPLIES[request.target] + return Reply(body=json.dumps(reply).encode()) + + provider: Final = resources.enter_context(wire_server(upstream)) if upstream_target is not None else None + model: Final = ( + scenario.model() + if provider is None + else scenario.model(model="openai/gpt-5.6", api_base=provider.url + "/v1") + ) + user: Final = scenario.user(tpm_limit=100000, rpm_limit=1000) + key: Final = scenario.key(user_id=user, models=[model]) + + gateway_warm: Final = _route_request(gateway, route, model, key, stream) + peer_warm: Final = _route_request(peer, route, model, key, stream) + assert gateway_warm.status_code == 200, f"Gateway rejected {route}: {gateway_warm.text}" + assert peer_warm.status_code == 200, f"Peer rejected {route}: {peer_warm.text}" + assert expected_rpm_header is None or ( + peer_warm.headers.get("x-ratelimit-user-limit-requests") == expected_rpm_header + ), f"Peer returned unexpected user RPM headers for {route}: {dict(peer_warm.headers)!r}" + targets: Final = tuple(request.target for request in provider.drain()) if provider is not None else () + expected_targets: Final = (upstream_target, upstream_target) if upstream_target is not None else () + assert targets == expected_targets, targets + + gateway.post("/user/update", {"user_id": user, "rpm_limit": 1}) + rows: Final = read_rows( + 'SELECT tpm_limit, rpm_limit FROM "LiteLLM_UserTable" WHERE user_id = %s', + (user,), + ) + assert rows == [{"tpm_limit": 100000, "rpm_limit": 1}], f"User RPM update changed the wrong limits: {rows!r}" + + peer_limited: Final = eventually( + lambda: _route_request(peer, route, model, key, stream), + lambda response: response.status_code == 429, + seconds=10, + return_last_on_timeout=True, + ) + _assert_route_user_requests_limit_error(peer_limited, user, route) + + +def test_internal_user_cannot_clear_own_rpm_limit(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model() + user: Final = scenario.user(tpm_limit=100000, rpm_limit=1, user_role="internal_user") + key: Final = scenario.key(user_id=user, models=[model]) + denied: Final = gateway.request( + "POST", + "/user/update", + {"user_id": user, "rpm_limit": None}, + key=key, + ) + context: Final = f"Expected internal-user route denial, received HTTP {denied.status_code}: {denied.text}" + assert denied.status_code == 401, context + assert "Only proxy admin can be used to generate" in denied.text, context + assert "Route=/user/update" in denied.text, context + + rows: Final = read_rows( + 'SELECT tpm_limit, rpm_limit FROM "LiteLLM_UserTable" WHERE user_id = %s', + (user,), + ) + assert rows == [{"tpm_limit": 100000, "rpm_limit": 1}], f"Denied self-update changed user limits: {rows!r}" + + first_chat: Final = _chat(gateway, model, key) + assert first_chat.status_code == 200, f"Internal-user first chat was rejected: {first_chat.text}" + second_chat: Final = _chat(gateway, model, key) + _assert_user_rate_limit_error(second_chat, user, "requests") + + +def test_bulk_update_lowered_rpm_is_enforced_on_every_proxy(gateway: Gateway, peer: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model() + first_user: Final = scenario.user(tpm_limit=100000, rpm_limit=1000) + first_key: Final = scenario.key(user_id=first_user, models=[model]) + second_user: Final = scenario.user(tpm_limit=100000, rpm_limit=1000) + second_key: Final = scenario.key(user_id=second_user, models=[model]) + + warm_responses: Final = ( + _chat(gateway, model, first_key), + _chat(peer, model, first_key), + _chat(gateway, model, second_key), + _chat(peer, model, second_key), + ) + assert tuple(response.status_code for response in warm_responses) == (200, 200, 200, 200), ( + f"Expected both users to warm on gateway and peer: {tuple(response.text for response in warm_responses)!r}" + ) + + bulk_update: Final = gateway.post( + "/user/bulk_update", + { + "users": [ + {"user_id": first_user, "rpm_limit": 1}, + {"user_id": second_user, "rpm_limit": 1}, + ] + }, + ) + assert ( + bulk_update["total_requested"], + bulk_update["successful_updates"], + bulk_update["failed_updates"], + ) == (2, 2, 0), bulk_update + results_json: Final = bulk_update.get("results") + assert isinstance(results_json, list), bulk_update + results: Final = tuple(object_value(result) for result in results_json) + assert tuple((string_value(result["user_id"]), result["success"]) for result in results) == ( + (first_user, True), + (second_user, True), + ), results + + rows: Final = read_rows( + 'SELECT user_id, tpm_limit, rpm_limit FROM "LiteLLM_UserTable" WHERE user_id IN (%s, %s) ORDER BY user_id', + (first_user, second_user), + ) + expected_rows: Final = tuple( + {"user_id": user, "tpm_limit": 100000, "rpm_limit": 1} for user in sorted((first_user, second_user)) + ) + assert tuple(rows) == expected_rows, f"Bulk RPM update changed unexpected limits: {rows!r}" + + _assert_user_rate_limit_on_every_proxy(gateway, peer, model, first_user, first_key) + _assert_user_rate_limit_on_every_proxy(gateway, peer, model, second_user, second_key) 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 633c6738521..712107e4a32 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -5,7 +5,7 @@ import logging from collections.abc import Mapping, Sequence from datetime import datetime, timezone from types import SimpleNamespace -from typing import Final +from typing import TYPE_CHECKING, Final, cast from unittest.mock import AsyncMock, MagicMock import httpx @@ -47,6 +47,10 @@ from tests.unit.proxy.management_endpoints.jwt_key_mapping_doubles import ( client = TestClient(app) +if TYPE_CHECKING: + from litellm.caching.redis_cache import RedisCache + from litellm.proxy.utils import PrismaClient + @pytest.mark.asyncio async def test_ui_view_users_with_null_email(mocker, caplog): @@ -2250,6 +2254,18 @@ 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"]) +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" + + assert result[field] is None + assert other_field not in result + + def test_update_internal_user_params_keeps_original_max_budget_when_not_provided(): """ Test that _update_internal_user_params does not include max_budget @@ -2478,6 +2494,178 @@ async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker: broadcast.assert_awaited_once_with(cache_key=saved_user.user_id) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("field", "new_limit"), + [("tpm_limit", 100), ("rpm_limit", 1), ("tpm_limit", None), ("rpm_limit", None)], + ids=["tpm_limit", "rpm_limit", "tpm_limit-cleared", "rpm_limit-cleared"], +) +@pytest.mark.parametrize("all_users", [False, True], ids=["single-user", "bulk-all-users"]) +async def test_user_rate_limit_update_reaches_cached_user_on_every_worker( + mocker: MockerFixture, field: str, new_limit: int | None, all_users: bool +) -> None: + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.auth.auth_checks import get_user_object + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.internal_user_endpoints import bulk_user_update, user_update + from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( + BulkUpdateUserRequest, + UpdateUserRequestNoUserIDorEmail, + ) + + published: Final[list[tuple[str, str]]] = [] # mutable-ok: captures messages from the async Redis publisher + + class _RecordingRedisClient: + async def publish(self, channel: str, message: str) -> int: + published.append((channel, message)) + return 1 + + class _FakeRedisCache: + namespace: str | None = None + + def init_pubsub_client(self) -> _RecordingRedisClient: + return _RecordingRedisClient() + + class _UserTableMocks: + def __init__( + self, + find_first: AsyncMock, + find_unique: AsyncMock, + find_many: AsyncMock, + update_many: AsyncMock, + ) -> None: + self.find_first = find_first + self.find_unique = find_unique + self.find_many = find_many + self.update_many = update_many + + class _DatabaseMocks: + def __init__(self, litellm_usertable: _UserTableMocks) -> None: + self.litellm_usertable = litellm_usertable + + class _PrismaClientMock: + def __init__( + self, + db: _DatabaseMocks, + get_data: AsyncMock, + updated_user: LiteLLM_UserTable, + ) -> None: + self.db = db + self.get_data = get_data + self.updated_user = updated_user + self.update_data_payload: dict[str, object] | None = None + + async def update_data(self, user_id: str, data: dict[str, object], table_name: str) -> dict[str, object]: + self.update_data_payload = data + return {"user_id": user_id, "data": self.updated_user} + + saved_user: Final = LiteLLM_UserTable(user_id="user-spruce", tpm_limit=100000, rpm_limit=1000) + updated_user: Final = saved_user.model_copy(update={field: new_limit}) + old_limit: Final = 100000 if field == "tpm_limit" else 1000 + + prisma_client: Final = _PrismaClientMock( + db=_DatabaseMocks( + litellm_usertable=_UserTableMocks( + find_first=mocker.AsyncMock(return_value=saved_user), + find_unique=mocker.AsyncMock(return_value=updated_user), + find_many=mocker.AsyncMock(return_value=[saved_user]), + update_many=mocker.AsyncMock(return_value=1), + ) + ), + get_data=mocker.AsyncMock(return_value=saved_user), + updated_user=updated_user, + ) + prisma_client_for_auth: Final = cast("PrismaClient", prisma_client) + mocker.patch( # test-quality-ok: substitute the database dependency + "litellm.proxy.proxy_server.prisma_client", prisma_client + ) + + handling_worker_cache: Final = UserApiKeyCache() + other_worker_cache: Final = UserApiKeyCache() + await handling_worker_cache.async_set_cache( + key=saved_user.user_id, + value=saved_user, + model_type=LiteLLM_UserTable, + ) + await other_worker_cache.async_set_cache( + key=saved_user.user_id, + value=saved_user, + model_type=LiteLLM_UserTable, + ) + mocker.patch( # test-quality-ok: exercise a real isolated cache for the endpoint's worker + "litellm.proxy.proxy_server.user_api_key_cache", handling_worker_cache + ) + mocker.patch( # test-quality-ok: inject an in-memory pub/sub client without live Redis + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(), + ) + + handling_user_before: Final = await get_user_object( + user_id=saved_user.user_id, + prisma_client=prisma_client_for_auth, + user_api_key_cache=handling_worker_cache, + user_id_upsert=False, + ) + other_user_before: Final = await get_user_object( + user_id=saved_user.user_id, + prisma_client=prisma_client_for_auth, + user_api_key_cache=other_worker_cache, + user_id_upsert=False, + ) + assert handling_user_before is not None + assert handling_user_before.model_dump()[field] == old_limit + assert other_user_before is not None + assert other_user_before.model_dump()[field] == old_limit + + admin: Final = UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN) + if all_users: + await bulk_user_update( + data=BulkUpdateUserRequest( + all_users=True, + user_updates=UpdateUserRequestNoUserIDorEmail.model_validate({field: new_limit}), + ), + user_api_key_dict=admin, + litellm_changed_by=None, + ) + prisma_client.db.litellm_usertable.update_many.assert_awaited_once_with(where={}, data={field: new_limit}) + else: + await user_update( + data=UpdateUserRequest.model_validate({"user_id": saved_user.user_id, field: new_limit}), + user_api_key_dict=admin, + ) + assert prisma_client.update_data_payload is not None + assert prisma_client.update_data_payload[field] == new_limit + + remote_subscriber: Final = AuthCacheInvalidationSubscriber( + redis_cache=cast("RedisCache", _FakeRedisCache()), + user_api_key_cache=other_worker_cache, + ) + for _, message in published: + remote_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API + {"type": "message", "data": message} + ) + + handling_user_after: Final = await get_user_object( + user_id=saved_user.user_id, + prisma_client=prisma_client_for_auth, + user_api_key_cache=handling_worker_cache, + user_id_upsert=False, + ) + other_user_after: Final = await get_user_object( + user_id=saved_user.user_id, + prisma_client=prisma_client_for_auth, + user_api_key_cache=other_worker_cache, + user_id_upsert=False, + ) + assert handling_user_after is not None + assert handling_user_after.model_dump()[field] == new_limit + assert other_user_after is not None + assert other_user_after.model_dump()[field] == new_limit, ( + "another worker still enforces the old limit; the update was never broadcast" + ) + + def test_generate_request_base_validator(): """ Test that GenerateRequestBase validator converts empty string to None for max_budget @@ -2888,49 +3076,65 @@ async def test_user_update_rejects_silent_create_for_non_proxy_admin(mocker): @pytest.mark.asyncio -async def test_user_info_v2_proxy_admin_can_query_any_user(mocker): +async def test_user_info_v2_proxy_admin_can_query_any_user(mocker: MockerFixture) -> None: """ Test that proxy admin can query any user via /v2/user/info. """ from fastapi import Request - from litellm.proxy._types import UserInfoV2Response + from litellm.proxy._types import LiteLLM_UserTable, UserInfoV2Response from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2 - mock_prisma_client = mocker.MagicMock() + mock_user_row: Final = LiteLLM_UserTable( + user_id="target-user-123", + user_email="target@example.com", + user_alias="Target User", + user_role="internal_user", + spend=42.5, + max_budget=100.0, + tpm_limit=100000, + rpm_limit=1000, + models=["gpt-4"], + budget_duration="30d", + budget_reset_at=None, + metadata={"team": "engineering"}, + created_at=datetime(2024, 1, 1, tzinfo=timezone.utc), + updated_at=datetime(2024, 6, 1, tzinfo=timezone.utc), + sso_user_id="sso-abc", + teams=["team-1", "team-2"], + ) - mock_user_row = mocker.MagicMock() - mock_user_row.model_dump.return_value = { - "user_id": "target-user-123", - "user_email": "target@example.com", - "user_alias": "Target User", - "user_role": "internal_user", - "spend": 42.5, - "max_budget": 100.0, - "models": ["gpt-4"], - "budget_duration": "30d", - "budget_reset_at": None, - "metadata": {"team": "engineering"}, - "created_at": datetime(2024, 1, 1, tzinfo=timezone.utc), - "updated_at": datetime(2024, 6, 1, tzinfo=timezone.utc), - "sso_user_id": "sso-abc", - "teams": ["team-1", "team-2"], - } + class _UserTable: + def __init__(self, find_unique: AsyncMock) -> None: + self.find_unique = find_unique - async def mock_find_unique(*args, **kwargs): - if kwargs.get("where", {}).get("user_id") == "target-user-123": + class _Database: + def __init__(self, litellm_usertable: _UserTable) -> None: + self.litellm_usertable = litellm_usertable + + class _PrismaClient: + def __init__(self, db: _Database) -> None: + self.db = db + + async def mock_find_unique(*_args: object, **kwargs: object) -> LiteLLM_UserTable | None: + where: Final = kwargs.get("where") + if isinstance(where, Mapping) and where.get("user_id") == "target-user-123": return mock_user_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) + mock_prisma_client: Final = _PrismaClient( + db=_Database( + litellm_usertable=_UserTable(find_unique=mocker.AsyncMock(side_effect=mock_find_unique)) + ) + ) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - mock_request = mocker.MagicMock(spec=Request) + mock_request: Final = mocker.MagicMock(spec=Request) - admin_key = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) + admin_key: Final = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) - response = await user_info_v2( + response: Final = await user_info_v2( request=mock_request, user_id="target-user-123", user_api_key_dict=admin_key, @@ -2943,6 +3147,8 @@ async def test_user_info_v2_proxy_admin_can_query_any_user(mocker): assert response.user_role == "internal_user" assert response.spend == 42.5 assert response.max_budget == 100.0 + assert response.tpm_limit == 100000 + assert response.rpm_limit == 1000 assert response.models == ["gpt-4"] assert response.teams == ["team-1", "team-2"] assert response.sso_user_id == "sso-abc" @@ -3273,6 +3479,8 @@ async def test_user_info_v2_response_shape(mocker): "user_role", "spend", "max_budget", + "tpm_limit", + "rpm_limit", "models", "budget_duration", "budget_reset_at", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx index 2571eb344f5..d1f309c9f99 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.test.tsx @@ -457,6 +457,151 @@ describe("UserEditView", () => { expect(checkbox).toBeChecked(); }); }); + + describe("user rate limits", () => { + const userDataWithRateLimits = () => ({ + ...MOCK_USER_DATA, + user_info: { + ...MOCK_USER_DATA.user_info, + tpm_limit: 100000, + rpm_limit: 50, + }, + }); + + it("seeds the TPM and RPM inputs from the selected user", async () => { + renderWithProviders(); + + expect(await screen.findByRole("spinbutton", { name: /tpm limit/i })).toHaveValue(100000); + expect(await screen.findByRole("spinbutton", { name: /rpm limit/i })).toHaveValue(50); + }); + + it("keeps unset rate limits empty and omits them from an untouched save", async () => { + const onSubmit = vi.fn(); + const userDataWithNullRateLimits = { + ...MOCK_USER_DATA, + user_info: { + ...MOCK_USER_DATA.user_info, + tpm_limit: null, + rpm_limit: null, + }, + }; + renderWithProviders(); + + expect(await screen.findByRole("spinbutton", { name: /tpm limit/i })).toHaveValue(null); + expect(await screen.findByRole("spinbutton", { name: /rpm limit/i })).toHaveValue(null); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmit).toHaveBeenCalled(); + }); + expect(onSubmit.mock.calls[0][0]).not.toHaveProperty("tpm_limit"); + expect(onSubmit.mock.calls[0][0]).not.toHaveProperty("rpm_limit"); + }); + + it("omits unchanged rate limits from the submit payload", async () => { + const onSubmit = vi.fn(); + renderWithProviders(); + + await userEvent.click(await screen.findByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmit).toHaveBeenCalled(); + }); + expect(onSubmit.mock.calls[0][0]).not.toHaveProperty("tpm_limit"); + expect(onSubmit.mock.calls[0][0]).not.toHaveProperty("rpm_limit"); + }); + + it("submits zero when the stored TPM limit changes to zero", async () => { + const onSubmit = vi.fn(); + renderWithProviders(); + + fireEvent.change(await screen.findByRole("spinbutton", { name: /tpm limit/i }), { + target: { value: "0" }, + }); + await userEvent.click(await screen.findByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmit).toHaveBeenCalled(); + }); + expect(onSubmit.mock.calls[0][0].tpm_limit).toBe(0); + }); + + it("omits the TPM limit when the stored value is re-entered", async () => { + const onSubmit = vi.fn(); + renderWithProviders(); + + fireEvent.change(await screen.findByRole("spinbutton", { name: /tpm limit/i }), { + target: { value: "100000" }, + }); + await userEvent.click(await screen.findByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmit).toHaveBeenCalled(); + }); + expect(onSubmit.mock.calls[0][0]).not.toHaveProperty("tpm_limit"); + }); + + it("sends null only for a deliberately cleared TPM limit", async () => { + const onSubmit = vi.fn(); + renderWithProviders(); + + fireEvent.change(await screen.findByRole("spinbutton", { name: /tpm limit/i }), { + target: { value: "" }, + }); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmit).toHaveBeenCalled(); + }); + expect(onSubmit.mock.calls[0][0].tpm_limit).toBeNull(); + expect(onSubmit.mock.calls[0][0]).not.toHaveProperty("rpm_limit"); + }); + + it("submits a new RPM limit as a number", async () => { + const onSubmit = vi.fn(); + renderWithProviders(); + + fireEvent.change(await screen.findByRole("spinbutton", { name: /rpm limit/i }), { + target: { value: "1" }, + }); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmit).toHaveBeenCalled(); + }); + expect(onSubmit.mock.calls[0][0].rpm_limit).toBe(1); + expect(typeof onSubmit.mock.calls[0][0].rpm_limit).toBe("number"); + }); + + it.each(["-1", "1.5"])("rejects an invalid TPM limit of %s", async (value) => { + const onSubmit = vi.fn(); + renderWithProviders(); + + fireEvent.change(await screen.findByRole("spinbutton", { name: /tpm limit/i }), { + target: { value }, + }); + const submitButton = screen.getByRole("button", { name: /save changes/i }) as HTMLButtonElement; + const form = submitButton.form; + if (!form) { + throw new Error("User edit form was not rendered"); + } + fireEvent.submit(form); + + expect( + await screen.findByText("Enter a non-negative whole number, or leave empty for unlimited"), + ).toBeInTheDocument(); + expect(onSubmit).not.toHaveBeenCalled(); + }); + + it("hides both rate-limit inputs in bulk edit mode", async () => { + renderWithProviders(); + + await screen.findByRole("button", { name: /save changes/i }); + expect(screen.queryByRole("spinbutton", { name: /tpm limit/i })).not.toBeInTheDocument(); + expect(screen.queryByRole("spinbutton", { name: /rpm limit/i })).not.toBeInTheDocument(); + }); + }); + describe("submit payload parity", () => { const submittedPayload = async (props: Partial[0]> = {}) => { const onSubmit = vi.fn(); @@ -483,7 +628,7 @@ describe("UserEditView", () => { "user_id", "user_role", ]); - expect(payload).toStrictEqual({ + const expectedPayload = { user_id: "user-123", user_email: "test@example.com", user_alias: "Test User", @@ -494,7 +639,8 @@ describe("UserEditView", () => { metadata: { key1: "value1", key2: "value2" }, mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: [] }, mcp_tool_permissions: {}, - }); + }; + expect(payload).toStrictEqual(expectedPayload); expect(typeof payload.max_budget).toBe("number"); }); @@ -567,23 +713,29 @@ describe("UserEditView", () => { await waitFor(() => { expect(onSubmit).toHaveBeenCalled(); }); - expect(onSubmit.mock.calls[0][0]).toMatchObject({ + const expectedPayload = { user_id: "user-null", user_email: "null@example.com", user_alias: null, user_role: null, budget_duration: null, max_budget: null, - }); + }; + expect(onSubmit.mock.calls[0][0]).toMatchObject(expectedPayload); }); it("should keep the budget input's native step constraint armed", async () => { renderWithProviders(); const budgetInput = await screen.findByRole("spinbutton", { name: /max budget/i }); + const submitButton = screen.getByRole("button", { name: /save changes/i }) as HTMLButtonElement; + const form = submitButton.form; + if (!form) { + throw new Error("User edit form was not rendered"); + } expect(budgetInput).toHaveAttribute("step", "0.01"); expect(budgetInput).not.toHaveAttribute("min"); - expect(budgetInput.closest("form")).not.toHaveAttribute("novalidate"); + expect(form).not.toHaveAttribute("novalidate"); }); it("shows the tool matrix for servers the user reaches only through an access group or toolset", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx index 8254134f956..14a0401714f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx @@ -21,6 +21,15 @@ import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/comp import { useZodForm } from "@/lib/forms/useZodForm"; import { CircleHelp } from "lucide-react"; +const RATE_LIMIT_ERROR = "Enter a non-negative whole number, or leave empty for unlimited"; +const isBlank = (value: string | number | null | undefined): boolean => + value === null || value === undefined || String(value).trim() === ""; +const rateLimitField = z + .union([z.string(), z.number()]) + .nullish() + .transform((value) => (isBlank(value) ? null : Number(value))) + .pipe(z.number({ error: RATE_LIMIT_ERROR }).int(RATE_LIMIT_ERROR).nonnegative(RATE_LIMIT_ERROR).nullable()); + interface UserEditViewProps { userData: any; onCancel: () => void; @@ -53,23 +62,23 @@ const userEditShape = { models: z.array(z.string()), budget_duration: z.string().nullish(), metadata: z.string().nullish(), + tpm_limit: rateLimitField, + rpm_limit: rateLimitField, mcp_servers_and_groups: MCP_SELECTION_SHAPE.optional(), mcp_tool_permissions: z.record(z.string(), z.array(z.string())).optional(), }; -const budgetSchema = (unlimitedBudget: boolean) => +const userEditSchema = (unlimitedBudget: boolean) => z.object({ ...userEditShape, max_budget: z .union([z.string(), z.number()]) .nullish() - .refine( - (value) => unlimitedBudget || (value !== "" && value !== null && value !== undefined), - "Please enter a budget or select Unlimited Budget", - ), + .refine((value) => unlimitedBudget || !isBlank(value), "Please enter a budget or select Unlimited Budget"), }); -type UserEditFormValues = z.infer>; +type UserEditFormInput = z.input>; +type UserEditFormValues = z.output>; const buildMcpFieldValues = (objectPermission: ObjectPermission | null | undefined) => ({ mcp_servers_and_groups: { @@ -88,11 +97,18 @@ const toFormValues = ( objectPermission: ObjectPermission | null | undefined, isBulkEdit: boolean, canEditMcpPermissions: boolean, -): UserEditFormValues => { +): UserEditFormInput => { const maxBudget = userData.user_info?.max_budget; const isUnlimited = maxBudget === null || maxBudget === undefined; return { - ...(isBulkEdit ? {} : { user_id: userData.user_id, user_email: userData.user_info?.user_email }), + ...(isBulkEdit + ? {} + : { + user_id: userData.user_id, + user_email: userData.user_info?.user_email, + tpm_limit: userData.user_info?.tpm_limit ?? "", + rpm_limit: userData.user_info?.rpm_limit ?? "", + }), user_alias: userData.user_info?.user_alias, user_role: userData.user_info?.user_role, models: userData.user_info?.models || [], @@ -117,6 +133,9 @@ const parseMetadata = (metadata: string | null | undefined): ParsedMetadata => { } }; +const changedLimit = (value: number | null, stored: number | null | undefined): number | null | undefined => + value === (stored ?? null) ? undefined : value; + const labelWithHint = (label: string, hint: string): React.ReactNode => ( <> {label} @@ -147,7 +166,7 @@ export function UserEditView({ userData.user_id, () => userData.user_info?.model_max_budget ?? {}, ); - const schema = useMemo(() => budgetSchema(unlimitedBudget), [unlimitedBudget]); + const schema = useMemo(() => userEditSchema(unlimitedBudget), [unlimitedBudget]); const form = useZodForm(schema, { defaultValues: toFormValues(userData, objectPermission, isBulkEdit, canEditMcpPermissions), }); @@ -171,14 +190,20 @@ export function UserEditView({ return; } + const { tpm_limit: tpmLimitInput, rpm_limit: rpmLimitInput, ...formValues } = values; const modelBudgets = modelMaxBudgetUpdate(modelMaxBudget, userData.user_info?.model_max_budget); - onSubmit({ - ...values, + const tpmLimit = changedLimit(tpmLimitInput, userData.user_info?.tpm_limit); + const rpmLimit = changedLimit(rpmLimitInput, userData.user_info?.rpm_limit); + const payload = { + ...formValues, ...("metadata" in values ? { metadata: metadata.value } : {}), ...(modelBudgets !== undefined && { model_max_budget: modelBudgets }), + ...(tpmLimit !== undefined && { tpm_limit: tpmLimit }), + ...(rpmLimit !== undefined && { rpm_limit: rpmLimit }), max_budget: unlimitedBudget || values.max_budget === "" || values.max_budget === undefined ? null : values.max_budget, - }); + }; + onSubmit(payload); }; const modelOptions = [ @@ -293,6 +318,56 @@ export function UserEditView({ {({ id, value, onChange }) => } + {!isBulkEdit && ( + <> + + {({ ref, value, onChange, ...control }) => ( + onChange(event.target.value)} + onWheel={(event) => event.currentTarget.blur()} + placeholder="Unlimited" + /> + )} + + + + {({ ref, value, onChange, ...control }) => ( + onChange(event.target.value)} + onWheel={(event) => event.currentTarget.blur()} + placeholder="Unlimited" + /> + )} + + + )} + {/* Bulk edit forwards a fixed field list and has no single stored budget to diff against, so the editor would silently discard whatever was typed. */} {!isBulkEdit && ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx index 0d8505ffbc8..3eae1dbdfc5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx @@ -1,4 +1,4 @@ -import { render, screen, waitFor, within } from "@testing-library/react"; +import { fireEvent, render, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi, beforeEach } from "vitest"; import UserInfoView from "./user_info_view"; @@ -130,6 +130,42 @@ describe("UserInfoView", () => { expect(aliases.length).toBeGreaterThan(0); }); + it("seeds the user rate limits when opening the edit form", async () => { + mockUserGetInfoV2.mockResolvedValue({ + ...MOCK_USER_DATA, + tpm_limit: 100000, + rpm_limit: 50, + }); + + render(); + + expect(await screen.findByRole("spinbutton", { name: /tpm limit/i })).toHaveValue(100000); + expect(await screen.findByRole("spinbutton", { name: /rpm limit/i })).toHaveValue(50); + }); + + it("keeps the updated TPM and stored RPM when reopening the edit form", async () => { + mockUserGetInfoV2.mockResolvedValue({ + ...MOCK_USER_DATA, + tpm_limit: 100000, + rpm_limit: 50, + }); + + render(); + + fireEvent.change(await screen.findByRole("spinbutton", { name: /tpm limit/i }), { + target: { value: "" }, + }); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + await waitFor(() => { + expect(mockUserUpdateUserCall).toHaveBeenCalledTimes(1); + }); + + await userEvent.click(await screen.findByRole("button", { name: /edit settings/i })); + + expect(await screen.findByRole("spinbutton", { name: /tpm limit/i })).toHaveValue(null); + expect(await screen.findByRole("spinbutton", { name: /rpm limit/i })).toHaveValue(50); + }); + it("should render overview spend and budget with two decimal places", async () => { mockUserGetInfoV2.mockResolvedValue({ ...MOCK_USER_DATA, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx index c95badc587a..2056142b50a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx @@ -341,6 +341,8 @@ export default function UserInfoView({ user_alias: formValues.user_alias ?? userData.user_alias, models: formValues.models ?? userData.models, max_budget: formValues.max_budget === undefined ? userData.max_budget : formValues.max_budget, + tpm_limit: formValues.tpm_limit === undefined ? userData.tpm_limit : formValues.tpm_limit, + rpm_limit: formValues.rpm_limit === undefined ? userData.rpm_limit : formValues.rpm_limit, budget_duration: formValues.budget_duration === undefined ? userData.budget_duration : formValues.budget_duration, metadata: formValues.metadata ?? userData.metadata, @@ -401,6 +403,8 @@ export default function UserInfoView({ user_role: userData.user_role, models: userData.models, max_budget: userData.max_budget, + tpm_limit: userData.tpm_limit, + rpm_limit: userData.rpm_limit, budget_duration: userData.budget_duration, metadata: userData.metadata, // Without these the per-model budget editor mounts empty and a save diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index a7e9387cbe6..36c34f8f9c4 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -1091,6 +1091,8 @@ export interface UserInfoV2Response { user_role: string | null; spend: number; max_budget: number | null; + tpm_limit?: number | null; + rpm_limit?: number | null; models: string[]; budget_duration: string | null; budget_reset_at: string | null; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index df2f394efa4..c1107a7e910 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -49790,6 +49790,8 @@ export interface components { */ models: string[]; object_permission?: components["schemas"]["LiteLLM_ObjectPermissionTable"] | null; + /** Rpm Limit */ + rpm_limit?: number | null; /** * Spend * @default 0 @@ -49802,6 +49804,8 @@ export interface components { * @default [] */ teams: string[]; + /** Tpm Limit */ + tpm_limit?: number | null; /** Updated At */ updated_at?: string | null; /** User Alias */