Register managed-config MCP servers on launch

src/ucode/cli.py

@@ -81,6 +81,7 @@
 from ucode.mcp import (
     MCP_CLIENTS,
     SKILLS_MCP_KIND,
+    apply_managed_mcp_servers,
     configure_mcp_command,
     configure_skills_mcp_command,
     purge_cross_workspace_mcp_residue,
@@ -1363,6 +1364,37 @@ def _print_budget_panel(recommendation: dict, tool: str, managed: dict | None =
         console.print(panel)
 
 
+def _register_managed_mcp_servers(managed: dict, tool: str, state: dict) -> None:
+    """Apply the managed config's MCP servers to ``tool`` and persist what was registered.
+
+    Persisting under ``managed_mcp_servers`` lets the next launch diff against it, so a server the
+    admin later removes from the config is unregistered rather than left behind. A failure here never
+    blocks the launch — the agent still starts, just without the workspace's MCP servers.
+    """
+    try:
+        registered = apply_managed_mcp_servers(
+            managed,
+            tool,
+            state["workspace"],
+            state.get("profile"),
+            use_pat=bool(state.get("use_pat")),
+        )
+    except RuntimeError as exc:
+        print_warning(f"Could not register your workspace's MCP servers: {exc}")
+        return
+    # Persist even when empty so a config that dropped its last server clears the prior registration.
+    others = [
+        server
+        for server in (state.get("managed_mcp_servers") or [])
+        if isinstance(server, dict) and tool not in (server.get("clients") or [])
+    ]
+    state["managed_mcp_servers"] = others + registered
+    save_state(state)
+    if registered:
+        names = ", ".join(str(server["name"]) for server in registered)
+        print_note(f"Registered workspace MCP server(s) for {TOOL_SPECS[tool]['display']}: {names}")
+
+
 def _launch_tool(
     tool_name: str,
     ctx: typer.Context,
@@ -1594,6 +1626,12 @@ def _launch_tool(
             )
         if recommendation is not None:
             _print_budget_panel(recommendation, tool, managed)
+        # Register the managed config's MCP servers so they reach the agent's `/mcp` list. Nothing
+        # else on this path does it — the config only lists them — so without this a
+        # workspace-published server never shows up. Skipped under --skip-preflight (deliberately
+        # unmanaged) and --dry-run (writes nothing).
+        if managed is not None and not skip_preflight and not is_dry_run():
+            _register_managed_mcp_servers(managed, tool, state)
         print_success(f"Starting {TOOL_SPECS[tool]['display']}")
         launch_agent(tool, state, ctx.args)
     except RuntimeError as exc:

src/ucode/managed_wizard.py

@@ -95,29 +95,47 @@ def _tracing_table_from_state(state: dict) -> str | None:
     return destination if isinstance(destination, str) and destination else None
 
 
-def _mcp_type_for_url(url: str) -> str | None:
-    """Classify a registered MCP server's URL into a managed-config type tag.
+def _mcp_server_from_url(url: str) -> tuple[str, str] | None:
+    """Derive a managed-config ``(name, type)`` entry from a registered server's resolved URL.
 
     ``state.json`` stores each MCP server's resolved URL but not its type, while the managed config
-    stores ``{name, type}`` and lets the developer's ucode rebuild the URL. The URL shape is the only
-    signal available, so map it back. Returns None for a URL that matches nothing known, so unknown
-    servers are skipped rather than published with a guessed type.
+    stores ``{name, type}`` and lets the developer's ucode rebuild the URL. So map the URL back to the
+    type *and* the identifier the ai-gateway ``McpServer.name`` field is meant to hold for that type
+    (a UC name for a UC service, a Genie space id for a genie space, a `<catalog>.<schema>` for
+    vector-search / uc-functions, a connection name for external). Deriving ``name`` from the URL —
+    rather than reusing the local display slug — is what lets the developer's ucode reconstruct the
+    URL on launch. Returns None for a URL that matches nothing reconstructable (e.g. an app's
+    off-workspace host), so those are skipped rather than published unusably.
     """
-    if "/ai-gateway/mcp-services/" in url:
-        return "mcp-service"
+    stripped = url.rstrip("/")
+    marker = "/ai-gateway/mcp-services/"
+    if marker in url:
+        # `.../mcp-services/<catalog>.<schema>.<svc>` — store the dash form the launch path expects.
+        service = url.split(marker, 1)[1].split("/", 1)[0]
+        return service.replace(".", "-"), "mcp-service"
     for fragment, tag in (
         ("/api/2.0/mcp/external/", "external"),
         ("/api/2.0/mcp/genie/", "genie-space"),
+    ):
+        if fragment in url:
+            # external -> connection name; genie -> space id. Both are the single trailing segment.
+            return url.split(fragment, 1)[1].split("/", 1)[0], tag
+    for fragment, tag in (
         ("/api/2.0/mcp/vector-search/", "vector-search"),
         ("/api/2.0/mcp/functions/", "uc-functions"),
     ):
         if fragment in url:
-            return tag
-    if url.rstrip("/").endswith("/api/2.0/mcp/sql"):
-        return "sql"
-    # Databricks apps are the residual case: an arbitrary app host with a /mcp suffix.
-    if url.rstrip("/").endswith("/mcp"):
-        return "app"
+            # `.../<catalog>/<schema>` — store the `<catalog>.<schema>` the launch path splits back.
+            rest = url.split(fragment, 1)[1].split("/")
+            if len(rest) >= 2 and rest[0] and rest[1]:
+                return f"{rest[0]}.{rest[1]}", tag
+            return None
+    if stripped.endswith("/api/2.0/mcp/sql"):
+        return "databricks-sql", "sql"
+    # Databricks apps are the residual case: an arbitrary app host with a /mcp suffix. Its host isn't
+    # reconstructable from the workspace + an id, so it can't be published to the managed config yet.
+    if stripped.endswith("/mcp"):
+        return None
     return None
 
 
@@ -130,18 +148,26 @@ def _mcp_servers_from_state(state: dict) -> list[dict]:
     from ucode.mcp import SKILLS_MCP_KIND
 
     servers: list[dict] = []
+    seen: set[str] = set()
     for entry in state.get("mcp_servers") or []:
         if not isinstance(entry, dict) or entry.get("kind") == SKILLS_MCP_KIND:
             continue
         name = entry.get("name")
         url = entry.get("url")
         if not isinstance(name, str) or not name or not isinstance(url, str):
             continue
-        tag = _mcp_type_for_url(url)
-        if tag is None:
-            print_warning(f"Skipping MCP server '{name}': unrecognized URL shape ({url}).")
+        resolved = _mcp_server_from_url(url)
+        if resolved is None:
+            print_warning(
+                f"Skipping MCP server '{name}': ucode can't publish it to a managed config "
+                f"(unrecognized or app-hosted URL: {url})."
+            )
+            continue
+        config_name, tag = resolved
+        if config_name in seen:
             continue
-        servers.append({"name": name, "type": tag})
+        seen.add(config_name)
+        servers.append({"name": config_name, "type": tag})
     return servers
 
 
@@ -935,7 +961,10 @@ def setup_command(from_file: str | None = None) -> int:
     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()
+        # Managed configs can't carry a Databricks app (its host isn't reconstructable from the
+        # workspace), so hide apps from the picker rather than let an admin pick one that is then
+        # dropped from the published config.
+        configure_mcp_command(exclude_sources={"apps"})
         mcp_servers = _mcp_servers_from_state(load_state())
         if mcp_servers:
             manifest["mcp_servers"] = mcp_servers

src/ucode/mcp.py

@@ -351,7 +351,12 @@ def remove_client_mcp_server(client: str, name: str) -> list[str]:
 
 def revert_mcp_configs(state: dict) -> dict[str, bool]:
     results: dict[str, bool] = {}
-    for server in state.get("mcp_servers") or []:
+    # Both the developer's own servers and any registered from the workspace's managed config, so a
+    # revert leaves no ucode-added MCP server behind in an agent's config.
+    all_servers = list(state.get("mcp_servers") or []) + list(
+        state.get("managed_mcp_servers") or []
+    )
+    for server in all_servers:
         name = server.get("name")
         if not isinstance(name, str) or not name:
             continue
@@ -1025,6 +1030,131 @@ def _mcp_server_clients(server: dict) -> list[str]:
     return [client for client in (server.get("clients") or []) if client in MCP_CLIENTS]
 
 
+def _is_app_mcp_server(server: dict) -> bool:
+    """Whether a registered server points at a Databricks app (an off-workspace ``*/mcp`` host).
+
+    Apps are the residual ``/mcp`` URL shape — everything else ucode registers is a known
+    workspace-relative path. Used to hide already-registered apps from the picker where they can't be
+    published (``ucode setup``)."""
+    url = server.get("url")
+    if not isinstance(url, str):
+        return False
+    stripped = url.rstrip("/")
+    known = (
+        "/ai-gateway/mcp-services/",
+        "/api/2.0/mcp/external/",
+        "/api/2.0/mcp/genie/",
+        "/api/2.0/mcp/vector-search/",
+        "/api/2.0/mcp/functions/",
+    )
+    if any(fragment in url for fragment in known):
+        return False
+    if stripped.endswith("/api/2.0/mcp/sql"):
+        return False
+    return stripped.endswith("/mcp")
+
+
+def managed_mcp_server_entry(name: str, mcp_type: str, workspace: str) -> tuple[str, str] | None:
+    """Rebuild an ``(entry_name, url)`` pair from a managed config's ``{name, type}`` entry.
+
+    ``entry_name`` is the identifier the server is registered under with the agent (dots stripped,
+    since the agent CLIs reject them); ``url`` is what the proxy forwards to. Returns None for a
+    type/name this can't reconstruct, so the caller skips it rather than registering a broken server.
+    Mirrors the shapes :func:`_resolve_mcp_selection` builds for the interactive picker, so a managed
+    and a locally-configured copy of the same server land on the same name.
+
+    The ai-gateway ``McpServer.name`` field is interpreted per ``type`` (see the proto): a UC name for
+    a UC service, a Genie space id for a genie space, a connection name for external, and — as ucode
+    serializes them — a `<catalog>.<schema>` for vector-search / uc-functions.
+    """
+    if mcp_type == "sql":
+        return "databricks-sql", f"{workspace}/api/2.0/mcp/sql"
+    if mcp_type == "external":
+        return name, f"{workspace}/api/2.0/mcp/external/{name}"
+    if mcp_type == "mcp-service":
+        # Stored in dash form (`system-ai-dbsql`), which is already the registered name; the URL wants
+        # the UC dotted form. Only the catalog and schema separators (first two dashes) become dots —
+        # the service name keeps its own dashes/underscores.
+        parts = name.split("-", 2)
+        if len(parts) != 3:
+            return None
+        return name, build_mcp_service_url(workspace, ".".join(parts))
+    if mcp_type == "genie-space":
+        # `name` is the Genie space id (per the proto); register under the id-based name the
+        # interactive path falls back to, and point the URL at the space.
+        return f"databricks-genie-{name}", f"{workspace}/api/2.0/mcp/genie/{name}"
+    if mcp_type in ("vector-search", "uc-functions"):
+        # `name` is a `<catalog>.<schema>`; the URL is workspace-relative on that pair, and the
+        # registered name is the same dot-free slug the interactive path uses.
+        catalog, _, schema = name.partition(".")
+        if not catalog or not schema or "." in schema:
+            return None
+        url_path = "vector-search" if mcp_type == "vector-search" else "functions"
+        name_prefix = (
+            "databricks-vector-search" if mcp_type == "vector-search" else "databricks-functions"
+        )
+        entry_name = _catalog_schema_server_name(name_prefix, catalog, schema, set())
+        return entry_name, f"{workspace}/api/2.0/mcp/{url_path}/{catalog}/{schema}"
+    return None
+
+
+def apply_managed_mcp_servers(
+    managed: dict, tool: str, workspace: str, profile: str | None = None, *, use_pat: bool = False
+) -> list[dict]:
+    """Register the managed config's MCP servers with ``tool`` so they reach its `/mcp` list.
+
+    The managed config only lists ``{name, type}`` entries; nothing else on the launch path turns
+    them into agent MCP registrations, so without this a workspace-published server never shows up.
+    Reconstructs each entry's ``(name, url)`` (see :func:`managed_mcp_server_entry`), diffs against
+    what ucode previously registered, and applies the change for the launching tool only. Entries
+    whose URL can't be rebuilt (e.g. ``app``, which needs an off-workspace host) are skipped.
+
+    Returns the server dicts registered (for state persistence); an empty list when the config names
+    none, or names only types that can't yet be reconstructed.
+    """
+    if tool not in MCP_CLIENTS:
+        return []
+    entries = managed.get("mcp_servers")
+    if not isinstance(entries, list):
+        return []
+    working: list[dict] = []
+    seen: set[str] = set()
+    skipped: list[str] = []
+    for entry in entries:
+        if not isinstance(entry, dict):
+            continue
+        name = entry.get("name")
+        mcp_type = entry.get("type")
+        if not isinstance(name, str) or not name or not isinstance(mcp_type, str):
+            continue
+        resolved = managed_mcp_server_entry(name, mcp_type, workspace)
+        if resolved is None:
+            skipped.append(f"{name} ({mcp_type})")
+            continue
+        entry_name, url = resolved
+        if entry_name in seen:
+            continue
+        seen.add(entry_name)
+        working.append({"name": entry_name, "url": url, "auth": "proxy", "clients": [tool]})
+    if skipped:
+        print_warning(
+            "Skipping managed MCP server(s) ucode can't yet auto-register from the workspace "
+            f"config: {', '.join(skipped)}. Add them with `ucode configure mcp`."
+        )
+    if not working:
+        return []
+    # Diff against the managed servers ucode registered on a prior launch so a removed entry is
+    # unregistered and an unchanged one is a no-op. Only this tool's managed servers are considered.
+    state = load_state()
+    previous = [
+        server
+        for server in (state.get("managed_mcp_servers") or [])
+        if isinstance(server, dict) and tool in (server.get("clients") or [])
+    ]
+    apply_mcp_server_changes(previous, working, [tool], workspace, profile, use_pat=use_pat)
+    return working
+
+
 def _resolve_mcp_selection(
     selection: str,
     workspace: str,
@@ -1484,12 +1614,18 @@ def _resolve_location_mcp_servers(
 )
 
 
-def prompt_for_mcp_search_sources() -> set[str] | None:
+def prompt_for_mcp_search_sources(exclude_sources: set[str] | None = None) -> set[str] | None:
     """First wizard step: choose which sources to search. Returns the set of
-    selected source keys, or `None` if the user cancelled (Ctrl-C)."""
+    selected source keys, or `None` if the user cancelled (Ctrl-C).
+
+    ``exclude_sources`` drops source keys the caller can't use — e.g. `ucode setup` excludes
+    ``apps`` because a managed config can't carry an app's off-workspace host, so offering it would
+    let an admin pick a server that is then silently dropped."""
+    excluded = exclude_sources or set()
     choices = [
         questionary.Choice(title=label, value=key, checked=checked)
         for key, label, checked in MCP_SEARCH_SOURCES
+        if key not in excluded
     ]
     selection = _scrolling_checkbox(
         "Search for:",
@@ -1546,7 +1682,15 @@ def setup_mcp_clients(state: dict, section: str) -> tuple[str, str | None, list[
     return workspace, profile, clients
 
 
-def configure_mcp_command(location: str | None = None, services: set[str] | None = None) -> int:
+def configure_mcp_command(
+    location: str | None = None,
+    services: set[str] | None = None,
+    *,
+    exclude_sources: set[str] | None = None,
+) -> int:
+    """Interactive MCP picker. ``exclude_sources`` hides search sources the caller can't use —
+    `ucode setup` passes ``{"apps"}`` because a managed config can't carry an app's off-workspace
+    host, so an app picked here would be silently dropped from the published config."""
     if services is not None and location is None:
         # `--services` works standalone with full names (`system.ai.github`): the
         # `<catalog>.<schema>` to configure is derived from them. Bare short names
@@ -1586,18 +1730,23 @@ def configure_mcp_command(location: str | None = None, services: set[str] | None
             print_success("Saved")
         return 0
 
+    excluded_sources = exclude_sources or set()
     original_mcp_servers: list[dict] = list(state.get("mcp_servers") or [])
     # Skills connections are managed by `configure skills`, so keep them out of
     # the picker and carry them through untouched.
     skills_servers = _skills_entries(original_mcp_servers)
     picker_servers = [s for s in original_mcp_servers if s.get("kind") != SKILLS_MCP_KIND]
+    # Drop already-registered servers from an excluded source too (e.g. a previously-added app under
+    # `ucode setup`), so the picker never shows a server the caller couldn't re-add.
+    if "apps" in excluded_sources:
+        picker_servers = [s for s in picker_servers if not _is_app_mcp_server(s)]
     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()
+        sources = prompt_for_mcp_search_sources(exclude_sources=excluded_sources)
         if sources is None:
             return 0
         discovered = _discover_selected_mcp_sources(workspace, profile, sources)