diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 5466f2da..8f6c27a1 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1140,7 +1140,10 @@ def _configure_optional_setup(state: dict, tools: list[str]) -> None: return install_databricks_ai_tools_for_agents(tools, state) - configure_mcp_command() + # Register every MCP service the user can access rather than making them pick during setup; + # `configure mcp` (below) is the granular picker for anyone who wants to choose. + configure_mcp_command(all_services=True) + print_note("To pick specific MCP servers instead, run `ug configure mcp`.") @mcp_app.command("add") @@ -1177,6 +1180,14 @@ def mcp_add( "the server is registered for every already-configured agent.", ), ] = None, + all_services: Annotated[ + bool, + typer.Option( + "--all", + help="Register every MCP service you can access across the workspace, without the " + "picker. Can't be combined with --location or --services.", + ), + ] = False, ) -> None: """Add Databricks MCP servers to installed coding tools. @@ -1192,7 +1203,9 @@ def mcp_add( ) try: scope = _configure_agents_for_mcp(sorted(requested_agents)) if requested_agents else None - add_mcp_command(location=location, services=selected, agents=scope) + add_mcp_command( + location=location, services=selected, agents=scope, all_services=all_services + ) except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None @@ -3138,13 +3151,21 @@ def configure_mcp( "`app:` (workspace access required).", ), ] = None, + all_services: Annotated[ + bool, + typer.Option( + "--all", + help="Register every MCP service you can access across the workspace, without the " + "picker. Can't be combined with --location or --services.", + ), + ] = False, ) -> None: """Add Databricks MCP servers to installed coding tools.""" # `--services` absent -> None (whole schema); present (even empty) -> the # explicit subset, so `--services ""` deselects everything. selected = None if services is None else {s.strip() for s in services.split(",") if s.strip()} try: - configure_mcp_command(location=location, services=selected) + configure_mcp_command(location=location, services=selected, all_services=all_services) except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index a4ffdd9b..328faae1 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -17,13 +17,6 @@ import subprocess import time from collections.abc import Callable -from concurrent.futures import ( - ThreadPoolExecutor, - as_completed, -) -from concurrent.futures import ( - TimeoutError as FutureTimeoutError, -) from dataclasses import dataclass from decimal import Decimal, InvalidOperation from pathlib import Path @@ -2446,35 +2439,11 @@ def resolve_provider_launch_model(model: str | None, provider_models: dict[str, _UC_LIST_PAGE_SIZE = 200 _UC_LIST_MAX_PAGES = 50 -_UC_FUNCTION_PROBE_WORKERS = 16 _UC_LIST_HTTP_TIMEOUT = 10 -# Most MCP services live outside `system.ai`, so this workspace-wide walk needs -# enough time to enumerate them; a slow workspace still degrades to partial -# results once the budget is exceeded instead of hanging indefinitely. +# Safety valve for the workspace-wide (metastore-scope) MCP-services listing: it is a single +# paginated call, but a pathologically large or slow listing still degrades to partial results +# once the budget is exceeded instead of hanging indefinitely. _MCP_SERVICES_WALK_DEADLINE_SECONDS = 30.0 -# Skip UC catalogs whose schemas almost never carry user-callable functions -# you'd want to expose as agent tools. -_UC_FUNCTIONS_SKIP_CATALOGS = frozenset( - {"__databricks_internal", "hive_metastore", "samples", "system"} -) - - -def _drain_with_deadline(futures: dict, deadline: float, on_result) -> None: - """Iterate `futures` via `as_completed`, calling `on_result(value, key)` per - completed future, until either all are done or `deadline` passes. Per-task - exceptions are swallowed so one failure doesn't stop the rest.""" - remaining = max(0.0, deadline - time.monotonic()) - try: - for future in as_completed(futures, timeout=remaining): - try: - value = future.result() - except Exception: # noqa: BLE001 - continue - on_result(value, futures[future]) - if time.monotonic() > deadline: - break - except FutureTimeoutError: - pass def _paginated_json_items( @@ -2486,17 +2455,24 @@ def _paginated_json_items( page_size: int = _UC_LIST_PAGE_SIZE, max_pages: int = _UC_LIST_MAX_PAGES, timeout: int = 30, + deadline: float | None = None, ) -> tuple[list[dict], str | None]: """Walk a Databricks `next_page_token` listing and return all items. Returns (items, reason). Items are dicts; reason is None on success or a - short description of why the walk stopped early. + short description of why the walk stopped early. When ``deadline`` (a + ``time.monotonic()`` value) is given, pagination stops before fetching a + further page once it passes, returning the pages gathered so far — the first + page is always fetched so a bounded walk still makes forward progress. """ items: list[dict] = [] page_token: str | None = None seen_tokens: set[str] = set() last_reason: str | None = None - for _ in range(max_pages): + for page_index in range(max_pages): + if page_index and deadline is not None and time.monotonic() > deadline: + last_reason = last_reason or "deadline exceeded during listing" + break params: dict[str, str] = {"max_results": str(page_size)} if extra_params: params.update(extra_params) @@ -2528,108 +2504,36 @@ def list_all_mcp_services( on_progress: Callable[[int, int, int], None] | None = None, on_services: Callable[[list[str]], None] | None = None, ) -> tuple[list[str], str | None]: - """Return sorted unique MCP-service full names across every `.` - in the workspace. The mcp-services API is one-schema-per-call, so this walks - catalogs -> schemas -> mcp-services in parallel under a wall-clock budget, - returning partial results once `deadline_seconds` is exceeded. - - `on_progress`, if given, is called as each schema's listing completes with - `(schemas_done, schemas_total, services_found)` so callers can render a live - count. `on_services`, if given, is called with each schema's newly-found service - names (deduped against everything emitted so far) so callers can stream results - into a picker as the walk progresses instead of waiting for the full result. Both - are invoked serially from the draining thread (not the workers). - - This walk is the slow, workspace-wide counterpart to `list_mcp_services` - (single schema).""" + """Return sorted unique MCP-service full names across the whole workspace. + + The mcp-services API supports a metastore scope (``parent`` omitted) that lists every service + the caller can see in one paginated call. That replaces the old `catalogs -> schemas -> + mcp-services` walk, which couldn't enumerate a large metastore (thousands of catalogs) within + any usable budget and so surfaced nothing but the `system.ai` list. `deadline_seconds` bounds + pagination as a safety valve; a truncated listing returns whatever it gathered. + + `on_services`, if given, is called once with the discovered service names so a caller can + stream them into a picker; `on_progress`, if given, is called once as `(1, 1, count)` for a + live count. This is the workspace-wide counterpart to `list_mcp_services` (single schema).""" hostname = workspace_hostname(workspace) deadline = time.monotonic() + deadline_seconds - - catalogs, catalogs_reason = _paginated_json_items( - f"https://{hostname}/api/2.1/unity-catalog/catalogs", + # Metastore scope: no `parent` query param (see the API's own error text: parent must be either + # '' for metastore scope or 'schemas/.' for a single schema). + items, reason = _paginated_json_items( + f"https://{hostname}/api/2.1/unity-catalog/mcp-services", token, - items_key="catalogs", + items_key="mcp_services", timeout=_UC_LIST_HTTP_TIMEOUT, + deadline=deadline, ) - if not catalogs: - return [], catalogs_reason or "no UC catalogs found" - - catalog_names = [ - c["name"] - for c in catalogs - if isinstance(c.get("name"), str) - and c["name"] - and c["name"] not in _UC_FUNCTIONS_SKIP_CATALOGS - ] - if not catalog_names: - return [], "no user UC catalogs found" - if time.monotonic() > deadline: - return [], "deadline exceeded while listing UC catalogs" - - # Parallel per-catalog schema listing. - schema_refs: list[str] = [] - schema_workers = max(1, min(_UC_FUNCTION_PROBE_WORKERS, len(catalog_names))) - with ThreadPoolExecutor(max_workers=schema_workers) as pool: - schema_futures = { - pool.submit( - _paginated_json_items, - f"https://{hostname}/api/2.1/unity-catalog/schemas", - token, - items_key="schemas", - extra_params={"catalog_name": cat}, - timeout=_UC_LIST_HTTP_TIMEOUT, - ): cat - for cat in catalog_names - } - - def collect_schemas(result, catalog): - schemas, _ = result - for schema in schemas: - schema_name = schema.get("name") - if ( - isinstance(schema_name, str) - and schema_name - and schema_name != "information_schema" - ): - schema_refs.append(f"{catalog}.{schema_name}") - - _drain_with_deadline(schema_futures, deadline, collect_schemas) - pool.shutdown(wait=False, cancel_futures=True) - - if not schema_refs: - if time.monotonic() > deadline: - return [], "deadline exceeded while listing UC schemas" - return [], "no UC schemas found" - - # Parallel per-schema mcp-services listing. - names: set[str] = set() - schemas_total = len(schema_refs) - schemas_done = 0 - probe_workers = max(1, min(_UC_FUNCTION_PROBE_WORKERS, schemas_total)) - with ThreadPoolExecutor(max_workers=probe_workers) as pool: - service_futures = { - pool.submit(list_mcp_services, workspace, token, ref): ref for ref in schema_refs - } - - def collect_services(result, _ref): - nonlocal schemas_done - found, _ = result - new = [n for n in found if n not in names] - names.update(found) - schemas_done += 1 - if on_progress is not None: - on_progress(schemas_done, schemas_total, len(names)) - if on_services is not None and new: - on_services(sorted(new)) - - _drain_with_deadline(service_futures, deadline, collect_services) - pool.shutdown(wait=False, cancel_futures=True) - + names = sorted({full for svc in items if (full := _mcp_service_full_name(svc, ""))}) + if on_services is not None and names: + on_services(names) + if on_progress is not None: + on_progress(1, 1, len(names)) if not names: - if time.monotonic() > deadline: - return [], "deadline exceeded while listing MCP services" - return [], "no MCP services found" - return sorted(names), None + return [], reason or "no MCP services found" + return names, None def _get_anthropic_models_json(workspace: str, token: str) -> tuple[dict | list | None, str | None]: diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index bb927ce6..1f1251a4 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -458,6 +458,56 @@ def _catalog_schema_server_name(prefix: str, catalog: str, schema: str, taken: s return f"{candidate}-{counter}" +def _mcp_service_full_name_from_url(url: str) -> str | None: + """The UC ``..`` of a UC MCP-service entry, recovered from its URL + (see :func:`build_mcp_service_url`), or ``None`` when ``url`` isn't a UC MCP-service URL. + Only UC MCP services use the ``/ai-gateway/mcp-services/`` route, so this doubles as the + "is this entry a UC MCP service?" test the leaf-rename pass relies on.""" + marker = "/ai-gateway/mcp-services/" + idx = url.find(marker) + if idx == -1: + return None + tail = url[idx + len(marker) :].split("?", 1)[0].strip("/") + return tail or None + + +def _apply_mcp_service_leaf_names(servers: list[dict]) -> None: + """Rename UC MCP-service entries in ``servers`` in place so each registers under its bare + service id (the leaf of ``..``) instead of the full dashed path — e.g. + ``system.ai.github`` registers as ``github`` (the agent-visible ``mcp__github__`` tool prefix) + rather than ``system-ai-github``. The full UC name always stays in the entry's ``url`` (and is + rebuilt from the managed config), so only the agent-facing name changes; loading is untouched. + + The leaf is used only when unambiguous across the desired set: if two services share an id + across schemas, or a service's leaf clashes with a non-service entry's name, the colliding + services keep their full dashed path so every server still registers under a distinct name.""" + service_full: dict[int, str] = {} + for i, server in enumerate(servers): + url = server.get("url") + if isinstance(url, str): + full = _mcp_service_full_name_from_url(url) + if full and "." in full: + service_full[i] = full + + # Names owned by non-service entries (skills/external/app/...); never rename a service onto one. + reserved = { + server.get("name") + for i, server in enumerate(servers) + if i not in service_full and isinstance(server.get("name"), str) + } + leaf_counts: dict[str, int] = {} + for full in service_full.values(): + leaf = full.rsplit(".", 1)[-1] + leaf_counts[leaf] = leaf_counts.get(leaf, 0) + 1 + + for i, full in service_full.items(): + leaf = full.rsplit(".", 1)[-1] + if leaf and leaf_counts[leaf] == 1 and leaf not in reserved: + servers[i]["name"] = leaf + else: + servers[i]["name"] = full.replace(".", "-") + + def _picker_style() -> questionary.Style: return questionary.Style( [ @@ -558,8 +608,11 @@ def _mcp_service_choice(name: str, known_names: set[str], additive: bool) -> que (and dedupes by value against what's already shown). An already-registered service is a removable toggle under `configure mcp` and a non-toggleable note under `mcp add` (additive); an unregistered one is an add-choice.""" - registered_as = name.replace(".", "-") - display_title = f"MCP: {name}" + # Registered under the bare service id (`github`); older configs may still hold the full + # dashed path (`system-ai-github`), so an already-configured server is matched against either. + leaf = name.rsplit(".", 1)[-1] + registered_as = leaf if leaf in known_names else name.replace(".", "-") + display_title = f"MCP: {leaf}" if registered_as in known_names: if additive: return questionary.Choice( @@ -909,10 +962,13 @@ def known_choice(name: str, title: str | None = None) -> questionary.Choice: # at the end. The `managed:sql` selection value is still resolvable for managed configs. for name in available_mcp_service_names or []: - # Picker shows the dotted UC name; state/agents store the dashed form - # (see resolver). The shared helper is also used by the background walk that - # streams more services in, so up-front and streamed rows match exactly. + # Picker shows the bare service id; state/agents store that id too (or the full dashed + # path on collision — see `_apply_mcp_service_leaf_names`). Track both forms so the + # known-server fallback below never re-lists a service already shown here. The shared + # helper is also used by the background walk that streams more services in, so up-front + # and streamed rows match exactly. choices.append(_mcp_service_choice(name, known_names, additive)) + displayed_names.add(name.rsplit(".", 1)[-1]) displayed_names.add(name.replace(".", "-")) for name in available_external_names: @@ -1345,6 +1401,11 @@ def apply_mcp_server_changes( *, use_pat: bool = False, ) -> bool: + # Register each UC MCP service under its bare id (`github`, not `system-ai-github`) — the + # agent-visible `mcp____` tool prefix. Done here, the single chokepoint every path funnels + # through after its own assembly/append, so carried-over servers are renamed consistently too. + # Mutates `working_servers` in place so the caller persists the same (renamed) list to state. + _apply_mcp_service_leaf_names(working_servers) original_by_name = _servers_by_name(original_servers) working_by_name = _servers_by_name(working_servers) @@ -1543,11 +1604,26 @@ def _resolve_location_mcp_servers( if full_name in services or full_name.split(".")[-1] in services ] + working_servers = _mcp_service_entries(names, clients, original_servers, workspace) + return [*working_servers, *_skills_entries(original_servers)] + + +def _mcp_service_entries( + full_names: list[str], clients: list[str], original_servers: list[dict], workspace: str +) -> list[dict]: + """Build (or reuse) an MCP-service server entry per `..` full name. + + Shared by the `--location` and `--all` paths. An already-registered copy is matched under + either its bare id (`github`, the name `apply_mcp_server_changes` assigns) or the older full + dashed path, so its existing clients are preserved; the leaf-vs-dashed registered name is + finalized later by `apply_mcp_server_changes`.""" original_by_name = _servers_by_name(original_servers) - working_servers: list[dict] = [] - for full_name in names: + servers: list[dict] = [] + for full_name in full_names: entry_name = full_name.replace(".", "-") - original = original_by_name.get(entry_name) + original = original_by_name.get(full_name.rsplit(".", 1)[-1]) or original_by_name.get( + entry_name + ) original_clients = list((original or {}).get("clients") or []) merged_clients = original_clients + [c for c in clients if c not in original_clients] candidate = { @@ -1557,10 +1633,27 @@ def _resolve_location_mcp_servers( "clients": merged_clients, } if original is not None and original == candidate: - working_servers.append(original.copy()) + servers.append(original.copy()) else: - working_servers.append(candidate) - return [*working_servers, *_skills_entries(original_servers)] + servers.append(candidate) + return servers + + +def _resolve_all_mcp_servers( + workspace: str, profile: str | None, clients: list[str], original_servers: list[dict] +) -> tuple[list[dict], str | None]: + """Build the desired MCP server list for `--all`: every MCP service the caller can access + across the whole workspace (metastore-wide), plus any existing skills connection, preserved. + Returns ``(servers, reason)`` where ``reason`` is a non-None listing-failure description. + + The workspace-wide listing is permission-filtered server-side, so this is exactly the set the + user is entitled to. Like `--location`, it's a strict replacement — a previously-registered + service the user can no longer see is removed by `apply_mcp_server_changes`.""" + token = get_databricks_token(workspace, profile) + with spinner("Discovering MCP services you can access..."): + names, reason = list_all_mcp_services(workspace, token) + working_servers = _mcp_service_entries(names, clients, original_servers, workspace) + return [*working_servers, *_skills_entries(original_servers)], reason # The interactive picker searches a single source: MCP services (the `/ai-gateway/mcp-services/` @@ -1672,14 +1765,16 @@ def add_mcp_command( location: str | None = None, services: set[str] | None = None, agents: set[str] | None = None, + *, + all_services: bool = False, ) -> int: """`ucode mcp add`: register Databricks MCP servers WITHOUT removing any that are already configured. Uses the same discovery and options as `configure mcp` — the interactive - picker, or the non-interactive `--location`/`--services` paths — but is purely - additive: unlike `configure mcp`, it never removes servers outside the - selection. + picker, the non-interactive `--location`/`--services` paths, or `--all` (every + MCP service the caller can access) — but is purely additive: unlike + `configure mcp`, it never removes servers outside the selection. ``agents`` scopes the registration to that subset of configured MCP clients (the agents must already be configured — the `--agents` CLI option sets up any @@ -1690,7 +1785,13 @@ def add_mcp_command( # so it's a no-op (and doesn't need --location the way a real subset does). print_note("No MCP services given to add (empty --services); nothing to do.") return 0 - return configure_mcp_command(location=location, services=services, append=True, agents=agents) + return configure_mcp_command( + location=location, + services=services, + all_services=all_services, + append=True, + agents=agents, + ) def _configure_v2_mcp_selectors( @@ -1755,6 +1856,13 @@ def _configure_v2_mcp_selectors( if changed or original_mcp_servers != working_mcp_servers: state["mcp_servers"] = working_mcp_servers save_state(state) + # `apply_mcp_server_changes` may have renamed UC services to their leaf ids in place, so + # recompute the working names for an accurate added/removed count in the summary. + working_names = { + n + for s in working_mcp_servers + if s.get("kind") != SKILLS_MCP_KIND and (n := _server_name(s)) + } added = sorted(working_names - set(original_by_name)) removed = [] if append else sorted(set(original_by_name) - working_names) print_success(_mcp_change_summary(added, removed, clients)) @@ -1765,6 +1873,7 @@ def configure_mcp_command( location: str | None = None, services: set[str] | None = None, *, + all_services: bool = False, exclude_sources: set[str] | None = None, append: bool = False, agents: set[str] | None = None, @@ -1773,10 +1882,13 @@ def configure_mcp_command( `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. - ``append`` (used by `ucode mcp add`) makes the command purely additive: the - final server list is unioned with the already-configured servers, so nothing - outside the current selection is removed. ``agents`` scopes the operation to - that subset of configured MCP clients.""" + ``all_services`` skips the picker and registers every MCP service the caller can access across + the workspace (the `--all` / onboarding path). ``append`` (used by `ucode mcp add`) makes the + command purely additive: the final server list is unioned with the already-configured servers, + so nothing outside the current selection is removed. ``agents`` scopes the operation to that + subset of configured MCP clients.""" + if all_services and (location is not None or services is not None): + raise RuntimeError("--all can't be combined with --location or --services.") if services is not None: # A typed V2 MCP selector (`vector-search:main.docs`, `uc-functions:main.tools`, # `external:conn`, `genie-space:`, `app:`) names a server the picker no @@ -1810,15 +1922,30 @@ def configure_mcp_command( ) location = next(iter(schemas)) state = load_state() + if all_services and agents is None: + # `--all` targets the agents you actually set up with ucode (`available_tools`), not every + # installed MCP-capable CLI — so a Codex-only user isn't surprised by Cursor (an MCP-only + # client) getting configured too. Fall back to the full configured set only when no + # model-routing agent was configured (e.g. a Cursor-only user); `--agents` overrides either. + configured = configured_mcp_clients(state, available_mcp_clients()) + scoped = set(state.get("available_tools") or []) & set(configured) + agents = scoped or None workspace, profile, clients = setup_mcp_clients( state, "Add MCP Servers" if append else "MCP Servers", agents=agents ) original_mcp_servers_for_location: list[dict] = list(state.get("mcp_servers") or []) - if location is not None: - working_mcp_servers = _resolve_location_mcp_servers( - workspace, profile, clients, location, original_mcp_servers_for_location, services - ) + if all_services or location is not None: + all_reason: str | None = None + if all_services: + working_mcp_servers, all_reason = _resolve_all_mcp_servers( + workspace, profile, clients, original_mcp_servers_for_location + ) + else: + assert location is not None # guaranteed by the `all_services or location is not None` + working_mcp_servers = _resolve_location_mcp_servers( + workspace, profile, clients, location, original_mcp_servers_for_location, services + ) if append: working_mcp_servers = _union_missing( original_mcp_servers_for_location, working_mcp_servers @@ -1834,6 +1961,21 @@ def configure_mcp_command( if changed or original_mcp_servers_for_location != working_mcp_servers: state["mcp_servers"] = working_mcp_servers save_state(state) + if all_services: + # Report what's registered (not whether this run changed anything), so a re-run when + # everything is already registered doesn't misreport "none found". + service_count = sum( + 1 + for s in working_mcp_servers + if isinstance(s.get("url"), str) and _mcp_service_full_name_from_url(s["url"]) + ) + if service_count: + print_success(f"Registered {service_count} MCP server(s) you have access to") + elif all_reason: + print_warning(f"Couldn't list MCP services you can access: {all_reason}") + else: + print_note("No MCP servers you can access were found.") + elif changed or original_mcp_servers_for_location != working_mcp_servers: print_success("Saved") return 0 @@ -1928,6 +2070,13 @@ def configure_mcp_command( if changed or original_mcp_servers != working_mcp_servers: state["mcp_servers"] = working_mcp_servers save_state(state) + # `apply_mcp_server_changes` may have renamed UC services to their leaf ids in place, so + # recompute the working names for an accurate added/removed count in the summary. + working_names = { + n + for s in working_mcp_servers + if s.get("kind") != SKILLS_MCP_KIND and (n := _server_name(s)) + } added = sorted(working_names - set(original_by_name)) # `add` never removes; the union above re-keeps unselected servers. removed = [] if append else sorted(set(original_by_name) - working_names) diff --git a/tests/test_cli.py b/tests/test_cli.py index cd2d3f63..b930d54c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2035,7 +2035,8 @@ def test_optional_setup_installs_ai_tools_and_configures_mcp(self): assert state["databricks_ai_tools_enabled"] is True mock_save.assert_called_once_with(state) mock_install.assert_called_once_with(["claude", "codex"], state) - mock_mcp.assert_called_once_with() + # Onboarding registers every accessible MCP service rather than opening the picker. + mock_mcp.assert_called_once_with(all_services=True) def test_optional_setup_decline_does_nothing(self): import ucode.cli as cli_mod diff --git a/tests/test_databricks.py b/tests/test_databricks.py index d402fe66..281bd08a 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -5,6 +5,7 @@ import json import os import subprocess +import time from decimal import Decimal from urllib.parse import parse_qs @@ -1100,106 +1101,162 @@ def test_http_404_reason_surfaces_for_invalid_parent(self, monkeypatch): class TestListAllMcpServices: - """Workspace-wide walk: catalogs -> schemas -> per-schema mcp-services.""" + """Workspace-wide listing via a single metastore-scope mcp-services call.""" - def _fake_http(self, catalogs, schemas_by_catalog, services_by_schema): - """Route `_http_get_json` by URL to the right stubbed payload.""" + def _metastore_http(self, services_by_page, *, capture=None): + """Route `_http_get_json` to metastore-scope mcp-services pages. + + `services_by_page` is a list of pages; each page is a list of entries (raw dicts, or bare + full-name strings turned into `{"name": "mcp-services/"}`). Pages after the first are + served on the matching `page_token`. `capture`, if given, records every requested URL. + """ + + def as_entry(item): + return item if isinstance(item, dict) else {"name": f"mcp-services/{item}"} def fake_get(url, token, timeout=30): - if "unity-catalog/catalogs" in url: - return {"catalogs": [{"name": c} for c in catalogs]}, None - if "unity-catalog/schemas" in url: - cat = url.split("catalog_name=")[1].split("&")[0] - return {"schemas": [{"name": s} for s in schemas_by_catalog.get(cat, [])]}, None - if "unity-catalog/mcp-services" in url: - # parent is url-encoded as `schemas%2F.` - parent = url.split("parent=")[1].split("&")[0] - schema_ref = parent.replace("schemas%2F", "").replace("schemas/", "") - return { - "mcp_services": [ - {"name": f"mcp-services/{full}"} - for full in services_by_schema.get(schema_ref, []) - ] - }, None - return None, "unexpected url" + if capture is not None: + capture.append(url) + if "unity-catalog/mcp-services" not in url: + return None, "unexpected url" + page = 0 + if "page_token=" in url: + page = int(url.split("page_token=")[1].split("&")[0]) + body = {"mcp_services": [as_entry(i) for i in services_by_page[page]]} + if page + 1 < len(services_by_page): + body["next_page_token"] = str(page + 1) + return body, None return fake_get - def test_aggregates_services_across_catalogs_and_schemas(self, monkeypatch): + def test_lists_every_service_via_metastore_scope(self, monkeypatch): + urls: list[str] = [] monkeypatch.setattr( db_mod, "_http_get_json", - self._fake_http( - catalogs=["mycat", "other"], - schemas_by_catalog={"mycat": ["myschema", "information_schema"], "other": ["ops"]}, - services_by_schema={ - "mycat.myschema": ["mycat.myschema.weather", "mycat.myschema.news"], - "other.ops": ["other.ops.pager"], - }, + self._metastore_http( + [ + [ + "main.default.sanjay_tavily", + "users.someone.my_mcp", + "system.ai.github", + "main.default.sanjay_tavily", # duplicate: de-duplicated + ] + ], + capture=urls, ), ) names, reason = db_mod.list_all_mcp_services(WS, "token") assert reason is None - # information_schema is skipped; results are sorted and de-duplicated. + # Every catalog/schema is returned in one call, sorted and de-duplicated. assert names == [ - "mycat.myschema.news", - "mycat.myschema.weather", - "other.ops.pager", + "main.default.sanjay_tavily", + "system.ai.github", + "users.someone.my_mcp", ] + # Metastore scope: the request carries no `parent` (that would scope it to one schema). + assert urls and all("parent=" not in u for u in urls) - def test_reports_progress_per_schema(self, monkeypatch): + def test_paginates_the_listing(self, monkeypatch): monkeypatch.setattr( db_mod, "_http_get_json", - self._fake_http( - catalogs=["mycat"], - schemas_by_catalog={"mycat": ["a", "b"]}, - services_by_schema={"mycat.a": ["mycat.a.one"], "mycat.b": ["mycat.b.two"]}, + self._metastore_http( + [["main.a.one"], ["other.b.two"]], # two pages via next_page_token ), ) + + names, reason = db_mod.list_all_mcp_services(WS, "token") + + assert reason is None + assert names == ["main.a.one", "other.b.two"] + + def test_streams_results_and_reports_progress_once(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_get_json", + self._metastore_http([["main.a.one", "main.a.two"]]), + ) + streamed: list[str] = [] progress: list[tuple[int, int, int]] = [] - names, reason = db_mod.list_all_mcp_services( + names, _reason = db_mod.list_all_mcp_services( WS, "token", + on_services=lambda new: streamed.extend(new), on_progress=lambda done, total, found: progress.append((done, total, found)), ) - assert reason is None - assert names == ["mycat.a.one", "mycat.b.two"] - # One callback per schema; the total is fixed and done/found climb. - assert len(progress) == 2 - assert [p[1] for p in progress] == [2, 2] - assert progress[-1][0] == 2 - assert progress[-1][2] == 2 - - def test_skips_internal_catalogs(self, monkeypatch): + assert names == ["main.a.one", "main.a.two"] + assert streamed == ["main.a.one", "main.a.two"] # streamed once, in full + assert progress == [(1, 1, 2)] + + def test_skips_inactive_services(self, monkeypatch): + # A service whose connection is not ACTIVE is excluded (see `_mcp_service_full_name`). monkeypatch.setattr( db_mod, "_http_get_json", - self._fake_http( - catalogs=["system", "hive_metastore", "samples", "__databricks_internal"], - schemas_by_catalog={}, - services_by_schema={}, + self._metastore_http( + [ + [ + {"name": "mcp-services/main.default.live"}, + { + "name": "mcp-services/main.default.pending", + "config": {"connection": {"status": "PENDING"}}, + }, + ] + ] ), ) names, reason = db_mod.list_all_mcp_services(WS, "token") + assert reason is None + assert names == ["main.default.live"] + + def test_returns_reason_when_no_services(self, monkeypatch): + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=30: ({"mcp_services": []}, None) + ) + + names, reason = db_mod.list_all_mcp_services(WS, "token") + assert names == [] - assert reason == "no user UC catalogs found" + assert reason == "no MCP services found" - def test_returns_reason_when_no_catalogs(self, monkeypatch): + def test_surfaces_listing_failure_reason_when_empty(self, monkeypatch): monkeypatch.setattr( - db_mod, "_http_get_json", lambda url, token, timeout=30: ({"catalogs": []}, None) + db_mod, "_http_get_json", lambda url, token, timeout=30: (None, "HTTP 500 Server Error") ) names, reason = db_mod.list_all_mcp_services(WS, "token") assert names == [] - assert reason == "no UC catalogs found" + assert reason == "HTTP 500 Server Error" + + def test_paginated_json_items_stops_at_deadline_after_first_page(self, monkeypatch): + # A past deadline still fetches page 1 (forward progress) but skips further pages, + # returning the partial result and a reason. This bounds the metastore listing as a safety. + calls = {"n": 0} + + def fake_get(url, token, timeout=30): + calls["n"] += 1 + return {"items": [{"name": f"i{calls['n']}"}], "next_page_token": "more"}, None + + monkeypatch.setattr(db_mod, "_http_get_json", fake_get) + + items, reason = db_mod._paginated_json_items( + "https://x/api/2.1/unity-catalog/mcp-services", + "token", + items_key="items", + deadline=time.monotonic() - 1, + ) + + assert calls["n"] == 1 # a further page was available but the deadline stopped it + assert [i["name"] for i in items] == ["i1"] + assert reason == "deadline exceeded during listing" def _foundation_models_payload(names): diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 985efd57..497aa377 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -547,6 +547,70 @@ def test_no_ops_returns_false_without_spinner(self, monkeypatch): assert mcp.apply_mcp_server_changes(servers, servers, ["claude"], WS) is False +class TestMcpServiceLeafNames: + """UC MCP services register under their bare id (`github`), not the full dashed path + (`system-ai-github`), so the agent-visible tool prefix is `mcp__github__`.""" + + def _service(self, full_name, name=None): + return { + "name": name or full_name.replace(".", "-"), + "url": f"{WS}/ai-gateway/mcp-services/{full_name}", + "auth": "proxy", + "clients": ["claude"], + } + + def test_full_name_from_url_only_matches_mcp_service_urls(self): + assert ( + mcp._mcp_service_full_name_from_url(f"{WS}/ai-gateway/mcp-services/system.ai.github") + == "system.ai.github" + ) + # Other MCP routes are not UC services and must not be renamed. + assert mcp._mcp_service_full_name_from_url(f"{WS}/api/2.0/mcp/sql") is None + assert mcp._mcp_service_full_name_from_url(f"{WS}/ai-gateway/skills/?schema=a.b") is None + + def test_unique_leaf_wins(self): + servers = [self._service("system.ai.github"), self._service("system.ai.slack")] + mcp._apply_mcp_service_leaf_names(servers) + assert [s["name"] for s in servers] == ["github", "slack"] + # The full UC name stays in the URL so loading/round-trip is unaffected. + assert servers[0]["url"].endswith("/mcp-services/system.ai.github") + + def test_leaf_collision_across_schemas_keeps_full_dashed_path(self): + # Two services share the id `github` in different schemas: both keep the full path so + # each still registers under a distinct name. + servers = [self._service("system.ai.github"), self._service("main.tools.github")] + mcp._apply_mcp_service_leaf_names(servers) + assert sorted(s["name"] for s in servers) == ["main-tools-github", "system-ai-github"] + + def test_leaf_clashing_with_non_service_entry_keeps_full_path(self): + # An external connection literally named `github` blocks the service from taking the leaf. + servers = [ + {"name": "github", "url": f"{WS}/api/2.0/mcp/external/github", "clients": ["claude"]}, + self._service("system.ai.github"), + ] + mcp._apply_mcp_service_leaf_names(servers) + assert servers[0]["name"] == "github" # the external connection is untouched + assert servers[1]["name"] == "system-ai-github" + + def test_service_choice_shows_and_registers_the_leaf(self): + # Unconfigured: an add-choice whose title is the bare id and whose value carries the + # full UC name for the resolver. + add = mcp._mcp_service_choice("system.ai.github", set(), additive=False) + assert add.title == "MCP: github" + assert ( + add.value == f"{mcp.MCP_ADD_PREFIX}{mcp.MCP_SERVICE_SELECTION_PREFIX}system.ai.github" + ) + + # Already configured under the leaf: a pre-checked, removable toggle keyed by the leaf. + toggle = mcp._mcp_service_choice("system.ai.github", {"github"}, additive=False) + assert toggle.title == "MCP: github" + assert toggle.value == "github" + + # A legacy config still holding the full dashed name is matched too (back-compat). + legacy = mcp._mcp_service_choice("system.ai.github", {"system-ai-github"}, additive=False) + assert legacy.value == "system-ai-github" + + class TestApplySkillsMcpChanges: def _entry(self, by_client): return mcp._build_skills_entry(WS, by_client, list(by_client)) @@ -1359,17 +1423,17 @@ def fake_list(workspace, token, parent): assert seen == {"parent": "system.ai"} assert picker_called == [] - assert [c[1] for c in configured] == ["system-ai-github", "system-ai-slack"] + assert [c[1] for c in configured] == ["github", "slack"] assert configured[0][2] == f"{WS}/ai-gateway/mcp-services/system.ai.github" assert saved_states[-1]["mcp_servers"] == [ { - "name": "system-ai-github", + "name": "github", "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", "auth": "proxy", "clients": ["claude"], }, { - "name": "system-ai-slack", + "name": "slack", "url": f"{WS}/ai-gateway/mcp-services/system.ai.slack", "auth": "proxy", "clients": ["claude"], @@ -1410,10 +1474,10 @@ def test_replaces_servers_outside_location(self, monkeypatch): assert mcp.configure_mcp_command(location="system.ai") == 0 assert removed == [("claude", "databricks-sql")] - assert [c[1] for c in configured] == ["system-ai-github"] + assert [c[1] for c in configured] == ["github"] assert saved_states[-1]["mcp_servers"] == [ { - "name": "system-ai-github", + "name": "github", "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", "auth": "proxy", "clients": ["claude"], @@ -1461,7 +1525,7 @@ def test_existing_entry_gets_reconfigured_for_newly_added_clients(self, monkeypa saved_states: list[dict] = [] configured: list[tuple[str, str, str, dict]] = [] existing = { - "name": "system-ai-github", + "name": "github", "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", "auth": "proxy", "clients": ["claude"], @@ -1492,7 +1556,7 @@ def test_existing_entry_gets_reconfigured_for_newly_added_clients(self, monkeypa assert [c[0] for c in configured] == ["claude", "codex"] assert saved_states[-1]["mcp_servers"] == [ { - "name": "system-ai-github", + "name": "github", "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", "auth": "proxy", "clients": ["claude", "codex"], @@ -1500,6 +1564,144 @@ def test_existing_entry_gets_reconfigured_for_newly_added_clients(self, monkeypa ] +class TestConfigureAllMcpServices: + """`configure mcp --all` / onboarding: register every MCP service the user can access, + metastore-wide, with no picker.""" + + def _capture_messages(self, monkeypatch): + msgs: dict[str, list[str]] = {"success": [], "note": [], "warning": []} + monkeypatch.setattr(mcp, "print_success", lambda m: msgs["success"].append(m)) + monkeypatch.setattr(mcp, "print_note", lambda m: msgs["note"].append(m)) + monkeypatch.setattr(mcp, "print_warning", lambda m: msgs["warning"].append(m)) + return msgs + + def test_registers_every_accessible_service_under_leaf_names(self, monkeypatch): + saved_states: list[dict] = [] + configured: list[tuple[str, str, str]] = [] + picker_called: list[bool] = [] + _stub_location_base(monkeypatch, {**CLAUDE_STATE}) + # The whole workspace listing (permission-filtered server-side), across several schemas. + monkeypatch.setattr( + mcp, + "list_all_mcp_services", + lambda workspace, token, **kw: ( + ["main.default.sanjay_tavily", "system.ai.github", "users.someone.my_mcp"], + None, + ), + ) + monkeypatch.setattr( + mcp, + "prompt_for_mcp_server_choices", + lambda *a, **kw: picker_called.append(True) or [], + ) + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, name, url, *a, **kw: configured.append((client, name, url)) or [], + ) + monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) + + assert mcp.configure_mcp_command(all_services=True) == 0 + + assert picker_called == [] # no picker: it just registers everything + # Each is registered under its bare id (the `mcp____` tool prefix). + assert sorted(c[1] for c in configured) == ["github", "my_mcp", "sanjay_tavily"] + assert sorted(s["name"] for s in saved_states[-1]["mcp_servers"]) == [ + "github", + "my_mcp", + "sanjay_tavily", + ] + # The full UC name is preserved in each URL, so loading is unaffected. + by_name = {s["name"]: s for s in saved_states[-1]["mcp_servers"]} + assert by_name["sanjay_tavily"]["url"].endswith("/mcp-services/main.default.sanjay_tavily") + + def test_rejects_combination_with_location_or_services(self, monkeypatch): + _stub_location_base(monkeypatch, {**CLAUDE_STATE}) + for kwargs in ({"location": "main.default"}, {"services": {"main.default.x"}}): + try: + mcp.configure_mcp_command(all_services=True, **kwargs) + except RuntimeError as exc: + assert "--all" in str(exc) + else: + raise AssertionError(f"expected RuntimeError for all_services with {kwargs}") + + def test_scopes_to_configured_agents_not_every_installed_cli(self, monkeypatch): + # A Codex-only user (available_tools=["codex"]) with Cursor also installed: `--all` must + # target Codex only, not sweep in Cursor just because it's an installed MCP-only client. + configured: list[tuple[str, str]] = [] + _stub_location_base(monkeypatch, {"workspace": WS, "available_tools": ["codex"]}) + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["codex", "cursor"]) + monkeypatch.setattr( + mcp, + "list_all_mcp_services", + lambda workspace, token, **kw: (["system.ai.github"], None), + ) + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, name, url, *a, **kw: configured.append((client, name)) or [], + ) + monkeypatch.setattr(mcp, "save_state", lambda state: None) + + assert mcp.configure_mcp_command(all_services=True) == 0 + assert {c[0] for c in configured} == {"codex"} # cursor is not configured + + def test_rerun_when_already_registered_reports_count_not_none_found(self, monkeypatch): + # The bug: a no-op re-run (everything already registered) reported "none found". It must + # report what's registered instead, since the services are all there. + existing = { + "name": "github", + "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", + "auth": "proxy", + "clients": ["claude"], + } + _stub_location_base(monkeypatch, {**CLAUDE_STATE, "mcp_servers": [existing]}) + monkeypatch.setattr( + mcp, + "list_all_mcp_services", + lambda workspace, token, **kw: (["system.ai.github"], None), + ) + monkeypatch.setattr(mcp, "configure_client_mcp_server", lambda *a, **kw: []) + monkeypatch.setattr(mcp, "save_state", lambda state: None) + msgs = self._capture_messages(monkeypatch) + + assert mcp.configure_mcp_command(all_services=True) == 0 + assert any("1 MCP server" in m for m in msgs["success"]) + assert not any("No MCP servers" in m for m in msgs["note"]) + + def test_no_accessible_services_says_none_found(self, monkeypatch): + configured: list[tuple[str, str, str]] = [] + _stub_location_base(monkeypatch, {**CLAUDE_STATE}) + monkeypatch.setattr(mcp, "list_all_mcp_services", lambda workspace, token, **kw: ([], None)) + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, name, url, *a, **kw: configured.append((client, name, url)) or [], + ) + monkeypatch.setattr(mcp, "save_state", lambda state: None) + msgs = self._capture_messages(monkeypatch) + + assert mcp.configure_mcp_command(all_services=True) == 0 + assert configured == [] + assert any("No MCP servers" in m for m in msgs["note"]) + + def test_listing_failure_surfaces_the_reason(self, monkeypatch): + _stub_location_base(monkeypatch, {**CLAUDE_STATE}) + monkeypatch.setattr( + mcp, + "list_all_mcp_services", + lambda workspace, token, **kw: ([], "HTTP 403 Forbidden"), + ) + monkeypatch.setattr(mcp, "configure_client_mcp_server", lambda *a, **kw: []) + monkeypatch.setattr(mcp, "save_state", lambda state: None) + msgs = self._capture_messages(monkeypatch) + + assert mcp.configure_mcp_command(all_services=True) == 0 + # Not a bare "none found" — the actual failure is surfaced. + assert any("HTTP 403" in m for m in msgs["warning"]) + assert not any("No MCP servers" in m for m in msgs["note"]) + + class TestAddMcpCommand: """`ucode mcp add` (append) registers new servers without removing existing ones.""" @@ -1540,10 +1742,10 @@ def test_keeps_servers_outside_location(self, monkeypatch): # Nothing is removed; the new service is added and the outside one kept. assert removed == [] - assert [c[1] for c in configured] == ["system-ai-github"] + assert [c[1] for c in configured] == ["github"] assert saved_states[-1]["mcp_servers"] == [ { - "name": "system-ai-github", + "name": "github", "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", "auth": "proxy", "clients": ["claude"], @@ -1557,7 +1759,7 @@ def test_services_subset_keeps_others_in_location(self, monkeypatch): saved_states: list[dict] = [] removed: list[tuple[str, str]] = [] existing = { - "name": "system-ai-slack", + "name": "slack", "url": f"{WS}/ai-gateway/mcp-services/system.ai.slack", "auth": "proxy", "clients": ["claude"], @@ -1583,7 +1785,7 @@ def test_services_subset_keeps_others_in_location(self, monkeypatch): assert removed == [] names = [s["name"] for s in saved_states[-1]["mcp_servers"]] - assert names == ["system-ai-github", "system-ai-slack"] + assert names == ["github", "slack"] def test_empty_services_is_a_noop(self, monkeypatch): """`mcp add --services ""` has nothing to add, so it's a no-op that never @@ -1616,7 +1818,7 @@ def test_agents_scopes_registration_to_named_agent(self, monkeypatch): assert mcp.add_mcp_command(location="system.ai", agents={"claude"}) == 0 - assert configured == [("claude", "system-ai-github")] + assert configured == [("claude", "github")] assert saved_states[-1]["mcp_servers"][0]["clients"] == ["claude"] def test_agents_not_configured_raises(self, monkeypatch): @@ -1797,10 +1999,10 @@ def test_configures_only_the_requested_subset(self, monkeypatch): ) # slack is dropped; only the two requested services are configured. - assert sorted(c[1] for c in configured) == ["system-ai-github", "system-ai-gmail"] + assert sorted(c[1] for c in configured) == ["github", "gmail"] assert sorted(s["name"] for s in saved_states[-1]["mcp_servers"]) == [ - "system-ai-github", - "system-ai-gmail", + "github", + "gmail", ] def test_matches_bare_short_names(self, monkeypatch): @@ -1820,7 +2022,7 @@ def test_matches_bare_short_names(self, monkeypatch): assert mcp.configure_mcp_command(location="system.ai", services={"github"}) == 0 - assert [c[1] for c in configured] == ["system-ai-github"] + assert [c[1] for c in configured] == ["github"] def test_unknown_requested_service_warns_and_skips(self, monkeypatch): configured: list[tuple[str, str, str, dict]] = [] @@ -1847,7 +2049,7 @@ def test_unknown_requested_service_warns_and_skips(self, monkeypatch): ) # The known service is still configured; the unknown one is reported, not fatal. - assert [c[1] for c in configured] == ["system-ai-github"] + assert [c[1] for c in configured] == ["github"] assert any("system.ai.ghost" in w for w in warnings) def test_empty_services_removes_everything(self, monkeypatch): @@ -1888,13 +2090,13 @@ def test_adds_and_removes_to_match_new_selection(self, monkeypatch): # The live case teammates want mid-session: started with github+slack, # then the user deselects slack and selects gmail. github = { - "name": "system-ai-github", + "name": "github", "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", "auth": "proxy", "clients": ["claude"], } slack = { - "name": "system-ai-slack", + "name": "slack", "url": f"{WS}/ai-gateway/mcp-services/system.ai.slack", "auth": "proxy", "clients": ["claude"], @@ -1931,11 +2133,11 @@ def test_adds_and_removes_to_match_new_selection(self, monkeypatch): ) # slack removed, gmail added, github untouched (entry unchanged). - assert removed == [("claude", "system-ai-slack")] - assert [c[1] for c in configured] == ["system-ai-gmail"] + assert removed == [("claude", "slack")] + assert [c[1] for c in configured] == ["gmail"] assert sorted(s["name"] for s in saved_states[-1]["mcp_servers"]) == [ - "system-ai-github", - "system-ai-gmail", + "github", + "gmail", ] def test_full_names_without_location_derive_the_schema(self, monkeypatch): @@ -1959,7 +2161,7 @@ def fake_list(workspace, token, parent): assert mcp.configure_mcp_command(services={"system.ai.github", "system.ai.slack"}) == 0 assert seen == {"parent": "system.ai"} - assert sorted(c[1] for c in configured) == ["system-ai-github", "system-ai-slack"] + assert sorted(c[1] for c in configured) == ["github", "slack"] def test_short_name_without_location_raises(self): try: @@ -2137,7 +2339,7 @@ def test_replaces_scope_for_configured_clients_only(self, monkeypatch): def test_preserves_mcp_service_entries_across_set(self, monkeypatch): saved_states: list[dict] = [] service_entry = { - "name": "system-ai-github", + "name": "github", "url": f"{WS}/ai-gateway/mcp-services/system.ai.github", "auth": "env:OAUTH_TOKEN", "clients": ["claude"], @@ -2152,7 +2354,7 @@ def test_preserves_mcp_service_entries_across_set(self, monkeypatch): assert mcp.configure_skills_mcp_command(["B.b"]) == 0 names = [s["name"] for s in saved_states[-1]["mcp_servers"]] - assert "system-ai-github" in names + assert "github" in names assert names.count(mcp.SKILLS_MCP_SERVER_NAME) == 1