diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d1c083529192..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( @@ -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=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 fb64914f6f51..13080a6cf839 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -447,6 +447,36 @@ async def enforce_all_proxy_mcp_servers_grant_is_admin_only( ) +async def _get_grandfathered_key_mcp_server_ids( + existing_object_permission: Optional["LiteLLM_ObjectPermissionTable"], + 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 existing_object_permission is None or prisma_client is None: + return frozenset() + 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_object_permission.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 +557,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: Optional["LiteLLM_ObjectPermissionTable"] = None, ) -> ObjectPermissionDict | None: """ Validate that MCP servers requested on a key are within the allowed scope. + 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. + 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 +625,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( + existing_object_permission=existing_key_object_permission, + 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..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 @@ -730,6 +730,68 @@ 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 row 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_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", + object_permission=existing_row, + ) + + 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, @@ -6552,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 b129ad0f659c..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 @@ -1213,6 +1213,124 @@ 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_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_mcpservertable.find_many = AsyncMock(return_value=[]) + return mock_prisma, existing_row + + +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, 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=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=existing_row, + ) + 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, 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=existing_row, + ) + 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 (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_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=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, 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=existing_row, + ) + 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, existing_row = _make_grandfather_fixtures( + 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=existing_row, + ) + 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