diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 6c5c3b20..210b13c5 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -101,6 +101,7 @@ available_mcp_clients, configure_mcp_command, configure_skills_mcp_command, + configure_skills_mcp_picker_command, configured_mcp_clients, purge_cross_workspace_mcp_residue, remove_mcp_command, @@ -1308,8 +1309,8 @@ def skills_add( to user-level skill directories when omitted, keeping already-downloaded skills. ``--location`` downloads whole ``.`` schemas; ``--skills`` downloads a named set of fully-qualified skills that may span schemas (and takes - no ``--location``). With none of ``--mcp``/``--location``/``--skills`` on an - interactive terminal, opens a picker of the workspace's skills to download. + no ``--location``). With no ``--location``/``--skills`` on an interactive terminal, + opens a picker of the workspace's schemas to scope (``--mcp``) or skills to download. """ try: requested_skills = ( @@ -1340,15 +1341,23 @@ def skills_add( return locations = _parse_skill_locations(location) if not locations: - if not mcp and _stdin_is_interactive(): - configure_skills_download_picker_command(path=path) + if _stdin_is_interactive(): + if mcp: + configured_agents = ( + _configure_agents_for_mcp(sorted(requested_agents)) + if requested_agents + else None + ) + configure_skills_mcp_picker_command(agents=configured_agents) + else: + configure_skills_download_picker_command(path=path) return raise RuntimeError("--location is required for `ucode skill add`.") if mcp: - scope = ( + configured_agents = ( _configure_agents_for_mcp(sorted(requested_agents)) if requested_agents else None ) - add_skills_command(locations, agents=scope) + add_skills_command(locations, agents=configured_agents) else: configure_location_skills_download_command(locations, path=path) except (RuntimeError, ValueError) as exc: diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index b52ce58c..c890fae5 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -29,6 +29,12 @@ list_mcp_services, workspace_hostname, ) +from ucode.skills_api import ( + _SKILLS_WALK_DEADLINE_SECONDS, + _SKILLS_WALK_TIMEOUT_REASON, + SkillRef, + list_all_skills, +) from ucode.state import load_full_state, load_state, save_state from ucode.ui import ( _BACK, @@ -1975,6 +1981,28 @@ def _union_locations(base: list[str], new: list[str]) -> list[str]: return merged +def add_skill_locations_to_mcp( + state: dict, + workspace: str, + profile: str | None, + clients: list[str], + locations: list[str], +) -> None: + """Add ``locations`` to each client's skill MCP scope, keeping any already configured.""" + locations_by_client = _skill_locations_by_client_from_state(state) + for client in clients: + locations_by_client[client] = _union_locations( + locations_by_client.get(client, []), locations + ) + _update_skills_mcp(state, workspace, profile, clients, locations_by_client) + + +def configured_skill_locations(state: dict, clients: list[str]) -> set[str]: + """The union of skill schemas already in the MCP scope across ``clients``.""" + locations_by_client = _skill_locations_by_client_from_state(state) + return {location for client in clients for location in locations_by_client.get(client, [])} + + def add_skills_command(locations: list[str], agents: set[str] | None = None) -> int: """Add ``locations`` to each targeted client's skill scope, keeping any already configured. @@ -1983,12 +2011,84 @@ def add_skills_command(locations: list[str], agents: set[str] | None = None) -> the only thing ``--agents`` changes.""" state = load_state() workspace, profile, clients = setup_mcp_clients(state, "Add Skills MCP", agents=agents) - locations_by_client = _skill_locations_by_client_from_state(state) - for client in clients: - locations_by_client[client] = _union_locations( - locations_by_client.get(client, []), locations - ) - _update_skills_mcp(state, workspace, profile, clients, locations_by_client) + add_skill_locations_to_mcp(state, workspace, profile, clients, locations) + return 0 + + +def _skill_schema_choice(location: str, skill_count: int, in_scope: bool) -> questionary.Choice: + """Picker row for one schema: value is ``.``, title carries the skill count. + + An already-scoped schema is flagged and stays selectable; re-selecting it is a no-op, since + adding to the MCP scope is additive (removal is ``ug skill remove --mcp``). + """ + noun = "skill" if skill_count == 1 else "skills" + scope_flag = " (already in skill MCP)" if in_scope else "" + return questionary.Choice( + title=f"{location} ({skill_count} {noun}){scope_flag}", value=location + ) + + +def _skill_schema_background_loader( + workspace: str, token: str, in_scope: set[str] +) -> Callable[[Callable[[list[questionary.Choice]], None]], str | None]: + """A picker ``background_loader`` that streams the workspace-wide skill walk in as schema rows. + + ``list_all_skills`` probes one schema per call, so each ``on_skills`` batch is that schema's + complete skill set: one row per schema, carrying its exact skill count. + """ + + def loader(append: Callable[[list[questionary.Choice]], None]) -> str | None: + def on_skills(refs: list[SkillRef]) -> None: + location = f"{refs[0].catalog}.{refs[0].schema}" + append([_skill_schema_choice(location, len(refs), location in in_scope)]) + + found, reason = list_all_skills(workspace, token, on_skills=on_skills) + if reason == _SKILLS_WALK_TIMEOUT_REASON: + schemas = len({(ref.catalog, ref.schema) for ref in found}) + return ( + f"⚠️ Timed out after {int(_SKILLS_WALK_DEADLINE_SECONDS)}s, " + f"found {schemas} skill schemas" + ) + return None + + return loader + + +def prompt_for_skill_schema_choices( + background_loader: Callable[[Callable[[list[questionary.Choice]], None]], str | None], +) -> list[str] | None: + """Show the skill-schema picker, returning the selected schemas or None on Ctrl-C.""" + selection = scrolling_checkbox( + "Skill schemas:", + choices=[], + instruction="(space to toggle, ctrl-a all, enter to save, type to filter)", + style=picker_style(), + background_loader=background_loader, + loading_noun="skill schemas", + ).ask() + if selection is None: + return None + return [str(value) for value in selection] + + +def configure_skills_mcp_picker_command(agents: set[str] | None = None) -> int: + """Pick skill schemas from an interactive workspace-wide list and add them to the MCP scope. + + Opens the picker immediately and streams schemas in as discovery finds them. Ctrl-C changes + nothing. ``agents`` scopes the addition to that subset of configured clients. + """ + state = load_state() + workspace, profile, clients = setup_mcp_clients(state, "Add Skills MCP", agents=agents) + token = get_databricks_token(workspace, profile) + + loader = _skill_schema_background_loader( + workspace, token, configured_skill_locations(state, clients) + ) + locations = prompt_for_skill_schema_choices(loader) + if not locations: + return 0 + + add_skill_locations_to_mcp(state, workspace, profile, clients, locations) return 0 diff --git a/src/ucode/skills_api.py b/src/ucode/skills_api.py new file mode 100644 index 00000000..847fb102 --- /dev/null +++ b/src/ucode/skills_api.py @@ -0,0 +1,281 @@ +"""Read-only client for the Unity Catalog skills API: list and get skills, and fetch bundles. + +The workspace-facing read layer shared by the download flow (``skills_download``) and the +skills MCP scope picker (``mcp``); it depends only on ``databricks`` so neither of those has +to import the other. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from urllib.parse import urlencode + +from ucode.databricks import ( + _http_get_bytes, + _http_get_json, + walk_catalog_schemas, + workspace_hostname, +) +from ucode.ui import print_warning + +SKILL_FILES_API_PREFIX = "Skills" + +# Wall-clock budget for the workspace-wide skill walk; a slow workspace degrades +# to partial results instead of hanging the picker. +_SKILLS_WALK_DEADLINE_SECONDS = 30.0 +_SKILLS_WALK_TIMEOUT_REASON = "deadline exceeded while listing skills" + + +@dataclass(frozen=True) +class SkillRef: + """A downloadable skill's UC location plus its two non-interchangeable names. + + ``catalog``/``schema``/``securable_name`` are the parts of ``skills/..``: + ``securable_name`` is the leaf, the only name the Files API resolves, and the + three together fully qualify the skill (``fqn``). ``bundle_name`` is the + ``name:`` an agent reads from the bundle's SKILL.md frontmatter, so it names + the on-disk directory. Finalize does not require the securable and bundle name + to match, so a skill created under a securable that differs from its + frontmatter carries both. ``description`` is the skill's UC description, used only + to preview a skill in the interactive picker. ``metastore_id``, ``skill_id``, and + ``uc_update_time`` are UC attribution metadata recorded when the skill is + downloaded (see ``skills_state``); the download flow never reads them. + """ + + catalog: str + schema: str + securable_name: str + bundle_name: str + description: str | None = None + metastore_id: str | None = None + skill_id: str | None = None + uc_update_time: str | None = None + + @property + def fqn(self) -> str: + return f"{self.catalog}.{self.schema}.{self.securable_name}" + + +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 + + +def _is_safe_bundle_name(bundle_name: str) -> bool: + path = Path(bundle_name) + return len(path.parts) == 1 and path.parts[0] != ".." and not path.is_absolute() + + +def _skill_ref(skill: dict) -> SkillRef | None: + """A finalized skill's ``SkillRef``, or None if it cannot be downloaded. + + A skill without a ``finalize_time`` has no bundle content yet and is skipped + quietly, since that is a normal in-progress state. + + A finalized skill is expected to carry both names: ``name`` is immutable from + creation, and finalize is the sole writer of ``bundle_name``. One missing is + therefore an anomaly, so warn and skip rather than substituting the other + name -- the two are not interchangeable, and guessing a directory name that + doesn't match the bundle's SKILL.md ``name:`` would hide the skill from the + agent meant to load it. + """ + if not skill.get("finalize_time"): + return None + + name = _non_empty_str(skill.get("name")) + bundle_name = _non_empty_str(skill.get("bundle_name")) + if name is None or bundle_name is None: + missing = " or ".join( + field + for field, value in (("name", name), ("bundle_name", bundle_name)) + if value is None + ) + print_warning( + f"Skipping `{name or ''}`: the skills API returned no {missing}." + ) + return None + + if not _is_safe_bundle_name(bundle_name): + print_warning(f"Skipping `{name}`: unsafe bundle name `{bundle_name}`.") + return None + + parts = name.split("/", 1)[-1].split(".") + if len(parts) != 3: + print_warning(f"Skipping `{name}`: expected a `catalog.schema.name` skill name.") + return None + catalog, schema, securable_name = parts + return SkillRef( + catalog=catalog, + schema=schema, + securable_name=securable_name, + bundle_name=bundle_name, + description=_non_empty_str(skill.get("description")), + metastore_id=_non_empty_str(skill.get("metastore_id")), + skill_id=_non_empty_str(skill.get("id")), + uc_update_time=_non_empty_str(skill.get("update_time")), + ) + + +def list_schema_skills( + workspace: str, token: str, catalog: str, schema: str +) -> tuple[list[SkillRef], str | None]: + """List the finalized skills in ``.``. + + A non-None reason indicates the listing call itself failed. + """ + hostname = workspace_hostname(workspace) + base_url = f"https://{hostname}/api/2.1/unity-catalog/skills" + query = {"parent": f"schemas/{catalog}.{schema}"} + + refs: list[SkillRef] = [] + page_token: str | None = None + while True: + if page_token: + query["page_token"] = page_token + payload, reason = _http_get_json(f"{base_url}?{urlencode(query)}", token, timeout=30) + if payload is None: + return [], reason + data = payload if isinstance(payload, dict) else {} + for skill in data.get("skills") or []: + ref = _skill_ref(skill) if isinstance(skill, dict) else None + if ref: + refs.append(ref) + page_token = data.get("next_page_token") + if not page_token: + return refs, None + + +def get_skill(workspace: str, token: str, fqn: str) -> SkillRef | None: + """The finalized skill named by ``fqn``, or None if it cannot be downloaded. + + ``GetSkill`` returns the same shape as a ``ListSkills`` entry, so the response + runs through ``_skill_ref``; a missing, unfinalized, or malformed skill is None. + """ + hostname = workspace_hostname(workspace) + payload, _ = _http_get_json( + f"https://{hostname}/api/2.1/unity-catalog/skills/{fqn}", token, timeout=30 + ) + return _skill_ref(payload) if isinstance(payload, dict) else None + + +def list_all_skills( + workspace: str, + token: str, + *, + deadline_seconds: float = _SKILLS_WALK_DEADLINE_SECONDS, + on_progress: Callable[[int, int, int], None] | None = None, + on_skills: Callable[[list[SkillRef]], None] | None = None, +) -> tuple[list[SkillRef], str | None]: + """Return every finalized skill across all ``.`` in the workspace, by FQN. + + The skills API is one-schema-per-call, so this walks catalogs -> schemas -> + skills in parallel under a wall-clock budget, returning partial results once + ``deadline_seconds`` is exceeded. ``on_progress`` is called as each schema + completes with ``(schemas_done, schemas_total, skills_found)``, and + ``on_skills`` with each schema's newly-found refs (deduped by FQN against + everything emitted so far) so a picker can stream them in as the walk runs. + The workspace-wide counterpart to ``list_schema_skills``. + """ + deadline = time.monotonic() + deadline_seconds + by_fqn: dict[str, SkillRef] = {} + + def probe(catalog: str, schema: str) -> tuple[list[SkillRef], str | None]: + return list_schema_skills(workspace, token, catalog, schema) + + def collect(result: tuple[list[SkillRef], str | None], done: int, total: int) -> None: + found, _ = result + new = [ref for ref in found if ref.fqn not in by_fqn] + for ref in new: + by_fqn[ref.fqn] = ref + if on_progress is not None: + on_progress(done, total, len(by_fqn)) + if on_skills is not None and new: + on_skills(sorted(new, key=lambda ref: ref.fqn)) + + reason = walk_catalog_schemas(workspace, token, deadline=deadline, probe=probe, collect=collect) + if reason is not None: + return [], reason + refs = sorted(by_fqn.values(), key=lambda ref: ref.fqn) + if time.monotonic() > deadline: + return refs, _SKILLS_WALK_TIMEOUT_REASON + if not refs: + return [], "no skills found" + return refs, None + + +def list_skill_files( + workspace: str, token: str, catalog: str, schema: str, securable: str +) -> tuple[list[str], str | None]: + """List a skill bundle's files, as paths relative to the skill directory. + + Recursively walks the skill's Files API directory (including ``SKILL.md``). + Takes the securable leaf, the only name the Files API resolves. A non-None + reason indicates the listing call itself failed. + """ + hostname = workspace_hostname(workspace) + dirs_base = f"https://{hostname}/api/2.0/fs/directories" + skill_prefix = f"/{SKILL_FILES_API_PREFIX}/{catalog}/{schema}/{securable}/" + + relative_paths: list[str] = [] + pending = [f"{SKILL_FILES_API_PREFIX}/{catalog}/{schema}/{securable}"] + while pending: + directory = pending.pop() + page_token: str | None = None + while True: + url = f"{dirs_base}/{directory}" + 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 entry in data.get("contents") or []: + path = entry.get("path") if isinstance(entry, dict) else None + if not isinstance(path, str): + continue + if entry.get("is_directory"): + pending.append(path.strip("/")) + else: + relative_paths.append(path.removeprefix(skill_prefix)) + page_token = data.get("next_page_token") + if not page_token: + break + return relative_paths, None + + +def fetch_skill_file( + workspace: str, token: str, catalog: str, schema: str, securable: str, relative_path: str +) -> tuple[bytes | None, str | None]: + """Fetch one skill bundle file's raw bytes from the Files API.""" + hostname = workspace_hostname(workspace) + url = ( + f"https://{hostname}/api/2.0/fs/files/" + f"{SKILL_FILES_API_PREFIX}/{catalog}/{schema}/{securable}/{relative_path}" + ) + return _http_get_bytes(url, token, timeout=30) + + +def fetch_skill_bundle( + workspace: str, token: str, catalog: str, schema: str, securable: str +) -> tuple[dict[str, bytes] | None, str | None]: + """Fetch a whole skill bundle as ``{relative_path: bytes}``. + + Lists the skill's files then fetches each one. All-or-nothing: a non-None + reason (and None bundle) means the listing or any file fetch failed, so a + partially-downloaded skill is never written to disk. + """ + relative_paths, reason = list_skill_files(workspace, token, catalog, schema, securable) + if reason: + return None, reason + bundle: dict[str, bytes] = {} + for relative_path in relative_paths: + content, reason = fetch_skill_file( + workspace, token, catalog, schema, securable, relative_path + ) + if content is None: + return None, reason + bundle[relative_path] = content + return bundle, None diff --git a/src/ucode/skills_download.py b/src/ucode/skills_download.py index 988d59ba..21de7a8b 100644 --- a/src/ucode/skills_download.py +++ b/src/ucode/skills_download.py @@ -2,24 +2,23 @@ from __future__ import annotations -import time from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor, as_completed -from dataclasses import dataclass from pathlib import Path -from urllib.parse import urlencode import questionary -from ucode.databricks import ( - _http_get_bytes, - _http_get_json, - get_databricks_token, - walk_catalog_schemas, - workspace_hostname, - workspace_org_id, -) +from ucode.databricks import get_databricks_token, workspace_org_id from ucode.mcp import register_schemaless_skills_connection, setup_mcp_clients +from ucode.skills_api import ( + _SKILLS_WALK_DEADLINE_SECONDS, + _SKILLS_WALK_TIMEOUT_REASON, + SkillRef, + fetch_skill_bundle, + get_skill, + list_all_skills, + list_schema_skills, +) from ucode.skills_state import ( SkillInstall, list_downloaded, @@ -43,213 +42,9 @@ # `.claude/skills` (Claude) + `.agents/skills` (the alias other agents read). SKILL_BASE_DIR_NAMES = (".claude/skills", ".agents/skills") -SKILL_FILES_API_PREFIX = "Skills" - # Parallel skill fetches per schema; writes stay sequential (they prompt). _MAX_FETCH_WORKERS = 8 -# Wall-clock budget for the workspace-wide skill walk; a slow workspace degrades -# to partial results instead of hanging the picker. -_SKILLS_WALK_DEADLINE_SECONDS = 30.0 -_SKILLS_WALK_TIMEOUT_REASON = "deadline exceeded while listing skills" - - -# --- Download client (UC skills API + Files API) --------------------------- - - -@dataclass(frozen=True) -class SkillRef: - """A downloadable skill's UC location plus its two non-interchangeable names. - - ``catalog``/``schema``/``securable_name`` are the parts of ``skills/..``: - ``securable_name`` is the leaf, the only name the Files API resolves, and the - three together fully qualify the skill (``fqn``). ``bundle_name`` is the - ``name:`` an agent reads from the bundle's SKILL.md frontmatter, so it names - the on-disk directory. Finalize does not require the securable and bundle name - to match, so a skill created under a securable that differs from its - frontmatter carries both. ``description`` is the skill's UC description, used only - to preview a skill in the interactive picker. ``metastore_id``, ``skill_id``, and - ``uc_update_time`` are UC attribution metadata recorded when the skill is - downloaded (see ``skills_state``); the download flow never reads them. - """ - - catalog: str - schema: str - securable_name: str - bundle_name: str - description: str | None = None - metastore_id: str | None = None - skill_id: str | None = None - uc_update_time: str | None = None - - @property - def fqn(self) -> str: - return f"{self.catalog}.{self.schema}.{self.securable_name}" - - -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 - - -def _is_safe_bundle_name(bundle_name: str) -> bool: - path = Path(bundle_name) - return len(path.parts) == 1 and path.parts[0] != ".." and not path.is_absolute() - - -def _skill_ref(skill: dict) -> SkillRef | None: - """A finalized skill's ``SkillRef``, or None if it cannot be downloaded. - - A skill without a ``finalize_time`` has no bundle content yet and is skipped - quietly, since that is a normal in-progress state. - - A finalized skill is expected to carry both names: ``name`` is immutable from - creation, and finalize is the sole writer of ``bundle_name``. One missing is - therefore an anomaly, so warn and skip rather than substituting the other - name -- the two are not interchangeable, and guessing a directory name that - doesn't match the bundle's SKILL.md ``name:`` would hide the skill from the - agent meant to load it. - """ - if not skill.get("finalize_time"): - return None - - name = _non_empty_str(skill.get("name")) - bundle_name = _non_empty_str(skill.get("bundle_name")) - if name is None or bundle_name is None: - missing = " or ".join( - field - for field, value in (("name", name), ("bundle_name", bundle_name)) - if value is None - ) - print_warning( - f"Skipping `{name or ''}`: the skills API returned no {missing}." - ) - return None - - if not _is_safe_bundle_name(bundle_name): - print_warning(f"Skipping `{name}`: unsafe bundle name `{bundle_name}`.") - return None - - parts = name.split("/", 1)[-1].split(".") - if len(parts) != 3: - print_warning(f"Skipping `{name}`: expected a `catalog.schema.name` skill name.") - return None - catalog, schema, securable_name = parts - return SkillRef( - catalog=catalog, - schema=schema, - securable_name=securable_name, - bundle_name=bundle_name, - description=_non_empty_str(skill.get("description")), - metastore_id=_non_empty_str(skill.get("metastore_id")), - skill_id=_non_empty_str(skill.get("id")), - uc_update_time=_non_empty_str(skill.get("update_time")), - ) - - -def list_schema_skills( - workspace: str, token: str, catalog: str, schema: str -) -> tuple[list[SkillRef], str | None]: - """List the finalized skills in ``.``. - - A non-None reason indicates the listing call itself failed. - """ - hostname = workspace_hostname(workspace) - base_url = f"https://{hostname}/api/2.1/unity-catalog/skills" - query = {"parent": f"schemas/{catalog}.{schema}"} - - refs: list[SkillRef] = [] - page_token: str | None = None - while True: - if page_token: - query["page_token"] = page_token - payload, reason = _http_get_json(f"{base_url}?{urlencode(query)}", token, timeout=30) - if payload is None: - return [], reason - data = payload if isinstance(payload, dict) else {} - for skill in data.get("skills") or []: - ref = _skill_ref(skill) if isinstance(skill, dict) else None - if ref: - refs.append(ref) - page_token = data.get("next_page_token") - if not page_token: - return refs, None - - -def list_skill_files( - workspace: str, token: str, catalog: str, schema: str, securable: str -) -> tuple[list[str], str | None]: - """List a skill bundle's files, as paths relative to the skill directory. - - Recursively walks the skill's Files API directory (including ``SKILL.md``). - Takes the securable leaf, the only name the Files API resolves. A non-None - reason indicates the listing call itself failed. - """ - hostname = workspace_hostname(workspace) - dirs_base = f"https://{hostname}/api/2.0/fs/directories" - skill_prefix = f"/{SKILL_FILES_API_PREFIX}/{catalog}/{schema}/{securable}/" - - relative_paths: list[str] = [] - pending = [f"{SKILL_FILES_API_PREFIX}/{catalog}/{schema}/{securable}"] - while pending: - directory = pending.pop() - page_token: str | None = None - while True: - url = f"{dirs_base}/{directory}" - 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 entry in data.get("contents") or []: - path = entry.get("path") if isinstance(entry, dict) else None - if not isinstance(path, str): - continue - if entry.get("is_directory"): - pending.append(path.strip("/")) - else: - relative_paths.append(path.removeprefix(skill_prefix)) - page_token = data.get("next_page_token") - if not page_token: - break - return relative_paths, None - - -def fetch_skill_file( - workspace: str, token: str, catalog: str, schema: str, securable: str, relative_path: str -) -> tuple[bytes | None, str | None]: - """Fetch one skill bundle file's raw bytes from the Files API.""" - hostname = workspace_hostname(workspace) - url = ( - f"https://{hostname}/api/2.0/fs/files/" - f"{SKILL_FILES_API_PREFIX}/{catalog}/{schema}/{securable}/{relative_path}" - ) - return _http_get_bytes(url, token, timeout=30) - - -def fetch_skill_bundle( - workspace: str, token: str, catalog: str, schema: str, securable: str -) -> tuple[dict[str, bytes] | None, str | None]: - """Fetch a whole skill bundle as ``{relative_path: bytes}``. - - Lists the skill's files then fetches each one. All-or-nothing: a non-None - reason (and None bundle) means the listing or any file fetch failed, so a - partially-downloaded skill is never written to disk. - """ - relative_paths, reason = list_skill_files(workspace, token, catalog, schema, securable) - if reason: - return None, reason - bundle: dict[str, bytes] = {} - for relative_path in relative_paths: - content, reason = fetch_skill_file( - workspace, token, catalog, schema, securable, relative_path - ) - if content is None: - return None, reason - bundle[relative_path] = content - return bundle, None - # --- On-disk writer -------------------------------------------------------- @@ -473,19 +268,6 @@ def download_skills_from_schema_locations( ) -def get_skill(workspace: str, token: str, fqn: str) -> SkillRef | None: - """The finalized skill named by ``fqn``, or None if it cannot be downloaded. - - ``GetSkill`` returns the same shape as a ``ListSkills`` entry, so the response - runs through ``_skill_ref``; a missing, unfinalized, or malformed skill is None. - """ - hostname = workspace_hostname(workspace) - payload, _ = _http_get_json( - f"https://{hostname}/api/2.1/unity-catalog/skills/{fqn}", token, timeout=30 - ) - return _skill_ref(payload) if isinstance(payload, dict) else None - - def download_selected_skills(workspace: str, token: str, fqns: list[str], path: str | None) -> None: """Download the skills named by ``fqns`` (``..``) to disk. @@ -585,52 +367,7 @@ def configure_selected_skills_download_command(fqns: list[str], path: str | None return 0 -# --- Interactive picker (workspace discovery + selective download) ---------- - - -def list_all_skills( - workspace: str, - token: str, - *, - deadline_seconds: float = _SKILLS_WALK_DEADLINE_SECONDS, - on_progress: Callable[[int, int, int], None] | None = None, - on_skills: Callable[[list[SkillRef]], None] | None = None, -) -> tuple[list[SkillRef], str | None]: - """Return every finalized skill across all ``.`` in the workspace, by FQN. - - The skills API is one-schema-per-call, so this walks catalogs -> schemas -> - skills in parallel under a wall-clock budget, returning partial results once - ``deadline_seconds`` is exceeded. ``on_progress`` is called as each schema - completes with ``(schemas_done, schemas_total, skills_found)``, and - ``on_skills`` with each schema's newly-found refs (deduped by FQN against - everything emitted so far) so a picker can stream them in as the walk runs. - The workspace-wide counterpart to ``list_schema_skills``. - """ - deadline = time.monotonic() + deadline_seconds - by_fqn: dict[str, SkillRef] = {} - - def probe(catalog: str, schema: str) -> tuple[list[SkillRef], str | None]: - return list_schema_skills(workspace, token, catalog, schema) - - def collect(result: tuple[list[SkillRef], str | None], done: int, total: int) -> None: - found, _ = result - new = [ref for ref in found if ref.fqn not in by_fqn] - for ref in new: - by_fqn[ref.fqn] = ref - if on_progress is not None: - on_progress(done, total, len(by_fqn)) - if on_skills is not None and new: - on_skills(sorted(new, key=lambda ref: ref.fqn)) - - reason = walk_catalog_schemas(workspace, token, deadline=deadline, probe=probe, collect=collect) - if reason is not None: - return [], reason - refs = sorted(by_fqn.values(), key=lambda ref: ref.fqn) - if time.monotonic() > deadline: - return refs, _SKILLS_WALK_TIMEOUT_REASON - if not refs: - return [], "no skills found" - return refs, None +# --- Interactive picker (selective download) -------------------------------- def _skill_download_choice(ref: SkillRef, roots: list[Path]) -> questionary.Choice: diff --git a/tests/test_cli.py b/tests/test_cli.py index c7c00980..9ef4d081 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1562,6 +1562,40 @@ def test_interactive_picker_passes_path(self): assert result.exit_code == 0, result.output mock_picker.assert_called_once_with(path="/tmp/s") + def test_mcp_no_location_interactive_opens_schema_picker(self): + with ( + patch("ucode.cli._stdin_is_interactive", return_value=True), + patch("ucode.cli.configure_skills_mcp_picker_command") as mock_picker, + patch("ucode.cli.configure_skills_download_picker_command") as mock_download, + ): + result = runner.invoke(app, ["skill", "add", "--mcp"]) + assert result.exit_code == 0, result.output + mock_picker.assert_called_once_with(agents=None) + mock_download.assert_not_called() + + def test_mcp_no_location_non_interactive_exit_1(self): + with ( + patch("ucode.cli._stdin_is_interactive", return_value=False), + patch("ucode.cli.configure_skills_mcp_picker_command") as mock_picker, + ): + result = runner.invoke(app, ["skill", "add", "--mcp"]) + assert result.exit_code == 1 + assert "--location is required" in _strip_ansi(result.output) + mock_picker.assert_not_called() + + def test_mcp_picker_with_agents_forwards_scope(self): + with ( + patch("ucode.cli._stdin_is_interactive", return_value=True), + patch( + "ucode.cli._configure_agents_for_mcp", return_value={"claude", "codex"} + ) as configure, + patch("ucode.cli.configure_skills_mcp_picker_command") as mock_picker, + ): + result = runner.invoke(app, ["skill", "add", "--mcp", "--agents", "codex,claude"]) + assert result.exit_code == 0, result.output + configure.assert_called_once_with(["claude", "codex"]) + mock_picker.assert_called_once_with(agents={"claude", "codex"}) + def test_skills_bypass_picker_even_when_interactive(self): with ( patch("ucode.cli._stdin_is_interactive", return_value=True), diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 63773bbb..a2756ca2 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -2306,6 +2306,134 @@ def test_agents_add_matching_existing_scope_is_a_noop(self, monkeypatch): assert configured == [] +class TestConfiguredSkillLocations: + def test_unions_locations_across_targeted_clients(self): + state = _skills_state( + mcp._resolve_skills_mcp_servers( + WS, ["claude", "codex"], {"claude": ["A.a", "B.b"], "codex": ["C.c"]}, [] + ) + ) + assert mcp.configured_skill_locations(state, ["claude", "codex"]) == {"A.a", "B.b", "C.c"} + assert mcp.configured_skill_locations(state, ["claude"]) == {"A.a", "B.b"} + + def test_empty_when_nothing_configured(self): + assert mcp.configured_skill_locations(_skills_state(), ["claude"]) == set() + + +class _FakePrompt: + def __init__(self, result): + self._result = result + + def ask(self): + return self._result + + +def _skill_ref(securable, *, catalog="main", schema="default"): + from ucode.skills_api import SkillRef + + return SkillRef(catalog=catalog, schema=schema, securable_name=securable, bundle_name=securable) + + +class TestSkillSchemaPicker: + def test_choice_value_is_location_and_shows_count(self): + choice = mcp._skill_schema_choice("main.default", 3, in_scope=False) + assert choice.value == "main.default" + assert "3 skills" in choice.title + assert "already in skill MCP" not in choice.title + + def test_choice_singular_count_and_in_scope_flag(self): + choice = mcp._skill_schema_choice("ml.prod", 1, in_scope=True) + assert "1 skill" in choice.title and "1 skills" not in choice.title + assert "already in skill MCP" in choice.title + + def test_background_loader_streams_one_row_per_schema(self, monkeypatch): + def fake_list_all(ws, tok, *, on_skills=None, **kwargs): + on_skills([_skill_ref("triage"), _skill_ref("pii")]) + on_skills([_skill_ref("scoring", catalog="ml", schema="prod")]) + return [], None + + monkeypatch.setattr(mcp, "list_all_skills", fake_list_all) + appended = [] + + message = mcp._skill_schema_background_loader(WS, "token", {"ml.prod"})(appended.extend) + + assert message is None + assert [c.value for c in appended] == ["main.default", "ml.prod"] + assert "2 skills" in appended[0].title and "already in skill MCP" not in appended[0].title + assert "already in skill MCP" in appended[1].title + + def test_background_loader_reports_timeout_message(self, monkeypatch): + def fake_list_all(ws, tok, *, on_skills=None, **kwargs): + on_skills([_skill_ref("triage")]) + on_skills([_skill_ref("scoring", catalog="ml", schema="prod")]) + return ( + [_skill_ref("triage"), _skill_ref("scoring", catalog="ml", schema="prod")], + mcp._SKILLS_WALK_TIMEOUT_REASON, + ) + + monkeypatch.setattr(mcp, "list_all_skills", fake_list_all) + + message = mcp._skill_schema_background_loader(WS, "token", set())(lambda choices: None) + + assert message == "⚠️ Timed out after 30s, found 2 skill schemas" + + def test_prompt_returns_selected_locations(self, monkeypatch): + loader = lambda append: None # noqa: E731 + captured = {} + + def fake_checkbox(message, *, choices, instruction, style, background_loader, **kwargs): + captured.update(background_loader=background_loader, **kwargs) + return _FakePrompt(["main.default", "ml.prod"]) + + monkeypatch.setattr(mcp, "scrolling_checkbox", fake_checkbox) + + assert mcp.prompt_for_skill_schema_choices(loader) == ["main.default", "ml.prod"] + assert captured["loading_noun"] == "skill schemas" + assert captured["background_loader"] is loader + + def test_prompt_returns_none_on_cancel(self, monkeypatch): + monkeypatch.setattr(mcp, "scrolling_checkbox", lambda *a, **k: _FakePrompt(None)) + assert mcp.prompt_for_skill_schema_choices(lambda append: None) is None + + +class TestConfigureSkillsMcpPickerCommand: + def _stub(self, monkeypatch, locations): + calls: dict[str, object] = {} + monkeypatch.setattr(mcp, "load_state", lambda: {"state": True}) + monkeypatch.setattr( + mcp, + "setup_mcp_clients", + lambda state, section, agents=None: (WS, "profile", ["claude"]), + ) + monkeypatch.setattr(mcp, "get_databricks_token", lambda ws, profile=None: "token") + monkeypatch.setattr(mcp, "configured_skill_locations", lambda state, clients: {"A.a"}) + monkeypatch.setattr( + mcp, "_skill_schema_background_loader", lambda ws, token, in_scope: "loader" + ) + monkeypatch.setattr(mcp, "prompt_for_skill_schema_choices", lambda loader: locations) + + def fake_add(state, ws, profile, clients, locs): + calls["added"] = (ws, profile, clients, locs) + + monkeypatch.setattr(mcp, "add_skill_locations_to_mcp", fake_add) + return calls + + def test_adds_selected_locations(self, monkeypatch): + calls = self._stub(monkeypatch, ["main.default", "ml.prod"]) + assert mcp.configure_skills_mcp_picker_command() == 0 + assert calls["added"] == (WS, "profile", ["claude"], ["main.default", "ml.prod"]) + + def test_cancel_adds_nothing(self, monkeypatch): + calls = self._stub(monkeypatch, None) + assert mcp.configure_skills_mcp_picker_command() == 0 + assert "added" not in calls + + def test_empty_selection_adds_nothing(self, monkeypatch): + calls = self._stub(monkeypatch, []) + assert mcp.configure_skills_mcp_picker_command() == 0 + assert "added" not in calls + + class TestRemoveSkillsCommand: def _state(self, by_client=None): by_client = by_client or _by_client(["claude", "codex"], ["A.a", "B.b"]) diff --git a/tests/test_skills_api.py b/tests/test_skills_api.py new file mode 100644 index 00000000..f5a43d2a --- /dev/null +++ b/tests/test_skills_api.py @@ -0,0 +1,482 @@ +"""Tests for skills_api.py -- the read-only UC skills API client and workspace walk.""" + +from __future__ import annotations + +import pytest + +import ucode.skills_api as sa +from ucode.skills_api import SkillRef + +WS = "https://example.databricks.com" + + +def ref( + securable_name: str, + bundle_name: str | None = None, + *, + catalog: str = "main", + schema: str = "default", + description: str | None = None, +) -> SkillRef: + """A SkillRef whose two names match unless a differing bundle name is given.""" + return SkillRef( + catalog=catalog, + schema=schema, + securable_name=securable_name, + bundle_name=bundle_name or securable_name, + description=description, + ) + + +class TestListSchemaSkills: + def test_keeps_finalized_skills_only(self, monkeypatch): + payload = { + "skills": [ + { + "name": "skills/main.default.pii-handling", + "bundle_name": "pii-handling", + "finalize_time": "2026-06-26T05:58:25Z", + }, + { + "name": "skills/main.default.triage", + "bundle_name": "triage", + "finalize_time": "2026-06-26T05:58:26Z", + }, + {"name": "skills/main.default.draft", "bundle_name": "draft"}, + ] + } + monkeypatch.setattr(sa, "_http_get_json", lambda url, token, timeout=30: (payload, None)) + + refs, reason = sa.list_schema_skills(WS, "token", "main", "default") + + assert reason is None + assert refs == [ref("pii-handling"), ref("triage")] + + def test_carries_description_when_present(self, monkeypatch): + payload = { + "skills": [ + { + "name": "skills/main.default.triage", + "bundle_name": "triage", + "finalize_time": "2026-06-26T05:58:25Z", + "description": "Routes tickets by severity.", + } + ] + } + monkeypatch.setattr(sa, "_http_get_json", lambda url, token, timeout=30: (payload, None)) + + refs, _ = sa.list_schema_skills(WS, "token", "main", "default") + + assert refs == [ref("triage", description="Routes tickets by severity.")] + + def test_keeps_both_names_when_bundle_differs_from_securable(self, monkeypatch): + # bundle_name comes from the bundle's SKILL.md frontmatter, so it can + # differ from the securable it was created under. + payload = { + "skills": [ + { + "name": "skills/main.default.task-prioritizer", + "bundle_name": "task-triage", + "finalize_time": "2026-06-26T05:58:25Z", + } + ] + } + monkeypatch.setattr(sa, "_http_get_json", lambda url, token, timeout=30: (payload, None)) + + refs, reason = sa.list_schema_skills(WS, "token", "main", "default") + + assert reason is None + assert refs == [ref("task-prioritizer", "task-triage")] + + @pytest.mark.parametrize( + ("skill", "expected_missing"), + [ + ({"name": "skills/main.default.pii-handling"}, "bundle_name"), + ({"name": "skills/main.default.pii-handling", "bundle_name": ""}, "bundle_name"), + ({"bundle_name": "orphan"}, "name"), + ({}, "name or bundle_name"), + ], + ids=["no-bundle-name", "blank-bundle-name", "no-resource-name", "neither"], + ) + def test_skips_and_warns_when_a_name_is_missing(self, skill, expected_missing, monkeypatch): + # Finalize owns bundle_name and `name` is immutable from creation, so a + # finalized skill missing either is an anomaly worth surfacing. + payload = {"skills": [{**skill, "finalize_time": "2026-06-26T05:58:25Z"}]} + monkeypatch.setattr(sa, "_http_get_json", lambda url, token, timeout=30: (payload, None)) + warnings = [] + monkeypatch.setattr(sa, "print_warning", warnings.append) + + refs, reason = sa.list_schema_skills(WS, "token", "main", "default") + + assert reason is None + assert refs == [] + assert len(warnings) == 1 + assert f"no {expected_missing}." in warnings[0] + + def test_unfinalized_skill_is_skipped_without_a_warning(self, monkeypatch): + # An unfinalized skill simply has no bundle yet, which is not an anomaly. + payload = {"skills": [{"name": "skills/main.default.draft"}]} + monkeypatch.setattr(sa, "_http_get_json", lambda url, token, timeout=30: (payload, None)) + warnings = [] + monkeypatch.setattr(sa, "print_warning", warnings.append) + + refs, reason = sa.list_schema_skills(WS, "token", "main", "default") + + assert reason is None + assert refs == [] + assert warnings == [] + + @pytest.mark.parametrize( + "bundle_name", + ["..", "../escape", "nested/../escape", "a/b", "/abs"], + ids=["dotdot", "parent-traversal", "embedded-traversal", "separator", "absolute"], + ) + def test_skips_and_warns_on_unsafe_bundle_name(self, bundle_name, monkeypatch): + payload = { + "skills": [ + { + "name": "skills/main.default.pii-handling", + "bundle_name": bundle_name, + "finalize_time": "2026-06-26T05:58:25Z", + } + ] + } + monkeypatch.setattr(sa, "_http_get_json", lambda url, token, timeout=30: (payload, None)) + warnings = [] + monkeypatch.setattr(sa, "print_warning", warnings.append) + + refs, reason = sa.list_schema_skills(WS, "token", "main", "default") + + assert reason is None + assert refs == [] + assert len(warnings) == 1 + assert "unsafe bundle name" in warnings[0] + + def test_follows_pagination(self, monkeypatch): + pages = [ + { + "skills": [ + {"name": "skills/main.default.a", "bundle_name": "a", "finalize_time": "t"} + ], + "next_page_token": "tok", + }, + { + "skills": [ + {"name": "skills/main.default.b", "bundle_name": "b", "finalize_time": "t"} + ] + }, + ] + captured_tokens = [] + + def fake_get(url, token, timeout=30): + captured_tokens.append("page_token=tok" in url) + return pages.pop(0), None + + monkeypatch.setattr(sa, "_http_get_json", fake_get) + + refs, reason = sa.list_schema_skills(WS, "token", "main", "default") + + assert reason is None + assert refs == [ref("a"), ref("b")] + assert captured_tokens == [False, True] + + def test_targets_uc_skills_api_for_the_schema(self, monkeypatch): + captured = {} + + def fake_get(url, token, timeout=30): + captured["url"] = url + return {"skills": []}, None + + monkeypatch.setattr(sa, "_http_get_json", fake_get) + + sa.list_schema_skills(WS, "token", "main", "default") + + assert "/api/2.1/unity-catalog/skills?" in captured["url"] + assert "parent=schemas%2Fmain.default" in captured["url"] + + def test_http_failure_propagates_reason(self, monkeypatch): + monkeypatch.setattr( + sa, "_http_get_json", lambda url, token, timeout=30: (None, "HTTP 500 Server Error") + ) + + leaves, reason = sa.list_schema_skills(WS, "token", "main", "default") + + assert leaves == [] + assert reason == "HTTP 500 Server Error" + + +class TestListSkillFiles: + def test_lists_under_the_skills_place(self, monkeypatch): + captured = {} + + def fake_get(url, token, timeout=30): + captured["url"] = url + return {"contents": []}, None + + monkeypatch.setattr(sa, "_http_get_json", fake_get) + + sa.list_skill_files(WS, "token", "main", "default", "triage") + + assert captured["url"] == f"{WS}/api/2.0/fs/directories/Skills/main/default/triage" + + def test_walks_nested_directories_into_relative_paths(self, monkeypatch): + # The Files API returns absolute paths. + skill = "/Skills/main/default/triage" + listings = { + "Skills/main/default/triage": { + "contents": [ + {"path": f"{skill}/SKILL.md", "is_directory": False}, + {"path": f"{skill}/references/", "is_directory": True}, + ] + }, + "Skills/main/default/triage/references": { + "contents": [{"path": f"{skill}/references/primary.md", "is_directory": False}] + }, + } + + def fake_get(url, token, timeout=30): + directory = url.split("/api/2.0/fs/directories/", 1)[1] + return listings[directory], None + + monkeypatch.setattr(sa, "_http_get_json", fake_get) + + paths, reason = sa.list_skill_files(WS, "token", "main", "default", "triage") + + assert reason is None + assert sorted(paths) == ["SKILL.md", "references/primary.md"] + + def test_follows_pagination(self, monkeypatch): + skill = "/Skills/main/default/triage" + pages = [ + { + "contents": [{"path": f"{skill}/a.md", "is_directory": False}], + "next_page_token": "tok", + }, + {"contents": [{"path": f"{skill}/b.md", "is_directory": False}]}, + ] + + monkeypatch.setattr( + sa, "_http_get_json", lambda url, token, timeout=30: (pages.pop(0), None) + ) + + paths, reason = sa.list_skill_files(WS, "token", "main", "default", "triage") + + assert reason is None + assert sorted(paths) == ["a.md", "b.md"] + + def test_http_failure_propagates_reason(self, monkeypatch): + monkeypatch.setattr( + sa, "_http_get_json", lambda url, token, timeout=30: (None, "HTTP 404 Not Found") + ) + + paths, reason = sa.list_skill_files(WS, "token", "main", "default", "triage") + + assert paths == [] + assert reason == "HTTP 404 Not Found" + + +class TestFetchSkillFile: + def test_returns_raw_bytes_from_files_api(self, monkeypatch): + captured = {} + + def fake_get_bytes(url, token, timeout=30): + captured["url"] = url + return b"# SKILL\n", None + + monkeypatch.setattr(sa, "_http_get_bytes", fake_get_bytes) + + body, reason = sa.fetch_skill_file(WS, "token", "main", "default", "triage", "SKILL.md") + + assert reason is None + assert body == b"# SKILL\n" + assert captured["url"] == f"{WS}/api/2.0/fs/files/Skills/main/default/triage/SKILL.md" + + def test_http_failure_propagates_reason(self, monkeypatch): + monkeypatch.setattr( + sa, "_http_get_bytes", lambda url, token, timeout=30: (None, "HTTP 404 Not Found") + ) + + body, reason = sa.fetch_skill_file(WS, "token", "main", "default", "triage", "gone.md") + + assert body is None + assert reason == "HTTP 404 Not Found" + + +class TestFetchSkillBundle: + def test_assembles_relpath_to_bytes_map(self, monkeypatch): + contents = {"SKILL.md": b"# skill", "references/a.md": b"aaa"} + monkeypatch.setattr(sa, "list_skill_files", lambda *a, **k: (list(contents), None)) + monkeypatch.setattr( + sa, "fetch_skill_file", lambda ws, tok, c, s, leaf, rel: (contents[rel], None) + ) + + bundle, reason = sa.fetch_skill_bundle(WS, "token", "main", "default", "triage") + + assert reason is None + assert bundle == contents + + def test_listing_failure_propagates_reason(self, monkeypatch): + monkeypatch.setattr(sa, "list_skill_files", lambda *a, **k: ([], "HTTP 404 Not Found")) + + bundle, reason = sa.fetch_skill_bundle(WS, "token", "main", "default", "triage") + + assert bundle is None + assert reason == "HTTP 404 Not Found" + + def test_file_failure_aborts_whole_bundle(self, monkeypatch): + monkeypatch.setattr( + sa, "list_skill_files", lambda *a, **k: (["SKILL.md", "broken.md"], None) + ) + monkeypatch.setattr( + sa, + "fetch_skill_file", + lambda ws, tok, c, s, leaf, rel: ( + (b"ok", None) if rel == "SKILL.md" else (None, "HTTP 500 Server Error") + ), + ) + + bundle, reason = sa.fetch_skill_bundle(WS, "token", "main", "default", "triage") + + assert bundle is None + assert reason == "HTTP 500 Server Error" + + +class TestGetSkill: + def test_returns_ref_with_location_parsed_from_fqn(self, monkeypatch): + captured = {} + + def fake_get(url, token, timeout=30): + captured["url"] = url + return { + "name": "skills/ml.prod.pii-handling", + "bundle_name": "pii-handling", + "finalize_time": "2026-06-26T05:58:25Z", + }, None + + monkeypatch.setattr(sa, "_http_get_json", fake_get) + + result = sa.get_skill(WS, "token", "ml.prod.pii-handling") + + assert result == ref("pii-handling", catalog="ml", schema="prod") + assert captured["url"] == f"{WS}/api/2.1/unity-catalog/skills/ml.prod.pii-handling" + + def test_not_found_returns_none(self, monkeypatch): + monkeypatch.setattr( + sa, "_http_get_json", lambda url, token, timeout=30: (None, "HTTP 404 Not Found") + ) + + assert sa.get_skill(WS, "token", "main.default.gone") is None + + def test_unfinalized_skill_returns_none(self, monkeypatch): + monkeypatch.setattr( + sa, + "_http_get_json", + lambda url, token, timeout=30: ({"name": "skills/main.default.draft"}, None), + ) + + assert sa.get_skill(WS, "token", "main.default.draft") is None + + +class TestSkillRefMetadata: + def test_captures_uc_attribution_fields(self, monkeypatch): + payload = { + "skills": [ + { + "name": "skills/main.default.triage", + "bundle_name": "triage", + "finalize_time": "2026-06-26T05:58:25Z", + "id": "skill-uuid", + "metastore_id": "metastore-uuid", + "update_time": "2026-06-26T05:58:25Z", + } + ] + } + monkeypatch.setattr(sa, "_http_get_json", lambda url, token, timeout=30: (payload, None)) + + (skill,), reason = sa.list_schema_skills(WS, "token", "main", "default") + + assert reason is None + assert skill.metastore_id == "metastore-uuid" + assert skill.skill_id == "skill-uuid" + assert skill.uc_update_time == "2026-06-26T05:58:25Z" + + +def _walk_stub(schemas, reason=None): + """A ``walk_catalog_schemas`` stub that probes each (catalog, schema) in order.""" + + def walk(workspace, token, *, deadline, probe, collect, **kwargs): + total = len(schemas) + for done, (catalog, schema) in enumerate(schemas, start=1): + collect(probe(catalog, schema), done, total) + return reason + + return walk + + +class TestListAllSkills: + def test_flattens_and_streams_across_schemas(self, monkeypatch): + monkeypatch.setattr( + sa, "walk_catalog_schemas", _walk_stub([("main", "default"), ("ml", "prod")]) + ) + by_schema = { + "main.default": [ref("triage"), ref("pii")], + "ml.prod": [ref("scoring", catalog="ml", schema="prod")], + } + monkeypatch.setattr( + sa, "list_schema_skills", lambda ws, tok, c, s: (by_schema[f"{c}.{s}"], None) + ) + streamed = [] + progress = [] + + refs, reason = sa.list_all_skills( + WS, + "token", + on_skills=streamed.append, + on_progress=lambda done, total, found: progress.append((done, total, found)), + ) + + assert reason is None + assert [r.fqn for r in refs] == [ + "main.default.pii", + "main.default.triage", + "ml.prod.scoring", + ] + assert [[r.fqn for r in batch] for batch in streamed] == [ + ["main.default.pii", "main.default.triage"], + ["ml.prod.scoring"], + ] + assert progress == [(1, 2, 2), (2, 2, 3)] + + def test_dedupes_repeated_fqns(self, monkeypatch): + monkeypatch.setattr( + sa, "walk_catalog_schemas", _walk_stub([("main", "default"), ("main", "default")]) + ) + monkeypatch.setattr(sa, "list_schema_skills", lambda ws, tok, c, s: ([ref("triage")], None)) + streamed = [] + + refs, reason = sa.list_all_skills(WS, "token", on_skills=streamed.append) + + assert [r.fqn for r in refs] == ["main.default.triage"] + assert streamed == [[ref("triage")]] + + def test_walk_failure_returns_its_reason(self, monkeypatch): + monkeypatch.setattr( + sa, "walk_catalog_schemas", _walk_stub([], reason="no UC catalogs found") + ) + + assert sa.list_all_skills(WS, "token") == ([], "no UC catalogs found") + + def test_empty_walk_reports_no_skills(self, monkeypatch): + monkeypatch.setattr(sa, "walk_catalog_schemas", _walk_stub([("main", "default")])) + monkeypatch.setattr(sa, "list_schema_skills", lambda *a, **k: ([], None)) + + assert sa.list_all_skills(WS, "token") == ([], "no skills found") + + def test_timeout_returns_partial_results_with_reason(self, monkeypatch): + monkeypatch.setattr(sa, "walk_catalog_schemas", _walk_stub([("main", "default")])) + monkeypatch.setattr(sa, "list_schema_skills", lambda ws, tok, c, s: ([ref("triage")], None)) + + refs, reason = sa.list_all_skills(WS, "token", deadline_seconds=-1) + + assert [r.fqn for r in refs] == ["main.default.triage"] + assert reason == sa._SKILLS_WALK_TIMEOUT_REASON diff --git a/tests/test_skills_download.py b/tests/test_skills_download.py index 6cae0ce0..3e05fd12 100644 --- a/tests/test_skills_download.py +++ b/tests/test_skills_download.py @@ -36,319 +36,6 @@ def ref( ) -class TestListSchemaSkills: - def test_keeps_finalized_skills_only(self, monkeypatch): - payload = { - "skills": [ - { - "name": "skills/main.default.pii-handling", - "bundle_name": "pii-handling", - "finalize_time": "2026-06-26T05:58:25Z", - }, - { - "name": "skills/main.default.triage", - "bundle_name": "triage", - "finalize_time": "2026-06-26T05:58:26Z", - }, - {"name": "skills/main.default.draft", "bundle_name": "draft"}, - ] - } - monkeypatch.setattr(sd, "_http_get_json", lambda url, token, timeout=30: (payload, None)) - - refs, reason = sd.list_schema_skills(WS, "token", "main", "default") - - assert reason is None - assert refs == [ref("pii-handling"), ref("triage")] - - def test_carries_description_when_present(self, monkeypatch): - payload = { - "skills": [ - { - "name": "skills/main.default.triage", - "bundle_name": "triage", - "finalize_time": "2026-06-26T05:58:25Z", - "description": "Routes tickets by severity.", - } - ] - } - monkeypatch.setattr(sd, "_http_get_json", lambda url, token, timeout=30: (payload, None)) - - refs, _ = sd.list_schema_skills(WS, "token", "main", "default") - - assert refs == [ref("triage", description="Routes tickets by severity.")] - - def test_keeps_both_names_when_bundle_differs_from_securable(self, monkeypatch): - # bundle_name comes from the bundle's SKILL.md frontmatter, so it can - # differ from the securable it was created under. - payload = { - "skills": [ - { - "name": "skills/main.default.task-prioritizer", - "bundle_name": "task-triage", - "finalize_time": "2026-06-26T05:58:25Z", - } - ] - } - monkeypatch.setattr(sd, "_http_get_json", lambda url, token, timeout=30: (payload, None)) - - refs, reason = sd.list_schema_skills(WS, "token", "main", "default") - - assert reason is None - assert refs == [ref("task-prioritizer", "task-triage")] - - @pytest.mark.parametrize( - ("skill", "expected_missing"), - [ - ({"name": "skills/main.default.pii-handling"}, "bundle_name"), - ({"name": "skills/main.default.pii-handling", "bundle_name": ""}, "bundle_name"), - ({"bundle_name": "orphan"}, "name"), - ({}, "name or bundle_name"), - ], - ids=["no-bundle-name", "blank-bundle-name", "no-resource-name", "neither"], - ) - def test_skips_and_warns_when_a_name_is_missing(self, skill, expected_missing, monkeypatch): - # Finalize owns bundle_name and `name` is immutable from creation, so a - # finalized skill missing either is an anomaly worth surfacing. - payload = {"skills": [{**skill, "finalize_time": "2026-06-26T05:58:25Z"}]} - monkeypatch.setattr(sd, "_http_get_json", lambda url, token, timeout=30: (payload, None)) - warnings = [] - monkeypatch.setattr(sd, "print_warning", warnings.append) - - refs, reason = sd.list_schema_skills(WS, "token", "main", "default") - - assert reason is None - assert refs == [] - assert len(warnings) == 1 - assert f"no {expected_missing}." in warnings[0] - - def test_unfinalized_skill_is_skipped_without_a_warning(self, monkeypatch): - # An unfinalized skill simply has no bundle yet, which is not an anomaly. - payload = {"skills": [{"name": "skills/main.default.draft"}]} - monkeypatch.setattr(sd, "_http_get_json", lambda url, token, timeout=30: (payload, None)) - warnings = [] - monkeypatch.setattr(sd, "print_warning", warnings.append) - - refs, reason = sd.list_schema_skills(WS, "token", "main", "default") - - assert reason is None - assert refs == [] - assert warnings == [] - - @pytest.mark.parametrize( - "bundle_name", - ["..", "../escape", "nested/../escape", "a/b", "/abs"], - ids=["dotdot", "parent-traversal", "embedded-traversal", "separator", "absolute"], - ) - def test_skips_and_warns_on_unsafe_bundle_name(self, bundle_name, monkeypatch): - payload = { - "skills": [ - { - "name": "skills/main.default.pii-handling", - "bundle_name": bundle_name, - "finalize_time": "2026-06-26T05:58:25Z", - } - ] - } - monkeypatch.setattr(sd, "_http_get_json", lambda url, token, timeout=30: (payload, None)) - warnings = [] - monkeypatch.setattr(sd, "print_warning", warnings.append) - - refs, reason = sd.list_schema_skills(WS, "token", "main", "default") - - assert reason is None - assert refs == [] - assert len(warnings) == 1 - assert "unsafe bundle name" in warnings[0] - - def test_follows_pagination(self, monkeypatch): - pages = [ - { - "skills": [ - {"name": "skills/main.default.a", "bundle_name": "a", "finalize_time": "t"} - ], - "next_page_token": "tok", - }, - { - "skills": [ - {"name": "skills/main.default.b", "bundle_name": "b", "finalize_time": "t"} - ] - }, - ] - captured_tokens = [] - - def fake_get(url, token, timeout=30): - captured_tokens.append("page_token=tok" in url) - return pages.pop(0), None - - monkeypatch.setattr(sd, "_http_get_json", fake_get) - - refs, reason = sd.list_schema_skills(WS, "token", "main", "default") - - assert reason is None - assert refs == [ref("a"), ref("b")] - assert captured_tokens == [False, True] - - def test_targets_uc_skills_api_for_the_schema(self, monkeypatch): - captured = {} - - def fake_get(url, token, timeout=30): - captured["url"] = url - return {"skills": []}, None - - monkeypatch.setattr(sd, "_http_get_json", fake_get) - - sd.list_schema_skills(WS, "token", "main", "default") - - assert "/api/2.1/unity-catalog/skills?" in captured["url"] - assert "parent=schemas%2Fmain.default" in captured["url"] - - def test_http_failure_propagates_reason(self, monkeypatch): - monkeypatch.setattr( - sd, "_http_get_json", lambda url, token, timeout=30: (None, "HTTP 500 Server Error") - ) - - leaves, reason = sd.list_schema_skills(WS, "token", "main", "default") - - assert leaves == [] - assert reason == "HTTP 500 Server Error" - - -class TestListSkillFiles: - def test_lists_under_the_skills_place(self, monkeypatch): - captured = {} - - def fake_get(url, token, timeout=30): - captured["url"] = url - return {"contents": []}, None - - monkeypatch.setattr(sd, "_http_get_json", fake_get) - - sd.list_skill_files(WS, "token", "main", "default", "triage") - - assert captured["url"] == f"{WS}/api/2.0/fs/directories/Skills/main/default/triage" - - def test_walks_nested_directories_into_relative_paths(self, monkeypatch): - # The Files API returns absolute paths. - skill = "/Skills/main/default/triage" - listings = { - "Skills/main/default/triage": { - "contents": [ - {"path": f"{skill}/SKILL.md", "is_directory": False}, - {"path": f"{skill}/references/", "is_directory": True}, - ] - }, - "Skills/main/default/triage/references": { - "contents": [{"path": f"{skill}/references/primary.md", "is_directory": False}] - }, - } - - def fake_get(url, token, timeout=30): - directory = url.split("/api/2.0/fs/directories/", 1)[1] - return listings[directory], None - - monkeypatch.setattr(sd, "_http_get_json", fake_get) - - paths, reason = sd.list_skill_files(WS, "token", "main", "default", "triage") - - assert reason is None - assert sorted(paths) == ["SKILL.md", "references/primary.md"] - - def test_follows_pagination(self, monkeypatch): - skill = "/Skills/main/default/triage" - pages = [ - { - "contents": [{"path": f"{skill}/a.md", "is_directory": False}], - "next_page_token": "tok", - }, - {"contents": [{"path": f"{skill}/b.md", "is_directory": False}]}, - ] - - monkeypatch.setattr( - sd, "_http_get_json", lambda url, token, timeout=30: (pages.pop(0), None) - ) - - paths, reason = sd.list_skill_files(WS, "token", "main", "default", "triage") - - assert reason is None - assert sorted(paths) == ["a.md", "b.md"] - - def test_http_failure_propagates_reason(self, monkeypatch): - monkeypatch.setattr( - sd, "_http_get_json", lambda url, token, timeout=30: (None, "HTTP 404 Not Found") - ) - - paths, reason = sd.list_skill_files(WS, "token", "main", "default", "triage") - - assert paths == [] - assert reason == "HTTP 404 Not Found" - - -class TestFetchSkillFile: - def test_returns_raw_bytes_from_files_api(self, monkeypatch): - captured = {} - - def fake_get_bytes(url, token, timeout=30): - captured["url"] = url - return b"# SKILL\n", None - - monkeypatch.setattr(sd, "_http_get_bytes", fake_get_bytes) - - body, reason = sd.fetch_skill_file(WS, "token", "main", "default", "triage", "SKILL.md") - - assert reason is None - assert body == b"# SKILL\n" - assert captured["url"] == f"{WS}/api/2.0/fs/files/Skills/main/default/triage/SKILL.md" - - def test_http_failure_propagates_reason(self, monkeypatch): - monkeypatch.setattr( - sd, "_http_get_bytes", lambda url, token, timeout=30: (None, "HTTP 404 Not Found") - ) - - body, reason = sd.fetch_skill_file(WS, "token", "main", "default", "triage", "gone.md") - - assert body is None - assert reason == "HTTP 404 Not Found" - - -class TestFetchSkillBundle: - def test_assembles_relpath_to_bytes_map(self, monkeypatch): - contents = {"SKILL.md": b"# skill", "references/a.md": b"aaa"} - monkeypatch.setattr(sd, "list_skill_files", lambda *a, **k: (list(contents), None)) - monkeypatch.setattr( - sd, "fetch_skill_file", lambda ws, tok, c, s, leaf, rel: (contents[rel], None) - ) - - bundle, reason = sd.fetch_skill_bundle(WS, "token", "main", "default", "triage") - - assert reason is None - assert bundle == contents - - def test_listing_failure_propagates_reason(self, monkeypatch): - monkeypatch.setattr(sd, "list_skill_files", lambda *a, **k: ([], "HTTP 404 Not Found")) - - bundle, reason = sd.fetch_skill_bundle(WS, "token", "main", "default", "triage") - - assert bundle is None - assert reason == "HTTP 404 Not Found" - - def test_file_failure_aborts_whole_bundle(self, monkeypatch): - monkeypatch.setattr( - sd, "list_skill_files", lambda *a, **k: (["SKILL.md", "broken.md"], None) - ) - monkeypatch.setattr( - sd, - "fetch_skill_file", - lambda ws, tok, c, s, leaf, rel: ( - (b"ok", None) if rel == "SKILL.md" else (None, "HTTP 500 Server Error") - ), - ) - - bundle, reason = sd.fetch_skill_bundle(WS, "token", "main", "default", "triage") - - assert bundle is None - assert reason == "HTTP 500 Server Error" - - class TestSkillDirRoots: def test_roots_under_project_dir(self, tmp_path): roots = skill_dir_roots(str(tmp_path)) @@ -668,42 +355,6 @@ def test_failed_fetch_counts_toward_total_but_not_written(self, tmp_path, monkey assert not (tmp_path / ".claude/skills/bad").exists() -class TestGetSkill: - def test_returns_ref_with_location_parsed_from_fqn(self, monkeypatch): - captured = {} - - def fake_get(url, token, timeout=30): - captured["url"] = url - return { - "name": "skills/ml.prod.pii-handling", - "bundle_name": "pii-handling", - "finalize_time": "2026-06-26T05:58:25Z", - }, None - - monkeypatch.setattr(sd, "_http_get_json", fake_get) - - result = sd.get_skill(WS, "token", "ml.prod.pii-handling") - - assert result == ref("pii-handling", catalog="ml", schema="prod") - assert captured["url"] == f"{WS}/api/2.1/unity-catalog/skills/ml.prod.pii-handling" - - def test_not_found_returns_none(self, monkeypatch): - monkeypatch.setattr( - sd, "_http_get_json", lambda url, token, timeout=30: (None, "HTTP 404 Not Found") - ) - - assert sd.get_skill(WS, "token", "main.default.gone") is None - - def test_unfinalized_skill_returns_none(self, monkeypatch): - monkeypatch.setattr( - sd, - "_http_get_json", - lambda url, token, timeout=30: ({"name": "skills/main.default.draft"}, None), - ) - - assert sd.get_skill(WS, "token", "main.default.draft") is None - - class TestDownloadSelectedSkills: def test_downloads_each_resolved_fqn(self, tmp_path, monkeypatch): by_fqn = { @@ -781,30 +432,6 @@ def test_records_workspace_id_when_known(self, tmp_path, monkeypatch): assert record["workspace_id"] == "org-42" -class TestSkillRefMetadata: - def test_captures_uc_attribution_fields(self, monkeypatch): - payload = { - "skills": [ - { - "name": "skills/main.default.triage", - "bundle_name": "triage", - "finalize_time": "2026-06-26T05:58:25Z", - "id": "skill-uuid", - "metastore_id": "metastore-uuid", - "update_time": "2026-06-26T05:58:25Z", - } - ] - } - monkeypatch.setattr(sd, "_http_get_json", lambda url, token, timeout=30: (payload, None)) - - (skill,), reason = sd.list_schema_skills(WS, "token", "main", "default") - - assert reason is None - assert skill.metastore_id == "metastore-uuid" - assert skill.skill_id == "skill-uuid" - assert skill.uc_update_time == "2026-06-26T05:58:25Z" - - class TestDownloadManagedSkillsOnLaunch: def test_writes_missing_skills_and_returns_their_bundle_names(self, tmp_path, monkeypatch): monkeypatch.setattr( @@ -969,18 +596,6 @@ def test_none_path_threads_through(self, monkeypatch): assert calls["register"] == (WS, "profile", ["claude"]) -def _walk_stub(schemas, reason=None): - """A ``walk_catalog_schemas`` stub that probes each (catalog, schema) in order.""" - - def walk(workspace, token, *, deadline, probe, collect, **kwargs): - total = len(schemas) - for done, (catalog, schema) in enumerate(schemas, start=1): - collect(probe(catalog, schema), done, total) - return reason - - return walk - - class _FakePrompt: def __init__(self, result): self._result = result @@ -989,75 +604,6 @@ def ask(self): return self._result -class TestListAllSkills: - def test_flattens_and_streams_across_schemas(self, monkeypatch): - monkeypatch.setattr( - sd, "walk_catalog_schemas", _walk_stub([("main", "default"), ("ml", "prod")]) - ) - by_schema = { - "main.default": [ref("triage"), ref("pii")], - "ml.prod": [ref("scoring", catalog="ml", schema="prod")], - } - monkeypatch.setattr( - sd, "list_schema_skills", lambda ws, tok, c, s: (by_schema[f"{c}.{s}"], None) - ) - streamed = [] - progress = [] - - refs, reason = sd.list_all_skills( - WS, - "token", - on_skills=streamed.append, - on_progress=lambda done, total, found: progress.append((done, total, found)), - ) - - assert reason is None - assert [r.fqn for r in refs] == [ - "main.default.pii", - "main.default.triage", - "ml.prod.scoring", - ] - assert [[r.fqn for r in batch] for batch in streamed] == [ - ["main.default.pii", "main.default.triage"], - ["ml.prod.scoring"], - ] - assert progress == [(1, 2, 2), (2, 2, 3)] - - def test_dedupes_repeated_fqns(self, monkeypatch): - monkeypatch.setattr( - sd, "walk_catalog_schemas", _walk_stub([("main", "default"), ("main", "default")]) - ) - monkeypatch.setattr(sd, "list_schema_skills", lambda ws, tok, c, s: ([ref("triage")], None)) - streamed = [] - - refs, reason = sd.list_all_skills(WS, "token", on_skills=streamed.append) - - assert [r.fqn for r in refs] == ["main.default.triage"] - assert streamed == [[ref("triage")]] - - def test_walk_failure_returns_its_reason(self, monkeypatch): - monkeypatch.setattr( - sd, "walk_catalog_schemas", _walk_stub([], reason="no UC catalogs found") - ) - - assert sd.list_all_skills(WS, "token") == ([], "no UC catalogs found") - - def test_empty_walk_reports_no_skills(self, monkeypatch): - monkeypatch.setattr(sd, "walk_catalog_schemas", _walk_stub([("main", "default")])) - monkeypatch.setattr(sd, "list_schema_skills", lambda *a, **k: ([], None)) - - assert sd.list_all_skills(WS, "token") == ([], "no skills found") - - def test_timeout_returns_partial_results_with_reason(self, monkeypatch): - monkeypatch.setattr(sd, "walk_catalog_schemas", _walk_stub([("main", "default")])) - monkeypatch.setattr(sd, "list_schema_skills", lambda ws, tok, c, s: ([ref("triage")], None)) - - refs, reason = sd.list_all_skills(WS, "token", deadline_seconds=-1) - - assert [r.fqn for r in refs] == ["main.default.triage"] - assert reason == sd._SKILLS_WALK_TIMEOUT_REASON - - class TestSkillDownloadPicker: def test_choice_value_is_fqn_flags_on_disk_and_carries_description(self, tmp_path): roots = skill_dir_roots(str(tmp_path))