Merge 20e5241cc4 into 65b0557f80

litellm/proxy/_types.py

@@ -1457,7 +1457,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")

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -117,6 +117,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,
@@ -235,6 +236,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]": ...
@@ -242,6 +247,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``."""
@@ -1810,6 +1818,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,
@@ -2429,6 +2481,110 @@ async def _apply_soft_budget_update(
     return remaining
 
 
+async def _apply_object_permission_update(
+    data_json: Mapping[str, object],
+    existing_object_permission_id: str | None,
+    prisma_client: PrismaClient,
+    table: "TableActions[prisma_models.LiteLLM_ObjectPermissionTable]",
+) -> Mapping[str, object]:
+    """Upsert the requested object permission row inside the caller's writer transaction."""
+    new_object_permission: Final = data_json.get("object_permission")
+    if new_object_permission is None:
+        return data_json
+    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},
+    )
+    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(
+    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,
+            "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,
+    )
+    if updated_count == 0:
+        raise HTTPException(
+            status_code=400,
+            detail="Project assignment failed because the key was updated concurrently. Retry the request.",
+        )
+    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,
+    non_default_values: Mapping[str, object],
+    existing_key_row: LiteLLM_VerificationToken,
+) -> _KeyUpdateResult:
+    """Assign a project in one writer transaction, guarded on the key's validated state."""
+    hashed_token: Final = _hash_token_if_needed(key)
+    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[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,
+        )
+        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({})
+    )
+    result: Final[_KeyUpdateResult] = {"token": hashed_token, "data": updated_data}
+    return result
+
+
 async def _update_key_row_with_soft_budget(
     prisma_client: PrismaClient,
     key: str,
@@ -2436,25 +2592,46 @@ 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 = False,
 ) -> _KeyUpdateResult:
     hashed_token: Final = _hash_token_if_needed(key)
     key_where: Final[_KeyRowWhere] = {"token": hashed_token}
     tx: _KeyUpdateTx
     async with prisma_client.tx() as tx:
+        update_input: 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,
+            )
+            if expect_unassigned_project
+            else non_default_values
+        )
         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,
         )
         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,
+                existing_key_row=existing_key_row,
+                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({})
@@ -3111,10 +3288,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
 
     acting_as_team_admin: Final = await _acting_as_team_admin_for_key_update(
@@ -3383,7 +3562,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`.
     - end_user_budget_id: Optional[str] - Proxy admin only. Budget id applied to end users first seen through this key that carry no budget of their own. Omit to keep the current value, pass an empty string to clear it.
@@ -3530,10 +3709,20 @@ async def update_key_fn(
         if prisma_client is None:
             raise Exception("Not connected to DB!")
 
-        update_values: Final = 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
         response: Final = (
@@ -3544,8 +3733,18 @@ 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,
+                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}))
         )
 

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,
@@ -20546,6 +20547,194 @@ 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()
+    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
+
+
+@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"])
+    data: Final = UpdateKeyRequest(key=existing.token, project_id="project-orbit")
+
+    await _validate_update_key_data(
+        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():
+    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_project_assignment_write_requires_row_still_unassigned():
+    database = MagicMock()
+    database.jsonify_object = lambda data: dict(data)
+    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"})
+    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"]
+    )
+
+    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"}
+
+
+@pytest.mark.asyncio
+async def test_project_assignment_write_rejects_when_row_gained_project():
+    database = MagicMock()
+    database.jsonify_object = lambda data: dict(data)
+    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",
+            non_default_values={"project_id": "project-orbit"},
+            existing_key_row=existing,
+        )
+
+    assert exc.value.status_code == 400
+    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
 async def test_key_creator_cannot_detach_project_without_admin_access():
     existing: Final = LiteLLM_VerificationToken(

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

@@ -8697,7 +8697,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`.
          *     - end_user_budget_id: Optional[str] - Proxy admin only. Budget id applied to end users first seen through this key that carry no budget of their own. Omit to keep the current value, pass an empty string to clear it.
@@ -47586,7 +47586,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 */