[BR-2167] Switch MCP registrations with the active workspace

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 <url>` 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:
 

src/ucode/mcp.py

@@ -262,7 +262,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
@@ -628,13 +628,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
@@ -646,19 +646,19 @@ 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():
         if ws == current_workspace or not isinstance(bucket, dict) or ALIAS_OF_KEY in bucket:
             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
 
 
@@ -1626,6 +1626,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."""
 
@@ -1889,6 +1917,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.
 
@@ -1911,7 +1940,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:
@@ -2005,7 +2035,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
@@ -2116,7 +2149,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))

tests/README.md

@@ -702,7 +702,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 <Sonnet id>` (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`. |

tests/integration/README.md

@@ -944,6 +944,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:

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
     )

tests/test_agent_vscode.py

@@ -478,7 +478,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",