From 9c890f658d3e63a554e0dfae341705e7a261c68d Mon Sep 17 00:00:00 2001 From: lei_lei Date: Sun, 13 Sep 2026 10:06:40 +0000 Subject: [PATCH 01/13] fix(proxy): warn when key limits exceed team caps on generate/update Compare rpm/tpm/max_parallel_requests/max_budget against the team on /key/generate and /key/update, and return field-level warnings when the key asks higher than the team cap, without rejecting the write yet --- litellm/proxy/_types.py | 12 + .../key_management_endpoints.py | 66 +++++- .../test_key_management_endpoints.py | 217 ++++++++++++++++++ 3 files changed, 288 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b315a2beac9..767e5d7b5ff 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1190,6 +1190,17 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): return v +KeyTeamLimitField = Literal["rpm_limit", "tpm_limit", "max_parallel_requests", "max_budget"] + + +class KeyTeamLimitWarning(TypedDict): + """Non-blocking warning when a key limit exceeds its team's effective cap.""" + + field: ReadOnly[KeyTeamLimitField] + requested: ReadOnly[float | int] + effective_team_cap: ReadOnly[float | int] + + class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase): index_name: str index_permissions: list[Literal["read", "write"]] @@ -1267,6 +1278,7 @@ class GenerateKeyResponse(KeyRequestBase): updated_by: str | None = None created_at: datetime | None = None updated_at: datetime | None = None + warnings: list[KeyTeamLimitWarning] | None = None @model_validator(mode="before") @classmethod diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 95ccb7bbe0b..a66fc35d443 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1501,6 +1501,33 @@ def check_team_key_rpm_tpm_limits( ) +def _collect_key_team_limit_warnings( + data: GenerateKeyRequest | UpdateKeyRequest, + team_table: LiteLLM_TeamTable | LiteLLM_TeamTableCachedObj, +) -> tuple[KeyTeamLimitWarning, ...]: + """ + Compare key rpm/tpm/max_parallel_requests/max_budget against the team's caps. + + Returns warnings when the key requests a higher value than the team allows. + Does not reject; runtime still applies the stricter team limit. + """ + comparisons: Final[tuple[tuple[KeyTeamLimitField, float | int | None, float | int | None], ...]] = ( + ("rpm_limit", data.rpm_limit, team_table.rpm_limit), + ("tpm_limit", data.tpm_limit, team_table.tpm_limit), + ("max_parallel_requests", data.max_parallel_requests, team_table.max_parallel_requests), + ("max_budget", data.max_budget, team_table.max_budget), + ) + return tuple( + KeyTeamLimitWarning( + field=field_name, + requested=requested, + effective_team_cap=team_cap, + ) + for field_name, requested, team_cap in comparisons + if requested is not None and team_cap is not None and requested > team_cap + ) + + async def _check_team_key_limits( team_table: LiteLLM_TeamTableCachedObj, data: GenerateKeyRequest | UpdateKeyRequest, @@ -1924,12 +1951,18 @@ async def generate_key_fn( user_api_key_cache=user_api_key_cache, ) - return await _common_key_generation_helper( + team_limit_warnings: Final = ( + _collect_key_team_limit_warnings(data=data, team_table=team_table) if team_table is not None else () + ) + response: Final = await _common_key_generation_helper( data=data, user_api_key_dict=user_api_key_dict, litellm_changed_by=litellm_changed_by, team_table=team_table, ) + if team_limit_warnings: + response.warnings = list(team_limit_warnings) + return response except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.generate_key_fn(): Exception occured - %s", e) @@ -2090,12 +2123,18 @@ async def generate_service_account_key_fn( data.user_id = None # do not allow user_id to be set for service account keys - return await _common_key_generation_helper( + team_limit_warnings: Final = ( + _collect_key_team_limit_warnings(data=data, team_table=team_table) if team_table is not None else () + ) + response: Final = await _common_key_generation_helper( data=data, user_api_key_dict=user_api_key_dict, litellm_changed_by=litellm_changed_by, team_table=team_table, ) + if team_limit_warnings: + response.warnings = list(team_limit_warnings) + return response def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_metadata: dict) -> dict: @@ -2538,9 +2577,10 @@ async def _process_single_key_update( # Get team object and check team limits if team_id is provided team_obj: LiteLLM_TeamTableCachedObj | None = None - if update_key_request.team_id is not None: + _team_id_to_check: Final = update_key_request.team_id or getattr(existing_key_row, "team_id", None) + if _team_id_to_check is not None: team_obj = await get_team_object( - team_id=update_key_request.team_id, + team_id=_team_id_to_check, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, check_db_only=True, @@ -2631,6 +2671,11 @@ async def _process_single_key_update( updated_key_info.pop("token", None) + team_limit_warnings: Final = ( + _collect_key_team_limit_warnings(data=update_key_request, team_table=team_obj) if team_obj is not None else () + ) + if team_limit_warnings: + return {**updated_key_info, "warnings": list(team_limit_warnings)} return updated_key_info @@ -2688,7 +2733,7 @@ async def _validate_update_key_data( premium_user: bool, prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, -) -> None: +) -> tuple[KeyTeamLimitWarning, ...]: """Validate permissions and constraints for key update.""" checked_prisma_client: Final = _require_prisma_client(prisma_client) @@ -2946,6 +2991,10 @@ async def _validate_update_key_data( if normalized_object_permission is not None: data.object_permission = LiteLLM_ObjectPermissionBase(**normalized_object_permission) + if team_obj is None: + return () + return _collect_key_team_limit_warnings(data=data, team_table=team_obj) + @router.post("/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper @@ -3067,7 +3116,7 @@ async def update_key_fn( key: Final = _resolve_token_to_update(data=data, existing_key_row=existing_key_row) data.key = key - await _validate_update_key_data( + team_limit_warnings: Final = await _validate_update_key_data( data=data, existing_key_row=existing_key_row, user_api_key_dict=user_api_key_dict, @@ -3177,7 +3226,10 @@ async def update_key_fn( if response is None: raise ValueError("Failed to update key got response = None") - return {"key": key, **response["data"]} + updated_key_info: Final[dict[str, object]] = {"key": key, **response["data"]} + if team_limit_warnings: + return {**updated_key_info, "warnings": list(team_limit_warnings)} + return updated_key_info # update based on remaining passed in values except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.update_key_fn(): Exception occured - %s", e) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 2ac52da57df..217ae0fb32f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_org_key_limits, _check_project_key_limits, _check_team_key_limits, + _collect_key_team_limit_warnings, _common_key_generation_helper, _enforce_upperbound_key_params, _get_and_validate_existing_key, @@ -3177,6 +3178,222 @@ async def test_update_key_fn_auto_rotate_disable(): assert result["auto_rotate"] is False +def test_collect_key_team_limit_warnings_above_team_caps(): + """Key limits above team caps produce field-level warnings without rejection.""" + team_table = LiteLLM_TeamTableCachedObj( + team_id="capped-team", + team_alias="capped-team", + rpm_limit=60, + tpm_limit=1000, + max_parallel_requests=5, + max_budget=10.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + data = GenerateKeyRequest( + team_id="capped-team", + rpm_limit=600, + tpm_limit=5000, + max_parallel_requests=20, + max_budget=100.0, + ) + + warnings = _collect_key_team_limit_warnings(data=data, team_table=team_table) + + assert warnings == ( + { + "field": "rpm_limit", + "requested": 600, + "effective_team_cap": 60, + }, + { + "field": "tpm_limit", + "requested": 5000, + "effective_team_cap": 1000, + }, + { + "field": "max_parallel_requests", + "requested": 20, + "effective_team_cap": 5, + }, + { + "field": "max_budget", + "requested": 100.0, + "effective_team_cap": 10.0, + }, + ) + + +def test_collect_key_team_limit_warnings_within_or_unset_caps(): + """No warning when key is within team caps, team has no cap, or key omits the field.""" + team_table = LiteLLM_TeamTableCachedObj( + team_id="partial-caps", + team_alias="partial-caps", + rpm_limit=60, + tpm_limit=None, + max_parallel_requests=5, + max_budget=None, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + data = UpdateKeyRequest( + key="sk-test-key-123456", + rpm_limit=60, + tpm_limit=999999, + max_parallel_requests=3, + max_budget=50.0, + ) + + warnings = _collect_key_team_limit_warnings(data=data, team_table=team_table) + + assert warnings == () + + +@pytest.mark.asyncio +async def test_generate_key_fn_attaches_team_limit_warnings(monkeypatch): + """/key/generate succeeds and returns warnings when key rpm exceeds team rpm.""" + team_table = LiteLLM_TeamTableCachedObj( + team_id="warn-team", + team_alias="warn-team", + rpm_limit=60, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + data = GenerateKeyRequest(team_id="warn-team", rpm_limit=600, key_alias="over-cap") + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-1234", + user_id="admin", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + AsyncMock(return_value=team_table), + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.key_generation_check", + MagicMock(), + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._check_team_key_limits", + AsyncMock(), + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.check_org_admin_can_generate_keys", + AsyncMock(), + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + MagicMock(), + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", + MagicMock(), + ) + + from litellm.proxy._types import GenerateKeyResponse + + generated = GenerateKeyResponse( + key="sk-generated-key-123456", + token_id="hashed", + team_id="warn-team", + rpm_limit=600, + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper", + AsyncMock(return_value=generated), + ) + + result = await generate_key_fn(data=data, user_api_key_dict=user_api_key_dict) + + assert result.key == "sk-generated-key-123456" + assert result.rpm_limit == 600 + assert result.warnings == [ + { + "field": "rpm_limit", + "requested": 600, + "effective_team_cap": 60, + } + ] + + +@pytest.mark.asyncio +async def test_validate_update_key_data_returns_team_limit_warnings(monkeypatch): + """/key/update validation returns warnings when updated limits exceed team caps.""" + existing_key = LiteLLM_VerificationToken( + token="hashed-token", + team_id="warn-team", + user_id="user-1", + models=[], + ) + team_table = LiteLLM_TeamTableCachedObj( + team_id="warn-team", + team_alias="warn-team", + rpm_limit=60, + max_budget=10.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + data = UpdateKeyRequest(key="sk-test-key-123456", rpm_limit=600, max_budget=100.0) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-1234", + user_id="admin", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + AsyncMock(return_value=team_table), + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._check_team_key_limits", + AsyncMock(), + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint", + AsyncMock(), + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.common_key_access_checks", + MagicMock(), + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.enforce_member_can_assign_access_groups", + MagicMock(), + ) + + warnings = await _validate_update_key_data( + data=data, + existing_key_row=existing_key, + user_api_key_dict=user_api_key_dict, + llm_router=None, + premium_user=True, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + ) + + assert warnings == ( + { + "field": "rpm_limit", + "requested": 600, + "effective_team_cap": 60, + }, + { + "field": "max_budget", + "requested": 100.0, + "effective_team_cap": 10.0, + }, + ) + + @pytest.mark.asyncio async def test_check_team_key_limits_no_existing_keys(): """ From 8c8ae1799b387858f93d05b9a8ebc9276562bcdf Mon Sep 17 00:00:00 2001 From: lei_lei Date: Sun, 13 Sep 2026 10:50:54 +0000 Subject: [PATCH 02/13] chore(ui): sync schema.d.ts for key team limit warnings --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 21 +++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..8b29eb33720 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -28026,6 +28024,8 @@ export interface components { updated_by?: string | null; /** User Id */ user_id?: string | null; + /** Warnings */ + warnings?: components["schemas"]["KeyTeamLimitWarning"][] | null; }; /** GenericGuardrailAPIInputs */ GenericGuardrailAPIInputs: { @@ -28478,6 +28478,21 @@ export interface components { /** Keys */ keys?: string[] | null; }; + /** + * KeyTeamLimitWarning + * @description Non-blocking warning when a key limit exceeds its team's effective cap. + */ + KeyTeamLimitWarning: { + /** Effective Team Cap */ + effective_team_cap: number; + /** + * Field + * @enum {string} + */ + field: "rpm_limit" | "tpm_limit" | "max_parallel_requests" | "max_budget"; + /** Requested */ + requested: number; + }; /** * KeyUpdateFields * @description Allowlist of bulk-broadcastable fields for /team/key/bulk_update; `extra="forbid"` blocks RBAC/ownership/scope mutations even by team admins. @@ -33091,6 +33106,8 @@ export interface components { user_id?: string | null; /** User Role */ user_role?: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; + /** Warnings */ + warnings?: components["schemas"]["KeyTeamLimitWarning"][] | null; }; /** * OAuth2SecurityScheme From 7b3b657d2307ef419df3226374e85dbb8af3a90f Mon Sep 17 00:00:00 2001 From: lei_lei Date: Sun, 13 Sep 2026 10:50:54 +0000 Subject: [PATCH 03/13] refactor(proxy): extract helper to restore C901 budget --- .../key_management_endpoints.py | 24 +++++++++++++++---- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a66fc35d443..f5b7e744b3b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1528,6 +1528,21 @@ def _collect_key_team_limit_warnings( ) + +def _maybe_add_key_team_limit_warnings( + payload: dict[str, object], + data: GenerateKeyRequest | UpdateKeyRequest, + team_table: LiteLLM_TeamTable | LiteLLM_TeamTableCachedObj | None, +) -> dict[str, object]: + """Attach team-limit warnings to a key update payload when caps are exceeded.""" + if team_table is None: + return payload + warnings = _collect_key_team_limit_warnings(data=data, team_table=team_table) + if not warnings: + return payload + return {**payload, "warnings": list(warnings)} + + async def _check_team_key_limits( team_table: LiteLLM_TeamTableCachedObj, data: GenerateKeyRequest | UpdateKeyRequest, @@ -2671,12 +2686,11 @@ async def _process_single_key_update( updated_key_info.pop("token", None) - team_limit_warnings: Final = ( - _collect_key_team_limit_warnings(data=update_key_request, team_table=team_obj) if team_obj is not None else () + return _maybe_add_key_team_limit_warnings( + updated_key_info, + update_key_request, + team_obj, ) - if team_limit_warnings: - return {**updated_key_info, "warnings": list(team_limit_warnings)} - return updated_key_info async def _validate_mcp_servers_for_key_update( From 1ced5b45f2a6db971803d1414b5d546dddb2b28f Mon Sep 17 00:00:00 2001 From: lei_lei Date: Sun, 13 Sep 2026 13:47:47 +0000 Subject: [PATCH 04/13] style(proxy): ruff format key_management_endpoints --- litellm/proxy/management_endpoints/key_management_endpoints.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f5b7e744b3b..7ef0d7ca736 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1528,7 +1528,6 @@ def _collect_key_team_limit_warnings( ) - def _maybe_add_key_team_limit_warnings( payload: dict[str, object], data: GenerateKeyRequest | UpdateKeyRequest, From b9a8c838b07e8d4e2075529a1eb4075b368b0fa9 Mon Sep 17 00:00:00 2001 From: lei_lei Date: Mon, 14 Sep 2026 01:52:24 +0000 Subject: [PATCH 05/13] fix(proxy): drop LIT001 mutable annotations on team-limit warnings --- litellm/proxy/_types.py | 4 ++-- .../proxy/management_endpoints/key_management_endpoints.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 767e5d7b5ff..49fe39a6d17 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1,7 +1,7 @@ import enum import json import os -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple @@ -1278,7 +1278,7 @@ class GenerateKeyResponse(KeyRequestBase): updated_by: str | None = None created_at: datetime | None = None updated_at: datetime | None = None - warnings: list[KeyTeamLimitWarning] | None = None + warnings: Sequence[KeyTeamLimitWarning] | None = None @model_validator(mode="before") @classmethod diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 7ef0d7ca736..fa08a316ef7 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1529,10 +1529,10 @@ def _collect_key_team_limit_warnings( def _maybe_add_key_team_limit_warnings( - payload: dict[str, object], + payload: Mapping[str, object], data: GenerateKeyRequest | UpdateKeyRequest, team_table: LiteLLM_TeamTable | LiteLLM_TeamTableCachedObj | None, -) -> dict[str, object]: +) -> Mapping[str, object]: """Attach team-limit warnings to a key update payload when caps are exceeded.""" if team_table is None: return payload @@ -3239,7 +3239,7 @@ async def update_key_fn( if response is None: raise ValueError("Failed to update key got response = None") - updated_key_info: Final[dict[str, object]] = {"key": key, **response["data"]} + updated_key_info: Final[Mapping[str, object]] = {"key": key, **response["data"]} if team_limit_warnings: return {**updated_key_info, "warnings": list(team_limit_warnings)} return updated_key_info From 15e8e5839487ec9a1ab5ae6e79abefd6a545ca17 Mon Sep 17 00:00:00 2001 From: lei_lei Date: Mon, 14 Sep 2026 04:49:14 +0000 Subject: [PATCH 06/13] test(proxy): cover team-limit warning attach helpers Exercise _maybe_add_key_team_limit_warnings and service-account /key/generate warning attachment so patch coverage includes those paths. --- .../test_key_management_endpoints.py | 123 ++++++++++++++++++ 1 file changed, 123 insertions(+) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 217ae0fb32f..f1a68112d03 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -37,6 +37,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_project_key_limits, _check_team_key_limits, _collect_key_team_limit_warnings, + _maybe_add_key_team_limit_warnings, _common_key_generation_helper, _enforce_upperbound_key_params, _get_and_validate_existing_key, @@ -53,6 +54,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( check_team_key_model_specific_limits, delete_verification_tokens, generate_key_fn, + generate_service_account_key_fn, generate_key_helper_fn, key_aliases, key_generation_check, @@ -3394,6 +3396,127 @@ async def test_validate_update_key_data_returns_team_limit_warnings(monkeypatch) ) +def test_maybe_add_key_team_limit_warnings_passthrough_and_attach(): + """Attach warnings to update payloads only when caps are exceeded.""" + team_table = LiteLLM_TeamTableCachedObj( + team_id="warn-team", + team_alias="warn-team", + rpm_limit=60, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + payload = {"key": "sk-test-key-123456", "rpm_limit": 600} + + assert ( + _maybe_add_key_team_limit_warnings( + payload, + UpdateKeyRequest(key="sk-test-key-123456", rpm_limit=600), + None, + ) + is payload + ) + assert ( + _maybe_add_key_team_limit_warnings( + payload, + UpdateKeyRequest(key="sk-test-key-123456", rpm_limit=30), + team_table, + ) + is payload + ) + + with_warnings = _maybe_add_key_team_limit_warnings( + payload, + UpdateKeyRequest(key="sk-test-key-123456", rpm_limit=600), + team_table, + ) + assert with_warnings is not payload + assert with_warnings["key"] == "sk-test-key-123456" + assert with_warnings["warnings"] == [ + { + "field": "rpm_limit", + "requested": 600, + "effective_team_cap": 60, + } + ] + + +@pytest.mark.asyncio +async def test_generate_service_account_key_fn_attaches_team_limit_warnings(monkeypatch): + """Service-account generate also surfaces team-limit warnings.""" + team_table = LiteLLM_TeamTableCachedObj( + team_id="warn-team", + team_alias="warn-team", + rpm_limit=60, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + data = GenerateKeyRequest(team_id="warn-team", rpm_limit=600, key_alias="sa-over-cap") + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-1234", + user_id="admin", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.check_org_admin_can_generate_keys", + AsyncMock(), + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_team_id_used_in_service_account_request", + AsyncMock(), + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + AsyncMock(return_value=team_table), + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._check_team_key_limits", + AsyncMock(), + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.key_generation_check", + MagicMock(), + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + MagicMock(), + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", + MagicMock(), + ) + + from litellm.proxy._types import GenerateKeyResponse + + generated = GenerateKeyResponse( + key="sk-service-account-123456", + token_id="hashed", + team_id="warn-team", + rpm_limit=600, + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper", + AsyncMock(return_value=generated), + ) + + result = await generate_service_account_key_fn( + data=data, user_api_key_dict=user_api_key_dict + ) + + assert result.key == "sk-service-account-123456" + assert result.warnings == [ + { + "field": "rpm_limit", + "requested": 600, + "effective_team_cap": 60, + } + ] + + @pytest.mark.asyncio async def test_check_team_key_limits_no_existing_keys(): """ From 1c69fe21cad39c71898827ee903008b402072682 Mon Sep 17 00:00:00 2001 From: leilei3167 Date: Wed, 16 Sep 2026 15:09:21 +0000 Subject: [PATCH 07/13] fix(proxy): warn on retained key limits during team reassignment /key/update often omits rpm/tpm/concurrency/budget when only team_id changes. Compare those retained values against the effective team caps so reassignment cannot silently keep an over-cap limit. Addresses greptile P1 on #40998. --- .../key_management_endpoints.py | 33 +++++++- .../test_key_management_endpoints.py | 75 +++++++++++++++++++ 2 files changed, 107 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index fa08a316ef7..bc22102a8c6 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1501,6 +1501,34 @@ def check_team_key_rpm_tpm_limits( ) +_KEY_TEAM_LIMIT_WARNING_FIELDS: Final[tuple[str, ...]] = ( + "rpm_limit", + "tpm_limit", + "max_parallel_requests", + "max_budget", +) + + +def _update_request_with_retained_team_limits( + data: UpdateKeyRequest, + existing_key_row: LiteLLM_VerificationToken, +) -> UpdateKeyRequest: + """Fill omitted limit fields from the existing key for team-cap warnings. + + /key/update often changes only team_id (or a subset of limits). Retained + rpm/tpm/concurrency/budget must still be compared against the effective + team caps so reassignment cannot silently keep an over-cap value. + """ + retained: dict[str, object] = {} + for field_name in _KEY_TEAM_LIMIT_WARNING_FIELDS: + if field_name in data.model_fields_set: + continue + retained[field_name] = getattr(existing_key_row, field_name, None) + if not retained: + return data + return data.model_copy(update=retained) + + def _collect_key_team_limit_warnings( data: GenerateKeyRequest | UpdateKeyRequest, team_table: LiteLLM_TeamTable | LiteLLM_TeamTableCachedObj, @@ -3006,7 +3034,10 @@ async def _validate_update_key_data( if team_obj is None: return () - return _collect_key_team_limit_warnings(data=data, team_table=team_obj) + return _collect_key_team_limit_warnings( + data=_update_request_with_retained_team_limits(data=data, existing_key_row=existing_key_row), + team_table=team_obj, + ) @router.post("/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)]) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index f1a68112d03..ac56e9aa007 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -3396,6 +3396,81 @@ async def test_validate_update_key_data_returns_team_limit_warnings(monkeypatch) ) +async def test_validate_update_key_data_warns_on_retained_limits_team_change(monkeypatch): + """Team reassignment without limit fields still warns on retained over-cap values.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import ( + KeyTeamLimitWarning, + LiteLLM_TeamTableCachedObj, + LiteLLM_VerificationToken, + LitellmUserRoles, + UpdateKeyRequest, + UserAPIKeyAuth, + ) + + existing = LiteLLM_VerificationToken( + token="hashed-token", + team_id="team-old", + user_id="user-1", + max_budget=100.0, + max_parallel_requests=50, + rpm_limit=None, + tpm_limit=None, + ) + new_team = LiteLLM_TeamTableCachedObj( + team_id="team-new", + max_budget=10.0, + max_parallel_requests=5, + rpm_limit=None, + tpm_limit=None, + ) + data = UpdateKeyRequest(key="sk-test", team_id="team-new") + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + AsyncMock(return_value=new_team), + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._check_team_key_limits", + AsyncMock(return_value=None), + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint", + AsyncMock(return_value=None), + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.enforce_member_can_assign_access_groups", + MagicMock(return_value=None), + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_team_change", + AsyncMock(return_value=None), + ) + # Avoid router requirement inside team-change validation path when mocked above + user = UserAPIKeyAuth( + user_id="user-1", + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + ) + warnings = await _validate_update_key_data( + data=data, + existing_key_row=existing, + user_api_key_dict=user, + llm_router=MagicMock(), + premium_user=True, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + ) + assert KeyTeamLimitWarning( + field="max_budget", requested=100.0, effective_team_cap=10.0 + ) in warnings + assert KeyTeamLimitWarning( + field="max_parallel_requests", requested=50, effective_team_cap=5 + ) in warnings + + + def test_maybe_add_key_team_limit_warnings_passthrough_and_attach(): """Attach warnings to update payloads only when caps are exceeded.""" team_table = LiteLLM_TeamTableCachedObj( From 0db9be28b50c49caf17508126ce517e734a54195 Mon Sep 17 00:00:00 2001 From: leilei3167 Date: Wed, 16 Sep 2026 15:11:02 +0000 Subject: [PATCH 08/13] test(proxy): align retained-limit warning regression with existing fixtures --- .../test_key_management_endpoints.py | 68 ++++++++++--------- 1 file changed, 35 insertions(+), 33 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index ac56e9aa007..c1a8d61c2b9 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -3398,34 +3398,30 @@ async def test_validate_update_key_data_returns_team_limit_warnings(monkeypatch) async def test_validate_update_key_data_warns_on_retained_limits_team_change(monkeypatch): """Team reassignment without limit fields still warns on retained over-cap values.""" - from unittest.mock import AsyncMock, MagicMock - - from litellm.proxy._types import ( - KeyTeamLimitWarning, - LiteLLM_TeamTableCachedObj, - LiteLLM_VerificationToken, - LitellmUserRoles, - UpdateKeyRequest, - UserAPIKeyAuth, - ) - existing = LiteLLM_VerificationToken( token="hashed-token", team_id="team-old", user_id="user-1", + models=[], max_budget=100.0, max_parallel_requests=50, - rpm_limit=None, - tpm_limit=None, ) new_team = LiteLLM_TeamTableCachedObj( team_id="team-new", + team_alias="team-new", max_budget=10.0, max_parallel_requests=5, - rpm_limit=None, - tpm_limit=None, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + data = UpdateKeyRequest(key="sk-test-key-123456", team_id="team-new") + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-1234", + user_id="admin", + user_role=LitellmUserRoles.PROXY_ADMIN, ) - data = UpdateKeyRequest(key="sk-test", team_id="team-new") monkeypatch.setattr( "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", @@ -3433,41 +3429,47 @@ async def test_validate_update_key_data_warns_on_retained_limits_team_change(mon ) monkeypatch.setattr( "litellm.proxy.management_endpoints.key_management_endpoints._check_team_key_limits", - AsyncMock(return_value=None), + AsyncMock(), ) monkeypatch.setattr( "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint", - AsyncMock(return_value=None), + AsyncMock(), + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.common_key_access_checks", + MagicMock(), ) monkeypatch.setattr( "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.enforce_member_can_assign_access_groups", - MagicMock(return_value=None), + MagicMock(), ) monkeypatch.setattr( "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_team_change", - AsyncMock(return_value=None), - ) - # Avoid router requirement inside team-change validation path when mocked above - user = UserAPIKeyAuth( - user_id="user-1", - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key="sk-admin", + AsyncMock(), ) + warnings = await _validate_update_key_data( data=data, existing_key_row=existing, - user_api_key_dict=user, + user_api_key_dict=user_api_key_dict, llm_router=MagicMock(), premium_user=True, prisma_client=MagicMock(), user_api_key_cache=MagicMock(), ) - assert KeyTeamLimitWarning( - field="max_budget", requested=100.0, effective_team_cap=10.0 - ) in warnings - assert KeyTeamLimitWarning( - field="max_parallel_requests", requested=50, effective_team_cap=5 - ) in warnings + + assert warnings == ( + { + "field": "max_parallel_requests", + "requested": 50, + "effective_team_cap": 5, + }, + { + "field": "max_budget", + "requested": 100.0, + "effective_team_cap": 10.0, + }, + ) From f76bb11bd80dd5a1b50033bfddb485bd2dbdd305 Mon Sep 17 00:00:00 2001 From: leilei3167 Date: Wed, 16 Sep 2026 15:25:35 +0000 Subject: [PATCH 09/13] fix(proxy): freeze team-limit warning payloads for LIT002 --- .../key_management_endpoints.py | 22 ++++++++++--------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index bc22102a8c6..d1ab423ed0e 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1519,11 +1519,13 @@ def _update_request_with_retained_team_limits( rpm/tpm/concurrency/budget must still be compared against the effective team caps so reassignment cannot silently keep an over-cap value. """ - retained: dict[str, object] = {} - for field_name in _KEY_TEAM_LIMIT_WARNING_FIELDS: - if field_name in data.model_fields_set: - continue - retained[field_name] = getattr(existing_key_row, field_name, None) + retained: Final = MappingProxyType( + { + field_name: getattr(existing_key_row, field_name, None) + for field_name in _KEY_TEAM_LIMIT_WARNING_FIELDS + if field_name not in data.model_fields_set + } + ) if not retained: return data return data.model_copy(update=retained) @@ -1567,7 +1569,7 @@ def _maybe_add_key_team_limit_warnings( warnings = _collect_key_team_limit_warnings(data=data, team_table=team_table) if not warnings: return payload - return {**payload, "warnings": list(warnings)} + return MappingProxyType({**payload, "warnings": warnings}) async def _check_team_key_limits( @@ -2003,7 +2005,7 @@ async def generate_key_fn( team_table=team_table, ) if team_limit_warnings: - response.warnings = list(team_limit_warnings) + response.warnings = team_limit_warnings return response except Exception as e: @@ -2175,7 +2177,7 @@ async def generate_service_account_key_fn( team_table=team_table, ) if team_limit_warnings: - response.warnings = list(team_limit_warnings) + response.warnings = team_limit_warnings return response @@ -3270,9 +3272,9 @@ async def update_key_fn( if response is None: raise ValueError("Failed to update key got response = None") - updated_key_info: Final[Mapping[str, object]] = {"key": key, **response["data"]} + updated_key_info: Final[Mapping[str, object]] = MappingProxyType({"key": key, **response["data"]}) if team_limit_warnings: - return {**updated_key_info, "warnings": list(team_limit_warnings)} + return MappingProxyType({**updated_key_info, "warnings": team_limit_warnings}) return updated_key_info # update based on remaining passed in values except Exception as e: From 10dce77b58f9df49499c6ced68697cf95e7c2e35 Mon Sep 17 00:00:00 2001 From: leilei3167 Date: Wed, 16 Sep 2026 15:39:54 +0000 Subject: [PATCH 10/13] fix(proxy): expose team-limit warnings as lists API/tests expect list payloads; keep internal collection as a tuple for LIT002. --- .../management_endpoints/key_management_endpoints.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d1ab423ed0e..db3bb10e9ee 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1569,7 +1569,7 @@ def _maybe_add_key_team_limit_warnings( warnings = _collect_key_team_limit_warnings(data=data, team_table=team_table) if not warnings: return payload - return MappingProxyType({**payload, "warnings": warnings}) + return MappingProxyType({**payload, "warnings": list(warnings)}) async def _check_team_key_limits( @@ -2005,7 +2005,7 @@ async def generate_key_fn( team_table=team_table, ) if team_limit_warnings: - response.warnings = team_limit_warnings + response.warnings = list(team_limit_warnings) return response except Exception as e: @@ -2177,7 +2177,7 @@ async def generate_service_account_key_fn( team_table=team_table, ) if team_limit_warnings: - response.warnings = team_limit_warnings + response.warnings = list(team_limit_warnings) return response @@ -3274,7 +3274,7 @@ async def update_key_fn( updated_key_info: Final[Mapping[str, object]] = MappingProxyType({"key": key, **response["data"]}) if team_limit_warnings: - return MappingProxyType({**updated_key_info, "warnings": team_limit_warnings}) + return MappingProxyType({**updated_key_info, "warnings": list(team_limit_warnings)}) return updated_key_info # update based on remaining passed in values except Exception as e: From 52cf58288320823a54eb0a738c8c30695dbdda7f Mon Sep 17 00:00:00 2001 From: leilei3167 Date: Wed, 16 Sep 2026 15:41:31 +0000 Subject: [PATCH 11/13] fix(proxy): annotate list() warnings boundary for LIT002 mutable-ok at API boundary so list payloads stay test-compatible. --- .../management_endpoints/key_management_endpoints.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index db3bb10e9ee..e29c6cab3ad 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1569,7 +1569,7 @@ def _maybe_add_key_team_limit_warnings( warnings = _collect_key_team_limit_warnings(data=data, team_table=team_table) if not warnings: return payload - return MappingProxyType({**payload, "warnings": list(warnings)}) + return MappingProxyType({**payload, "warnings": list(warnings)}) # mutable-ok: GenerateKeyResponse/tests expect list warnings async def _check_team_key_limits( @@ -2005,7 +2005,7 @@ async def generate_key_fn( team_table=team_table, ) if team_limit_warnings: - response.warnings = list(team_limit_warnings) + response.warnings = list(team_limit_warnings) # mutable-ok: GenerateKeyResponse/tests expect list warnings return response except Exception as e: @@ -2177,7 +2177,7 @@ async def generate_service_account_key_fn( team_table=team_table, ) if team_limit_warnings: - response.warnings = list(team_limit_warnings) + response.warnings = list(team_limit_warnings) # mutable-ok: GenerateKeyResponse/tests expect list warnings return response @@ -3274,7 +3274,7 @@ async def update_key_fn( updated_key_info: Final[Mapping[str, object]] = MappingProxyType({"key": key, **response["data"]}) if team_limit_warnings: - return MappingProxyType({**updated_key_info, "warnings": list(team_limit_warnings)}) + return MappingProxyType({**updated_key_info, "warnings": list(team_limit_warnings)}) # mutable-ok: key/update response tests expect list warnings return updated_key_info # update based on remaining passed in values except Exception as e: From 9771996d0414d4d41e67e5e77942cb475dda8107 Mon Sep 17 00:00:00 2001 From: leilei3167 Date: Wed, 16 Sep 2026 16:08:06 +0000 Subject: [PATCH 12/13] fix(proxy): soft-resolve existing team for key-update warnings Hard get_team_object only when the request sets team_id (limits check). Existing-key team is looked up for warnings and 404 is ignored so custom key policy and updates are not blocked. Also ruff-format list() boundaries. --- .../key_management_endpoints.py | 29 +++++++++++++++---- 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index e29c6cab3ad..e3cce7096cb 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1569,7 +1569,9 @@ def _maybe_add_key_team_limit_warnings( warnings = _collect_key_team_limit_warnings(data=data, team_table=team_table) if not warnings: return payload - return MappingProxyType({**payload, "warnings": list(warnings)}) # mutable-ok: GenerateKeyResponse/tests expect list warnings + return MappingProxyType( + {**payload, "warnings": list(warnings)} + ) # mutable-ok: GenerateKeyResponse/tests expect list warnings async def _check_team_key_limits( @@ -2619,12 +2621,13 @@ async def _process_single_key_update( # Enforce upperbound key params on update (don't fill defaults) _enforce_upperbound_key_params(update_key_request, fill_defaults=False) - # Get team object and check team limits if team_id is provided + # Get team object and check team limits if team_id is provided on the request. + # Existing-key team is soft-resolved below for warnings only — a missing team + # must not block an otherwise valid update (custom key policy runs later). team_obj: LiteLLM_TeamTableCachedObj | None = None - _team_id_to_check: Final = update_key_request.team_id or getattr(existing_key_row, "team_id", None) - if _team_id_to_check is not None: + if update_key_request.team_id is not None: team_obj = await get_team_object( - team_id=_team_id_to_check, + team_id=update_key_request.team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, check_db_only=True, @@ -2636,6 +2639,18 @@ async def _process_single_key_update( data=update_key_request, prisma_client=prisma_client, ) + elif getattr(existing_key_row, "team_id", None) is not None and prisma_client is not None: + try: + team_obj = await get_team_object( + team_id=existing_key_row.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + except HTTPException as e: + if e.status_code != status.HTTP_404_NOT_FOUND: + raise + team_obj = None # Validate team change if team is being changed if is_different_team(data=update_key_request, existing_key_row=existing_key_row): @@ -3274,7 +3289,9 @@ async def update_key_fn( updated_key_info: Final[Mapping[str, object]] = MappingProxyType({"key": key, **response["data"]}) if team_limit_warnings: - return MappingProxyType({**updated_key_info, "warnings": list(team_limit_warnings)}) # mutable-ok: key/update response tests expect list warnings + return MappingProxyType( + {**updated_key_info, "warnings": list(team_limit_warnings)} + ) # mutable-ok: key/update response tests expect list warnings return updated_key_info # update based on remaining passed in values except Exception as e: From a25d2926b0a3f86dc837419e5754f745c5ea7492 Mon Sep 17 00:00:00 2001 From: leilei3167 Date: Wed, 16 Sep 2026 16:15:00 +0000 Subject: [PATCH 13/13] refactor(proxy): extract soft team resolve to restore C901 budget Move existing-key team soft-resolve (404-tolerant) out of _process_single_key_update so McCabe stays within the strict C901 ceiling. --- .../key_management_endpoints.py | 44 +++++++++++++------ 1 file changed, 30 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index e3cce7096cb..722b42bbd61 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1574,6 +1574,28 @@ def _maybe_add_key_team_limit_warnings( ) # mutable-ok: GenerateKeyResponse/tests expect list warnings +async def _soft_resolve_existing_key_team_for_warnings( + existing_key_row: LiteLLM_VerificationToken, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, +) -> LiteLLM_TeamTableCachedObj | None: + """Resolve the existing key's team for warning payloads only; ignore missing teams.""" + team_id = getattr(existing_key_row, "team_id", None) + if team_id is None or prisma_client is None: + return None + try: + return await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + except HTTPException as e: + if e.status_code != status.HTTP_404_NOT_FOUND: + raise + return None + + async def _check_team_key_limits( team_table: LiteLLM_TeamTableCachedObj, data: GenerateKeyRequest | UpdateKeyRequest, @@ -2622,8 +2644,8 @@ async def _process_single_key_update( _enforce_upperbound_key_params(update_key_request, fill_defaults=False) # Get team object and check team limits if team_id is provided on the request. - # Existing-key team is soft-resolved below for warnings only — a missing team - # must not block an otherwise valid update (custom key policy runs later). + # Existing-key team is soft-resolved for warnings only — a missing team must not + # block an otherwise valid update (custom key policy runs later). team_obj: LiteLLM_TeamTableCachedObj | None = None if update_key_request.team_id is not None: team_obj = await get_team_object( @@ -2639,18 +2661,12 @@ async def _process_single_key_update( data=update_key_request, prisma_client=prisma_client, ) - elif getattr(existing_key_row, "team_id", None) is not None and prisma_client is not None: - try: - team_obj = await get_team_object( - team_id=existing_key_row.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=True, - ) - except HTTPException as e: - if e.status_code != status.HTTP_404_NOT_FOUND: - raise - team_obj = None + else: + team_obj = await _soft_resolve_existing_key_team_for_warnings( + existing_key_row=existing_key_row, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) # Validate team change if team is being changed if is_different_team(data=update_key_request, existing_key_row=existing_key_row):