fix(key_management): allow /key/update to keep or shrink MCP server grants the key already holds

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,

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

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

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