From ccbdc5ca754fceb48a225a1d075cfa1de5afa1fd Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Sun, 20 Sep 2026 21:33:12 -0400 Subject: [PATCH 1/2] Handle stale MCP cleanup failures during workspace switches --- .github/workflows/integration.yml | 4 +- scripts/run_integration.py | 11 +- src/ucode/mcp.py | 27 +++-- tests/README.md | 5 +- tests/integration/README.md | 35 +++++- tests/integration/conftest.py | 19 +++ tests/integration/pytest.ini | 1 + ...st_ug_configure_claude_workspace_switch.py | 112 ++++++++++++++++++ tests/integration/utils/harness.py | 2 +- tests/test_cli.py | 43 +++++++ tests/test_integration_contract.py | 2 +- tests/test_mcp.py | 78 ++++++++++++ 12 files changed, 320 insertions(+), 19 deletions(-) create mode 100644 tests/integration/test_ug_configure_claude_workspace_switch.py diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index bd7e5d55..f6de84f0 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -227,6 +227,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 +239,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 diff --git a/scripts/run_integration.py b/scripts/run_integration.py index b7b418e3..59d53657 100644 --- a/scripts/run_integration.py +++ b/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, "") 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: diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index b36c24bc..3b6a7cdb 100644 --- a/src/ucode/mcp.py +++ b/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: diff --git a/tests/README.md b/tests/README.md index f73dd721..f45a8a83 100644 --- a/tests/README.md +++ b/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 | diff --git a/tests/integration/README.md b/tests/integration/README.md index 1923d49e..10b3a221 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -103,6 +103,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 +173,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 +201,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 diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 7f14fd53..23b168f6 100644 --- a/tests/integration/conftest.py +++ b/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 diff --git a/tests/integration/pytest.ini b/tests/integration/pytest.ini index ee9e4c58..2eb428be 100644 --- a/tests/integration/pytest.ini +++ b/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 diff --git a/tests/integration/test_ug_configure_claude_workspace_switch.py b/tests/integration/test_ug_configure_claude_workspace_switch.py new file mode 100644 index 00000000..4ec3a47d --- /dev/null +++ b/tests/integration/test_ug_configure_claude_workspace_switch.py @@ -0,0 +1,112 @@ +"""CUJ: switch workspaces after registering the Databricks skills MCP connection.""" + +import os +from urllib.parse import urlparse + +import pytest +from utils.evidence import FileTask + +pytestmark = [pytest.mark.workspace_switch, pytest.mark.claude] + +SKILLS_SERVER = "databricks-skill-registry" + + +def _assert_claude_reports_missing(result) -> None: + assert result.returncode != 0, result.stdout + output = (result.stdout + result.stderr).lower() + expected = ( + "no mcp server configured with that name", + "no mcp server found with name", + "no mcp server named", + "mcp server not found", + ) + assert any(message in output for message in expected), output + + +def test_ug_configure_claude_cleans_stale_skills_mcp_on_workspace_switch( + live_session, workspace, second_workspace +): + """Scenario: configure the managed workspace, register its skills MCP connection, then + configure a second real workspace in the same fresh home and use Claude there. + + Expected: the switch warns once, removes the old skills server from Claude's real config and + the new workspace state, preserves the old workspace bucket, remains stable on repeat + configure, and Claude completes a real task against the second workspace. This covers the + normal black-box cleanup path; injected subprocess failures stay in unit/component coverage. + """ + session = live_session + task = FileTask(session) + + session.run( + "configure", + "--agents", + "claude", + "--workspace", + workspace, + "--skip-upgrade", + "--disable-databricks-ai-tools", + timeout=240, + ) + session.run("skills", timeout=240) + + primary_state = session.workspace_state() + primary_entry = next( + server for server in primary_state["mcp_servers"] if server.get("name") == SKILLS_SERVER + ) + assert urlparse(primary_entry["url"]).hostname == urlparse(workspace).hostname + assert "claude" in primary_entry["clients"] + configured = session.run("mcp", "get", SKILLS_SERVER, binary="claude", timeout=60) + assert urlparse(workspace).hostname in (configured.stdout + configured.stderr) + + second_url = second_workspace + session.env["DATABRICKS_BEARER"] = os.environ["DATABRICKS_SECOND_BEARER"] + switched = session.run( + "configure", + "--agents", + "claude", + "--workspace", + second_url, + "--skip-upgrade", + "--disable-databricks-ai-tools", + timeout=240, + ) + assert switched.stdout.count("Dropping 1 stale MCP entry") == 1, switched.stdout + assert SKILLS_SERVER in switched.stdout, switched.stdout + missing = session.run("mcp", "get", SKILLS_SERVER, binary="claude", timeout=60, ok=False) + _assert_claude_reports_missing(missing) + + state = session.state() + assert state["current_workspace"] == second_url + assert all( + server.get("name") != SKILLS_SERVER + for server in state["workspaces"][second_url].get("mcp_servers", []) + ) + assert primary_entry in state["workspaces"][workspace]["mcp_servers"] + assert second_url in session.run("status").stdout + + repeated = session.run( + "configure", + "--agents", + "claude", + "--workspace", + second_url, + "--skip-upgrade", + "--disable-databricks-ai-tools", + timeout=240, + ) + assert "Dropping 1 stale MCP entry" not in repeated.stdout, repeated.stdout + still_missing = session.run("mcp", "get", SKILLS_SERVER, binary="claude", timeout=60, ok=False) + _assert_claude_reports_missing(still_missing) + + result = session.run( + "claude", + "--", + "-p", + task.prompt, + "--output-format", + "json", + "--allowedTools", + "Read", + timeout=180, + ) + task.assert_headless_answer("claude", result) diff --git a/tests/integration/utils/harness.py b/tests/integration/utils/harness.py index 9644bf4c..b72827f6 100644 --- a/tests/integration/utils/harness.py +++ b/tests/integration/utils/harness.py @@ -89,7 +89,7 @@ def __init__(self, root: Path, project_root: Path, binary: Path, artifacts: Path def redact(self, text: str, *, strip_ansi: bool = True) -> str: # Also scrub the relayed launch's subscription OAuth token, not just the bearer. - for name in ("DATABRICKS_BEARER", "CLAUDE_CODE_OAUTH_TOKEN"): + for name in ("DATABRICKS_BEARER", "DATABRICKS_SECOND_BEARER", "CLAUDE_CODE_OAUTH_TOKEN"): for token in (os.environ.get(name), self.env.get(name)): if token: text = text.replace(token, "") diff --git a/tests/test_cli.py b/tests/test_cli.py index 5bb6be07..eb0b0efb 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -4346,6 +4346,49 @@ def _stub_external_deps(monkeypatch): monkeypatch.setattr(cli_mod, "discover_codex_models", lambda w, t: ([], None)) monkeypatch.setattr(cli_mod, "build_shared_base_urls", lambda w: {}) + def test_workspace_switch_continues_after_skills_cleanup_timeout(self, monkeypatch, capsys): + from ucode import mcp + from ucode import state as state_mod + + self._stub_external_deps(monkeypatch) + monkeypatch.setattr(state_mod, "APP_DIR", state_mod.STATE_PATH.parent) + monkeypatch.setattr(state_mod, "build_agent_state", lambda state: {}) + # Use the real writer against the global fixture's temporary state file. + monkeypatch.setattr(cli_mod, "save_state", mcp.save_state) + old_workspace = "https://old.databricks.com" + new_workspace = "https://new.databricks.com" + entry = { + "name": "databricks-skill-registry", + "kind": "skills", + "url": f"{old_workspace}/ai-gateway/skills/", + "clients": ["claude"], + } + mcp.save_state({"workspace": old_workspace, "mcp_servers": [entry]}) + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude"]) + calls = [] + + def time_out(args, **kwargs): + assert args[:4] == ["claude", "mcp", "remove", "databricks-skill-registry"] + calls.append(args) + raise subprocess.TimeoutExpired(args, kwargs["timeout"]) + + monkeypatch.setattr(mcp.subprocess, "run", time_out) + + state = cli_mod.configure_shared_state(new_workspace, force_login=True) + + assert state["workspace"] == new_workspace + assert state["mcp_servers"] == [] + assert "_discovery_reasons" in state + assert len(calls) == 1 + full = state_mod.load_full_state() + assert full["current_workspace"] == new_workspace + assert full["workspaces"][new_workspace]["mcp_servers"] == [] + assert full["workspaces"][old_workspace]["mcp_servers"] == [entry] + output = " ".join(_strip_ansi(capsys.readouterr().out).split()) + assert "Unity Gateway connected" in output + assert "Dropping 1 stale MCP entry" in output + assert "Failed to remove `databricks-skill-registry` from Claude Code" in output + def test_purges_residue_when_workspace_changes(self, monkeypatch): import ucode.cli as cli_mod diff --git a/tests/test_integration_contract.py b/tests/test_integration_contract.py index def610a9..5e408b2c 100644 --- a/tests/test_integration_contract.py +++ b/tests/test_integration_contract.py @@ -67,7 +67,7 @@ def test_live_integration_cases_belong_to_exactly_one_ci_agent(): for node in tree.body: if isinstance(node, ast.FunctionDef) and node.name.startswith("test_"): marks = module_marks | _markers(node.decorator_list) - if marks & {"live", "managed"}: + if marks & {"live", "managed", "workspace_switch"}: assert len(marks & {"claude", "codex"}) == 1, node.name diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 051b8d70..bf04f696 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -2709,6 +2709,84 @@ def test_removes_skills_registry_across_its_clients(self, monkeypatch): class TestPurgeCrossWorkspaceSkillsEntry: + @pytest.mark.parametrize("copied_clients", [["claude"], ["claude", "codex"]]) + def test_workspace_switch_removes_each_client_once(self, monkeypatch, copied_clients): + from ucode import state as state_mod + + monkeypatch.setattr(state_mod, "APP_DIR", state_mod.STATE_PATH.parent) + monkeypatch.setattr(state_mod, "build_agent_state", lambda state: {}) + foreign = "https://other.databricks.com" + skills_entry = mcp._resolve_skills_mcp_servers( + foreign, ["claude", "codex"], _by_client(["claude", "codex"], ["a.b"]), [] + )[0] + state_mod.save_state({"workspace": foreign, "mcp_servers": [skills_entry]}) + # configure_shared_state carries the previous workspace's entries into + # the new bucket before invoking cleanup. + state = state_mod.load_state() + state["workspace"] = WS + state["mcp_servers"] = [{**skills_entry, "clients": copied_clients}] + state_mod.save_state(state) + removed: list[tuple[str, str]] = [] + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude", "codex"]) + monkeypatch.setattr( + mcp, + "remove_client_mcp_server", + lambda client, name: removed.append((client, name)) or ["user"], + ) + + mcp.purge_cross_workspace_mcp_residue(state, WS) + + assert removed == [ + ("claude", mcp.SKILLS_MCP_SERVER_NAME), + ("codex", mcp.SKILLS_MCP_SERVER_NAME), + ] + full = state_mod.load_full_state() + assert full["workspaces"][WS]["mcp_servers"] == [] + assert full["workspaces"][foreign]["mcp_servers"] == [skills_entry] + + @pytest.mark.parametrize("copied_to_current", [True, False], ids=["copied", "orphan"]) + @pytest.mark.parametrize("failure", ["timeout", "missing-binary", "command-error"]) + def test_removal_failure_warns_and_continues( + self, monkeypatch, capsys, copied_to_current, failure + ): + from ucode import state as state_mod + + monkeypatch.setattr(state_mod, "APP_DIR", state_mod.STATE_PATH.parent) + monkeypatch.setattr(state_mod, "build_agent_state", lambda state: {}) + foreign = "https://other.databricks.com" + skills_entry = mcp._resolve_skills_mcp_servers( + foreign, ["claude", "codex"], _by_client(["claude", "codex"], ["a.b"]), [] + )[0] + state_mod.save_state({"workspace": foreign, "mcp_servers": [skills_entry]}) + state = {"workspace": WS, "mcp_servers": [skills_entry] if copied_to_current else []} + state_mod.save_state(state) + calls: list[list[str]] = [] + + def run_removal(args, **kwargs): + calls.append(args) + assert args[1:4] == ["mcp", "remove", mcp.SKILLS_MCP_SERVER_NAME] + if args[0] == "claude": + if failure == "timeout": + raise mcp.subprocess.TimeoutExpired(args, kwargs["timeout"]) + if failure == "missing-binary": + raise FileNotFoundError(2, "No such file or directory", "claude") + raise mcp.subprocess.CalledProcessError(1, args, stderr="config is locked") + assert args[0] == "codex" + return mcp.subprocess.CompletedProcess(args, 0, stdout="Removed", stderr="") + + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude", "codex"]) + # Exercise the real client dispatch and exception conversion, stopping + # only at the subprocess boundary so no installed agent is modified. + monkeypatch.setattr(mcp.subprocess, "run", run_removal) + + mcp.purge_cross_workspace_mcp_residue(state, WS) + + assert sorted(args[0] for args in calls) == ["claude", "codex"] + output = _unwrap(capsys.readouterr().out) + assert output.count("Failed to remove `databricks-skill-registry` from Claude Code") == 1 + assert state_mod.load_state()["mcp_servers"] == [] + assert state_mod.load_full_state()["workspaces"][foreign]["mcp_servers"] == [skills_entry] + def test_drops_foreign_workspace_skills_entry(self, monkeypatch): removed: list[tuple[str, str]] = [] saved_states: list[dict] = [] From 4364e858d08f83b78d3f1404b1cdd33445b4ee90 Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Sun, 20 Sep 2026 21:53:08 -0400 Subject: [PATCH 2/2] ci: pin managed integration Databricks CLI for skills journeys --- .github/workflows/integration.yml | 2 ++ tests/integration/README.md | 3 +++ tests/test_integration_contract.py | 16 ++++++++++++++++ 3 files changed, 21 insertions(+) diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index f6de84f0..73692905 100644 --- a/.github/workflows/integration.yml +++ b/.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: diff --git a/tests/integration/README.md b/tests/integration/README.md index 10b3a221..2ac6a4dd 100644 --- a/tests/integration/README.md +++ b/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. diff --git a/tests/test_integration_contract.py b/tests/test_integration_contract.py index 5e408b2c..3ab5450b 100644 --- a/tests/test_integration_contract.py +++ b/tests/test_integration_contract.py @@ -1,6 +1,7 @@ """Keep the black-box suite independent of application internals and test doubles.""" import ast +import re from pathlib import Path @@ -17,6 +18,21 @@ def _markers(nodes): } +def test_integration_ci_pins_a_skills_capable_databricks_cli(): + from ucode.databricks import SKILLS_MCP_MIN_DATABRICKS_CLI_VERSION + + workflow = Path(__file__).parent.parent / ".github/workflows/integration.yml" + setup_blocks = re.findall( + r"(?m)^ - uses: databricks/setup-cli@[^\n]+\n((?: [^\n]*\n)*)", + workflow.read_text(), + ) + assert setup_blocks, "Integration CI must install the Databricks CLI explicitly" + for block in setup_blocks: + version = re.search(r"(?m)^ version: (\d+)\.(\d+)\.(\d+)\s*$", block) + assert version, "Every integration setup-cli step must pin an exact CLI version" + assert tuple(map(int, version.groups())) >= SKILLS_MCP_MIN_DATABRICKS_CLI_VERSION + + def test_integration_suite_uses_only_public_process_boundaries(): violations = [] for path in (Path(__file__).parent / "integration").rglob("*.py"):