Repository navigation
Surface user-owned MCP services in the configure picker - #1162
nachill-databricks wants to merge 4 commits into
Conversation
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 <no-reply@databricks.com>
|
/ug-review |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| ) | ||
| if reason is not None: | ||
| return [], reason | ||
| timed_out = time.monotonic() > deadline |
There was a problem hiding this comment.
what happens if every schema drains but the clock passes the deadline right after? the full list then gets flagged incomplete. can the walk report whether items were left unprobed instead of re-checking the 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.""" |
There was a problem hiding this comment.
the docstring says the footer explains timeouts and errors, but only timeouts show a warning (test_walk_error_degrades_to_system_ai_only checks that errors stay silent). can we fix the docstring to match?
| cancel_event=cancel_event, | ||
| on_error=failure.append, | ||
| ) | ||
| if failure and "deadline exceeded" in failure[0]: |
There was a problem hiding this comment.
can we avoid matching on the free-form "deadline exceeded" text from databricks.py? if that wording changes, the warning stops showing without any error, and the test wouldn't catch it because it passes in the same literal. a shared constant or a timed-out flag would fix both.
| 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`" |
There was a problem hiding this comment.
this picker is often already ug mcp add, and running it again just reopens the same walk and times out the same way. can we point to ug mcp add --names <catalog.schema.name> instead?
Suggestion:
| return "⚠ Timed out listing your MCP services; to add one not shown, use `ug mcp add`" | |
| return ( | |
| "⚠ Timed out listing your MCP services; to add one not shown, use " | |
| "`ug mcp add --names <catalog.schema.name>`" | |
| ) |
| 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" |
There was a problem hiding this comment.
can this be flaky? probes run on several worker threads, so main.alpha can append before main.default even when dequeued second. can we assert on the dequeue order (e.g. max_workers=1) instead?
Co-authored-by: Isaac <no-reply@databricks.com>
|
/ug-review |
| if not drained and on_incomplete is not None: | ||
| on_incomplete() |
There was a problem hiding this comment.
should_stop also covers cancel_event, so a cancel (e.g. the picker closing) calls on_incomplete and gets reported as "deadline exceeded". can we only flag when time.monotonic() > deadline, or name the reason so it covers both?
| return sorted(names), None | ||
| 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 |
There was a problem hiding this comment.
(names, reason) now means "failure" when names is empty but "warning" when it isn't, and callers can only tell by checking both. could we return a small dataclass with an explicit incomplete flag?
| 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): |
There was a problem hiding this comment.
both new partial-result tests replace walk_catalog_schemas with a stub, so nothing tests the real path where _collect_concurrently returns False and that triggers on_incomplete. can we run the real walk with a probe that blocks until the deadline?
Polly AI ReviewBlocking issues
Security vulnerabilitiesNone identified in the changed code. Non-blocking notes
ApproachThe approach fits the existing walker and picker-message interfaces. Tracking whether collection drained avoids falsely reporting a timeout merely because the clock passed the deadline afterward. No broader refactor is needed. SummaryDiff size (computed): 4 files; +187 / -17 lines (204 changed text lines); 0 binary files.
Shares use additions + deletions, including generated text and lockfiles. Binary files have no line count. Tests include fixtures/helpers under test paths and colocated test/spec files; categories are path-based, not coverage measurements. ChangesPrioritizes The size is proportional to the outcome. One concrete simplification is to exercise probe priority directly through the existing walker test setup with TestsTest-by-test assessment
No tests were removed. Existing walker deadline, cancellation, successful-drain, pagination, and probe-error coverage provides relevant foundations. Extend the existing deadline/full-success cases to verify the new callback. A small picker-rendering test at 80 columns would catch the blocking truncation; live-workspace E2E coverage is unnecessary for that assertion. The PR reports passing focused tests and lint; these were not independently executed. Review used the checksum-verified frozen diff and trusted default-branch context because the supplied base/head commit objects were unavailable locally. Human verification: open ScopeAll changes support discovery priority and partial-result reporting. The shared walker also changes skills-discovery probe order, but no incompatible behavior or independent outcome was identified. No clearly unrelated changes or unresolved scope questions were found. Automated review by Polly · workflow run |
…-walk test Co-authored-by: Isaac <no-reply@databricks.com>
|
/ug-review |
| # 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" |
There was a problem hiding this comment.
list_all_mcp_services doesn't return this anymore. only mcp.py builds it from timed_out. can we move it to mcp.py, or drop it from the reasons set so the two timeout signals don't overlap?
| 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) |
There was a problem hiding this comment.
can we pass the walk's timed_out flag through instead of turning it back into a string? the loader then string-matches against MCP_SERVICES_WALK_TIMEOUT_REASONS, so any new or renamed timeout reason silently drops the warning
| def test_flags_partial_results_past_deadline(self, monkeypatch): | ||
| def slow_walk(workspace, token, *, deadline, probe, collect, on_incomplete, **kwargs): |
There was a problem hiding this comment.
are the stubbed slow_walk/cancelled_walk tests necessary? the real-walk tests below cover the same paths. also, the sleep loop here does nothing because the stub calls on_incomplete(True) either way
| 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): |
There was a problem hiding this comment.
these tests all monkeypatch discover_all_mcp_service_names, so its timed_out→on_error mapping is never run. can we stub list_all_mcp_services returning McpServicesWalk(names, timed_out=True) instead?
Polly AI ReviewBlocking issuesNone identified. Security vulnerabilitiesNone identified in the changed code. Non-blocking notes
ApproachReusing the shared walker and existing picker-message channel is sound. A shorter footer avoids changing the shared picker layout. The ordering applies to queued probes after schema enumeration; it cannot guarantee discovery of SummaryDiff size (computed): 4 files; +298 / -37 lines (335 changed text lines); 0 binary files.
Shares use additions + deletions, including generated text and lockfiles. Binary files have no line count. Tests include fixtures/helpers under test paths and colocated test/spec files; categories are path-based, not coverage measurements. ChangesThe change prioritizes The implementation is proportional to the stated outcome. Most avoidable bulk is overlapping test setup, rather than unnecessary production changes. TestsTest-by-test assessmentThe PR reports 654 passing tests and a passing lint gate. These results were not independently executed. Assessment used the frozen diff and existing tests in the trusted checkout.
Nine test functions are added and five are modified; none are removed. No shared fixtures or helpers change. Added and modified local helpers are assessed with their parent cases above. Existing tests cover generic walker cancellation, deadline return, complete draining, raised probes, and picker deduplication. Existing assertions are retained. The concrete gaps are discovery-result-to-footer propagation, the remaining priority tiers, and recovery-hint visibility at standard terminal width. Unit/component coverage is sufficient for scheduling and callbacks; terminal rendering is the distinct UI boundary. A live workspace E2E is not necessary for those assertions. For manual verification, open the picker in an 80-column terminal against a workspace where discovery times out. Confirm the complete ScopeAll changed files support the stated discovery and timeout-reporting outcome. No clearly unrelated changes were found. The ordering change in Context came from the trusted default-branch checkout, which differs from the supplied base/head revisions. Target changes were assessed exclusively from the checksum-verified frozen diff. Automated review by Polly · workflow run |
Co-authored-by: Isaac <no-reply@databricks.com>
|
/ug-review |
| if not drained and on_incomplete is not None: | ||
| on_incomplete(time.monotonic() > deadline) |
There was a problem hiding this comment.
can we key this off cancel_event.is_set() instead of re-reading the clock? if the user cancels and the deadline passes before this line runs, the cancel gets reported as a timeout.
Suggestion:
| if not drained and on_incomplete is not None: | |
| on_incomplete(time.monotonic() > deadline) | |
| if not drained and on_incomplete is not None: | |
| on_incomplete(not cancel_event.is_set()) |
|
|
||
| assert not walk.timed_out | ||
|
|
||
| def test_fully_drained_walk_is_not_flagged_even_if_deadline_passes_after(self, monkeypatch): |
There was a problem hiding this comment.
is this test necessary? the stub never calls on_incomplete, so it passes whether or not the real walk re-checks the clock. same goes for the sleep loop in test_flags_partial_results_past_deadline. can we drop them or use the real walk?
Polly AI ReviewBlocking issuesNone identified in the frozen diff. Security vulnerabilitiesNone identified. Token validation remains ahead of timeout notification; no new credential handling, network destinations, or configuration writes are introduced. Non-blocking notes
ApproachThe approach follows existing patterns: structured walk results carry timeout metadata, and the loader returns a footer message through the existing picker interface. Prioritization starts after catalog/schema enumeration, so it cannot guarantee discovery of SummaryDiff size (computed): 4 files; +317 / -37 lines (354 changed text lines); 0 binary files.
Shares use additions + deletions, including generated text and lockfiles. Binary files have no line count. Tests include fixtures/helpers under test paths and colocated test/spec files; categories are path-based, not coverage measurements. ChangesThe implementation prioritizes schema probes, preserves partial service results, distinguishes timeouts from cancellation, and supplies an actionable picker warning. The remaining changes adapt existing tests and add ordering, timeout, cancellation, and loader coverage. No documentation, dependency, or generated-file changes appear. The implementation size is reasonable. The concrete reduction opportunity is overlapping test setup, assessed below. TestsTest-by-test assessmentStatic assessment only; tests and lint were not run. The reported passing results remain author-provided evidence. The frozen base/head commits were unavailable locally, so surrounding context was checked against the trusted default-branch checkout rather than the exact target tree. Databricks cases belong to
Existing picker message handling provides the rendering contract; a new end-to-end test is not necessary for the loader-return assertion. The concrete missing coverage is real early-deadline notification and pre-probe cancellation. For verification in a trusted test environment, run ScopeNo clearly unrelated changes were found. The priority sort lives in the shared catalog walker, so skills discovery also receives the new probe order; no separate skills behavior change was identified. The implementation and test changes otherwise support the stated discovery and partial-results outcome. Automated review by Polly · workflow run |
2 participants
Problem
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.exacould be missing with no hint why.Change
Probe the
maincatalog anddefaultschemas first, flag partial results past the deadline, and show a timeout notice in the picker footer pointing toug mcp add.Before: services silently missing. After: user-owned services surface earlier, and a timeout notice appears when the listing is partial.
Testing
tests/test_databricks.pyandtests/test_mcp.py.uv run pytest -q tests/test_databricks.py tests/test_mcp.py: 654 passed.uvx --from rust-just just lint: passes (ruff check, ruff format --check, ty check).tests/test_agents_init.pyandtests/test_e2e_user_agent.pyfail onmaintoo, independent of this change.mainsilently dropped.Not yet verified
main.defaultservice.Before / after
Real runs in a sandbox HOME against a live workspace; workspace hostname and catalog name replaced with and . Picker filtered to "figma".
Before (main):
After (this branch): listing timed out on this workspace, so the new fallback note shows. The test service lives outside
main.default, so the new probe order isn't visible here; the unit tests cover it.This pull request and its description were written by Isaac.