diff --git a/src/ucode/cli.py b/src/ucode/cli.py index a0ef7c4..daf2560 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -80,6 +80,7 @@ prompt_for_selection, prompt_for_tools, prompt_for_workspace, + prompt_yes_no, set_verbosity, spinner, status_badge, @@ -1167,6 +1168,9 @@ def configure( # target claude instead of dropping into the interactive agent picker. if enable_fable is not None and agent is None and agents is None: agent = "claude" + # Set True only in the fully-interactive branch below; gates the optional + # MCP setup prompt so flag-driven / scripted runs are never interrupted. + fully_interactive = False if agent is not None: tool = normalize_tool(agent) install_tool_binary( @@ -1212,6 +1216,10 @@ def configure( prompt_optional_updates=prompt_optional_updates, **skip_kwargs, ) + # Only the no-agent, no-workspace path is truly interactive (the user + # picked agents/workspace via prompts); that's where we offer the MCP + # step below. Flag-driven runs stay scriptable. + fully_interactive = workspace_entries is None if tracing: # The workspaces were just configured, so enable tracing for them # directly instead of re-prompting. Fall back to the workspace that @@ -1222,6 +1230,12 @@ def configure( tracing_workspaces = [(current, None)] if current else None if tracing_workspaces: configure_tracing_command(workspaces=tracing_workspaces) + # Offer MCP setup as the natural next step of interactive configuration, + # so users discover it without needing to know `configure mcp` exists. + # Skipped in dry-run and non-interactive/flag-driven runs (which stay + # scriptable), and when --dry-run is set. + if fully_interactive and not dry_run and prompt_yes_no("Configure MCP servers now?"): + configure_mcp_command() 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 b38bca9..c0894a6 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -15,6 +15,7 @@ import shutil import subprocess import time +from collections.abc import Callable from concurrent.futures import ( ThreadPoolExecutor, as_completed, @@ -242,6 +243,12 @@ def _http_get_json( except urllib_error.URLError as exc: _debug(f"GET {url}", f"URLError: {exc.reason}") return None, f"network error: {exc.reason}" + except OSError as exc: + # A socket read timeout raises a bare TimeoutError (an OSError), not a + # URLError, so it must be caught explicitly or it escapes the whole + # discovery flow. Surface it as a reason like every other failure. + _debug(f"GET {url}", f"OSError: {exc}") + return None, f"network error: {exc}" def _http_post_json( @@ -287,6 +294,11 @@ def _http_post_json( except urllib_error.URLError as exc: _debug(f"POST {url}", f"URLError: {exc.reason}") return None, f"network error: {exc.reason}" + except OSError as exc: + # See `_http_get_json`: a bare socket timeout is an OSError, not a + # URLError, and would otherwise escape the caller's error handling. + _debug(f"POST {url}", f"OSError: {exc}") + return None, f"network error: {exc}" def get_current_user_name(workspace: str, token: str) -> str | None: @@ -1542,6 +1554,9 @@ def map_bedrock_claude_models(targets: list[str]) -> dict[str, str]: _UC_FUNCTION_PROBE_TIMEOUT = 5 _VECTOR_SEARCH_DEADLINE_SECONDS = 15.0 _UC_FUNCTIONS_DEADLINE_SECONDS = 20.0 +# This walk runs on every `configure mcp` (not opt-in), so keep the budget tight: +# a slow workspace degrades to partial/instant results instead of a long wait. +_MCP_SERVICES_WALK_DEADLINE_SECONDS = 8.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( @@ -1630,11 +1645,16 @@ def list_vector_search_catalog_schemas( token: str, *, deadline_seconds: float = _VECTOR_SEARCH_DEADLINE_SECONDS, + on_progress: Callable[[int, int, int], None] | None = None, ) -> tuple[list[tuple[str, str]], str | None]: """Return sorted unique `(catalog, schema)` pairs that contain at least one Databricks Vector Search index. Walks the per-endpoint index listings in parallel under a wall-clock budget; returns partial results once - `deadline_seconds` is exceeded.""" + `deadline_seconds` is exceeded. + + `on_progress`, if given, is called as each endpoint's listing completes with + `(endpoints_done, endpoints_total, pairs_found)` for live count reporting. + It is invoked serially from the draining thread (not the workers).""" hostname = workspace_hostname(workspace) deadline = time.monotonic() + deadline_seconds endpoints, reason = _paginated_json_items( @@ -1651,7 +1671,9 @@ def list_vector_search_catalog_schemas( return [], "no vector search endpoints with names" pairs: set[tuple[str, str]] = set() - workers = max(1, min(_UC_FUNCTION_PROBE_WORKERS, len(endpoint_names))) + endpoints_total = len(endpoint_names) + endpoints_done = 0 + workers = max(1, min(_UC_FUNCTION_PROBE_WORKERS, endpoints_total)) with ThreadPoolExecutor(max_workers=workers) as pool: futures = { pool.submit( @@ -1666,11 +1688,15 @@ def list_vector_search_catalog_schemas( } def collect(result, _endpoint): + nonlocal endpoints_done indexes, _ = result for index in indexes: pair = _vector_index_catalog_schema(index) if pair: pairs.add(pair) + endpoints_done += 1 + if on_progress is not None: + on_progress(endpoints_done, endpoints_total, len(pairs)) _drain_with_deadline(futures, deadline, collect) pool.shutdown(wait=False, cancel_futures=True) @@ -1698,9 +1724,14 @@ def list_uc_functions_catalog_schemas( token: str, *, deadline_seconds: float = _UC_FUNCTIONS_DEADLINE_SECONDS, + on_progress: Callable[[int, int, int], None] | None = None, ) -> tuple[list[tuple[str, str]], str | None]: """Return sorted unique `(catalog, schema)` pairs containing at least one - user-defined UC function.""" + user-defined UC function. + + `on_progress`, if given, is called during the function-probe phase with + `(schemas_done, schemas_total, pairs_found)` for live count reporting. It is + invoked serially from the draining thread (not the workers).""" hostname = workspace_hostname(workspace) deadline = time.monotonic() + deadline_seconds @@ -1764,6 +1795,8 @@ def collect_schemas(result, catalog): # Parallel function-existence probes. pairs: set[tuple[str, str]] = set() + schemas_total = len(candidate_pairs) + schemas_done = 0 with ThreadPoolExecutor(max_workers=_UC_FUNCTION_PROBE_WORKERS) as pool: probe_futures = { pool.submit(_schema_has_user_function, hostname, token, cat, schema): (cat, schema) @@ -1771,8 +1804,12 @@ def collect_schemas(result, catalog): } def collect_pair(has_fn, pair): + nonlocal schemas_done if has_fn: pairs.add(pair) + schemas_done += 1 + if on_progress is not None: + on_progress(schemas_done, schemas_total, len(pairs)) _drain_with_deadline(probe_futures, deadline, collect_pair) pool.shutdown(wait=False, cancel_futures=True) @@ -1784,6 +1821,111 @@ def collect_pair(has_fn, pair): return sorted(pairs), None +def list_all_mcp_services( + workspace: str, + token: str, + *, + deadline_seconds: float = _MCP_SERVICES_WALK_DEADLINE_SECONDS, + on_progress: Callable[[int, int, int], 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. It is invoked serially from the draining thread (not the workers). + + This walk is the slow, 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", + token, + items_key="catalogs", + timeout=_UC_LIST_HTTP_TIMEOUT, + ) + 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 + names.update(found) + schemas_done += 1 + if on_progress is not None: + on_progress(schemas_done, schemas_total, len(names)) + + _drain_with_deadline(service_futures, deadline, collect_services) + pool.shutdown(wait=False, cancel_futures=True) + + if not names: + if time.monotonic() > deadline: + return [], "deadline exceeded while listing MCP services" + return [], "no MCP services found" + return sorted(names), None + + def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str], str | None]: """Discover Claude families on this workspace's AI Gateway. diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 6257458..c0b48b0 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -7,7 +7,9 @@ import shutil import string import subprocess +import threading from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any from urllib.parse import urlparse @@ -31,6 +33,7 @@ build_mcp_service_url, ensure_databricks_auth, get_databricks_token, + list_all_mcp_services, list_databricks_apps, list_databricks_connections, list_genie_spaces, @@ -52,6 +55,15 @@ MCP_USER_SCOPE = "user" MCP_CLEANUP_SCOPES = ("local", "project", MCP_USER_SCOPE) MCP_PICKER_VISIBLE_ROWS = 10 + + +class _Back: + """Sentinel type: a wizard step returns the `_BACK` instance when the user + presses Left (←) to go back. Distinct from None (cancel) and [] (empty).""" + + +# Singleton instance used everywhere; compare with `is _BACK`. +_BACK = _Back() MCP_CLIENTS = { "claude": { "binary": "claude", @@ -381,6 +393,20 @@ def discover_mcp_service_names(workspace: str, profile: str | None = None) -> li return names +def discover_all_mcp_service_names( + workspace: str, + profile: str | None = None, + on_progress: Callable[[int, int, int], None] | None = None, +) -> list[str]: + """All MCP services across every `.` in the workspace. This + walks the workspace (see `list_all_mcp_services`) and is the workspace-wide + counterpart to `discover_mcp_service_names`. `on_progress` is forwarded to + the walk for live count reporting.""" + token = get_databricks_token(workspace, profile) + names, _reason = list_all_mcp_services(workspace, token, on_progress=on_progress) + return names + + def _normalize_workspace_title(text: str) -> str: """Collapse a Databricks workspace title to lowercase alphanumerics joined by single hyphens, trimmed at the edges. Output is safe to use as an MCP @@ -500,9 +526,13 @@ def vector_search_mcp_servers(pairs: list[tuple[str, str]], workspace: str) -> l return sorted(servers, key=lambda server: str(server["title"]).lower()) -def discover_vector_search_mcp_servers(workspace: str, profile: str | None = None) -> list[dict]: +def discover_vector_search_mcp_servers( + workspace: str, + profile: str | None = None, + on_progress: Callable[[int, int, int], None] | None = None, +) -> list[dict]: token = get_databricks_token(workspace, profile) - pairs, _reason = list_vector_search_catalog_schemas(workspace, token) + pairs, _reason = list_vector_search_catalog_schemas(workspace, token, on_progress=on_progress) return vector_search_mcp_servers(pairs, workspace) @@ -526,9 +556,13 @@ def uc_functions_mcp_servers(pairs: list[tuple[str, str]], workspace: str) -> li return sorted(servers, key=lambda server: str(server["title"]).lower()) -def discover_uc_functions_mcp_servers(workspace: str, profile: str | None = None) -> list[dict]: +def discover_uc_functions_mcp_servers( + workspace: str, + profile: str | None = None, + on_progress: Callable[[int, int, int], None] | None = None, +) -> list[dict]: token = get_databricks_token(workspace, profile) - pairs, _reason = list_uc_functions_catalog_schemas(workspace, token) + pairs, _reason = list_uc_functions_catalog_schemas(workspace, token, on_progress=on_progress) return uc_functions_mcp_servers(pairs, workspace) @@ -629,6 +663,7 @@ def _scrolling_checkbox( choices: list[questionary.Choice | questionary.Separator], instruction: str, style: questionary.Style, + allow_back: bool = False, ) -> Question: merged_style = merge_styles_default( [ @@ -659,7 +694,6 @@ def perform_validation() -> bool: control.error_message = None return True - prompt_session: PromptSession = PromptSession(get_prompt_tokens, reserve_space_for_menu=0) visible_rows = min(MCP_PICKER_VISIBLE_ROWS, max(1, len(choices))) has_more_choices = len(choices) > MCP_PICKER_VISIBLE_ROWS @@ -668,10 +702,16 @@ def has_search_string() -> bool: return control.get_search_string_tokens() is not None validation_prompt: PromptSession = PromptSession(bottom_toolbar=lambda: control.error_message) + # Render the prompt as a fixed 1-row window rather than a PromptSession + # container: the latter expands to fill the terminal height, which in a tall + # window pushes the choices list to the very bottom (a large blank gap). layout = Layout( HSplit( [ - prompt_session.layout.container, + Window( + height=Dimension.exact(1), + content=FormattedTextControl(get_prompt_tokens), + ), ConditionalContainer( Window(control, height=Dimension(preferred=visible_rows, max=visible_rows)), filter=~IsDone(), @@ -716,6 +756,21 @@ def _(_event: Any) -> None: control.selected_options.append(pointed_choice) perform_validation() + @bindings.add(Keys.ControlA, eager=True) + def _(_event: Any) -> None: + # Toggle-all: select every selectable choice, or clear the selection if + # everything is already selected. `a` alone is reserved for type-to-filter. + selectable = [ + choice.value + for choice in control.choices + if not isinstance(choice, questionary.Separator) and not choice.disabled + ] + if all(value in control.selected_options for value in selectable): + control.selected_options = [] + else: + control.selected_options = list(selectable) + perform_validation() + def move_cursor_down(event: Any) -> None: control.select_next() while not control.is_selection_valid(): @@ -747,6 +802,15 @@ def _(event: Any) -> None: control.is_answered = True event.app.exit(result=get_selected_values()) + if allow_back: + + @bindings.add(Keys.Left, eager=True) + def _(event: Any) -> None: + # Wizard back-navigation: exit this step with the _BACK sentinel so + # the caller re-shows the previous step. Left arrow is otherwise + # unused in this multi-select (cursor moves with up/down). + event.app.exit(result=_BACK) + @bindings.add(Keys.Any) def _(_event: Any) -> None: """Ignore other text input.""" @@ -786,17 +850,19 @@ def build_mcp_picker_choices( # (see resolver). Compare against the dashed form when checking what's # already registered. registered_as = name.replace(".", "-") + display_title = f"MCP: {name}" if registered_as in known_names: - choices.append(_server_choice(registered_as, True, name)) + choices.append(_server_choice(registered_as, True, display_title)) else: - choices.append(_add_choice(f"{MCP_SERVICE_SELECTION_PREFIX}{name}", name)) + choices.append(_add_choice(f"{MCP_SERVICE_SELECTION_PREFIX}{name}", display_title)) displayed_names.add(registered_as) for name in available_external_names: + display_title = f"Connection: {name}" if name in known_names: - choices.append(_server_choice(name, True, name)) + choices.append(_server_choice(name, True, display_title)) else: - choices.append(_add_choice(f"{EXTERNAL_MCP_SELECTION_PREFIX}{name}", name)) + choices.append(_add_choice(f"{EXTERNAL_MCP_SELECTION_PREFIX}{name}", display_title)) displayed_names.add(name) for server in available_genie_servers: @@ -882,7 +948,14 @@ def prompt_for_mcp_server_choices( available_mcp_service_names: list[str] | None = None, available_vector_search_servers: list[dict] | None = None, available_uc_functions_servers: list[dict] | None = None, -) -> list[str] | None: + allow_back: bool = False, +) -> list[str] | None | _Back: + """Show the MCP server picker. Returns the list of selected values, `None` + if cancelled (Ctrl-C), or `_BACK` if `allow_back` and the user pressed Left + to return to the previous wizard step.""" + instruction = "(space to toggle, ctrl-a all, enter to save, type to filter)" + if allow_back: + instruction = "(space to toggle, ctrl-a all, ← back, enter to save, type to filter)" selection = _scrolling_checkbox( "MCP:", choices=build_mcp_picker_choices( @@ -895,10 +968,13 @@ def prompt_for_mcp_server_choices( available_uc_functions_servers, ), style=_picker_style(), - instruction="(space to toggle, enter to save, type to filter)", + instruction=instruction, + allow_back=allow_back, ).ask() if selection is None: return None + if selection is _BACK: + return _BACK return [str(value) for value in selection] @@ -1010,11 +1086,123 @@ def _discover_mcp_source(label: str, discover: Callable[[], list[Any]]) -> list[ try: with spinner(f"Discovering {label}..."): return discover() - except RuntimeError: - print_warning(f"Skipped {label}.") + except (RuntimeError, OSError) as exc: + # Discovery is best-effort: a failure here (auth error, network timeout) + # skips just this source so the rest of the picker still works. + print_warning(f"Skipped {label} ({exc}).") return [] +def _discover_mcp_source_with_progress( + label: str, + unit: str, + discover: Callable[[Callable[[int, int, int], None]], list[Any]], +) -> list[Any]: + """Run a walk-based discovery behind a spinner whose message shows a live + count (e.g. `Searching Vector Search... 3/8 endpoints, 2 found`). `discover` + receives an `on_progress(done, total, found)` callback and `unit` names what + is being counted. Best-effort like `_discover_mcp_source`: any failure is + warned and yields an empty list.""" + progress = {"done": 0, "total": 0, "found": 0} + + def on_progress(done: int, total: int, found: int) -> None: + progress.update(done=done, total=total, found=found) + + def message() -> str: + if progress["total"]: + return ( + f"Searching {label}... {progress['done']}/{progress['total']} {unit}, " + f"{progress['found']} found" + ) + return f"Searching {label}..." + + try: + with spinner(message): + return discover(on_progress) + except (RuntimeError, OSError) as exc: + print_warning(f"Skipped {label} ({exc}).") + return [] + + +def _discover_selected_mcp_sources( + workspace: str, profile: str | None, sources: set[str] +) -> dict[str, list]: + """Run discovery for the sources the user selected on the search screen. + Returns a dict keyed by picker argument (external/apps/services/genie/ + vector_search/uc_functions); unselected sources yield empty lists so the + picker still renders (and can still remove already-registered servers).""" + external = ( + _discover_mcp_source( + "external connections", + lambda: discover_external_mcp_connection_names(workspace, profile), + ) + if "external" in sources + else [] + ) + apps = ( + _discover_mcp_source( + "Databricks apps", + lambda: discover_app_mcp_servers(workspace, profile), + ) + if "apps" in sources + else [] + ) + # MCP services: curated `system.ai` list plus the workspace-wide walk, + # merged and de-duplicated (the walk skips the `system` catalog). + services: list[str] = [] + if "mcp-services" in sources: + curated = _discover_mcp_source( + "MCP services", + lambda: discover_mcp_service_names(workspace, profile), + ) + walked = _discover_mcp_source_with_progress( + "all MCP services", + "schemas", + lambda on_progress: discover_all_mcp_service_names( + workspace, profile, on_progress=on_progress + ), + ) + services = list(dict.fromkeys(curated + walked)) + genie = ( + _discover_mcp_source( + "Genie spaces", + lambda: discover_genie_mcp_servers(workspace, profile), + ) + if "genie" in sources + else [] + ) + vector_search = ( + _discover_mcp_source_with_progress( + "Vector Search", + "endpoints", + lambda on_progress: discover_vector_search_mcp_servers( + workspace, profile, on_progress=on_progress + ), + ) + if "vector-search" in sources + else [] + ) + uc_functions = ( + _discover_mcp_source_with_progress( + "UC functions", + "schemas", + lambda on_progress: discover_uc_functions_mcp_servers( + workspace, profile, on_progress=on_progress + ), + ) + if "uc-functions" in sources + else [] + ) + return { + "external": external, + "apps": apps, + "services": services, + "genie": genie, + "vector_search": vector_search, + "uc_functions": uc_functions, + } + + def apply_mcp_server_changes( original_servers: list[dict], working_servers: list[dict], @@ -1022,12 +1210,22 @@ def apply_mcp_server_changes( ) -> bool: original_by_name = _servers_by_name(original_servers) working_by_name = _servers_by_name(working_servers) + + # Build the per-client work lists. Each add/remove shells out to a CLI or + # rewrites a config file, so a large diff means hundreds of operations; we + # run them concurrently ACROSS clients but SERIALLY within a client, since + # every operation for one client mutates that client's single shared config + # (`claude mcp add-json` edits ~/.claude.json, etc.) and concurrent + # read-modify-writes would clobber each other. + work: dict[str, list[Callable[[], object]]] = {client: [] for client in clients} changed = False for name, server in original_by_name.items(): if name not in working_by_name: for client in _mcp_server_clients(server): - remove_client_mcp_server(client, name) + work.setdefault(client, []).append( + lambda c=client, n=name: remove_client_mcp_server(c, n) + ) changed = True for name, server in working_by_name.items(): @@ -1039,12 +1237,51 @@ def apply_mcp_server_changes( continue entry = build_mcp_http_entry(url) for client in clients: - configure_client_mcp_server(client, name, url, entry) + work[client].append( + lambda c=client, n=name, u=url, e=entry: configure_client_mcp_server(c, n, u, e) + ) changed = True + total_ops = sum(len(ops) for ops in work.values()) + if total_ops == 0: + return changed + + completed = _Counter() + + def run_client_ops(ops: list[Callable[[], object]]) -> None: + for op in ops: + op() + completed.increment() + + def message() -> str: + return f"Configuring MCP servers... {completed.value()}/{total_ops}" + + with spinner(message): + with ThreadPoolExecutor(max_workers=max(1, len(work))) as pool: + futures = [pool.submit(run_client_ops, ops) for ops in work.values() if ops] + # Surface the first failure (if any) once all client threads finish. + for future in as_completed(futures): + future.result() + return changed +class _Counter: + """Thread-safe monotonic counter for cross-thread progress reporting.""" + + def __init__(self) -> None: + self._value = 0 + self._lock = threading.Lock() + + def increment(self) -> None: + with self._lock: + self._value += 1 + + def value(self) -> int: + with self._lock: + return self._value + + def purge_cross_workspace_mcp_residue(state: dict, workspace: str) -> None: installed = set(available_mcp_clients()) @@ -1177,6 +1414,38 @@ def _resolve_location_mcp_servers( return working_servers +# The first wizard step lets the user choose which sources to search. Each is a +# (key, label, default_checked) triple. Vector Search and UC functions default +# off because they walk the workspace (endpoints/catalogs/schemas) and are slow; +# everything else is a cheap listing and defaults on. +MCP_SEARCH_SOURCES = ( + ("external", "External connections", True), + ("apps", "Databricks apps", True), + ("mcp-services", "MCP services", True), + ("genie", "Genie spaces", True), + ("vector-search", "Vector Search indexes (slower)", False), + ("uc-functions", "UC functions (slower)", False), +) + + +def prompt_for_mcp_search_sources() -> 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).""" + choices = [ + questionary.Choice(title=label, value=key, checked=checked) + for key, label, checked in MCP_SEARCH_SOURCES + ] + selection = _scrolling_checkbox( + "Search for:", + choices=choices, + style=_picker_style(), + instruction="(space to toggle, ctrl-a all, enter to search)", + ).ask() + if selection is None: + return None + return {str(value) for value in selection} + + 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 @@ -1246,48 +1515,38 @@ def configure_mcp_command(location: str | None = None, services: set[str] | None print_success("Saved") return 0 - available_external_mcp_names = _discover_mcp_source( - "external connections", - lambda: discover_external_mcp_connection_names(workspace, profile), - ) - available_genie_mcp_servers = _discover_mcp_source( - "Genie spaces", - lambda: discover_genie_mcp_servers(workspace, profile), - ) - available_app_mcp_servers = _discover_mcp_source( - "Databricks apps", - lambda: discover_app_mcp_servers(workspace, profile), - ) - # Curated `system.ai.*` MCP services live behind a separate UC API. Like - # the other sources this is best-effort — `_discover_mcp_source` swallows - # failures and returns [] so workspaces without them just see nothing extra. - available_mcp_service_names = _discover_mcp_source( - "MCP services", - lambda: discover_mcp_service_names(workspace, profile), - ) - # Per-(catalog, schema) managed MCP servers (Vector Search + UC Functions). - available_vector_search_servers = _discover_mcp_source( - "Vector Search", - lambda: discover_vector_search_mcp_servers(workspace, profile), - ) - available_uc_functions_servers = _discover_mcp_source( - "UC functions", - lambda: discover_uc_functions_mcp_servers(workspace, profile), - ) - original_mcp_servers: list[dict] = list(state.get("mcp_servers") or []) original_by_name = _servers_by_name(original_mcp_servers) - selections = prompt_for_mcp_server_choices( - available_external_mcp_names, - available_genie_mcp_servers, - available_app_mcp_servers, - original_mcp_servers, - available_mcp_service_names, - available_vector_search_servers, - available_uc_functions_servers, - ) - if selections is None: - return 0 + + # 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"], + original_mcp_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] = [] working_names: set[str] = set() @@ -1330,8 +1589,28 @@ def configure_mcp_command(location: str | None = None, services: set[str] | None if changed or original_mcp_servers != working_mcp_servers: state["mcp_servers"] = working_mcp_servers save_state(state) - print_success("Saved") + 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: # 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 + + +def _mcp_change_summary(added: list[str], removed: list[str], clients: list[str]) -> str: + """Human-readable one-liner describing what `configure mcp` just saved, e.g. + `Added 2, removed 1 MCP server across Claude Code, Codex`. Falls back to a + plain `Saved` when only client bindings changed (no add/remove).""" + client_names = ", ".join(str(MCP_CLIENTS[c]["display"]) for c in clients if c in MCP_CLIENTS) + parts: list[str] = [] + if added: + parts.append(f"added {len(added)}") + if removed: + parts.append(f"removed {len(removed)}") + if not parts: + return "Saved" + total = len(added) + len(removed) + noun = "MCP server" if total == 1 else "MCP servers" + summary = ", ".join(parts).capitalize() + return f"{summary} {noun} across {client_names}" if client_names else f"{summary} {noun}" diff --git a/src/ucode/ui.py b/src/ucode/ui.py index cb55797..7900114 100644 --- a/src/ucode/ui.py +++ b/src/ucode/ui.py @@ -7,6 +7,7 @@ import textwrap import threading import time +from collections.abc import Callable from contextlib import contextmanager from datetime import timedelta @@ -87,21 +88,34 @@ def status_badge(text: str, kind: str) -> str: @contextmanager -def spinner(message: str): +def spinner(message: str | Callable[[], str]): + """Show a spinner while the block runs. `message` may be a callable, which + is re-evaluated on every frame so callers can render live progress (e.g. a + running count) during a long operation.""" if not sys.stdout.isatty(): yield return + if isinstance(message, str): + static_message = message + + def current_message() -> str: + return static_message + else: + current_message = message + stop_event = threading.Event() def spin() -> None: for frame in itertools.cycle("|/-\\"): if stop_event.is_set(): break - sys.stdout.write(f"\r\033[2m{frame}\033[0m {message}") + # `\033[K` erases to end of line so a shrinking dynamic message + # doesn't leave stale characters behind. + sys.stdout.write(f"\r\033[2m{frame}\033[0m {current_message()}\033[K") sys.stdout.flush() time.sleep(0.1) - sys.stdout.write("\r" + " " * (len(message) + 4) + "\r") + sys.stdout.write("\r\033[K") sys.stdout.flush() thread = threading.Thread(target=spin, daemon=True) diff --git a/tests/test_cli.py b/tests/test_cli.py index 76a2ec5..b6ecb10 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -493,10 +493,41 @@ def test_no_flag_calls_configure_all(self): patch("ucode.cli.install_databricks_cli"), patch("ucode.cli.install_tool_binary"), patch("ucode.cli.configure_workspace_command") as mock_cfg, + # Fully-interactive configure ends by offering the MCP step; decline it. + patch("ucode.cli.prompt_yes_no", return_value=False) as mock_mcp_prompt, + patch("ucode.cli.configure_mcp_command") as mock_mcp, ): result = runner.invoke(app, ["configure"]) assert result.exit_code == 0, result.output mock_cfg.assert_called_once_with(prompt_optional_updates=True) + mock_mcp_prompt.assert_called_once() + mock_mcp.assert_not_called() + + def test_interactive_accepting_mcp_prompt_runs_mcp_config(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.install_tool_binary"), + patch("ucode.cli.configure_workspace_command"), + patch("ucode.cli.prompt_yes_no", return_value=True), + patch("ucode.cli.configure_mcp_command") as mock_mcp, + ): + result = runner.invoke(app, ["configure"]) + assert result.exit_code == 0, result.output + mock_mcp.assert_called_once_with() + + def test_agents_flag_skips_mcp_prompt(self): + # Flag-driven (non-interactive) runs must stay scriptable: no MCP prompt. + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.install_tool_binary"), + patch("ucode.cli.configure_workspace_command"), + patch("ucode.cli.prompt_yes_no") as mock_prompt, + patch("ucode.cli.configure_mcp_command") as mock_mcp, + ): + result = runner.invoke(app, ["configure", "--agents", "claude,codex"]) + assert result.exit_code == 0, result.output + mock_prompt.assert_not_called() + mock_mcp.assert_not_called() def test_agents_flag_calls_configure_with_tools(self): with ( @@ -643,6 +674,9 @@ def test_skip_upgrade_flag_disables_optional_update_prompt(self): patch("ucode.cli.install_databricks_cli"), patch("ucode.cli.install_tool_binary"), patch("ucode.cli.configure_workspace_command") as mock_cfg, + # Fully-interactive configure ends by offering the MCP step; decline it. + patch("ucode.cli.prompt_yes_no", return_value=False), + patch("ucode.cli.configure_mcp_command"), ): result = runner.invoke(app, ["configure", "--skip-upgrade"]) assert result.exit_code == 0, result.output diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 40af048..bcc8dad 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -676,6 +676,109 @@ def test_http_404_reason_surfaces_for_invalid_parent(self, monkeypatch): assert reason and reason.startswith("HTTP 404") +class TestListAllMcpServices: + """Workspace-wide walk: catalogs -> schemas -> per-schema mcp-services.""" + + def _fake_http(self, catalogs, schemas_by_catalog, services_by_schema): + """Route `_http_get_json` by URL to the right stubbed payload.""" + + 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" + + return fake_get + + def test_aggregates_services_across_catalogs_and_schemas(self, monkeypatch): + 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"], + }, + ), + ) + + names, reason = db_mod.list_all_mcp_services(WS, "token") + + assert reason is None + # information_schema is skipped; results are sorted and de-duplicated. + assert names == [ + "mycat.myschema.news", + "mycat.myschema.weather", + "other.ops.pager", + ] + + def test_reports_progress_per_schema(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"]}, + ), + ) + progress: list[tuple[int, int, int]] = [] + + names, reason = db_mod.list_all_mcp_services( + WS, + "token", + 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): + monkeypatch.setattr( + db_mod, + "_http_get_json", + self._fake_http( + catalogs=["system", "hive_metastore", "samples", "__databricks_internal"], + schemas_by_catalog={}, + services_by_schema={}, + ), + ) + + names, reason = db_mod.list_all_mcp_services(WS, "token") + + assert names == [] + assert reason == "no user UC catalogs found" + + def test_returns_reason_when_no_catalogs(self, monkeypatch): + monkeypatch.setattr( + db_mod, "_http_get_json", lambda url, token, timeout=30: ({"catalogs": []}, None) + ) + + names, reason = db_mod.list_all_mcp_services(WS, "token") + + assert names == [] + assert reason == "no UC catalogs found" + + def _foundation_models_payload(names): return { "endpoints": [ @@ -1699,3 +1802,33 @@ def test_falls_through_for_unrelated_permission_error(self, monkeypatch): db_mod.run_usage_query(WS, "/sql/1.0/warehouses/abc", "tok", "SELECT 1") assert "Ask your workspace admin" not in str(exc_info.value) assert str(exc_info.value).startswith("Usage query failed:") + + +class TestHttpGetJsonTimeout: + """A socket read timeout raises a bare TimeoutError (an OSError), not a + URLError. It must be returned as a reason, not propagated — otherwise it + escapes the best-effort MCP discovery flow and crashes the command.""" + + def test_read_timeout_returns_reason_instead_of_raising(self, monkeypatch): + def raise_timeout(request, timeout=None): + raise TimeoutError("The read operation timed out") + + monkeypatch.setattr(db_mod.urllib_request, "urlopen", raise_timeout) + + payload, reason = db_mod._http_get_json(f"{WS}/api/2.0/anything", "tok") + + assert payload is None + assert reason is not None + assert "timed out" in reason + + def test_post_read_timeout_returns_reason_instead_of_raising(self, monkeypatch): + def raise_timeout(request, timeout=None): + raise TimeoutError("The read operation timed out") + + monkeypatch.setattr(db_mod.urllib_request, "urlopen", raise_timeout) + + payload, reason = db_mod._http_post_json(f"{WS}/api/2.0/anything", "tok", {"k": "v"}) + + assert payload is None + assert reason is not None + assert "timed out" in reason diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 81333c8..24fbdb3 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -4,8 +4,11 @@ import json import subprocess +import threading from unittest.mock import MagicMock +import pytest + from ucode import mcp WS = "https://example.databricks.com" @@ -13,6 +16,23 @@ ALL_MCP_CLIENTS = ["claude", "codex", "gemini", "opencode", "copilot"] +class TestMcpChangeSummary: + def test_added_and_removed_with_clients(self): + summary = mcp._mcp_change_summary(["a", "b"], ["c"], ["claude", "codex"]) + assert summary == "Added 2, removed 1 MCP servers across Claude Code, Codex" + + def test_single_add_uses_singular_noun(self): + summary = mcp._mcp_change_summary(["a"], [], ["claude"]) + assert summary == "Added 1 MCP server across Claude Code" + + def test_only_removed(self): + summary = mcp._mcp_change_summary([], ["a"], ["claude"]) + assert summary == "Removed 1 MCP server across Claude Code" + + def test_no_changes_falls_back_to_saved(self): + assert mcp._mcp_change_summary([], [], ["claude"]) == "Saved" + + class TestBuildMcpHttpEntry: def test_uses_http_url(self): entry = mcp.build_mcp_http_entry(f"{WS}/api/2.0/mcp/external/github") @@ -249,11 +269,11 @@ def fake_checkbox(*args, **kwargs): assert "Custom servers" not in choice_text assert choice_text == [ "Databricks SQL", - "github-mcp", + "Connection: github-mcp", ] assert "Built-in AI tools" not in choice_text assert checkbox_calls[0]["kwargs"]["instruction"] == ( - "(space to toggle, enter to save, type to filter)" + "(space to toggle, ctrl-a all, enter to save, type to filter)" ) def test_prompt_returns_none_when_cancelled(self, monkeypatch): @@ -273,7 +293,7 @@ def test_picker_marks_configured_servers(self): [{"name": "github-mcp", "url": f"{WS}/api/2.0/mcp/external/github-mcp"}], ) choices_by_title = {choice.title: choice for choice in choices} - assert choices_by_title["github-mcp"].checked is True + assert choices_by_title["Connection: github-mcp"].checked is True assert choices_by_title["Databricks SQL"].checked is False def test_picker_keeps_databricks_sql_when_nothing_discovered(self): @@ -467,23 +487,156 @@ def test_picker_keeps_saved_legacy_servers_for_removal(self): assert choices_by_title["databricks-vector-search-main-search-docs"].checked is True -def _patch_mcp_choices(monkeypatch, *values: str) -> None: +def _patch_mcp_choices(monkeypatch, *values: str, categories: set[str] | None = None) -> None: monkeypatch.setattr( mcp, "prompt_for_mcp_server_choices", lambda *args, **kwargs: list(values), ) + # The first wizard step chooses which sources to search. Default to the + # fast pre-checked ones (external, apps, MCP services, genie); tests that + # exercise the slow walks (vector-search / uc-functions) pass those keys via + # `categories`, which are unioned in. + default_sources = {"external", "apps", "mcp-services", "genie"} + selected_sources = default_sources | (categories or set()) + monkeypatch.setattr(mcp, "prompt_for_mcp_search_sources", lambda: selected_sources) # Stub the always-on discoveries so configure_mcp_command tests don't hit # real APIs. Individual tests override these after calling the helper. monkeypatch.setattr(mcp, "discover_mcp_service_names", lambda workspace, profile=None: []) monkeypatch.setattr( - mcp, "discover_vector_search_mcp_servers", lambda workspace, profile=None: [] + mcp, + "discover_all_mcp_service_names", + lambda workspace, profile=None, on_progress=None: [], + ) + monkeypatch.setattr( + mcp, + "discover_vector_search_mcp_servers", + lambda workspace, profile=None, on_progress=None: [], ) monkeypatch.setattr( - mcp, "discover_uc_functions_mcp_servers", lambda workspace, profile=None: [] + mcp, + "discover_uc_functions_mcp_servers", + lambda workspace, profile=None, on_progress=None: [], ) +class TestApplyMcpServerChanges: + def _server(self, name, clients): + return {"name": name, "url": f"{WS}/api/2.0/mcp/external/{name}", "clients": clients} + + def test_adds_across_clients_serial_within_client(self, monkeypatch): + # Record (client, name) for every add. Each client's ops must stay in + # order even though clients run concurrently. + recorded: list[tuple[str, str]] = [] + lock = threading.Lock() + + def fake_configure(client, name, url, entry): + with lock: + recorded.append((client, name)) + return [] + + monkeypatch.setattr(mcp, "configure_client_mcp_server", fake_configure) + + working = [ + self._server("a", ["claude", "codex"]), + self._server("b", ["claude", "codex"]), + self._server("c", ["claude", "codex"]), + ] + + changed = mcp.apply_mcp_server_changes([], working, ["claude", "codex"]) + + assert changed is True + # 3 servers x 2 clients = 6 operations. + assert len(recorded) == 6 + # Within each client, the server order is preserved. + assert [n for c, n in recorded if c == "claude"] == ["a", "b", "c"] + assert [n for c, n in recorded if c == "codex"] == ["a", "b", "c"] + + def test_removes_servers_dropped_from_working_set(self, monkeypatch): + removed: list[tuple[str, str]] = [] + monkeypatch.setattr( + mcp, + "remove_client_mcp_server", + lambda client, name: removed.append((client, name)) or [], + ) + monkeypatch.setattr(mcp, "configure_client_mcp_server", lambda *a, **k: []) + + original = [self._server("gone", ["claude"])] + + changed = mcp.apply_mcp_server_changes(original, [], ["claude"]) + + assert changed is True + assert removed == [("claude", "gone")] + + def test_no_ops_returns_false_without_spinner(self, monkeypatch): + # Identical original/working, so nothing to do. + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda *a, **k: pytest.fail("should not configure"), + ) + servers = [self._server("a", ["claude"])] + + assert mcp.apply_mcp_server_changes(servers, servers, ["claude"]) is False + + +class TestConfigureMcpWizardNavigation: + def test_back_reshows_source_screen(self, monkeypatch): + """Pressing ← in the picker (returns _BACK) re-runs the source screen, + then the picker again; a real selection on the second pass proceeds.""" + monkeypatch.setattr(mcp, "load_state", lambda: {**CLAUDE_STATE}) + monkeypatch.setattr(mcp.shutil, "which", lambda binary: f"/usr/bin/{binary}") + monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace, profile=None: None) + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude"]) + monkeypatch.setattr( + mcp, "discover_external_mcp_connection_names", lambda workspace, profile=None: [] + ) + monkeypatch.setattr(mcp, "discover_app_mcp_servers", lambda workspace, profile=None: []) + monkeypatch.setattr(mcp, "discover_genie_mcp_servers", lambda workspace, profile=None: []) + monkeypatch.setattr(mcp, "discover_mcp_service_names", lambda workspace, profile=None: []) + monkeypatch.setattr( + mcp, + "discover_all_mcp_service_names", + lambda workspace, profile=None, on_progress=None: [], + ) + + source_calls: list[int] = [] + + def fake_sources(): + source_calls.append(1) + return {"external", "apps", "mcp-services", "genie"} + + monkeypatch.setattr(mcp, "prompt_for_mcp_search_sources", fake_sources) + + # First picker press = back, second = submit nothing. + picker_results = [mcp._BACK, []] + monkeypatch.setattr( + mcp, + "prompt_for_mcp_server_choices", + lambda *a, **k: picker_results.pop(0), + ) + monkeypatch.setattr(mcp, "save_state", lambda state: None) + + assert mcp.configure_mcp_command() == 0 + # Source screen shown twice (initial + after back). + assert len(source_calls) == 2 + + def test_cancel_on_source_screen_exits(self, monkeypatch): + monkeypatch.setattr(mcp, "load_state", lambda: {**CLAUDE_STATE}) + monkeypatch.setattr(mcp.shutil, "which", lambda binary: f"/usr/bin/{binary}") + monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace, profile=None: None) + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude"]) + # Cancelling the first screen (None) returns without discovering anything. + monkeypatch.setattr(mcp, "prompt_for_mcp_search_sources", lambda: None) + monkeypatch.setattr( + mcp, + "prompt_for_mcp_server_choices", + lambda *a, **k: pytest.fail("picker should not run after cancel"), + ) + + assert mcp.configure_mcp_command() == 0 + + class TestConfigureMcpCommand: def test_skips_existing_server_state_by_name(self, monkeypatch): saved_states: list[dict] = [] @@ -603,7 +756,9 @@ def test_registers_discovered_genie_space_server(self, monkeypatch): ], ) monkeypatch.setattr(mcp, "discover_app_mcp_servers", lambda workspace, profile=None: []) - _patch_mcp_choices(monkeypatch, f"{mcp.MCP_ADD_PREFIX}genie-space:space-123") + _patch_mcp_choices( + monkeypatch, f"{mcp.MCP_ADD_PREFIX}genie-space:space-123", categories={"genie"} + ) monkeypatch.setattr( mcp, "configure_client_mcp_server", @@ -648,12 +803,14 @@ def test_registers_discovered_vector_search_server(self, monkeypatch): monkeypatch.setattr(mcp, "discover_genie_mcp_servers", lambda workspace, profile=None: []) monkeypatch.setattr(mcp, "discover_app_mcp_servers", lambda workspace, profile=None: []) _patch_mcp_choices( - monkeypatch, f"{mcp.MCP_ADD_PREFIX}{mcp.VECTOR_SEARCH_SELECTION_PREFIX}main.search" + monkeypatch, + f"{mcp.MCP_ADD_PREFIX}{mcp.VECTOR_SEARCH_SELECTION_PREFIX}main.search", + categories={"vector-search"}, ) monkeypatch.setattr( mcp, "discover_vector_search_mcp_servers", - lambda workspace, profile=None: mcp.vector_search_mcp_servers( + lambda workspace, profile=None, on_progress=None: mcp.vector_search_mcp_servers( [("main", "search")], workspace ), ) @@ -703,11 +860,12 @@ def test_registers_discovered_uc_functions_server(self, monkeypatch): _patch_mcp_choices( monkeypatch, f"{mcp.MCP_ADD_PREFIX}{mcp.UC_FUNCTIONS_SELECTION_PREFIX}analytics.tools", + categories={"uc-functions"}, ) monkeypatch.setattr( mcp, "discover_uc_functions_mcp_servers", - lambda workspace, profile=None: mcp.uc_functions_mcp_servers( + lambda workspace, profile=None, on_progress=None: mcp.uc_functions_mcp_servers( [("analytics", "tools")], workspace ), ) @@ -741,6 +899,96 @@ def test_registers_discovered_uc_functions_server(self, monkeypatch): } ] + def test_registers_mcp_service_from_workspace_wide_walk(self, monkeypatch): + """The workspace-wide walk runs by default and folds its services into + the picker via the same mcp-service path as the curated system.ai list.""" + saved_states: list[dict] = [] + configured: list[tuple[str, str, str, dict]] = [] + walk_calls: list[str] = [] + + monkeypatch.setattr(mcp, "load_state", lambda: {**CLAUDE_STATE}) + monkeypatch.setattr(mcp.shutil, "which", lambda binary: f"/usr/bin/{binary}") + monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace, profile=None: None) + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude"]) + monkeypatch.setattr( + mcp, "discover_external_mcp_connection_names", lambda workspace, profile=None: [] + ) + monkeypatch.setattr(mcp, "discover_genie_mcp_servers", lambda workspace, profile=None: []) + monkeypatch.setattr(mcp, "discover_app_mcp_servers", lambda workspace, profile=None: []) + _patch_mcp_choices( + monkeypatch, + f"{mcp.MCP_ADD_PREFIX}{mcp.MCP_SERVICE_SELECTION_PREFIX}mycat.myschema.weather", + ) + + def fake_walk(workspace, profile=None, on_progress=None): + walk_calls.append(workspace) + return ["mycat.myschema.weather"] + + monkeypatch.setattr(mcp, "discover_all_mcp_service_names", fake_walk) + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + ) + monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) + + assert mcp.configure_mcp_command() == 0 + + assert walk_calls == [WS] + assert configured == [ + ( + "claude", + "mycat-myschema-weather", + f"{WS}/ai-gateway/mcp-services/mycat.myschema.weather", + { + "type": "http", + "url": f"{WS}/ai-gateway/mcp-services/mycat.myschema.weather", + "headers": {"Authorization": "Bearer ${OAUTH_TOKEN}"}, + }, + ) + ] + assert saved_states[-1]["mcp_servers"] == [ + { + "name": "mycat-myschema-weather", + "url": f"{WS}/ai-gateway/mcp-services/mycat.myschema.weather", + "auth": "env:OAUTH_TOKEN", + "clients": ["claude"], + } + ] + + def test_skips_slow_walks_unless_source_selected(self, monkeypatch): + """Vector Search and UC functions walk the workspace and are OFF by + default on the search-sources screen, so their discovery must not run + unless the user selects them. Genie is a fast default and may run.""" + called: list[str] = [] + + monkeypatch.setattr(mcp, "load_state", lambda: {**CLAUDE_STATE}) + monkeypatch.setattr(mcp.shutil, "which", lambda binary: f"/usr/bin/{binary}") + monkeypatch.setattr(mcp, "ensure_databricks_auth", lambda workspace, profile=None: None) + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude"]) + monkeypatch.setattr( + mcp, "discover_external_mcp_connection_names", lambda workspace, profile=None: [] + ) + monkeypatch.setattr(mcp, "discover_app_mcp_servers", lambda workspace, profile=None: []) + monkeypatch.setattr(mcp, "discover_genie_mcp_servers", lambda workspace, profile=None: []) + # Default sources only (no vector-search / uc-functions). Set the + # trackers AFTER `_patch_mcp_choices` since it stubs the same discoveries. + _patch_mcp_choices(monkeypatch) + + def track(name): + def _discover(workspace, profile=None, on_progress=None): + called.append(name) + return [] + + return _discover + + monkeypatch.setattr(mcp, "discover_vector_search_mcp_servers", track("vector-search")) + monkeypatch.setattr(mcp, "discover_uc_functions_mcp_servers", track("uc-functions")) + monkeypatch.setattr(mcp, "save_state", lambda state: None) + + assert mcp.configure_mcp_command() == 0 + assert called == [] + def test_registers_discovered_app_mcp_server(self, monkeypatch): saved_states: list[dict] = [] configured: list[tuple[str, str, str, dict]] = [] @@ -1067,7 +1315,7 @@ def test_continues_when_optional_discovery_fails(self, monkeypatch, capsys): RuntimeError("permission denied") ), ) - _patch_mcp_choices(monkeypatch, f"{mcp.MCP_ADD_PREFIX}managed:sql") + _patch_mcp_choices(monkeypatch, f"{mcp.MCP_ADD_PREFIX}managed:sql", categories={"genie"}) monkeypatch.setattr( mcp, "configure_client_mcp_server", @@ -1078,9 +1326,9 @@ def test_continues_when_optional_discovery_fails(self, monkeypatch, capsys): assert mcp.configure_mcp_command() == 0 output = capsys.readouterr().out - assert "Skipped external connections." in output - assert "Skipped Genie spaces." in output - assert "Skipped Databricks apps." in output + assert "Skipped external connections" in output + assert "Skipped Genie spaces" in output + assert "Skipped Databricks apps" in output assert configured[0][1] == "databricks-sql" assert saved_states[-1]["mcp_servers"][0]["name"] == "databricks-sql" @@ -1112,7 +1360,7 @@ def fake_apps(workspace, profile=None): monkeypatch.setattr(mcp, "discover_external_mcp_connection_names", fake_external) monkeypatch.setattr(mcp, "discover_genie_mcp_servers", fake_genie) monkeypatch.setattr(mcp, "discover_app_mcp_servers", fake_apps) - _patch_mcp_choices(monkeypatch) + _patch_mcp_choices(monkeypatch, categories={"genie"}) monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) assert mcp.configure_mcp_command() == 0