3446 - Enable public MCP objects on team Servers

Signed-off-by: Gabriel Costa <gabrielcg@proton.me>

Author: gcgoncalves

mcpgateway/admin.py

@@ -1383,6 +1383,42 @@ def _owner_access_condition(owner_column, team_column, *, user_email: str, team_
     return owner_column == user_email
 
 
+def _merge_select_all_ids(form: Any, flag_key: str, all_ids_key: str, checked_list: list[str]) -> list[str]:
+    """Merge server-fetched IDs with UI-checked IDs when "Select All" is active.
+
+    When the user clicks "Select All" in a paginated list, the browser populates
+    *all_ids_key* with IDs fetched from the corresponding /ids endpoint. Because
+    that endpoint may be team-scoped, it can miss platform-public items that are
+    still visible (and checked) in the UI. Taking the union of both sources
+    ensures every explicitly selected item is preserved.
+
+    Note: both sources are client-supplied form values. Downstream persistence
+    code is responsible for enforcing final access control on the merged IDs.
+
+    Args:
+        form: Starlette form object.
+        flag_key (str): Form field that signals "Select All" mode (e.g. ``"selectAllTools"``).
+        all_ids_key (str): Form field holding the JSON-encoded server-fetched IDs.
+        checked_list (list[str]): IDs collected from checked checkboxes in the form.
+
+    Returns:
+        list[str]: Merged, deduplicated list of string IDs; or *checked_list* unchanged
+        when Select All is not active or the JSON payload cannot be parsed.
+    """
+    if form.get(flag_key) != "true":
+        return checked_list
+    raw = form.get(all_ids_key) or "[]"
+    try:
+        server_ids = orjson.loads(raw)
+        # Normalise to str to avoid silent int/str duplicates from different sources.
+        merged = list({str(i) for i in server_ids} | set(checked_list))
+        LOGGER.info("Select All (%s): %d items after merge", all_ids_key, len(merged))
+        return merged
+    except orjson.JSONDecodeError:
+        LOGGER.warning("Failed to parse %s JSON, falling back to checked items", all_ids_key)
+        return checked_list
+
+
 async def _has_permission(
     *,
     db: Session,
@@ -2531,39 +2567,13 @@ async def admin_add_server(request: Request, db: Session = Depends(get_db), user
     try:
         LOGGER.debug(f"User {get_user_email(user)} is adding a new server with name: {form['name']}")
 
-        # Handle "Select All" for tools
-        associated_tools_list = form.getlist("associatedTools")
-        if form.get("selectAllTools") == "true":
-            # User clicked "Select All" - get all tool IDs from hidden field
-            all_tool_ids_json = str(form.get("allToolIds", "[]"))
-            try:
-                all_tool_ids = orjson.loads(all_tool_ids_json)
-                associated_tools_list = all_tool_ids
-                LOGGER.info(f"Select All tools enabled: {len(all_tool_ids)} tools selected")
-            except orjson.JSONDecodeError:
-                LOGGER.warning("Failed to parse allToolIds JSON, falling back to checked tools")
-
-        # Handle "Select All" for resources
-        associated_resources_list = form.getlist("associatedResources")
-        if form.get("selectAllResources") == "true":
-            all_resource_ids_json = str(form.get("allResourceIds", "[]"))
-            try:
-                all_resource_ids = orjson.loads(all_resource_ids_json)
-                associated_resources_list = all_resource_ids
-                LOGGER.info(f"Select All resources enabled: {len(all_resource_ids)} resources selected")
-            except orjson.JSONDecodeError:
-                LOGGER.warning("Failed to parse allResourceIds JSON, falling back to checked resources")
-
-        # Handle "Select All" for prompts
-        associated_prompts_list = form.getlist("associatedPrompts")
-        if form.get("selectAllPrompts") == "true":
-            all_prompt_ids_json = str(form.get("allPromptIds", "[]"))
-            try:
-                all_prompt_ids = orjson.loads(all_prompt_ids_json)
-                associated_prompts_list = all_prompt_ids
-                LOGGER.info(f"Select All prompts enabled: {len(all_prompt_ids)} prompts selected")
-            except orjson.JSONDecodeError:
-                LOGGER.warning("Failed to parse allPromptIds JSON, falling back to checked prompts")
+        # Handle "Select All" for tools, resources, and prompts.
+        # _merge_select_all_ids takes the union of the server-fetched paginated IDs
+        # (allToolIds etc.) with the explicitly checked form values so that
+        # platform-public items visible in the UI are never silently dropped.
+        associated_tools_list = _merge_select_all_ids(form, "selectAllTools", "allToolIds", form.getlist("associatedTools"))
+        associated_resources_list = _merge_select_all_ids(form, "selectAllResources", "allResourceIds", form.getlist("associatedResources"))
+        associated_prompts_list = _merge_select_all_ids(form, "selectAllPrompts", "allPromptIds", form.getlist("associatedPrompts"))
 
         # Handle OAuth 2.0 configuration (RFC 9728)
         oauth_enabled = form.get("oauth_enabled") == "on"
@@ -2709,39 +2719,13 @@ async def admin_edit_server(
 
         mod_metadata = MetadataCapture.extract_modification_metadata(request, user, 0)
 
-        # Handle "Select All" for tools
-        associated_tools_list = form.getlist("associatedTools")
-        if form.get("selectAllTools") == "true":
-            # User clicked "Select All" - get all tool IDs from hidden field
-            all_tool_ids_json = str(form.get("allToolIds", "[]"))
-            try:
-                all_tool_ids = orjson.loads(all_tool_ids_json)
-                associated_tools_list = all_tool_ids
-                LOGGER.info(f"Select All tools enabled for edit: {len(all_tool_ids)} tools selected")
-            except orjson.JSONDecodeError:
-                LOGGER.warning("Failed to parse allToolIds JSON, falling back to checked tools")
-
-        # Handle "Select All" for resources
-        associated_resources_list = form.getlist("associatedResources")
-        if form.get("selectAllResources") == "true":
-            all_resource_ids_json = str(form.get("allResourceIds", "[]"))
-            try:
-                all_resource_ids = orjson.loads(all_resource_ids_json)
-                associated_resources_list = all_resource_ids
-                LOGGER.info(f"Select All resources enabled for edit: {len(all_resource_ids)} resources selected")
-            except orjson.JSONDecodeError:
-                LOGGER.warning("Failed to parse allResourceIds JSON, falling back to checked resources")
-
-        # Handle "Select All" for prompts
-        associated_prompts_list = form.getlist("associatedPrompts")
-        if form.get("selectAllPrompts") == "true":
-            all_prompt_ids_json = str(form.get("allPromptIds", "[]"))
-            try:
-                all_prompt_ids = orjson.loads(all_prompt_ids_json)
-                associated_prompts_list = all_prompt_ids
-                LOGGER.info(f"Select All prompts enabled for edit: {len(all_prompt_ids)} prompts selected")
-            except orjson.JSONDecodeError:
-                LOGGER.warning("Failed to parse allPromptIds JSON, falling back to checked prompts")
+        # Handle "Select All" for tools, resources, and prompts.
+        # _merge_select_all_ids takes the union of the server-fetched paginated IDs
+        # (allToolIds etc.) with the explicitly checked form values so that
+        # platform-public items visible in the UI are never silently dropped.
+        associated_tools_list = _merge_select_all_ids(form, "selectAllTools", "allToolIds", form.getlist("associatedTools"))
+        associated_resources_list = _merge_select_all_ids(form, "selectAllResources", "allResourceIds", form.getlist("associatedResources"))
+        associated_prompts_list = _merge_select_all_ids(form, "selectAllPrompts", "allPromptIds", form.getlist("associatedPrompts"))
 
         # Handle OAuth 2.0 configuration (RFC 9728)
         oauth_enabled = form.get("oauth_enabled") == "on"
@@ -8288,14 +8272,19 @@ async def admin_get_all_tool_ids(
                 LOGGER.debug(f"Filtering tools by gateway IDs: {non_null_ids}")
 
     # Build access conditions
-    # When team_id is specified, show ONLY items from that team (team-scoped view)
-    # Otherwise, show all accessible items (All Teams view)
+    # When team_id is specified, show items from that team plus all platform-public tools
+    # (visibility="public") so the "Select All" count and payload match what is actually
+    # visible in the edit UI. Public visibility is platform-wide regardless of team ownership.
+    # Otherwise, show all accessible items (All Teams view).
     if team_id:
         if team_id in team_ids:
             # Apply visibility check: team/public resources + user's own resources (including private)
+            # Also include all platform-public tools so they can be associated with team-owned
+            # virtual servers.
             team_access = [
                 and_(DbTool.team_id == team_id, DbTool.visibility.in_(["team", "public"])),
                 and_(DbTool.team_id == team_id, DbTool.owner_email == user_email),
+                DbTool.visibility == "public",
             ]
             query = query.where(or_(*team_access))
             LOGGER.debug(f"Filtering tool IDs by team_id: {team_id}")
@@ -9526,15 +9515,20 @@ async def admin_get_all_prompt_ids(
         query = query.where(DbPrompt.enabled.is_(True))
 
     # Build access conditions
-    # When team_id is specified, show ONLY items from that team (team-scoped view)
-    # Otherwise, show all accessible items (All Teams view)
+    # When team_id is specified, show items from that team plus all platform-public prompts
+    # (visibility="public") so the "Select All" count and payload match what is actually
+    # visible in the edit UI. Public visibility is platform-wide regardless of team ownership.
+    # Otherwise, show all accessible items (All Teams view).
     if team_id:
-        # Team-specific view: only show prompts from the specified team
+        # Team-specific view: show prompts from the specified team plus platform-public prompts
         if team_id in team_ids:
             # Apply visibility check: team/public resources + user's own resources (including private)
+            # Also include all platform-public prompts so they can be associated with team-owned
+            # virtual servers.
             team_access = [
                 and_(DbPrompt.team_id == team_id, DbPrompt.visibility.in_(["team", "public"])),
                 and_(DbPrompt.team_id == team_id, DbPrompt.owner_email == user_email),
+                DbPrompt.visibility == "public",
             ]
             query = query.where(or_(*team_access))
             LOGGER.debug(f"Filtering prompt IDs by team_id: {team_id}")
@@ -9606,15 +9600,20 @@ async def admin_get_all_resource_ids(
         query = query.where(DbResource.enabled.is_(True))
 
     # Build access conditions
-    # When team_id is specified, show ONLY items from that team (team-scoped view)
-    # Otherwise, show all accessible items (All Teams view)
+    # When team_id is specified, show items from that team plus all platform-public resources
+    # (visibility="public") so the "Select All" count and payload match what is actually
+    # visible in the edit UI. Public visibility is platform-wide regardless of team ownership.
+    # Otherwise, show all accessible items (All Teams view).
     if team_id:
-        # Team-specific view: only show resources from the specified team
+        # Team-specific view: show resources from the specified team plus platform-public resources
         if team_id in team_ids:
             # Apply visibility check: team/public resources + user's own resources (including private)
+            # Also include all platform-public resources so they can be associated with team-owned
+            # virtual servers.
             team_access = [
                 and_(DbResource.team_id == team_id, DbResource.visibility.in_(["team", "public"])),
                 and_(DbResource.team_id == team_id, DbResource.owner_email == user_email),
+                DbResource.visibility == "public",
             ]
             query = query.where(or_(*team_access))
             LOGGER.debug(f"Filtering resource IDs by team_id: {team_id}")

tests/unit/mcpgateway/test_admin.py

@@ -606,9 +606,10 @@ async def test_admin_add_server_select_all_parses_json(self, mock_register_serve
         assert result.status_code == 200
 
         server_create = mock_register_server.call_args.args[1]
-        assert server_create.associated_tools == ["tool-1", "tool-2"]
-        assert server_create.associated_resources == ["res-1"]
-        assert server_create.associated_prompts == ["prompt-1", "prompt-2"]
+        # Merge of allToolIds + associatedTools (public items from UI are preserved)
+        assert set(server_create.associated_tools) == {"tool-1", "tool-2", "tool-x"}
+        assert set(server_create.associated_resources) == {"res-1", "res-x"}
+        assert set(server_create.associated_prompts) == {"prompt-1", "prompt-2", "prompt-x"}
         assert server_create.oauth_enabled is True
         assert server_create.oauth_config["authorization_servers"] == ["https://idp.example.com"]
         assert server_create.oauth_config["scopes_supported"] == ["openid", "profile"]
@@ -880,9 +881,10 @@ async def test_admin_edit_server_select_all_parses_json(self, mock_update_server
         assert result.status_code == 200
 
         server_update = mock_update_server.call_args[0][2]
-        assert server_update.associated_tools == ["tool-1", "tool-2"]
-        assert server_update.associated_resources == ["res-1"]
-        assert server_update.associated_prompts == ["prompt-1", "prompt-2"]
+        # Merge of allToolIds + associatedTools (public items from UI are preserved)
+        assert set(server_update.associated_tools) == {"tool-1", "tool-2", "tool-x"}
+        assert set(server_update.associated_resources) == {"res-1", "res-x"}
+        assert set(server_update.associated_prompts) == {"prompt-1", "prompt-2", "prompt-x"}
 
     @patch.object(ServerService, "update_server")
     async def test_admin_edit_server_select_all_json_decode_error(self, mock_update_server, mock_request, mock_db, monkeypatch):
@@ -10323,6 +10325,38 @@ async def test_admin_get_all_tool_ids_gateway_and_team_filters(monkeypatch, mock
     assert result["count"] == 0
 
 
+@pytest.mark.asyncio
+async def test_admin_get_all_tool_ids_team_scoped_includes_public(monkeypatch, mock_db):
+    """When team_id is set, the SQL query must include a standalone visibility='public'
+    condition (not gated by team_id) so that platform-public tools from public MCP
+    servers appear in team-scoped Select All fetches and can be associated with
+    team-owned virtual servers.  Regression test for issue #3446."""
+    import re
+    from sqlalchemy.dialects import sqlite as sqlite_dialect
+
+    setup_team_service(monkeypatch, ["team-1"])
+    mock_db.execute.return_value.all.return_value = []
+
+    await admin_get_all_tool_ids(
+        include_inactive=False,
+        gateway_id=None,
+        team_id="team-1",
+        db=mock_db,
+        user={"email": "user@example.com", "db": mock_db},
+    )
+
+    executed_query = mock_db.execute.call_args[0][0]
+    sql = str(executed_query.compile(dialect=sqlite_dialect.dialect(), compile_kwargs={"literal_binds": True}))
+
+    # A standalone `visibility = 'public'` condition must be present as a top-level
+    # OR alternative — not wrapped inside `team_id = '...' AND visibility IN (...)`.
+    # This is what makes platform-public tools visible to team-scoped queries.
+    assert re.search(r"tools\.visibility\s*=\s*'public'", sql), (
+        "Expected a standalone visibility='public' condition in team-scoped tool IDs query. "
+        "Platform-public tools must be accessible when associating with team-owned virtual servers."
+    )
+
+
 @pytest.mark.asyncio
 async def test_admin_get_all_prompt_ids_gateway_and_team_filters(monkeypatch, mock_db):
     """Cover non-null gateway filters and team membership checks for prompt IDs helper."""