remove stale uc securable mcp entries

src/ucode/cli.py

@@ -48,6 +48,7 @@
     MCP_CLIENTS,
     configure_mcp_command,
     purge_cross_workspace_mcp_residue,
+    purge_uc_mcp_residue,
     revert_mcp_configs,
 )
 from ucode.state import STATE_PATH, clear_state, load_full_state, load_state, save_state
@@ -173,17 +174,13 @@ def configure_shared_state(
     # launch: target workspace's persisted state). Use *target* state on the
     # launch path so the flag is sticky per-workspace and doesn't leak
     # across workspace switches.
-    # TODO: when this flips uc_enabled True->False, prune any
-    # `system.ai.*` MCP services from state["mcp_servers"] (and their
-    # cross-tool registrations). Today they linger as orphans pointing at
-    # /ai-gateway/mcp-services/* until the user re-runs `configure mcp`
-    # or switches workspaces.
+    target_ws_state = load_full_state().get("workspaces", {}).get(workspace) or {}
+    previous_uc_enabled = bool(target_ws_state.get("uc_enabled"))
     if enable_uc is None:
         if reset_uc:
             enable_uc = uc_enabled(default=False)
         else:
-            target_ws_state = load_full_state().get("workspaces", {}).get(workspace) or {}
-            enable_uc = uc_enabled(default=bool(target_ws_state.get("uc_enabled")))
+            enable_uc = uc_enabled(default=previous_uc_enabled)
     fetch_all = tools is None
     if force_login:
         run_databricks_login(workspace, profile)
@@ -260,6 +257,10 @@ def configure_shared_state(
     # workspace's agent configs aren't stale.
     if previous_workspace and previous_workspace != workspace:
         purge_cross_workspace_mcp_residue(state, workspace)
+    # When the user flips UC discovery off (typically by re-running plain
+    # `ucode configure`), drop any `system.ai.*` MCP service entries
+    if previous_uc_enabled and not enable_uc:
+        purge_uc_mcp_residue(state)
     # Diagnostic reasons are transient — attach after save_state so they don't
     # land on disk but are available to the caller for this run.
     state["_discovery_reasons"] = {

src/ucode/mcp.py

@@ -511,6 +511,15 @@ def _partition_mcp_entries_by_workspace(
     return current, foreign
 
 
+_UC_MCP_SERVICE_URL_FRAGMENT = "/ai-gateway/mcp-services/"
+
+
+def _is_uc_mcp_service_entry(entry: dict) -> bool:
+    """An MCP entry pointing at a curated `system.ai.*` UC MCP service."""
+    url = entry.get("url")
+    return isinstance(url, str) and _UC_MCP_SERVICE_URL_FRAGMENT in url
+
+
 def _mcp_entries_only_in_other_workspaces(current_workspace: str) -> dict[str, set[str]]:
     """Return ``{name: {client, ...}}`` for MCPs ucode tracks only in workspaces other than ``current_workspace``."""
     full_state = load_full_state()
@@ -939,6 +948,43 @@ def purge_cross_workspace_mcp_residue(state: dict, workspace: str) -> None:
         )
 
 
+def purge_uc_mcp_residue(state: dict) -> None:
+    """Drop `system.ai.*` UC MCP service entries when UC discovery is disabled."""
+    raw_mcp_servers = list(state.get("mcp_servers") or [])
+    keep: list[dict] = []
+    drop: list[dict] = []
+    for entry in raw_mcp_servers:
+        if _is_uc_mcp_service_entry(entry):
+            drop.append(entry)
+        else:
+            keep.append(entry)
+    if not drop:
+        return
+
+    installed = set(available_mcp_clients())
+    dropped_names = ", ".join((_server_name(server) or "(unnamed)") for server in drop)
+    noun = "entry" if len(drop) == 1 else "entries"
+    print_warning(
+        f"Dropping {len(drop)} `system.ai.*` MCP {noun} after UC discovery was "
+        f"disabled: {dropped_names}."
+    )
+    for server in drop:
+        name = _server_name(server)
+        if not name:
+            continue
+        for client in server.get("clients") or []:
+            if client not in installed or client not in MCP_CLIENTS:
+                continue
+            try:
+                remove_client_mcp_server(client, name)
+            except RuntimeError as exc:
+                print_warning(
+                    f"Failed to remove `{name}` from {MCP_CLIENTS[client]['display']}: {exc}"
+                )
+    state["mcp_servers"] = keep
+    save_state(state)
+
+
 def configure_mcp_command() -> int:
     state = load_state()
     workspace = state.get("workspace")

tests/test_e2e_uc.py

@@ -161,6 +161,47 @@ def test_plain_configure_resets_persisted_uc_enabled(
             f"Reset run still pulled UC ids: {[m for m in ids if m.startswith('system.ai.')][:5]}"
         )
 
+    def test_plain_configure_purges_system_ai_mcp_entries(
+        self, monkeypatch, e2e_workspace, e2e_token
+    ):
+        """The True->False reset must also drop any `system.ai.*` MCP
+        service entries from state["mcp_servers"] (and from each tool's
+        registration) so they don't linger as orphans after UC is off.
+        Connection-based entries are left intact."""
+        from ucode.state import save_state
+
+        if not _has_uc_models(e2e_workspace, e2e_token):
+            pytest.skip("Workspace has no system.ai.* model services.")
+        monkeypatch.delenv("UCODE_ENABLE_UC", raising=False)
+
+        # Pretend the user just ran `--enable-uc` and then `configure mcp`,
+        # registering a UC MCP service entry alongside a regular connection
+        # entry. We seed the state directly to avoid driving the picker.
+        configure_shared_state(e2e_workspace, force_login=False, enable_uc=True)
+        seeded = load_state()
+        seeded["mcp_servers"] = [
+            {
+                "name": "system-ai-github",
+                "url": f"{e2e_workspace}/ai-gateway/mcp-services/system.ai.github",
+                "clients": [],  # no installed clients => no per-tool removals
+            },
+            {
+                "name": "regular-connection-mcp",
+                "url": f"{e2e_workspace}/api/2.0/mcp/external/regular_connection_mcp",
+                "clients": [],
+            },
+        ]
+        save_state(seeded)
+
+        # Plain `ucode configure` -> reset_uc=True, no flag, no env -> False.
+        configure_shared_state(e2e_workspace, force_login=False, enable_uc=None, reset_uc=True)
+
+        after = load_state()
+        names = sorted((s.get("name") or "") for s in (after.get("mcp_servers") or []))
+        assert names == ["regular-connection-mcp"], (
+            f"Expected only the connection-based entry to remain, got: {names}"
+        )
+
     def test_launch_path_preserves_persisted_uc_enabled(
         self, monkeypatch, e2e_workspace, e2e_token
     ):

tests/test_mcp.py

@@ -1107,3 +1107,149 @@ def test_removes_cli_registered_servers_and_restores_copilot_config(self, monkey
             "opencode": True,
             "copilot": True,
         }
+
+
+class TestPurgeUcMcpResidue:
+    """`purge_uc_mcp_residue` runs when uc_enabled flips True->False on a
+    workspace; it drops `system.ai.*` MCP entries from state["mcp_servers"]
+    and de-registers them from each agent CLI, while leaving
+    connection-based (`/api/2.0/mcp/external/`) entries alone."""
+
+    def _patch_clients_installed(self, monkeypatch, clients):
+        monkeypatch.setattr(mcp, "available_mcp_clients", lambda: clients)
+
+    def test_drops_only_system_ai_entries_keeps_connections(self, monkeypatch):
+        removed: list[tuple[str, str]] = []
+        saved: list[dict] = []
+        self._patch_clients_installed(monkeypatch, ["claude", "codex"])
+        monkeypatch.setattr(
+            mcp,
+            "remove_client_mcp_server",
+            lambda client, name: removed.append((client, name)) or ["user"],
+        )
+        monkeypatch.setattr(mcp, "save_state", lambda s: saved.append(dict(s)))
+
+        state = {
+            "workspace": WS,
+            "mcp_servers": [
+                {
+                    "name": "system-ai-github",
+                    "url": f"{WS}/ai-gateway/mcp-services/system.ai.github",
+                    "clients": ["claude", "codex"],
+                },
+                {
+                    "name": "asda-github-mcp",
+                    "url": f"{WS}/api/2.0/mcp/external/asda_github_mcp",
+                    "clients": ["claude"],
+                },
+                {
+                    "name": "system-ai-slack",
+                    "url": f"{WS}/ai-gateway/mcp-services/system.ai.slack",
+                    "clients": ["claude"],
+                },
+            ],
+        }
+
+        mcp.purge_uc_mcp_residue(state)
+
+        assert removed == [
+            ("claude", "system-ai-github"),
+            ("codex", "system-ai-github"),
+            ("claude", "system-ai-slack"),
+        ]
+        assert state["mcp_servers"] == [
+            {
+                "name": "asda-github-mcp",
+                "url": f"{WS}/api/2.0/mcp/external/asda_github_mcp",
+                "clients": ["claude"],
+            }
+        ]
+        assert saved and saved[-1]["mcp_servers"] == state["mcp_servers"]
+
+    def test_no_uc_entries_is_a_noop(self, monkeypatch):
+        saved: list[dict] = []
+        monkeypatch.setattr(
+            mcp,
+            "remove_client_mcp_server",
+            lambda client, name: (_ for _ in ()).throw(AssertionError("should not be called")),
+        )
+        monkeypatch.setattr(mcp, "save_state", lambda s: saved.append(dict(s)))
+        self._patch_clients_installed(monkeypatch, ["claude"])
+
+        state = {
+            "workspace": WS,
+            "mcp_servers": [
+                {
+                    "name": "asda-github-mcp",
+                    "url": f"{WS}/api/2.0/mcp/external/asda_github_mcp",
+                    "clients": ["claude"],
+                }
+            ],
+        }
+        mcp.purge_uc_mcp_residue(state)
+        assert state["mcp_servers"] == [
+            {
+                "name": "asda-github-mcp",
+                "url": f"{WS}/api/2.0/mcp/external/asda_github_mcp",
+                "clients": ["claude"],
+            }
+        ]
+        assert saved == []
+
+    def test_skips_clients_not_installed(self, monkeypatch):
+        # If gemini isn't on the box (e.g. uninstalled after the entry was
+        # registered), we still drop the entry from state but skip the
+        # de-registration call so we don't crash.
+        removed: list[tuple[str, str]] = []
+        self._patch_clients_installed(monkeypatch, ["claude"])
+        monkeypatch.setattr(
+            mcp,
+            "remove_client_mcp_server",
+            lambda client, name: removed.append((client, name)) or ["user"],
+        )
+        monkeypatch.setattr(mcp, "save_state", lambda s: None)
+
+        state = {
+            "workspace": WS,
+            "mcp_servers": [
+                {
+                    "name": "system-ai-github",
+                    "url": f"{WS}/ai-gateway/mcp-services/system.ai.github",
+                    "clients": ["claude", "gemini"],
+                },
+            ],
+        }
+        mcp.purge_uc_mcp_residue(state)
+        assert removed == [("claude", "system-ai-github")]
+        assert state["mcp_servers"] == []
+
+    def test_swallows_per_client_removal_errors(self, monkeypatch):
+        # One client failing to remove must not stop the rest.
+        attempts: list[tuple[str, str]] = []
+
+        def fake_remove(client, name):
+            attempts.append((client, name))
+            if client == "claude":
+                raise RuntimeError("claude config locked")
+            return ["user"]
+
+        self._patch_clients_installed(monkeypatch, ["claude", "codex"])
+        monkeypatch.setattr(mcp, "remove_client_mcp_server", fake_remove)
+        monkeypatch.setattr(mcp, "save_state", lambda s: None)
+
+        state = {
+            "workspace": WS,
+            "mcp_servers": [
+                {
+                    "name": "system-ai-github",
+                    "url": f"{WS}/ai-gateway/mcp-services/system.ai.github",
+                    "clients": ["claude", "codex"],
+                },
+            ],
+        }
+        mcp.purge_uc_mcp_residue(state)
+        assert attempts == [
+            ("claude", "system-ai-github"),
+            ("codex", "system-ai-github"),
+        ]
+        assert state["mcp_servers"] == []