Surface user-owned MCP services in the configure picker

src/ucode/databricks.py

@@ -3097,6 +3097,10 @@ def resolve_provider_launch_model(model: str | None, provider_models: dict[str,
 # enough time to enumerate them; a slow workspace still degrades to partial
 # results once the budget is exceeded instead of hanging indefinitely.
 _MCP_SERVICES_WALK_DEADLINE_SECONDS = 30.0
+# Reasons the walk reports when the deadline cuts it short before any schema is probed. Callers
+# learn about a timeout from `on_incomplete(True)` / `McpServicesWalk.timed_out`, not these strings.
+_WALK_TIMEOUT_CATALOGS_REASON = "deadline exceeded while listing UC catalogs"
+_WALK_TIMEOUT_SCHEMAS_REASON = "deadline exceeded while listing UC schemas"
 # Skip UC catalogs whose schemas almost never carry user-callable functions
 # you'd want to expose as agent tools.
 _UC_FUNCTIONS_SKIP_CATALOGS = frozenset(
@@ -3111,10 +3115,11 @@ def _collect_concurrently[T, R](
     *,
     max_workers: int,
     should_stop: Callable[[], bool],
-) -> None:
+) -> bool:
     """Run `run` over `items` on daemon workers that are never joined, so a slow or stuck
     call can't block process exit. Results go to `on_result` on the calling thread until
-    every item is drained or `should_stop()` returns True."""
+    every item is drained or `should_stop()` returns True. Returns True if every item
+    finished, False if it stopped with items left unprobed."""
     pending: queue.Queue[T] = queue.Queue()
     for item in items:
         pending.put(item)
@@ -3143,6 +3148,7 @@ def worker() -> None:
         remaining -= 1
         if result is not _PROBE_FAILED:
             on_result(result, item)
+    return remaining == 0
 
 
 def _paginated_json_items(
@@ -3198,6 +3204,7 @@ def walk_catalog_schemas[T](
     skip_catalogs: frozenset[str] = _UC_FUNCTIONS_SKIP_CATALOGS,
     max_workers: int = _SCHEMA_PROBE_WORKERS,
     cancel_event: threading.Event | None = None,
+    on_incomplete: Callable[[bool], None] | None = None,
 ) -> str | None:
     """Discover every user `<catalog>.<schema>` in the workspace and probe each one in parallel.
 
@@ -3210,6 +3217,11 @@ def walk_catalog_schemas[T](
       - `collect(result, done, total)`: handle each probe result as it lands — accumulating,
         de-duping, streaming — where `done`/`total` are the completed and total schema counts.
 
+    `on_incomplete(timed_out)`, if given, is called when the deadline or `cancel_event` stopped the
+    walk before every schema was probed (including before any probe ran), so callers can flag partial results without re-checking the clock. `timed_out` is
+    True when the deadline passed and False when `cancel_event` stopped the walk (e.g. the picker
+    closed), so a cancel is not reported as a timeout.
+
     Returns None once the probes run, or a short reason string if there are no catalogs or schemas
     to probe."""
     cancel_event = cancel_event or threading.Event()
@@ -3236,7 +3248,9 @@ def should_stop() -> bool:
     if not catalog_names:
         return "no user UC catalogs found"
     if time.monotonic() > deadline:
-        return "deadline exceeded while listing UC catalogs"
+        if on_incomplete is not None:
+            on_incomplete(True)
+        return _WALK_TIMEOUT_CATALOGS_REASON
 
     schema_refs: list[tuple[str, str]] = []
 
@@ -3263,9 +3277,16 @@ def collect_schemas(result, catalog):
 
     if not schema_refs:
         if time.monotonic() > deadline:
-            return "deadline exceeded while listing UC schemas"
+            if on_incomplete is not None:
+                on_incomplete(True)
+            return _WALK_TIMEOUT_SCHEMAS_REASON
         return "no UC schemas found"
 
+    # Probe the likeliest homes of a user's own services first: the `main` catalog and any
+    # `default` schema (where `CREATE MCP SERVICE` lands by default). Under the wall-clock
+    # deadline a big workspace returns partial results, so order decides what is found.
+    schema_refs.sort(key=lambda ref: (ref[0] != "main", ref[1] != "default"))
+
     schemas_total = len(schema_refs)
     schemas_done = 0
 
@@ -3274,17 +3295,32 @@ def collect_probe(result, _ref):
         schemas_done += 1
         collect(result, schemas_done, schemas_total)
 
-    _collect_concurrently(
+    drained = _collect_concurrently(
         schema_refs,
         lambda ref: probe(*ref),
         collect_probe,
         max_workers=max_workers,
         should_stop=should_stop,
     )
+    if not drained and on_incomplete is not None:
+        on_incomplete(time.monotonic() > deadline)
 
     return None
 
 
+@dataclass(frozen=True)
+class McpServicesWalk:
+    """Result of `list_all_mcp_services`.
+
+    `names`: sorted unique service full names found (possibly partial).
+    `error`: why the walk failed outright (no catalogs, bad token, early deadline); None otherwise.
+    `timed_out`: the deadline cut the probe phase short, so `names` may be incomplete."""
+
+    names: list[str]
+    error: str | None = None
+    timed_out: bool = False
+
+
 def list_all_mcp_services(
     workspace: str,
     token: str,
@@ -3293,7 +3329,7 @@ def list_all_mcp_services(
     on_progress: Callable[[int, int, int], None] | None = None,
     on_services: Callable[[list[str]], None] | None = None,
     cancel_event: threading.Event | None = None,
-) -> tuple[list[str], str | None]:
+) -> McpServicesWalk:
     """Return sorted unique MCP-service full names across every `<catalog>.<schema>`
     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,
@@ -3310,6 +3346,11 @@ def list_all_mcp_services(
     (single schema)."""
     deadline = time.monotonic() + deadline_seconds
     names: set[str] = set()
+    timed_out = False
+
+    def mark_incomplete(deadline_passed: bool) -> None:
+        nonlocal timed_out
+        timed_out = deadline_passed
 
     def probe(catalog, schema):
         return list_mcp_services(workspace, token, f"{catalog}.{schema}")
@@ -3323,16 +3364,23 @@ def collect(result, schemas_done, schemas_total):
         if on_services is not None and new:
             on_services(sorted(new))
 
-    reason = walk_catalog_schemas(
-        workspace, token, deadline=deadline, probe=probe, collect=collect, cancel_event=cancel_event
+    error = walk_catalog_schemas(
+        workspace,
+        token,
+        deadline=deadline,
+        probe=probe,
+        collect=collect,
+        cancel_event=cancel_event,
+        on_incomplete=mark_incomplete,
+    )
+    if error is not None:
+        return McpServicesWalk([], error=error, timed_out=timed_out)
+    # Partial results are still returned, flagged so callers can say the list is incomplete.
+    return McpServicesWalk(
+        sorted(names),
+        error=None if names or timed_out else "no MCP services found",
+        timed_out=timed_out,
     )
-    if reason is not None:
-        return [], reason
-    if not names:
-        if time.monotonic() > deadline:
-            return [], "deadline exceeded while listing MCP services"
-        return [], "no MCP services found"
-    return sorted(names), None
 
 
 def _get_anthropic_models_json(

src/ucode/mcp.py

@@ -484,22 +484,26 @@ def discover_all_mcp_service_names(
     on_progress: Callable[[int, int, int], None] | None = None,
     on_services: Callable[[list[str]], None] | None = None,
     cancel_event: threading.Event | None = None,
+    on_timeout: Callable[[], None] | None = None,
 ) -> list[str]:
     """All MCP services across every `<catalog>.<schema>` 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, and `on_services` to stream newly-found
-    service names into the picker as the walk progresses."""
+    service names into the picker as the walk progresses. `on_timeout` is called when the walk's
+    deadline cut it short (partial results were already streamed via `on_services`)."""
     token = get_databricks_token(workspace, profile)
-    names, reason = list_all_mcp_services(
+    walk = list_all_mcp_services(
         workspace,
         token,
         on_progress=on_progress,
         on_services=on_services,
         cancel_event=cancel_event,
     )
-    raise_for_invalid_access_token(workspace, reason)
-    return names
+    raise_for_invalid_access_token(workspace, walk.error)
+    if walk.timed_out and on_timeout is not None:
+        on_timeout()
+    return walk.names
 
 
 def _normalize_workspace_title(text: str) -> str:
@@ -808,7 +812,9 @@ def prompt_for_mcp_server_choices(
     allow_back: bool = False,
     additive: bool = False,
     managed_names: set[str] | None = None,
-    background_loader: Callable[[Callable[[list[questionary.Choice]], None], threading.Event], None]
+    background_loader: Callable[
+        [Callable[[list[questionary.Choice]], None], threading.Event], str | None
+    ]
     | None = None,
 ) -> list[str] | None | _Back:
     """Show the MCP server picker. Returns the list of selected values, `None`
@@ -1317,15 +1323,23 @@ def _mcp_services_background_loader(
     known_names: set[str],
     additive: bool,
     managed_names: set[str] | None = None,
-) -> Callable[[Callable[[list[questionary.Choice]], None], threading.Event], None]:
+) -> Callable[[Callable[[list[questionary.Choice]], None], threading.Event], str | None]:
     """Return a picker `background_loader` that runs the workspace-wide MCP-services walk and
     streams each schema's newly-found services into the open picker as choices, so the walk
     never blocks the picker from opening. Deduping against already-shown rows (e.g. the fast
-    `system.ai` list) is handled by the picker's append."""
+    `system.ai` list) is handled by the picker's append. The walk is bounded by a wall-clock
+    deadline; on a timeout or error the picker keeps what streamed in (at least `system.ai`), and
+    only a timeout adds a footer warning that the list may be incomplete (other errors stay silent)."""
 
     def loader(
         append: Callable[[list[questionary.Choice]], None], cancel_event: threading.Event
-    ) -> None:
+    ) -> str | None:
+        timed_out = False
+
+        def on_timeout() -> None:
+            nonlocal timed_out
+            timed_out = True
+
         def on_services(new_names: list[str]) -> None:
             # A streamed service can be hidden (managed, or already configured under `mcp add`),
             # in which case its choice is None — drop those so only real rows reach the picker.
@@ -1339,8 +1353,18 @@ def on_services(new_names: list[str]) -> None:
                 append(new_choices)
 
         discover_all_mcp_service_names(
-            workspace, profile, on_services=on_services, cancel_event=cancel_event
+            workspace,
+            profile,
+            on_services=on_services,
+            cancel_event=cancel_event,
+            on_timeout=on_timeout,
         )
+        if timed_out:
+            return (
+                "⚠ Timed out listing your MCP services; to add one not shown, use "
+                "`ug mcp add --names <catalog.schema.name>`"
+            )
+        return None
 
     return loader