From a20ffb0df0bfefaef55883e4df4a10e9d5b43bb6 Mon Sep 17 00:00:00 2001 From: Tien Le Date: Fri, 7 Aug 2026 21:53:18 +0000 Subject: [PATCH] setup: author MCP and skills into the manifest, not local state MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ucode setup`'s MCP and skills sections delegated to `configure_mcp_command` / `configure_skills_mcp_command`, which are built for `ucode configure` — they pre-check the picker from the admin's *own* registered servers, mutate the admin's machine (`claude mcp add-json`, `~/.claude.json`, `save_state`), then read the result back out of `state.json`. So authoring a workspace-wide config both showed the wrong pre-selections (the reported bug) and silently reconfigured the admin's local agents. The picker's interaction half is extracted into `mcp.pick_mcp_servers` — source picker, discovery, checkbox, resolve each pick to `{name, url}`. Persistence stays with each caller: `configure_mcp_command` keeps its diff-apply-save against local state (behaviour unchanged, 99 mcp tests green), while the new `_author_mcp_servers` maps picks to the manifest's `{name, type}` and writes nothing to the machine. It starts the picker empty — an admin's own servers are irrelevant to a workspace declaration, and an empty start also keeps every pick resolvable to a type (a kept row would carry a state-only URL the manifest omits). Skills likewise now author `catalog.schema` names straight into the manifest instead of registering a live skills MCP connection. Scope: author side only. The pull side (a developer's ucode rebuilding these into agent configs) is still shown as "pending" and unimplemented — a separate change. Tests: +TestAuthorMcpServers (starts empty, no machine/state mutation, cancel); TestMcpServersFromState reworked to TestMcpEntriesToManifest (now takes a list). Mutation-verified: seeding the picker from local state fails the empty-start test. Co-authored-by: Isaac --- src/ucode/managed_wizard.py | 50 ++++++----- src/ucode/mcp.py | 155 ++++++++++++++++++++--------------- tests/test_managed_wizard.py | 117 ++++++++++++++++---------- 3 files changed, 193 insertions(+), 129 deletions(-) diff --git a/src/ucode/managed_wizard.py b/src/ucode/managed_wizard.py index d1d40c11..03310bb5 100644 --- a/src/ucode/managed_wizard.py +++ b/src/ucode/managed_wizard.py @@ -120,16 +120,19 @@ def _mcp_type_for_url(url: str) -> str | None: return None -def _mcp_servers_from_state(state: dict) -> list[dict]: - """The registered MCP servers, as managed-config ``{name, type}`` entries. - - Skips the skills registry connection: skills are published under the manifest's own ``skills`` - field, so including its MCP entry would configure it twice. +def _mcp_entries_to_manifest(entries: list[dict]) -> list[dict]: + """Map picker/state ``{name, url, ...}`` entries to managed-config ``{name, type}``. + + The manifest stores only ``name`` + ``type``; the URL is a local-machine concept (it is what + ``ucode configure --mcp`` writes into agent configs, and what a developer's ucode rebuilds from + ``{name, type}`` at pull time). Here the URL is derived-and-dropped: it exists only to recover + the type. Skips the skills registry connection, which is published under the manifest's own + ``skills`` field, and any entry whose URL shape isn't recognized. """ from ucode.mcp import SKILLS_MCP_KIND servers: list[dict] = [] - for entry in state.get("mcp_servers") or []: + for entry in entries: if not isinstance(entry, dict) or entry.get("kind") == SKILLS_MCP_KIND: continue name = entry.get("name") @@ -144,11 +147,21 @@ def _mcp_servers_from_state(state: dict) -> list[dict]: return servers -def _skill_names_from_state(state: dict) -> list[str]: - """Skill schemas registered on the skills MCP connection (``catalog.schema`` entries).""" - from ucode.mcp import _skill_mcp_locations +def _author_mcp_servers(workspace: str, profile: str | None) -> list[dict]: + """Pick MCP servers for the managed manifest, as ``{name, type}`` — no machine changes. + + Runs the shared picker starting from an empty selection (so it does not pre-check the admin's + own registered servers — those are irrelevant to a workspace-wide declaration) and maps the + result to the manifest shape. Unlike ``configure_mcp_command``, nothing is written to agent + config files or ``state.json``: authoring a managed config must not reconfigure the admin's own + machine. + """ + from ucode.mcp import pick_mcp_servers - return [name for name in _skill_mcp_locations(state) if isinstance(name, str) and name] + picked = pick_mcp_servers(workspace, profile) + if not picked: + return [] + return _mcp_entries_to_manifest(picked) def provider_service_model_options(service: dict) -> list[str]: @@ -893,13 +906,12 @@ def setup_command(from_file: str | None = None) -> int: print_section("MCP servers") if prompt_yes_no_default("Set up managed MCP servers for this workspace?", default=False): - from ucode.mcp import configure_mcp_command - - configure_mcp_command() - mcp_servers = _mcp_servers_from_state(load_state()) + mcp_servers = _author_mcp_servers(workspace, profile) if mcp_servers: manifest["mcp_servers"] = mcp_servers print_success(f"{len(mcp_servers)} MCP server(s) added to the managed config") + else: + print_note("No MCP servers selected.") print_section("Skills") if prompt_yes_no_default("Set up managed skills for this workspace?", default=False): @@ -907,14 +919,10 @@ def setup_command(from_file: str | None = None) -> int: "Skill schemas to publish, comma-separated `catalog.schema` (blank to skip)", default="", ) - parsed: list[str] = [item.strip() for item in (locations or "").split(",") if item.strip()] + parsed = [item.strip() for item in (locations or "").split(",") if item.strip()] if parsed: - from ucode.mcp import configure_skills_mcp_command - - configure_skills_mcp_command(parsed) - skill_names = _skill_names_from_state(load_state()) or parsed - manifest["skills"] = {"names": skill_names} - print_success(f"{len(skill_names)} skill schema(s) added to the managed config") + manifest["skills"] = {"names": parsed} + print_success(f"{len(parsed)} skill schema(s) added to the managed config") budget_policy = _prompt_budget_policy(workspace, token, enabled_agents, state) if budget_policy: diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index fb11bd4b..a73ea268 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -1546,6 +1546,88 @@ def setup_mcp_clients(state: dict, section: str) -> tuple[str, str | None, list[ return workspace, profile, clients +def pick_mcp_servers( + workspace: str, + profile: str | None, + *, + preselected: list[dict] | None = None, + clients: list[str] | None = None, +) -> list[dict] | None: + """Run the two-step MCP picker and return the chosen servers as resolved entries. + + The interaction only: source picker -> discovery -> server checkbox -> resolve each pick to a + ``{name, url}`` entry. Persistence is the caller's job — ``configure_mcp_command`` diffs the + result against local state and applies it to the machine, while ``ucode setup`` maps it to the + managed manifest's ``{name, type}`` and writes nothing to the machine. + + Returns the working list (each entry ``{name, url, auth, clients}``), or None if the user + cancelled at either step. ``preselected`` seeds which rows start checked; the manifest-authoring + caller passes ``[]`` (or nothing) so the picker starts empty and every pick resolves through the + add-path — a kept row would carry a state-only URL the manifest can't supply. + """ + preselected_servers = list(preselected or []) + original_by_name = _servers_by_name(preselected_servers) + + # Two-step wizard: (1) choose which sources to search, (2) pick servers from the results. + # Pressing Left (←) in the picker returns to step 1, so the user can revise their source + # selection without restarting. + while True: + sources = prompt_for_mcp_search_sources() + if sources is None: + return None + discovered = _discover_selected_mcp_sources(workspace, profile, sources) + + selections = prompt_for_mcp_server_choices( + discovered["external"], + discovered["genie"], + discovered["apps"], + preselected_servers, + discovered["services"], + discovered["vector_search"], + discovered["uc_functions"], + allow_back=True, + ) + if selections is None: + return None + if isinstance(selections, _Back): + continue + break + + working_mcp_servers: list[dict] = [] + working_names: set[str] = set() + add_selections: list[str] = [] + for selection in selections: + if selection.startswith(MCP_ADD_PREFIX): + add_selections.append(selection.removeprefix(MCP_ADD_PREFIX)) + continue + original = original_by_name.get(selection) + if original and selection not in working_names: + working_mcp_servers.append(original.copy()) + working_names.add(selection) + + for selection in add_selections: + try: + entry_name, url = _resolve_mcp_selection( + selection, + workspace, + discovered["apps"], + discovered["genie"], + discovered["vector_search"], + discovered["uc_functions"], + ) + except RuntimeError as exc: + print_warning(f"Skipped MCP selection `{selection}`: {exc}.") + continue + if entry_name in working_names: + continue + working_mcp_servers.append( + {"name": entry_name, "url": url, "auth": "proxy", "clients": clients or []} + ) + working_names.add(entry_name) + + return working_mcp_servers + + def configure_mcp_command(location: str | None = None, services: set[str] | None = None) -> int: if services is not None and location is None: # `--services` works standalone with full names (`system.ai.github`): the @@ -1593,72 +1675,13 @@ def configure_mcp_command(location: str | None = None, services: set[str] | None picker_servers = [s for s in original_mcp_servers if s.get("kind") != SKILLS_MCP_KIND] original_by_name = _servers_by_name(picker_servers) - # Two-step wizard: (1) choose which sources to search, (2) pick servers from - # the results. Pressing Left (←) in the picker returns to step 1, so the user - # can revise their source selection without restarting the command. - while True: - sources = prompt_for_mcp_search_sources() - if sources is None: - return 0 - discovered = _discover_selected_mcp_sources(workspace, profile, sources) - - selections = prompt_for_mcp_server_choices( - discovered["external"], - discovered["genie"], - discovered["apps"], - picker_servers, - discovered["services"], - discovered["vector_search"], - discovered["uc_functions"], - allow_back=True, - ) - if selections is None: - return 0 - if isinstance(selections, _Back): - continue - break - - available_app_mcp_servers = discovered["apps"] - available_genie_mcp_servers = discovered["genie"] - available_vector_search_servers = discovered["vector_search"] - available_uc_functions_servers = discovered["uc_functions"] - - working_mcp_servers: list[dict] = list(skills_servers) - working_names: set[str] = set() - add_selections: list[str] = [] - for selection in selections: - if selection.startswith(MCP_ADD_PREFIX): - add_selections.append(selection.removeprefix(MCP_ADD_PREFIX)) - continue - original = original_by_name.get(selection) - if original and selection not in working_names: - working_mcp_servers.append(original.copy()) - working_names.add(selection) + picked = pick_mcp_servers(workspace, profile, preselected=picker_servers, clients=clients) + if picked is None: + return 0 - for selection in add_selections: - try: - entry_name, url = _resolve_mcp_selection( - selection, - workspace, - available_app_mcp_servers, - available_genie_mcp_servers, - available_vector_search_servers, - available_uc_functions_servers, - ) - except RuntimeError as exc: - print_warning(f"Skipped MCP selection `{selection}`: {exc}.") - continue - if entry_name in working_names: - continue - working_mcp_servers.append( - { - "name": entry_name, - "url": url, - "auth": "proxy", - "clients": clients, - } - ) - working_names.add(entry_name) + # Carry the skills connections through untouched; the picker never sees them. + working_mcp_servers: list[dict] = list(skills_servers) + picked + working_names = {s["name"] for s in picked} changed = apply_mcp_server_changes( original_mcp_servers, @@ -1674,7 +1697,7 @@ def configure_mcp_command(location: str | None = None, services: set[str] | None added = sorted(working_names - set(original_by_name)) removed = sorted(set(original_by_name) - working_names) print_success(_mcp_change_summary(added, removed, clients)) - elif not selections and not original_mcp_servers: + elif not picked and not original_mcp_servers: # User submitted the picker without toggling anything --> make it clear nothing was selected print_note("No MCP servers selected. Press space to toggle an item, then enter to save.") return 0 diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index f5e92577..d240efd3 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -100,18 +100,19 @@ def test_sql_is_not_confused_for_app(self): assert wizard._mcp_type_for_url("https://ws.example.com/api/2.0/mcp/sql") == "sql" -class TestMcpServersFromState: - def test_maps_registered_servers_to_name_and_type(self): - state = { - "mcp_servers": [ - { - "name": "databricks-github", - "url": f"{WORKSPACE}/ai-gateway/mcp-services/system.ai.github", - }, - {"name": "databricks-sql", "url": f"{WORKSPACE}/api/2.0/mcp/sql"}, - ] - } - assert wizard._mcp_servers_from_state(state) == [ +class TestMcpEntriesToManifest: + """The picker returns state-shape `{name, url}` entries; the manifest stores `{name, type}`. + The URL is derived to a type and dropped — it is a local-machine concept the manifest omits.""" + + def test_maps_entries_to_name_and_type(self): + entries = [ + { + "name": "databricks-github", + "url": f"{WORKSPACE}/ai-gateway/mcp-services/system.ai.github", + }, + {"name": "databricks-sql", "url": f"{WORKSPACE}/api/2.0/mcp/sql"}, + ] + assert wizard._mcp_entries_to_manifest(entries) == [ {"name": "databricks-github", "type": "mcp-service"}, {"name": "databricks-sql", "type": "sql"}, ] @@ -121,45 +122,77 @@ def test_skips_the_skills_registry_entry(self): # would configure them twice. from ucode.mcp import SKILLS_MCP_KIND - state = { - "mcp_servers": [ - { - "name": "databricks-skill-registry", - "kind": SKILLS_MCP_KIND, - "url": f"{WORKSPACE}/api/2.0/mcp/sql", - }, - {"name": "databricks-sql", "url": f"{WORKSPACE}/api/2.0/mcp/sql"}, - ] - } - assert wizard._mcp_servers_from_state(state) == [{"name": "databricks-sql", "type": "sql"}] + entries = [ + { + "name": "databricks-skill-registry", + "kind": SKILLS_MCP_KIND, + "url": f"{WORKSPACE}/api/2.0/mcp/sql", + }, + {"name": "databricks-sql", "url": f"{WORKSPACE}/api/2.0/mcp/sql"}, + ] + assert wizard._mcp_entries_to_manifest(entries) == [ + {"name": "databricks-sql", "type": "sql"} + ] - def test_skips_unclassifiable_servers(self): - state = {"mcp_servers": [{"name": "mystery", "url": "https://example.com/nope"}]} - assert wizard._mcp_servers_from_state(state) == [] + def test_skips_unclassifiable_entries(self): + assert wizard._mcp_entries_to_manifest([{"name": "mystery", "url": "https://x/nope"}]) == [] def test_skips_entries_missing_name_or_url(self): - state = { - "mcp_servers": [ - {"url": f"{WORKSPACE}/api/2.0/mcp/sql"}, - {"name": "no-url"}, - "not-a-dict", - ] - } - assert wizard._mcp_servers_from_state(state) == [] + entries = [ + {"url": f"{WORKSPACE}/api/2.0/mcp/sql"}, + {"name": "no-url"}, + "not-a-dict", + ] + assert wizard._mcp_entries_to_manifest(entries) == [] - def test_empty_state_yields_nothing(self): - assert wizard._mcp_servers_from_state({}) == [] + def test_empty_yields_nothing(self): + assert wizard._mcp_entries_to_manifest([]) == [] def test_output_validates_as_a_manifest(self): - state = { - "mcp_servers": [ - {"name": "databricks-sql", "url": f"{WORKSPACE}/api/2.0/mcp/sql"}, - ] - } - servers = wizard._mcp_servers_from_state(state) + entries = [{"name": "databricks-sql", "url": f"{WORKSPACE}/api/2.0/mcp/sql"}] + servers = wizard._mcp_entries_to_manifest(entries) assert validate_manifest({"mcp_servers": servers}) == [] +class TestAuthorMcpServers: + """`ucode setup` authors MCP into the manifest via the shared picker, touching no local state.""" + + def test_starts_the_picker_empty_and_returns_manifest_shape(self): + # The bug: the picker pre-checked the admin's own registered servers. Authoring a workspace + # declaration must start empty, so `preselected` is not passed (defaults to none). + seen = {} + + def fake_pick(workspace, profile, **kwargs): + seen["kwargs"] = kwargs + return [{"name": "databricks-sql", "url": f"{WORKSPACE}/api/2.0/mcp/sql"}] + + with patch("ucode.mcp.pick_mcp_servers", side_effect=fake_pick): + servers = wizard._author_mcp_servers(WORKSPACE, "profile") + assert servers == [{"name": "databricks-sql", "type": "sql"}] + # `preselected` must not be passed at all — not merely passed empty. Seeding it from local + # state (the bug) would pass it, so asserting the key is absent catches that regression even + # when the caller's own state happens to hold no servers. + assert "preselected" not in seen["kwargs"] + + def test_does_not_touch_local_state_or_the_machine(self): + # Authoring a managed config must not register servers on the admin's machine or save state. + with ( + patch( + "ucode.mcp.pick_mcp_servers", + return_value=[{"name": "databricks-sql", "url": f"{WORKSPACE}/api/2.0/mcp/sql"}], + ), + patch("ucode.mcp.apply_mcp_server_changes") as apply, + patch("ucode.mcp.save_state") as save, + ): + wizard._author_mcp_servers(WORKSPACE, "profile") + assert not apply.called + assert not save.called + + def test_cancelled_picker_yields_no_servers(self): + with patch("ucode.mcp.pick_mcp_servers", return_value=None): + assert wizard._author_mcp_servers(WORKSPACE, "profile") == [] + + class TestAdminGate: def test_non_admin_is_rejected(self): with patch.object(wizard, "is_workspace_admin", return_value=False):