Handle stale MCP cleanup failures when switching workspaces

.github/workflows/integration.yml

@@ -219,6 +219,8 @@ jobs:
         with:
           version: 0.9.8
       - uses: databricks/setup-cli@bdb89f81c11a5bd647fd55b585b7c396ec68a25a  # v1.0.0
+        with:
+          version: 1.17.0
       - name: Run the managed cases against the managed e2e workspace
         shell: bash
         env:
@@ -227,6 +229,8 @@ jobs:
           UCODE_TEST_WORKSPACE: ${{ secrets.E2E_ADMIN_WORKSPACE }}
           DATABRICKS_CLIENT_ID: ${{ secrets.E2E_ADMIN_SP_CLIENT_ID }}
           DATABRICKS_CLIENT_SECRET: ${{ secrets.E2E_ADMIN_SP_CLIENT_SECRET }}
+          UCODE_TEST_SECOND_WORKSPACE: ${{ secrets.UCODE_TEST_WORKSPACE }}
+          DATABRICKS_SECOND_BEARER: ${{ secrets.DATABRICKS_BEARER }}
         run: |
           # A managed config enables both agents and `ug configure` applies it to every enabled
           # agent, so both CLIs must be installed even though this lane asserts only one agent.
@@ -237,7 +241,7 @@ jobs:
           uv run --no-project --python 3.12 python scripts/run_integration.py \
             --python 3.12 --ug-version "$UG_VERSION" --entry-point "$ENTRY_POINT" \
             --default-index "$PACKAGE_INDEX" --output "$RUNNER_TEMP/ug-integration" \
-            "${args[@]}" -- -m "(managed or managed_fixture) and $AGENT"
+            "${args[@]}" -- -m "(managed or managed_fixture or workspace_switch) and $AGENT"
       - name: Upload managed test evidence
         if: ${{ !cancelled() }}
         uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02  # v4.6.2

scripts/run_integration.py

@@ -132,6 +132,11 @@ def arguments():
     parser.add_argument("--npm-registry", default="https://registry.npmjs.org")
     parser.add_argument("--profile", help="Explicit Databricks profile to mint the live bearer.")
     parser.add_argument("--workspace", default=os.environ.get("UCODE_TEST_WORKSPACE"))
+    parser.add_argument(
+        "--second-workspace",
+        default=os.environ.get("UCODE_TEST_SECOND_WORKSPACE"),
+        help="Second real workspace for workspace_switch CUJs; requires DATABRICKS_SECOND_BEARER.",
+    )
     parser.add_argument("--output", type=Path, help="New results directory; never reused.")
     parser.add_argument("--installation-only", action="store_true", help="No workspace calls.")
     parser.add_argument(
@@ -250,10 +255,11 @@ def terminate(signum, frame):
     base_env["UV_CACHE_DIR"] = str(output / "cache")
     base_env["UV_DEFAULT_INDEX"] = args.default_index
     bearer = os.environ.get("DATABRICKS_BEARER", "").strip()
+    second_bearer = os.environ.get("DATABRICKS_SECOND_BEARER", "").strip()
     oauth_token = os.environ.get("CLAUDE_CODE_OAUTH_TOKEN", "").strip()
 
     def redact(value: str) -> str:
-        for secret in (bearer, oauth_token):
+        for secret in (bearer, second_bearer, oauth_token):
             if secret:
                 value = value.replace(secret, "<redacted>")
         return value
@@ -298,6 +304,7 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str:
             "parent_schema": args.parent_schema,
             "dependencies": args.dependency,
             "workspace": args.workspace,
+            "second_workspace": args.second_workspace,
         },
         "platform": platform.platform(),
         "installation_only": args.installation_only,
@@ -539,6 +546,8 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str:
                 "UG_INTEGRATION_PARENT_SCHEMA": args.parent_schema,
                 "UCODE_TEST_WORKSPACE": args.workspace or "",
                 "DATABRICKS_BEARER": bearer,
+                "UCODE_TEST_SECOND_WORKSPACE": args.second_workspace or "",
+                "DATABRICKS_SECOND_BEARER": second_bearer,
             }
         )
         for agent in agents:

src/ucode/mcp.py

@@ -1221,6 +1221,18 @@ def message() -> str:
 
 def purge_cross_workspace_mcp_residue(state: dict, workspace: str) -> None:
     installed = set(available_mcp_clients())
+    attempted_removals: set[tuple[str, str]] = set()
+
+    def remove_stale_server(client: str, name: str) -> list[str] | None:
+        key = (client, name)
+        if key in attempted_removals:
+            return None
+        attempted_removals.add(key)
+        try:
+            return remove_client_mcp_server(client, name)
+        except (RuntimeError, subprocess.TimeoutExpired, OSError) as exc:
+            print_warning(f"Failed to remove `{name}` from {MCP_CLIENTS[client]['display']}: {exc}")
+            return None
 
     raw_mcp_servers = list(state.get("mcp_servers") or [])
     current_mcp_servers, foreign_mcp_servers = _partition_mcp_entries_by_workspace(
@@ -1242,12 +1254,7 @@ def purge_cross_workspace_mcp_residue(state: dict, workspace: str) -> None:
             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}"
-                    )
+                remove_stale_server(client, name)
         state["mcp_servers"] = current_mcp_servers
         save_state(state)
 
@@ -1258,13 +1265,7 @@ def purge_cross_workspace_mcp_residue(state: dict, workspace: str) -> None:
         for client in other_ws_mcps[name]:
             if client not in installed or client not in MCP_CLIENTS:
                 continue
-            try:
-                removed_scopes = remove_client_mcp_server(client, name)
-            except RuntimeError as exc:
-                print_warning(
-                    f"Failed to remove `{name}` from {MCP_CLIENTS[client]['display']}: {exc}"
-                )
-                continue
+            removed_scopes = remove_stale_server(client, name)
             if removed_scopes:
                 any_removed = True
         if any_removed:

tests/README.md

@@ -51,6 +51,7 @@ All tests live directly in `integration/`; shared mechanics live in `utils/`.
 | `test_smart_routing_claude_route_subagent_hook`, `test_smart_routing_codex_route_subagent_hook` | Pipe a real PreToolUse spawn payload to the installed route-subagent hook with subagent-only routing enabled | Allow decision against the live router; requested model replaced by a routed agent definition (Claude) or bundled catalog slug (Codex) from the offered models; one audited decision matching the session and task |
 | `test_smart_routing_claude_subagent_only_launch_shows_no_first_prompt_banner`, `test_smart_routing_codex_subagent_only_launch_shows_no_first_prompt_banner` | Configure, then launch the real TUI with both the full and subagent-only routing flags set and submit one file prompt | Subagent-only takes precedence: the prompt completes with no smart-routing banner and no first-prompt routing wrapper (PTY/interposer); Claude's SessionStart canary proves the routing hooks armed; normal exit |
 | `test_ug_configure_claude_repeat_and_revert`, `test_ug_configure_codex_repeat_and_revert` | Configure twice over user settings; complete a task; revert twice | Settings preserved; no bearer in ug state; generated config removed; status unconfigured |
+| `test_ug_configure_claude_cleans_stale_skills_mcp_on_workspace_switch` | Configure the first workspace, register its skills MCP, switch to a second real workspace, and use Claude | Old registration removed from Claude and the new workspace state; old workspace bucket preserved; repeat configure stays clean; real file task completes on the second workspace |
 | `test_ug_configure_claude_rejects_invalid_credentials`, `test_ug_configure_codex_rejects_invalid_credentials` | Configure with a rejected bearer against the real workspace | Authentication failure; no successful saved setup |
 | `test_ug_configure_managed_claude`, `test_ug_configure_managed_codex` | Configure against a workspace that publishes a managed CodingAgentConfig | No agent selector; each agent's generated config exposes exactly the admin's static model_services; real gateway prompt on launch |
 | `test_case_01_*`, `test_case_03_*` | Launch managed Claude after configure and from fresh state, with personal discovery enabled and disabled | Claude receives the admin MPS header, caches native discovery results, and opens its real model picker |
@@ -70,7 +71,8 @@ All tests live directly in `integration/`; shared mechanics live in `utils/`.
 
 With both agents selected there are **46 live cases** (6 interactive TUI cases),
 **4 managed-workspace cases** (marker `managed`, run against a separate workspace that
-publishes a CodingAgentConfig), **36 managed-fixture cases** (marker `managed_fixture`, with only
+publishes a CodingAgentConfig), **1 two-workspace case** (marker `workspace_switch`),
+**36 managed-fixture cases** (marker `managed_fixture`, with only
 the CodingAgentConfig input injected), and **5 installation checks**. The 24 numbered scenarios
 cover configured and fresh state across the Claude and Codex managed-discovery matrix; twelve
 existing collected cases cover focused model, MCP, skills, and lifecycle shapes. Parametrization
@@ -131,6 +133,7 @@ pending. The descriptive jobs provide the actual coverage and diagnostics.
 | --- | --- |
 | Live MCP and skills functionality | Deferred; installation tests cover the local web-search MCP handshake and tool listing, not upstream proxying or a real search request |
 | Broad configure flags, tracing, multiple workspaces, and PAT flows | Deferred while focusing on basic CUJs |
+| Workspace-switch MCP cleanup | The `workspace_switch` CUJ covers real registration, cleanup, repeat configure, and a completed Claude task. Unit/component tests cover duplicate attempts and injected removal failures; the CUJ does not force an agent timeout. It runs in the existing non-blocking managed CI lane. |
 | Provider switching, relayed/subscription MPS | Not covered by the four provider journeys |
 | TUI initial prompt supplied on the launch command line | Not yet covered; headless prompt arguments are covered |
 | Follow-up turns and conversation resume | Not covered; reopen proves startup, not conversation resume |

tests/integration/README.md

@@ -16,6 +16,9 @@ installs the requested agents into a new npm prefix and ug into a new virtualenv
 Pytest and the PTY/screen libraries (pexpect and pyte) live in a different virtualenv, so they cannot accidentally supply a
 missing application dependency. No packages are installed into your existing
 agent installations or checkout's `.venv`.
+CI pins Databricks CLI 1.17.0 in the live and managed integration lanes. The
+runner's isolated `PATH` exposes that selected CLI, so skills journeys meet ug's
+CLI minimum without falling back to another version installed on the machine.
 Native live runs refuse existing machine-wide Claude/Codex configuration, which
 could override the selected workspace even with a fresh home. Use a clean VM
 in that case; the runner never edits or bypasses those managed settings.
@@ -103,6 +106,7 @@ test_ug_codex_commands.py               # command help and parser error forwardi
 test_ug_codex_app_server.py             # actual client/server initialize exchange
 test_ug_smart_routing_hooks.py           # route-subagent hook contract against the live router
 test_ug_configure_claude_lifecycle.py   # repeat setup, revert, rejected credentials
+test_ug_configure_claude_workspace_switch.py # real skills MCP cleanup across two workspaces
 test_ug_configure_codex_lifecycle.py    # repeat setup, revert, rejected credentials
 test_ug_claude_managed_model_discovery.py # fetched/reused Claude MPS policy cases
 test_ug_codex_managed_model_discovery.py  # fetched/reused Codex MPS policy cases
@@ -172,7 +176,9 @@ fails the selected CUJ, rather than skipping it.
 There are **46 live cases** (including 6 TUI journeys) and **5 installation
 checks** with both agents. A separate **4 managed-workspace cases** (one per agent, an idempotent
 re-configure, and a cache-TTL journey; marker `managed`) run against a workspace that publishes a
-CodingAgentConfig; see "Managed-workspace journeys" below. A further **36 `managed_fixture`
+CodingAgentConfig; see "Managed-workspace journeys" below. One **`workspace_switch` case**
+uses two real workspaces and checks skills MCP cleanup and a completed Claude task.
+A further **36 `managed_fixture`
 cases** use `UCODE_MANAGED_CONFIG_STUB`. Twenty-four explicit configured/fresh Claude and Codex
 discovery and source-override journeys fetch the published config once per agent, replace that
 agent's static source with its dedicated MPS, and reuse the result. Twelve existing collected cases
@@ -198,6 +204,36 @@ config left after revert and banners on app-server stdout, remain assertions.
 Live MCP/skills functionality, tracing, the broad configure-option matrix, and other
 agents are outside this focused revision.
 
+The workspace-switch CUJ is an exception to that deferred multi-workspace scope:
+it configures the first workspace and registers its skills MCP through `ug skills`,
+switches to a second real host using that host's bearer, and verifies the stale
+registration is removed from Claude and the target workspace state. It preserves
+the first workspace's saved bucket, repeats configure, and requires a completed
+Claude file task on the second workspace. This exercises real commands, state,
+and agents without injected configuration. Deterministic duplicate-attempt and
+timeout/missing-executable regressions remain in `../test_mcp.py` and `../test_cli.py`;
+the CUJ does not force an agent failure.
+
+Run this case with both agents installed if either workspace's managed config
+enables both. Supply the second workspace and its bearer explicitly:
+
+```bash
+# DATABRICKS_SECOND_BEARER must already contain a token for SECOND_WORKSPACE_URL.
+python3.12 scripts/run_integration.py \
+  --ug-version checkout --claude-version 2.1.268 --codex-version 0.154.0 \
+  --workspace FIRST_WORKSPACE_URL --profile FIRST_WORKSPACE_PROFILE \
+  --second-workspace SECOND_WORKSPACE_URL -- -m workspace_switch
+```
+
+`UCODE_TEST_SECOND_WORKSPACE` is the environment equivalent of `--second-workspace`.
+Missing credentials or equal workspace hosts fail the selected test. The runner
+records both URLs in `versions.json`, redacts both bearers in evidence, and passes
+only the active workspace's bearer to each tested command. CI runs the case in
+the existing **Managed config · Claude** lane: the first host uses
+`E2E_ADMIN_WORKSPACE` and its service-principal credentials; the second uses the
+existing `UCODE_TEST_WORKSPACE` / `DATABRICKS_BEARER` secrets. That lane remains
+non-blocking under its existing policy. Collection or lint success is not a live pass.
+
 The configure terminal helper recognizes `[✓]` / `[ ]` agent checkboxes as well
 as legacy markers in older pinned ug releases. It explicitly toggles
 the requested agent on and all others off before submitting; the existing live

tests/integration/conftest.py

@@ -7,6 +7,7 @@
 import shutil
 import tempfile
 from pathlib import Path
+from urllib.parse import urlparse
 
 import pytest
 from utils.harness import UserSession
@@ -44,6 +45,24 @@ def workspace():
     return value
 
 
+@pytest.fixture(scope="session")
+def second_workspace(workspace):
+    value = os.environ.get("UCODE_TEST_SECOND_WORKSPACE", "").strip().rstrip("/")
+    parsed = urlparse(value)
+    if (
+        parsed.scheme != "https"
+        or not parsed.hostname
+        or not os.environ.get("DATABRICKS_SECOND_BEARER", "").strip()
+    ):
+        pytest.fail(
+            "Workspace-switch CUJs require --second-workspace (or UCODE_TEST_SECOND_WORKSPACE) "
+            "and DATABRICKS_SECOND_BEARER for that workspace."
+        )
+    if parsed.hostname == urlparse(workspace).hostname:
+        pytest.fail("Workspace-switch CUJs require two distinct workspace hosts.")
+    return value
+
+
 @pytest.fixture
 def session(request, installed_binary):
     # Codex rejects helper installation beneath /tmp. Keep the disposable home

tests/integration/pytest.ini

@@ -6,6 +6,7 @@ markers =
     live: requires the real workspace used by the existing e2e suite
     managed: requires the managed e2e workspace that publishes a CodingAgentConfig
     managed_fixture: real ug/TUI against a real workspace, but the managed CodingAgentConfig is injected via UCODE_MANAGED_CONFIG_STUB
+    workspace_switch: requires two real workspace hosts and their respective credentials
     smoke: Databricks Hosted, custom OAuth CLI TUI, and headless prompt for each agent
     tui: real interactive terminal boot, keyboard input, exit and reopen
     claude: only runs when Claude Code is explicitly selected