From b746ac44563589a0e2b407b064470c2ee18b4b27 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 03:44:07 -0700 Subject: [PATCH] fix(proxy): accept object_permission on /key/bulk_update items instead of 422 --- .../key_management_endpoints.py | 3 +- .../key_management_endpoints.py | 12 ++++-- .../test_key_management_endpoints.py | 41 +++++++++---------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 ++- 4 files changed, 34 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f28101f7072..959fee7b010 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3517,9 +3517,10 @@ 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 an - explicit null clears it, the same as /key/update. An item carrying any other field is rejected with 422. + explicit null clears it, the same as /key/update. Returns: - total_requested: int - Total number of keys requested for update diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 3e193956d30..001bc3c0d51 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/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,15 +30,14 @@ class KeySearchWhere(TypedDict): class BulkUpdateKeyRequestItem(BaseModel): - """One /key/bulk_update item; only the fields it carries are written, and unknown fields are rejected.""" - - model_config = ConfigDict(extra="forbid") + """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): 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 9d023e129aa..1bf5018900a 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 @@ -7100,7 +7100,7 @@ async def test_list_key_helper_applies_search_to_prisma_where(): _BULK_UPDATE_TOKEN: Final = "1f2e3d4c5b6a79880123456789abcdef0123456789abcdef0123456789abcdef" -async def _bulk_update_one_key(monkeypatch, item_payload: Mapping[str, object]) -> Mapping[str, object]: +async def _bulk_update_one_key(monkeypatch, item_payload: Mapping[str, object]) -> AsyncMock: from litellm.proxy.management_endpoints.key_management_endpoints import bulk_update_keys from litellm.types.proxy.management_endpoints.key_management_endpoints import BulkUpdateKeyRequest @@ -7109,6 +7109,10 @@ async def _bulk_update_one_key(monkeypatch, item_payload: Mapping[str, object]) ) mock_prisma_client = AsyncMock() mock_prisma_client.db.litellm_verificationtoken.find_unique = 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) @@ -7135,14 +7139,18 @@ async def _bulk_update_one_key(monkeypatch, item_payload: Mapping[str, object]) ) assert response.failed_updates == [] - return mock_prisma_client.update_data.call_args.kwargs["data"] + return mock_prisma_client + + +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 = await _bulk_update_one_key(monkeypatch, {"tags": ["team-a"]}) + 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() @@ -7151,32 +7159,23 @@ async def test_bulk_update_keys_item_without_a_field_leaves_that_column_alone(mo @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 = await _bulk_update_one_key(monkeypatch, {"max_budget": None}) + 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() -def test_bulk_update_keys_rejects_a_field_the_bulk_path_cannot_apply(): +@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.""" - from fastapi import FastAPI + prisma = await _bulk_update_one_key(monkeypatch, {"object_permission": {"vector_stores": ["vs-1"]}}) - from litellm.proxy.auth.user_api_key_auth import user_api_key_auth - from litellm.proxy.management_endpoints.key_management_endpoints import router - - test_app = FastAPI() - test_app.include_router(router) - test_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" - ) - response = TestClient(test_app).post( - "/key/bulk_update", - json={"keys": [{"key": _BULK_UPDATE_TOKEN, "object_permission": {"vector_stores": ["vs-1"]}}]}, - ) - - assert response.status_code == 422, response.text - assert "object_permission" in response.text + 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 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a9174603290..4ada2c3372b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7693,9 +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 an - * explicit null clears it, the same as /key/update. An item carrying any other field is rejected with 422. + * explicit null clears it, the same as /key/update. * * Returns: * - total_requested: int - Total number of keys requested for update @@ -25240,7 +25241,7 @@ export interface components { }; /** * BulkUpdateKeyRequestItem - * @description One /key/bulk_update item; only the fields it carries are written, and unknown fields are rejected. + * @description One /key/bulk_update item; only the fields it carries are written. */ BulkUpdateKeyRequestItem: { /** Budget Id */ @@ -25249,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 */