From 9dabcfb6036abb7633593550ea634342e5ac7054 Mon Sep 17 00:00:00 2001 From: Achille Negrier Date: Fri, 16 Oct 2026 13:51:20 +0000 Subject: [PATCH 1/4] Surface user-owned MCP services in the configure picker The picker's background workspace walk (30s budget) probed schemas in arbitrary order, so on large workspaces a user's own service such as main.default.exa could still be missing, with no hint why. Probe the `main` catalog and `default` schemas first, flag partial results past the deadline, and show a "timed out" footer so users know to use `ug mcp add`. Co-authored-by: Isaac --- src/ucode/databricks.py | 18 +++++++++--- src/ucode/mcp.py | 29 ++++++++++++++++---- tests/test_databricks.py | 35 ++++++++++++++++++++++++ tests/test_mcp.py | 59 +++++++++++++++++++++++++++++++++++++++- 4 files changed, 130 insertions(+), 11 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 9ff58f313..45f88987d 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3266,6 +3266,11 @@ def collect_schemas(result, catalog): return "deadline exceeded while listing UC schemas" 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 @@ -3328,11 +3333,16 @@ def collect(result, schemas_done, schemas_total): ) if reason is not None: return [], reason + timed_out = time.monotonic() > deadline if not names: - if time.monotonic() > deadline: - return [], "deadline exceeded while listing MCP services" - return [], "no MCP services found" - return sorted(names), None + return ( + [], + "deadline exceeded while listing MCP services" + if timed_out + else "no MCP services found", + ) + # Partial results are still returned, but flagged so callers can say the list is incomplete. + return sorted(names), "deadline exceeded while listing MCP services" if timed_out else None def _get_anthropic_models_json( diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 079871830..414638a7c 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -484,12 +484,14 @@ 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_error: Callable[[str], 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_error` receives the walk's + failure/timeout reason (partial results were already streamed via `on_services`).""" token = get_databricks_token(workspace, profile) names, reason = list_all_mcp_services( workspace, @@ -499,6 +501,8 @@ def discover_all_mcp_service_names( cancel_event=cancel_event, ) raise_for_invalid_access_token(workspace, reason) + if reason is not None and on_error is not None: + on_error(reason) return names @@ -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,19 @@ 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 + the footer says why the list may be incomplete.""" def loader( append: Callable[[list[questionary.Choice]], None], cancel_event: threading.Event - ) -> None: + ) -> str | None: + failure: list[str] = [] + 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 +1349,15 @@ 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_error=failure.append, ) + if failure and "deadline exceeded" in failure[0]: + return "⚠ Timed out listing your MCP services; to add one not shown, use `ug mcp add`" + return None return loader diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 17eb7239c..d8f41f510 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1743,6 +1743,41 @@ def test_returns_reason_when_no_catalogs(self, monkeypatch): assert names == [] assert reason == "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), + ) + + 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, **kwargs): + collect((["main.default.exa"], None), 1, 2) + while time.monotonic() <= deadline: + time.sleep(0.01) + + monkeypatch.setattr(db_mod, "walk_catalog_schemas", slow_walk) + + names, reason = db_mod.list_all_mcp_services(WS, "token", deadline_seconds=0.05) + + assert names == ["main.default.exa"] + assert reason == "deadline exceeded while listing MCP services" + def _foundation_models_payload(names): return { diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 68f1fb39d..3860cae67 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1094,7 +1094,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_error=None, ): walk_calls.append(workspace) if on_services is not None: @@ -1112,6 +1117,58 @@ 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): + def fake_walk(workspace, profile=None, on_error=None, **kwargs): + on_error("deadline exceeded while listing MCP services") + return [] + + 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 = [] + + message = loader(streamed.extend, threading.Event()) + + assert streamed == [] + assert message is not None and "Timed out" 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.""" + + def fake_walk(workspace, profile=None, on_error=None, **kwargs): + on_error("no UC catalogs found") + return [] + + monkeypatch.setattr(mcp, "discover_all_mcp_service_names", fake_walk) + 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. From 56442acc0307c412cbac3457f254c2caf867f98e Mon Sep 17 00:00:00 2001 From: Achille Negrier Date: Sat, 17 Oct 2026 12:23:14 +0000 Subject: [PATCH 2/4] Address review: timed-out flag, clearer footer, stable ordering test Co-authored-by: Isaac --- src/ucode/databricks.py | 52 +++++++++++++++++++++++++++++----------- src/ucode/mcp.py | 12 ++++++---- tests/test_databricks.py | 25 +++++++++++++++++-- tests/test_mcp.py | 4 +++- 4 files changed, 72 insertions(+), 21 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 45f88987d..33e5ec682 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3097,6 +3097,17 @@ 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; callers match on these, not on text. +_WALK_TIMEOUT_CATALOGS_REASON = "deadline exceeded while listing UC catalogs" +_WALK_TIMEOUT_SCHEMAS_REASON = "deadline exceeded while listing UC schemas" +MCP_SERVICES_WALK_TIMEOUT_REASON = "deadline exceeded while listing MCP services" +MCP_SERVICES_WALK_TIMEOUT_REASONS = frozenset( + { + _WALK_TIMEOUT_CATALOGS_REASON, + _WALK_TIMEOUT_SCHEMAS_REASON, + MCP_SERVICES_WALK_TIMEOUT_REASON, + } +) # 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 +3122,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 +3155,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 +3211,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[[], None] | None = None, ) -> str | None: """Discover every user `.` in the workspace and probe each one in parallel. @@ -3210,6 +3224,9 @@ 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`, if given, is called when the probe phase stopped (deadline or cancel) with + schemas left unprobed, so callers can flag partial results without re-checking the clock. + 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 +3253,7 @@ 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" + return _WALK_TIMEOUT_CATALOGS_REASON schema_refs: list[tuple[str, str]] = [] @@ -3263,7 +3280,7 @@ def collect_schemas(result, catalog): if not schema_refs: if time.monotonic() > deadline: - return "deadline exceeded while listing UC schemas" + 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 @@ -3279,13 +3296,15 @@ 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() return None @@ -3315,6 +3334,11 @@ def list_all_mcp_services( (single schema).""" deadline = time.monotonic() + deadline_seconds names: set[str] = set() + incomplete = False + + def mark_incomplete() -> None: + nonlocal incomplete + incomplete = True def probe(catalog, schema): return list_mcp_services(workspace, token, f"{catalog}.{schema}") @@ -3329,20 +3353,20 @@ def collect(result, schemas_done, schemas_total): on_services(sorted(new)) reason = walk_catalog_schemas( - workspace, token, deadline=deadline, probe=probe, collect=collect, cancel_event=cancel_event + workspace, + token, + deadline=deadline, + probe=probe, + collect=collect, + cancel_event=cancel_event, + on_incomplete=mark_incomplete, ) if reason is not None: return [], reason - timed_out = time.monotonic() > deadline if not names: - return ( - [], - "deadline exceeded while listing MCP services" - if timed_out - else "no MCP services found", - ) + return [], MCP_SERVICES_WALK_TIMEOUT_REASON if incomplete else "no MCP services found" # Partial results are still returned, but flagged so callers can say the list is incomplete. - return sorted(names), "deadline exceeded while listing MCP services" if timed_out else None + return sorted(names), MCP_SERVICES_WALK_TIMEOUT_REASON if incomplete else None def _get_anthropic_models_json( diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 414638a7c..2d3ef410c 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -20,6 +20,7 @@ from ucode.constants import MCP_CLEANUP_SCOPES, MCP_USER_SCOPE from ucode.databricks import ( AIGW_MCP_SERVICES_SEGMENT, + MCP_SERVICES_WALK_TIMEOUT_REASONS, AuthTokenError, PermissionDeniedError, apply_pat_environment, @@ -1328,8 +1329,8 @@ def _mcp_services_background_loader( 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. 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 - the footer says why the list may be incomplete.""" + 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 @@ -1355,8 +1356,11 @@ def on_services(new_names: list[str]) -> None: cancel_event=cancel_event, on_error=failure.append, ) - if failure and "deadline exceeded" in failure[0]: - return "⚠ Timed out listing your MCP services; to add one not shown, use `ug mcp add`" + if failure and failure[0] in MCP_SERVICES_WALK_TIMEOUT_REASONS: + 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 d8f41f510..f6e6cd7ea 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1760,23 +1760,44 @@ def test_probes_main_default_first(self, monkeypatch): 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, **kwargs): + 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() monkeypatch.setattr(db_mod, "walk_catalog_schemas", slow_walk) names, reason = db_mod.list_all_mcp_services(WS, "token", deadline_seconds=0.05) assert names == ["main.default.exa"] - assert reason == "deadline exceeded while listing MCP services" + assert reason == db_mod.MCP_SERVICES_WALK_TIMEOUT_REASON + + 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) + + names, reason = db_mod.list_all_mcp_services(WS, "token", deadline_seconds=0.05) + + assert names == ["main.default.exa"] + assert reason is None def _foundation_models_payload(names): diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 3860cae67..7b360cca5 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 MCP_SERVICES_WALK_TIMEOUT_REASON WS = "https://example.databricks.com" CLAUDE_STATE = {"workspace": WS, "available_tools": ["claude"]} @@ -1139,7 +1140,7 @@ def fake_walk(workspace, profile=None, on_services=None, **kwargs): def test_walk_timeout_keeps_system_ai_and_warns(self, monkeypatch): def fake_walk(workspace, profile=None, on_error=None, **kwargs): - on_error("deadline exceeded while listing MCP services") + on_error(MCP_SERVICES_WALK_TIMEOUT_REASON) return [] monkeypatch.setattr(mcp, "discover_all_mcp_service_names", fake_walk) @@ -1151,6 +1152,7 @@ def fake_walk(workspace, profile=None, on_error=None, **kwargs): assert streamed == [] 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" ] From 1425dbc96e170ec34683d44ba5d28a5ea1c6b1bd Mon Sep 17 00:00:00 2001 From: Achille Negrier Date: Sat, 17 Oct 2026 13:28:01 +0000 Subject: [PATCH 3/4] Address review round 2: cancel vs timeout, explicit result type, real-walk test Co-authored-by: Isaac --- src/ucode/databricks.py | 49 +++++++++++------ src/ucode/mcp.py | 8 +-- tests/test_databricks.py | 110 ++++++++++++++++++++++++++++++++------- 3 files changed, 129 insertions(+), 38 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 33e5ec682..3e61cd45f 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3211,7 +3211,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[[], None] | None = None, + on_incomplete: Callable[[bool], None] | None = None, ) -> str | None: """Discover every user `.` in the workspace and probe each one in parallel. @@ -3224,8 +3224,10 @@ 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`, if given, is called when the probe phase stopped (deadline or cancel) with - schemas left unprobed, so callers can flag partial results without re-checking the clock. + `on_incomplete(timed_out)`, if given, is called when the probe phase stopped with schemas left + unprobed, 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.""" @@ -3304,11 +3306,24 @@ def collect_probe(result, _ref): should_stop=should_stop, ) if not drained and on_incomplete is not None: - on_incomplete() + 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, @@ -3317,7 +3332,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, @@ -3334,11 +3349,11 @@ def list_all_mcp_services( (single schema).""" deadline = time.monotonic() + deadline_seconds names: set[str] = set() - incomplete = False + timed_out = False - def mark_incomplete() -> None: - nonlocal incomplete - incomplete = True + 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}") @@ -3352,7 +3367,7 @@ def collect(result, schemas_done, schemas_total): if on_services is not None and new: on_services(sorted(new)) - reason = walk_catalog_schemas( + error = walk_catalog_schemas( workspace, token, deadline=deadline, @@ -3361,12 +3376,14 @@ def collect(result, schemas_done, schemas_total): cancel_event=cancel_event, on_incomplete=mark_incomplete, ) - if reason is not None: - return [], reason - if not names: - return [], MCP_SERVICES_WALK_TIMEOUT_REASON if incomplete else "no MCP services found" - # Partial results are still returned, but flagged so callers can say the list is incomplete. - return sorted(names), MCP_SERVICES_WALK_TIMEOUT_REASON if incomplete else None + if error is not None: + return McpServicesWalk([], error=error) + # 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, + ) def _get_anthropic_models_json( diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 2d3ef410c..34e62b65e 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -20,6 +20,7 @@ from ucode.constants import MCP_CLEANUP_SCOPES, MCP_USER_SCOPE from ucode.databricks import ( AIGW_MCP_SERVICES_SEGMENT, + MCP_SERVICES_WALK_TIMEOUT_REASON, MCP_SERVICES_WALK_TIMEOUT_REASONS, AuthTokenError, PermissionDeniedError, @@ -494,17 +495,18 @@ def discover_all_mcp_service_names( service names into the picker as the walk progresses. `on_error` receives the walk's failure/timeout reason (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) + raise_for_invalid_access_token(workspace, walk.error) + reason = walk.error or (MCP_SERVICES_WALK_TIMEOUT_REASON if walk.timed_out else None) if reason is not None and on_error is not None: on_error(reason) - return names + return walk.names def _normalize_workspace_title(text: str) -> str: diff --git a/tests/test_databricks.py b/tests/test_databricks.py index f6e6cd7ea..5a9e40285 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,20 @@ 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( @@ -1777,14 +1778,84 @@ def slow_walk(workspace, token, *, deadline, probe, collect, on_incomplete, **kw collect((["main.default.exa"], None), 1, 2) while time.monotonic() <= deadline: time.sleep(0.01) - on_incomplete() + on_incomplete(True) monkeypatch.setattr(db_mod, "walk_catalog_schemas", slow_walk) - names, reason = db_mod.list_all_mcp_services(WS, "token", deadline_seconds=0.05) + 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_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 names == ["main.default.exa"] - assert reason == db_mod.MCP_SERVICES_WALK_TIMEOUT_REASON + 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): @@ -1794,10 +1865,11 @@ def drained_walk(workspace, token, *, deadline, probe, collect, **kwargs): monkeypatch.setattr(db_mod, "walk_catalog_schemas", drained_walk) - names, reason = db_mod.list_all_mcp_services(WS, "token", deadline_seconds=0.05) + walk = db_mod.list_all_mcp_services(WS, "token", deadline_seconds=0.05) - assert names == ["main.default.exa"] - assert reason is None + assert walk.names == ["main.default.exa"] + assert not walk.timed_out + assert walk.error is None def _foundation_models_payload(names): From a1dd7f81a9f2da0da8d759e1daca43cc5c6f8530 Mon Sep 17 00:00:00 2001 From: Achille Negrier Date: Sat, 17 Oct 2026 14:15:15 +0000 Subject: [PATCH 4/4] Address review round 3: pass timed_out through to the picker Co-authored-by: Isaac --- src/ucode/databricks.py | 21 +++++++++------------ src/ucode/mcp.py | 23 ++++++++++++----------- tests/test_databricks.py | 15 +++++++++++++++ tests/test_mcp.py | 30 ++++++++++++++++++------------ 4 files changed, 54 insertions(+), 35 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 3e61cd45f..c598480e0 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3097,17 +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; callers match on these, not on text. +# 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" -MCP_SERVICES_WALK_TIMEOUT_REASON = "deadline exceeded while listing MCP services" -MCP_SERVICES_WALK_TIMEOUT_REASONS = frozenset( - { - _WALK_TIMEOUT_CATALOGS_REASON, - _WALK_TIMEOUT_SCHEMAS_REASON, - MCP_SERVICES_WALK_TIMEOUT_REASON, - } -) # Skip UC catalogs whose schemas almost never carry user-callable functions # you'd want to expose as agent tools. _UC_FUNCTIONS_SKIP_CATALOGS = frozenset( @@ -3224,8 +3217,8 @@ 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 probe phase stopped with schemas left - unprobed, so callers can flag partial results without re-checking the clock. `timed_out` is + `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. @@ -3255,6 +3248,8 @@ def should_stop() -> bool: if not catalog_names: return "no user UC catalogs found" if time.monotonic() > deadline: + if on_incomplete is not None: + on_incomplete(True) return _WALK_TIMEOUT_CATALOGS_REASON schema_refs: list[tuple[str, str]] = [] @@ -3282,6 +3277,8 @@ def collect_schemas(result, catalog): if not schema_refs: if time.monotonic() > deadline: + if on_incomplete is not None: + on_incomplete(True) return _WALK_TIMEOUT_SCHEMAS_REASON return "no UC schemas found" @@ -3377,7 +3374,7 @@ def collect(result, schemas_done, schemas_total): on_incomplete=mark_incomplete, ) if error is not None: - return McpServicesWalk([], error=error) + 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), diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 34e62b65e..f05ca5279 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -20,8 +20,6 @@ from ucode.constants import MCP_CLEANUP_SCOPES, MCP_USER_SCOPE from ucode.databricks import ( AIGW_MCP_SERVICES_SEGMENT, - MCP_SERVICES_WALK_TIMEOUT_REASON, - MCP_SERVICES_WALK_TIMEOUT_REASONS, AuthTokenError, PermissionDeniedError, apply_pat_environment, @@ -486,14 +484,14 @@ 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_error: Callable[[str], None] | 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. `on_error` receives the walk's - failure/timeout reason (partial results were already streamed via `on_services`).""" + 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) walk = list_all_mcp_services( workspace, @@ -503,9 +501,8 @@ def discover_all_mcp_service_names( cancel_event=cancel_event, ) raise_for_invalid_access_token(workspace, walk.error) - reason = walk.error or (MCP_SERVICES_WALK_TIMEOUT_REASON if walk.timed_out else None) - if reason is not None and on_error is not None: - on_error(reason) + if walk.timed_out and on_timeout is not None: + on_timeout() return walk.names @@ -1337,7 +1334,11 @@ def _mcp_services_background_loader( def loader( append: Callable[[list[questionary.Choice]], None], cancel_event: threading.Event ) -> str | None: - failure: list[str] = [] + 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`), @@ -1356,9 +1357,9 @@ def on_services(new_names: list[str]) -> None: profile, on_services=on_services, cancel_event=cancel_event, - on_error=failure.append, + on_timeout=on_timeout, ) - if failure and failure[0] in MCP_SERVICES_WALK_TIMEOUT_REASONS: + if timed_out: return ( "⚠ Timed out listing your MCP services; to add one not shown, use " "`ug mcp add --names `" diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 5a9e40285..3d2de72d5 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1788,6 +1788,21 @@ def slow_walk(workspace, token, *, deadline, probe, collect, on_incomplete, **kw 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) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 7b360cca5..685bccadb 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -9,7 +9,7 @@ from ucode import mcp from ucode.agents import claude -from ucode.databricks import MCP_SERVICES_WALK_TIMEOUT_REASON +from ucode.databricks import McpServicesWalk WS = "https://example.databricks.com" CLAUDE_STATE = {"workspace": WS, "available_tools": ["claude"]} @@ -1100,7 +1100,7 @@ def fake_walk( on_progress=None, on_services=None, cancel_event=None, - on_error=None, + on_timeout=None, ): walk_calls.append(workspace) if on_services is not None: @@ -1139,18 +1139,23 @@ def fake_walk(workspace, profile=None, on_services=None, **kwargs): assert sorted(values) == [f"{prefix}main.default.exa", f"{prefix}system.ai.github"] def test_walk_timeout_keeps_system_ai_and_warns(self, monkeypatch): - def fake_walk(workspace, profile=None, on_error=None, **kwargs): - on_error(MCP_SERVICES_WALK_TIMEOUT_REASON) - return [] + """A timed-out walk (flag set on `McpServicesWalk`) streams partial results and warns.""" + monkeypatch.setattr(mcp, "get_databricks_token", lambda workspace, profile=None: "token") - monkeypatch.setattr(mcp, "discover_all_mcp_service_names", fake_walk) + 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 streamed == [] + 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] == [ @@ -1160,11 +1165,12 @@ def fake_walk(workspace, profile=None, on_error=None, **kwargs): 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.""" - def fake_walk(workspace, profile=None, on_error=None, **kwargs): - on_error("no UC catalogs found") - return [] - - monkeypatch.setattr(mcp, "discover_all_mcp_service_names", fake_walk) + 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 = []