diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 9ff58f313..c598480e0 100644 --- a/src/ucode/databricks.py +++ b/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 `.` 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 `.` 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( diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 079871830..f05ca5279 100644 --- a/src/ucode/mcp.py +++ b/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 `.` 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 `" + ) + return None return loader diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 17eb7239c..3d2de72d5 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1681,11 +1681,12 @@ def test_aggregates_services_across_catalogs_and_schemas(self, monkeypatch): ), ) - names, reason = db_mod.list_all_mcp_services(WS, "token") + walk = db_mod.list_all_mcp_services(WS, "token") - assert reason is None + assert walk.error is None + assert not walk.timed_out # information_schema is skipped; results are sorted and de-duplicated. - assert names == [ + assert walk.names == [ "mycat.myschema.news", "mycat.myschema.weather", "other.ops.pager", @@ -1703,14 +1704,14 @@ def test_reports_progress_per_schema(self, monkeypatch): ) progress: list[tuple[int, int, int]] = [] - names, reason = db_mod.list_all_mcp_services( + walk = 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"] + assert walk.error is None + assert walk.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] @@ -1728,20 +1729,162 @@ def test_skips_internal_catalogs(self, monkeypatch): ), ) - names, reason = db_mod.list_all_mcp_services(WS, "token") + walk = db_mod.list_all_mcp_services(WS, "token") - assert names == [] - assert reason == "no user UC catalogs found" + assert walk.names == [] + assert walk.error == "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") + walk = db_mod.list_all_mcp_services(WS, "token") - assert names == [] - assert reason == "no UC catalogs found" + assert walk.names == [] + assert walk.error == "no UC catalogs found" + + def test_probes_main_default_first(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_get_json", + self._fake_http( + catalogs=["zeta", "main"], + schemas_by_catalog={"zeta": ["default", "x"], "main": ["alpha", "default"]}, + services_by_schema={}, + ), + ) + probed: list[str] = [] + monkeypatch.setattr( + db_mod, + "list_mcp_services", + lambda workspace, token, parent="system.ai": probed.append(parent) or ([], None), + ) + + real_walk = db_mod.walk_catalog_schemas + monkeypatch.setattr( + db_mod, + "walk_catalog_schemas", + lambda *a, **kw: real_walk(*a, **{**kw, "max_workers": 1}), + ) + + db_mod.list_all_mcp_services(WS, "token") + + # The user's likeliest schema leads, so a deadline cut-off still finds it. + assert probed[0] == "main.default" + + def test_flags_partial_results_past_deadline(self, monkeypatch): + def slow_walk(workspace, token, *, deadline, probe, collect, on_incomplete, **kwargs): + collect((["main.default.exa"], None), 1, 2) + while time.monotonic() <= deadline: + time.sleep(0.01) + on_incomplete(True) + + monkeypatch.setattr(db_mod, "walk_catalog_schemas", slow_walk) + + walk = db_mod.list_all_mcp_services(WS, "token", deadline_seconds=0.05) + + assert walk.names == ["main.default.exa"] + assert walk.timed_out + assert walk.error is None + + def test_timeout_before_any_probe_sets_timed_out_with_error(self, monkeypatch): + """The deadline expiring while listing catalogs/schemas is still flagged as a timeout.""" + + def early_timeout_walk(workspace, token, *, on_incomplete, **kwargs): + on_incomplete(True) + return "deadline exceeded while listing UC schemas" + + monkeypatch.setattr(db_mod, "walk_catalog_schemas", early_timeout_walk) + + walk = db_mod.list_all_mcp_services(WS, "token") + + assert walk.names == [] + assert walk.timed_out + assert walk.error == "deadline exceeded while listing UC schemas" + + def test_cancel_is_incomplete_but_not_a_timeout(self, monkeypatch): + def cancelled_walk(workspace, token, *, deadline, probe, collect, on_incomplete, **kwargs): + collect((["main.default.exa"], None), 1, 2) + on_incomplete(False) + + monkeypatch.setattr(db_mod, "walk_catalog_schemas", cancelled_walk) + + walk = db_mod.list_all_mcp_services(WS, "token") + + assert walk.names == ["main.default.exa"] + assert not walk.timed_out + assert walk.error is None + + def test_real_walk_flags_timeout_when_probe_blocks_past_deadline(self, monkeypatch): + """No `walk_catalog_schemas` stub: the real `_collect_concurrently` returns False + because the probe outlives the deadline, which must fire `on_incomplete`.""" + monkeypatch.setattr( + db_mod, + "_http_get_json", + self._fake_http( + catalogs=["main"], + schemas_by_catalog={"main": ["fast", "slow"]}, + services_by_schema={}, + ), + ) + release = threading.Event() + + def probe(workspace, token, parent="system.ai"): + if parent == "main.fast": + return (["main.fast.exa"], None) + release.wait(timeout=2) + return ([], None) + + monkeypatch.setattr(db_mod, "list_mcp_services", probe) + try: + walk = db_mod.list_all_mcp_services(WS, "token", deadline_seconds=0.3) + finally: + release.set() + + assert walk.names == ["main.fast.exa"] + assert walk.timed_out + assert walk.error is None + + def test_real_walk_cancel_is_not_a_timeout(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_get_json", + self._fake_http( + catalogs=["main"], + schemas_by_catalog={"main": ["slow"]}, + services_by_schema={}, + ), + ) + release = threading.Event() + cancel = threading.Event() + + def probe(workspace, token, parent="system.ai"): + cancel.set() + release.wait(timeout=2) + return ([], None) + + monkeypatch.setattr(db_mod, "list_mcp_services", probe) + try: + walk = db_mod.list_all_mcp_services(WS, "token", cancel_event=cancel) + finally: + release.set() + + assert not walk.timed_out + + def test_fully_drained_walk_is_not_flagged_even_if_deadline_passes_after(self, monkeypatch): + def drained_walk(workspace, token, *, deadline, probe, collect, **kwargs): + collect((["main.default.exa"], None), 1, 1) + while time.monotonic() <= deadline: + time.sleep(0.01) + + monkeypatch.setattr(db_mod, "walk_catalog_schemas", drained_walk) + + walk = db_mod.list_all_mcp_services(WS, "token", deadline_seconds=0.05) + + assert walk.names == ["main.default.exa"] + assert not walk.timed_out + assert walk.error is None def _foundation_models_payload(names): diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 68f1fb39d..685bccadb 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -9,6 +9,7 @@ from ucode import mcp from ucode.agents import claude +from ucode.databricks import McpServicesWalk WS = "https://example.databricks.com" CLAUDE_STATE = {"workspace": WS, "available_tools": ["claude"]} @@ -1094,7 +1095,12 @@ def test_workspace_wide_walk_streams_services_into_picker(self, monkeypatch): walk_calls: list[str] = [] def fake_walk( - workspace, profile=None, on_progress=None, on_services=None, cancel_event=None + workspace, + profile=None, + on_progress=None, + on_services=None, + cancel_event=None, + on_timeout=None, ): walk_calls.append(workspace) if on_services is not None: @@ -1112,6 +1118,65 @@ def fake_walk( f"{mcp.MCP_ADD_PREFIX}{mcp.MCP_SERVICE_SELECTION_PREFIX}mycat.myschema.weather" ] + def test_picker_includes_user_mcp_service_alongside_system_ai(self, monkeypatch): + """A user's own `main.default.exa` streams into a picker that already shows + the `system.ai` list, and a service seen by both lists appears once.""" + from ucode.ui.interactive_picker import merge_new_choices + + def fake_walk(workspace, profile=None, on_services=None, **kwargs): + on_services(["main.default.exa", "system.ai.github"]) + return ["main.default.exa", "system.ai.github"] + + monkeypatch.setattr(mcp, "discover_all_mcp_service_names", fake_walk) + shown = mcp.build_mcp_picker_choices([], [], [], [], ["system.ai.github"], additive=True) + loader = mcp._mcp_services_background_loader(WS, None, set(), additive=True) + streamed: list = [] + assert loader(streamed.extend, threading.Event()) is None + + merged = [*shown, *merge_new_choices(shown, streamed)] + values = [c.value for c in merged] + prefix = f"{mcp.MCP_ADD_PREFIX}{mcp.MCP_SERVICE_SELECTION_PREFIX}" + assert sorted(values) == [f"{prefix}main.default.exa", f"{prefix}system.ai.github"] + + def test_walk_timeout_keeps_system_ai_and_warns(self, monkeypatch): + """A timed-out walk (flag set on `McpServicesWalk`) streams partial results and warns.""" + monkeypatch.setattr(mcp, "get_databricks_token", lambda workspace, profile=None: "token") + + def fake_list_all(workspace, token, on_services=None, **kwargs): + on_services(["main.default.exa"]) + return McpServicesWalk(["main.default.exa"], timed_out=True) + + monkeypatch.setattr(mcp, "list_all_mcp_services", fake_list_all) + shown = mcp.build_mcp_picker_choices([], [], [], [], ["system.ai.github"], additive=True) + loader = mcp._mcp_services_background_loader(WS, None, set(), additive=True) + streamed: list = [] + + message = loader(streamed.extend, threading.Event()) + + assert [c.value for c in streamed] == [ + f"{mcp.MCP_ADD_PREFIX}{mcp.MCP_SERVICE_SELECTION_PREFIX}main.default.exa" + ] + assert message is not None and "Timed out" in message + assert "ug mcp add --names " in message + assert [c.value for c in shown] == [ + f"{mcp.MCP_ADD_PREFIX}{mcp.MCP_SERVICE_SELECTION_PREFIX}system.ai.github" + ] + + def test_walk_error_degrades_to_system_ai_only(self, monkeypatch): + """A non-timeout walk failure (e.g. no catalogs) adds nothing and shows no warning.""" + + monkeypatch.setattr(mcp, "get_databricks_token", lambda workspace, profile=None: "token") + monkeypatch.setattr( + mcp, + "list_all_mcp_services", + lambda *args, **kwargs: McpServicesWalk([], error="no UC catalogs found"), + ) + loader = mcp._mcp_services_background_loader(WS, None, set(), additive=True) + streamed: list = [] + + assert loader(streamed.extend, threading.Event()) is None + assert streamed == [] + def test_mcp_service_choice_known_vs_unknown(self): # Unregistered -> an add-choice; already-registered -> a removable toggle in replace # mode, and hidden (None) under `mcp add` since there is nothing to add.