From bb49224f6356331397bac8fc83ca1dcce793c0c3 Mon Sep 17 00:00:00 2001 From: Xiang Shen Date: Sun, 13 Sep 2026 04:26:41 +0000 Subject: [PATCH 1/2] [skills] Add an interactive schema picker to `ug skill add --mcp` `ug skill add --mcp` with no --location errored asking for a typed schema. On an interactive terminal it now opens a picker of the workspace's schemas that contain skills, streamed in as discovery finds them, and adds the selected schemas to the skills MCP scope. The picker mirrors the download picker and reuses its machinery: the workspace walk (list_all_skills) and the streaming checkbox widget. Because list_all_skills probes one schema per call, each batch is one schema's skill set, so the loader renders one row per schema with its exact skill count. Already-scoped schemas are flagged and stay selectable, since adding is additive (removal is `ug skill remove --mcp`). To let both the picker (in mcp.py) and the download flow use the workspace walk without a circular import, the UC skills read layer (SkillRef, the API client, and list_all_skills) moves into a new skills_api.py that depends only on databricks.py; mcp.py and skills_download.py both import it. add_skills_command's apply body is extracted into add_skill_locations_to_mcp, shared by the command and the picker. Co-authored-by: Isaac --- src/ucode/cli.py | 21 +- src/ucode/mcp.py | 100 ++++++- src/ucode/skills_api.py | 281 ++++++++++++++++++++ src/ucode/skills_download.py | 285 +------------------- tests/test_cli.py | 34 +++ tests/test_mcp.py | 112 ++++++++ tests/test_skills_api.py | 482 ++++++++++++++++++++++++++++++++++ tests/test_skills_download.py | 454 -------------------------------- 8 files changed, 1029 insertions(+), 740 deletions(-) create mode 100644 src/ucode/skills_api.py create mode 100644 tests/test_skills_api.py diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 8cdf7023..5c9ed8d5 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -98,6 +98,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, @@ -1230,8 +1231,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 = ( @@ -1262,15 +1263,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 df22a8b0..cb8eaf5b 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -34,6 +34,7 @@ CURSOR_OAUTH_CLIENT_ID, oauth_client_available, ) +from ucode.skills_api import SkillRef, list_all_skills from ucode.state import load_full_state, load_state, save_state from ucode.ui import ( _BACK, @@ -1973,6 +1974,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. @@ -1981,12 +2004,77 @@ 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]], 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]) -> 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)]) + + list_all_skills(workspace, token, on_skills=on_skills) + + return loader + + +def prompt_for_skill_schema_choices( + background_loader: Callable[[Callable[[list[questionary.Choice]], None]], 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 beeaee98..a127313a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1597,6 +1597,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 a0bce3d5..8a36bb76 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -2283,6 +2283,118 @@ 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 = [] + + mcp._skill_schema_background_loader(WS, "token", {"ml.prod"})(appended.extend) + + 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_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)) From b85013e6c82f7567b0680adb885024f422b15ca2 Mon Sep 17 00:00:00 2001 From: xiang-shen_data Date: Tue, 15 Sep 2026 23:42:33 +0000 Subject: [PATCH 2/2] Address review: warn when the skill-schema picker walk times out The `ug skill add --mcp` schema picker's loader dropped list_all_skills' return, so a timed-out walk silently showed a partial schema list. Capture the deadline and surface it in the footer, mirroring the download picker. Co-authored-by: Isaac --- src/ucode/mcp.py | 22 +++++++++++++++++----- tests/test_mcp.py | 18 +++++++++++++++++- 2 files changed, 34 insertions(+), 6 deletions(-) diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index cb8eaf5b..ca82003b 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -34,7 +34,12 @@ CURSOR_OAUTH_CLIENT_ID, oauth_client_available, ) -from ucode.skills_api import SkillRef, list_all_skills +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, @@ -2023,25 +2028,32 @@ def _skill_schema_choice(location: str, skill_count: int, in_scope: bool) -> que def _skill_schema_background_loader( workspace: str, token: str, in_scope: set[str] -) -> Callable[[Callable[[list[questionary.Choice]], None]], None]: +) -> 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]) -> None: + 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)]) - list_all_skills(workspace, token, on_skills=on_skills) + 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]], None], + 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( diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 8a36bb76..1b3e323a 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -2332,12 +2332,28 @@ def fake_list_all(ws, tok, *, on_skills=None, **kwargs): monkeypatch.setattr(mcp, "list_all_skills", fake_list_all) appended = [] - mcp._skill_schema_background_loader(WS, "token", {"ml.prod"})(appended.extend) + 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 = {}