From 118fc7a7aec097707ea9b6886335784bc7e32e5e Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 27 Aug 2026 05:25:46 +0000 Subject: [PATCH 1/2] fix(key_management): allow /key/update to keep or shrink MCP server grants the key already holds Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 2 + .../object_permission_utils.py | 47 ++++++- .../test_key_management_endpoints.py | 64 +++++++++ .../test_object_permission_utils.py | 121 ++++++++++++++++++ 4 files changed, 233 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d1c083529192..9dda1a9443fd 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2442,11 +2442,13 @@ async def _validate_mcp_servers_for_key_update( check_db_only=True, ) object_permission_dict: Final = _object_permission_to_dict(data.object_permission) + team_unchanged: Final = data.team_id is None or data.team_id == existing_key_row.team_id normalized_object_permission: Final = await validate_key_mcp_servers_against_team( object_permission=object_permission_dict, team_obj=effective_team_obj, prisma_client=prisma_client, is_proxy_admin=is_proxy_admin, + existing_key_object_permission_id=existing_key_row.object_permission_id if team_unchanged else None, ) await validate_key_search_tools_against_team( object_permission=object_permission_dict, diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index fb64914f6f51..488ff285504a 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -447,6 +447,41 @@ async def enforce_all_proxy_mcp_servers_grant_is_admin_only( ) +async def _get_grandfathered_key_mcp_server_ids( + object_permission_id: str | None, + prisma_client: PrismaClient | None, +) -> frozenset[str]: + """ + Resolve the canonical MCP server IDs a key's stored object_permission already + grants. Updates that keep or shrink those grants stay valid even when the + team allowlist has since changed; sentinels are excluded so they cannot + grandfather anything. + """ + if not object_permission_id or prisma_client is None: + return frozenset() + existing: Final = await ObjectPermissionRepository(prisma_client).table.find_unique( + where={"object_permission_id": object_permission_id}, + ) + if existing is None: + return frozenset() + raw_tool_perms: Final = existing.mcp_tool_permissions or {} + tool_perm_keys: Final[frozenset[str]] = frozenset( + json.loads(raw_tool_perms).keys() if isinstance(raw_tool_perms, str) else raw_tool_perms.keys() + ) + identifiers: Final = (frozenset(existing.mcp_servers or []) | tool_perm_keys) - { + SpecialMCPServerNames.no_mcp_servers.value, + SpecialMCPServerName.all_proxy_servers.value, + } + return frozenset( + _flatten_resolved_mcp_server_ids( + await _resolve_mcp_server_identifiers_to_ids( + identifiers=set(identifiers), + prisma_client=prisma_client, + ) + ) + ) + + async def _get_team_allowed_mcp_servers( team_obj: Optional["LiteLLM_TeamTableCachedObj"], prisma_client: PrismaClient | None = None, @@ -527,10 +562,16 @@ async def validate_key_mcp_servers_against_team( team_obj: Optional["LiteLLM_TeamTableCachedObj"], prisma_client: PrismaClient | None = None, is_proxy_admin: bool = False, + existing_key_object_permission_id: str | None = None, ) -> ObjectPermissionDict | None: """ Validate that MCP servers requested on a key are within the allowed scope. + When ``existing_key_object_permission_id`` is provided (key updates), servers + the key already holds are grandfathered: keeping or removing them stays valid + even if the team allowlist has since shrunk, while adding new servers outside + the allowlist is still rejected. + Rules: - If key is in a team: key's mcp_servers must be a subset of (team's allowed servers + allow_all_keys servers) @@ -589,7 +630,11 @@ async def validate_key_mcp_servers_against_team( if teamless_admin_assignment: allowed_servers = all_allowed_servers | active_requested_servers - disallowed_servers: Final = active_requested_servers - allowed_servers + grandfathered_servers: Final = await _get_grandfathered_key_mcp_server_ids( + object_permission_id=existing_key_object_permission_id, + prisma_client=prisma_client, + ) + disallowed_servers: Final = active_requested_servers - allowed_servers - grandfathered_servers if disallowed_servers: if team_obj is not None: team_id = team_obj.team_id 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 a37c4f72b3d9..34869efb4912 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 @@ -730,6 +730,70 @@ async def test_update_key_personal_non_admin_denied_vector_stores(monkeypatch): assert "Vector stores" in str(exc.value.detail) +@pytest.mark.asyncio +async def test_update_key_grandfathers_existing_mcp_servers(monkeypatch): + """/key/update on a team key that already holds MCP servers outside the + team allowlist must accept re-sent or shrunk grants (LIT-6062). The wrapper + must pass the existing key's object_permission_id into the validator when + the team is unchanged.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import ( + LiteLLM_ObjectPermissionBase, + UpdateKeyRequest, + ) + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_mcp_servers_for_key_update, + ) + + existing_row = MagicMock() + existing_row.mcp_servers = ["server-a", "server-b"] + existing_row.mcp_tool_permissions = {} + mock_prisma = MagicMock() + mock_prisma.db.litellm_objectpermissiontable.find_unique = AsyncMock( + return_value=existing_row + ) + mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + + team_obj = MagicMock() + team_obj.team_id = "team-1" + team_obj.object_permission = None + + existing_key_row = MagicMock( + team_id="team-1", + object_permission_id="perm-1", + ) + + mock_server_a = MagicMock() + mock_server_a.server_id = "server-a" + mock_server_b = MagicMock() + mock_server_b.server_id = "server-b" + mock_mgr = MagicMock() + mock_mgr.get_registry.return_value = { + "server-a": mock_server_a, + "server-b": mock_server_b, + } + mock_mgr.get_allow_all_keys_server_ids.return_value = [] + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_mgr, + ) + + result = await _validate_mcp_servers_for_key_update( + data=UpdateKeyRequest( + key="sk-team-key", + object_permission=LiteLLM_ObjectPermissionBase(mcp_servers=["server-a"]), + ), + team_obj=team_obj, + existing_key_row=existing_key_row, + prisma_client=mock_prisma, + user_api_key_cache=MagicMock(), + is_proxy_admin=False, + ) + assert result is not None + assert result["mcp_servers"] == ["server-a"] + + @pytest.mark.asyncio async def test_update_key_personal_non_admin_denied_access_groups( monkeypatch, diff --git a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py index b129ad0f659c..01f3eff6b5d5 100644 --- a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py @@ -1213,6 +1213,127 @@ async def test_empty_object_permission_passes_for_personal_non_admin(): ) +# ---- Tests for grandfathering existing key MCP servers on /key/update (LIT-6062) ---- + + +def _make_grandfather_prisma(mcp_servers=None, mcp_tool_permissions=None): + """Mock prisma client whose object permission table returns the key's existing row.""" + existing_row = MagicMock() + existing_row.mcp_servers = mcp_servers or [] + existing_row.mcp_tool_permissions = mcp_tool_permissions or {} + mock_prisma = MagicMock() + mock_prisma.db.litellm_objectpermissiontable.find_unique = AsyncMock( + return_value=existing_row + ) + mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + return mock_prisma + + +def _patch_grandfather_env(monkeypatch, mock_mgr): + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_mgr, + ) + monkeypatch.setattr( + "litellm.proxy.management_helpers.object_permission_utils._get_allow_all_keys_server_ids", + lambda: set(), + ) + + +@pytest.mark.asyncio +async def test_validate_key_update_grandfathers_existing_servers(monkeypatch): + """A key already holding servers outside the team allowlist can re-send or + shrink those grants on /key/update without a 403 (LIT-6062).""" + _patch_grandfather_env(monkeypatch, _make_mock_mcp_manager("server-a", "server-b")) + team_obj = _make_team_obj(mcp_servers=[]) + mock_prisma = _make_grandfather_prisma(mcp_servers=["server-a", "server-b"]) + resend = await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["server-a", "server-b"]}, + team_obj=team_obj, + prisma_client=mock_prisma, + existing_key_object_permission_id="perm-1", + ) + assert sorted(resend["mcp_servers"]) == ["server-a", "server-b"] + shrink = await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["server-a"]}, + team_obj=team_obj, + prisma_client=mock_prisma, + existing_key_object_permission_id="perm-1", + ) + assert shrink["mcp_servers"] == ["server-a"] + + +@pytest.mark.asyncio +async def test_validate_key_update_grandfather_does_not_allow_new_servers(monkeypatch): + """Grandfathering only covers servers the key already holds; adding a new + server outside the team allowlist still raises 403.""" + _patch_grandfather_env(monkeypatch, _make_mock_mcp_manager("server-a", "server-new")) + team_obj = _make_team_obj(mcp_servers=[]) + mock_prisma = _make_grandfather_prisma(mcp_servers=["server-a"]) + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["server-a", "server-new"]}, + team_obj=team_obj, + prisma_client=mock_prisma, + existing_key_object_permission_id="perm-1", + ) + assert exc_info.value.status_code == 403 + assert "server-new" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_validate_key_update_without_existing_permission_still_raises(monkeypatch): + """Without an existing permission row id (new grants or team change) the + subset check stays strict.""" + _patch_grandfather_env(monkeypatch, _make_mock_mcp_manager("server-a")) + team_obj = _make_team_obj(mcp_servers=[]) + mock_prisma = _make_grandfather_prisma(mcp_servers=["server-a"]) + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["server-a"]}, + team_obj=team_obj, + prisma_client=mock_prisma, + existing_key_object_permission_id=None, + ) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_validate_key_update_grandfathers_tool_permission_keys(monkeypatch): + """Servers granted only via mcp_tool_permissions keys on the existing row + (stored as a JSON string) are grandfathered too.""" + _patch_grandfather_env(monkeypatch, _make_mock_mcp_manager("server-a")) + team_obj = _make_team_obj(mcp_servers=[]) + mock_prisma = _make_grandfather_prisma( + mcp_tool_permissions=json.dumps({"server-a": ["tool1"]}) + ) + result = await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["server-a"]}, + team_obj=team_obj, + prisma_client=mock_prisma, + existing_key_object_permission_id="perm-1", + ) + assert result["mcp_servers"] == ["server-a"] + + +@pytest.mark.asyncio +async def test_validate_key_update_sentinels_do_not_grandfather(monkeypatch): + """Sentinels stored on the existing row must not grandfather anything.""" + _patch_grandfather_env(monkeypatch, _make_mock_mcp_manager("server-a")) + team_obj = _make_team_obj(mcp_servers=[]) + mock_prisma = _make_grandfather_prisma( + mcp_servers=[SpecialMCPServerName.all_proxy_servers.value, "no-mcp-servers"] + ) + with pytest.raises(HTTPException) as exc_info: + await validate_key_mcp_servers_against_team( + object_permission={"mcp_servers": ["server-a"]}, + team_obj=team_obj, + prisma_client=mock_prisma, + existing_key_object_permission_id="perm-1", + ) + assert exc_info.value.status_code == 403 + + def test_object_permission_dict_mirrors_pydantic_model(): """ObjectPermissionDict must stay field-for-field aligned with LiteLLM_ObjectPermissionBase. If a new field is added to the Pydantic From ede05eb1d023bf48f3655b817bfcfa6a6c5394cd Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 27 Aug 2026 05:57:10 +0000 Subject: [PATCH 2/2] fix(key_management): reuse key row's included object_permission instead of a second lookup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 4 +-- .../object_permission_utils.py | 19 ++++------- .../test_key_management_endpoints.py | 8 ++--- .../test_object_permission_utils.py | 33 +++++++++---------- 4 files changed, 27 insertions(+), 37 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 9dda1a9443fd..97999cbb6d72 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2190,7 +2190,7 @@ async def _get_and_validate_existing_key( existing_key_row: Final[LiteLLM_VerificationToken | None] = await _prisma_table( VerificationTokenRepository(prisma_client) - ).find_unique(where={"token": hashed_token}) + ).find_unique(where={"token": hashed_token}, include={"object_permission": True}) if existing_key_row is None: raise ProxyException( @@ -2448,7 +2448,7 @@ async def _validate_mcp_servers_for_key_update( team_obj=effective_team_obj, prisma_client=prisma_client, is_proxy_admin=is_proxy_admin, - existing_key_object_permission_id=existing_key_row.object_permission_id if team_unchanged else None, + existing_key_object_permission=existing_key_row.object_permission if team_unchanged else None, ) await validate_key_search_tools_against_team( object_permission=object_permission_dict, diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 488ff285504a..13080a6cf839 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -448,7 +448,7 @@ async def enforce_all_proxy_mcp_servers_grant_is_admin_only( async def _get_grandfathered_key_mcp_server_ids( - object_permission_id: str | None, + existing_object_permission: Optional["LiteLLM_ObjectPermissionTable"], prisma_client: PrismaClient | None, ) -> frozenset[str]: """ @@ -457,18 +457,13 @@ async def _get_grandfathered_key_mcp_server_ids( team allowlist has since changed; sentinels are excluded so they cannot grandfather anything. """ - if not object_permission_id or prisma_client is None: - return frozenset() - existing: Final = await ObjectPermissionRepository(prisma_client).table.find_unique( - where={"object_permission_id": object_permission_id}, - ) - if existing is None: + if existing_object_permission is None or prisma_client is None: return frozenset() - raw_tool_perms: Final = existing.mcp_tool_permissions or {} + raw_tool_perms: Final = existing_object_permission.mcp_tool_permissions or {} tool_perm_keys: Final[frozenset[str]] = frozenset( json.loads(raw_tool_perms).keys() if isinstance(raw_tool_perms, str) else raw_tool_perms.keys() ) - identifiers: Final = (frozenset(existing.mcp_servers or []) | tool_perm_keys) - { + identifiers: Final = (frozenset(existing_object_permission.mcp_servers or []) | tool_perm_keys) - { SpecialMCPServerNames.no_mcp_servers.value, SpecialMCPServerName.all_proxy_servers.value, } @@ -562,12 +557,12 @@ async def validate_key_mcp_servers_against_team( team_obj: Optional["LiteLLM_TeamTableCachedObj"], prisma_client: PrismaClient | None = None, is_proxy_admin: bool = False, - existing_key_object_permission_id: str | None = None, + existing_key_object_permission: Optional["LiteLLM_ObjectPermissionTable"] = None, ) -> ObjectPermissionDict | None: """ Validate that MCP servers requested on a key are within the allowed scope. - When ``existing_key_object_permission_id`` is provided (key updates), servers + When ``existing_key_object_permission`` is provided (key updates), servers the key already holds are grandfathered: keeping or removing them stays valid even if the team allowlist has since shrunk, while adding new servers outside the allowlist is still rejected. @@ -631,7 +626,7 @@ async def validate_key_mcp_servers_against_team( allowed_servers = all_allowed_servers | active_requested_servers grandfathered_servers: Final = await _get_grandfathered_key_mcp_server_ids( - object_permission_id=existing_key_object_permission_id, + existing_object_permission=existing_key_object_permission, prisma_client=prisma_client, ) disallowed_servers: Final = active_requested_servers - allowed_servers - grandfathered_servers 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 34869efb4912..6045b64023de 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 @@ -734,7 +734,7 @@ async def test_update_key_personal_non_admin_denied_vector_stores(monkeypatch): async def test_update_key_grandfathers_existing_mcp_servers(monkeypatch): """/key/update on a team key that already holds MCP servers outside the team allowlist must accept re-sent or shrunk grants (LIT-6062). The wrapper - must pass the existing key's object_permission_id into the validator when + must pass the existing key's object_permission row into the validator when the team is unchanged.""" from unittest.mock import AsyncMock, MagicMock @@ -750,9 +750,6 @@ async def test_update_key_grandfathers_existing_mcp_servers(monkeypatch): existing_row.mcp_servers = ["server-a", "server-b"] existing_row.mcp_tool_permissions = {} mock_prisma = MagicMock() - mock_prisma.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=existing_row - ) mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) team_obj = MagicMock() @@ -762,6 +759,7 @@ async def test_update_key_grandfathers_existing_mcp_servers(monkeypatch): existing_key_row = MagicMock( team_id="team-1", object_permission_id="perm-1", + object_permission=existing_row, ) mock_server_a = MagicMock() @@ -6616,7 +6614,7 @@ async def test_get_and_validate_existing_key(): assert result == mock_key mock_prisma_client.db.litellm_verificationtoken.find_unique.assert_called_once_with( - where={"token": "hashed-test-key-123"} + where={"token": "hashed-test-key-123"}, include={"object_permission": True} ) # Test Case 2: Key not found raises ProxyException diff --git a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py index 01f3eff6b5d5..5ef83344c1a5 100644 --- a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py @@ -1216,17 +1216,14 @@ async def test_empty_object_permission_passes_for_personal_non_admin(): # ---- Tests for grandfathering existing key MCP servers on /key/update (LIT-6062) ---- -def _make_grandfather_prisma(mcp_servers=None, mcp_tool_permissions=None): - """Mock prisma client whose object permission table returns the key's existing row.""" +def _make_grandfather_fixtures(mcp_servers=None, mcp_tool_permissions=None): + """Mock prisma client plus the key's existing object permission row.""" existing_row = MagicMock() existing_row.mcp_servers = mcp_servers or [] existing_row.mcp_tool_permissions = mcp_tool_permissions or {} mock_prisma = MagicMock() - mock_prisma.db.litellm_objectpermissiontable.find_unique = AsyncMock( - return_value=existing_row - ) mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) - return mock_prisma + return mock_prisma, existing_row def _patch_grandfather_env(monkeypatch, mock_mgr): @@ -1246,19 +1243,19 @@ async def test_validate_key_update_grandfathers_existing_servers(monkeypatch): shrink those grants on /key/update without a 403 (LIT-6062).""" _patch_grandfather_env(monkeypatch, _make_mock_mcp_manager("server-a", "server-b")) team_obj = _make_team_obj(mcp_servers=[]) - mock_prisma = _make_grandfather_prisma(mcp_servers=["server-a", "server-b"]) + mock_prisma, existing_row = _make_grandfather_fixtures(mcp_servers=["server-a", "server-b"]) resend = await validate_key_mcp_servers_against_team( object_permission={"mcp_servers": ["server-a", "server-b"]}, team_obj=team_obj, prisma_client=mock_prisma, - existing_key_object_permission_id="perm-1", + existing_key_object_permission=existing_row, ) assert sorted(resend["mcp_servers"]) == ["server-a", "server-b"] shrink = await validate_key_mcp_servers_against_team( object_permission={"mcp_servers": ["server-a"]}, team_obj=team_obj, prisma_client=mock_prisma, - existing_key_object_permission_id="perm-1", + existing_key_object_permission=existing_row, ) assert shrink["mcp_servers"] == ["server-a"] @@ -1269,13 +1266,13 @@ async def test_validate_key_update_grandfather_does_not_allow_new_servers(monkey server outside the team allowlist still raises 403.""" _patch_grandfather_env(monkeypatch, _make_mock_mcp_manager("server-a", "server-new")) team_obj = _make_team_obj(mcp_servers=[]) - mock_prisma = _make_grandfather_prisma(mcp_servers=["server-a"]) + mock_prisma, existing_row = _make_grandfather_fixtures(mcp_servers=["server-a"]) with pytest.raises(HTTPException) as exc_info: await validate_key_mcp_servers_against_team( object_permission={"mcp_servers": ["server-a", "server-new"]}, team_obj=team_obj, prisma_client=mock_prisma, - existing_key_object_permission_id="perm-1", + existing_key_object_permission=existing_row, ) assert exc_info.value.status_code == 403 assert "server-new" in str(exc_info.value.detail) @@ -1283,17 +1280,17 @@ async def test_validate_key_update_grandfather_does_not_allow_new_servers(monkey @pytest.mark.asyncio async def test_validate_key_update_without_existing_permission_still_raises(monkeypatch): - """Without an existing permission row id (new grants or team change) the + """Without an existing permission row (new grants or team change) the subset check stays strict.""" _patch_grandfather_env(monkeypatch, _make_mock_mcp_manager("server-a")) team_obj = _make_team_obj(mcp_servers=[]) - mock_prisma = _make_grandfather_prisma(mcp_servers=["server-a"]) + mock_prisma, _ = _make_grandfather_fixtures(mcp_servers=["server-a"]) with pytest.raises(HTTPException) as exc_info: await validate_key_mcp_servers_against_team( object_permission={"mcp_servers": ["server-a"]}, team_obj=team_obj, prisma_client=mock_prisma, - existing_key_object_permission_id=None, + existing_key_object_permission=None, ) assert exc_info.value.status_code == 403 @@ -1304,14 +1301,14 @@ async def test_validate_key_update_grandfathers_tool_permission_keys(monkeypatch (stored as a JSON string) are grandfathered too.""" _patch_grandfather_env(monkeypatch, _make_mock_mcp_manager("server-a")) team_obj = _make_team_obj(mcp_servers=[]) - mock_prisma = _make_grandfather_prisma( + mock_prisma, existing_row = _make_grandfather_fixtures( mcp_tool_permissions=json.dumps({"server-a": ["tool1"]}) ) result = await validate_key_mcp_servers_against_team( object_permission={"mcp_servers": ["server-a"]}, team_obj=team_obj, prisma_client=mock_prisma, - existing_key_object_permission_id="perm-1", + existing_key_object_permission=existing_row, ) assert result["mcp_servers"] == ["server-a"] @@ -1321,7 +1318,7 @@ async def test_validate_key_update_sentinels_do_not_grandfather(monkeypatch): """Sentinels stored on the existing row must not grandfather anything.""" _patch_grandfather_env(monkeypatch, _make_mock_mcp_manager("server-a")) team_obj = _make_team_obj(mcp_servers=[]) - mock_prisma = _make_grandfather_prisma( + mock_prisma, existing_row = _make_grandfather_fixtures( mcp_servers=[SpecialMCPServerName.all_proxy_servers.value, "no-mcp-servers"] ) with pytest.raises(HTTPException) as exc_info: @@ -1329,7 +1326,7 @@ async def test_validate_key_update_sentinels_do_not_grandfather(monkeypatch): object_permission={"mcp_servers": ["server-a"]}, team_obj=team_obj, prisma_client=mock_prisma, - existing_key_object_permission_id="perm-1", + existing_key_object_permission=existing_row, ) assert exc_info.value.status_code == 403