From bdbe265c7020a6533b8dd729708ff01725bac4ad Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 03:14:39 -0700 Subject: [PATCH] fix(proxy): /key/bulk_update writes only the fields each item carries A bulk item that carried only tags reached the DB with max_budget, team_id, and budget_id as explicit nulls, wiping the key's budget and detaching it from its team. The per-key update is now built from the fields the item actually set, so a field left out keeps its value and an explicit null still clears it, the same as /key/update. Items carrying a field the bulk path cannot apply (object_permission and the like) are rejected with 422 instead of being silently dropped. --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- .../key_management_endpoints.py | 14 ++-- .../key_management_endpoints.py | 4 +- .../test_key_management_endpoints.py | 83 +++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 +- 5 files changed, 96 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 73ea8cf1991..18849ef5b64 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19622,7 +19622,7 @@ } } }, - "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " + "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" }, "500": { "content": { diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 6c195d713c8..f28101f7072 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3517,7 +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 - + + 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. + 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 +3589,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, diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 63bbaa5ba4e..3e193956d30 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -25,7 +25,9 @@ class KeySearchWhere(TypedDict): class BulkUpdateKeyRequestItem(BaseModel): - """Individual key update request item""" + """One /key/bulk_update item; only the fields it carries are written, and unknown fields are rejected.""" + + model_config = ConfigDict(extra="forbid") key: str # Key identifier (token) budget_id: str | None = None # Budget ID associated with the key 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 c852307b051..2f0f605b839 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 @@ -7097,6 +7097,89 @@ 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" + + +async def _bulk_update_one_key(monkeypatch, item_payload: Mapping[str, object]) -> Mapping[str, object]: + """Runs /key/bulk_update with one item against a budgeted team key and returns the row written to the DB.""" + from litellm.proxy.management_endpoints.key_management_endpoints import bulk_update_keys + from litellm.types.proxy.management_endpoints.key_management_endpoints import BulkUpdateKeyRequest + + 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.update_data = AsyncMock(return_value={"data": {"token": _BULK_UPDATE_TOKEN}}) + _setup_update_key_mocks(monkeypatch, mock_prisma_client) + + 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, + ) + + assert response.failed_updates == [] + return mock_prisma_client.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"]}) + + 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 = 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(): + """`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 + + 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 + + @pytest.mark.asyncio async def test_generate_key_negative_max_budget(): """ diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a5ede63bf7f..a9174603290 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7694,6 +7694,9 @@ export interface paths { * - team_id: Optional[str] - Team ID associated with key * - tags: Optional[List[str]] - Tags for organizing keys * + * 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. + * * Returns: * - total_requested: int - Total number of keys requested for update * - successful_updates: List[SuccessfulKeyUpdate] - List of successfully updated keys with their updated info @@ -25237,7 +25240,7 @@ export interface components { }; /** * BulkUpdateKeyRequestItem - * @description Individual key update request item + * @description One /key/bulk_update item; only the fields it carries are written, and unknown fields are rejected. */ BulkUpdateKeyRequestItem: { /** Budget Id */