From a27bb1088e630cd295a65ee44d6514c08c156818 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 3 Oct 2026 01:58:34 +0000 Subject: [PATCH] fix(proxy): reject key project ownership before permission and budget writes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/project_endpoints.py | 86 ++-- .../key_management_endpoints.py | 113 ++++-- .../management/test_project_lifecycle.py | 102 ++++- .../test_project_endpoints_prisma.py | 19 +- .../test_key_management_endpoints.py | 384 +++++++++++------- 5 files changed, 462 insertions(+), 242 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index 796b4acdc46..5a754141063 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -790,8 +790,16 @@ async def update_project( data, existing_project, _router_access_group_names(llm_router) ) + object_permission_data: Final = ( + data.object_permission.model_dump(exclude_none=True) if data.object_permission is not None else None + ) + object_permission_payload: Final = ( + _OBJECT_PERMISSION_PAYLOAD.validate_python(object_permission_data) if object_permission_data else None + ) + # Prepare update data update_data = _jsonified(prisma_client, data.model_dump(exclude_none=True, exclude={"project_id"})) + update_data.pop("object_permission", None) update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name # Handle budget updates @@ -801,48 +809,6 @@ async def update_project( **({"max_budget": None} if "max_budget" in data.model_fields_set and data.max_budget is None else {}), } - if budget_updates and existing_project.budget_id: - # Update existing budget - await _budget_table(prisma_client).update( - where={"budget_id": existing_project.budget_id}, - data={ - **budget_updates, - "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - }, - ) - # Remove budget fields from project update - for field in budget_updates.keys(): - update_data.pop(field, None) - - # Handle object permissions - if "object_permission" in update_data: - object_permission_data = update_data.pop("object_permission") - if object_permission_data: - object_permission_payload: Final = _OBJECT_PERMISSION_PAYLOAD.validate_python(object_permission_data) - if existing_project.object_permission_id: - # Update existing permission - await _object_permission_table(prisma_client).update( - where={"object_permission_id": existing_project.object_permission_id}, - data=object_permission_payload, - ) - else: - # Create new permission - created_permission: Final = await _object_permission_table(prisma_client).create( - data=object_permission_payload, - ) - update_data["object_permission_id"] = created_permission.object_permission_id - - # Handle metadata fields - for field in LiteLLM_ManagementEndpoint_MetadataFields: - if field in update_data: - existing_metadata = update_data.get("metadata") - metadata_dict: dict[str, object] = existing_metadata if isinstance(existing_metadata, dict) else {} - metadata_dict[field] = update_data.pop(field) - update_data["metadata"] = metadata_dict - - # Remove budget fields (following organization_endpoints.py pattern) - update_data = _remove_budget_fields_from_project_data(update_data) - if data.team_id is not None: current_project_record: Final = await _writer_project_table(prisma_client).find_unique( where={"project_id": data.project_id} @@ -870,6 +836,42 @@ async def update_project( }, ) + if budget_updates and existing_project.budget_id: + # Update existing budget + await _budget_table(prisma_client).update( + where={"budget_id": existing_project.budget_id}, + data={ + **budget_updates, + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + }, + ) + # Remove budget fields from project update + for field in budget_updates.keys(): + update_data.pop(field, None) + + if object_permission_payload is not None: + if existing_project.object_permission_id: + await _object_permission_table(prisma_client).update( + where={"object_permission_id": existing_project.object_permission_id}, + data=object_permission_payload, + ) + else: + created_permission: Final = await _object_permission_table(prisma_client).create( + data=object_permission_payload, + ) + update_data["object_permission_id"] = created_permission.object_permission_id + + # Handle metadata fields + for field in LiteLLM_ManagementEndpoint_MetadataFields: + if field in update_data: + existing_metadata = update_data.get("metadata") + metadata_dict: dict[str, object] = existing_metadata if isinstance(existing_metadata, dict) else {} + metadata_dict[field] = update_data.pop(field) + update_data["metadata"] = metadata_dict + + # Remove budget fields (following organization_endpoints.py pattern) + update_data = _remove_budget_fields_from_project_data(update_data) + # Update project updated_project: Final = await _project_table(prisma_client).update( where={"project_id": data.project_id}, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index c68316871f4..3b5b7264d1a 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -113,10 +113,11 @@ from litellm.proxy.management_helpers.access_group_key_sync import ( ) from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at from litellm.proxy.management_helpers.object_permission_utils import ( + ObjectPermissionUpsert, _set_object_permission, 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, @@ -139,6 +140,7 @@ from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.config_repository import ConfigParam, ConfigRepository from litellm.repositories.credentials_repository import CredentialsRepository from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ( DeletedVerificationTokenRepository, @@ -2502,11 +2504,6 @@ async def _update_key_row_with_soft_budget( existing_key_row=existing_key_row, changed_by=changed_by, ) - await _check_key_project_team_on_mutation( - data=data, - existing_key_row=existing_key_row, - prisma_client=prisma_client, - ) include_object_permission: Final[prisma.types.LiteLLM_VerificationTokenInclude] = {"object_permission": True} updated_row: Final = await tx.litellm_verificationtoken.update( where=key_where, @@ -2522,19 +2519,12 @@ async def _update_key_row_with_soft_budget( return result -async def _update_key_row_with_project_team_check( +async def _update_key_row( prisma_client: PrismaClient, key: str, - data: UpdateKeyRequest, update_values: Mapping[str, object], - existing_key_row: LiteLLM_VerificationToken, ) -> _KeyUpdateResult | None: key_update_data: Final = MappingProxyType({**update_values, "token": key}) - await _check_key_project_team_on_mutation( - data=data, - existing_key_row=existing_key_row, - prisma_client=prisma_client, - ) response: Final = await prisma_client.update_data(token=key, data=key_update_data) if response is None: return None @@ -2647,27 +2637,42 @@ async def prepare_key_update_data( return non_default_values -async def _handle_update_object_permission( - data_json: dict, +async def _prepare_key_update_object_permission( + object_permission_data: object, existing_key_row: LiteLLM_VerificationToken, prisma_client: PrismaClient, -) -> dict: - """Persist the requested object permission row and swap it for its id, only after the key policy allowed the write.""" - if "object_permission" not in data_json: - return data_json +) -> ObjectPermissionUpsert | None: + if object_permission_data is None: + return None - object_permission_id: Final = await handle_update_object_permission_common( - data_json=data_json, + parsed_object_permission: Final[object] = ( + json.loads(object_permission_data) if isinstance(object_permission_data, str) else object_permission_data + ) + permission_data: Final[dict[str, object]] = ( + TypeAdapter(dict[str, object]).validate_python(parsed_object_permission) + if isinstance(parsed_object_permission, dict) + else {} + ) + return await prepare_object_permission_upsert( + new_object_permission=permission_data, existing_object_permission_id=existing_key_row.object_permission_id, prisma_client=prisma_client, ) - # Add the object_permission_id to data_json if one was created/updated - if object_permission_id is not None: - data_json["object_permission_id"] = object_permission_id - verbose_proxy_logger.debug("updated object_permission_id: %s", object_permission_id) - return data_json +async def _write_prepared_key_update_object_permission( + data_json: Mapping[str, object], + upsert: ObjectPermissionUpsert | None, + prisma_client: PrismaClient, +) -> Mapping[str, object]: + if upsert is None: + return data_json + + await ObjectPermissionRepository(prisma_client).table.upsert( + where={"object_permission_id": upsert.object_permission_id}, + data={"create": upsert.record, "update": upsert.record}, + ) + return MappingProxyType({**data_json, "object_permission_id": upsert.object_permission_id}) def is_different_team(data: UpdateKeyRequest, existing_key_row: LiteLLM_VerificationToken) -> bool: @@ -2949,17 +2954,25 @@ async def _process_single_key_update( detail={"error": "Database not connected"}, ) - update_values: Final = await _handle_update_object_permission( - data_json=non_default_values, + object_permission_upsert: Final = await _prepare_key_update_object_permission( + object_permission_data=non_default_values.get("object_permission"), existing_key_row=existing_key_row, prisma_client=prisma_client, ) - _data: Final = {**update_values, "token": key_request.key} + key_update_values: Final = MappingProxyType( + {field: value for field, value in non_default_values.items() if field != "object_permission"} + ) await _check_key_project_team_on_mutation( data=key_request, existing_key_row=existing_key_row, prisma_client=prisma_client, ) + update_values: Final = await _write_prepared_key_update_object_permission( + data_json=key_update_values, + upsert=object_permission_upsert, + prisma_client=prisma_client, + ) + _data: Final = {**update_values, "token": key_request.key} response: Final[Mapping[str, object] | None] = cast( # cast-ok: every update_data branch returns a str-keyed dict "Mapping[str, object] | None", await prisma_client.update_data(token=key_request.key, data=_data), @@ -3627,11 +3640,24 @@ 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, + object_permission_upsert: Final = await _prepare_key_update_object_permission( + object_permission_data=non_default_values.get("object_permission"), existing_key_row=existing_key_row, prisma_client=prisma_client, ) + key_update_values: Final = MappingProxyType( + {field: value for field, value in non_default_values.items() if field != "object_permission"} + ) + await _check_key_project_team_on_mutation( + data=data, + existing_key_row=existing_key_row, + prisma_client=prisma_client, + ) + update_values: Final = await _write_prepared_key_update_object_permission( + data_json=key_update_values, + upsert=object_permission_upsert, + prisma_client=prisma_client, + ) changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name response: Final = ( await _update_key_row_with_soft_budget( @@ -3643,12 +3669,10 @@ async def update_key_fn( changed_by=changed_by, ) if "soft_budget" in data.model_fields_set - else await _update_key_row_with_project_team_check( + else await _update_key_row( prisma_client=prisma_client, key=key, - data=data, update_values=update_values, - existing_key_row=existing_key_row, ) ) @@ -3657,7 +3681,7 @@ async def update_key_fn( await invalidate_cached_object_permissions( object_permission_ids=( existing_key_row.object_permission_id, - non_default_values.get("object_permission_id"), + update_values.get("object_permission_id"), ), user_api_key_cache=user_api_key_cache, ) @@ -5710,7 +5734,6 @@ async def _execute_virtual_key_regeneration( new_token: Final = await get_new_token(data=data) new_token_hash: Final = hash_token(new_token) new_token_key_name: Final = abbreviate_api_key(api_key=new_token) - update_data = {"token": new_token_hash, "key_name": new_token_key_name} non_default_values = {} if data is not None: @@ -5736,18 +5759,28 @@ async def _execute_virtual_key_regeneration( request=data if data is not None else RegenerateKeyRequest(), ), ) - update_values: Final = await _handle_update_object_permission( - data_json=non_default_values, + object_permission_upsert: Final = await _prepare_key_update_object_permission( + object_permission_data=non_default_values.get("object_permission"), existing_key_row=key_in_db, prisma_client=prisma_client, ) - update_data.update(update_values) + key_update_values: Final = MappingProxyType( + {field: value for field, value in non_default_values.items() if field != "object_permission"} + ) if data is not None: await _check_key_project_team_on_mutation( data=data, existing_key_row=key_in_db, prisma_client=prisma_client, ) + update_values: Final = await _write_prepared_key_update_object_permission( + data_json=key_update_values, + upsert=object_permission_upsert, + prisma_client=prisma_client, + ) + update_data: Final = MappingProxyType( + {"token": new_token_hash, "key_name": new_token_key_name, **update_values} + ) jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data) diff --git a/tests/integration/management/test_project_lifecycle.py b/tests/integration/management/test_project_lifecycle.py index e86a989271a..25b8da7875d 100644 --- a/tests/integration/management/test_project_lifecycle.py +++ b/tests/integration/management/test_project_lifecycle.py @@ -33,6 +33,49 @@ def _key_rows(key: str) -> list[dict[str, JsonValue]]: ) +def _key_permission_and_budget_ids(key: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT object_permission_id, budget_id FROM "LiteLLM_VerificationToken" WHERE token = %s', + (sha256(key.encode()).hexdigest(),), + ) + + +def _object_permission_rows(permission_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT * FROM "LiteLLM_ObjectPermissionTable" WHERE object_permission_id = %s', + (permission_id,), + ) + + +def _budget_rows(budget_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT to_jsonb(b) AS row FROM "LiteLLM_BudgetTable" AS b WHERE budget_id = %s', + (budget_id,), + ) + + +def _clear_key_object_permission(key: str, permission_id: str) -> None: + write_rows( + 'UPDATE "LiteLLM_VerificationToken" SET object_permission_id = NULL WHERE token = %s', + (sha256(key.encode()).hexdigest(),), + ) + write_rows( + 'DELETE FROM "LiteLLM_ObjectPermissionTable" WHERE object_permission_id = %s', + (permission_id,), + ) + + +def _clear_project_object_permission(project_id: str, permission_id: str) -> None: + write_rows( + 'UPDATE "LiteLLM_ProjectTable" SET object_permission_id = NULL WHERE project_id = %s', + (project_id,), + ) + write_rows( + 'DELETE FROM "LiteLLM_ObjectPermissionTable" WHERE object_permission_id = %s', + (permission_id,), + ) + + def _deleted_key_rows(key: str) -> list[dict[str, JsonValue]]: return read_rows( 'SELECT token FROM "LiteLLM_DeletedVerificationToken" WHERE token = %s', @@ -253,7 +296,22 @@ def test_key_update_rejects_team_change_and_allows_unchanged_values(ownership_ga team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) project_a: Final = scenario.project(team_a, models=[model]) - key: Final = scenario.key(team_id=team_a, project_id=project_a, models=[model]) + key: Final = scenario.key(team_id=team_a, project_id=project_a, models=[model], soft_budget=3.0) + permission_seed: Final = ownership_gateway.request( + "POST", + "/key/update", + {"key": key, "object_permission": {"vector_stores": ["existing-store"]}}, + ) + assert permission_seed.status_code == 200, permission_seed.text + key_permission_state: Final = _key_permission_and_budget_ids(key) + assert len(key_permission_state) == 1 + permission_id: Final = string_value(key_permission_state[0]["object_permission_id"]) + budget_id: Final = string_value(key_permission_state[0]["budget_id"]) + scenario.cleanups.callback(_clear_key_object_permission, key, permission_id) + permission_before: Final = _object_permission_rows(permission_id) + budget_before: Final = _budget_rows(budget_id) + assert len(permission_before) == 1 + assert len(budget_before) == 1 aliased: Final = ownership_gateway.request("POST", "/key/update", {"key": key, "key_alias": "updated"}) assert aliased.status_code == 200, aliased.text after_alias: Final = _key_rows(key) @@ -267,9 +325,20 @@ def test_key_update_rejects_team_change_and_allows_unchanged_values(ownership_ga assert after_unchanged_team == after_alias before_reassignment: Final = _key_rows(key) assert len(before_reassignment) == 1 - reassigned: Final = ownership_gateway.request("POST", "/key/update", {"key": key, "team_id": team_b}) + reassigned: Final = ownership_gateway.request( + "POST", + "/key/update", + { + "key": key, + "team_id": team_b, + "soft_budget": 5.0, + "object_permission": {"vector_stores": ["replacement-store"]}, + }, + ) assert reassigned.status_code == 400, reassigned.text assert _key_rows(key) == before_reassignment + assert _object_permission_rows(permission_id) == permission_before + assert _budget_rows(budget_id) == budget_before def test_key_update_can_detach_project_and_change_team(ownership_gateway: Gateway) -> None: @@ -588,16 +657,41 @@ def test_project_update_rejects_moving_project_with_attached_key(ownership_gatew model: Final = scenario.model() team_a: Final = scenario.team(models=[model]) team_b: Final = scenario.team(models=[model]) - attached_project: Final = scenario.project(team_a, models=[model]) + budget: Final = scenario.budget(max_budget=3) + attached_project: Final = scenario.project(team_a, budget_id=budget, models=[model]) key: Final = scenario.key(team_id=team_a, project_id=attached_project, models=[model]) + permission_seed: Final = ownership_gateway.request( + "POST", + "/project/update", + {"project_id": attached_project, "object_permission": {"vector_stores": ["existing-store"]}}, + ) + assert permission_seed.status_code == 200, permission_seed.text + project_permission: Final = read_rows( + 'SELECT object_permission_id FROM "LiteLLM_ProjectTable" WHERE project_id = %s', + (attached_project,), + ) + assert len(project_permission) == 1 + permission_id: Final = string_value(project_permission[0]["object_permission_id"]) + scenario.cleanups.callback(_clear_project_object_permission, attached_project, permission_id) project_before: Final = _project_rows(attached_project) key_before: Final = _key_rows(key) + permission_before: Final = _object_permission_rows(permission_id) + budget_before: Final = _budget_rows(budget) moved_with_key: Final = ownership_gateway.request( - "POST", "/project/update", {"project_id": attached_project, "team_id": team_b} + "POST", + "/project/update", + { + "project_id": attached_project, + "team_id": team_b, + "max_budget": 11, + "object_permission": {"vector_stores": ["replacement-store"]}, + }, ) assert moved_with_key.status_code == 400, moved_with_key.text assert _project_rows(attached_project) == project_before assert _key_rows(key) == key_before + assert _object_permission_rows(permission_id) == permission_before + assert _budget_rows(budget) == budget_before def test_project_update_allows_moving_project_without_keys(ownership_gateway: Gateway) -> None: diff --git a/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py index 98bb2969201..cd185179ee9 100644 --- a/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py +++ b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py @@ -1277,6 +1277,15 @@ async def test_update_project_rejects_move_when_attached_teamless_key_exists( destination_team_id: Final = "team-b" mock_prisma: Final = _project_update_mocks(monkeypatch, {}) mock_prisma.db.litellm_projecttable.find_unique.return_value.team_id = "team-a" + mock_prisma.db.litellm_projecttable.find_unique.return_value.budget_id = "budget-project" + mock_prisma.db.litellm_projecttable.find_unique.return_value.object_permission_id = "permission-project" + budget_table: Final = mock.MagicMock() + budget_table.update = mock.AsyncMock() + permission_table: Final = mock.MagicMock() + permission_table.update = mock.AsyncMock() + permission_table.create = mock.AsyncMock() + mock_prisma.db.litellm_budgettable = budget_table + mock_prisma.db.litellm_objectpermissiontable = permission_table mock_prisma.db.litellm_teamtable.find_unique = mock.AsyncMock( return_value=LiteLLM_TeamTable(team_id=destination_team_id) ) @@ -1293,7 +1302,12 @@ async def test_update_project_rejects_move_when_attached_teamless_key_exists( mock_prisma.writer_db.litellm_verificationtoken.count = mock.AsyncMock(side_effect=count_teamless_keys) with pytest.raises(ProxyException) as error: - await _run_project_update(project_id, team_id=destination_team_id) + await _run_project_update( + project_id, + team_id=destination_team_id, + max_budget=50, + object_permission={"vector_stores": ["replacement-store"]}, + ) expected_detail: Final = { "error": ( @@ -1305,6 +1319,9 @@ async def test_update_project_rejects_move_when_attached_teamless_key_exists( assert expected_detail["error"] in error.value.message mock_prisma.writer_db.litellm_verificationtoken.count.assert_awaited_once() mock_prisma.db.litellm_verificationtoken.count.assert_not_awaited() + budget_table.update.assert_not_awaited() + permission_table.update.assert_not_awaited() + permission_table.create.assert_not_awaited() mock_prisma.db.litellm_projecttable.update.assert_not_awaited() 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 7e5b4987953..80c51d20f1e 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -1076,208 +1076,146 @@ async def test_key_generation_with_mcp_tool_permissions(monkeypatch): @pytest.mark.asyncio -async def test_key_update_object_permissions_existing_permission(): - """ - Test updating object permissions when a key already has an existing object_permission_id. - - This test verifies that when updating vector stores for a key that already has an - object_permission_id, the existing LiteLLM_ObjectPermissionTable record is updated - with the new permissions and the object_permission_id remains the same. - """ +async def test_key_update_prepares_existing_object_permission_without_writing(): from unittest.mock import AsyncMock, MagicMock - import pytest - - from litellm.proxy._types import ( - LiteLLM_ObjectPermissionBase, - LiteLLM_VerificationToken, - ) + from litellm.proxy._types import LiteLLM_ObjectPermissionBase, LiteLLM_VerificationToken from litellm.proxy.management_endpoints.key_management_endpoints import ( - _handle_update_object_permission, + _prepare_key_update_object_permission, + _write_prepared_key_update_object_permission, ) mock_prisma_client = AsyncMock() - - # Mock existing key with object_permission_id existing_key_row = LiteLLM_VerificationToken( token="test_token_hash", object_permission_id="existing_perm_id_123", user_id="user123", team_id=None, ) - - # Mock existing object permission record - existing_object_permission = MagicMock() - existing_object_permission.model_dump.return_value = { + existing_permission = MagicMock() + existing_permission.model_dump.return_value = { "object_permission_id": "existing_perm_id_123", - "vector_stores": ["old_store_1", "old_store_2"], + "vector_stores": ["old_store"], } + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=existing_permission) + permission_upsert = AsyncMock() + mock_prisma_client.db.litellm_objectpermissiontable.upsert = permission_upsert - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=existing_object_permission + object_permission_data = LiteLLM_ObjectPermissionBase(vector_stores=["new_store"]).model_dump( + exclude_unset=True, exclude_none=True ) - - # Mock upsert operation - updated_permission = MagicMock() - updated_permission.object_permission_id = "existing_perm_id_123" - mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( - return_value=updated_permission - ) - - # Test data with new object permission - data_json = { - "object_permission": LiteLLM_ObjectPermissionBase( - vector_stores=["new_store_1", "new_store_2", "new_store_3"] - ).model_dump(exclude_unset=True, exclude_none=True), - "user_id": "user123", - } - - # Call the function - result = await _handle_update_object_permission( - data_json=data_json, + upsert = await _prepare_key_update_object_permission( + object_permission_data=object_permission_data, existing_key_row=existing_key_row, prisma_client=mock_prisma_client, ) - # Verify the object_permission was removed from data_json and object_permission_id was set - assert "object_permission" not in result - assert result["object_permission_id"] == "existing_perm_id_123" - - # Verify database operations were called correctly - mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with( + assert upsert is not None + assert upsert.object_permission_id == "existing_perm_id_123" + assert upsert.record["vector_stores"] == ["new_store"] + permission_upsert.assert_not_awaited() + mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_awaited_once_with( where={"object_permission_id": "existing_perm_id_123"} ) - mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once() + + result = await _write_prepared_key_update_object_permission( + data_json={"user_id": "user123"}, + upsert=upsert, + prisma_client=mock_prisma_client, + ) + + assert result == {"user_id": "user123", "object_permission_id": "existing_perm_id_123"} + permission_upsert.assert_awaited_once_with( + where={"object_permission_id": "existing_perm_id_123"}, + data={"create": upsert.record, "update": upsert.record}, + ) @pytest.mark.asyncio -async def test_key_update_object_permissions_no_existing_permission(): - """ - Test creating object permissions when a key has no existing object_permission_id. +async def test_key_update_prepares_json_object_permission_and_upserts_new_row(): + import json + from unittest.mock import AsyncMock - This test verifies that when updating object permissions for a key that has - object_permission_id set to None, a new entry is created in the - LiteLLM_ObjectPermissionTable and the key is updated with the new object_permission_id. - """ - from unittest.mock import AsyncMock, MagicMock - - import pytest - - from litellm.proxy._types import ( - LiteLLM_ObjectPermissionBase, - LiteLLM_VerificationToken, - ) + from litellm.proxy._types import LiteLLM_VerificationToken from litellm.proxy.management_endpoints.key_management_endpoints import ( - _handle_update_object_permission, + _prepare_key_update_object_permission, + _write_prepared_key_update_object_permission, ) mock_prisma_client = AsyncMock() - - existing_key_row_no_perm = LiteLLM_VerificationToken( + existing_key_row = LiteLLM_VerificationToken( token="test_token_hash_2", object_permission_id=None, user_id="user456", team_id=None, ) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) + permission_upsert = AsyncMock() + mock_prisma_client.db.litellm_objectpermissiontable.upsert = permission_upsert - # Mock find_unique to return None (no existing permission) - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=None - ) - - # Mock upsert to create new record - new_permission = MagicMock() - new_permission.object_permission_id = "new_perm_id_456" - mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( - return_value=new_permission - ) - - data_json = { - "object_permission": LiteLLM_ObjectPermissionBase( - vector_stores=["brand_new_store"] - ).model_dump(exclude_unset=True, exclude_none=True), - "user_id": "user456", - } - - result = await _handle_update_object_permission( - data_json=data_json, - existing_key_row=existing_key_row_no_perm, + upsert = await _prepare_key_update_object_permission( + object_permission_data=json.dumps({"vector_stores": ["brand_new_store"]}), + existing_key_row=existing_key_row, prisma_client=mock_prisma_client, ) - # Verify new object_permission_id was set - assert "object_permission" not in result - assert result["object_permission_id"] == "new_perm_id_456" - # Verify upsert was called to create new record - mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once() + assert upsert is not None + assert upsert.record["vector_stores"] == ["brand_new_store"] + permission_upsert.assert_not_awaited() + result = await _write_prepared_key_update_object_permission( + data_json={}, + upsert=upsert, + prisma_client=mock_prisma_client, + ) + + assert result["object_permission_id"] == upsert.object_permission_id + permission_upsert.assert_awaited_once_with( + where={"object_permission_id": upsert.object_permission_id}, + data={"create": upsert.record, "update": upsert.record}, + ) @pytest.mark.asyncio -async def test_key_update_object_permissions_missing_permission_record(): - """ - Test creating object permissions when existing object_permission_id record is not found. +async def test_key_update_recreates_missing_object_permission_with_existing_id(): + from unittest.mock import AsyncMock - This test verifies that when updating object permissions for a key that has an - object_permission_id but the corresponding record cannot be found in the database, - a new entry is created in the LiteLLM_ObjectPermissionTable with the new permissions. - """ - from unittest.mock import AsyncMock, MagicMock - - import pytest - - from litellm.proxy._types import ( - LiteLLM_ObjectPermissionBase, - LiteLLM_VerificationToken, - ) + from litellm.proxy._types import LiteLLM_ObjectPermissionBase, LiteLLM_VerificationToken from litellm.proxy.management_endpoints.key_management_endpoints import ( - _handle_update_object_permission, + _prepare_key_update_object_permission, + _write_prepared_key_update_object_permission, ) mock_prisma_client = AsyncMock() - - existing_key_row_missing_perm = LiteLLM_VerificationToken( + existing_key_row = LiteLLM_VerificationToken( token="test_token_hash_3", object_permission_id="missing_perm_id_789", user_id="user789", team_id=None, ) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) + permission_upsert = AsyncMock() + mock_prisma_client.db.litellm_objectpermissiontable.upsert = permission_upsert - # Mock find_unique to return None (permission record not found) - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=None - ) - - # Mock upsert to create new record - new_permission = MagicMock() - new_permission.object_permission_id = "recreated_perm_id_789" - mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( - return_value=new_permission - ) - - data_json = { - "object_permission": LiteLLM_ObjectPermissionBase( - vector_stores=["recreated_store"] - ).model_dump(exclude_unset=True, exclude_none=True), - "user_id": "user789", - } - - result = await _handle_update_object_permission( - data_json=data_json, - existing_key_row=existing_key_row_missing_perm, + upsert = await _prepare_key_update_object_permission( + object_permission_data=LiteLLM_ObjectPermissionBase(vector_stores=["recreated_store"]).model_dump( + exclude_unset=True, exclude_none=True + ), + existing_key_row=existing_key_row, prisma_client=mock_prisma_client, ) - # Verify new object_permission_id was set - assert "object_permission" not in result - assert result["object_permission_id"] == "recreated_perm_id_789" - - # Verify find_unique was called with the missing permission ID - mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with( - where={"object_permission_id": "missing_perm_id_789"} + assert upsert is not None + assert upsert.object_permission_id == "missing_perm_id_789" + permission_upsert.assert_not_awaited() + await _write_prepared_key_update_object_permission( + data_json={}, + upsert=upsert, + prisma_client=mock_prisma_client, + ) + permission_upsert.assert_awaited_once_with( + where={"object_permission_id": "missing_perm_id_789"}, + data={"create": upsert.record, "update": upsert.record}, ) - - # Verify upsert was called to create new record - mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once() @pytest.mark.asyncio @@ -7291,7 +7229,7 @@ async def test_bulk_update_keys_object_permission_is_granted_not_dropped(monkeyp 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 written["object_permission_id"] == upserted["object_permission_id"] assert not {"max_budget", "team_id", "budget_id"} & written.keys() @@ -13170,6 +13108,7 @@ async def test_execute_virtual_key_regeneration_hides_the_untouched_modal_expiry _POLICY_DENIAL_MESSAGE = "key duration must be 7d or less" _POLICY_HASHED_TOKEN = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b" _POLICY_GENERATED_KEY = {"key": "sk-test-key", "expires": None, "user_id": "test-user", "team_id": None} +_OBJECT_PERMISSION_ID_AFTER_POLICY = "perm-after-policy" def _seven_day_policy(received: list[CustomKeyPolicyRequest]): @@ -13311,7 +13250,11 @@ async def test_regenerate_without_changes_still_runs_custom_key_policy(data): def _policy_existing_team_key() -> LiteLLM_VerificationToken: return LiteLLM_VerificationToken( - token=_POLICY_HASHED_TOKEN, user_id="test-user", team_id="team-a", max_budget=200.0 + token=_POLICY_HASHED_TOKEN, + user_id="test-user", + team_id="team-a", + max_budget=200.0, + object_permission_id=_OBJECT_PERMISSION_ID_AFTER_POLICY, ) @@ -13457,9 +13400,6 @@ async def test_process_single_key_update_rejects_when_custom_key_policy_denies() assert [policy_request.operation for policy_request in received] == ["update"] -_OBJECT_PERMISSION_ID_AFTER_POLICY = "perm-after-policy" - - def _record_object_permission_writes(mock_prisma_client: AsyncMock, events: list[str]) -> None: mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) @@ -13579,9 +13519,12 @@ async def test_regenerate_writes_the_object_permission_row_only_after_the_policy events: list[str] = [] _record_object_permission_writes(mock_prisma_client, events) data = RegenerateKeyRequest(max_budget=50.0, object_permission=LiteLLM_ObjectPermissionBase(vector_stores=["vs-1"])) + existing_key: Final = _make_regenerate_existing_key().model_copy( + update={"object_permission_id": _OBJECT_PERMISSION_ID_AFTER_POLICY} + ) with _regenerate_policy_mocks(_recording_policy(events, allowed=True), AsyncMock(), AsyncMock()): - await _regenerate_under_policy(mock_prisma_client, _make_regenerate_existing_key(), data) + await _regenerate_under_policy(mock_prisma_client, existing_key, data) _assert_permission_row_written_after_policy( events, mock_prisma_client.db.litellm_verificationtoken.update.await_args.kwargs["data"] @@ -20960,11 +20903,25 @@ async def test_key_update_rejects_team_change_for_project_bound_key(monkeypatch: token="hashed-key", team_id=_OWNERSHIP_KEY_TEAM, project_id=_OWNED_PROJECT, + object_permission_id="permission-update", + budget_id="budget-update", ) mock_prisma_client: Final = _wire_update_key_fn(monkeypatch, existing_key_row) + existing_permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="permission-update", + vector_stores=["existing-store"], + ) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=existing_permission) + permission_upsert: Final = AsyncMock() + mock_prisma_client.db.litellm_objectpermissiontable.upsert = permission_upsert + permission_before: Final = existing_permission.model_dump() monkeypatch.setattr( "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", - AsyncMock(return_value=LiteLLM_TeamTable(team_id=_OWNERSHIP_DESTINATION_TEAM, team_members=[])), + AsyncMock( + return_value=LiteLLM_TeamTable(team_id=_OWNERSHIP_DESTINATION_TEAM).model_copy( + update={"team_members": []} + ) + ), ) mock_prisma_client.writer_db.litellm_projecttable.find_unique = AsyncMock( return_value=LiteLLM_ProjectTable(project_id=_OWNED_PROJECT, team_id=_OWNERSHIP_PROJECT_TEAM) @@ -20976,7 +20933,12 @@ async def test_key_update_rejects_team_change_for_project_bound_key(monkeypatch: with pytest.raises((HTTPException, ProxyException)) as error: await update_key_fn( request=mock_request, - data=UpdateKeyRequest(key="sk-key", team_id=_OWNERSHIP_DESTINATION_TEAM), + data=UpdateKeyRequest( + key="sk-key", + team_id=_OWNERSHIP_DESTINATION_TEAM, + soft_budget=5.0, + object_permission=LiteLLM_ObjectPermissionBase(vector_stores=["replacement-store"]), + ), user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin"), litellm_changed_by=None, ) @@ -20987,6 +20949,80 @@ async def test_key_update_rejects_team_change_for_project_bound_key(monkeypatch: f"but the key belongs to {_OWNERSHIP_DESTINATION_TEAM}" ) assert expected_detail in str(getattr(error.value, "detail", None) or getattr(error.value, "message", None)) + assert existing_permission.model_dump() == permission_before + permission_upsert.assert_not_awaited() + mock_prisma_client.tx.assert_not_called() + mock_prisma_client.update_data.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_key_update_ambiguous_permission_error_precedes_project_ownership( + monkeypatch: pytest.MonkeyPatch, +) -> None: + existing_key_row: Final = LiteLLM_VerificationToken( + token="hashed-key", + team_id=_OWNERSHIP_KEY_TEAM, + project_id=_OWNED_PROJECT, + object_permission_id="permission-ambiguous", + ) + mock_prisma_client: Final = _wire_update_key_fn(monkeypatch, existing_key_row) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( + return_value=LiteLLM_ObjectPermissionTable( + object_permission_id="permission-ambiguous", + mcp_tool_permissions={}, + ) + ) + permission_upsert: Final = AsyncMock() + mock_prisma_client.db.litellm_objectpermissiontable.upsert = permission_upsert + mock_prisma_client.db.litellm_mcpservertable.find_many = AsyncMock( + return_value=[ + MagicMock(server_id="wiki-a-id", alias="wiki", server_name="wiki-a"), + MagicMock(server_id="wiki-b-id", alias="wiki", server_name="wiki-b"), + ] + ) + mock_prisma_client.writer_db.litellm_projecttable.find_unique = AsyncMock( + return_value=LiteLLM_ProjectTable(project_id=_OWNED_PROJECT, team_id=_OWNERSHIP_PROJECT_TEAM) + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + AsyncMock( + return_value=LiteLLM_TeamTable( + team_id=_OWNERSHIP_DESTINATION_TEAM, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="permission-team", + mcp_servers=["wiki-a-id", "wiki-b-id"], + ), + ) + ), + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock()) + + with pytest.raises((HTTPException, ProxyException)) as error: + await update_key_fn( + request=MagicMock(query_params={}), + data=UpdateKeyRequest( + key="sk-key", + team_id=_OWNERSHIP_DESTINATION_TEAM, + object_permission=LiteLLM_ObjectPermissionBase( + mcp_tool_permissions={"wiki": ["read_wiki_structure"]} + ), + ), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin"), + litellm_changed_by=None, + ) + + assert str(getattr(error.value, "status_code", None) or getattr(error.value, "code", None)) == "400" + expected_detail: Final = { + "error": ( + "Ambiguous mcp_tool_permissions key: 'wiki' matches MCP servers ['wiki-a-id', 'wiki-b-id']. " + "Key tool permissions by server_id when servers share a name or alias." + ) + } + assert getattr(error.value, "detail", None) == expected_detail or getattr(error.value, "message", None) == str( + expected_detail + ) + permission_upsert.assert_not_awaited() + mock_prisma_client.writer_db.litellm_projecttable.find_unique.assert_not_awaited() mock_prisma_client.update_data.assert_not_awaited() @@ -21070,11 +21106,23 @@ async def test_bulk_key_update_rejects_project_team_change_and_allows_other_fiel models=[], team_id=_OWNERSHIP_PROJECT_TEAM, project_id=_OWNED_PROJECT, + object_permission_id="permission-bulk", + budget_id="budget-bulk", ) mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing_key_row) mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( return_value=LiteLLM_TeamTable(team_id=_OWNERSHIP_DESTINATION_TEAM) ) + existing_permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="permission-bulk", + vector_stores=["existing-store"], + ) + permission_before: Final = existing_permission.model_dump() + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=existing_permission) + permission_upsert: Final = AsyncMock() + mock_prisma_client.db.litellm_objectpermissiontable.upsert = permission_upsert + budget_update: Final = AsyncMock() + mock_prisma_client.db.litellm_budgettable.update = budget_update mock_prisma_client.get_data = AsyncMock(return_value=existing_key_row) updated_key: Final = MagicMock() updated_key.model_dump.return_value = { @@ -21089,7 +21137,14 @@ async def test_bulk_key_update_rejects_project_team_change_and_allows_other_fiel request: Final = BulkUpdateKeyRequest( keys=[ - BulkUpdateKeyRequestItem(key="sk-bulk-key", team_id=_OWNERSHIP_DESTINATION_TEAM), + BulkUpdateKeyRequestItem.model_validate( + { + "key": "sk-bulk-key", + "team_id": _OWNERSHIP_DESTINATION_TEAM, + "soft_budget": 5.0, + "object_permission": LiteLLM_ObjectPermissionBase(vector_stores=["replacement-store"]), + } + ), BulkUpdateKeyRequestItem(key="sk-bulk-key", max_budget=10.0, tags=["bulk-update"]), ] ) @@ -21136,6 +21191,10 @@ async def test_bulk_key_update_rejects_project_team_change_and_allows_other_fiel assert response.successful_updates[0].key == "sk-bulk-key" assert response.successful_updates[0].key_info["max_budget"] == 10.0 assert response.successful_updates[0].key_info["tags"] == ["bulk-update"] + assert existing_permission.model_dump() == permission_before + permission_upsert.assert_not_awaited() + budget_update.assert_not_awaited() + mock_prisma_client.tx.assert_not_called() mock_prisma_client.update_data.assert_awaited_once() @@ -21199,7 +21258,11 @@ async def test_regenerate_checks_project_team_ownership( expected_status: int | None, expected_updates: int, ) -> None: - existing_key: Final = LiteLLM_VerificationToken(token="abc123", team_id=key_team_id) + existing_key: Final = LiteLLM_VerificationToken( + token="abc123", + team_id=key_team_id, + object_permission_id="permission-regenerate", + ) mock_prisma_client: Final = _make_regenerate_mock_prisma() mock_prisma_client.writer_db = MagicMock() mock_prisma_client.writer_db.litellm_projecttable = MagicMock() @@ -21212,6 +21275,14 @@ async def test_regenerate_checks_project_team_ownership( deprecated_key_table: Final = MagicMock() deprecated_key_table.upsert = AsyncMock() mock_prisma_client.db.litellm_deprecatedverificationtoken = deprecated_key_table + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( + return_value=LiteLLM_ObjectPermissionTable( + object_permission_id="permission-regenerate", + vector_stores=["existing-store"], + ) + ) + permission_upsert: Final = AsyncMock() + mock_prisma_client.db.litellm_objectpermissiontable.upsert = permission_upsert user_api_key_cache: Final = await _cache_with_project(_OWNED_PROJECT, [], team_id="team-b") async def regenerate() -> None: @@ -21234,6 +21305,7 @@ async def test_regenerate_checks_project_team_ownership( data=RegenerateKeyRequest( project_id=_OWNED_PROJECT, grace_period="1h" if expected_status is not None else None, + object_permission=LiteLLM_ObjectPermissionBase(vector_stores=["replacement-store"]), ), user_api_key_dict=_make_regenerate_user_api_key_dict(), litellm_changed_by=None, @@ -21248,8 +21320,10 @@ async def test_regenerate_checks_project_team_ownership( assert "belongs to team team-b, but the key belongs to team-a" in str(error.value.detail) deleted_history_table.create_many.assert_not_awaited() deprecated_key_table.upsert.assert_not_awaited() + permission_upsert.assert_not_awaited() else: await regenerate() + permission_upsert.assert_awaited_once() assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == expected_updates