From 5844e4b64d20586ad15bdc2591145f0bfb98bded Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 1 Aug 2026 16:40:16 -0700 Subject: [PATCH] fix(rate-limiter): key model rpm/tpm override takes precedence over team limit (LIT-2792) The v3 limiter emitted independent model_per_key and model_per_team descriptors, so the strictest value won and a key's model_rpm_limit or model_tpm_limit override was clamped by the team's, contradicting the documented resolution order (key metadata > key model_max_budget > team metadata). The team descriptor was also emitted twice per request, once inline in _create_rate_limit_descriptors and once via _add_team_model_rate_limit_descriptor_from_metadata, double counting every team per-model limit. The team descriptor is now emitted exactly once and only for the metrics the key does not override for the requested model, resolved per model and per metric through the same helper that feeds the model_per_key descriptor. Post-call TPM accounting honors the same suppression decision so an overriding key's tokens no longer drain the shared team counter used by inheriting keys. --- litellm/proxy/auth/auth_utils.py | 68 +++-- litellm/proxy/hooks/batch_rate_limiter.py | 14 +- .../hooks/parallel_request_limiter_v3.py | 160 +++++----- .../proxy/auth/test_auth_utils.py | 50 ++++ .../hooks/test_parallel_request_limiter_v3.py | 278 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 +- 6 files changed, 476 insertions(+), 100 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 644253ceac7..ca0fb91118e 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -901,6 +901,40 @@ def _get_deployment_default_tpm_limit(model_name: str) -> Optional[int]: return _get_deployment_default_limit(model_name, "default_api_key_tpm_limit") +def get_key_own_model_rate_limits( + user_api_key_dict: UserAPIKeyAuth, + rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"], +) -> dict[str, int] | None: + """ + 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 + + def get_key_model_rpm_limit( user_api_key_dict: UserAPIKeyAuth, model_name: Optional[str] = None, @@ -914,20 +948,9 @@ def get_key_model_rpm_limit( 3. Team metadata (model_rpm_limit) 4. Deployment default_api_key_rpm_limit (when model_name is provided) """ - # 1. Check key metadata first (takes priority) - if user_api_key_dict.metadata: - result = user_api_key_dict.metadata.get("model_rpm_limit") - if result: - return result - - # 2. Check model_max_budget - if user_api_key_dict.model_max_budget: - model_rpm_limit: Dict[str, Any] = {} - for model, budget in user_api_key_dict.model_max_budget.items(): - if isinstance(budget, dict) and budget.get("rpm_limit") is not None: - model_rpm_limit[model] = budget["rpm_limit"] - if model_rpm_limit: - return model_rpm_limit + key_own_limits = get_key_own_model_rate_limits(user_api_key_dict, "model_rpm_limit") + if key_own_limits is not None: + return key_own_limits # 3. Fallback to team metadata if user_api_key_dict.team_metadata: @@ -957,20 +980,9 @@ def get_key_model_tpm_limit( 3. Team metadata (model_tpm_limit) 4. Deployment default_api_key_tpm_limit (when model_name is provided) """ - # 1. Check key metadata first (takes priority) - if user_api_key_dict.metadata: - result = user_api_key_dict.metadata.get("model_tpm_limit") - if result: - return result - - # 2. Check model_max_budget (iterate per-model like RPM does) - if user_api_key_dict.model_max_budget: - model_tpm_limit: Dict[str, Any] = {} - for model, budget in user_api_key_dict.model_max_budget.items(): - if isinstance(budget, dict) and budget.get("tpm_limit") is not None: - model_tpm_limit[model] = budget["tpm_limit"] - if model_tpm_limit: - return model_tpm_limit + key_own_limits = get_key_own_model_rate_limits(user_api_key_dict, "model_tpm_limit") + if key_own_limits is not None: + return key_own_limits # 3. Fallback to team metadata if user_api_key_dict.team_metadata: diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 3477af36285..f5c1609bcf7 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -209,13 +209,25 @@ class _PROXY_BatchRateLimiter(CustomLogger): user_api_key_dict: UserAPIKeyAuth, data: Dict, ) -> List["RateLimitDescriptor"]: - return self.parallel_request_limiter._create_rate_limit_descriptors( + """Batch descriptors, mirroring what ``async_pre_call_hook`` enforces for a chat request. + + Team per-model limits live in their own helper (the same one the chat + path calls), so they have to be added explicitly here too. + """ + descriptors = self.parallel_request_limiter._create_rate_limit_descriptors( user_api_key_dict=user_api_key_dict, data=data, rpm_limit_type=None, tpm_limit_type=None, model_has_failures=False, ) + requested_model = data.get("model") + self.parallel_request_limiter._add_team_model_rate_limit_descriptor_from_metadata( + user_api_key_dict=user_api_key_dict, + requested_model=requested_model if isinstance(requested_model, str) else None, + descriptors=descriptors, + ) + return descriptors def _should_skip_batch_input_file_processing( self, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 3e7edd8cd06..8be2ed10e9c 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -35,6 +35,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( + get_key_own_model_rate_limits, get_key_tag_rpm_limit, get_model_rate_limit_from_metadata, ) @@ -376,6 +377,7 @@ class RequestRateLimiterStash: reserved_tokens: int = 0 reserved_model: Optional[str] = None reserved_scopes: FrozenSet[Tuple[str, str]] = field(default_factory=frozenset) + suppressed_tpm_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset) reservation_released: bool = False @@ -2031,11 +2033,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Returns list of descriptors for API key, user, team, team member, end user, model-specific, agent, and agent-session limits. """ - from litellm.proxy.auth.auth_utils import ( - get_team_model_rpm_limit, - get_team_model_tpm_limit, - ) - descriptors = [] # API Key rate limits @@ -2161,37 +2158,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptors=descriptors, ) - if ( - get_team_model_rpm_limit(user_api_key_dict) is not None - or get_team_model_tpm_limit(user_api_key_dict) is not None - ): - _tpm_limit_for_team_model = get_team_model_tpm_limit(user_api_key_dict) or {} - _rpm_limit_for_team_model = get_team_model_rpm_limit(user_api_key_dict) or {} - should_check_rate_limit = False - if requested_model in _tpm_limit_for_team_model: - should_check_rate_limit = True - elif requested_model in _rpm_limit_for_team_model: - should_check_rate_limit = True - - if should_check_rate_limit: - model_specific_tpm_limit = None - model_specific_rpm_limit = None - if requested_model in _tpm_limit_for_team_model: - model_specific_tpm_limit = _tpm_limit_for_team_model[requested_model] - if requested_model in _rpm_limit_for_team_model: - model_specific_rpm_limit = _rpm_limit_for_team_model[requested_model] - descriptors.append( - RateLimitDescriptor( - key="model_per_team", - value=f"{user_api_key_dict.team_id}:{requested_model}", - rate_limit={ - "requests_per_unit": model_specific_rpm_limit, - "tokens_per_unit": model_specific_tpm_limit, - "window_size": self.window_size, - }, - ) - ) - # Agent-level and session-level rate limits resolved_agent_id = self._get_resolved_agent_id(user_api_key_dict, data) @@ -2266,6 +2232,58 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return batch_limiter return None + def _team_model_limit_for_metric( + self, + user_api_key_dict: UserAPIKeyAuth, + requested_model: str, + rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"], + ) -> int | None: + """ + Team per-model limit that still applies to this key for one metric. + + A key that sets its own limit for this model (key metadata or + model_max_budget) already gets it enforced through the model_per_key + descriptor, and documented precedence is key > team. Enforcing the team + counter too would clamp the key back to the stricter of the two, so the + team limit is dropped for that metric only - a key overriding RPM alone + still inherits the team's TPM ceiling. + """ + team_limits = get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", rate_limit_key) or {} + team_limit = team_limits.get(requested_model) + if team_limit is None: + return None + + key_own_limits = get_key_own_model_rate_limits(user_api_key_dict, rate_limit_key) or {} + if key_own_limits.get(requested_model) is not None: + return None + + return team_limit + + def _suppressed_team_model_tpm_scopes( + self, + user_api_key_dict: UserAPIKeyAuth, + requested_model: str | None, + ) -> frozenset[tuple[str, str]]: + """ + Scopes whose team per-model TPM ceiling this key's own override replaced. + + Post-call reconciliation charges every candidate scope regardless of what + was enforced pre-call, so without this an overriding key keeps filling the + team counter it is no longer gated by and starves the keys that do inherit + the team limit. RPM needs no equivalent: its counters are incremented only + for descriptors that were actually emitted. + """ + if requested_model is None or not user_api_key_dict.team_id: + return frozenset() + + team_limits = get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_tpm_limit") or {} + if team_limits.get(requested_model) is None: + return frozenset() + if self._team_model_limit_for_metric(user_api_key_dict, requested_model, "model_tpm_limit") is not None: + return frozenset() + + return frozenset({("model_per_team", f"{user_api_key_dict.team_id}:{requested_model}")}) + def _add_team_model_rate_limit_descriptor_from_metadata( self, user_api_key_dict: UserAPIKeyAuth, @@ -2273,34 +2291,29 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptors: List[RateLimitDescriptor], ) -> None: """Add team model rate limit descriptor from team_metadata if applicable.""" - if ( - get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_rpm_limit") is not None - or get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_tpm_limit") is not None - ): - _tpm_limit_for_team_model = ( - get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_tpm_limit") or {} - ) - _rpm_limit_for_team_model = ( - get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", "model_rpm_limit") or {} - ) - should_check_rate_limit = ( - requested_model in _tpm_limit_for_team_model or requested_model in _rpm_limit_for_team_model - ) + if requested_model is None: + return - if should_check_rate_limit and requested_model is not None: - model_specific_tpm_limit = _tpm_limit_for_team_model.get(requested_model) - model_specific_rpm_limit = _rpm_limit_for_team_model.get(requested_model) - descriptors.append( - RateLimitDescriptor( - key="model_per_team", - value=f"{user_api_key_dict.team_id}:{requested_model}", - rate_limit={ - "requests_per_unit": model_specific_rpm_limit, - "tokens_per_unit": model_specific_tpm_limit, - "window_size": self.window_size, - }, - ) - ) + model_specific_rpm_limit = self._team_model_limit_for_metric( + user_api_key_dict, requested_model, "model_rpm_limit" + ) + model_specific_tpm_limit = self._team_model_limit_for_metric( + user_api_key_dict, requested_model, "model_tpm_limit" + ) + if model_specific_rpm_limit is None and model_specific_tpm_limit is None: + return + + descriptors.append( + RateLimitDescriptor( + key="model_per_team", + value=f"{user_api_key_dict.team_id}:{requested_model}", + rate_limit={ + "requests_per_unit": model_specific_rpm_limit, + "tokens_per_unit": model_specific_tpm_limit, + "window_size": self.window_size, + }, + ) + ) def _add_project_model_rate_limit_descriptor_from_metadata( self, @@ -2447,6 +2460,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model=requested_model, descriptors=descriptors, ) + stash.suppressed_tpm_scopes = self._suppressed_team_model_tpm_scopes( + user_api_key_dict=user_api_key_dict, + requested_model=requested_model, + ) # Project Level Rate Limits self._add_project_model_rate_limit_descriptor_from_metadata( @@ -2941,6 +2958,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): reserved_tokens = stash.reserved_tokens if stash is not None else 0 reserved_model = stash.reserved_model if stash is not None else None reserved_scopes: FrozenSet[Tuple[str, str]] = stash.reserved_scopes if stash is not None else frozenset() + suppressed_scopes: frozenset[tuple[str, str]] = ( + stash.suppressed_tpm_scopes if stash is not None else frozenset() + ) # Reconciliation must target the same model-scoped counter that the # pre-call reservation incremented. If a reservation was made, # ``reserved_model`` is authoritative; otherwise fall back to the @@ -2960,11 +2980,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # is empty, so every scope falls through the unreserved branch and # gets the full actual charge — matching pre-PR behavior. # ---------------------------------------------------------------- - targets = self._collect_tpm_scope_targets( - standard_logging_metadata=standard_logging_metadata, - kwargs=kwargs, - model_group=reconcile_model, - ) + targets = [ + target + for target in self._collect_tpm_scope_targets( + standard_logging_metadata=standard_logging_metadata, + kwargs=kwargs, + model_group=reconcile_model, + ) + if target not in suppressed_scopes + ] if reserved_tokens > 0 and total_tokens < reserved_tokens: verbose_proxy_logger.debug( f"Releasing unused TPM budget on success: " diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 1610d76efb7..183d1b7acda 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -20,6 +20,7 @@ from litellm.proxy.auth.auth_utils import ( get_key_mcp_rpm_limit, get_key_model_rpm_limit, get_key_model_tpm_limit, + get_key_own_model_rate_limits, get_key_tag_rpm_limit, get_model_from_request, get_project_model_rpm_limit, @@ -169,6 +170,55 @@ class TestGetKeyModelRpmLimit: assert result == {} +class TestGetKeyOwnModelRateLimits: + """Only limits set on the key itself, so callers can tell an override from an inherited limit.""" + + def test_metadata_limits_are_the_keys_own(self): + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-123", + metadata={"model_rpm_limit": {"gpt-4": 100}}, + team_metadata={"model_rpm_limit": {"gpt-4": 5}}, + ) + assert get_key_own_model_rate_limits(user_api_key_dict, "model_rpm_limit") == { + "gpt-4": 100 + } + + def test_model_max_budget_limits_are_the_keys_own(self): + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-123", + model_max_budget={ + "gpt-4": {"rpm_limit": 100, "tpm_limit": 2000}, + "gpt-3.5-turbo": {"tpm_limit": 3000}, + }, + ) + assert get_key_own_model_rate_limits(user_api_key_dict, "model_rpm_limit") == { + "gpt-4": 100 + } + assert get_key_own_model_rate_limits(user_api_key_dict, "model_tpm_limit") == { + "gpt-4": 2000, + "gpt-3.5-turbo": 3000, + } + + def test_team_metadata_is_not_the_keys_own(self): + team_only = UserAPIKeyAuth( + api_key="sk-123", + team_metadata={ + "model_rpm_limit": {"gpt-4": 5}, + "model_tpm_limit": {"gpt-4": 500}, + }, + ) + assert get_key_own_model_rate_limits(team_only, "model_rpm_limit") is None + assert get_key_own_model_rate_limits(team_only, "model_tpm_limit") is None + + def test_returns_none_when_nothing_configured_on_the_key(self): + assert ( + get_key_own_model_rate_limits( + UserAPIKeyAuth(api_key="sk-123"), "model_rpm_limit" + ) + is None + ) + + class TestGetKeyMcpRpmLimit: def test_empty_dict_limits_are_returned(self): key_override = UserAPIKeyAuth( diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 1c29c287c3a..9c79ddf350a 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -5100,3 +5100,281 @@ async def test_reserve_tpm_tokens_never_evaluates_the_requests_dimension(): f"reservation pass, got: {response}" ) assert [s["rate_limit_type"] for s in response["statuses"]] == ["tokens"] + + +def _make_team_model_limit_handler(): + """Limiter backed only by in-memory cache, plus the descriptors it enforced.""" + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + captured: List[Dict[str, Any]] = [] + + original_should_rate_limit = handler.should_rate_limit + + async def capturing_should_rate_limit(descriptors, **kwargs): + captured.extend(descriptors) + return await original_should_rate_limit(descriptors, **kwargs) + + handler.should_rate_limit = capturing_should_rate_limit + return handler, captured + + +async def _run_pre_call_hook(handler, user_api_key_dict, model="gpt-4"): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"model": model, "messages": [{"role": "user", "content": "hi"}]}, + call_type="", + ) + + +def _descriptors_with_key(captured, key): + return [d for d in captured if d["key"] == key] + + +@pytest.mark.asyncio +async def test_key_model_rpm_override_wins_over_team_model_rpm(): + """Key metadata model_rpm_limit beats the team's for the same model, so the + team counter must not also cap the key at the (lower) team value.""" + handler, captured = _make_team_model_limit_handler() + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-key-rpm-override"), + team_id="team-precedence", + metadata={"model_rpm_limit": {"gpt-4": 100}}, + team_metadata={"model_rpm_limit": {"gpt-4": 3}}, + ) + + await _run_pre_call_hook(handler, user_api_key_dict) + + key_descriptors = _descriptors_with_key(captured, "model_per_key") + assert len(key_descriptors) == 1 + assert key_descriptors[0]["rate_limit"]["requests_per_unit"] == 100 + + team_descriptors = _descriptors_with_key(captured, "model_per_team") + assert ( + team_descriptors == [] + ), f"the team's only limit is suppressed by the key override, so no team descriptor: {team_descriptors}" + + +@pytest.mark.asyncio +async def test_key_rpm_override_still_inherits_team_model_tpm(): + """Overriding RPM on the key must not release the team's TPM ceiling.""" + handler, captured = _make_team_model_limit_handler() + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-key-rpm-only"), + team_id="team-precedence", + metadata={"model_rpm_limit": {"gpt-4": 100}}, + team_metadata={ + "model_rpm_limit": {"gpt-4": 3}, + "model_tpm_limit": {"gpt-4": 500}, + }, + ) + + await _run_pre_call_hook(handler, user_api_key_dict) + + team_descriptors = _descriptors_with_key(captured, "model_per_team") + assert len(team_descriptors) == 1, f"expected one team descriptor, got: {team_descriptors}" + assert team_descriptors[0]["rate_limit"]["requests_per_unit"] is None + assert team_descriptors[0]["rate_limit"]["tokens_per_unit"] == 500 + + +@pytest.mark.asyncio +async def test_key_model_max_budget_rpm_override_wins_over_team_model_rpm(): + """The override may be expressed as model_max_budget[model].rpm_limit.""" + handler, captured = _make_team_model_limit_handler() + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-key-budget-override"), + team_id="team-precedence", + model_max_budget={"gpt-4": {"rpm_limit": 100}}, + team_metadata={"model_rpm_limit": {"gpt-4": 3}}, + ) + + await _run_pre_call_hook(handler, user_api_key_dict) + + key_descriptors = _descriptors_with_key(captured, "model_per_key") + assert len(key_descriptors) == 1 + assert key_descriptors[0]["rate_limit"]["requests_per_unit"] == 100 + + team_descriptors = _descriptors_with_key(captured, "model_per_team") + assert ( + team_descriptors == [] + ), f"the team's only limit is suppressed by the key override, so no team descriptor: {team_descriptors}" + + +@pytest.mark.asyncio +async def test_key_override_on_one_model_leaves_team_limit_on_other_model(): + """Suppression is per model: an override on gpt-4 must not free claude-3.""" + handler, captured = _make_team_model_limit_handler() + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-key-other-model"), + team_id="team-precedence", + metadata={"model_rpm_limit": {"gpt-4": 100}}, + team_metadata={"model_rpm_limit": {"gpt-4": 3, "claude-3": 7}}, + ) + + await _run_pre_call_hook(handler, user_api_key_dict, model="claude-3") + + team_descriptors = _descriptors_with_key(captured, "model_per_team") + assert len(team_descriptors) == 1, f"expected one team descriptor, got: {team_descriptors}" + assert team_descriptors[0]["value"] == "team-precedence:claude-3" + assert team_descriptors[0]["rate_limit"]["requests_per_unit"] == 7 + + +@pytest.mark.asyncio +async def test_team_model_descriptor_is_emitted_exactly_once(): + """Two code paths used to emit the same model_per_team descriptor, so every + request burned two slots of the team's per-model budget.""" + handler, captured = _make_team_model_limit_handler() + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-team-only"), + team_id="team-precedence", + team_metadata={"model_rpm_limit": {"gpt-4": 3}}, + ) + + await _run_pre_call_hook(handler, user_api_key_dict) + + team_descriptors = _descriptors_with_key(captured, "model_per_team") + assert len(team_descriptors) == 1, f"expected one team descriptor, got: {team_descriptors}" + assert team_descriptors[0]["rate_limit"]["requests_per_unit"] == 3 + + +@pytest.mark.asyncio +async def test_team_model_rpm_limit_admits_exactly_the_configured_requests(): + """A team model_rpm_limit of N admits N requests, not N/2.""" + handler, _ = _make_team_model_limit_handler() + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-team-budget-burn"), + team_id="team-burn", + team_metadata={"model_rpm_limit": {"gpt-4": 4}}, + ) + + for _ in range(4): + await _run_pre_call_hook(handler, user_api_key_dict) + + with pytest.raises(HTTPException) as exc_info: + await _run_pre_call_hook(handler, user_api_key_dict) + assert exc_info.value.status_code == 429 + + +@pytest.mark.asyncio +async def test_key_model_rpm_override_admits_its_own_higher_limit(): + """A key raising its per-model RPM above the team's gets the key's budget.""" + handler, _ = _make_team_model_limit_handler() + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-key-budget-burn"), + team_id="team-burn", + metadata={"model_rpm_limit": {"gpt-4": 5}}, + team_metadata={"model_rpm_limit": {"gpt-4": 2}}, + ) + + for _ in range(5): + await _run_pre_call_hook(handler, user_api_key_dict) + + with pytest.raises(HTTPException) as exc_info: + await _run_pre_call_hook(handler, user_api_key_dict) + assert exc_info.value.status_code == 429 + + +async def _increments_for_completed_call(handler, user_api_key_dict, model="gpt-4"): + """Drive one full request (pre-call gate, then success logging) and return the + counter increments the post-call reconciliation emitted.""" + await _run_pre_call_hook(handler, user_api_key_dict, model=model) + + increments: List[Dict[str, Any]] = [] + + async def mock_increment(increment_list, **kwargs): + increments.extend( + {"key": op["key"], "increment": op["increment_value"]} for op in increment_list + ) + + handler.internal_usage_cache.dual_cache.async_increment_cache_pipeline = mock_increment + + await handler.async_log_success_event( + kwargs={ + "litellm_params": {"metadata": {"model_group": model}}, + "standard_logging_object": { + "metadata": { + "user_api_key_hash": user_api_key_dict.api_key, + "user_api_key_team_id": user_api_key_dict.team_id, + } + }, + }, + response_obj=ModelResponse( + usage=Usage(prompt_tokens=40, completion_tokens=60, total_tokens=100) + ), + start_time=datetime.now(), + end_time=datetime.now(), + ) + return increments + + +@pytest.mark.asyncio +async def test_key_tpm_override_does_not_charge_the_team_token_counter(): + """A key exempt from the team's per-model TPM ceiling must not fill it either, + or its traffic starves the keys that do inherit the limit.""" + handler, _ = _make_team_model_limit_handler() + team_id = "team-tpm-accounting" + overriding_key = UserAPIKeyAuth( + api_key=hash_token("sk-tpm-override"), + team_id=team_id, + metadata={"model_tpm_limit": {"gpt-4": 100000}}, + team_metadata={"model_tpm_limit": {"gpt-4": 500}}, + ) + + increments = await _increments_for_completed_call(handler, overriding_key) + + team_tokens_key = handler.create_rate_limit_keys("model_per_team", f"{team_id}:gpt-4", "tokens") + assert not [ + i for i in increments if i["key"] == team_tokens_key + ], f"the team token counter must not be charged by an exempt key, got: {increments}" + + key_tokens_key = handler.create_rate_limit_keys( + "model_per_key", f"{overriding_key.api_key}:gpt-4", "tokens" + ) + assert [ + i for i in increments if i["key"] == key_tokens_key + ], f"the key's own token counter must still be charged, got: {increments}" + + +@pytest.mark.asyncio +async def test_inheriting_key_still_charges_the_team_token_counter(): + """The suppression is scoped to keys that override; an inheriting key still + spends the team's per-model token budget.""" + handler, _ = _make_team_model_limit_handler() + team_id = "team-tpm-accounting" + inheriting_key = UserAPIKeyAuth( + api_key=hash_token("sk-tpm-inherits"), + team_id=team_id, + team_metadata={"model_tpm_limit": {"gpt-4": 500}}, + ) + + increments = await _increments_for_completed_call(handler, inheriting_key) + + team_tokens_key = handler.create_rate_limit_keys("model_per_team", f"{team_id}:gpt-4", "tokens") + assert [ + i for i in increments if i["key"] == team_tokens_key + ], f"an inheriting key must charge the team token counter, got: {increments}" + + +@pytest.mark.asyncio +async def test_key_rpm_override_still_charges_the_team_token_counter(): + """Suppression is per metric: overriding RPM leaves the team's TPM ceiling in + force, so its counter must keep being charged.""" + handler, _ = _make_team_model_limit_handler() + team_id = "team-tpm-accounting" + rpm_only_override = UserAPIKeyAuth( + api_key=hash_token("sk-rpm-override-tpm-inherits"), + team_id=team_id, + metadata={"model_rpm_limit": {"gpt-4": 100}}, + team_metadata={ + "model_rpm_limit": {"gpt-4": 3}, + "model_tpm_limit": {"gpt-4": 500}, + }, + ) + + increments = await _increments_for_completed_call(handler, rpm_only_override) + + team_tokens_key = handler.create_rate_limit_keys("model_per_team", f"{team_id}:gpt-4", "tokens") + assert [ + i for i in increments if i["key"] == team_tokens_key + ], f"an RPM-only override must still charge the team token counter, got: {increments}" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 109d638fb9c..90552b6eae1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -23578,7 +23578,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 @@ -26172,7 +26172,7 @@ export interface components { [key: string]: unknown; } | null; /** Stream Timeout */ - stream_timeout?: number | string | null; + stream_timeout?: string | number | null; /** Tag Regex */ tag_regex?: string[] | null; /** Tags */ @@ -34304,7 +34304,7 @@ export interface components { [key: string]: unknown; } | null; /** Stream Timeout */ - stream_timeout?: number | string | null; + stream_timeout?: string | number | null; /** Tag Regex */ tag_regex?: string[] | null; /** Tags */