scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

feat(key-management): grant-time enforcement for team per-model limit allocation

Browse source 
The guaranteed_throughput allocation check missed per-model limits granted
through model_max_budget, on both the incoming request and existing team
keys, so a guaranteed team could still be overallocated through that
spelling. Requested and allocated limits now resolve through
resolve_own_model_rate_limits, the same helper the rate limiter uses at
request time, so grant-time math cannot drift from enforcement.

New opt-in general_settings.enforce_team_model_limit_allocation runs the
model-specific allocation check on every team key create/update, not just
guaranteed_throughput requests. Since a key's own per-model limit
deliberately wins over the team's at request time, this flag is how an
admin makes the team per-model limit a true cap: overallocation is
rejected with a 400 at grant time. The aggregate rpm/tpm check keeps its
guaranteed_throughput-only trigger because the team descriptor still
binds every key at request time. Over-allocation 400s now name the model.
This commit is contained in:
ryan-crabbe-berri 2026-08-01 23:18:32 -07:00
parent c957e39465
commit 982e8076dd
5 changed files with 341 additions and 45 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/_types.py
View file

@ -2412,6 +2412,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="By default, the user calling /team/new is automatically added to the new team as a team admin. If True, proxy admins are no longer auto-added; members explicitly listed in members_with_roles are unaffected. Default is False.",
)
enforce_team_model_limit_allocation: bool | None = Field(
None,
description="If True, a team key's per-model rpm/tpm limits are rejected at create/update time when they would overallocate the team's per-model limits, regardless of rpm_limit_type/tpm_limit_type. A key's own per-model limit wins over the team's at request time, so this is how the team limit becomes a true cap. Default is False.",
)
maximum_spend_logs_retention_period: str | None = Field(
None,
description="Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted.",

61
litellm/proxy/auth/auth_utils.py
View file

@ -902,6 +902,42 @@ def _get_deployment_default_tpm_limit(model_name: str) -> int | None:
return _get_deployment_default_limit(model_name, "default_api_key_tpm_limit")
def resolve_own_model_rate_limits(
metadata: object,
model_max_budget: object,
rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"],
) -> dict[str, int] | None:
"""
Per-model limits an entity sets on itself, from the two places one can carry
them.
Priority order (returns first found, per metric):
1. A ``model_rpm_limit`` / ``model_tpm_limit`` map in metadata
2. The per-model ``rpm_limit`` / ``tpm_limit`` entries of a model_max_budget
Takes the raw maps rather than a typed row so runtime enforcement
(UserAPIKeyAuth) and grant-time allocation math (verification-token rows,
key request bodies) resolve a key's own limits identically; both sources are
untyped JSON, so every level is shape-checked before use.
"""
if isinstance(metadata, dict):
metadata_limits = metadata.get(rate_limit_key)
if metadata_limits and isinstance(metadata_limits, dict):
return metadata_limits
if isinstance(model_max_budget, dict):
budget_field = "rpm_limit" if rate_limit_key == "model_rpm_limit" else "tpm_limit"
budget_limits = {
model: budget[budget_field]
for model, budget in model_max_budget.items()
if isinstance(budget, dict) and budget.get(budget_field) is not None
}
if budget_limits:
return budget_limits
return None
def get_key_own_model_rate_limits(
user_api_key_dict: UserAPIKeyAuth,
rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"],
@ -910,30 +946,15 @@ def get_key_own_model_rate_limits(
Get the per-model limits configured on the key itself, ignoring anything it
only inherits (team metadata, deployment defaults).
Priority order (returns first found):
1. Key metadata (model_rpm_limit / model_tpm_limit)
2. Key model_max_budget (rpm_limit / tpm_limit per model)
Callers that must know whether the key *overrides* an inherited limit use
this; callers that want the effective limit use get_key_model_rpm_limit /
get_key_model_tpm_limit, which continue the chain past the key.
"""
if user_api_key_dict.metadata:
metadata_limits = user_api_key_dict.metadata.get(rate_limit_key)
if metadata_limits:
return metadata_limits
if user_api_key_dict.model_max_budget:
budget_field = "rpm_limit" if rate_limit_key == "model_rpm_limit" else "tpm_limit"
budget_limits = {
model: budget[budget_field]
for model, budget in user_api_key_dict.model_max_budget.items()
if isinstance(budget, dict) and budget.get(budget_field) is not None
}
if budget_limits:
return budget_limits
return None
return resolve_own_model_rate_limits(
metadata=user_api_key_dict.metadata,
model_max_budget=user_api_key_dict.model_max_budget,
rate_limit_key=rate_limit_key,
)
def get_key_model_rpm_limit(

100
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -55,7 +55,10 @@ from litellm.proxy.auth.auth_checks import (
get_project_object,
get_team_object,
)
from litellm.proxy.auth.auth_utils import abbreviate_api_key
from litellm.proxy.auth.auth_utils import (
abbreviate_api_key,
resolve_own_model_rate_limits,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.callback_utils import (
decrypt_callback_vars,
@ -1158,6 +1161,44 @@ async def _common_key_generation_helper(
return response
def _requested_model_specific_limits(
data: Union[GenerateKeyRequest, UpdateKeyRequest],
rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"],
) -> Mapping[str, int] | None:
"""
Per-model limits this request would grant the key, from any of the shapes a
caller can express them in: the top-level field, the same map under
metadata, or per-model limits inside model_max_budget.
"""
explicit_limits = data.model_rpm_limit if rate_limit_key == "model_rpm_limit" else data.model_tpm_limit
if explicit_limits:
return explicit_limits
return resolve_own_model_rate_limits(
metadata=data.metadata,
model_max_budget=data.model_max_budget,
rate_limit_key=rate_limit_key,
)
def _allocated_model_specific_limits(
keys: Sequence[LiteLLM_VerificationToken],
rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"],
) -> Mapping[str, int]:
"""Per-model total already handed out across an entity's existing keys."""
per_key_limits = tuple(
resolve_own_model_rate_limits(
metadata=key.metadata,
model_max_budget=key.model_max_budget,
rate_limit_key=rate_limit_key,
)
or {}
for key in keys
)
allocated_models = frozenset(model for limits in per_key_limits for model in limits)
return {model: sum(limits.get(model, 0) for limits in per_key_limits) for model in allocated_models}
def _check_key_model_specific_limits(
keys: list[LiteLLM_VerificationToken],
data: Union[GenerateKeyRequest, UpdateKeyRequest],
@ -1171,33 +1212,20 @@ def _check_key_model_specific_limits(
Generic function to check if a key is allocating model specific limits.
Raises an error if we're overallocating.
"""
model_rpm_limit = getattr(data, "model_rpm_limit", None) or (
data.metadata.get("model_rpm_limit", None) if data.metadata else None
)
model_tpm_limit = getattr(data, "model_tpm_limit", None) or (
data.metadata.get("model_tpm_limit", None) if data.metadata else None
)
model_rpm_limit = _requested_model_specific_limits(data, "model_rpm_limit")
model_tpm_limit = _requested_model_specific_limits(data, "model_tpm_limit")
if model_rpm_limit is None and model_tpm_limit is None:
return
# get total model specific tpm/rpm limit
model_specific_rpm_limit: dict[str, int] = {}
model_specific_tpm_limit: dict[str, int] = {}
for key in keys:
if key.metadata.get("model_rpm_limit", None) is not None:
for model, rpm_limit in key.metadata.get("model_rpm_limit", {}).items():
model_specific_rpm_limit[model] = model_specific_rpm_limit.get(model, 0) + rpm_limit
if key.metadata.get("model_tpm_limit", None) is not None:
for model, tpm_limit in key.metadata.get("model_tpm_limit", {}).items():
model_specific_tpm_limit[model] = model_specific_tpm_limit.get(model, 0) + tpm_limit
model_specific_rpm_limit = _allocated_model_specific_limits(keys, "model_rpm_limit")
model_specific_tpm_limit = _allocated_model_specific_limits(keys, "model_tpm_limit")
if model_rpm_limit is not None:
for model, rpm_limit in model_rpm_limit.items():
if entity_rpm_limit is not None and model_specific_rpm_limit.get(model, 0) + rpm_limit > entity_rpm_limit:
raise HTTPException(
status_code=400,
detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Key RPM limit={rpm_limit} is greater than {entity_type} RPM limit={entity_rpm_limit}",
detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Key RPM limit={rpm_limit} is greater than {entity_type} RPM limit={entity_rpm_limit} for model={model}",
)
elif entity_model_rpm_limit_dict:
entity_model_specific_rpm_limit = entity_model_rpm_limit_dict.get(model)
@ -1207,7 +1235,7 @@ def _check_key_model_specific_limits(
):
raise HTTPException(
status_code=400,
detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Key RPM limit={rpm_limit} is greater than {entity_type} RPM limit={entity_model_specific_rpm_limit}",
detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Key RPM limit={rpm_limit} is greater than {entity_type} RPM limit={entity_model_specific_rpm_limit} for model={model}",
)
if model_tpm_limit is not None:
@ -1215,7 +1243,7 @@ def _check_key_model_specific_limits(
if entity_tpm_limit is not None and model_specific_tpm_limit.get(model, 0) + tpm_limit > entity_tpm_limit:
raise HTTPException(
status_code=400,
detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Key TPM limit={tpm_limit} is greater than {entity_type} TPM limit={entity_tpm_limit}",
detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Key TPM limit={tpm_limit} is greater than {entity_type} TPM limit={entity_tpm_limit} for model={model}",
)
elif entity_model_tpm_limit_dict:
entity_model_specific_tpm_limit = entity_model_tpm_limit_dict.get(model)
@ -1225,7 +1253,7 @@ def _check_key_model_specific_limits(
):
raise HTTPException(
status_code=400,
detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Key TPM limit={tpm_limit} is greater than {entity_type} TPM limit={entity_model_specific_tpm_limit}",
detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Key TPM limit={tpm_limit} is greater than {entity_type} TPM limit={entity_model_specific_tpm_limit} for model={model}",
)
@ -1309,6 +1337,21 @@ def check_team_key_rpm_tpm_limits(
)
def _is_team_model_limit_allocation_enforced() -> bool:
"""
Whether every team key must fit inside the team's per-model limits, not just
keys asking for guaranteed throughput.
Off by default: at runtime a key's own per-model limit deliberately wins over
the team's, so without this an admin can hand one key a ceiling above the
team's. Operators who want the team limit to be a true cap opt in here, and
the rejection happens at grant time rather than silently at request time.
"""
from litellm.proxy.proxy_server import general_settings
return general_settings.get("enforce_team_model_limit_allocation") is True
async def _check_team_key_limits(
team_table: LiteLLM_TeamTableCachedObj,
data: Union[GenerateKeyRequest, UpdateKeyRequest],
@ -1317,9 +1360,16 @@ async def _check_team_key_limits(
"""
Check if the team key is allocating guaranteed throughput limits. If so, raise an error if we're overallocating.
Only runs check if tpm_limit_type or rpm_limit_type is "guaranteed_throughput"
Only runs check if tpm_limit_type or rpm_limit_type is "guaranteed_throughput".
The model-specific half also runs for every team key when
general_settings.enforce_team_model_limit_allocation is on; the aggregate
half needs no such gate, since the team's own rpm/tpm descriptor still binds
every key at request time regardless of what each key was granted.
"""
if data.tpm_limit_type != "guaranteed_throughput" and data.rpm_limit_type != "guaranteed_throughput":
is_guaranteed_throughput = (
data.tpm_limit_type == "guaranteed_throughput" or data.rpm_limit_type == "guaranteed_throughput"
)
if not is_guaranteed_throughput and not _is_team_model_limit_allocation_enforced():
return
# get all team keys
# calculate allocated tpm/rpm limit
@ -1338,6 +1388,8 @@ async def _check_team_key_limits(
team_table=team_table,
data=data,
)
if not is_guaranteed_throughput:
return
check_team_key_rpm_tpm_limits(
keys=keys,
team_table=team_table,

214
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -15256,3 +15256,217 @@ async def test_rotate_master_key_rotates_sso_identity_assertions(
prisma_client=mock_prisma_client,
new_master_key="sk-new-master-key",
)
def _team_with_model_limits(
team_id: str,
metadata: dict,
tpm_limit: int | None = None,
rpm_limit: int | None = None,
) -> LiteLLM_TeamTableCachedObj:
return LiteLLM_TeamTableCachedObj(
team_id=team_id,
team_alias="test-team",
tpm_limit=tpm_limit,
rpm_limit=rpm_limit,
max_budget=100.0,
spend=0.0,
models=[],
blocked=False,
members_with_roles=[],
metadata=metadata,
)
@pytest.mark.parametrize(
"team_metadata_key, budget_field, expected_detail",
[
(
"model_rpm_limit",
"rpm_limit",
"Allocated RPM limit=0 + Key RPM limit=500 is greater than team RPM limit=100",
),
(
"model_tpm_limit",
"tpm_limit",
"Allocated TPM limit=0 + Key TPM limit=500 is greater than team TPM limit=100",
),
],
)
def test_check_team_key_model_specific_limits_counts_request_model_max_budget(
team_metadata_key, budget_field, expected_detail
):
"""A per-model limit granted through model_max_budget is the same override the
limiter enforces at runtime, so it must count against the team's allocation."""
team_table = _team_with_model_limits(
team_id="test-team-mmb-request",
metadata={team_metadata_key: {"gpt-4": 100}},
)
data = GenerateKeyRequest(model_max_budget={"gpt-4": {budget_field: 500}})
with pytest.raises(HTTPException) as exc_info:
check_team_key_model_specific_limits(
keys=[],
team_table=team_table,
data=data,
)
assert exc_info.value.status_code == 400
assert expected_detail in str(exc_info.value.detail)
def test_check_team_key_model_specific_limits_counts_existing_key_model_max_budget():
"""Existing keys hold their per-model limits in model_max_budget just as often
as in metadata; both must count toward what the team has already handed out."""
keys = [
LiteLLM_VerificationToken(
token="test-token-mmb-1",
team_id="test-team-mmb-existing",
model_max_budget={"gpt-4": {"rpm_limit": 80}},
),
]
team_table = _team_with_model_limits(
team_id="test-team-mmb-existing",
metadata={"model_rpm_limit": {"gpt-4": 100}},
)
data = GenerateKeyRequest(model_rpm_limit={"gpt-4": 30})
with pytest.raises(HTTPException) as exc_info:
check_team_key_model_specific_limits(
keys=keys,
team_table=team_table,
data=data,
)
assert exc_info.value.status_code == 400
assert (
"Allocated RPM limit=80 + Key RPM limit=30 is greater than team RPM limit=100"
in str(exc_info.value.detail)
)
def test_check_team_key_model_specific_limits_prefers_metadata_over_model_max_budget():
"""Mirrors runtime precedence: a metadata map wins outright, so a
model_max_budget entry for a model the map omits is not additionally counted."""
keys = [
LiteLLM_VerificationToken(
token="test-token-precedence",
team_id="test-team-precedence",
metadata={"model_rpm_limit": {"gpt-4": 40}},
model_max_budget={"gpt-4": {"rpm_limit": 900}},
),
]
team_table = _team_with_model_limits(
team_id="test-team-precedence",
metadata={"model_rpm_limit": {"gpt-4": 100}},
)
check_team_key_model_specific_limits(
keys=keys,
team_table=team_table,
data=GenerateKeyRequest(model_rpm_limit={"gpt-4": 60}),
)
def _prisma_client_returning(keys: list) -> AsyncMock:
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=keys)
return mock_prisma_client
@pytest.mark.asyncio
async def test_check_team_key_limits_enforces_model_allocation_when_flag_on(monkeypatch):
"""With enforce_team_model_limit_allocation on, a plain create (no limit_type)
can no longer be granted a per-model limit above the team's."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"enforce_team_model_limit_allocation": True},
)
mock_prisma_client = _prisma_client_returning([])
team_table = _team_with_model_limits(
team_id="test-team-flag-on",
metadata={"model_rpm_limit": {"gpt-4": 100}},
)
with pytest.raises(HTTPException) as exc_info:
await _check_team_key_limits(
team_table=team_table,
data=GenerateKeyRequest(model_rpm_limit={"gpt-4": 500}),
prisma_client=mock_prisma_client,
)
assert exc_info.value.status_code == 400
assert "is greater than team RPM limit=100" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_check_team_key_limits_allows_model_overallocation_by_default(monkeypatch):
"""Default stays permissive: without the flag and without guaranteed_throughput,
the same create is granted and the team's keys are never even read."""
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
mock_prisma_client = _prisma_client_returning([])
team_table = _team_with_model_limits(
team_id="test-team-flag-off",
metadata={"model_rpm_limit": {"gpt-4": 100}},
)
await _check_team_key_limits(
team_table=team_table,
data=GenerateKeyRequest(model_rpm_limit={"gpt-4": 500}),
prisma_client=mock_prisma_client,
)
mock_prisma_client.db.litellm_verificationtoken.find_many.assert_not_called()
@pytest.mark.asyncio
async def test_check_team_key_limits_allows_model_limit_within_remaining_allocation(monkeypatch):
"""The flag rejects overallocation, not allocation: a key that fits in what the
team has left is still granted."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"enforce_team_model_limit_allocation": True},
)
mock_prisma_client = _prisma_client_returning(
[
LiteLLM_VerificationToken(
token="test-token-allocated",
team_id="test-team-flag-within",
metadata={"model_rpm_limit": {"gpt-4": 60}},
)
]
)
team_table = _team_with_model_limits(
team_id="test-team-flag-within",
metadata={"model_rpm_limit": {"gpt-4": 100}},
)
await _check_team_key_limits(
team_table=team_table,
data=GenerateKeyRequest(model_rpm_limit={"gpt-4": 30}),
prisma_client=mock_prisma_client,
)
@pytest.mark.asyncio
async def test_check_team_key_limits_flag_does_not_gate_aggregate_limits(monkeypatch):
"""The flag covers per-model allocation only. The team's aggregate rpm/tpm
descriptor still binds every key at request time, so a plain create asking for
more than the team's aggregate stays a runtime concern, not a grant-time one."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"enforce_team_model_limit_allocation": True},
)
mock_prisma_client = _prisma_client_returning([])
team_table = _team_with_model_limits(
team_id="test-team-flag-aggregate",
metadata={},
tpm_limit=1000,
rpm_limit=100,
)
await _check_team_key_limits(
team_table=team_table,
data=GenerateKeyRequest(tpm_limit=5000, rpm_limit=500),
prisma_client=mock_prisma_client,
)

7
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -22789,6 +22789,11 @@ export interface components {
* @default false
*/
enable_public_model_hub: boolean;
/**
* Enforce Team Model Limit Allocation
* @description If True, a team key's per-model rpm/tpm limits are rejected at create/update time when they would overallocate the team's per-model limits, regardless of rpm_limit_type/tpm_limit_type. A key's own per-model limit wins over the team's at request time, so this is how the team limit becomes a true cap. Default is False.
*/
enforce_team_model_limit_allocation?: boolean | null;
/**
* Forward Client Headers To Llm Api
* @description If True, forwards client headers (e.g. Authorization) to the LLM API. Required for Claude Code with Max subscription.
@ -23584,7 +23589,7 @@ export interface components {
* @description Default role assigned to new users created
* @default internal_user_viewer
*/
user_role: ("internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer") | null;
user_role: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null;
};
/**
* DefaultTeamSSOParams