From 55f730ac0569a7532bd02b1a8a0894b3cdfa5f50 Mon Sep 17 00:00:00 2001 From: Anjali Sujithan Date: Mon, 15 Jun 2026 21:16:04 +0000 Subject: [PATCH] configure mcp surface vector search and UC functions --- src/ucode/databricks.py | 261 ++++++++++++++++++++++++++++++++++++++++ src/ucode/mcp.py | 179 ++++++++++++++++++++++++++- tests/test_mcp.py | 183 +++++++++++++++++++++++++++- 3 files changed, 619 insertions(+), 4 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 574d906..4b01120 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -14,6 +14,14 @@ import shlex import shutil import subprocess +import time +from concurrent.futures import ( + ThreadPoolExecutor, + as_completed, +) +from concurrent.futures import ( + TimeoutError as FutureTimeoutError, +) from pathlib import Path from typing import Literal, cast, overload from urllib import error as urllib_error @@ -1251,6 +1259,259 @@ def build_mcp_service_url(workspace: str, full_name: str) -> str: return f"{workspace}/ai-gateway/mcp-services/{full_name}" +# `list_vector_search_catalog_schemas` walks Vector Search endpoints+indexes. +# `list_uc_functions_catalog_schemas` walks UC catalogs+schemas in parallel and +# keeps only schemas with at least one user function. + +_UC_LIST_PAGE_SIZE = 200 +_UC_LIST_MAX_PAGES = 50 +_UC_FUNCTION_PROBE_WORKERS = 16 +_UC_LIST_HTTP_TIMEOUT = 10 +_UC_FUNCTION_PROBE_TIMEOUT = 5 +_VECTOR_SEARCH_DEADLINE_SECONDS = 15.0 +_UC_FUNCTIONS_DEADLINE_SECONDS = 20.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( + base_url: str, + token: str, + *, + items_key: str, + extra_params: dict[str, str] | None = None, + page_size: int = _UC_LIST_PAGE_SIZE, + max_pages: int = _UC_LIST_MAX_PAGES, + timeout: int = 30, +) -> 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. + """ + items: list[dict] = [] + page_token: str | None = None + seen_tokens: set[str] = set() + last_reason: str | None = None + for _ in range(max_pages): + params: dict[str, str] = {"max_results": str(page_size)} + if extra_params: + params.update(extra_params) + if page_token: + params["page_token"] = page_token + url = f"{base_url}?{urlencode(params)}" + payload, reason = _http_get_json(url, token, timeout=timeout) + if payload is None: + last_reason = reason + break + data = cast(dict, payload) if isinstance(payload, dict) else {} + raw = data.get(items_key) or [] + if isinstance(raw, list): + for item in raw: + if isinstance(item, dict): + items.append(item) + page_token = data.get("next_page_token") or None + if not page_token or page_token in seen_tokens: + break + seen_tokens.add(page_token) + return items, last_reason + + +def _vector_index_catalog_schema(index: dict) -> tuple[str, str] | None: + """Pull (catalog, schema) from one vector-search index entry.""" + catalog = index.get("catalog_name") + schema = index.get("schema_name") + if isinstance(catalog, str) and isinstance(schema, str) and catalog and schema: + return catalog, schema + # Fallback: `name` is the fully-qualified UC name `catalog.schema.index`. + name = index.get("name") + if isinstance(name, str): + parts = name.split(".") + if len(parts) >= 3 and parts[0] and parts[1]: + return parts[0], parts[1] + return None + + +def list_vector_search_catalog_schemas( + workspace: str, + token: str, + *, + deadline_seconds: float = _VECTOR_SEARCH_DEADLINE_SECONDS, +) -> 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.""" + hostname = workspace_hostname(workspace) + deadline = time.monotonic() + deadline_seconds + endpoints, reason = _paginated_json_items( + f"https://{hostname}/api/2.0/vector-search/endpoints", + token, + items_key="endpoints", + timeout=_UC_LIST_HTTP_TIMEOUT, + ) + if not endpoints: + return [], reason or "no vector search endpoints found" + + endpoint_names = [e["name"] for e in endpoints if isinstance(e.get("name"), str) and e["name"]] + if not endpoint_names: + return [], "no vector search endpoints with names" + + pairs: set[tuple[str, str]] = set() + workers = max(1, min(_UC_FUNCTION_PROBE_WORKERS, len(endpoint_names))) + with ThreadPoolExecutor(max_workers=workers) as pool: + futures = { + pool.submit( + _paginated_json_items, + f"https://{hostname}/api/2.0/vector-search/indexes", + token, + items_key="vector_indexes", + extra_params={"endpoint_name": name}, + timeout=_UC_LIST_HTTP_TIMEOUT, + ): name + for name in endpoint_names + } + + def collect(result, _endpoint): + indexes, _ = result + for index in indexes: + pair = _vector_index_catalog_schema(index) + if pair: + pairs.add(pair) + + _drain_with_deadline(futures, deadline, collect) + pool.shutdown(wait=False, cancel_futures=True) + + if not pairs: + return [], "no vector search indexes found" + return sorted(pairs), None + + +def _schema_has_user_function(hostname: str, token: str, catalog: str, schema: str) -> bool: + """One-shot probe: does `{catalog}.{schema}` expose any UC function?""" + url = ( + f"https://{hostname}/api/2.1/unity-catalog/functions" + f"?{urlencode({'catalog_name': catalog, 'schema_name': schema, 'max_results': '1'})}" + ) + payload, _reason = _http_get_json(url, token, timeout=_UC_FUNCTION_PROBE_TIMEOUT) + if not isinstance(payload, dict): + return False + functions = payload.get("functions") or [] + return isinstance(functions, list) and any(isinstance(item, dict) for item in functions) + + +def list_uc_functions_catalog_schemas( + workspace: str, + token: str, + *, + deadline_seconds: float = _UC_FUNCTIONS_DEADLINE_SECONDS, +) -> tuple[list[tuple[str, str]], str | None]: + """Return sorted unique `(catalog, schema)` pairs containing at least one + user-defined UC function.""" + 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. + candidate_pairs: list[tuple[str, 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") + # `information_schema` is auto-attached to every catalog and + # never holds user functions. + if ( + isinstance(schema_name, str) + and schema_name + and schema_name != "information_schema" + ): + candidate_pairs.append((catalog, schema_name)) + + _drain_with_deadline(schema_futures, deadline, collect_schemas) + pool.shutdown(wait=False, cancel_futures=True) + + if not candidate_pairs: + if time.monotonic() > deadline: + return [], "deadline exceeded while listing UC schemas" + return [], "no UC schemas found" + + # Parallel function-existence probes. + pairs: set[tuple[str, str]] = set() + with ThreadPoolExecutor(max_workers=_UC_FUNCTION_PROBE_WORKERS) as pool: + probe_futures = { + pool.submit(_schema_has_user_function, hostname, token, cat, schema): (cat, schema) + for cat, schema in candidate_pairs + } + + def collect_pair(has_fn, pair): + if has_fn: + pairs.add(pair) + + _drain_with_deadline(probe_futures, deadline, collect_pair) + pool.shutdown(wait=False, cancel_futures=True) + + if not pairs: + if time.monotonic() > deadline: + return [], "deadline exceeded probing UC schemas for functions" + return [], "no UC schemas with user functions found" + return sorted(pairs), 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 e2912a2..5845149 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -35,6 +35,8 @@ list_databricks_connections, list_genie_spaces, list_mcp_services, + list_uc_functions_catalog_schemas, + list_vector_search_catalog_schemas, workspace_hostname, ) from ucode.state import load_full_state, load_state, save_state @@ -43,6 +45,7 @@ print_section, print_success, print_warning, + spinner, ) MCP_AUTH_TOKEN_ENV_VAR = "OAUTH_TOKEN" @@ -81,6 +84,8 @@ GENIE_SPACE_SELECTION_PREFIX = "genie-space:" APP_MCP_SELECTION_PREFIX = "app:" MCP_SERVICE_SELECTION_PREFIX = "mcp-service:" +VECTOR_SEARCH_SELECTION_PREFIX = "vector-search:" +UC_FUNCTIONS_SELECTION_PREFIX = "uc-functions:" MCP_ADD_PREFIX = "add:" MCP_CONNECTION_MARKERS = ( "is_mcp", @@ -460,6 +465,73 @@ def discover_app_mcp_servers(workspace: str, profile: str | None = None) -> list return app_mcp_servers(list_databricks_apps(workspace, profile)) +def _catalog_schema_server_name(prefix: str, catalog: str, schema: str, taken: set[str]) -> str: + """Stable server name for a per-(catalog, schema) managed MCP entry. + + Prefers the lowercase alphanumeric slug; falls back to a numeric suffix on + collision so two schemas that slug to the same value still both render.""" + slug = f"{_normalize_workspace_title(catalog)}-{_normalize_workspace_title(schema)}".strip("-") + candidate = f"{prefix}-{slug}" if slug else prefix + if candidate not in taken: + return candidate + counter = 2 + while f"{candidate}-{counter}" in taken: + counter += 1 + return f"{candidate}-{counter}" + + +def vector_search_mcp_servers(pairs: list[tuple[str, str]], workspace: str) -> list[dict]: + servers: list[dict] = [] + seen_names: set[str] = set() + for catalog, schema in pairs: + if not catalog or not schema: + continue + name = _catalog_schema_server_name("databricks-vector-search", catalog, schema, seen_names) + seen_names.add(name) + servers.append( + { + "name": name, + "title": f"{catalog}.{schema}", + "catalog": catalog, + "schema": schema, + "url": f"{workspace}/api/2.0/mcp/vector-search/{catalog}/{schema}", + } + ) + return sorted(servers, key=lambda server: str(server["title"]).lower()) + + +def discover_vector_search_mcp_servers(workspace: str, profile: str | None = None) -> list[dict]: + token = get_databricks_token(workspace, profile) + pairs, _reason = list_vector_search_catalog_schemas(workspace, token) + return vector_search_mcp_servers(pairs, workspace) + + +def uc_functions_mcp_servers(pairs: list[tuple[str, str]], workspace: str) -> list[dict]: + servers: list[dict] = [] + seen_names: set[str] = set() + for catalog, schema in pairs: + if not catalog or not schema: + continue + name = _catalog_schema_server_name("databricks-functions", catalog, schema, seen_names) + seen_names.add(name) + servers.append( + { + "name": name, + "title": f"{catalog}.{schema}", + "catalog": catalog, + "schema": schema, + "url": f"{workspace}/api/2.0/mcp/functions/{catalog}/{schema}", + } + ) + return sorted(servers, key=lambda server: str(server["title"]).lower()) + + +def discover_uc_functions_mcp_servers(workspace: str, profile: str | None = None) -> list[dict]: + token = get_databricks_token(workspace, profile) + pairs, _reason = list_uc_functions_catalog_schemas(workspace, token) + return uc_functions_mcp_servers(pairs, workspace) + + def _picker_style() -> questionary.Style: return questionary.Style( [ @@ -694,6 +766,8 @@ def build_mcp_picker_choices( available_app_servers: list[dict], original_servers: list[dict], 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[questionary.Choice | questionary.Separator]: original_by_name = _servers_by_name(original_servers) known_names = set(original_by_name) @@ -759,6 +833,42 @@ def build_mcp_picker_choices( ) displayed_names.add(name) + for server in available_vector_search_servers or []: + name = _server_name(server) + catalog = server.get("catalog") + schema = server.get("schema") + if not name or not isinstance(catalog, str) or not isinstance(schema, str): + continue + display_title = f"Vector Search: {catalog}.{schema}" + if name in known_names: + choices.append(_server_choice(name, True, display_title)) + else: + choices.append( + _add_choice( + f"{VECTOR_SEARCH_SELECTION_PREFIX}{catalog}.{schema}", + display_title, + ) + ) + displayed_names.add(name) + + for server in available_uc_functions_servers or []: + name = _server_name(server) + catalog = server.get("catalog") + schema = server.get("schema") + if not name or not isinstance(catalog, str) or not isinstance(schema, str): + continue + display_title = f"UC Functions: {catalog}.{schema}" + if name in known_names: + choices.append(_server_choice(name, True, display_title)) + else: + choices.append( + _add_choice( + f"{UC_FUNCTIONS_SELECTION_PREFIX}{catalog}.{schema}", + display_title, + ) + ) + displayed_names.add(name) + for name in sorted(known_names - displayed_names): choices.append(_server_choice(name, True)) return choices @@ -770,6 +880,8 @@ def prompt_for_mcp_server_choices( available_app_servers: list[dict], original_servers: list[dict], 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: selection = _scrolling_checkbox( "MCP:", @@ -779,6 +891,8 @@ def prompt_for_mcp_server_choices( available_app_servers, original_servers, available_mcp_service_names, + available_vector_search_servers, + available_uc_functions_servers, ), style=_picker_style(), instruction="(space to toggle, enter to save, type to filter)", @@ -797,6 +911,8 @@ def _resolve_mcp_selection( workspace: str, available_app_servers: list[dict] | None = None, available_genie_servers: list[dict] | None = None, + available_vector_search_servers: list[dict] | None = None, + available_uc_functions_servers: list[dict] | None = None, ) -> tuple[str, str]: if selection.startswith(APP_MCP_SELECTION_PREFIX): app_name = selection.removeprefix(APP_MCP_SELECTION_PREFIX) @@ -837,15 +953,63 @@ def _resolve_mcp_selection( # URL keeps the UC `..` form; entry name uses dashes. return full_name.replace(".", "-"), build_mcp_service_url(workspace, full_name) + if selection.startswith(VECTOR_SEARCH_SELECTION_PREFIX): + return _resolve_catalog_schema_selection( + selection.removeprefix(VECTOR_SEARCH_SELECTION_PREFIX), + kind="vector search", + url_path="vector-search", + name_prefix="databricks-vector-search", + workspace=workspace, + available_servers=available_vector_search_servers, + ) + + if selection.startswith(UC_FUNCTIONS_SELECTION_PREFIX): + return _resolve_catalog_schema_selection( + selection.removeprefix(UC_FUNCTIONS_SELECTION_PREFIX), + kind="UC functions", + url_path="functions", + name_prefix="databricks-functions", + workspace=workspace, + available_servers=available_uc_functions_servers, + ) + if selection == SQL_MCP_VALUE: return "databricks-sql", f"{workspace}/api/2.0/mcp/sql" raise RuntimeError(f"unrecognized selection prefix in `{selection}`") +def _resolve_catalog_schema_selection( + payload: str, + *, + kind: str, + url_path: str, + name_prefix: str, + workspace: str, + available_servers: list[dict] | None, +) -> tuple[str, str]: + """Map a `catalog.schema` picker value back to the discovered server's name + and URL, falling back to a deterministic slug when discovery has been lost + (e.g. picker reopened on a stale workspace).""" + if not payload or "." not in payload: + raise RuntimeError(f"missing catalog.schema for {kind}") + catalog, _, schema = payload.partition(".") + if not catalog or not schema: + raise RuntimeError(f"missing catalog.schema for {kind}") + for server in available_servers or []: + if server.get("catalog") == catalog and server.get("schema") == schema: + name = _server_name(server) + url = server.get("url") + if name and isinstance(url, str) and url: + return name, url + name = _catalog_schema_server_name(name_prefix, catalog, schema, set()) + return name, f"{workspace}/api/2.0/mcp/{url_path}/{catalog}/{schema}" + + def _discover_mcp_source(label: str, discover: Callable[[], list[Any]]) -> list[Any]: try: - return discover() + with spinner(f"Discovering {label}..."): + return discover() except RuntimeError: print_warning(f"Skipped {label}.") return [] @@ -996,6 +1160,15 @@ def configure_mcp_command() -> int: "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) @@ -1005,6 +1178,8 @@ def configure_mcp_command() -> int: 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 @@ -1028,6 +1203,8 @@ def configure_mcp_command() -> int: workspace, available_app_mcp_servers, available_genie_mcp_servers, + available_vector_search_servers, + available_uc_functions_servers, ) except RuntimeError as exc: print_warning(f"Skipped MCP selection `{selection}`: {exc}.") diff --git a/tests/test_mcp.py b/tests/test_mcp.py index a0b4eba..08b9705 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -385,6 +385,71 @@ def test_picker_lists_discovered_app_mcps(self): choices_by_title = {choice.title: choice for choice in choices} assert choices_by_title["App: mcp-my-app"].value == f"{mcp.MCP_ADD_PREFIX}app:mcp-my-app" + def test_vector_search_mcp_servers_emit_managed_url_per_pair(self): + servers = mcp.vector_search_mcp_servers( + [("main", "search"), ("Marketing", "Docs")], + WS, + ) + assert servers == [ + { + "name": "databricks-vector-search-main-search", + "title": "main.search", + "catalog": "main", + "schema": "search", + "url": f"{WS}/api/2.0/mcp/vector-search/main/search", + }, + { + "name": "databricks-vector-search-marketing-docs", + "title": "Marketing.Docs", + "catalog": "Marketing", + "schema": "Docs", + "url": f"{WS}/api/2.0/mcp/vector-search/Marketing/Docs", + }, + ] + + def test_uc_functions_mcp_servers_emit_managed_url_per_pair(self): + servers = mcp.uc_functions_mcp_servers( + [("analytics", "tools"), ("ml", "udfs")], + WS, + ) + assert servers == [ + { + "name": "databricks-functions-analytics-tools", + "title": "analytics.tools", + "catalog": "analytics", + "schema": "tools", + "url": f"{WS}/api/2.0/mcp/functions/analytics/tools", + }, + { + "name": "databricks-functions-ml-udfs", + "title": "ml.udfs", + "catalog": "ml", + "schema": "udfs", + "url": f"{WS}/api/2.0/mcp/functions/ml/udfs", + }, + ] + + def test_picker_lists_discovered_vector_search_and_uc_functions(self): + choices = mcp.build_mcp_picker_choices( + [], + [], + [], + [], + available_vector_search_servers=mcp.vector_search_mcp_servers([("main", "search")], WS), + available_uc_functions_servers=mcp.uc_functions_mcp_servers( + [("analytics", "tools")], WS + ), + ) + choices_by_title = {choice.title: choice for choice in choices} + assert ( + choices_by_title["Vector Search: main.search"].value + == f"{mcp.MCP_ADD_PREFIX}vector-search:main.search" + ) + assert ( + choices_by_title["UC Functions: analytics.tools"].value + == f"{mcp.MCP_ADD_PREFIX}uc-functions:analytics.tools" + ) + def test_picker_keeps_saved_legacy_servers_for_removal(self): choices = mcp.build_mcp_picker_choices( [], @@ -408,10 +473,15 @@ def _patch_mcp_choices(monkeypatch, *values: str) -> None: "prompt_for_mcp_server_choices", lambda *args, **kwargs: list(values), ) - # Curated system.ai.* MCP-services discovery now always runs; stub it so - # configure_mcp_command tests don't shell out to the `databricks` CLI. - # Tests that exercise it override this after calling the helper. + # 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: [] + ) + monkeypatch.setattr( + mcp, "discover_uc_functions_mcp_servers", lambda workspace, profile=None: [] + ) class TestConfigureMcpCommand: @@ -564,6 +634,113 @@ def test_registers_discovered_genie_space_server(self, monkeypatch): } ] + def test_registers_discovered_vector_search_server(self, monkeypatch): + saved_states: list[dict] = [] + configured: list[tuple[str, str, str, dict]] = [] + + 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.VECTOR_SEARCH_SELECTION_PREFIX}main.search" + ) + monkeypatch.setattr( + mcp, + "discover_vector_search_mcp_servers", + lambda workspace, profile=None: mcp.vector_search_mcp_servers( + [("main", "search")], workspace + ), + ) + 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 configured == [ + ( + "claude", + "databricks-vector-search-main-search", + f"{WS}/api/2.0/mcp/vector-search/main/search", + { + "type": "http", + "url": f"{WS}/api/2.0/mcp/vector-search/main/search", + "headers": {"Authorization": "Bearer ${OAUTH_TOKEN}"}, + }, + ) + ] + assert saved_states[-1]["mcp_servers"] == [ + { + "name": "databricks-vector-search-main-search", + "url": f"{WS}/api/2.0/mcp/vector-search/main/search", + "auth": "env:OAUTH_TOKEN", + "clients": ["claude"], + } + ] + + def test_registers_discovered_uc_functions_server(self, monkeypatch): + saved_states: list[dict] = [] + configured: list[tuple[str, str, str, dict]] = [] + + 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.UC_FUNCTIONS_SELECTION_PREFIX}analytics.tools", + ) + monkeypatch.setattr( + mcp, + "discover_uc_functions_mcp_servers", + lambda workspace, profile=None: mcp.uc_functions_mcp_servers( + [("analytics", "tools")], workspace + ), + ) + 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 configured == [ + ( + "claude", + "databricks-functions-analytics-tools", + f"{WS}/api/2.0/mcp/functions/analytics/tools", + { + "type": "http", + "url": f"{WS}/api/2.0/mcp/functions/analytics/tools", + "headers": {"Authorization": "Bearer ${OAUTH_TOKEN}"}, + }, + ) + ] + assert saved_states[-1]["mcp_servers"] == [ + { + "name": "databricks-functions-analytics-tools", + "url": f"{WS}/api/2.0/mcp/functions/analytics/tools", + "auth": "env:OAUTH_TOKEN", + "clients": ["claude"], + } + ] + def test_registers_discovered_app_mcp_server(self, monkeypatch): saved_states: list[dict] = [] configured: list[tuple[str, str, str, dict]] = []