From 93d579fc9c16b15113eb4ea3cecbd6a7aa30c958 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 2 Oct 2026 04:27:41 +0000 Subject: [PATCH] fix(keys): write project assignment only while key row is still unassigned Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 71 +++++++++++++++++-- .../test_key_management_endpoints.py | 36 ++++++++++ 2 files changed, 101 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 419e1ca5891..22286347bd9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2467,6 +2467,47 @@ async def _apply_soft_budget_update( return remaining +async def _write_guarded_project_assignment( + table: "TableActions[prisma_models.LiteLLM_VerificationToken]", + hashed_token: str, + data: Mapping[str, object], +) -> "prisma_models.LiteLLM_VerificationToken | None": + updated_count: Final = await table.update_many( + where={"token": hashed_token, "project_id": None}, + data=data, + ) + if updated_count == 0: + raise HTTPException( + status_code=400, + detail="Project reassignment is not supported. Use null to detach the key.", + ) + return await table.find_unique( + where={"token": hashed_token}, + include={"object_permission": True}, + ) + + +async def _update_key_row_assigning_project( + prisma_client: PrismaClient, + key: str, + update_values: Mapping[str, object], +) -> _KeyUpdateResult: + """Write a key update that assigns a project only while the row still has none.""" + hashed_token: Final = _hash_token_if_needed(key) + updated_row: Final = await _write_guarded_project_assignment( + table=_prisma_table(VerificationTokenRepository(prisma_client)), + hashed_token=hashed_token, + data=with_settings_updated_at( + prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token})) + ), + ) + updated_data: Final[Mapping[str, object]] = ( + updated_row.model_dump() if updated_row is not None else MappingProxyType({}) + ) + result: Final[_KeyUpdateResult] = {"token": hashed_token, "data": updated_data} + return result + + async def _update_key_row_with_soft_budget( prisma_client: PrismaClient, key: str, @@ -2474,6 +2515,7 @@ async def _update_key_row_with_soft_budget( non_default_values: Mapping[str, object], existing_key_row: LiteLLM_VerificationToken, changed_by: str, + expect_unassigned_project: bool, ) -> _KeyUpdateResult: hashed_token: Final = _hash_token_if_needed(key) key_where: Final[_KeyRowWhere] = {"token": hashed_token} @@ -2487,12 +2529,21 @@ async def _update_key_row_with_soft_budget( changed_by=changed_by, ) include_object_permission: Final[prisma.types.LiteLLM_VerificationTokenInclude] = {"object_permission": True} - updated_row: Final = await tx.litellm_verificationtoken.update( - where=key_where, - data=with_settings_updated_at( - prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token})) - ), - include=include_object_permission, + update_payload: Final = with_settings_updated_at( + prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token})) + ) + updated_row: Final = ( + await _write_guarded_project_assignment( + table=tx.litellm_verificationtoken, + hashed_token=hashed_token, + data=update_payload, + ) + if expect_unassigned_project + else await tx.litellm_verificationtoken.update( + where=key_where, + data=update_payload, + include=include_object_permission, + ) ) updated_data: Final[Mapping[str, object]] = ( updated_row.model_dump() if updated_row is not None else MappingProxyType({}) @@ -3576,6 +3627,7 @@ async def update_key_fn( prisma_client=prisma_client, ) changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name + is_project_assignment: Final = data.project_id is not None and existing_key_row.project_id is None response: Final = ( await _update_key_row_with_soft_budget( prisma_client=prisma_client, @@ -3584,8 +3636,15 @@ async def update_key_fn( non_default_values=update_values, existing_key_row=existing_key_row, changed_by=changed_by, + expect_unassigned_project=is_project_assignment, ) if "soft_budget" in data.model_fields_set + else await _update_key_row_assigning_project( + prisma_client=prisma_client, + key=key, + update_values=update_values, + ) + if is_project_assignment else await prisma_client.update_data(token=key, data=MappingProxyType({**update_values, "token": key})) ) diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 8f2060cc8a8..d1caaf6447c 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -61,6 +61,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _requested_end_user_budget_id, _save_deleted_verification_token_records, _transform_verification_tokens_to_deleted_records, + _update_key_row_assigning_project, _validate_end_user_budget_id_change, _validate_max_budget, _validate_reset_spend_value, @@ -20595,6 +20596,41 @@ async def test_prepare_key_update_data_includes_new_project_assignment(): assert result["project_id"] == "project-orbit" +@pytest.mark.asyncio +async def test_project_assignment_write_requires_row_still_unassigned(): + database = MagicMock() + database.jsonify_object = lambda data: dict(data) + database.db.litellm_verificationtoken.update_many = AsyncMock(return_value=1) + row = MagicMock() + row.model_dump = MagicMock(return_value={"project_id": "project-orbit"}) + database.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=row) + + result: Final = await _update_key_row_assigning_project( + prisma_client=database, key="sk-assign", update_values={"project_id": "project-orbit"} + ) + + update_where: Final = database.db.litellm_verificationtoken.update_many.await_args.kwargs["where"] + assert update_where["project_id"] is None + assert update_where["token"] == result["token"] + assert result["data"] == {"project_id": "project-orbit"} + + +@pytest.mark.asyncio +async def test_project_assignment_write_rejects_when_row_gained_project(): + database = MagicMock() + database.jsonify_object = lambda data: dict(data) + database.db.litellm_verificationtoken.update_many = AsyncMock(return_value=0) + + with pytest.raises(HTTPException) as exc: + await _update_key_row_assigning_project( + prisma_client=database, key="sk-assign", update_values={"project_id": "project-orbit"} + ) + + assert exc.value.status_code == 400 + assert "reassignment" in str(exc.value.detail) + database.db.litellm_verificationtoken.find_unique.assert_not_called() + + @pytest.mark.asyncio async def test_key_creator_cannot_detach_project_without_admin_access(): existing: Final = LiteLLM_VerificationToken(