scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

[Feat] Add endpoint for bulk key updates for team (#26468)

Browse source 
Squash-merged by litellm-agent from Michael-RZ-Berri's PR.
This commit is contained in:
Michael-RZ-Berri 2026-05-09 12:32:16 -07:00 • committed by GitHub
parent 9380940ced
commit b834817785
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 905 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/_types.py
View file

@ -239,6 +239,7 @@ class KeyManagementRoutes(str, enum.Enum):
KEY_BLOCK = "/key/block"
KEY_UNBLOCK = "/key/unblock"
KEY_BULK_UPDATE = "/key/bulk_update"
TEAM_KEY_BULK_UPDATE = "/team/key/bulk_update"
KEY_RESET_SPEND = "/key/{key_id}/reset_spend"
# info and health routes
@ -540,6 +541,7 @@ class LiteLLMRoutes(enum.Enum):
KeyManagementRoutes.KEY_BLOCK.value,
KeyManagementRoutes.KEY_UNBLOCK.value,
KeyManagementRoutes.KEY_BULK_UPDATE.value,
KeyManagementRoutes.TEAM_KEY_BULK_UPDATE.value,
KeyManagementRoutes.TEAM_DAILY_ACTIVITY.value,
KeyManagementRoutes.SPEND_LOGS.value,
KeyManagementRoutes.KEY_RESET_SPEND.value,

2
litellm/proxy/auth/route_checks.py
View file

@ -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",
]
)

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

@ -88,8 +88,8 @@ from litellm.router import Router
from litellm.secret_managers.main import get_secret
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateKeyRequest,
BulkUpdateKeyRequestItem,
BulkUpdateKeyResponse,
BulkUpdateTeamKeysRequest,
FailedKeyUpdate,
SuccessfulKeyUpdate,
)
@ -1881,7 +1881,7 @@ async def _get_and_validate_existing_key(
async def _process_single_key_update(
key_update_item: BulkUpdateKeyRequestItem,
update_key_request: UpdateKeyRequest,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str],
prisma_client: Optional[PrismaClient],
@ -1889,6 +1889,7 @@ async def _process_single_key_update(
proxy_logging_obj: Any,
llm_router: Optional[Router],
user_custom_key_update: Optional[Callable] = None,
existing_key_row: Optional[LiteLLM_VerificationToken] = None,
) -> Dict[str, Any]:
"""
Process a single key update with all validations and checks.
@ -1897,13 +1898,14 @@ async def _process_single_key_update(
including validation, permission checks, team checks, and database updates.
Args:
key_update_item: The key update request item
update_key_request: Fully-constructed UpdateKeyRequest for the target key
user_api_key_dict: The authenticated user's API key info
litellm_changed_by: Optional header for tracking who made the change
prisma_client: Prisma client instance
user_api_key_cache: User API key cache
proxy_logging_obj: Proxy logging object
llm_router: LLM router instance
existing_key_row: Optional pre-fetched key row to avoid redundant lookups
Returns:
Dict containing the updated key information
@ -1912,13 +1914,14 @@ async def _process_single_key_update(
HTTPException: For various validation and permission errors
"""
# Validate max_budget
_validate_max_budget(key_update_item.max_budget)
_validate_max_budget(update_key_request.max_budget)
# Get and validate existing key
existing_key_row = await _get_and_validate_existing_key(
token=key_update_item.key,
prisma_client=prisma_client,
)
if existing_key_row is None:
existing_key_row = await _get_and_validate_existing_key(
token=update_key_request.key,
prisma_client=prisma_client,
)
# Check team member permissions
if prisma_client is not None:
@ -1930,15 +1933,6 @@ async def _process_single_key_update(
user_api_key_cache=user_api_key_cache,
)
# Create UpdateKeyRequest from BulkUpdateKeyRequestItem
update_key_request = UpdateKeyRequest(
key=key_update_item.key,
budget_id=key_update_item.budget_id,
max_budget=key_update_item.max_budget,
team_id=key_update_item.team_id,
tags=key_update_item.tags,
)
# Custom key update hook
if user_custom_key_update is not None:
if inspect.iscoroutinefunction(user_custom_key_update):
@ -2003,12 +1997,12 @@ async def _process_single_key_update(
detail={"error": "Database not connected"},
)
_data = {**non_default_values, "token": key_update_item.key}
response = await prisma_client.update_data(token=key_update_item.key, data=_data)
_data = {**non_default_values, "token": update_key_request.key}
response = await prisma_client.update_data(token=update_key_request.key, data=_data)
# Delete cache
await _delete_cache_key_object(
hashed_token=_hash_token_if_needed(key_update_item.key),
hashed_token=_hash_token_if_needed(update_key_request.key),
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
@ -2598,9 +2592,15 @@ async def bulk_update_keys(
for key_update_item in data.keys:
try:
# Process single key update using reusable function
update_key_request = UpdateKeyRequest(
key=key_update_item.key,
budget_id=key_update_item.budget_id,
max_budget=key_update_item.max_budget,
team_id=key_update_item.team_id,
tags=key_update_item.tags,
)
updated_key_info = await _process_single_key_update(
key_update_item=key_update_item,
update_key_request=update_key_request,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
prisma_client=prisma_client,
@ -2665,6 +2665,223 @@ async def bulk_update_keys(
)
def _build_failed_team_key_update(
token: str,
exception: Exception,
existing_key_row: Optional[LiteLLM_VerificationToken],
) -> FailedKeyUpdate:
"""Normalize an exception from the per-key update loop into a FailedKeyUpdate."""
if isinstance(exception, HTTPException):
detail = exception.detail
if isinstance(detail, dict):
error_message = detail.get("error", str(exception))
else:
error_message = str(detail)
elif isinstance(exception, ProxyException):
error_message = exception.message
else:
error_message = str(exception)
key_info: Optional[Dict[str, Any]] = None
if existing_key_row is not None:
if hasattr(existing_key_row, "model_dump"):
key_info = existing_key_row.model_dump()
elif hasattr(existing_key_row, "dict"):
key_info = existing_key_row.dict()
if key_info:
key_info.pop("token", None)
return FailedKeyUpdate(key=token, key_info=key_info, failed_reason=error_message)
@router.post(
"/team/key/bulk_update",
tags=["key management"],
dependencies=[Depends(user_api_key_auth)],
response_model=BulkUpdateKeyResponse,
)
@management_endpoint_wrapper
async def bulk_update_team_keys(
data: BulkUpdateTeamKeysRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
litellm_changed_by: Optional[str] = Header(
None,
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
),
):
"""
Apply one update payload to many keys inside a single team.
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.
"""
from litellm.proxy.proxy_server import (
llm_router,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
user_custom_key_update,
)
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": "Database not connected"},
)
if not data.team_id:
raise HTTPException(
status_code=400,
detail={"error": "team_id is required"},
)
MAX_BATCH_SIZE = 500
if data.key_ids is not None and len(data.key_ids) > MAX_BATCH_SIZE:
raise HTTPException(
status_code=400,
detail={
"error": f"Maximum {MAX_BATCH_SIZE} keys can be updated at once. Found {len(data.key_ids)} key_ids."
},
)
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,
"AND": [
{"OR": [{"blocked": False}, {"blocked": None}]},
{"OR": [{"expires": None}, {"expires": {"gt": now}}]},
],
},
order={"token": "asc"},
take=MAX_BATCH_SIZE + 1,
)
if len(existing_keys) > MAX_BATCH_SIZE:
raise HTTPException(
status_code=400,
detail={
"error": f"Team {data.team_id} has more than {MAX_BATCH_SIZE} keys. Use `key_ids` to update in batches of {MAX_BATCH_SIZE}."
},
)
requested_tokens = [row.token for row in existing_keys]
else:
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": hashed_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(
status_code=404,
detail={"error": f"No keys found for team {data.team_id}"},
)
existing_by_token = {row.token: row for row in existing_keys}
update_field_dict = data.update_fields.model_dump(exclude_unset=True)
successful_updates: List[SuccessfulKeyUpdate] = []
failed_updates: List[FailedKeyUpdate] = []
for token in requested_tokens:
db_token = _hash_token_if_needed(token)
try:
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(
update_key_request=update_key_request,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
user_custom_key_update=user_custom_key_update,
existing_key_row=existing_by_token[db_token],
)
successful_updates.append(
SuccessfulKeyUpdate(key=token, key_info=updated_key_info)
)
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 {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(db_token),
)
)
return BulkUpdateKeyResponse(
total_requested=len(requested_tokens),
successful_updates=successful_updates,
failed_updates=failed_updates,
)
async def validate_key_team_change(
key: LiteLLM_VerificationToken,
team: LiteLLM_TeamTable,

80
litellm/types/proxy/management_endpoints/key_management_endpoints.py
View file

@ -1,6 +1,7 @@
from typing import Any, Dict, List, Optional
from datetime import datetime
from typing import Any, Dict, List, Literal, Optional
from pydantic import BaseModel
from pydantic import BaseModel, ConfigDict, model_validator
class BulkUpdateKeyRequestItem(BaseModel):
@ -40,3 +41,78 @@ class BulkUpdateKeyResponse(BaseModel):
total_requested: int
successful_updates: List[SuccessfulKeyUpdate]
failed_updates: List[FailedKeyUpdate]
class KeyUpdateFields(BaseModel):
"""Allowlist of bulk-broadcastable fields for /team/key/bulk_update; `extra="forbid"` blocks RBAC/ownership/scope mutations even by team admins."""
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
# 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":
if self.temp_budget_increase is not None or self.temp_budget_expiry is not None:
if self.temp_budget_increase is None or self.temp_budget_expiry is None:
raise ValueError(
"temp_budget_increase and temp_budget_expiry must be set together"
)
return self
@model_validator(mode="after")
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):
"""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
all_keys_in_team: bool = False
update_fields: KeyUpdateFields
@model_validator(mode="after")
def validate_selection(self) -> "BulkUpdateTeamKeysRequest":
has_key_ids = self.key_ids is not None and len(self.key_ids) > 0
if has_key_ids and self.all_keys_in_team:
raise ValueError(
"Provide either `key_ids` or `all_keys_in_team=True`, not both."
)
if not has_key_ids and not self.all_keys_in_team:
raise ValueError(
"Must provide either `key_ids` (non-empty) or `all_keys_in_team=True`."
)
return self

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

@ -5689,7 +5689,7 @@ async def test_process_single_key_update():
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook"
):
# Create update request
key_update_item = BulkUpdateKeyRequestItem(
update_key_request = UpdateKeyRequest(
key="test-key-123",
max_budget=100.0,
tags=["production"],
@ -5703,7 +5703,7 @@ async def test_process_single_key_update():
# Call the function
result = await _process_single_key_update(
key_update_item=key_update_item,
update_key_request=update_key_request,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
prisma_client=mock_prisma_client,
@ -9855,9 +9855,6 @@ async def test_process_single_key_update_cache_invalidation_with_token_hash():
from litellm.proxy.management_endpoints.key_management_endpoints import (
_process_single_key_update,
)
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateKeyRequestItem,
)
token_hash = "abc123def456"
@ -9900,7 +9897,7 @@ async def test_process_single_key_update_cache_invalidation_with_token_hash():
new_callable=AsyncMock,
),
):
key_update_item = BulkUpdateKeyRequestItem(
update_key_request = UpdateKeyRequest(
key=token_hash,
max_budget=100.0,
)
@ -9912,7 +9909,7 @@ async def test_process_single_key_update_cache_invalidation_with_token_hash():
)
await _process_single_key_update(
key_update_item=key_update_item,
update_key_request=update_key_request,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
prisma_client=mock_prisma_client,
@ -10019,3 +10016,583 @@ async def test_execute_virtual_key_regeneration_cache_invalidation_with_token_ha
call_kwargs = mock_delete_cache.call_args.kwargs
# The token hash should be passed as-is, NOT double-hashed
assert call_kwargs["hashed_token"] == token_hash
# ---------------------------------------------------------------------------
# /team/key/bulk_update tests
# ---------------------------------------------------------------------------
_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,
user_id="user-123",
models=[],
team_id=team_id,
max_budget=None,
)
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(
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(
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(
f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
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):
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
keys = [_make_team_key("tok-a"), _make_team_key("tok-b")]
find_unique = AsyncMock(side_effect=keys)
mock = _setup_team_keys_mocks(
monkeypatch,
find_many=keys,
find_unique=find_unique,
update_data=AsyncMock(
side_effect=[{"data": _updated({"max_budget": 50.0})}] * 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
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"]}
find_unique.assert_not_called()
@pytest.mark.asyncio
async def test_bulk_update_team_keys_success_all_keys_in_team(monkeypatch):
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
keys = [_make_team_key(f"tok-{i}") for i in range(3)]
find_unique = AsyncMock(side_effect=keys)
mock = _setup_team_keys_mocks(
monkeypatch,
find_many=keys,
find_unique=find_unique,
update_data=AsyncMock(
side_effect=[{"data": _updated({"max_budget": 50.0})}] * 3
),
)
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)
)
find_unique.assert_not_called()
@pytest.mark.asyncio
async def test_bulk_update_team_keys_key_not_in_team(monkeypatch):
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
in_team = _make_team_key("tok-a")
_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})}),
)
response = await _call_as_admin(
BulkUpdateTeamKeysRequest(
team_id="team-abc",
key_ids=["tok-a", "tok-foreign"],
update_fields=KeyUpdateFields(max_budget=50.0),
)
)
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_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,
)
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
key_a = _make_team_key("tok-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})}),
)
auth_check = AsyncMock()
monkeypatch.setattr(
f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
auth_check,
)
response = await bulk_update_team_keys(
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 + 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_member_no_permission(monkeypatch):
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=[_make_team_key("tok-a")])
monkeypatch.setattr(
f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
AsyncMock(
side_effect=ProxyException(
message="not in team",
type="team_member_permission_error",
param="/key/update",
code=401,
)
),
)
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=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,
)
# 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_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,
)
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],
update_data=AsyncMock(side_effect=RuntimeError("boom")),
hash_identity=False,
)
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_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,
)
_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_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,
)
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateTeamKeysRequest,
KeyUpdateFields,
)
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(
metadata={"allowed_passthrough_routes": ["/admin/*"]}
),
)
with pytest.raises(HTTPException) as exc:
await bulk_update_team_keys(
data=request,
user_api_key_dict=_internal_user(),
litellm_changed_by=None,
)
assert exc.value.status_code == 403
assert "allowed_passthrough_routes" in str(exc.value.detail)
mock.update_data.assert_not_called()