From 79c96354a539f0887076600b1f585d4e5c41981c Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 17 Sep 2026 23:40:41 +0000 Subject: [PATCH 01/12] feat(keys): allow assigning a project to an unassigned key on /key/update Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 2 +- .../key_management_endpoints.py | 56 +++++++++- .../test_key_management_endpoints.py | 102 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 4 files changed, 156 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index fa81f2ab6f4..8dc9e31740a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1325,7 +1325,7 @@ class UpdateKeyRequest(KeyRequestBase): project_id: str | None = Field( default=None, - description="Omit to retain the project, or send null to detach. Assigning a different project is not supported.", + description="Omit to retain the project, send null to detach, or send a project id to assign an unassigned key to a project on the key's team. Moving a key between projects is not supported.", ) @model_validator(mode="before") diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 802a7c3e469..8649c3a690e 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1767,6 +1767,50 @@ async def _check_project_key_limits( ) +async def _validate_project_assignment( + data: UpdateKeyRequest, + existing_key_row: LiteLLM_VerificationToken, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> None: + """Validate assigning a project to a key that does not have one yet.""" + if data.project_id is None or data.project_id == existing_key_row.project_id: + return + if existing_key_row.project_id is not None: + raise HTTPException( + status_code=400, detail="Project reassignment is not supported. Use null to detach the key." + ) + project_obj: Final = await get_project_object( + project_id=data.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + if project_obj is None: + raise HTTPException( + status_code=404, + detail={"error": f"Project not found, project_id={data.project_id}"}, + ) + team: Final = data.team_id if "team_id" in data.model_fields_set else existing_key_row.team_id + if team is None or team != project_obj.team_id: + raise HTTPException( + status_code=400, + detail={ + "error": f"Project {data.project_id} belongs to team {project_obj.team_id}. Keys can only be assigned to a project owned by their own team (key team: {team})" + }, + ) + await _check_project_key_limits( + project_id=data.project_id, + data=data.model_copy( + update={ + "models": data.models if "models" in data.model_fields_set else existing_key_row.models, + "max_budget": data.max_budget if data.max_budget is not None else existing_key_row.max_budget, + } + ), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + def check_org_key_model_specific_limits( keys: Sequence[LiteLLM_VerificationToken], org_table: LiteLLM_OrganizationTable, @@ -2925,10 +2969,12 @@ async def _validate_update_key_data( user_api_key_dict=user_api_key_dict, ) - if data.project_id is not None and data.project_id != existing_key_row.project_id: - raise HTTPException( - status_code=400, detail="Project reassignment is not supported. Use null to detach the key." - ) + await _validate_project_assignment( + data=data, + existing_key_row=existing_key_row, + prisma_client=checked_prisma_client, + user_api_key_cache=user_api_key_cache, + ) is_project_change: Final = "project_id" in data.model_fields_set and data.project_id != existing_key_row.project_id common_key_access_checks( @@ -3179,7 +3225,7 @@ async def update_key_fn( - user_id: Optional[str] - User ID associated with key - team_id: Optional[str] - Team ID associated with key - agent_id: Optional[str] - The agent id associated with the key. - - project_id: Optional[str] - Omit to retain the project, or send null to detach. A different project ID is rejected. + - project_id: Optional[str] - Omit to retain the project, send null to detach, or send a project id to assign an unassigned key to a project on the key's team. Moving a key between projects is rejected. - organization_id: Optional[str] - The organization id of the key. - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. - models: Optional[list] - Model_name's a user is allowed to call 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 cc0a7631b59..bd8bf4316fb 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 @@ -19683,6 +19683,108 @@ async def test_project_detachment_uses_effective_project_for_validation(project_ assert expected in str(exc.value.detail) +def _project_assignment_prisma() -> MagicMock: + database = MagicMock() + database.db.litellm_teamtable.find_unique = AsyncMock( + return_value=LiteLLM_TeamTable(team_id="team-lit-5823", members=[]) + ) + return database + + +@pytest.mark.asyncio +async def test_project_assignment_to_unassigned_key_on_same_team(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + + +@pytest.mark.asyncio +async def test_project_assignment_rejects_project_on_another_team(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-other", models=["model-orbit"] + ) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + assert exc.value.status_code == 400 + assert "team" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_rejects_key_with_no_team(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id=None, models=["model-orbit"] + ) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + assert exc.value.status_code == 400 + assert "team" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_rejects_nonexistent_project(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + database: Final = _project_assignment_prisma() + database.db.litellm_projecttable.find_unique = AsyncMock(return_value=None) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-ghost"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, database, UserApiKeyCache(), + ) + assert exc.value.status_code == 404 + assert "Project not found" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_validates_existing_models_against_project(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-elsewhere"] + ) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + assert exc.value.status_code == 400 + assert "not in project's allowed models" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_includes_new_project_assignment(): + existing: Final = LiteLLM_VerificationToken(token="project-assign-token", project_id=None) + + result: Final = await prepare_key_update_data( + data=UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing_key_row=existing, + ) + + assert result["project_id"] == "project-orbit" + + @pytest.mark.asyncio async def test_key_creator_cannot_detach_project_without_admin_access(): existing: Final = LiteLLM_VerificationToken( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 558c10e4f10..adff9fed806 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8169,7 +8169,7 @@ export interface paths { * - user_id: Optional[str] - User ID associated with key * - team_id: Optional[str] - Team ID associated with key * - agent_id: Optional[str] - The agent id associated with the key. - * - project_id: Optional[str] - Omit to retain the project, or send null to detach. A different project ID is rejected. + * - project_id: Optional[str] - Omit to retain the project, send null to detach, or send a project id to assign an unassigned key to a project on the key's team. Moving a key between projects is rejected. * - organization_id: Optional[str] - The organization id of the key. * - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. * - models: Optional[list] - Model_name's a user is allowed to call @@ -39011,7 +39011,7 @@ export interface components { policies?: string[] | null; /** * Project Id - * @description Omit to retain the project, or send null to detach. Assigning a different project is not supported. + * @description Omit to retain the project, send null to detach, or send a project id to assign an unassigned key to a project on the key's team. Moving a key between projects is not supported. */ project_id?: string | null; /** Prompts */ From 9e9b937cb2ff838aa880c6e3fa84aa184582fc85 Mon Sep 17 00:00:00 2001 From: jesus Date: Thu, 17 Sep 2026 23:53:47 +0000 Subject: [PATCH 02/12] chore(keys): suppress LIT002 on project assignment error payloads Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/management_endpoints/key_management_endpoints.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 8649c3a690e..5ab5ebe0371 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1788,20 +1788,20 @@ async def _validate_project_assignment( if project_obj is None: raise HTTPException( status_code=404, - detail={"error": f"Project not found, project_id={data.project_id}"}, + detail={"error": f"Project not found, project_id={data.project_id}"}, # mutable-ok: HTTPException.detail has no immutable form ) team: Final = data.team_id if "team_id" in data.model_fields_set else existing_key_row.team_id if team is None or team != project_obj.team_id: raise HTTPException( status_code=400, - detail={ + detail={ # mutable-ok: HTTPException.detail has no immutable form "error": f"Project {data.project_id} belongs to team {project_obj.team_id}. Keys can only be assigned to a project owned by their own team (key team: {team})" }, ) await _check_project_key_limits( project_id=data.project_id, data=data.model_copy( - update={ + update={ # mutable-ok: model_copy takes a dict-shaped update payload "models": data.models if "models" in data.model_fields_set else existing_key_row.models, "max_budget": data.max_budget if data.max_budget is not None else existing_key_row.max_budget, } From 22c91139877676f276f71f5a22c0e600364980c4 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 18 Sep 2026 00:05:49 +0000 Subject: [PATCH 03/12] chore(keys): format project assignment error detail Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/management_endpoints/key_management_endpoints.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 5ab5ebe0371..b73871976cf 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1788,7 +1788,9 @@ async def _validate_project_assignment( if project_obj is None: raise HTTPException( status_code=404, - detail={"error": f"Project not found, project_id={data.project_id}"}, # mutable-ok: HTTPException.detail has no immutable form + detail={ # mutable-ok: HTTPException.detail has no immutable form + "error": f"Project not found, project_id={data.project_id}" + }, ) team: Final = data.team_id if "team_id" in data.model_fields_set else existing_key_row.team_id if team is None or team != project_obj.team_id: From 2858a74f3843f39071629c550dcdc1e5007ffbd3 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 18 Sep 2026 00:18:21 +0000 Subject: [PATCH 04/12] test(keys): assert accepted project assignment flows into update payload Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/test_key_management_endpoints.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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 bd8bf4316fb..1b9642fcb94 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 @@ -19697,13 +19697,17 @@ async def test_project_assignment_to_unassigned_key_on_same_team(): token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-orbit"] ) cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + data: Final = UpdateKeyRequest(key=existing.token, project_id="project-orbit") await _validate_update_key_data( - UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + data, existing, UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), None, False, _project_assignment_prisma(), cache, ) + update: Final = await prepare_key_update_data(data=data, existing_key_row=existing) + assert update["project_id"] == "project-orbit" + @pytest.mark.asyncio async def test_project_assignment_rejects_project_on_another_team(): From 0a3557e7cad1dbce1cc5b57736ee37d100abcdc0 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 2 Oct 2026 02:40:01 +0000 Subject: [PATCH 05/12] chore(keys): drop stale LIT002 suppressions removed by main Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/key_management_endpoints.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 9903f3afeba..419e1ca5891 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1825,22 +1825,20 @@ async def _validate_project_assignment( if project_obj is None: raise HTTPException( status_code=404, - detail={ # mutable-ok: HTTPException.detail has no immutable form - "error": f"Project not found, project_id={data.project_id}" - }, + detail={"error": f"Project not found, project_id={data.project_id}"}, ) team: Final = data.team_id if "team_id" in data.model_fields_set else existing_key_row.team_id if team is None or team != project_obj.team_id: raise HTTPException( status_code=400, - detail={ # mutable-ok: HTTPException.detail has no immutable form + detail={ "error": f"Project {data.project_id} belongs to team {project_obj.team_id}. Keys can only be assigned to a project owned by their own team (key team: {team})" }, ) await _check_project_key_limits( project_id=data.project_id, data=data.model_copy( - update={ # mutable-ok: model_copy takes a dict-shaped update payload + update={ "models": data.models if "models" in data.model_fields_set else existing_key_row.models, "max_budget": data.max_budget if data.max_budget is not None else existing_key_row.max_budget, } From 1917700a576c6bc9c3f94fd1042694fc95065da2 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 2 Oct 2026 03:01:46 +0000 Subject: [PATCH 06/12] test(keys): mock writer_db team lookup used by check_db_only path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/test_key_management_endpoints.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 a18a2833f36..8f2060cc8a8 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -20491,9 +20491,9 @@ async def test_project_detachment_uses_effective_project_for_validation(project_ def _project_assignment_prisma() -> MagicMock: database = MagicMock() - database.db.litellm_teamtable.find_unique = AsyncMock( - return_value=LiteLLM_TeamTable(team_id="team-lit-5823", members=[]) - ) + team_row: Final = LiteLLM_TeamTable(team_id="team-lit-5823", members=[]) + database.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + database.writer_db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) return database From 93d579fc9c16b15113eb4ea3cecbd6a7aa30c958 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 2 Oct 2026 04:27:41 +0000 Subject: [PATCH 07/12] 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( From f508500b7aa74f47a718796d26e7040a4523ee1b Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 2 Oct 2026 04:37:45 +0000 Subject: [PATCH 08/12] chore(keys): satisfy type gate on guarded project assignment write Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 22286347bd9..a537afcec54 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -232,6 +232,10 @@ class _BudgetRowSoftBudgetCreate(TypedDict): updated_by: ReadOnly[str] +class _KeyRowDumpable(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + class _KeyUpdateTx(Protocol): @property def litellm_verificationtoken(self) -> "TableActions[prisma_models.LiteLLM_VerificationToken]": ... @@ -2468,10 +2472,10 @@ async def _apply_soft_budget_update( async def _write_guarded_project_assignment( - table: "TableActions[prisma_models.LiteLLM_VerificationToken]", + table: "TableActions[_KeyRowDumpable]", hashed_token: str, data: Mapping[str, object], -) -> "prisma_models.LiteLLM_VerificationToken | None": +) -> "_KeyRowDumpable | None": updated_count: Final = await table.update_many( where={"token": hashed_token, "project_id": None}, data=data, @@ -3621,7 +3625,7 @@ async def update_key_fn( if prisma_client is None: raise Exception("Not connected to DB!") - update_values: Final = await _handle_update_object_permission( + update_values: Final[Mapping[str, object]] = await _handle_update_object_permission( data_json=non_default_values, existing_key_row=existing_key_row, prisma_client=prisma_client, @@ -3642,7 +3646,9 @@ async def update_key_fn( else await _update_key_row_assigning_project( prisma_client=prisma_client, key=key, - update_values=update_values, + update_values=cast( # cast-ok: _handle_update_object_permission returns a bare dict + "Mapping[str, object]", update_values + ), ) if is_project_assignment else await prisma_client.update_data(token=key, data=MappingProxyType({**update_values, "token": key})) From b18987f6e8dba6f5caffc2b0ffaeaaa9a5153c71 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 2 Oct 2026 04:58:02 +0000 Subject: [PATCH 09/12] fix(keys): default soft-budget update path to unguarded write Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/management_endpoints/key_management_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a537afcec54..24fd9a32952 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2519,7 +2519,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, + expect_unassigned_project: bool = False, ) -> _KeyUpdateResult: hashed_token: Final = _hash_token_if_needed(key) key_where: Final[_KeyRowWhere] = {"token": hashed_token} From 4686180e474efb86db99a1528d0532df508221bb Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 2 Oct 2026 07:03:50 +0000 Subject: [PATCH 10/12] fix(keys): run guarded project assignment in one writer transaction Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 119 +++++++++++++++--- .../test_key_management_endpoints.py | 67 ++++++++-- 2 files changed, 157 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 24fd9a32952..f66481b0b54 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -116,6 +116,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( attach_object_permission_to_dict, handle_update_object_permission_common, invalidate_cached_object_permissions, + prepare_object_permission_upsert, validate_key_mcp_servers_against_team, validate_key_search_tools_against_team, validate_key_vector_stores_against_team, @@ -243,6 +244,9 @@ class _KeyUpdateTx(Protocol): @property def litellm_budgettable(self) -> "TableActions[prisma_models.LiteLLM_BudgetTable]": ... + @property + def litellm_objectpermissiontable(self) -> "TableActions[prisma_models.LiteLLM_ObjectPermissionTable]": ... + class _ConfigTableActions(Protocol): """Config table surface this module needs; the shared repository seam exposes no ``update``.""" @@ -2471,19 +2475,64 @@ async def _apply_soft_budget_update( return remaining +async def _apply_object_permission_update( + data_json: dict[str, object], + existing_object_permission_id: str | None, + prisma_client: PrismaClient, + table: "TableActions[prisma_models.LiteLLM_ObjectPermissionTable]", +) -> None: + """Upsert the requested object permission row inside the caller's writer transaction.""" + new_object_permission: Final = data_json.pop("object_permission", None) + if new_object_permission is None: + return + loaded: Final[object] = ( + json.loads(new_object_permission) if isinstance(new_object_permission, str) else new_object_permission + ) + grants: Final[Mapping[str, object]] = ( + cast( # cast-ok: object_permission payloads are grant-name to grant-list mappings + "Mapping[str, object]", loaded + ) + if isinstance(loaded, dict) + else MappingProxyType({}) + ) + upsert: Final = await prepare_object_permission_upsert( + new_object_permission=grants, + existing_object_permission_id=existing_object_permission_id, + prisma_client=prisma_client, + ) + row: Final = await table.upsert( + where={"object_permission_id": upsert.object_permission_id}, + data={"create": upsert.record, "update": upsert.record}, + ) + data_json["object_permission_id"] = row.object_permission_id + + async def _write_guarded_project_assignment( table: "TableActions[_KeyRowDumpable]", hashed_token: str, + existing_key_row: LiteLLM_VerificationToken, data: Mapping[str, object], ) -> "_KeyRowDumpable | None": + validated_models: Final[list[str]] = ( + cast( # cast-ok: LiteLLM_VerificationToken.models is a bare list + "list[str] | None", existing_key_row.models + ) + or [] + ) updated_count: Final = await table.update_many( - where={"token": hashed_token, "project_id": None}, + where={ + "token": hashed_token, + "project_id": None, + "team_id": existing_key_row.team_id, + "models": {"equals": validated_models}, + "max_budget": existing_key_row.max_budget, + }, data=data, ) if updated_count == 0: raise HTTPException( status_code=400, - detail="Project reassignment is not supported. Use null to detach the key.", + detail="Project assignment failed because the key was updated concurrently. Retry the request.", ) return await table.find_unique( where={"token": hashed_token}, @@ -2494,17 +2543,30 @@ async def _write_guarded_project_assignment( async def _update_key_row_assigning_project( prisma_client: PrismaClient, key: str, - update_values: Mapping[str, object], + non_default_values: Mapping[str, object], + existing_key_row: LiteLLM_VerificationToken, ) -> _KeyUpdateResult: - """Write a key update that assigns a project only while the row still has none.""" + """Assign a project in one writer transaction, guarded on the key's validated state.""" 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})) - ), - ) + async with prisma_client.tx() as tx_ctx: + tx: Final[_KeyUpdateTx] = cast( # cast-ok: the transaction object exposes the same table actions + "_KeyUpdateTx", tx_ctx + ) + update_values: Final[dict[str, object]] = dict(non_default_values) + await _apply_object_permission_update( + data_json=update_values, + existing_object_permission_id=existing_key_row.object_permission_id, + prisma_client=prisma_client, + table=tx.litellm_objectpermissiontable, + ) + updated_row: Final = await _write_guarded_project_assignment( + table=tx.litellm_verificationtoken, + hashed_token=hashed_token, + existing_key_row=existing_key_row, + 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({}) ) @@ -2525,9 +2587,17 @@ async def _update_key_row_with_soft_budget( key_where: Final[_KeyRowWhere] = {"token": hashed_token} tx: _KeyUpdateTx async with prisma_client.tx() as tx: + update_input: Final[dict[str, object]] = dict(non_default_values) + if expect_unassigned_project: + await _apply_object_permission_update( + data_json=update_input, + existing_object_permission_id=existing_key_row.object_permission_id, + prisma_client=prisma_client, + table=tx.litellm_objectpermissiontable, + ) update_values: Final = await _apply_soft_budget_update( data=data, - non_default_values=non_default_values, + non_default_values=update_input, db=tx, existing_key_row=existing_key_row, changed_by=changed_by, @@ -2540,6 +2610,7 @@ async def _update_key_row_with_soft_budget( await _write_guarded_project_assignment( table=tx.litellm_verificationtoken, hashed_token=hashed_token, + existing_key_row=existing_key_row, data=update_payload, ) if expect_unassigned_project @@ -3625,13 +3696,22 @@ async def update_key_fn( if prisma_client is None: raise Exception("Not connected to DB!") - update_values: Final[Mapping[str, object]] = await _handle_update_object_permission( - data_json=non_default_values, - existing_key_row=existing_key_row, - prisma_client=prisma_client, + is_project_assignment: Final = data.project_id is not None and existing_key_row.project_id is None + update_values: Final[Mapping[str, object]] = ( + cast( # cast-ok: prepare_key_update_data returns a bare dict + "Mapping[str, object]", non_default_values + ) + if is_project_assignment + else cast( # cast-ok: _handle_update_object_permission returns a bare dict + "Mapping[str, object]", + await _handle_update_object_permission( + data_json=non_default_values, + existing_key_row=existing_key_row, + 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, @@ -3646,9 +3726,10 @@ async def update_key_fn( else await _update_key_row_assigning_project( prisma_client=prisma_client, key=key, - update_values=cast( # cast-ok: _handle_update_object_permission returns a bare dict - "Mapping[str, object]", update_values + non_default_values=cast( # cast-ok: prepare_key_update_data returns a bare dict + "Mapping[str, object]", non_default_values ), + existing_key_row=existing_key_row, ) 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 d1caaf6447c..dfd16c23631 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -20600,17 +20600,26 @@ async def test_prepare_key_update_data_includes_new_project_assignment(): 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) + tx: Final = database.tx.return_value.__aenter__.return_value + tx.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"} + tx.litellm_verificationtoken.find_unique = AsyncMock(return_value=row) + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-orbit"] ) - update_where: Final = database.db.litellm_verificationtoken.update_many.await_args.kwargs["where"] + result: Final = await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + update_where: Final = tx.litellm_verificationtoken.update_many.await_args.kwargs["where"] assert update_where["project_id"] is None + assert update_where["team_id"] == "team-lit-5823" + assert update_where["models"] == {"equals": ["model-orbit"]} assert update_where["token"] == result["token"] assert result["data"] == {"project_id": "project-orbit"} @@ -20619,16 +20628,54 @@ async def test_project_assignment_write_requires_row_still_unassigned(): 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) + tx: Final = database.tx.return_value.__aenter__.return_value + tx.litellm_verificationtoken.update_many = AsyncMock(return_value=0) + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) with pytest.raises(HTTPException) as exc: await _update_key_row_assigning_project( - prisma_client=database, key="sk-assign", update_values={"project_id": "project-orbit"} + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, ) assert exc.value.status_code == 400 - assert "reassignment" in str(exc.value.detail) - database.db.litellm_verificationtoken.find_unique.assert_not_called() + assert "concurrently" in str(exc.value.detail) + tx.litellm_verificationtoken.find_unique.assert_not_called() + + +@pytest.mark.asyncio +async def test_project_assignment_write_upserts_object_permission_in_same_tx(): + database = MagicMock() + database.jsonify_object = lambda data: dict(data) + database.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) + tx: Final = database.tx.return_value.__aenter__.return_value + tx.litellm_verificationtoken.update_many = AsyncMock(return_value=1) + permission_row = MagicMock() + permission_row.object_permission_id = "perm-lit-01" + tx.litellm_objectpermissiontable.upsert = AsyncMock(return_value=permission_row) + row = MagicMock() + row.model_dump = MagicMock(return_value={"project_id": "project-orbit"}) + tx.litellm_verificationtoken.find_unique = AsyncMock(return_value=row) + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823" + ) + + result: Final = await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit", "object_permission": {"agents": ["agent-a"]}}, + existing_key_row=existing, + ) + + tx.litellm_objectpermissiontable.upsert.assert_awaited_once() + update_data: Final = tx.litellm_verificationtoken.update_many.await_args.kwargs["data"] + assert update_data["object_permission_id"] == "perm-lit-01" + assert "object_permission" not in update_data + assert result["data"] == {"project_id": "project-orbit"} @pytest.mark.asyncio From 31a08dad482f33710528d43e340229169670b5a0 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 2 Oct 2026 07:34:12 +0000 Subject: [PATCH 11/12] fix(keys): include object permission link in assignment guard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/management_endpoints/key_management_endpoints.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f66481b0b54..f6c229a1d90 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2526,6 +2526,7 @@ async def _write_guarded_project_assignment( "team_id": existing_key_row.team_id, "models": {"equals": validated_models}, "max_budget": existing_key_row.max_budget, + "object_permission_id": existing_key_row.object_permission_id, }, data=data, ) From 20e5241cc4344daa766fd4bbc20596eaddd274a6 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 2 Oct 2026 08:04:55 +0000 Subject: [PATCH 12/12] refactor(keys): return permission-linked payload without mutating the update dict Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 28 +++++++++++-------- 1 file changed, 17 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f6c229a1d90..ad8b423279c 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2476,15 +2476,15 @@ async def _apply_soft_budget_update( async def _apply_object_permission_update( - data_json: dict[str, object], + data_json: Mapping[str, object], existing_object_permission_id: str | None, prisma_client: PrismaClient, table: "TableActions[prisma_models.LiteLLM_ObjectPermissionTable]", -) -> None: +) -> Mapping[str, object]: """Upsert the requested object permission row inside the caller's writer transaction.""" - new_object_permission: Final = data_json.pop("object_permission", None) + new_object_permission: Final = data_json.get("object_permission") if new_object_permission is None: - return + return data_json loaded: Final[object] = ( json.loads(new_object_permission) if isinstance(new_object_permission, str) else new_object_permission ) @@ -2504,7 +2504,12 @@ async def _apply_object_permission_update( where={"object_permission_id": upsert.object_permission_id}, data={"create": upsert.record, "update": upsert.record}, ) - data_json["object_permission_id"] = row.object_permission_id + return MappingProxyType( + { + **{k: v for k, v in data_json.items() if k != "object_permission"}, + "object_permission_id": row.object_permission_id, + } + ) async def _write_guarded_project_assignment( @@ -2553,9 +2558,8 @@ async def _update_key_row_assigning_project( tx: Final[_KeyUpdateTx] = cast( # cast-ok: the transaction object exposes the same table actions "_KeyUpdateTx", tx_ctx ) - update_values: Final[dict[str, object]] = dict(non_default_values) - await _apply_object_permission_update( - data_json=update_values, + update_values: Final[Mapping[str, object]] = await _apply_object_permission_update( + data_json=non_default_values, existing_object_permission_id=existing_key_row.object_permission_id, prisma_client=prisma_client, table=tx.litellm_objectpermissiontable, @@ -2588,14 +2592,16 @@ async def _update_key_row_with_soft_budget( key_where: Final[_KeyRowWhere] = {"token": hashed_token} tx: _KeyUpdateTx async with prisma_client.tx() as tx: - update_input: Final[dict[str, object]] = dict(non_default_values) - if expect_unassigned_project: + update_input: Final[Mapping[str, object]] = ( await _apply_object_permission_update( - data_json=update_input, + data_json=non_default_values, existing_object_permission_id=existing_key_row.object_permission_id, prisma_client=prisma_client, table=tx.litellm_objectpermissiontable, ) + if expect_unassigned_project + else non_default_values + ) update_values: Final = await _apply_soft_budget_update( data=data, non_default_values=update_input,