Merge pull request #41949 from BerriAI/litellm_bulk_update_keys_keep_unset_fields

fix(proxy): /key/bulk_update writes only the fields each item carries

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -2798,9 +2798,18 @@ async def _process_single_key_update(
             llm_router=llm_router,
         )
 
+    key_request: Final = await _with_validated_object_permission(
+        update_key_request=update_key_request,
+        team_obj=team_obj,
+        existing_key_row=existing_key_row,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        user_api_key_dict=user_api_key_dict,
+    )
+
     # Prepare update data
     non_default_values = await prepare_key_update_data(
-        data=update_key_request, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router
+        data=key_request, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router
     )
 
     await _enforce_custom_key_policy(
@@ -2809,7 +2818,7 @@ async def _process_single_key_update(
             operation="update",
             existing_key_row=existing_key_row,
             non_default_values=non_default_values,
-            request=update_key_request,
+            request=key_request,
         ),
     )
 
@@ -2825,15 +2834,15 @@ async def _process_single_key_update(
         existing_key_row=existing_key_row,
         prisma_client=prisma_client,
     )
-    _data: Final = {**update_values, "token": update_key_request.key}
+    _data: Final = {**update_values, "token": key_request.key}
     response: Final[Mapping[str, object] | None] = cast(  # cast-ok: every update_data branch returns a str-keyed dict
         "Mapping[str, object] | None",
-        await prisma_client.update_data(token=update_key_request.key, data=_data),
+        await prisma_client.update_data(token=key_request.key, data=_data),
     )
 
     # Delete cache
     await _delete_cache_key_object(
-        hashed_token=_hash_token_if_needed(update_key_request.key),
+        hashed_token=_hash_token_if_needed(key_request.key),
         user_api_key_cache=user_api_key_cache,
         proxy_logging_obj=proxy_logging_obj,
     )
@@ -2842,17 +2851,15 @@ async def _process_single_key_update(
     # authenticating against the access groups it just lost.
     await sync_key_update_access_group_membership(
         prisma_client=prisma_client,
-        key_token=_hash_token_if_needed(
-            _resolve_token_to_update(data=update_key_request, existing_key_row=existing_key_row)
-        ),
-        data=update_key_request,
+        key_token=_hash_token_if_needed(_resolve_token_to_update(data=key_request, existing_key_row=existing_key_row)),
+        data=key_request,
         existing_key_row=existing_key_row,
     )
 
     # Trigger async hook
     asyncio.create_task(
         KeyManagementEventHooks.async_key_updated_hook(
-            data=update_key_request,
+            data=key_request,
             existing_key_row=existing_key_row,
             response=response,
             user_api_key_dict=user_api_key_dict,
@@ -2875,6 +2882,31 @@ async def _process_single_key_update(
     return updated_key_info
 
 
+async def _with_validated_object_permission(
+    update_key_request: UpdateKeyRequest,
+    team_obj: LiteLLM_TeamTableCachedObj | None,
+    existing_key_row: LiteLLM_VerificationToken,
+    prisma_client: PrismaClient | None,
+    user_api_key_cache: UserApiKeyCache,
+    user_api_key_dict: UserAPIKeyAuth,
+) -> UpdateKeyRequest:
+    if update_key_request.object_permission is None:
+        return update_key_request
+    normalized_object_permission: Final = await _validate_mcp_servers_for_key_update(
+        data=update_key_request,
+        team_obj=team_obj,
+        existing_key_row=existing_key_row,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        is_proxy_admin=user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value,
+    )
+    if normalized_object_permission is None:
+        return update_key_request
+    return update_key_request.model_copy(
+        update=MappingProxyType({"object_permission": LiteLLM_ObjectPermissionBase(**normalized_object_permission)})
+    )
+
+
 async def _validate_mcp_servers_for_key_update(
     data: "UpdateKeyRequest",
     team_obj: Optional["LiteLLM_TeamTableCachedObj"],
@@ -3517,7 +3549,11 @@ async def bulk_update_keys(
         - max_budget: Optional[float] - Max budget for key
         - team_id: Optional[str] - Team ID associated with key
         - tags: Optional[List[str]] - Tags for organizing keys
-    
+        - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission, as on /key/update
+
+    Only the fields an item carries are written: a field left out keeps its current value, and a field
+    sent explicitly, null included, is applied exactly as /key/update applies it.
+
     Returns:
     - total_requested: int - Total number of keys requested for update
     - successful_updates: List[SuccessfulKeyUpdate] - List of successfully updated keys with their updated info
@@ -3586,15 +3622,8 @@ async def bulk_update_keys(
 
     for key_update_item in data.keys:
         try:
-            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(
-                update_key_request=update_key_request,
+                update_key_request=UpdateKeyRequest.model_validate(key_update_item.model_dump(exclude_unset=True)),
                 user_api_key_dict=user_api_key_dict,
                 litellm_changed_by=litellm_changed_by,
                 prisma_client=prisma_client,

litellm/types/proxy/management_endpoints/key_management_endpoints.py

@@ -5,7 +5,12 @@ from pydantic import BaseModel, ConfigDict, model_validator
 from typing_extensions import ReadOnly, TypedDict
 
 from litellm.models.verification_token import LiteLLM_VerificationToken
-from litellm.proxy._types import GenerateKeyRequest, RegenerateKeyRequest, UpdateKeyRequest
+from litellm.proxy._types import (
+    GenerateKeyRequest,
+    LiteLLM_ObjectPermissionBase,
+    RegenerateKeyRequest,
+    UpdateKeyRequest,
+)
 from litellm.types.llms.base import LiteLLMPydanticObjectBase
 from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains
 
@@ -25,13 +30,14 @@ class KeySearchWhere(TypedDict):
 
 
 class BulkUpdateKeyRequestItem(BaseModel):
-    """Individual key update request item"""
+    """One /key/bulk_update item; only the fields it carries are written."""
 
     key: str  # Key identifier (token)
     budget_id: str | None = None  # Budget ID associated with the key
     max_budget: float | None = None  # Max budget for key
     team_id: str | None = None  # Team ID associated with key
     tags: list[str] | None = None  # Tags for organizing keys
+    object_permission: LiteLLM_ObjectPermissionBase | None = None
 
 
 class BulkUpdateKeyRequest(BaseModel):

tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py

@@ -80,7 +80,11 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
     validate_key_team_change,
 )
 from litellm.proxy.proxy_server import app
-from litellm.types.proxy.management_endpoints.key_management_endpoints import CustomKeyPolicyRequest
+from litellm.types.proxy.management_endpoints.key_management_endpoints import (
+    BulkUpdateKeyRequest,
+    BulkUpdateKeyResponse,
+    CustomKeyPolicyRequest,
+)
 
 client = TestClient(app)
 
@@ -7097,6 +7101,115 @@ async def test_list_key_helper_applies_search_to_prisma_where():
     assert _search_clause("key-id-123", "key-id-123") in where["AND"], f"search not in Prisma where: {where}"
 
 
+_BULK_UPDATE_TOKEN: Final = "1f2e3d4c5b6a79880123456789abcdef0123456789abcdef0123456789abcdef"
+_BULK_UPDATE_TEAM: Final = LiteLLM_TeamTableCachedObj(team_id="team-1")
+
+
+async def _run_bulk_update_on_one_key(
+    monkeypatch, item_payload: Mapping[str, object], team: LiteLLM_TeamTableCachedObj = _BULK_UPDATE_TEAM
+) -> tuple[BulkUpdateKeyResponse, AsyncMock]:
+    from litellm.proxy.management_endpoints.key_management_endpoints import bulk_update_keys
+
+    key_in_db = LiteLLM_VerificationToken(
+        token=_BULK_UPDATE_TOKEN, user_id="test-user", team_id="team-1", max_budget=100.0, budget_id="budget-1"
+    )
+    mock_prisma_client = AsyncMock()
+    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=key_in_db)
+    mock_prisma_client.get_data = AsyncMock(return_value=key_in_db)
+    mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None)
+    mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(
+        return_value=MagicMock(object_permission_id="objperm-bulk")
+    )
+    mock_prisma_client.update_data = AsyncMock(return_value={"data": {"token": _BULK_UPDATE_TOKEN}})
+    _setup_update_key_mocks(monkeypatch, mock_prisma_client)
+    monkeypatch.setattr(
+        "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", AsyncMock(return_value=team)
+    )
+
+    with (
+        patch(  # test-quality-ok: the handler reads the cache and hook singletons from module globals, no injection seam
+            "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
+            new_callable=AsyncMock,
+        ),
+        patch(  # test-quality-ok: the permission check is a classmethod the handler calls directly, no injection seam
+            "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
+            new_callable=AsyncMock,
+        ),
+        patch(  # test-quality-ok: the audit hook is a classmethod the handler calls directly, no injection seam
+            "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook",
+            new_callable=AsyncMock,
+        ),
+    ):
+        response = await bulk_update_keys(
+            data=BulkUpdateKeyRequest.model_validate({"keys": [{"key": _BULK_UPDATE_TOKEN, **item_payload}]}),
+            user_api_key_dict=UserAPIKeyAuth(
+                user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
+            ),
+            litellm_changed_by=None,
+        )
+
+    return response, mock_prisma_client
+
+
+async def _bulk_update_one_key(monkeypatch, item_payload: Mapping[str, object]) -> AsyncMock:
+    response, prisma = await _run_bulk_update_on_one_key(monkeypatch, item_payload)
+    assert response.failed_updates == []
+    return prisma
+
+
+def _written_key_row(prisma: AsyncMock) -> Mapping[str, object]:
+    return prisma.update_data.call_args.kwargs["data"]
+
+
+@pytest.mark.asyncio
+async def test_bulk_update_keys_item_without_a_field_leaves_that_column_alone(monkeypatch):
+    """A tags-only item used to reach the DB with max_budget, team_id, and budget_id as explicit
+    nulls, so tagging a key wiped its budget and detached it from its team."""
+    written = _written_key_row(await _bulk_update_one_key(monkeypatch, {"tags": ["team-a"]}))
+
+    assert written["metadata"]["tags"] == ["team-a"]
+    assert not {"max_budget", "team_id", "budget_id"} & written.keys()
+
+
+@pytest.mark.asyncio
+async def test_bulk_update_keys_explicit_null_still_clears_the_field(monkeypatch):
+    """Sending `"max_budget": null` on an item is a request to remove the budget, as on /key/update."""
+    written = _written_key_row(await _bulk_update_one_key(monkeypatch, {"max_budget": None}))
+
+    assert written["max_budget"] is None
+    assert not {"team_id", "budget_id"} & written.keys()
+
+
+@pytest.mark.asyncio
+async def test_bulk_update_keys_object_permission_is_granted_not_dropped(monkeypatch):
+    """`object_permission` used to be accepted with 200 and dropped, leaving an item that carried
+    nothing but the key, so the call wiped the key's budget instead of granting the permission."""
+    prisma = await _bulk_update_one_key(monkeypatch, {"object_permission": {"vector_stores": ["vs-1"]}})
+
+    upserted = prisma.db.litellm_objectpermissiontable.upsert.call_args.kwargs["data"]["create"]
+    assert upserted["vector_stores"] == ["vs-1"]
+    written = _written_key_row(prisma)
+    assert written["object_permission_id"] == "objperm-bulk"
+    assert not {"max_budget", "team_id", "budget_id"} & written.keys()
+
+
+@pytest.mark.asyncio
+async def test_bulk_update_keys_object_permission_outside_the_team_allowlist_is_refused(monkeypatch):
+    """A bulk item's object_permission is checked against the key's team exactly as /key/update
+    checks it, so a team key cannot be granted a search tool its team does not allow."""
+    team = LiteLLM_TeamTableCachedObj(
+        team_id="team-1",
+        object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-team-1", search_tools=["team-search"]),
+    )
+    response, prisma = await _run_bulk_update_on_one_key(
+        monkeypatch, {"object_permission": {"search_tools": ["other-search"]}}, team=team
+    )
+
+    assert response.successful_updates == []
+    assert "not allowed by team 'team-1'" in response.failed_updates[0].failed_reason
+    prisma.update_data.assert_not_called()
+
+
 @pytest.mark.asyncio
 async def test_generate_key_negative_max_budget():
     """
@@ -13076,6 +13189,11 @@ async def _process_single_key_update_under_policy(prisma_client: AsyncMock, data
             "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook",
             new_callable=AsyncMock,
         ),
+        patch(  # test-quality-ok: the existing key's team is outside the policy path, as in the /key/update tests
+            "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
+            new_callable=AsyncMock,
+            return_value=None,
+        ),
     ):
         return await _process_single_key_update(
             update_key_request=data,

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -7693,6 +7693,10 @@ export interface paths {
          *         - max_budget: Optional[float] - Max budget for key
          *         - team_id: Optional[str] - Team ID associated with key
          *         - tags: Optional[List[str]] - Tags for organizing keys
+         *         - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission, as on /key/update
+         *
+         *     Only the fields an item carries are written: a field left out keeps its current value, and a field
+         *     sent explicitly, null included, is applied exactly as /key/update applies it.
          *
          *     Returns:
          *     - total_requested: int - Total number of keys requested for update
@@ -25237,7 +25241,7 @@ export interface components {
         };
         /**
          * BulkUpdateKeyRequestItem
-         * @description Individual key update request item
+         * @description One /key/bulk_update item; only the fields it carries are written.
          */
         BulkUpdateKeyRequestItem: {
             /** Budget Id */
@@ -25246,6 +25250,7 @@ export interface components {
             key: string;
             /** Max Budget */
             max_budget?: number | null;
+            object_permission?: components["schemas"]["LiteLLM_ObjectPermissionBase"] | null;
             /** Tags */
             tags?: string[] | null;
             /** Team Id */