From ccd783d907d8b8a88add48c35732371e8701febf Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 26 Oct 2026 17:16:15 -0400 Subject: [PATCH] fix(proxy): reject negative RPM and TPM request limits (#48444) --- litellm/proxy/_types.py | 12 ++- .../management_endpoints/key_requests.py | 4 + .../management_endpoints/user_requests.py | 4 + tests/unit/proxy/test__types.py | 83 ++++++++++++++++++- 4 files changed, 98 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 2b01e02bb06..db433f07f3b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1766,8 +1766,10 @@ class BudgetNewRequest(LiteLLMPydanticObjectBase): max_parallel_requests: int | None = Field( default=None, description="Max concurrent requests allowed for this budget id." ) - tpm_limit: int | None = Field(default=None, description="Max tokens per minute, allowed for this budget id.") - rpm_limit: int | None = Field(default=None, description="Max requests per minute, allowed for this budget id.") + tpm_limit: int | None = Field(default=None, ge=0, description="Max tokens per minute, allowed for this budget id.") + rpm_limit: int | None = Field( + default=None, ge=0, description="Max requests per minute, allowed for this budget id." + ) tpd_limit: int | None = Field( default=None, description="Max tokens per day, charged by batch submissions, allowed for this budget id." ) @@ -1878,6 +1880,8 @@ RouterSettingsDict = Annotated[ class NewTeamRequest(TeamBase): + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) router_settings: RouterSettingsDict | None = None model_aliases: dict | None = None model_max_budget: GenericBudgetConfigType | None = Field( @@ -1957,8 +1961,8 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): team_alias: str | None = None organization_id: str | None = None metadata: dict | None = None - tpm_limit: int | None = None - rpm_limit: int | None = None + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) tpd_limit: int | None = None max_budget: float | None = None soft_budget: float | None = None diff --git a/litellm/types/proxy/management_endpoints/key_requests.py b/litellm/types/proxy/management_endpoints/key_requests.py index baa02b68d98..172357d8ac4 100644 --- a/litellm/types/proxy/management_endpoints/key_requests.py +++ b/litellm/types/proxy/management_endpoints/key_requests.py @@ -55,6 +55,8 @@ class LiteLLMKeyType(str, enum.Enum): class GenerateKeyRequest(KeyRequestBase): + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) soft_budget: float | None = None send_invite_email: bool | None = None key_type: LiteLLMKeyType | None = Field( @@ -82,6 +84,8 @@ def _is_mapping(value: object) -> TypeIs[Mapping[object, object]]: # guard-ok: class UpdateKeyRequest(KeyRequestBase): + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) # Note: the defaults of all Params here MUST BE NONE # else they will get overwritten duration: str | None = None diff --git a/litellm/types/proxy/management_endpoints/user_requests.py b/litellm/types/proxy/management_endpoints/user_requests.py index 49dda5995bc..e53dc4152c3 100644 --- a/litellm/types/proxy/management_endpoints/user_requests.py +++ b/litellm/types/proxy/management_endpoints/user_requests.py @@ -16,6 +16,8 @@ class NewUserRequestTeam(LiteLLMPydanticObjectBase): class NewUserRequest(GenerateRequestBase): + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) max_budget: float | None = None user_email: str | None = None user_alias: str | None = None @@ -47,6 +49,8 @@ class NewUserRequest(GenerateRequestBase): class UpdateUserRequestNoUserIDorEmail(GenerateRequestBase): # shared with BulkUpdateUserRequest + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) # repr=False keeps the plaintext out of management-endpoint alerts, which str() the request model password: str | None = Field(default=None, repr=False) spend: float | None = None diff --git a/tests/unit/proxy/test__types.py b/tests/unit/proxy/test__types.py index 68bfe4b34e5..76ca49362b4 100644 --- a/tests/unit/proxy/test__types.py +++ b/tests/unit/proxy/test__types.py @@ -2,23 +2,35 @@ import json from typing import Final import pytest -from pydantic import ValidationError +from pydantic import BaseModel, ValidationError from litellm.proxy._types import ( ROLES_WITHIN_ORG, + BudgetNewRequest, ChangePasswordRequest, GenerateKeyRequest, + GenerateKeyResponse, KeyRequest, LiteLLM_AuditLogs, + LiteLLM_BudgetTable, + LiteLLM_EndUserTable, LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, LitellmUserRoles, + NewCustomerRequest, NewMCPServerRequest, + NewTeamRequest, NewUserRequest, + NewUserResponse, OrganizationMemberUpdateRequest, + RegenerateKeyRequest, ResetSpendRequest, UpdateKeyRequest, UpdateMCPServerRequest, + UpdateTeamRequest, UpdateUserRequest, + UpdateUserRequestNoUserIDorEmail, UserAPIKeyAuth, ) @@ -213,6 +225,75 @@ def test_a_max_budget_that_was_actually_sent_is_kept(sent): assert GenerateKeyRequest(max_budget=sent).max_budget == sent +RATE_LIMIT_REQUESTS: Final = ( + (GenerateKeyRequest, {}), + (RegenerateKeyRequest, {}), + (UpdateKeyRequest, {"key": "sk-test"}), + (NewUserRequest, {}), + (UpdateUserRequestNoUserIDorEmail, {}), + (UpdateUserRequest, {"user_id": "user-test"}), + (NewTeamRequest, {}), + (UpdateTeamRequest, {"team_id": "team-test"}), + (BudgetNewRequest, {}), + (NewCustomerRequest, {"user_id": "customer-test"}), +) + + +@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS) +@pytest.mark.parametrize("field", ("rpm_limit", "tpm_limit")) +@pytest.mark.parametrize("value", (-1, -100)) +def test_management_requests_reject_negative_rate_limits( + request_model: type[BaseModel], identifiers: dict[str, str], field: str, value: int +) -> None: + with pytest.raises(ValidationError) as error: + request_model.model_validate({**identifiers, field: value}) + assert error.value.errors()[0]["loc"] == (field,) + + +@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS) +@pytest.mark.parametrize("value", (None, 0, 100)) +def test_management_requests_preserve_explicit_rate_limits( + request_model: type[BaseModel], identifiers: dict[str, str], value: int | None +) -> None: + limits: Final = {"rpm_limit": value, "tpm_limit": value} + request: Final = request_model.model_validate({**identifiers, **limits}) + assert request.model_dump(include={"rpm_limit", "tpm_limit"}, exclude_unset=True) == limits + + +@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS) +def test_management_requests_leave_omitted_rate_limits_unset( + request_model: type[BaseModel], identifiers: dict[str, str] +) -> None: + request: Final = request_model.model_validate({**identifiers}) + assert request.model_dump(include={"rpm_limit", "tpm_limit"}) == {"rpm_limit": None, "tpm_limit": None} + assert request.model_dump(include={"rpm_limit", "tpm_limit"}, exclude_unset=True) == {} + + +@pytest.mark.parametrize(("response_model", "identifiers"), ( + (GenerateKeyResponse, {"key": "sk-test"}), + (NewUserResponse, {"key": "sk-test"}), + (LiteLLM_TeamTable, {"team_id": "team-test"}), + (LiteLLM_UserTable, {"user_id": "user-test"}), + (LiteLLM_BudgetTable, {"budget_id": "budget-test"}), +)) +def test_existing_negative_rate_limits_remain_readable( + response_model: type[BaseModel], identifiers: dict[str, str] +) -> None: + limits: Final = {"rpm_limit": -1, "tpm_limit": -100} + response: Final = response_model.model_validate({**identifiers, **limits}) + assert response.model_dump(include={"rpm_limit", "tpm_limit"}) == limits + + +def test_existing_customer_negative_rate_limits_remain_readable() -> None: + customer: Final = LiteLLM_EndUserTable.model_validate({ + "user_id": "customer-test", "blocked": False, + "litellm_budget_table": {"rpm_limit": -1, "tpm_limit": -100}, + }) + assert customer.litellm_budget_table is not None + assert customer.litellm_budget_table.rpm_limit == -1 + assert customer.litellm_budget_table.tpm_limit == -100 + + USER_IDENTIFIER_CASES = ( ({"user_id": "user-1"}, True), ({"user_email": "user@example.com"}, True),