diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 8bcfbb67539..a9c36dfc512 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -50,6 +50,7 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES = frozenset( KeyManagementRoutes.KEY_BLOCK.value, KeyManagementRoutes.KEY_UNBLOCK.value, KeyManagementRoutes.KEY_BULK_UPDATE.value, + KeyManagementRoutes.TEAM_KEY_BULK_UPDATE.value, ] ) @@ -671,6 +672,7 @@ class RouteChecks: "/key/service-account/generate", "/key/block", "/key/unblock", + "/team/key/bulk_update", ] ) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 411c02d6262..d6702cfbdf1 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2709,27 +2709,11 @@ async def bulk_update_team_keys( """ Apply one update payload to many keys inside a single team. - Scoped to a single team — pass `team_id` plus either `key_ids` (specific - keys in that team) or `all_keys_in_team=True`. The same `update_fields` - payload is applied to every selected key. + Pass `team_id` plus either `key_ids` or `all_keys_in_team=True`. The + `update_fields` payload is broadcast to every selected key. Per-key + failures are returned in `failed_updates` rather than aborting the batch. - Callable by proxy admins, or by team admins with `KEY_UPDATE` permission - for the target team. - - Each key update is processed independently — partial failures are returned - in `failed_updates` rather than aborting the whole batch. - - Example request: - ```bash - curl --location 'http://0.0.0.0:4000/team/key/bulk_update' \\ - --header 'Authorization: Bearer sk-1234' \\ - --header 'Content-Type: application/json' \\ - --data '{ - "team_id": "team-123", - "all_keys_in_team": true, - "update_fields": {"max_budget": 50.0, "budget_duration": "30d"} - }' - ``` + Callable by proxy admins, or by team admins with `KEY_UPDATE` permission. """ from litellm.proxy.proxy_server import ( llm_router, @@ -2760,10 +2744,19 @@ async def bulk_update_team_keys( }, ) - # Resolve which keys to update (single batched query, scoped to team). if data.all_keys_in_team: + # "all" excludes blocked/expired — bulk refresh shouldn't revive a key an admin disabled. + # `blocked` is Boolean? with no default; `/key/generate` writes NULL. Prisma's `NOT` + # excludes NULLs, so explicitly OR `false` with `null` to include them. + now = datetime.now(timezone.utc) existing_keys = await prisma_client.db.litellm_verificationtoken.find_many( - where={"team_id": data.team_id}, + where={ + "team_id": data.team_id, + "AND": [ + {"OR": [{"blocked": False}, {"blocked": None}]}, + {"OR": [{"expires": None}, {"expires": {"gt": now}}]}, + ], + }, order={"token": "asc"}, take=MAX_BATCH_SIZE + 1, ) @@ -2776,12 +2769,52 @@ async def bulk_update_team_keys( ) requested_tokens = [row.token for row in existing_keys] else: - # validator guarantees key_ids is set and non-empty when all_keys_in_team=False - assert data.key_ids is not None + if data.key_ids is None or len(data.key_ids) == 0: + raise HTTPException( + status_code=400, + detail={ + "error": "key_ids must be provided when all_keys_in_team is False" + }, + ) + # Dedupe by hashed form — duplicates collapse to one update. + requested_tokens = [] + hashed_key_ids = [] + seen_hashes = set() + for k in data.key_ids: + h = _hash_token_if_needed(k) + if h in seen_hashes: + continue + seen_hashes.add(h) + requested_tokens.append(k) + hashed_key_ids.append(h) existing_keys = await prisma_client.db.litellm_verificationtoken.find_many( - where={"team_id": data.team_id, "token": {"in": data.key_ids}} + where={"team_id": data.team_id, "token": {"in": hashed_key_ids}} ) - requested_tokens = list(data.key_ids) + + # Anchor membership check on data.team_id (not existing_keys[0]); empty result must still gate non-admins. + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: + auth_anchor = ( + existing_keys[0] + if existing_keys + else LiteLLM_VerificationToken( + token="__team_scope_auth_check__", + team_id=data.team_id, + models=[], + ) + ) + await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( + user_api_key_dict=user_api_key_dict, + route=KeyManagementRoutes.KEY_UPDATE, + prisma_client=prisma_client, + existing_key_row=auth_anchor, + user_api_key_cache=user_api_key_cache, + ) + + # Block metadata.allowed_passthrough_routes for non-admins — the runtime + # route checker reads it from key/team metadata to grant passthrough. + _check_passthrough_routes_caller_permission( + data=data.update_fields, user_api_key_dict=user_api_key_dict + ) if not requested_tokens: raise HTTPException( @@ -2789,22 +2822,6 @@ async def bulk_update_team_keys( detail={"error": f"No keys found for team {data.team_id}"}, ) - # Fail-fast auth: if caller is not proxy admin, verify they have KEY_UPDATE - # permission for this team once (using any in-team key as the auth subject). - # Per-key checks inside _process_single_key_update will short-circuit on the - # cached team object after this. - if ( - user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value - and existing_keys - ): - await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint( - user_api_key_dict=user_api_key_dict, - route=KeyManagementRoutes.KEY_UPDATE, - prisma_client=prisma_client, - existing_key_row=existing_keys[0], - user_api_key_cache=user_api_key_cache, - ) - existing_by_token = {row.token: row for row in existing_keys} update_field_dict = data.update_fields.model_dump(exclude_unset=True) @@ -2812,15 +2829,18 @@ async def bulk_update_team_keys( failed_updates: List[FailedKeyUpdate] = [] for token in requested_tokens: + db_token = _hash_token_if_needed(token) try: - if token not in existing_by_token: + if db_token not in existing_by_token: raise HTTPException( status_code=404, detail={"error": f"Key not found in team {data.team_id}"}, ) + # team_id from validated scope, never user payload — drives _check_team_key_limits. update_key_request = UpdateKeyRequest( key=token, + team_id=data.team_id, **update_field_dict, ) updated_key_info = await _process_single_key_update( @@ -2839,14 +2859,15 @@ async def bulk_update_team_keys( ) except Exception as e: + # Log the hashed prefix — `token` may be a raw sk-... and ERROR logs persist. verbose_proxy_logger.exception( - f"Failed to update key {token} in team {data.team_id}: {e}" + f"Failed to update key {db_token[:12]}... in team {data.team_id}: {e}" ) failed_updates.append( _build_failed_team_key_update( token=token, exception=e, - existing_key_row=existing_by_token.get(token), + existing_key_row=existing_by_token.get(db_token), ) ) diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 55e7a135533..d214cdb4f5d 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -1,9 +1,7 @@ from datetime import datetime -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Literal, Optional -from pydantic import BaseModel, model_validator - -from litellm.proxy._types import KeyRequestBase +from pydantic import BaseModel, ConfigDict, model_validator class BulkUpdateKeyRequestItem(BaseModel): @@ -45,21 +43,41 @@ class BulkUpdateKeyResponse(BaseModel): failed_updates: List[FailedKeyUpdate] -class KeyUpdateFields(KeyRequestBase): - """ - Mirror of UpdateKeyRequest minus per-key identifiers (`key`, `key_alias`) - and the scope guard (`team_id`). Used as the broadcast payload in - BulkUpdateTeamKeysRequest — one set of fields applied to many keys. - """ +class KeyUpdateFields(BaseModel): + """Allowlist of bulk-broadcastable fields for /team/key/bulk_update; `extra="forbid"` blocks RBAC/ownership/scope mutations even by team admins.""" - duration: Optional[str] = None - spend: Optional[float] = None - metadata: Optional[dict] = None + model_config = ConfigDict(extra="forbid", protected_namespaces=()) + + # Budgets + max_budget: Optional[float] = None + budget_id: Optional[str] = None + budget_duration: Optional[str] = None + budget_limits: Optional[List[Any]] = None + model_max_budget: Optional[Dict[str, Any]] = None + + # Rate limits + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + model_tpm_limit: Optional[Dict[str, Any]] = None + model_rpm_limit: Optional[Dict[str, Any]] = None + max_parallel_requests: Optional[int] = None + rpm_limit_type: Optional[ + Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] + ] = None + tpm_limit_type: Optional[ + Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] + ] = None + + # Temporary budget grants (auto-expire). `spend` deliberately omitted — bulk-zeroing it bypasses budget enforcement; admin-only via /key/update. temp_budget_increase: Optional[float] = None temp_budget_expiry: Optional[datetime] = None - auto_rotate: Optional[bool] = None - rotation_interval: Optional[str] = None - organization_id: Optional[str] = None + + # Expiry + duration: Optional[str] = None + + # Operational metadata + tags: Optional[List[str]] = None + metadata: Optional[Dict[str, Any]] = None @model_validator(mode="after") def validate_temp_budget(self) -> "KeyUpdateFields": @@ -71,28 +89,15 @@ class KeyUpdateFields(KeyRequestBase): return self @model_validator(mode="after") - def reject_per_key_or_scope_fields(self) -> "KeyUpdateFields": - forbidden = [ - name - for name in ("key", "key_alias", "team_id") - if getattr(self, name, None) is not None - ] - if forbidden: - raise ValueError( - f"Fields not allowed in update_fields for bulk team key updates: " - f"{forbidden}. `key`/`key_alias` are per-key identifiers; `team_id` " - f"is the scope guard set at the request top level." - ) + def require_at_least_one_field(self) -> "KeyUpdateFields": + # Reject empty payload — would iterate every key with no-op writes. + if not self.model_fields_set: + raise ValueError("update_fields must specify at least one field to update.") return self class BulkUpdateTeamKeysRequest(BaseModel): - """ - Request for applying one update payload to many keys inside a single team. - - Exactly one of `key_ids` (specific keys in the team) or `all_keys_in_team` - (every key in the team) must be provided. - """ + """Apply one update payload to many keys inside a team; provide either `key_ids` or `all_keys_in_team=True`.""" team_id: str key_ids: Optional[List[str]] = None 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 2a3719e3770..653cf161e06 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 @@ -10023,6 +10023,9 @@ async def test_execute_virtual_key_regeneration_cache_invalidation_with_token_ha # --------------------------------------------------------------------------- +_BULK_PKG = "litellm.proxy.management_endpoints.key_management_endpoints" + + def _make_team_key(token: str, team_id: str = "team-abc") -> LiteLLM_VerificationToken: return LiteLLM_VerificationToken( token=token, @@ -10033,222 +10036,228 @@ def _make_team_key(token: str, team_id: str = "team-abc") -> LiteLLM_Verificatio ) -def _patch_team_keys_helpers(monkeypatch, *, hash_lookup): - """Common monkeypatching for the bulk_update_team_keys path.""" +def _admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin" + ) + + +def _internal_user() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-iu", user_id="iu" + ) + + +def _updated(payload): + m = MagicMock() + m.model_dump.return_value = payload + return m + + +def _setup_team_keys_mocks( + monkeypatch, + *, + find_many=None, + find_unique=None, + update_data=None, + hash_identity=True, +): + """Set up mocks for bulk_update_team_keys; returns mock_prisma.""" + mock_prisma = AsyncMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[] if find_many is None else find_many + ) + if find_unique is not None: + mock_prisma.db.litellm_verificationtoken.find_unique = find_unique + if update_data is not None: + mock_prisma.update_data = update_data + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None) monkeypatch.setattr( - "litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data", + f"{_BULK_PKG}.prepare_key_update_data", AsyncMock(return_value={"max_budget": 50.0}), ) + monkeypatch.setattr(f"{_BULK_PKG}._delete_cache_key_object", AsyncMock()) monkeypatch.setattr( - "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed", - lambda token: hash_lookup[token], + f"{_BULK_PKG}.KeyManagementEventHooks.async_key_updated_hook", AsyncMock() ) + monkeypatch.setattr(f"{_BULK_PKG}.get_team_object", AsyncMock(return_value=None)) + monkeypatch.setattr(f"{_BULK_PKG}._check_team_key_limits", AsyncMock()) monkeypatch.setattr( - "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint", AsyncMock(), ) - monkeypatch.setattr( - "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook", - AsyncMock(), + if hash_identity: + # Tests use already-hashed tokens; the raw-sk regression opts out. + monkeypatch.setattr(f"{_BULK_PKG}._hash_token_if_needed", lambda token: token) + return mock_prisma + + +async def _call_as_admin(data): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + bulk_update_team_keys, ) + return await bulk_update_team_keys( + data=data, user_api_key_dict=_admin(), litellm_changed_by=None + ) + + +# ---- happy paths ---------------------------------------------------------- + @pytest.mark.asyncio async def test_bulk_update_team_keys_success_with_key_ids(monkeypatch): - """Proxy admin updates two explicitly listed keys; both succeed.""" - from litellm.proxy.management_endpoints.key_management_endpoints import ( - bulk_update_team_keys, - ) from litellm.types.proxy.management_endpoints.key_management_endpoints import ( BulkUpdateTeamKeysRequest, KeyUpdateFields, ) - key_a = _make_team_key("tok-a") - key_b = _make_team_key("tok-b") - - mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[key_a, key_b] - ) - mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( - side_effect=[key_a, key_b] - ) - updated_a = MagicMock() - updated_a.model_dump.return_value = {"max_budget": 50.0, "user_id": "user-123"} - updated_b = MagicMock() - updated_b.model_dump.return_value = {"max_budget": 50.0, "user_id": "user-123"} - mock_prisma_client.update_data = AsyncMock( - side_effect=[{"data": updated_a}, {"data": updated_b}] - ) - - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None) - _patch_team_keys_helpers( - monkeypatch, hash_lookup={"tok-a": "hashed-a", "tok-b": "hashed-b"} - ) - monkeypatch.setattr( - "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint", - AsyncMock(), - ) - - request = BulkUpdateTeamKeysRequest( - team_id="team-abc", - key_ids=["tok-a", "tok-b"], - update_fields=KeyUpdateFields(max_budget=50.0), - ) - response = await bulk_update_team_keys( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key="sk-admin", - user_id="admin", + keys = [_make_team_key("tok-a"), _make_team_key("tok-b")] + mock = _setup_team_keys_mocks( + monkeypatch, + find_many=keys, + find_unique=AsyncMock(side_effect=keys), + update_data=AsyncMock( + side_effect=[{"data": _updated({"max_budget": 50.0})}] * 2 ), - litellm_changed_by=None, ) - assert response.total_requested == 2 + response = await _call_as_admin( + BulkUpdateTeamKeysRequest( + team_id="team-abc", + key_ids=["tok-a", "tok-b"], + update_fields=KeyUpdateFields(max_budget=50.0), + ) + ) + assert len(response.successful_updates) == 2 assert len(response.failed_updates) == 0 - assert {u.key for u in response.successful_updates} == {"tok-a", "tok-b"} - - mock_prisma_client.db.litellm_verificationtoken.find_many.assert_awaited_once() - where_arg = ( - mock_prisma_client.db.litellm_verificationtoken.find_many.await_args.kwargs[ - "where" - ] - ) - assert where_arg["team_id"] == "team-abc" - assert where_arg["token"] == {"in": ["tok-a", "tok-b"]} + where = mock.db.litellm_verificationtoken.find_many.await_args.kwargs["where"] + assert where["team_id"] == "team-abc" + assert where["token"] == {"in": ["tok-a", "tok-b"]} @pytest.mark.asyncio async def test_bulk_update_team_keys_success_all_keys_in_team(monkeypatch): - """`all_keys_in_team=True` resolves via find_many and updates every key.""" - from litellm.proxy.management_endpoints.key_management_endpoints import ( - bulk_update_team_keys, - ) from litellm.types.proxy.management_endpoints.key_management_endpoints import ( BulkUpdateTeamKeysRequest, KeyUpdateFields, ) keys = [_make_team_key(f"tok-{i}") for i in range(3)] - - mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=keys - ) - mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( - side_effect=keys - ) - updated_objs = [MagicMock() for _ in keys] - for obj in updated_objs: - obj.model_dump.return_value = {"max_budget": 50.0} - mock_prisma_client.update_data = AsyncMock( - side_effect=[{"data": obj} for obj in updated_objs] - ) - - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None) - _patch_team_keys_helpers( - monkeypatch, hash_lookup={f"tok-{i}": f"hashed-{i}" for i in range(3)} - ) - monkeypatch.setattr( - "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint", - AsyncMock(), - ) - - request = BulkUpdateTeamKeysRequest( - team_id="team-abc", - all_keys_in_team=True, - update_fields=KeyUpdateFields(max_budget=50.0), - ) - response = await bulk_update_team_keys( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key="sk-admin", - user_id="admin", + mock = _setup_team_keys_mocks( + monkeypatch, + find_many=keys, + find_unique=AsyncMock(side_effect=keys), + update_data=AsyncMock( + side_effect=[{"data": _updated({"max_budget": 50.0})}] * 3 ), - litellm_changed_by=None, ) - assert response.total_requested == 3 - assert len(response.successful_updates) == 3 - assert len(response.failed_updates) == 0 - where_kwargs = ( - mock_prisma_client.db.litellm_verificationtoken.find_many.await_args.kwargs + response = await _call_as_admin( + BulkUpdateTeamKeysRequest( + team_id="team-abc", + all_keys_in_team=True, + update_fields=KeyUpdateFields(max_budget=50.0), + ) + ) + + assert len(response.successful_updates) == 3 + where = mock.db.litellm_verificationtoken.find_many.await_args.kwargs["where"] + # `blocked` is Boolean? with no default → /key/generate writes NULL. Prisma's + # NOT excludes NULLs, so the filter has to OR `false` with `null` explicitly. + blocked_or, expires_or = where["AND"][0]["OR"], where["AND"][1]["OR"] + assert {"blocked": False} in blocked_or and {"blocked": None} in blocked_or + assert {"expires": None} in expires_or + assert any( + "gt" in c.get("expires", {}) + for c in expires_or + if isinstance(c.get("expires"), dict) ) - assert where_kwargs["take"] == 501 - assert where_kwargs["where"] == {"team_id": "team-abc"} @pytest.mark.asyncio async def test_bulk_update_team_keys_key_not_in_team(monkeypatch): - """key_ids includes a token not in the team — that token fails, others succeed.""" - from litellm.proxy.management_endpoints.key_management_endpoints import ( - bulk_update_team_keys, - ) from litellm.types.proxy.management_endpoints.key_management_endpoints import ( BulkUpdateTeamKeysRequest, KeyUpdateFields, ) in_team = _make_team_key("tok-a") - - mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[in_team] - ) - mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( - return_value=in_team - ) - updated = MagicMock() - updated.model_dump.return_value = {"max_budget": 50.0} - mock_prisma_client.update_data = AsyncMock(return_value={"data": updated}) - - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None) - _patch_team_keys_helpers(monkeypatch, hash_lookup={"tok-a": "hashed-a"}) - monkeypatch.setattr( - "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint", - AsyncMock(), + _setup_team_keys_mocks( + monkeypatch, + find_many=[in_team], + find_unique=AsyncMock(return_value=in_team), + update_data=AsyncMock(return_value={"data": _updated({"max_budget": 50.0})}), ) - request = BulkUpdateTeamKeysRequest( - team_id="team-abc", - key_ids=["tok-a", "tok-foreign"], - update_fields=KeyUpdateFields(max_budget=50.0), + response = await _call_as_admin( + BulkUpdateTeamKeysRequest( + team_id="team-abc", + key_ids=["tok-a", "tok-foreign"], + update_fields=KeyUpdateFields(max_budget=50.0), + ) ) - response = await bulk_update_team_keys( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key="sk-admin", - user_id="admin", - ), - litellm_changed_by=None, - ) - - assert response.total_requested == 2 assert [u.key for u in response.successful_updates] == ["tok-a"] assert [u.key for u in response.failed_updates] == ["tok-foreign"] assert "not found in team" in response.failed_updates[0].failed_reason +# ---- error paths ---------------------------------------------------------- + + @pytest.mark.asyncio -async def test_bulk_update_team_keys_team_admin_authorized(monkeypatch): - """Non-admin caller passes the upfront team-permission check and succeeds.""" +async def test_bulk_update_team_keys_batch_size_cap(monkeypatch): + from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateTeamKeysRequest, + KeyUpdateFields, + ) + + _setup_team_keys_mocks( + monkeypatch, + find_many=[_make_team_key(f"tok-{i}") for i in range(501)], + ) + + with pytest.raises(HTTPException) as exc: + await _call_as_admin( + BulkUpdateTeamKeysRequest( + team_id="team-abc", + all_keys_in_team=True, + update_fields=KeyUpdateFields(max_budget=50.0), + ) + ) + assert exc.value.status_code == 400 + assert "more than 500" in exc.value.detail["error"] + + +@pytest.mark.asyncio +async def test_bulk_update_team_keys_empty_team_returns_404(monkeypatch): + from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateTeamKeysRequest, + KeyUpdateFields, + ) + + _setup_team_keys_mocks(monkeypatch, find_many=[]) + with pytest.raises(HTTPException) as exc: + await _call_as_admin( + BulkUpdateTeamKeysRequest( + team_id="team-empty", + all_keys_in_team=True, + update_fields=KeyUpdateFields(max_budget=50.0), + ) + ) + assert exc.value.status_code == 404 + + +# ---- auth ----------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_bulk_update_team_keys_team_member_with_permission(monkeypatch): from litellm.proxy.management_endpoints.key_management_endpoints import ( bulk_update_team_keys, ) @@ -10258,53 +10267,34 @@ async def test_bulk_update_team_keys_team_admin_authorized(monkeypatch): ) key_a = _make_team_key("tok-a") - mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[key_a] + _setup_team_keys_mocks( + monkeypatch, + find_many=[key_a], + find_unique=AsyncMock(return_value=key_a), + update_data=AsyncMock(return_value={"data": _updated({"max_budget": 50.0})}), ) - mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( - return_value=key_a - ) - updated = MagicMock() - updated.model_dump.return_value = {"max_budget": 50.0} - mock_prisma_client.update_data = AsyncMock(return_value={"data": updated}) - - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None) - _patch_team_keys_helpers(monkeypatch, hash_lookup={"tok-a": "hashed-a"}) - auth_check = AsyncMock() monkeypatch.setattr( - "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint", + f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint", auth_check, ) - request = BulkUpdateTeamKeysRequest( - team_id="team-abc", - all_keys_in_team=True, - update_fields=KeyUpdateFields(max_budget=50.0), - ) response = await bulk_update_team_keys( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, - api_key="sk-team-admin", - user_id="team-admin-user", + data=BulkUpdateTeamKeysRequest( + team_id="team-abc", + all_keys_in_team=True, + update_fields=KeyUpdateFields(max_budget=50.0), ), + user_api_key_dict=_internal_user(), litellm_changed_by=None, ) - assert len(response.successful_updates) == 1 - # upfront check (1) + per-key check inside _process_single_key_update (1) + # Upfront check + per-key check inside _process_single_key_update assert auth_check.await_count == 2 @pytest.mark.asyncio -async def test_bulk_update_team_keys_team_admin_no_permission(monkeypatch): - """Non-admin without KEY_UPDATE permission raises ProxyException upfront (no partial result).""" +async def test_bulk_update_team_keys_team_member_no_permission(monkeypatch): from litellm.proxy.management_endpoints.key_management_endpoints import ( bulk_update_team_keys, ) @@ -10313,20 +10303,9 @@ async def test_bulk_update_team_keys_team_admin_no_permission(monkeypatch): KeyUpdateFields, ) - key_a = _make_team_key("tok-a") - mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[key_a] - ) - - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None) - + mock = _setup_team_keys_mocks(monkeypatch, find_many=[_make_team_key("tok-a")]) monkeypatch.setattr( - "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint", + f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint", AsyncMock( side_effect=ProxyException( message="not in team", @@ -10337,67 +10316,255 @@ async def test_bulk_update_team_keys_team_admin_no_permission(monkeypatch): ), ) - request = BulkUpdateTeamKeysRequest( - team_id="team-abc", - all_keys_in_team=True, - update_fields=KeyUpdateFields(max_budget=50.0), + with pytest.raises(ProxyException): + await bulk_update_team_keys( + data=BulkUpdateTeamKeysRequest( + team_id="team-abc", + all_keys_in_team=True, + update_fields=KeyUpdateFields(max_budget=1.0), + ), + user_api_key_dict=_internal_user(), + litellm_changed_by=None, + ) + mock.update_data.assert_not_called() + + +# ---- pydantic-layer validation ------------------------------------------- + + +def test_bulk_update_team_keys_request_validation(): + """Allowlist (extra='forbid'), empty-payload rejection, and selection XOR.""" + from pydantic import ValidationError + + from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateTeamKeysRequest, + KeyUpdateFields, + ) + + forbidden = [ + "key", + "key_alias", + "team_id", + "allowed_routes", + "allowed_passthrough_routes", + "permissions", + "object_permission", + "access_group_ids", + "user_id", + "organization_id", + "blocked", + "key_type", + "models", + "config", + "router_settings", + "spend", + ] + for f in forbidden: + with pytest.raises(ValidationError, match=f): + KeyUpdateFields(**{f: True}) + + with pytest.raises(ValidationError, match="at least one"): + KeyUpdateFields() + + assert KeyUpdateFields(max_budget=50.0, tags=["x"]).max_budget == 50.0 + + valid = KeyUpdateFields(max_budget=10) + with pytest.raises(ValidationError): + BulkUpdateTeamKeysRequest( + team_id="t", key_ids=["k"], all_keys_in_team=True, update_fields=valid + ) + with pytest.raises(ValidationError): + BulkUpdateTeamKeysRequest(team_id="t", update_fields=valid) + + +# ---- security regressions ------------------------------------------------ + + +@pytest.mark.asyncio +async def test_bulk_update_team_keys_hashes_raw_sk_key_ids(monkeypatch): + """Regression: raw sk-... key_ids must be hashed before the find_many lookup.""" + from litellm.proxy._types import hash_token + from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateTeamKeysRequest, + KeyUpdateFields, + ) + + raw_sk = "sk-rawkey1234567890" + hashed = hash_token(raw_sk) + row = LiteLLM_VerificationToken( + token=hashed, user_id="u", models=[], team_id="team-abc", max_budget=None + ) + mock = _setup_team_keys_mocks( + monkeypatch, + find_many=[row], + find_unique=AsyncMock(return_value=row), + update_data=AsyncMock(return_value={"data": _updated({"max_budget": 50.0})}), + hash_identity=False, + ) + + response = await _call_as_admin( + BulkUpdateTeamKeysRequest( + team_id="team-abc", + key_ids=[raw_sk], + update_fields=KeyUpdateFields(max_budget=50.0), + ) + ) + where = mock.db.litellm_verificationtoken.find_many.await_args.kwargs["where"] + assert where["token"] == {"in": [hashed]} + # Response reports the user-supplied form, not the hash. + assert response.successful_updates[0].key == raw_sk + + +@pytest.mark.asyncio +async def test_bulk_update_team_keys_auth_check_runs_when_no_keys_match(monkeypatch): + """Regression: non-admin with bogus key_ids must still hit the membership gate.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + bulk_update_team_keys, + ) + from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateTeamKeysRequest, + KeyUpdateFields, + ) + + mock = _setup_team_keys_mocks(monkeypatch, find_many=[]) + auth_check = AsyncMock( + side_effect=ProxyException( + message="not in team", + type="team_member_permission_error", + param="/key/update", + code=401, + ) + ) + monkeypatch.setattr( + f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint", + auth_check, ) with pytest.raises(ProxyException): await bulk_update_team_keys( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, - api_key="sk-outsider", - user_id="outsider", + data=BulkUpdateTeamKeysRequest( + team_id="victim-team", + key_ids=["bogus-1", "bogus-2"], + update_fields=KeyUpdateFields(max_budget=1.0), ), + user_api_key_dict=_internal_user(), litellm_changed_by=None, ) - - mock_prisma_client.update_data.assert_not_called() + # Anchored on data.team_id, not existing_keys[0]. + assert auth_check.await_args.kwargs["existing_key_row"].team_id == "victim-team" + mock.update_data.assert_not_called() @pytest.mark.asyncio -async def test_bulk_update_team_keys_xor_validation(): - """Both selection modes set, or neither set, raises at request construction.""" +async def test_bulk_update_team_keys_does_not_log_raw_sk_token_on_failure( + monkeypatch, caplog +): + """Regression: per-key failure must not log the raw sk-... (ERROR-level logs persist).""" + import logging + + from litellm.proxy._types import hash_token from litellm.types.proxy.management_endpoints.key_management_endpoints import ( BulkUpdateTeamKeysRequest, KeyUpdateFields, ) - with pytest.raises(Exception): - BulkUpdateTeamKeysRequest( - team_id="t", - key_ids=["k"], - all_keys_in_team=True, - update_fields=KeyUpdateFields(max_budget=10.0), - ) + raw_sk = "sk-supersecret1234567890" + row = LiteLLM_VerificationToken( + token=hash_token(raw_sk), + user_id="u", + models=[], + team_id="team-abc", + max_budget=None, + ) + _setup_team_keys_mocks( + monkeypatch, + find_many=[row], + # Force the per-key fetch to raise so the exception logger fires. + find_unique=AsyncMock(side_effect=RuntimeError("boom")), + hash_identity=False, + ) - with pytest.raises(Exception): - BulkUpdateTeamKeysRequest( - team_id="t", - update_fields=KeyUpdateFields(max_budget=10.0), + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + response = await _call_as_admin( + BulkUpdateTeamKeysRequest( + team_id="team-abc", + key_ids=[raw_sk], + update_fields=KeyUpdateFields(max_budget=50.0), + ) ) + assert len(response.failed_updates) == 1 + log_text = "\n".join(r.getMessage() for r in caplog.records) + assert raw_sk not in log_text @pytest.mark.asyncio -async def test_bulk_update_team_keys_forbidden_fields(): - """update_fields.team_id / .key / .key_alias are rejected on construction.""" +async def test_bulk_update_team_keys_propagates_team_id_to_per_key_request(monkeypatch): + """Regression: per-key UpdateKeyRequest carries data.team_id (gates _check_team_key_limits).""" from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateTeamKeysRequest, KeyUpdateFields, ) - with pytest.raises(Exception): - KeyUpdateFields(team_id="other-team") - with pytest.raises(Exception): - KeyUpdateFields(key="sk-abc") - with pytest.raises(Exception): - KeyUpdateFields(key_alias="alias") + _setup_team_keys_mocks(monkeypatch, find_many=[_make_team_key("tok-a")]) + captured = [] + + async def fake_process(*, update_key_request, **kw): + captured.append(update_key_request) + return {"max_budget": update_key_request.max_budget} + + monkeypatch.setattr(f"{_BULK_PKG}._process_single_key_update", fake_process) + + await _call_as_admin( + BulkUpdateTeamKeysRequest( + team_id="team-abc", + key_ids=["tok-a"], + update_fields=KeyUpdateFields( + tpm_limit=10_000, tpm_limit_type="guaranteed_throughput" + ), + ) + ) + assert captured[0].team_id == "team-abc" + assert captured[0].tpm_limit_type == "guaranteed_throughput" @pytest.mark.asyncio -async def test_bulk_update_team_keys_batch_size_cap(monkeypatch): - """find_many returning > MAX_BATCH_SIZE rows raises 400.""" +async def test_bulk_update_team_keys_dedupes_key_ids(monkeypatch): + """Duplicate key_ids collapse to a single update (no redundant DB writes, no inflated counts).""" + from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateTeamKeysRequest, + KeyUpdateFields, + ) + + key_a = _make_team_key("tok-a") + update_data = AsyncMock(return_value={"data": _updated({"max_budget": 50.0})}) + _setup_team_keys_mocks( + monkeypatch, + find_many=[key_a], + find_unique=AsyncMock(return_value=key_a), + update_data=update_data, + ) + + response = await _call_as_admin( + BulkUpdateTeamKeysRequest( + team_id="team-abc", + key_ids=["tok-a", "tok-a", "tok-a"], + update_fields=KeyUpdateFields(max_budget=50.0), + ) + ) + + assert response.total_requested == 1 + assert len(response.successful_updates) == 1 + assert len(response.failed_updates) == 0 + update_data.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_bulk_update_team_keys_blocks_metadata_allowed_passthrough_routes( + monkeypatch, +): + """Non-admin can't grant passthrough access by smuggling allowed_passthrough_routes through metadata.""" + from fastapi import HTTPException + from litellm.proxy.management_endpoints.key_management_endpoints import ( bulk_update_team_keys, ) @@ -10406,76 +10573,23 @@ async def test_bulk_update_team_keys_batch_size_cap(monkeypatch): KeyUpdateFields, ) - too_many = [_make_team_key(f"tok-{i}") for i in range(501)] - mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=too_many - ) - - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None) + mock = _setup_team_keys_mocks(monkeypatch, find_many=[_make_team_key("tok-a")]) request = BulkUpdateTeamKeysRequest( team_id="team-abc", all_keys_in_team=True, - update_fields=KeyUpdateFields(max_budget=50.0), + update_fields=KeyUpdateFields( + metadata={"allowed_passthrough_routes": ["/admin/*"]} + ), ) with pytest.raises(HTTPException) as exc: await bulk_update_team_keys( data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key="sk-admin", - user_id="admin", - ), + user_api_key_dict=_internal_user(), litellm_changed_by=None, ) - assert exc.value.status_code == 400 - assert "more than 500" in exc.value.detail["error"] - - -@pytest.mark.asyncio -async def test_bulk_update_team_keys_no_keys_found(monkeypatch): - """find_many returns nothing → 404 (clear signal, not silent empty success).""" - from litellm.proxy.management_endpoints.key_management_endpoints import ( - bulk_update_team_keys, - ) - from litellm.types.proxy.management_endpoints.key_management_endpoints import ( - BulkUpdateTeamKeysRequest, - KeyUpdateFields, - ) - - mock_prisma_client = AsyncMock() - mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[] - ) - - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock()) - monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None) - - request = BulkUpdateTeamKeysRequest( - team_id="team-empty", - all_keys_in_team=True, - update_fields=KeyUpdateFields(max_budget=50.0), - ) - - with pytest.raises(HTTPException) as exc: - await bulk_update_team_keys( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key="sk-admin", - user_id="admin", - ), - litellm_changed_by=None, - ) - - assert exc.value.status_code == 404 + assert exc.value.status_code == 403 + assert "allowed_passthrough_routes" in str(exc.value.detail) + mock.update_data.assert_not_called()