From 76bc51ac1b29cda1aa713742420e86261ad24266 Mon Sep 17 00:00:00 2001 From: Arthur Jenoudet Date: Wed, 2 Sep 2026 15:42:15 +0000 Subject: [PATCH] Add interactive metastore skill picker --- README.md | 7 +- src/ucode/cli.py | 11 ++- src/ucode/skills_download.py | 95 +++++++++++++++++++++++++ tests/test_cli.py | 14 +++- tests/test_skills_download.py | 129 ++++++++++++++++++++++++++++++++++ 5 files changed, 250 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 82331003..2d905346 100644 --- a/README.md +++ b/README.md @@ -226,9 +226,13 @@ you to run `ucode ` (existing agent sessions need a restart before the MC `ucode skill add` registers skills additively, keeping anything already configured. With `--mcp` it adds the schemas to the connection's scope, otherwise it downloads their skills to disk. `--skills` -narrows a download to a subset of one schema's skills. +narrows a download to a subset of one schema's skills. With no selection flags, a searchable picker +lists finalized skills visible in the metastore. ```bash +# Browse the metastore and choose skills to download. +ucode skill add + # Add schemas to the skills MCP scope, keeping any already configured. ucode skill add --location main.default,ml.prod --mcp @@ -374,6 +378,7 @@ The output looks like: | `ucode configure skills --location main.default [--path ]` | Download a schema's skills to disk (under ``, or your home dir) and register a schema-less skills MCP connection | | `ucode configure skills --location main.default --skill my-skill` | Download only the named skill(s) from a schema (comma-separated for several) | | `ucode configure skills --location main.default --mcp` | Expose a schema's skills as MCP tools (override-only) instead of downloading | +| `ucode skill add` | Interactively choose finalized metastore skills to download | | `ucode skill add --location main.default --mcp` | Add schemas to the skills MCP scope, keeping any already configured (additive; never replaces) | | `ucode skill add --location main.default` | Download a schema's skills to disk without removing existing downloads | | `ucode skill add --skills main.default.my-skill` | Download a named subset of skills (bare names need `--location`; fully-qualified names stand alone) | diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 000f1db6..8b997ba5 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -109,6 +109,7 @@ ) from ucode.skills_download import ( configure_skills_download_command, + configure_skills_download_interactive_command, download_managed_skills_on_launch, ) from ucode.smart_routing import v2 as smart_routing_v2 @@ -1358,8 +1359,9 @@ def skills_add( With ``--mcp``, adds the given schemas to the skills MCP connection's scope. Otherwise downloads each schema's skills to project-level skill directories under ``--path``, or to user-level skill directories when omitted, keeping - already-downloaded skills. ``--skills`` narrows a download to a subset of one - schema's skills, by bare name (with ``--location``) or fully-qualified + already-downloaded skills. With no selection flags, opens a searchable + metastore picker. ``--skills`` narrows a download to a subset of one schema's skills, + by bare name (with ``--location``) or fully-qualified ``..``. """ try: @@ -1402,7 +1404,10 @@ def skills_add( ) locations = list(schemas) if not locations: - raise RuntimeError("--location is required for `ucode skill add`.") + if mcp: + raise RuntimeError("--location is required when using --mcp.") + configure_skills_download_interactive_command(path=path) + return if requested_skills is not None and len(locations) != 1: raise RuntimeError( f"--skills requires a single --location (got: {', '.join(locations)})." diff --git a/src/ucode/skills_download.py b/src/ucode/skills_download.py index 2168b522..dab2daf4 100644 --- a/src/ucode/skills_download.py +++ b/src/ucode/skills_download.py @@ -21,6 +21,7 @@ print_success, print_warning, progress_bar, + prompt_for_multi_selection, prompt_yes_no, ) @@ -52,6 +53,18 @@ class SkillRef: bundle_name: str +@dataclass(frozen=True) +class MetastoreSkill: + """A finalized skill discovered without a schema filter.""" + + full_name: str + ref: SkillRef + + @property + def location(self) -> str: + return self.full_name.rsplit(".", 1)[0] + + def _non_empty_str(value: object) -> str | None: """``value`` when it is a non-empty string, else None.""" return value if isinstance(value, str) and value else None @@ -118,6 +131,58 @@ def list_schema_skills( return refs, None +def list_metastore_skills(workspace: str, token: str) -> tuple[list[MetastoreSkill], str | None]: + """List finalized, downloadable skills visible in the current metastore.""" + hostname = workspace_hostname(workspace) + base_url = f"https://{hostname}/api/2.1/unity-catalog/skills" + + skills: list[MetastoreSkill] = [] + page_token: str | None = None + while True: + url = base_url + if page_token: + url = f"{url}?{urlencode({'page_token': page_token})}" + payload, reason = _http_get_json(url, token, timeout=30) + if payload is None: + return [], reason + data = payload if isinstance(payload, dict) else {} + for skill in data.get("skills") or []: + if not isinstance(skill, dict): + continue + resource_name = _non_empty_str(skill.get("name")) + ref = _skill_ref(skill) + if ref is None: + continue + full_name = resource_name.removeprefix("skills/") if resource_name else None + if full_name is None or full_name.count(".") != 2: + print_warning( + f"Skipping `{resource_name or ''}`: expected a fully-qualified " + "`..` from the skills API." + ) + continue + skills.append(MetastoreSkill(full_name=full_name, ref=ref)) + page_token = data.get("next_page_token") + if not page_token: + return sorted(skills, key=lambda skill: skill.full_name.lower()), None + + +def prompt_for_skill_download( + skills: list[MetastoreSkill], +) -> list[MetastoreSkill] | None: + """Select metastore skills to download, or return None when cancelled.""" + by_name = {skill.full_name: skill for skill in skills} + options = [] + for skill in skills: + label = skill.full_name + if skill.ref.bundle_name != skill.ref.securable_name: + label = f"{label} (bundle: {skill.ref.bundle_name})" + options.append((skill.full_name, label)) + selected = prompt_for_multi_selection("Skills:", options, searchable=True) + if selected is None: + return None + return [by_name[name] for name in selected if name in by_name] + + def list_skill_files( workspace: str, token: str, catalog: str, schema: str, securable: str ) -> tuple[list[str], str | None]: @@ -450,3 +515,33 @@ def configure_skills_download_command( register_schemaless_skills_connection(state, workspace, profile, clients) return 0 + + +def configure_skills_download_interactive_command(*, path: str | None) -> int: + """Discover metastore skills, let the user select some, and download them.""" + state = load_state() + workspace, profile, clients = setup_mcp_clients(state, "Add Skills") + token = get_databricks_token(workspace, profile) + + available, reason = list_metastore_skills(workspace, token) + if reason: + raise RuntimeError(f"Could not list workspace skills: {reason}.") + if not available: + print_note("No finalized skills are available to download in this metastore.") + return 0 + + selected = prompt_for_skill_download(available) + if selected is None: + return 0 + if not selected: + print_note("No skills selected. Press space to toggle an item, then enter to download.") + return 0 + + selected_by_location: dict[str, set[str]] = {} + for skill in selected: + selected_by_location.setdefault(skill.location, set()).add(skill.ref.securable_name) + for location, securable_names in selected_by_location.items(): + download_skills(workspace, token, [location], path, securable_names) + + register_schemaless_skills_connection(state, workspace, profile, clients) + return 0 diff --git a/tests/test_cli.py b/tests/test_cli.py index 110dd6a9..6efa8322 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1131,16 +1131,26 @@ def test_fully_qualified_skills_across_schemas_exit_1(self): assert "must all share one" in _strip_ansi(result.output) mock_download.assert_not_called() - def test_without_location_exit_1(self): + def test_without_location_opens_download_picker(self): with ( patch("ucode.cli.add_skills_command") as mock_add, patch("ucode.cli.configure_skills_download_command") as mock_download, + patch("ucode.cli.configure_skills_download_interactive_command") as mock_interactive, ): result = runner.invoke(app, ["skill", "add"]) + + assert result.exit_code == 0, result.output + mock_add.assert_not_called() + mock_download.assert_not_called() + mock_interactive.assert_called_once_with(path=None) + + def test_mcp_without_location_exit_1(self): + with patch("ucode.cli.add_skills_command") as mock_add: + result = runner.invoke(app, ["skill", "add", "--mcp"]) + assert result.exit_code == 1 assert "--location is required" in _strip_ansi(result.output) mock_add.assert_not_called() - mock_download.assert_not_called() def test_skill_with_mcp_exit_1(self): with ( diff --git a/tests/test_skills_download.py b/tests/test_skills_download.py index c7927019..f333fce0 100644 --- a/tests/test_skills_download.py +++ b/tests/test_skills_download.py @@ -156,6 +156,59 @@ def test_http_failure_propagates_reason(self, monkeypatch): assert reason == "HTTP 500 Server Error" +class TestListMetastoreSkills: + def test_lists_finalized_skills_across_schemas_and_follows_pagination(self, monkeypatch): + pages = [ + { + "skills": [ + { + "name": "skills/ml.prod.triage", + "bundle_name": "triage", + "finalize_time": "t", + }, + {"name": "skills/ml.prod.draft", "bundle_name": "draft"}, + ], + "next_page_token": "next", + }, + { + "skills": [ + { + "name": "skills/main.default.pii", + "bundle_name": "pii-handling", + "finalize_time": "t", + } + ] + }, + ] + urls = [] + + def fake_get(url, token, timeout=30): + urls.append(url) + return pages.pop(0), None + + monkeypatch.setattr(sd, "_http_get_json", fake_get) + + skills, reason = sd.list_metastore_skills(WS, "token") + + assert reason is None + assert skills == [ + sd.MetastoreSkill("main.default.pii", ref("pii", "pii-handling")), + sd.MetastoreSkill("ml.prod.triage", ref("triage")), + ] + assert urls == [ + f"{WS}/api/2.1/unity-catalog/skills", + f"{WS}/api/2.1/unity-catalog/skills?page_token=next", + ] + + def test_http_failure_propagates_reason(self, monkeypatch): + monkeypatch.setattr(sd, "_http_get_json", lambda *a, **k: (None, "HTTP 500 Server Error")) + + skills, reason = sd.list_metastore_skills(WS, "token") + + assert skills == [] + assert reason == "HTTP 500 Server Error" + + class TestListSkillFiles: def test_lists_under_the_skills_place(self, monkeypatch): captured = {} @@ -744,3 +797,79 @@ def test_skills_filter_threads_through(self, monkeypatch): assert calls["download"] == (WS, "token", ["a.b"], None, {"triage"}) assert calls["register"] == (WS, "profile", ["claude"]) + + +class TestConfigureSkillsDownloadInteractiveCommand: + def _stub(self, monkeypatch): + calls: dict[str, object] = {"downloads": []} + state = {"state": True} + monkeypatch.setattr(sd, "load_state", lambda: state) + monkeypatch.setattr( + sd, "setup_mcp_clients", lambda actual, section: (WS, "profile", ["claude"]) + ) + monkeypatch.setattr(sd, "get_databricks_token", lambda ws, profile: "token") + monkeypatch.setattr( + sd, + "download_skills", + lambda ws, token, locations, path, skills: calls["downloads"].append( + (ws, token, locations, path, skills) + ), + ) + monkeypatch.setattr( + sd, + "register_schemaless_skills_connection", + lambda actual, ws, profile, clients: calls.update( + register=(actual, ws, profile, clients) + ), + ) + return state, calls + + def test_downloads_picker_selection_grouped_by_schema(self, monkeypatch): + state, calls = self._stub(monkeypatch) + available = [ + sd.MetastoreSkill("main.default.pii", ref("pii")), + sd.MetastoreSkill("main.default.triage", ref("triage")), + sd.MetastoreSkill("ml.prod.eval", ref("eval")), + ] + monkeypatch.setattr(sd, "list_metastore_skills", lambda *a: (available, None)) + monkeypatch.setattr( + sd, "prompt_for_skill_download", lambda actual: [available[0], available[2]] + ) + + assert sd.configure_skills_download_interactive_command(path="/tmp/project") == 0 + + assert calls["downloads"] == [ + (WS, "token", ["main.default"], "/tmp/project", {"pii"}), + (WS, "token", ["ml.prod"], "/tmp/project", {"eval"}), + ] + assert calls["register"] == (state, WS, "profile", ["claude"]) + + def test_cancel_is_a_noop(self, monkeypatch): + _, calls = self._stub(monkeypatch) + available = [sd.MetastoreSkill("main.default.pii", ref("pii"))] + monkeypatch.setattr(sd, "list_metastore_skills", lambda *a: (available, None)) + monkeypatch.setattr(sd, "prompt_for_skill_download", lambda actual: None) + + assert sd.configure_skills_download_interactive_command(path=None) == 0 + + assert calls["downloads"] == [] + assert "register" not in calls + + def test_picker_is_searchable_and_shows_bundle_name_when_different(self, monkeypatch): + captured = {} + skill = sd.MetastoreSkill("main.default.task", ref("task", "task-triage")) + monkeypatch.setattr( + sd, + "prompt_for_multi_selection", + lambda prompt, options, searchable: ( + captured.update(prompt=prompt, options=options, searchable=searchable) + or ["main.default.task"] + ), + ) + + assert sd.prompt_for_skill_download([skill]) == [skill] + assert captured == { + "prompt": "Skills:", + "options": [("main.default.task", "main.default.task (bundle: task-triage)")], + "searchable": True, + }