diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 15973b0..45404d8 100644 --- a/src/ucode/cli.py +++ b/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"] = { diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index e961ff8..86db7ab 100644 --- a/src/ucode/mcp.py +++ b/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") diff --git a/tests/test_e2e_uc.py b/tests/test_e2e_uc.py index d0f994d..4213f78 100644 --- a/tests/test_e2e_uc.py +++ b/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 ): diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 333bb32..c8e65d6 100644 --- a/tests/test_mcp.py +++ b/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"] == []