mcp: --mcp / configure mcp --location no longer clobbers other MCP servers

src/ucode/mcp.py

@@ -1582,6 +1582,24 @@ def _skills_entries(servers: list[dict]) -> list[dict]:
     return [s for s in servers if s.get("kind") == SKILLS_MCP_KIND]
 
 
+_MCP_SERVICE_URL_MARKER = "/ai-gateway/mcp-services/"
+
+
+def _is_mcp_service_in_location(server: dict, location: str) -> bool:
+    """True if ``server`` is an mcp-service registered at ``location``.
+
+    mcp-service entries carry a URL ``.../ai-gateway/mcp-services/<cat>.<schema>.<name>``.
+    A server belongs to ``location`` when that full name is ``<location>.<leaf>``
+    (exactly one segment past the schema) — so ``system.ai`` matches
+    ``system.ai.slack`` but not ``system.ai.sub.x`` or ``system.aiX.y``."""
+    url = server.get("url")
+    if not isinstance(url, str) or _MCP_SERVICE_URL_MARKER not in url:
+        return False
+    full_name = url.split(_MCP_SERVICE_URL_MARKER, 1)[1].strip("/")
+    prefix = f"{location}."
+    return full_name.startswith(prefix) and "." not in full_name[len(prefix) :]
+
+
 def _resolve_location_mcp_servers(
     workspace: str,
     profile: str | None,
@@ -1592,19 +1610,23 @@ def _resolve_location_mcp_servers(
 ) -> list[dict]:
     """Build the desired MCP server list for ``--location <cat>.<schema>``.
 
-    Strict replacement for mcp-services: the returned list is exactly the ones
-    discovered at ``location`` (any previously-registered mcp-service outside it
-    is removed by ``apply_mcp_server_changes``), plus any existing skills
-    connection, preserved untouched. Raises ``RuntimeError`` for an invalid
-    location (HTTP 404 from the listing API) or any other listing failure.
+    Replacement is scoped to the mcp-services *in ``location``*: the returned
+    list is the ones discovered there, plus every other original server carried
+    through untouched — external connections, Genie/apps/Vector-Search/UC
+    functions, mcp-services in *other* schemas, and skills. So configuring one
+    location never disturbs servers registered elsewhere. (Previously this
+    returned only ``location``'s services + skills, so a one-shot
+    ``ucode configure --mcp system.ai.slack`` wiped every other MCP server —
+    #the custom-mcp-clobber bug.) Raises ``RuntimeError`` for an invalid location
+    (HTTP 404 from the listing API) or any other listing failure.
 
     When ``services`` is given, the discovered set is narrowed to exactly that
     subset (matched by full name like ``system.ai.github`` or bare short name
     like ``github``); names not found at ``location`` are skipped with a
     warning rather than failing, so a saved selection that references a
     since-removed service still configures the rest. An empty set selects
-    nothing (every previously-registered service in the location is removed).
-    ``None`` keeps the whole schema."""
+    nothing (every previously-registered service *in the location* is removed;
+    servers outside it are still preserved). ``None`` keeps the whole schema."""
     if location.count(".") != 1 or not all(part.strip() for part in location.split(".")):
         raise RuntimeError(f"--location must be `<catalog>.<schema>`, got `{location}`.")
 
@@ -1654,7 +1676,23 @@ def _resolve_location_mcp_servers(
             working_servers.append(original.copy())
         else:
             working_servers.append(candidate)
-    return [*working_servers, *_skills_entries(original_servers)]
+
+    # Carry through every original server that isn't an mcp-service in *this*
+    # location: other schemas' services, external/genie/app/vector-search/UC
+    # servers, etc. Only the location's own services are (re)placed above; skills
+    # are re-appended by name below. Without this, configuring one location would
+    # remove all unrelated MCP servers (`apply_mcp_server_changes` deletes any
+    # original not present in the returned set).
+    replaced_names = {s["name"] for s in working_servers}
+    skills_names = {s.get("name") for s in _skills_entries(original_servers)}
+    preserved = [
+        server
+        for server in original_servers
+        if server.get("name") not in replaced_names
+        and server.get("name") not in skills_names
+        and not _is_mcp_service_in_location(server, location)
+    ]
+    return [*working_servers, *preserved, *_skills_entries(original_servers)]
 
 
 # The first wizard step lets the user choose which sources to search. Each is a

tests/test_mcp.py

@@ -1595,6 +1595,39 @@ def _stub_location_base(monkeypatch, state):
     monkeypatch.setattr(mcp, "get_databricks_token", lambda workspace, profile=None: "token")
 
 
+class TestIsMcpServiceInLocation:
+    def _svc(self, full_name: str) -> dict:
+        return {
+            "name": full_name.replace(".", "-"),
+            "url": f"{WS}/ai-gateway/mcp-services/{full_name}",
+        }
+
+    def test_matches_a_service_directly_in_the_location(self):
+        assert mcp._is_mcp_service_in_location(self._svc("system.ai.slack"), "system.ai")
+
+    def test_does_not_match_a_nested_name(self):
+        # Only one segment past the schema counts; a deeper name isn't "in" it.
+        assert not mcp._is_mcp_service_in_location(self._svc("system.ai.sub.x"), "system.ai")
+
+    def test_does_not_match_a_schema_prefix_lookalike(self):
+        # `system.aiX` must not match location `system.ai`.
+        assert not mcp._is_mcp_service_in_location(self._svc("system.aiX.y"), "system.ai")
+
+    def test_does_not_match_another_schema(self):
+        assert not mcp._is_mcp_service_in_location(self._svc("main.tools.helper"), "system.ai")
+
+    def test_non_service_entries_are_never_in_a_location(self):
+        for url in (
+            f"{WS}/api/2.0/mcp/external/jira",
+            f"{WS}/api/2.0/mcp/genie/123",
+            f"{WS}/api/2.0/mcp/sql",
+        ):
+            assert not mcp._is_mcp_service_in_location({"name": "x", "url": url}, "system.ai")
+
+    def test_entry_without_url_is_not_in_a_location(self):
+        assert not mcp._is_mcp_service_in_location({"name": "x"}, "system.ai")
+
+
 class TestConfigureMcpFromLocation:
     def test_rejects_malformed_location(self, monkeypatch):
         _stub_location_base(monkeypatch, {**CLAUDE_STATE})
@@ -1682,19 +1715,36 @@ def fake_list(workspace, token, parent):
             },
         ]
 
-    def test_replaces_servers_outside_location(self, monkeypatch):
+    def test_preserves_servers_outside_the_location(self, monkeypatch):
+        # Regression: configuring one location must NOT remove unrelated MCP
+        # servers. Previously `ucode configure --mcp system.ai.slack` wiped the
+        # user's other servers (custom-mcp-clobber bug); they must be carried
+        # through untouched, with only the location's own services (re)placed.
         saved_states: list[dict] = []
-        configured: list[tuple[str, str, str, dict]] = []
+        configured: list[tuple[str, str, str]] = []
         removed: list[tuple[str, str]] = []
-        outside_entry = {
+        # A grab-bag of servers that live outside `system.ai`.
+        sql_entry = {
             "name": "databricks-sql",
             "url": f"{WS}/api/2.0/mcp/sql",
             "auth": "proxy",
             "clients": ["claude"],
         }
+        external_entry = {
+            "name": "jira-mcp",
+            "url": f"{WS}/api/2.0/mcp/external/jira-mcp",
+            "auth": "proxy",
+            "clients": ["claude"],
+        }
+        other_schema_entry = {
+            "name": "main-tools-helper",
+            "url": f"{WS}/ai-gateway/mcp-services/main.tools.helper",
+            "auth": "proxy",
+            "clients": ["claude"],
+        }
         _stub_location_base(
             monkeypatch,
-            {**CLAUDE_STATE, "mcp_servers": [outside_entry]},
+            {**CLAUDE_STATE, "mcp_servers": [sql_entry, external_entry, other_schema_entry]},
         )
         monkeypatch.setattr(
             mcp,
@@ -1715,16 +1765,60 @@ def test_replaces_servers_outside_location(self, monkeypatch):
 
         assert mcp.configure_mcp_command(location="system.ai") == 0
 
-        assert removed == [("claude", "databricks-sql")]
+        # Nothing outside system.ai was removed.
+        assert removed == []
+        # Only the location's own service was (re)configured.
         assert [c[1] for c in configured] == ["system-ai-github"]
-        assert saved_states[-1]["mcp_servers"] == [
-            {
-                "name": "system-ai-github",
-                "url": f"{WS}/ai-gateway/mcp-services/system.ai.github",
-                "auth": "proxy",
-                "clients": ["claude"],
-            },
-        ]
+        # Saved state keeps the three outside servers plus the new one.
+        saved_names = {s["name"] for s in saved_states[-1]["mcp_servers"]}
+        assert saved_names == {
+            "databricks-sql",
+            "jira-mcp",
+            "main-tools-helper",
+            "system-ai-github",
+        }
+
+    def test_replaces_stale_service_within_the_location(self, monkeypatch):
+        # The location's OWN services are still strictly replaced: a service
+        # previously registered under system.ai but no longer discovered there
+        # is removed, while out-of-location servers stay.
+        saved_states: list[dict] = []
+        removed: list[tuple[str, str]] = []
+        stale_in_location = {
+            "name": "system-ai-oldservice",
+            "url": f"{WS}/ai-gateway/mcp-services/system.ai.oldservice",
+            "auth": "proxy",
+            "clients": ["claude"],
+        }
+        outside = {
+            "name": "jira-mcp",
+            "url": f"{WS}/api/2.0/mcp/external/jira-mcp",
+            "auth": "proxy",
+            "clients": ["claude"],
+        }
+        _stub_location_base(
+            monkeypatch,
+            {**CLAUDE_STATE, "mcp_servers": [stale_in_location, outside]},
+        )
+        monkeypatch.setattr(
+            mcp,
+            "list_mcp_services",
+            lambda workspace, token, parent: (["system.ai.github"], None),
+        )
+        monkeypatch.setattr(mcp, "configure_client_mcp_server", lambda *a, **kw: [])
+        monkeypatch.setattr(
+            mcp,
+            "remove_client_mcp_server",
+            lambda client, name: removed.append((client, name)) or [],
+        )
+        monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy()))
+
+        assert mcp.configure_mcp_command(location="system.ai") == 0
+
+        # Stale in-location service removed; out-of-location server untouched.
+        assert removed == [("claude", "system-ai-oldservice")]
+        saved_names = {s["name"] for s in saved_states[-1]["mcp_servers"]}
+        assert saved_names == {"jira-mcp", "system-ai-github"}
 
     def test_preserves_skills_connection(self, monkeypatch):
         """A skills connection is owned by `configure skills`, so `configure mcp