From 10fe9973d6835b9caf1e8b0bdd103ab5d3de01ae Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Thu, 29 Oct 2026 23:33:03 +0000 Subject: [PATCH 1/6] [BR-2167] Keep MCP add scoped to the requested workspace and agents --- src/ucode/cli.py | 93 ++++++++++++++++------------------------------- src/ucode/mcp.py | 38 +++++++++++++++++-- tests/test_cli.py | 70 +++++++++++++++++++++++++---------- tests/test_mcp.py | 51 ++++++++++++++++++++++++++ 4 files changed, 168 insertions(+), 84 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 96e84088a..0fdcac0eb 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -138,7 +138,6 @@ ) from ucode.mcp import ( MCP_CLIENTS, - MCP_ONLY_CLIENTS, McpServiceListingRateLimited, add_mcp_command, add_skills_command, @@ -146,7 +145,6 @@ configure_bare_skills_mcp_command, configure_mcp_command, configure_skills_mcp_picker_command, - configured_mcp_clients, configured_mcp_servers_by_name, list_mcp_command, purge_cross_workspace_mcp_residue, @@ -1903,68 +1901,36 @@ def _configure_agents_for_mcp( requested: list[str], workspace_entries: list[tuple[str, str | None]] | None = None, ) -> set[str]: - """Ensure the named coding agents are set up (workspace + models) so a - subsequent `ug mcp add` / `ug skills add --via mcp` has them as targets, and - return the full canonical name set. Agents already configured are left as-is; - only the rest are bootstrapped. Model agents go through - configure_workspace_command (which installs binaries and configures models); - Cursor is MCP-only, so it just needs workspace state established and rides - along via MCP_ONLY_CLIENTS. ``workspace_entries`` (from `--workspace` / - `--profile`) makes that workspace current and skips the workspace prompt; - otherwise the prompt is shown on first run. If setup fails, the previous - current workspace is restored.""" + """Select MCP clients for one workspace without changing agent model settings.""" scope = {a if a == "cursor" else normalize_tool(a) for a in requested} installed = available_mcp_clients() - # Setup installs model agents but not MCP-only ones, so a missing MCP-only CLI - # would otherwise leave the agent permanently unready and re-enter setup every run. - not_installed = sorted(a for a in scope if a in MCP_ONLY_CLIENTS and a not in installed) + not_installed = sorted(scope - set(installed)) if not_installed: raise RuntimeError(_not_installed_message(not_installed)) previous_workspace = load_full_state().get("current_workspace") try: - _bootstrap_agents_for_mcp(scope, installed, workspace_entries) - except BaseException: - # Don't leave launches pointed at a workspace whose setup didn't finish. if workspace_entries: - set_current_workspace(previous_workspace) - raise - return scope - - -def _bootstrap_agents_for_mcp( - scope: set[str], - installed: list[str], - workspace_entries: list[tuple[str, str | None]] | None, -) -> None: - workspace_known = True - if workspace_entries: - url = workspace_entries[0][0] - set_current_workspace(url) - workspace_known = url in (load_full_state().get("workspaces") or {}) - ready = set(configured_mcp_clients(load_state(), installed)) if workspace_known else set() - to_bootstrap = scope - ready - model_agents = sorted(a for a in to_bootstrap if a != "cursor") - if model_agents: - if workspace_entries is None: - configure_workspace_command(selected_tools=model_agents) + workspace, profile = workspace_entries[0] + set_current_workspace(workspace) else: - configure_workspace_command(selected_tools=model_agents, workspaces=workspace_entries) - if "cursor" in to_bootstrap and not model_agents: - _configure_shared_workspace_states( - workspace_entries or [_prompt_for_configuration(None)], tools=[], force_login=True - ) - state = load_state() - unready = sorted(scope - set(configured_mcp_clients(state, available_mcp_clients()))) - if unready: - raise RuntimeError(_not_installed_message(unready)) - # Setup saves an explicit --profile, but it's skipped when every agent is already - # ready; save it here so registration authenticates as that profile. A saved - # --use-pat belonged to the old profile, and --profile means OAuth without it. - profile = workspace_entries[0][1] if workspace_entries else None - if profile and state.get("profile") != profile: - state["profile"] = profile - state.pop("use_pat", None) + state = load_state() + workspace = state.get("workspace") + if not workspace: + workspace, profile = _prompt_for_configuration(None) + set_current_workspace(workspace) + else: + profile = None + state = load_state() + state["workspace"] = workspace + state["mcp_clients"] = sorted(set(state.get("mcp_clients") or []) | scope) + if profile and state.get("profile") != profile: + state["profile"] = profile + state.pop("use_pat", None) save_state(state) + except BaseException: + set_current_workspace(previous_workspace) + raise + return scope def _not_installed_message(agents: list[str]) -> str: @@ -2040,8 +2006,7 @@ def mcp_add( typer.Option( "--agents", help="Comma-separated coding agents to register the server(s) for (e.g. " - "claude,codex,cursor). Any that aren't configured yet are set up first " - "(workspace + models), so this works as a one-command setup. Without --agents, " + "claude,codex,cursor). MCP registration is independent of model setup. Without --agents, " "the server is registered for every already-configured agent.", ), ] = None, @@ -2049,8 +2014,8 @@ def mcp_add( str | None, typer.Option( "--workspace", - help="Workspace URL to register --agents against (made the current workspace; set " - "up first if needed), skipping the workspace prompt. Defaults to the UG_WORKSPACE " + help="Workspace URL to register --agents against, without changing the launch " + "workspace or model settings. Defaults to the UG_WORKSPACE " "environment variable when set.", ), ] = None, @@ -2068,7 +2033,7 @@ def mcp_add( Purely additive: it never removes MCP servers that are already configured, only registers new ones (use `ug mcp remove` to remove). Pass --agents to target (and, if needed, set up) specific - agents. + agents. Explicit MCP setup never configures agent models or machine-wide settings. """ selected = None if names is None else {s.strip() for s in names.split(",") if s.strip()} requested_agents = ( @@ -2076,6 +2041,8 @@ def mcp_add( if agents is None else ({a.strip().lower() for a in agents.split(",") if a.strip()} or None) ) + previous_workspace = load_full_state().get("current_workspace") + workspace_entries = None try: # `ug --workspace mcp add ...` is parsed by the root command, which # returns early for subcommands, so forward its value here. @@ -2105,13 +2072,17 @@ def mcp_add( else None ) add_mcp_command(location=location, services=selected, agents=scope) - _sync_claude_cowork_plugin() + if not workspace_entries or workspace_entries[0][0] == previous_workspace: + _sync_claude_cowork_plugin() except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None except KeyboardInterrupt: print_err("Interrupted.") raise typer.Exit(130) from None + finally: + if requested_agents and workspace_entries: + set_current_workspace(previous_workspace) @mcp_app.command("remove") diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index d441c345d..7ec742cc9 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -261,7 +261,7 @@ def configured_mcp_clients(state: dict, installed_clients: list[str]) -> list[st configured_tools = state.get("available_tools") or [] if not isinstance(configured_tools, list): configured_tools = [] - configured = set(configured_tools) + configured = set(configured_tools) | set(state.get("mcp_clients") or []) return [ client for client in MCP_CLIENTS @@ -1882,6 +1882,7 @@ def setup_mcp_clients( action_note: str = "Configuring for", agents: set[str] | None = None, quiet: bool = False, + purge_foreign: bool = True, ) -> tuple[str, str | None, list[str]]: """Validate the workspace, resolve configured MCP clients, and prepare auth. @@ -1904,7 +1905,8 @@ def setup_mcp_clients( if not workspace: raise RuntimeError("Workspace is not configured. Run `ug configure` first.") - purge_cross_workspace_mcp_residue(state, workspace) + if purge_foreign: + purge_cross_workspace_mcp_residue(state, workspace) installed_clients = available_mcp_clients() if not installed_clients: @@ -1957,6 +1959,25 @@ def _union_missing(base: list[dict], selected: list[dict]) -> list[dict]: return [*selected, *extra] +def _namespace_foreign_mcp_names(servers: list[dict], workspace: str) -> list[dict]: + """Keep a second workspace from replacing an agent's same-named MCP server.""" + foreign_names = { + name + for other_workspace, bucket in (load_full_state().get("workspaces") or {}).items() + if other_workspace != workspace and isinstance(bucket, dict) + for key in ("mcp_servers", "managed_mcp_servers") + for entry in bucket.get(key) or [] + if (name := _server_name(entry)) + } + suffix = _normalize_workspace_title(workspace_hostname(workspace)) + return [ + {**server, "name": f"{server['name']}-{suffix}"} + if server.get("name") in foreign_names and not server["name"].endswith(f"-{suffix}") + else server + for server in servers + ] + + def add_mcp_command( location: str | None = None, services: set[str] | None = None, @@ -1998,7 +2019,10 @@ def _configure_v2_mcp_selectors( exact replacement otherwise (replace mode), always preserving the skills connection.""" state = load_state() workspace, profile, clients = setup_mcp_clients( - state, "Add MCP Servers" if append else "MCP Servers", agents=agents + state, + "Add MCP Servers" if append else "MCP Servers", + agents=agents, + purge_foreign=not append, ) # `app:` selectors need the app's off-workspace URL, which only discovery knows. A 403 here @@ -2040,6 +2064,7 @@ def _configure_v2_mcp_selectors( ) if append: + working_mcp_servers = _namespace_foreign_mcp_names(working_mcp_servers, workspace) working_mcp_servers = _union_missing(original_mcp_servers, working_mcp_servers) changed = apply_mcp_server_changes( @@ -2109,7 +2134,10 @@ def configure_mcp_command( location = next(iter(schemas)) state = load_state() workspace, profile, clients = setup_mcp_clients( - state, "Add MCP Servers" if append else "MCP Servers", agents=agents + state, + "Add MCP Servers" if append else "MCP Servers", + agents=agents, + purge_foreign=not append, ) managed_names = managed_mcp_server_names(state, set(clients)) @@ -2121,6 +2149,7 @@ def configure_mcp_command( ) working_mcp_servers = _drop_managed_servers(working_mcp_servers, managed_names) if append: + working_mcp_servers = _namespace_foreign_mcp_names(working_mcp_servers, workspace) working_mcp_servers = _union_missing( original_mcp_servers_for_location, working_mcp_servers ) @@ -2222,6 +2251,7 @@ def configure_mcp_command( working_names.add(entry_name) if append: + working_mcp_servers = _namespace_foreign_mcp_names(working_mcp_servers, workspace) working_mcp_servers = _union_missing(original_mcp_servers, working_mcp_servers) changed = apply_mcp_server_changes( diff --git a/tests/test_cli.py b/tests/test_cli.py index f460fc67b..bf83e3898 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3192,25 +3192,23 @@ def test_agents_is_rejected_for_download_mode(self): class TestConfigureAgentsForMcp: - def test_bootstraps_only_unconfigured_and_returns_full_scope(self): + def test_records_mcp_clients_without_configuring_models(self): with ( patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), patch("ucode.cli.available_mcp_clients", return_value=["claude", "codex"]), - patch( - "ucode.cli.configured_mcp_clients", side_effect=[["claude"], ["claude", "codex"]] - ), patch("ucode.cli.configure_workspace_command") as mock_cfg, + patch("ucode.cli.save_state") as mock_save, ): scope = cli_mod._configure_agents_for_mcp(["claude", "codex"]) assert scope == {"claude", "codex"} - mock_cfg.assert_called_once_with(selected_tools=["codex"]) + mock_cfg.assert_not_called() + assert mock_save.call_args.args[0]["mcp_clients"] == ["claude", "codex"] def test_all_configured_skips_bootstrap(self): with ( patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), patch("ucode.cli.available_mcp_clients", return_value=["claude", "codex"]), - patch("ucode.cli.configured_mcp_clients", return_value=["claude", "codex"]), patch("ucode.cli.configure_workspace_command") as mock_cfg, ): scope = cli_mod._configure_agents_for_mcp(["claude", "codex"]) @@ -3218,20 +3216,19 @@ def test_all_configured_skips_bootstrap(self): assert scope == {"claude", "codex"} mock_cfg.assert_not_called() - def test_workspace_entries_select_workspace_and_forward_to_bootstrap(self): + def test_workspace_entries_select_workspace_without_bootstrap(self): entries = [("https://ws", None)] with ( patch("ucode.cli.set_current_workspace") as mock_switch, patch("ucode.cli.load_full_state", return_value={"workspaces": {}}), patch("ucode.cli.load_state", return_value={}), patch("ucode.cli.available_mcp_clients", return_value=["claude"]), - patch("ucode.cli.configured_mcp_clients", return_value=["claude"]), patch("ucode.cli.configure_workspace_command") as mock_cfg, ): cli_mod._configure_agents_for_mcp(["claude"], entries) mock_switch.assert_called_once_with("https://ws") - mock_cfg.assert_called_once_with(selected_tools=["claude"], workspaces=entries) + mock_cfg.assert_not_called() def test_known_workspace_with_ready_agents_switches_without_bootstrap(self): entries = [("https://ws", None)] @@ -3240,7 +3237,6 @@ def test_known_workspace_with_ready_agents_switches_without_bootstrap(self): patch("ucode.cli.load_full_state", return_value={"workspaces": {"https://ws": {}}}), patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), patch("ucode.cli.available_mcp_clients", return_value=["claude", "cursor"]), - patch("ucode.cli.configured_mcp_clients", return_value=["claude", "cursor"]), patch("ucode.cli.configure_workspace_command") as mock_cfg, patch("ucode.cli._configure_shared_workspace_states") as mock_states, ): @@ -3251,14 +3247,13 @@ def test_known_workspace_with_ready_agents_switches_without_bootstrap(self): mock_cfg.assert_not_called() mock_states.assert_not_called() - def test_cursor_on_unseen_workspace_sets_it_up_without_prompting(self): + def test_cursor_on_unseen_workspace_is_mcp_only(self): entries = [("https://ws", None)] with ( patch("ucode.cli.set_current_workspace"), patch("ucode.cli.load_full_state", return_value={"workspaces": {}}), patch("ucode.cli.load_state", return_value={}), patch("ucode.cli.available_mcp_clients", return_value=["cursor"]), - patch("ucode.cli.configured_mcp_clients", return_value=["cursor"]), patch("ucode.cli._prompt_for_configuration") as mock_prompt, patch("ucode.cli._configure_shared_workspace_states") as mock_states, ): @@ -3266,7 +3261,7 @@ def test_cursor_on_unseen_workspace_sets_it_up_without_prompting(self): assert scope == {"cursor"} mock_prompt.assert_not_called() - mock_states.assert_called_once_with(entries, tools=[], force_login=True) + mock_states.assert_not_called() def test_missing_mcp_only_cli_errors_instead_of_prompting(self): with ( @@ -3283,7 +3278,6 @@ def test_agent_still_unready_after_setup_errors(self): with ( patch("ucode.cli.load_state", return_value={}), patch("ucode.cli.available_mcp_clients", return_value=[]), - patch("ucode.cli.configured_mcp_clients", return_value=[]), patch("ucode.cli.configure_workspace_command"), pytest.raises(RuntimeError, match="Not installed or not set up: Codex"), ): @@ -3406,12 +3400,51 @@ def test_new_profile_drops_saved_pat_mode(self): assert state["profile"] == "NEW" assert "use_pat" not in state + def test_second_workspace_add_preserves_launch_and_model_configuration(self): + from ucode.state import load_full_state + + self._save_workspace(self.WS_A, profile="A", mcp_servers=[{"name": "existing"}]) + before = load_full_state()["workspaces"][self.WS_A].copy() + with ( + patch("ucode.cli.available_mcp_clients", return_value=["claude", "codex"]), + patch("ucode.cli.configure_workspace_command") as configure_models, + patch("ucode.cli._configure_shared_workspace_states") as configure_workspace, + patch("ucode.cli.add_mcp_command") as add_mcp, + patch("ucode.cli._sync_claude_cowork_plugin") as sync_cowork, + ): + result = runner.invoke( + app, + [ + "mcp", + "add", + "--agents", + "claude,codex", + "--names", + "a.b.c", + "--workspace", + self.WS_B, + ], + ) + + assert result.exit_code == 0, result.output + configure_models.assert_not_called() + configure_workspace.assert_not_called() + sync_cowork.assert_not_called() + add_mcp.assert_called_once_with( + location=None, services={"a.b.c"}, agents={"claude", "codex"} + ) + full = load_full_state() + assert full["current_workspace"] == self.WS_A + assert full["workspaces"][self.WS_A] == before + assert full["workspaces"][self.WS_B]["mcp_clients"] == ["claude", "codex"] + assert full["workspaces"][self.WS_B].get("available_tools") is None + def test_failed_setup_with_no_previous_workspace_clears_current(self): from ucode.state import load_full_state with ( patch("ucode.cli.available_mcp_clients", return_value=["codex"]), - patch("ucode.cli.configure_workspace_command", side_effect=RuntimeError("boom")), + patch("ucode.cli.save_state", side_effect=RuntimeError("boom")), pytest.raises(RuntimeError, match="boom"), ): cli_mod._configure_agents_for_mcp(["codex"], [(self.WS_B, None)]) @@ -3427,8 +3460,7 @@ def test_failed_setup_restores_previous_workspace(self, error, exit_code): self._save_workspace(self.WS_A) with ( patch("ucode.cli.available_mcp_clients", return_value=["codex"]), - patch("ucode.cli.configure_workspace_command", side_effect=error), - patch("ucode.cli.add_mcp_command") as mock_add, + patch("ucode.cli.add_mcp_command", side_effect=error) as mock_add, ): result = runner.invoke( app, @@ -3436,10 +3468,10 @@ def test_failed_setup_restores_previous_workspace(self, error, exit_code): ) assert result.exit_code == exit_code, result.output - mock_add.assert_not_called() + mock_add.assert_called_once() full = load_full_state() assert full["current_workspace"] == self.WS_A - assert self.WS_B not in full["workspaces"] + assert full["workspaces"][self.WS_B]["mcp_clients"] == ["codex"] class TestMcpAddRootWorkspace: diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 87dac057b..7a023bcb4 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -33,6 +33,25 @@ def test_no_changes_falls_back_to_saved(self): assert mcp._mcp_change_summary([], [], ["claude"]) == "Saved" +def test_mcp_only_client_is_available_without_model_configuration(): + state = {"workspace": WS, "mcp_clients": ["claude"]} + assert mcp.configured_mcp_clients(state, ["claude", "codex"]) == ["claude"] + + +def test_second_workspace_mcp_name_does_not_replace_first(monkeypatch): + other = "https://first.cloud.databricks.com" + current = "https://second.cloud.databricks.com" + monkeypatch.setattr( + mcp, + "load_full_state", + lambda: {"workspaces": {other: {"mcp_servers": [{"name": "system-ai-github"}]}}}, + ) + selected = [{"name": "system-ai-github", "url": f"{current}/mcp"}] + result = mcp._namespace_foreign_mcp_names(selected, current) + assert result[0]["name"] == "system-ai-github-second-cloud-databricks-com" + assert selected[0]["name"] == "system-ai-github" + + # The proxy argv every client registers as a stdio command. The leading element # is the resolved `ug` binary path, so tests assert the tail (the stable part). GH_URL = f"{WS}/api/2.0/mcp/external/github" @@ -1894,6 +1913,38 @@ def test_existing_entry_gets_reconfigured_for_newly_added_clients(self, monkeypa class TestAddMcpCommand: """`ucode mcp add` (append) registers new servers without removing existing ones.""" + def test_second_workspace_keeps_first_workspace_server(self, monkeypatch): + other = "https://first.cloud.databricks.com" + saved: list[dict] = [] + registered: list[str] = [] + _stub_location_base(monkeypatch, {**CLAUDE_STATE, "mcp_servers": []}) + monkeypatch.setattr( + mcp, + "load_full_state", + lambda: { + "workspaces": { + other: {"mcp_servers": [{"name": "system-ai-github", "clients": ["claude"]}]} + } + }, + ) + monkeypatch.setattr( + mcp, + "purge_cross_workspace_mcp_residue", + lambda *a: pytest.fail("purged other workspace"), + ) + monkeypatch.setattr(mcp, "list_mcp_services", lambda *a, **kw: (["system.ai.github"], None)) + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, name, *a, **kw: registered.append(name) or [], + ) + monkeypatch.setattr(mcp, "save_state", lambda state: saved.append(state.copy())) + + assert mcp.add_mcp_command(location="system.ai", agents={"claude"}) == 0 + expected = "system-ai-github-example-databricks-com" + assert registered == [expected] + assert saved[-1]["mcp_servers"][0]["name"] == expected + def test_keeps_servers_outside_location(self, monkeypatch): """Unlike a replace-mode `--location` run, `mcp add --location` preserves any server outside the location instead of removing it.""" From 54bfd9dd06f74684b0a64cbc4392f4b743f892d2 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Thu, 29 Oct 2026 23:36:01 +0000 Subject: [PATCH 2/6] [BR-2167] Honor second-workspace MCP add without agents --- src/ucode/cli.py | 29 ++++++++++++++++++++++------- tests/test_cli.py | 32 +++++++++++++++++++++++++++++--- 2 files changed, 51 insertions(+), 10 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 0fdcac0eb..b50c598f4 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -145,6 +145,7 @@ configure_bare_skills_mcp_command, configure_mcp_command, configure_skills_mcp_picker_command, + configured_mcp_clients, configured_mcp_servers_by_name, list_mcp_command, purge_cross_workspace_mcp_residue, @@ -2014,7 +2015,7 @@ def mcp_add( str | None, typer.Option( "--workspace", - help="Workspace URL to register --agents against, without changing the launch " + help="Workspace URL to register MCP servers against, without changing the launch " "workspace or model settings. Defaults to the UG_WORKSPACE " "environment variable when set.", ), @@ -2023,7 +2024,7 @@ def mcp_add( str | None, typer.Option( "--profile", - help="Existing Databricks CLI profile to register --agents against, like " + help="Existing Databricks CLI profile to register MCP servers against, like " "--workspace. The profile's host supplies the workspace URL, and the profile is " "saved as that workspace's profile.", ), @@ -2057,19 +2058,33 @@ def mcp_add( workspace = root_workspace if workspace is not None and profile is not None: raise RuntimeError("Use either --workspace or --profile, not both.") - if (workspace is not None or profile is not None) and not requested_agents: - raise RuntimeError("--workspace and --profile only apply with --agents.") - if workspace is None and profile is None and requested_agents: + if workspace is None and profile is None: workspace = os.environ.get("UG_WORKSPACE") or None workspace_entries = None if profile is not None: workspace_entries = _parse_profile_option(profile) elif workspace is not None: workspace_entries = _parse_workspace_option(workspace) + existing_scope = None + if workspace_entries and not requested_agents: + target = workspace_entries[0][0] + bucket = (load_full_state().get("workspaces") or {}).get(target) + recorded = set((bucket or {}).get("available_tools") or []) | set( + (bucket or {}).get("mcp_clients") or [] + ) + existing_scope = ( + set(configured_mcp_clients(bucket or {}, available_mcp_clients())) & recorded + ) + if not existing_scope: + raise RuntimeError( + f"No MCP clients are configured for {target}. " + "Pass --agents to select the clients for this workspace." + ) + _configure_agents_for_mcp([], workspace_entries) scope = ( _configure_agents_for_mcp(sorted(requested_agents), workspace_entries) if requested_agents - else None + else existing_scope ) add_mcp_command(location=location, services=selected, agents=scope) if not workspace_entries or workspace_entries[0][0] == previous_workspace: @@ -2081,7 +2096,7 @@ def mcp_add( print_err("Interrupted.") raise typer.Exit(130) from None finally: - if requested_agents and workspace_entries: + if workspace_entries: set_current_workspace(previous_workspace) diff --git a/tests/test_cli.py b/tests/test_cli.py index bf83e3898..768aec868 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3319,14 +3319,14 @@ def test_ug_workspace_env_used_when_no_flag(self, monkeypatch): assert result.exit_code == 0, result.output configure.assert_called_once_with(["cursor"], [("https://ws", None)]) - def test_workspace_without_agents_exit_1(self): + def test_new_workspace_without_agents_explains_selection(self): with patch("ucode.cli.add_mcp_command") as mock_add: result = runner.invoke( app, ["mcp", "add", "--names", "a.b.c", "--workspace", "https://ws"] ) assert result.exit_code == 1 - assert "only apply with --agents" in _strip_ansi(result.output) + assert "Pass --agents to select" in _strip_ansi(result.output) mock_add.assert_not_called() def test_workspace_and_profile_together_exit_1(self): @@ -3439,6 +3439,32 @@ def test_second_workspace_add_preserves_launch_and_model_configuration(self): assert full["workspaces"][self.WS_B]["mcp_clients"] == ["claude", "codex"] assert full["workspaces"][self.WS_B].get("available_tools") is None + @pytest.mark.parametrize("use_env", [False, True]) + def test_second_workspace_without_agents_uses_its_existing_mcp_clients( + self, monkeypatch, use_env + ): + from ucode.state import load_full_state + + self._save_workspace(self.WS_B, mcp_clients=["codex"]) + self._save_workspace(self.WS_A) + if use_env: + monkeypatch.setenv("UG_WORKSPACE", self.WS_B) + args = ["mcp", "add", "--names", "a.b.c"] + if not use_env: + args += ["--workspace", self.WS_B] + with ( + patch("ucode.cli.available_mcp_clients", return_value=["codex"]), + patch("ucode.cli.add_mcp_command") as add_mcp, + patch("ucode.cli.configure_workspace_command") as configure_models, + ): + result = runner.invoke(app, args) + + assert result.exit_code == 0, result.output + add_mcp.assert_called_once_with(location=None, services={"a.b.c"}, agents={"codex"}) + configure_models.assert_not_called() + assert load_full_state()["current_workspace"] == self.WS_A + assert load_full_state()["workspaces"][self.WS_B]["mcp_clients"] == ["codex"] + def test_failed_setup_with_no_previous_workspace_clears_current(self): from ucode.state import load_full_state @@ -3526,7 +3552,7 @@ def test_same_workspace_in_both_positions_is_accepted(self): [ (["--agents", "codex", "--workspace", "https://other"], "two different --workspace"), (["--agents", "codex", "--profile", "P"], "not both"), - (["--names", "a.b.c"], "only apply with --agents"), + (["--names", "a.b.c"], "Pass --agents to select"), ], ) def test_root_workspace_conflicts_exit_1(self, extra, message): From 31dec717da0db7d554a9414abb362dfce12fc1d6 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Fri, 30 Oct 2026 01:28:58 +0000 Subject: [PATCH 3/6] [BR-2167] Switch MCP registrations with the active workspace --- README.md | 8 ++- src/ucode/cli.py | 59 +++++++++++++++---- src/ucode/mcp.py | 50 +++++++++------- tests/test_cli.py | 143 +++++++++++++++++++++++++++++++++++++++++----- tests/test_mcp.py | 58 ++++++++++--------- 5 files changed, 242 insertions(+), 76 deletions(-) diff --git a/README.md b/README.md index f3eb06598..217222712 100644 --- a/README.md +++ b/README.md @@ -224,9 +224,11 @@ ug mcp add --agents cursor --names system.ai.github --workspace https://first.da ug mcp add --agents codex --names system.ai.github --profile DEFAULT ``` -With `--agents`, agents that aren't set up yet are set up first. Pass `--workspace` -or `--profile` to pick the workspace instead of the picker; it becomes the current -workspace. +With `--agents`, installed agents can be registered for MCP without configuring their +models. `--workspace` or `--profile` selects the workspace on a fresh setup. If another +workspace is already active, run `ug configure --workspace ` to switch first; +the switch replaces the agent's active MCP registrations with those saved for the +selected workspace. Switching back restores the earlier workspace's saved servers. Remove configured servers: diff --git a/src/ucode/cli.py b/src/ucode/cli.py index b50c598f4..3b9d9d03c 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -153,6 +153,7 @@ remove_mcp_command, remove_skills_command, remove_skills_locations_command, + restore_workspace_mcp_servers, revert_mcp_configs, ) from ucode.mcp import managed_mcp_servers as load_managed_mcp_servers @@ -566,6 +567,15 @@ def configure_shared_state( # model lists (carried over by load_state, left untouched). state = load_state() state["workspace"] = workspace + if previous_workspace != workspace: + # Model discovery below refreshes the selected workspace. MCP selections are different: + # they are user choices, so carry the destination's saved registry across a switch. + destination = load_full_state().get("workspaces", {}).get(workspace, {}) + for key in ("mcp_servers", "mcp_clients", "managed_mcp_servers"): + if key in destination: + state[key] = destination[key] + else: + state.pop(key, None) # `state` starts from the current workspace's block, so its self-managed list belongs to the # source workspace. The opt-in is per workspace: replace it with the destination's own list. dest_self_managed = workspace_self_managed_agents(workspace) @@ -903,6 +913,7 @@ def _setup_single_agent( rejected). ``ug agents add`` passes ``apply_managed=False`` to set up a self-managed agent standalone, independent of the managed config. """ + previous_workspace = load_state().get("workspace") states = _configure_shared_workspace_states( workspace_entries, [tool], @@ -936,11 +947,14 @@ def _setup_single_agent( from ucode.agents.desktop.setup import configure_claude_desktop configure_claude_desktop(state, parent_schema=parent_schema) - if not is_dry_run(): - resolved_mcps = _configure_managed_mcp_servers(managed) - skills_configured = _configure_managed_skills(managed) - if resolved_mcps is not None and skills_configured: - _sync_claude_cowork_plugin(state, managed_mcp_servers=resolved_mcps) + if apply_managed and tool in ("claude", "codex") and not is_dry_run(): + resolved_mcps = _configure_managed_mcp_servers(managed) + skills_configured = _configure_managed_skills(managed) + _restore_mcp_after_workspace_switch(previous_workspace, state) + if tool == "claude" and resolved_mcps is not None and skills_configured: + _sync_claude_cowork_plugin(state, managed_mcp_servers=resolved_mcps) + else: + _restore_mcp_after_workspace_switch(previous_workspace, state) # Only the Claude and Codex branch above reads the managed config fresh. managed_refreshed = apply_managed and tool in ("claude", "codex") install_databricks_ai_tools_for_agents([tool], state, force_refresh=not managed_refreshed) @@ -959,6 +973,12 @@ def _setup_single_agent( return state +def _restore_mcp_after_workspace_switch(previous_workspace: str | None, state: dict) -> None: + """Put the selected workspace's saved MCP entries back in global agent configs.""" + if previous_workspace and previous_workspace != state.get("workspace") and not is_dry_run(): + restore_workspace_mcp_servers(state, set(MCP_CLIENTS)) + + def _record_unmanaged_configured_agents(state: dict, tools: list[str]) -> dict: """With no managed agent list, every agent the developer configures is theirs to manage, so record it as self-managed (shown by ``ug agents list``, kept if an admin later omits it).""" @@ -1135,6 +1155,7 @@ def _configure_workspace_command( offer_provider = tool is None and selected_tools is None workspace_entries = workspaces or [_prompt_for_configuration(tool)] + previous_workspace = load_state().get("workspace") if tool is not None: _setup_single_agent( @@ -1214,6 +1235,7 @@ def _configure_workspace_command( if not is_dry_run(): resolved_mcps = _configure_managed_mcp_servers(managed) skills_configured = _configure_managed_skills(managed) + _restore_mcp_after_workspace_switch(previous_workspace, state) registered_mcps = [ str(server["name"]) for server in (resolved_mcps or []) if server.get("name") ] @@ -1275,6 +1297,7 @@ def _configure_workspace_command( if not is_dry_run(): _configure_managed_mcp_servers(None) _configure_managed_skills(None) + _restore_mcp_after_workspace_switch(previous_workspace, state) # Prefer this run's outcome; available_tools is cumulative and can still list a # tool that failed this run from an earlier success. Fall back to it only when the @@ -2044,6 +2067,7 @@ def mcp_add( ) previous_workspace = load_full_state().get("current_workspace") workspace_entries = None + added = False try: # `ug --workspace mcp add ...` is parsed by the root command, which # returns early for subcommands, so forward its value here. @@ -2065,6 +2089,16 @@ def mcp_add( workspace_entries = _parse_profile_option(profile) elif workspace is not None: workspace_entries = _parse_workspace_option(workspace) + if ( + workspace_entries + and previous_workspace + and workspace_entries[0][0] != previous_workspace + ): + raise RuntimeError( + f"MCP clients are configured for {previous_workspace}. " + f"Run `ug configure --workspace {workspace_entries[0][0]}` to switch " + "the active workspace before adding MCP servers there." + ) existing_scope = None if workspace_entries and not requested_agents: target = workspace_entries[0][0] @@ -2087,8 +2121,8 @@ def mcp_add( else existing_scope ) add_mcp_command(location=location, services=selected, agents=scope) - if not workspace_entries or workspace_entries[0][0] == previous_workspace: - _sync_claude_cowork_plugin() + added = True + _sync_claude_cowork_plugin() except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None @@ -2096,7 +2130,7 @@ def mcp_add( print_err("Interrupted.") raise typer.Exit(130) from None finally: - if workspace_entries: + if workspace_entries and (previous_workspace is not None or not added): set_current_workspace(previous_workspace) @@ -4646,6 +4680,7 @@ def configure( # MCP setup prompt so flag-driven / scripted runs are never interrupted. fully_interactive = False combined_optional_setup = False + previous_workspace = load_state().get("workspace") if agent is not None: tool = normalize_tool(agent) install_tool_binary( @@ -4673,6 +4708,7 @@ def configure( wants_cursor = "cursor" in requested wants_vscode = "vscode" in requested model_agent_names = ",".join(a for a in requested if a not in ("cursor", "vscode")) + mcp_states = None if model_agent_names: selected_tools = _parse_agents_option(model_agent_names) if workspace_entries is None: @@ -4688,7 +4724,7 @@ def configure( ) elif wants_cursor or wants_vscode: # Cursor/VS Code only: establish workspace state without the model picker. - _configure_shared_workspace_states( + mcp_states = _configure_shared_workspace_states( workspace_entries or [_prompt_for_configuration(None)], tools=[], force_login=not use_pat, @@ -4709,6 +4745,8 @@ def configure( managed, _ = _read_policy_for_other_agent(state) refuse_user_scope_over_managed_config(state.get("workspace"), managed) vscode.configure(state) + if mcp_states is not None: + _restore_mcp_after_workspace_switch(previous_workspace, mcp_states[0]) elif mcp is not None: # MCP-only: `--mcp` without --agent(s) (e.g. Cursor, which isn't a # model agent, or adding MCP servers to an already-configured setup). @@ -4716,13 +4754,14 @@ def configure( # `--mcp` registration below has a current workspace to target. if workspace_entries is None: workspace_entries = [_prompt_for_configuration(None)] - _configure_shared_workspace_states( + mcp_states = _configure_shared_workspace_states( workspace_entries, tools=[], force_login=not use_pat, use_pat=use_pat, custom_oauth=custom_oauth, ) + _restore_mcp_after_workspace_switch(previous_workspace, mcp_states[0]) else: # Tool binaries are installed after the user picks which agents # they want, in configure_workspace_command. diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 7ec742cc9..e1521e166 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -1621,6 +1621,34 @@ def apply_mcp_server_changes( return changed +def restore_workspace_mcp_servers(state: dict, agents: set[str]) -> None: + """Reapply this workspace's saved user MCP entries after an explicit workspace switch.""" + workspace = state.get("workspace") + if not workspace: + return + installed = set(available_mcp_clients()) & agents + managed_names = managed_mcp_server_names(state, installed) + servers = list(state.get("mcp_servers") or []) + for client in sorted(installed): + selected = [ + server + for server in servers + if client in _mcp_server_clients(server) + and (name := _server_name(server)) + and name not in managed_names + ] + if selected: + apply_mcp_server_changes( + [], + selected, + [client], + workspace, + state.get("profile"), + use_pat=bool(state.get("use_pat")), + batch_agents=frozenset({client}), + ) + + class _Counter: """Thread-safe monotonic counter for cross-thread progress reporting.""" @@ -1959,25 +1987,6 @@ def _union_missing(base: list[dict], selected: list[dict]) -> list[dict]: return [*selected, *extra] -def _namespace_foreign_mcp_names(servers: list[dict], workspace: str) -> list[dict]: - """Keep a second workspace from replacing an agent's same-named MCP server.""" - foreign_names = { - name - for other_workspace, bucket in (load_full_state().get("workspaces") or {}).items() - if other_workspace != workspace and isinstance(bucket, dict) - for key in ("mcp_servers", "managed_mcp_servers") - for entry in bucket.get(key) or [] - if (name := _server_name(entry)) - } - suffix = _normalize_workspace_title(workspace_hostname(workspace)) - return [ - {**server, "name": f"{server['name']}-{suffix}"} - if server.get("name") in foreign_names and not server["name"].endswith(f"-{suffix}") - else server - for server in servers - ] - - def add_mcp_command( location: str | None = None, services: set[str] | None = None, @@ -2064,7 +2073,6 @@ def _configure_v2_mcp_selectors( ) if append: - working_mcp_servers = _namespace_foreign_mcp_names(working_mcp_servers, workspace) working_mcp_servers = _union_missing(original_mcp_servers, working_mcp_servers) changed = apply_mcp_server_changes( @@ -2149,7 +2157,6 @@ def configure_mcp_command( ) working_mcp_servers = _drop_managed_servers(working_mcp_servers, managed_names) if append: - working_mcp_servers = _namespace_foreign_mcp_names(working_mcp_servers, workspace) working_mcp_servers = _union_missing( original_mcp_servers_for_location, working_mcp_servers ) @@ -2251,7 +2258,6 @@ def configure_mcp_command( working_names.add(entry_name) if append: - working_mcp_servers = _namespace_foreign_mcp_names(working_mcp_servers, workspace) working_mcp_servers = _union_missing(original_mcp_servers, working_mcp_servers) changed = apply_mcp_server_changes( diff --git a/tests/test_cli.py b/tests/test_cli.py index 768aec868..139297368 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -3360,6 +3360,23 @@ def _real_state_writes(self): def _save_workspace(workspace: str, **extra) -> None: real_save_state({"workspace": workspace, "available_tools": ["codex"], **extra}) + def test_first_mcp_add_selects_its_workspace(self): + from ucode.state import load_full_state + + with ( + patch("ucode.cli.available_mcp_clients", return_value=["codex"]), + patch("ucode.cli.add_mcp_command") as add_mcp, + patch("ucode.cli._sync_claude_cowork_plugin"), + ): + result = runner.invoke( + app, + ["mcp", "add", "--agents", "codex", "--names", "a.b.c", "--workspace", self.WS_B], + ) + + assert result.exit_code == 0, result.output + add_mcp.assert_called_once() + assert load_full_state()["current_workspace"] == self.WS_B + def test_profile_used_for_auth_when_workspace_already_ready(self): from ucode.mcp import setup_mcp_clients from ucode.state import load_state @@ -3400,7 +3417,7 @@ def test_new_profile_drops_saved_pat_mode(self): assert state["profile"] == "NEW" assert "use_pat" not in state - def test_second_workspace_add_preserves_launch_and_model_configuration(self): + def test_second_workspace_add_requires_explicit_workspace_switch(self): from ucode.state import load_full_state self._save_workspace(self.WS_A, profile="A", mcp_servers=[{"name": "existing"}]) @@ -3426,18 +3443,16 @@ def test_second_workspace_add_preserves_launch_and_model_configuration(self): ], ) - assert result.exit_code == 0, result.output + assert result.exit_code == 1, result.output + assert f"configure --workspace {self.WS_B}" in _strip_ansi(result.output) configure_models.assert_not_called() configure_workspace.assert_not_called() sync_cowork.assert_not_called() - add_mcp.assert_called_once_with( - location=None, services={"a.b.c"}, agents={"claude", "codex"} - ) + add_mcp.assert_not_called() full = load_full_state() assert full["current_workspace"] == self.WS_A assert full["workspaces"][self.WS_A] == before - assert full["workspaces"][self.WS_B]["mcp_clients"] == ["claude", "codex"] - assert full["workspaces"][self.WS_B].get("available_tools") is None + assert self.WS_B not in full["workspaces"] @pytest.mark.parametrize("use_env", [False, True]) def test_second_workspace_without_agents_uses_its_existing_mcp_clients( @@ -3459,8 +3474,9 @@ def test_second_workspace_without_agents_uses_its_existing_mcp_clients( ): result = runner.invoke(app, args) - assert result.exit_code == 0, result.output - add_mcp.assert_called_once_with(location=None, services={"a.b.c"}, agents={"codex"}) + assert result.exit_code == 1, result.output + assert f"configure --workspace {self.WS_B}" in _strip_ansi(result.output) + add_mcp.assert_not_called() configure_models.assert_not_called() assert load_full_state()["current_workspace"] == self.WS_A assert load_full_state()["workspaces"][self.WS_B]["mcp_clients"] == ["codex"] @@ -3480,7 +3496,7 @@ def test_failed_setup_with_no_previous_workspace_clears_current(self): @pytest.mark.parametrize( ("error", "exit_code"), [(RuntimeError("boom"), 1), (KeyboardInterrupt(), 130)] ) - def test_failed_setup_restores_previous_workspace(self, error, exit_code): + def test_failed_same_workspace_add_preserves_workspace(self, error, exit_code): from ucode.state import load_full_state self._save_workspace(self.WS_A) @@ -3490,14 +3506,14 @@ def test_failed_setup_restores_previous_workspace(self, error, exit_code): ): result = runner.invoke( app, - ["mcp", "add", "--agents", "codex", "--names", "a.b.c", "--workspace", self.WS_B], + ["mcp", "add", "--agents", "codex", "--names", "a.b.c", "--workspace", self.WS_A], ) assert result.exit_code == exit_code, result.output mock_add.assert_called_once() full = load_full_state() assert full["current_workspace"] == self.WS_A - assert full["workspaces"][self.WS_B]["mcp_clients"] == ["codex"] + assert self.WS_B not in full["workspaces"] class TestMcpAddRootWorkspace: @@ -5481,6 +5497,34 @@ def test_single_copilot_agent_refreshes_ai_tools_since_it_only_reads_the_policy( install_ai_tools.assert_called_once_with(["copilot"], state, force_refresh=True) + def test_single_codex_switch_reconciles_managed_mcps_before_restoring_saved_entries( + self, monkeypatch + ): + state = {**MINIMAL_STATE, "workspace": "https://b.com"} + calls = [] + monkeypatch.setattr(cli_mod, "load_state", lambda: {"workspace": "https://a.com"}) + monkeypatch.setattr(cli_mod, "_configure_shared_workspace_states", lambda *a, **k: [state]) + monkeypatch.setattr(cli_mod, "_refresh_configured_policy", lambda state: (None, None)) + monkeypatch.setattr(cli_mod, "refuse_user_scope_over_managed_config", lambda *a: None) + monkeypatch.setattr(cli_mod, "configure_single_tool", lambda *a, **k: state) + monkeypatch.setattr(cli_mod, "_record_unmanaged_configured_agents", lambda s, t: s) + monkeypatch.setattr( + cli_mod, "_configure_managed_mcp_servers", lambda managed: calls.append("managed") or [] + ) + monkeypatch.setattr( + cli_mod, "_configure_managed_skills", lambda managed: calls.append("skills") or True + ) + monkeypatch.setattr( + cli_mod, + "_restore_mcp_after_workspace_switch", + lambda previous, current: calls.append((previous, current["workspace"])), + ) + monkeypatch.setattr(cli_mod, "install_databricks_ai_tools_for_agents", lambda *a, **k: None) + + cli_mod._setup_single_agent([("https://b.com", None)], "codex", apply_managed=True) + + assert calls == ["managed", "skills", ("https://a.com", "https://b.com")] + def test_managed_config_fails_when_no_enabled_agent_is_available(self, monkeypatch): import ucode.cli as cli_mod @@ -6689,6 +6733,75 @@ 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: {}) + @pytest.mark.parametrize("same_name", [False, True]) + def test_explicit_configure_switch_restores_saved_mcp_entries(self, monkeypatch, same_name): + 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: {}) + monkeypatch.setattr(cli_mod, "save_state", mcp.save_state) + monkeypatch.setattr(cli_mod, "_refresh_configured_policy", lambda state: (None, None)) + monkeypatch.setattr(cli_mod, "refuse_user_scope_over_managed_files", lambda tools: None) + monkeypatch.setattr(cli_mod, "refuse_user_scope_over_managed_config", lambda *args: None) + monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda state, tool: True) + monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *args, **kwargs: True) + monkeypatch.setattr( + cli_mod, + "configure_selected_tools", + lambda state, tools, **kwargs: { + **state, + "available_tools": tools, + "last_configured_tools": tools, + }, + ) + monkeypatch.setattr( + cli_mod, "_record_unmanaged_configured_agents", lambda state, tools: state + ) + monkeypatch.setattr(cli_mod, "_configure_managed_mcp_servers", lambda managed: []) + monkeypatch.setattr(cli_mod, "_configure_managed_skills", lambda managed: True) + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["codex"]) + monkeypatch.setattr(mcp, "managed_mcp_server_names", lambda state, agents: set()) + monkeypatch.setattr(mcp, "_oauth_http_client", lambda *args, **kwargs: None) + monkeypatch.setattr( + mcp, + "_managed_mcp_entry", + lambda client, url, *args, **kwargs: {"url": url}, + ) + actual: dict[str, dict] = {} + + def write(add: dict, remove: set[str]) -> set[str]: + removed = set(actual) & remove + for name in remove: + actual.pop(name, None) + actual.update(add) + return removed + + monkeypatch.setattr(mcp.codex, "write_user_mcp_servers", write) + a = "https://a.databricks.com" + b = "https://b.databricks.com" + a_entry = {"name": "a-tools", "url": f"{a}/mcp", "clients": ["codex"]} + b_entry = { + "name": "a-tools" if same_name else "b-tools", + "url": f"{b}/mcp", + "clients": ["codex"], + } + mcp.save_state({"workspace": b, "mcp_servers": [b_entry]}) + mcp.save_state({"workspace": a, "mcp_servers": [a_entry]}) + actual["a-tools"] = {"url": a_entry["url"]} + + cli_mod._configure_workspace_command(selected_tools=["codex"], workspaces=[(b, None)]) + assert actual == {b_entry["name"]: {"url": b_entry["url"]}} + assert state_mod.load_full_state()["current_workspace"] == b + + cli_mod._configure_workspace_command(selected_tools=["codex"], workspaces=[(a, None)]) + assert actual == {"a-tools": {"url": a_entry["url"]}} + full = state_mod.load_full_state() + assert full["current_workspace"] == a + assert full["workspaces"][a]["mcp_servers"] == [a_entry] + assert full["workspaces"][b]["mcp_servers"] == [b_entry] + def test_workspace_switch_continues_after_skills_cleanup_timeout(self, monkeypatch, capsys): from ucode import mcp from ucode import state as state_mod @@ -6719,18 +6832,18 @@ def time_out(add, remove): state = cli_mod.configure_shared_state(new_workspace, force_login=True) assert state["workspace"] == new_workspace - assert state["mcp_servers"] == [] + assert state.get("mcp_servers", []) == [] assert "_discovery_reasons" in state assert calls == [({}, {"databricks-skill-registry"})] full = state_mod.load_full_state() assert full["current_workspace"] == new_workspace - assert full["workspaces"][new_workspace]["mcp_servers"] == [] + assert full["workspaces"][new_workspace].get("mcp_servers", []) == [] # The previous workspace's own bucket is always preserved so switching back still recognizes # its configured 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 "Dropping 1 stale MCP entry" not in output assert "Failed to remove stale MCP entries from Claude Code" in output assert "left over from previously-configured workspaces" not in output diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 7a023bcb4..bf81c535e 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -38,18 +38,39 @@ def test_mcp_only_client_is_available_without_model_configuration(): assert mcp.configured_mcp_clients(state, ["claude", "codex"]) == ["claude"] -def test_second_workspace_mcp_name_does_not_replace_first(monkeypatch): - other = "https://first.cloud.databricks.com" - current = "https://second.cloud.databricks.com" +def test_restore_workspace_mcp_servers_uses_saved_name_and_destination_url(monkeypatch): + workspace = "https://second.cloud.databricks.com" + selected = { + "name": "system-ai-github", + "url": f"{workspace}/ai-gateway/mcp-services/system.ai.github", + "clients": ["claude"], + } + writes = [] + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude"]) + monkeypatch.setattr(mcp, "managed_mcp_server_names", lambda state, agents: set()) monkeypatch.setattr( mcp, - "load_full_state", - lambda: {"workspaces": {other: {"mcp_servers": [{"name": "system-ai-github"}]}}}, + "apply_mcp_server_changes", + lambda original, working, clients, ws, profile, **kwargs: writes.append( + (original, working, clients, ws, profile, kwargs) + ), ) - selected = [{"name": "system-ai-github", "url": f"{current}/mcp"}] - result = mcp._namespace_foreign_mcp_names(selected, current) - assert result[0]["name"] == "system-ai-github-second-cloud-databricks-com" - assert selected[0]["name"] == "system-ai-github" + + mcp.restore_workspace_mcp_servers( + {"workspace": workspace, "profile": "B", "mcp_servers": [selected]}, + {"claude", "codex"}, + ) + + assert writes == [ + ( + [], + [selected], + ["claude"], + workspace, + "B", + {"use_pat": False, "batch_agents": frozenset({"claude"})}, + ) + ] # The proxy argv every client registers as a stdio command. The leading element @@ -1913,25 +1934,10 @@ def test_existing_entry_gets_reconfigured_for_newly_added_clients(self, monkeypa class TestAddMcpCommand: """`ucode mcp add` (append) registers new servers without removing existing ones.""" - def test_second_workspace_keeps_first_workspace_server(self, monkeypatch): - other = "https://first.cloud.databricks.com" + def test_add_uses_original_server_name(self, monkeypatch): saved: list[dict] = [] registered: list[str] = [] _stub_location_base(monkeypatch, {**CLAUDE_STATE, "mcp_servers": []}) - monkeypatch.setattr( - mcp, - "load_full_state", - lambda: { - "workspaces": { - other: {"mcp_servers": [{"name": "system-ai-github", "clients": ["claude"]}]} - } - }, - ) - monkeypatch.setattr( - mcp, - "purge_cross_workspace_mcp_residue", - lambda *a: pytest.fail("purged other workspace"), - ) monkeypatch.setattr(mcp, "list_mcp_services", lambda *a, **kw: (["system.ai.github"], None)) monkeypatch.setattr( mcp, @@ -1941,7 +1947,7 @@ def test_second_workspace_keeps_first_workspace_server(self, monkeypatch): monkeypatch.setattr(mcp, "save_state", lambda state: saved.append(state.copy())) assert mcp.add_mcp_command(location="system.ai", agents={"claude"}) == 0 - expected = "system-ai-github-example-databricks-com" + expected = "system-ai-github" assert registered == [expected] assert saved[-1]["mcp_servers"][0]["name"] == expected From 4158a054d276947714d05cc9e0e960d4a5422203 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Fri, 30 Oct 2026 03:57:52 +0000 Subject: [PATCH 4/6] [BR-2167] Preserve managed MCP cleanup across workspace switches --- src/ucode/cli.py | 4 ++- tests/test_agent_vscode.py | 2 +- tests/test_cli.py | 51 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 55 insertions(+), 2 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 3b9d9d03c..bdd4c8831 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -571,7 +571,9 @@ def configure_shared_state( # Model discovery below refreshes the selected workspace. MCP selections are different: # they are user choices, so carry the destination's saved registry across a switch. destination = load_full_state().get("workspaces", {}).get(workspace, {}) - for key in ("mcp_servers", "mcp_clients", "managed_mcp_servers"): + # managed_mcp_servers records what the source workspace delivered to user scope. + # Keep it until reconciliation removes registrations the destination no longer wants. + for key in ("mcp_servers", "mcp_clients"): if key in destination: state[key] = destination[key] else: diff --git a/tests/test_agent_vscode.py b/tests/test_agent_vscode.py index 48c55402a..9b996d827 100644 --- a/tests/test_agent_vscode.py +++ b/tests/test_agent_vscode.py @@ -341,7 +341,7 @@ def test_configure_agents_vscode_sets_up_the_workspace_then_vscode(monkeypatch, shared: list[dict] = [] monkeypatch.setattr( "ucode.cli._configure_shared_workspace_states", - lambda *args, **kwargs: shared.append(kwargs) or [], + lambda *args, **kwargs: shared.append(kwargs) or [STATE], ) monkeypatch.setattr( "ucode.cli.configure_workspace_command", diff --git a/tests/test_cli.py b/tests/test_cli.py index 139297368..14b6f75ce 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -6733,6 +6733,57 @@ 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_switch_reconciles_source_managed_fallback_without_losing_destination_choices( + self, monkeypatch + ): + 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: {}) + monkeypatch.setattr(cli_mod, "save_state", mcp.save_state) + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude"]) + monkeypatch.setattr(mcp, "_MANAGED_FILE_AGENTS", {}) + monkeypatch.setattr( + mcp, + "setup_mcp_clients", + lambda state, *args, **kwargs: (state["workspace"], None, ["claude"]), + ) + changes = [] + monkeypatch.setattr( + mcp, + "apply_mcp_server_changes", + lambda original, working, clients, *args, **kwargs: changes.append( + (original, working, clients) + ), + ) + source = "https://source.databricks.com" + destination = "https://destination.databricks.com" + managed = { + "name": "managed-old", + "url": f"{source}/ai-gateway/mcp-services/system.ai.old", + "clients": ["claude"], + } + selected = { + "name": "user-choice", + "url": f"{destination}/ai-gateway/mcp-services/system.ai.choice", + "clients": ["claude"], + } + mcp.save_state({"workspace": destination, "mcp_servers": [selected]}) + mcp.save_state( + {"workspace": source, "available_tools": ["claude"], "managed_mcp_servers": [managed]} + ) + monkeypatch.setattr(cli_mod, "purge_cross_workspace_mcp_residue", lambda *args: None) + + switched = cli_mod.configure_shared_state(destination, force_login=True) + assert switched["mcp_servers"] == [selected] + assert switched["managed_mcp_servers"] == [managed] + + mcp.reconcile_managed_mcp_servers({}, {"claude"}) + assert changes == [([managed], [], ["claude"])] + assert mcp.load_state().get("managed_mcp_servers") == [] + @pytest.mark.parametrize("same_name", [False, True]) def test_explicit_configure_switch_restores_saved_mcp_entries(self, monkeypatch, same_name): from ucode import mcp From 73ac9d0bd5d149cad2b0f46d1d7369059a23d0a1 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Fri, 30 Oct 2026 04:09:51 +0000 Subject: [PATCH 5/6] [BR-2167] Align workspace-switch CUJ with scoped MCP cleanup --- tests/README.md | 2 +- tests/integration/README.md | 2 ++ .../test_ug_configure_claude_workspace_switch.py | 9 +++++---- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/tests/README.md b/tests/README.md index a78179707..245c564d5 100644 --- a/tests/README.md +++ b/tests/README.md @@ -697,7 +697,7 @@ integration utilities; only CUJ-specific evidence correlation stays in a test fi | `test_smart_routing_claude_route_subagent_hook`, `test_smart_routing_codex_route_subagent_hook` | Pipe a native spawn payload to the installed Claude mod command or Codex PreToolUse hook with legacy subagent routing or each supported routing preset | Claude returns model/rationale/updatedInput JSON naming the routed agent the mod labels; Codex returns an allow-hook decision; 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_router_routes_{claude,codex}_subagent_at_launch`, `test_smart_router_skill_turns_{claude,codex}_subagent_routing_off`, `test_smart_router_skill_turns_{claude,codex}_subagent_routing_back_on` | Configure, launch a real subagent-only TUI with legacy routing-only/orchestration flags or the three subagent-only presets, each in its own fresh session: spawn a tagged child at launch; invoke the installed Smart Router skill to turn routing off then spawn a child; or spawn a routed child, turn routing off and spawn an unrouted child, then turn routing back on and spawn a routed child | Routing-only cases install `smart-router`; orchestration cases also install `smart-router-orchestrator`; each tagged native child completes; routed children (launch, before off and after back on) show a native model/rationale row (Claude) or subagent banner (Codex) and produce a live routing decision correlated with the child, while each child after off shows neither and has no decision; Claude's recorded task inference payloads carry `smart_router_recipe_name` as the launch recipe (`at_launch`, `back_on`, where the second child must match the launch child) or DISABLED (after off), omit the recipe header, and receive successful gateway responses; opted-in Codex banners include `[orchestrator on]`; no first-prompt routing wrapper; the customer full-mode preset is intentionally outside this journey; 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_cleans_stale_skills_mcp_on_workspace_switch` | Configure the first workspace, register its skills MCP, switch to a second real workspace, and use Claude | One cross-workspace removal reported; 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 on the managed workspace with the stubbed `managed_workspace_default.json` CodingAgentConfig | No agent selector; each agent's generated config exposes exactly the admin's static model_services; real gateway prompt on launch. The Codex case also checks the shared catalog pointer, restart guidance, and a fresh bare app-server's visible model list | | `e2e_cuj/test_cuj5_budget_defaults.py` | Launch bare `ug` with separate low-spend and above-tier principals against the fixed 1% tier (`spending_percentage=0.01`); explicitly launch `ug claude`, and `ug claude --model`, above the tier | Bare launches check Claude/Sonnet below the tier and the Codex/Luna recommendation above it; `ug usage` agrees with backend spend, threshold, and percentage. Explicit `ug claude` displays the backend's Codex/Luna recommendation while its generated model setting and native header select Sonnet. `ug claude --model ` (the workspace's only Claude model) displays the same recommendation, selects Sonnet in the native header, and leaves the user and OS-managed settings files byte-identical to configure's (Sonnet `ANTHROPIC_MODEL`); a following plain launch selects Sonnet without rewriting them. A differing Claude budget pick, alone or with smart routing, is unit-covered only. No budget writes or inference tasks. Runs on its own `E2E CUJs · CUJ 5 · Budget-driven defaults` runner with both credential pairs documented in `integration/README.md`. | diff --git a/tests/integration/README.md b/tests/integration/README.md index e249c9e07..116ae99b2 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -937,6 +937,8 @@ 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. +The switch reports one removal from previously configured workspaces; the old +"Dropping stale entry" warning applied when the target inherited the source's MCP selection. Run this case with both agents installed if either workspace's managed config enables both. Supply the second workspace and its bearer explicitly: diff --git a/tests/integration/test_ug_configure_claude_workspace_switch.py b/tests/integration/test_ug_configure_claude_workspace_switch.py index 95289b81a..811797468 100644 --- a/tests/integration/test_ug_configure_claude_workspace_switch.py +++ b/tests/integration/test_ug_configure_claude_workspace_switch.py @@ -29,8 +29,8 @@ def test_ug_configure_claude_cleans_stale_skills_mcp_on_workspace_switch( """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 + Expected: the switch reports one removal, 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. """ @@ -72,7 +72,8 @@ def test_ug_configure_claude_cleans_stale_skills_mcp_on_workspace_switch( "--disable-databricks-ai-tools", timeout=240, ) - assert switched.stdout.count("Dropping 1 stale MCP entry") == 1, switched.stdout + removal = "Removed 1 MCP entry left over from previously-configured workspaces" + assert switched.stdout.count(removal) == 1, switched.stdout assert SKILLS_SERVER in switched.stdout, switched.stdout missing = session.run( "mcp", "get", SKILLS_SERVER, binary=session.which("claude"), timeout=60, ok=False @@ -98,7 +99,7 @@ def test_ug_configure_claude_cleans_stale_skills_mcp_on_workspace_switch( "--disable-databricks-ai-tools", timeout=240, ) - assert "Dropping 1 stale MCP entry" not in repeated.stdout, repeated.stdout + assert removal not in repeated.stdout, repeated.stdout still_missing = session.run( "mcp", "get", SKILLS_SERVER, binary=session.which("claude"), timeout=60, ok=False ) From f0866c14d2149a6723b13770073c7af47de255b9 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Fri, 30 Oct 2026 17:38:16 +0000 Subject: [PATCH 6/6] [BR-2167] Reconcile MCP entries per client on workspace switches --- src/ucode/cli.py | 6 ++++-- src/ucode/mcp.py | 14 +++++++------- tests/test_cli.py | 16 ++++++++++++++-- tests/test_mcp.py | 20 ++++++++++++++++++++ 4 files changed, 45 insertions(+), 11 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index bdd4c8831..f27bf7e81 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -977,8 +977,10 @@ def _setup_single_agent( def _restore_mcp_after_workspace_switch(previous_workspace: str | None, state: dict) -> None: """Put the selected workspace's saved MCP entries back in global agent configs.""" - if previous_workspace and previous_workspace != state.get("workspace") and not is_dry_run(): - restore_workspace_mcp_servers(state, set(MCP_CLIENTS)) + if previous_workspace != state.get("workspace") and not is_dry_run(): + # Managed reconciliation updates the persisted state, not the configure snapshot. + # Its freshly delivered names must take precedence over saved user selections. + restore_workspace_mcp_servers(load_state(), set(MCP_CLIENTS)) def _record_unmanaged_configured_agents(state: dict, tools: list[str]) -> dict: diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index e1521e166..8f5cf19f3 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -623,13 +623,13 @@ def _partition_mcp_entries_by_workspace( 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``.""" + """Return clients with a user MCP entry only in another workspace, grouped by name.""" full_state = load_full_state() workspaces = full_state.get("workspaces") if not isinstance(workspaces, dict): return {} - current_names: set[str] = set() + current_pairs: set[tuple[str, str]] = set() current_bucket = workspaces.get(current_workspace) if isinstance(current_bucket, dict): # Include the current workspace's managed servers too — the `system.ai` default lives in @@ -641,7 +641,7 @@ def _mcp_entries_only_in_other_workspaces(current_workspace: str) -> dict[str, s for entry in current_bucket.get(key) or []: name = _server_name(entry) if name: - current_names.add(name) + current_pairs.update((name, client) for client in _mcp_server_clients(entry)) external_entries: dict[str, set[str]] = {} for ws, bucket in workspaces.items(): @@ -649,11 +649,11 @@ def _mcp_entries_only_in_other_workspaces(current_workspace: str) -> dict[str, s continue for entry in bucket.get("mcp_servers") or []: name = _server_name(entry) - if not name or name in current_names: + if not name: continue - client_set = external_entries.setdefault(name, set()) - for client in entry.get("clients") or []: - client_set.add(client) + for client in _mcp_server_clients(entry): + if (name, client) not in current_pairs: + external_entries.setdefault(name, set()).add(client) return external_entries diff --git a/tests/test_cli.py b/tests/test_cli.py index 14b6f75ce..c496eae72 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -6761,12 +6761,12 @@ def test_switch_reconciles_source_managed_fallback_without_losing_destination_ch source = "https://source.databricks.com" destination = "https://destination.databricks.com" managed = { - "name": "managed-old", + "name": "shared-name", "url": f"{source}/ai-gateway/mcp-services/system.ai.old", "clients": ["claude"], } selected = { - "name": "user-choice", + "name": "shared-name", "url": f"{destination}/ai-gateway/mcp-services/system.ai.choice", "clients": ["claude"], } @@ -6784,6 +6784,11 @@ def test_switch_reconciles_source_managed_fallback_without_losing_destination_ch assert changes == [([managed], [], ["claude"])] assert mcp.load_state().get("managed_mcp_servers") == [] + # The original configure snapshot still has A's managed server. Restore must use + # reconciliation's persisted result so B's user choice can reuse that name. + cli_mod._restore_mcp_after_workspace_switch(source, switched) + assert changes == [([managed], [], ["claude"]), ([], [selected], ["claude"])] + @pytest.mark.parametrize("same_name", [False, True]) def test_explicit_configure_switch_restores_saved_mcp_entries(self, monkeypatch, same_name): from ucode import mcp @@ -6853,6 +6858,13 @@ def write(add: dict, remove: set[str]) -> set[str]: assert full["workspaces"][a]["mcp_servers"] == [a_entry] assert full["workspaces"][b]["mcp_servers"] == [b_entry] + # Revert clears only A's active bucket. B's saved choice should return when + # configure selects B again, even though there is no current workspace. + state_mod.clear_state() + actual.clear() + cli_mod._configure_workspace_command(selected_tools=["codex"], workspaces=[(b, None)]) + assert actual == {b_entry["name"]: {"url": b_entry["url"]}} + def test_workspace_switch_continues_after_skills_cleanup_timeout(self, monkeypatch, capsys): from ucode import mcp from ucode import state as state_mod diff --git a/tests/test_mcp.py b/tests/test_mcp.py index bf81c535e..dc1035008 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -3318,6 +3318,26 @@ def test_skips_names_the_current_workspace_manages(self, monkeypatch, capsys): capsys.readouterr().out ) + def test_same_name_in_destination_only_protects_its_selected_client(self, monkeypatch): + 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: {}) + other = "https://other.databricks.com" + old = {"name": "shared", "url": f"{other}/mcp", "clients": ["claude", "codex"]} + selected = {"name": "shared", "url": f"{WS}/mcp", "clients": ["codex"]} + state_mod.save_state({"workspace": other, "mcp_servers": [old]}) + state = {"workspace": WS, "mcp_servers": [selected]} + state_mod.save_state(state) + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude", "codex"]) + writes = self._capture_batched_writes(monkeypatch, ["claude", "codex"]) + + mcp.purge_cross_workspace_mcp_residue(state, WS) + + assert writes["claude"] == [({}, {"shared"})] + assert writes["codex"] == [] + assert state_mod.load_state()["mcp_servers"] == [selected] + @pytest.mark.parametrize( "failure", [RuntimeError("config is locked"), OSError("no such file"), None],