diff --git a/AGENTS.md b/AGENTS.md index deeb6c7..7810cb9 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -20,7 +20,7 @@ Tests live in `tests/`. - Use Python 3.12+. - Keep changes scoped to the requested behavior. -- Follow the existing module boundaries: CLI orchestration in `cli.py`, agent-specific behavior in `agents/.py`, shared agent dispatch in `agents/__init__.py`, Databricks calls in `databricks.py`, and presentation helpers in `ui.py`. +- Follow the existing module boundaries: CLI orchestration in `cli.py`, agent-specific behavior in `agents/.py`, shared agent dispatch in `agents/__init__.py`, Databricks calls in `databricks.py`, skill download (UC fetch client + on-disk writer + download orchestration) in `skills_download.py`, MCP-connection state glue in `mcp.py`, and presentation helpers in `ui.py`. Skill download persists no disk state — it writes files to `--path` (or the home dir) and registers only the schema-less skills MCP connection. - Prefer existing helpers for config file writes, state persistence, UI messages, and Databricks authentication. - Add or update focused tests for behavior changes. - Do not modify generated or lock files unless the dependency graph intentionally changes. diff --git a/README.md b/README.md index dad25e0..346fb80 100644 --- a/README.md +++ b/README.md @@ -102,6 +102,28 @@ ucode configure --agents claude --mcp system.ai.slack `--mcp` also works without `--agents` for MCP-only clients (it configures just the workspace, then registers the servers); pass a comma-separated list to register several at once. +### Skills (optional) + +Configure Unity Catalog Skills for your coding tools with `ucode configure skills`. It has two +mutually-exclusive modes, both scoped by `--location .` (comma-separated for +multiple schemas): + +```bash +# Download mode (default): fetch every skill in the schema to disk. +ucode configure skills --location main.default --path /abs/project/dir + +# MCP mode: expose the schema's skills as MCP tools instead of downloading. +ucode configure skills --location main.default,ml.prod --mcp +``` + +- **Download mode** writes each skill flat as `/SKILL.md` (plus its bundled files) into both + `.claude/skills/` and `.agents/skills/`. `--path` (an existing absolute directory) is optional; + when omitted, skills are written under your home directory. Any pre-existing skill dir prompts + before it's overwritten. It then registers a schema-less skills MCP connection (utility tools + only), leaving any prior `--mcp` scope untouched. +- **MCP mode** sets the connection's location set to exactly `` (override-only) and rebuilds + its `?schema=` URL; no files are downloaded and `--path` is rejected. + --- ## Other Commands @@ -118,6 +140,8 @@ then registers the servers); pass a comma-separated list to register several at | `ucode configure --profiles DEFAULT --use-pat` | Authenticate with the profile's personal access token — no browser login | | `ucode configure --skip-validate` | Write configs without sending a test message through each agent | | `ucode configure --agents claude --mcp system.ai.slack` | Configure an agent and register its Databricks MCP server(s) in one command | +| `ucode configure skills --location main.default [--path ]` | Download a schema's skills to disk (under ``, or your home dir) and register a schema-less skills MCP connection | +| `ucode configure skills --location main.default --mcp` | Expose a schema's skills as MCP tools (override-only) instead of downloading | ## Managed Local Files diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 865612d..975eb0b 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -59,6 +59,7 @@ purge_cross_workspace_mcp_residue, revert_mcp_configs, ) +from ucode.skills_download import configure_skills_download_command from ucode.state import ( STATE_PATH, clear_state, @@ -1403,21 +1404,31 @@ def configure_skills( ], mcp: Annotated[ bool, + typer.Option("--mcp", help="Mutate the skills MCP connection instead of downloading."), + ] = False, + path: Annotated[ + str | None, typer.Option( - "--mcp", help="Manage the skills MCP connection (required until download mode lands)." + "--path", + help="(download) Existing absolute dir to download into; defaults to your home dir.", ), - ] = False, + ] = None, ) -> None: """Configure Databricks Skills for your coding tools. - ``--location`` sets the skills MCP connection's scope to exactly the listed - schemas, replacing any previous set. + By default, downloads every skill in each ``--location`` schema to disk + (under ``--path``, or your home dir when omitted) and registers a schema-less + MCP connection. With ``--mcp``, instead sets the skills MCP connection's scope + to exactly the listed schemas. """ try: - if not mcp: - raise RuntimeError("Download mode is not available yet; pass --mcp for now.") - configure_skills_mcp_command(_parse_skill_locations(location)) - except RuntimeError as exc: + if mcp: + if path is not None: + raise RuntimeError("--path is not valid with --mcp.") + configure_skills_mcp_command(_parse_skill_locations(location)) + else: + configure_skills_download_command(_parse_skill_locations(location), path=path) + except (RuntimeError, ValueError) as exc: print_err(str(exc)) raise typer.Exit(1) from None except KeyboardInterrupt: diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 19f151e..f9e3f1b 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -290,6 +290,35 @@ def _http_post_json( return None, f"network error: {exc.reason}" +def _http_get_bytes(url: str, token: str, *, timeout: int = 10) -> tuple[bytes | None, str | None]: + """GET raw bytes. Returns (body, None) on success, (None, reason) on failure. + + Like `_http_get_json` but leaves the body undecoded, since skill bundles can + contain binary files. + """ + request = urllib_request.Request(url, headers={"Authorization": f"Bearer {token}"}) + try: + with urllib_request.urlopen(request, timeout=timeout) as response: + body = response.read() + _debug(f"GET {url}", f"HTTP 200, {len(body)} bytes") + return body, None + except urllib_error.HTTPError as exc: + detail = "" + try: + detail = exc.read().decode("utf-8", errors="replace") if exc.fp else "" + except Exception: + detail = "" + _debug(f"GET {url}", f"HTTP {exc.code} {exc.reason}") + reason = f"HTTP {exc.code} {exc.reason}" + excerpt = detail.strip()[:200] + if excerpt: + reason = f"{reason}: {excerpt}" + return None, reason + except urllib_error.URLError as exc: + _debug(f"GET {url}", f"URLError: {exc.reason}") + return None, f"network error: {exc.reason}" + + def get_current_user_name(workspace: str, token: str) -> str | None: """Return the current user's login (email) via SCIM `Me`, or None on failure. diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 9ab931f..5f57860 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -1190,7 +1190,7 @@ def _resolve_location_mcp_servers( return [*working_servers, *_skills_entries(original_servers)] -def _setup_mcp_clients(state: dict, section: str) -> tuple[str, str | None, list[str]]: +def setup_mcp_clients(state: dict, section: str) -> tuple[str, str | None, list[str]]: """Validate the workspace, resolve configured MCP clients, and prepare auth. Returns ``(workspace, profile, clients)`` and prints the section header, the @@ -1253,7 +1253,7 @@ def configure_mcp_command(location: str | None = None, services: set[str] | None ) location = next(iter(schemas)) state = load_state() - workspace, profile, clients = _setup_mcp_clients(state, "MCP Servers") + workspace, profile, clients = setup_mcp_clients(state, "MCP Servers") original_mcp_servers_for_location: list[dict] = list(state.get("mcp_servers") or []) if location is not None: @@ -1424,6 +1424,21 @@ def configure_skills_mcp_command(locations: list[str]) -> int: """Set the skills MCP connection's ``skill_locations`` to exactly ``locations``, replacing any previous set.""" state = load_state() - workspace, _profile, clients = _setup_mcp_clients(state, "Skills MCP") + workspace, _profile, clients = setup_mcp_clients(state, "Skills MCP") _update_skills_mcp(state, workspace, clients, locations) return 0 + + +def _skill_mcp_locations(state: dict) -> list[str]: + """The skills MCP connection's ``skill_locations``, or ``[]`` if none exists.""" + entry = next(iter(_skills_entries(list(state.get("mcp_servers") or []))), None) + return list((entry or {}).get("skill_locations") or []) + + +def register_schemaless_skills_connection(state: dict, workspace: str, clients: list[str]) -> None: + """Register/keep the skills MCP connection without changing its schema set. + + Download mode calls this after writing files: it preserves any prior + ``--mcp`` ``skill_locations`` and otherwise registers the bare schema-less + route (utility tools only).""" + _update_skills_mcp(state, workspace, clients, _skill_mcp_locations(state)) diff --git a/src/ucode/skills_download.py b/src/ucode/skills_download.py new file mode 100644 index 0000000..db43108 --- /dev/null +++ b/src/ucode/skills_download.py @@ -0,0 +1,286 @@ +"""Download Unity Catalog skills and write them to disk, one flat dir per skill.""" + +from __future__ import annotations + +import re +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path +from urllib.parse import urlencode + +from ucode.databricks import ( + _http_get_bytes, + _http_get_json, + get_databricks_token, + workspace_hostname, +) +from ucode.mcp import register_schemaless_skills_connection, setup_mcp_clients +from ucode.state import load_state +from ucode.ui import print_note, print_success, print_warning, progress_bar, prompt_yes_no + +# `.claude/skills` (Claude) + `.agents/skills` (the alias other agents read). +SKILL_BASE_DIR_NAMES = (".claude/skills", ".agents/skills") + +SKILL_NAME_PATTERN = re.compile(r"^[a-z0-9-]+$") + +# Parallel skill fetches per schema; writes stay sequential (they prompt). +_MAX_FETCH_WORKERS = 8 + + +# --- Download client (UC skills API + Files API) --------------------------- + + +def _skill_bundle_name(skill: dict) -> str | None: + """The downloadable leaf name of a skill, or None if it isn't finalized. + + Only finalized skills (those with a ``finalize_time``) have bundle content + to download. ``bundle_name`` is the leaf; fall back to the last dotted + segment of the resource ``name`` (``skills/..``). + """ + if not skill.get("finalize_time"): + return None + bundle_name = skill.get("bundle_name") + if isinstance(bundle_name, str) and bundle_name: + return bundle_name + name = skill.get("name") + return name.rsplit(".", 1)[-1] if isinstance(name, str) else None + + +def list_schema_skills( + workspace: str, token: str, catalog: str, schema: str +) -> tuple[list[str], str | None]: + """List the finalized skill leaf names 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}"} + + leaves: list[str] = [] + 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 []: + leaf = _skill_bundle_name(skill) if isinstance(skill, dict) else None + if leaf: + leaves.append(leaf) + page_token = data.get("next_page_token") + if not page_token: + return leaves, None + + +def list_skill_files( + workspace: str, token: str, catalog: str, schema: str, leaf: str +) -> tuple[list[str], str | None]: + """List a skill bundle's files, as paths relative to the skill directory. + + Recursively walks the skill's UC Volume directory (including ``SKILL.md``). + A non-None reason indicates the listing call itself failed. + """ + hostname = workspace_hostname(workspace) + dirs_base = f"https://{hostname}/api/2.0/fs/directories" + volume_prefix = f"/Volumes/{catalog}/{schema}/{leaf}/" + + relative_paths: list[str] = [] + pending = [f"Volumes/{catalog}/{schema}/{leaf}"] + 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(volume_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, leaf: str, relative_path: str +) -> tuple[bytes | None, str | None]: + """Fetch one skill bundle file's raw bytes from its UC Volume.""" + hostname = workspace_hostname(workspace) + url = f"https://{hostname}/api/2.0/fs/files/Volumes/{catalog}/{schema}/{leaf}/{relative_path}" + return _http_get_bytes(url, token, timeout=30) + + +def fetch_skill_bundle( + workspace: str, token: str, catalog: str, schema: str, leaf: 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, leaf) + if reason: + return None, reason + bundle: dict[str, bytes] = {} + for relative_path in relative_paths: + content, reason = fetch_skill_file(workspace, token, catalog, schema, leaf, relative_path) + if content is None: + return None, reason + bundle[relative_path] = content + return bundle, None + + +# --- On-disk writer -------------------------------------------------------- + + +def skill_dir_roots(project_dir: str | None) -> list[Path]: + """The ``.claude/skills`` and ``.agents/skills`` roots to download into. + + ``project_dir`` must be an existing absolute directory when given; when + omitted, roots default to the user's home directory (user scope). + """ + if project_dir is None: + base = Path.home() + else: + base = Path(project_dir) + if not base.is_absolute(): + raise ValueError(f"--path must be an absolute path, got `{project_dir}`.") + if not base.is_dir(): + raise ValueError(f"--path directory does not exist: `{project_dir}`.") + return [base / name for name in SKILL_BASE_DIR_NAMES] + + +def _is_valid_leaf(leaf: str) -> bool: + return bool(SKILL_NAME_PATTERN.match(leaf)) + + +def _safe_relative_path(relative_path: str) -> Path | None: + """A bundle file's path within its skill dir, or None if it escapes the dir. + + The Files API returns server-controlled paths, but ucode writes them to + disk, so reject absolute paths and any ``..`` traversal. + """ + path = Path(relative_path) + if path.is_absolute() or ".." in path.parts: + return None + return path + + +def _write_bundle(skill_dir: Path, leaf: str, files: dict[str, bytes]) -> None: + for relative_path, content in files.items(): + safe_path = _safe_relative_path(relative_path) + if safe_path is None: + print_warning(f"Skipping unsafe path in `{leaf}`: {relative_path}") + continue + destination = skill_dir / safe_path + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(content) + + +def write_skill(roots: list[Path], leaf: str, files: dict[str, bytes], *, location: str) -> bool: + """Write ``leaf``'s bundle (``{relpath: bytes}``) into every root. + + Prompts before overwriting an existing skill dir. ``location`` is the source + ``.``, shown in that prompt. Returns True if the skill was + written, False if it was skipped or kept. + """ + if not _is_valid_leaf(leaf): + print_warning(f"Skipping `{leaf}`: not a valid skill name (lowercase a-z, 0-9, -).") + return False + + already_on_disk = any((root / leaf).exists() for root in roots) + if already_on_disk and not prompt_yes_no( + f"A skill named `{leaf}` already exists. Overwrite it with `{location}.{leaf}`?" + ): + print_note(f"Kept existing `{leaf}`.") + return False + + for root in roots: + _write_bundle(root / leaf, leaf, files) + return True + + +# --- Orchestration --------------------------------------------------------- + + +def _fetch_bundles( + workspace: str, token: str, catalog: str, schema: str, leaves: list[str] +) -> dict[str, tuple[dict[str, bytes] | None, str | None]]: + """Fetch every leaf's bundle concurrently, keyed by leaf name. + + Renders a ``k/n`` progress bar that advances as each fetch completes. + """ + if not leaves: + return {} + results: dict[str, tuple[dict[str, bytes] | None, str | None]] = {} + with ( + progress_bar(f"Fetching skills from {catalog}.{schema}", len(leaves)) as advance, + ThreadPoolExecutor(max_workers=min(_MAX_FETCH_WORKERS, len(leaves))) as pool, + ): + futures = { + pool.submit(fetch_skill_bundle, workspace, token, catalog, schema, leaf): leaf + for leaf in leaves + } + for future in as_completed(futures): + results[futures[future]] = future.result() + advance() + return results + + +def download_skills(workspace: str, token: str, locations: list[str], path: str | None) -> None: + """Download every skill in each ``.`` location to disk. + + Bundles are fetched concurrently (with a progress bar) per schema, then + written sequentially so overwrite prompts don't interleave. A failure on one + skill warns and skips it without aborting the batch. + """ + roots = skill_dir_roots(path) + for location in locations: + catalog, schema = location.split(".") + leaves, reason = list_schema_skills(workspace, token, catalog, schema) + if reason: + print_warning(f"Skipping `{location}`: {reason}.") + continue + if not leaves: + print_note(f"No skills found in `{location}`.") + continue + + bundles = _fetch_bundles(workspace, token, catalog, schema, leaves) + written = 0 + for leaf in leaves: + files, reason = bundles[leaf] + if reason or files is None: + print_warning(f"Skipping `{location}.{leaf}`: {reason}.") + continue + if write_skill(roots, leaf, files, location=location): + written += 1 + print_success(f"Downloaded {written}/{len(leaves)} skill(s) from `{location}`.") + + +def configure_skills_download_command(locations: list[str], *, path: str | None) -> int: + """Download every skill in each schema to disk and register the skills connection. + + Downloads to ``path`` (or the home dir when None), then registers/keeps the + schema-less MCP connection. ``skill_locations`` is never touched, so a prior + ``--mcp`` set survives a download run.""" + state = load_state() + workspace, profile, clients = setup_mcp_clients(state, "Skills") + token = get_databricks_token(workspace, profile) + + download_skills(workspace, token, locations, path) + + register_schemaless_skills_connection(state, workspace, clients) + return 0 diff --git a/src/ucode/ui.py b/src/ucode/ui.py index c1b6ee3..cec7950 100644 --- a/src/ucode/ui.py +++ b/src/ucode/ui.py @@ -7,12 +7,14 @@ import textwrap import threading import time +from collections.abc import Callable, Iterator from contextlib import contextmanager from datetime import timedelta import questionary from rich.console import Console from rich.panel import Panel +from rich.progress import BarColumn, MofNCompleteColumn, Progress, TextColumn console = Console(highlight=False) err_console = Console(stderr=True, highlight=False) @@ -113,6 +115,27 @@ def spin() -> None: thread.join(timeout=1) +@contextmanager +def progress_bar(description: str, total: int) -> Iterator[Callable[[], None]]: + """Yield an ``advance()`` callback that drives a ``k/n`` progress bar. + + Falls back to no live bar off a tty (e.g. CI), so logs stay single-line. + """ + if total <= 0 or not sys.stdout.isatty(): + yield lambda: None + return + + with Progress( + TextColumn("[dim]{task.description}[/dim]"), + BarColumn(), + MofNCompleteColumn(), + console=console, + transient=True, + ) as progress: + task = progress.add_task(description, total=total) + yield lambda: progress.advance(task) + + def render_box_table( headers: list[str], rows: list[list[str]], diff --git a/tests/test_cli.py b/tests/test_cli.py index c66c9c1..8e7473d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -364,11 +364,32 @@ def test_comma_location_yields_multiple_schemas(self): assert result.exit_code == 0, result.output mock_mcp.assert_called_once_with(["a.b", "c.d"]) - def test_without_mcp_is_not_implemented_exit_1(self): - with patch("ucode.cli.configure_skills_mcp_command") as mock_mcp: + def test_default_mode_dispatches_download_with_path(self): + with patch("ucode.cli.configure_skills_download_command") as mock_download: + result = runner.invoke( + app, ["configure", "skills", "--location", "a.b", "--path", "/tmp/skills"] + ) + assert result.exit_code == 0, result.output + mock_download.assert_called_once_with(["a.b"], path="/tmp/skills") + + def test_default_mode_without_path_dispatches_download(self): + with patch("ucode.cli.configure_skills_download_command") as mock_download: result = runner.invoke(app, ["configure", "skills", "--location", "a.b"]) + assert result.exit_code == 0, result.output + mock_download.assert_called_once_with(["a.b"], path=None) + + def test_path_with_mcp_exit_1(self): + with ( + patch("ucode.cli.configure_skills_mcp_command") as mock_mcp, + patch("ucode.cli.configure_skills_download_command") as mock_download, + ): + result = runner.invoke( + app, ["configure", "skills", "--location", "a.b", "--mcp", "--path", "/tmp/skills"] + ) assert result.exit_code == 1 + assert "--path" in _strip_ansi(result.output) mock_mcp.assert_not_called() + mock_download.assert_not_called() def test_three_part_location_exit_1(self): with patch("ucode.cli.configure_skills_mcp_command") as mock_mcp: diff --git a/tests/test_mcp.py b/tests/test_mcp.py index a520afa..b428cea 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1876,6 +1876,45 @@ def test_preserves_mcp_service_entries_across_set(self, monkeypatch): assert names.count(mcp.SKILLS_MCP_SERVER_NAME) == 1 +class TestSkillMcpLocations: + def test_reads_locations_off_skills_entry(self): + state = _skills_state(mcp._resolve_skills_mcp_servers(WS, ["claude"], ["A.a", "B.b"], [])) + assert mcp._skill_mcp_locations(state) == ["A.a", "B.b"] + + def test_empty_when_no_skills_entry(self): + assert mcp._skill_mcp_locations(_skills_state([])) == [] + assert mcp._skill_mcp_locations(_skills_state()) == [] + + +class TestRegisterSchemalessSkillsConnection: + def _stub(self, monkeypatch): + saved_states: list[dict] = [] + monkeypatch.setattr(mcp, "configure_client_mcp_server", lambda *a, **kw: []) + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude"]) + monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) + return saved_states + + def test_registers_bare_route_when_none_exists(self, monkeypatch): + saved_states = self._stub(monkeypatch) + state = _skills_state([]) + + mcp.register_schemaless_skills_connection(state, WS, ["claude"]) + + skills = _find_skills(saved_states[-1]["mcp_servers"]) + assert len(skills) == 1 + assert skills[0]["skill_locations"] == [] + assert skills[0]["url"] == f"{WS}/ai-gateway/skills/" + + def test_preserves_prior_mcp_location_set(self, monkeypatch): + self._stub(monkeypatch) + prior = mcp._resolve_skills_mcp_servers(WS, ["claude"], ["X.x", "Y.y"], []) + state = _skills_state(prior) + + mcp.register_schemaless_skills_connection(state, WS, ["claude"]) + + assert _find_skills(state["mcp_servers"])[0]["skill_locations"] == ["X.x", "Y.y"] + + class TestRevertMcpConfigs: def test_removes_cli_registered_servers_and_restores_copilot_config(self, monkeypatch): removed: list[tuple[str, str]] = [] @@ -1922,3 +1961,45 @@ def test_removes_cli_registered_servers_and_restores_copilot_config(self, monkey "opencode": True, "copilot": True, } + + def test_removes_skills_registry_across_its_clients(self, monkeypatch): + removed: list[tuple[str, str]] = [] + monkeypatch.setattr( + mcp, + "remove_client_mcp_server", + lambda client, name: removed.append((client, name)) or ["user"], + ) + monkeypatch.setattr(mcp, "restore_file", lambda *a, **kw: False) + + skills_entry = mcp._resolve_skills_mcp_servers(WS, ["claude", "codex"], ["a.b"], [])[0] + mcp.revert_mcp_configs({"mcp_servers": [skills_entry]}) + + assert removed == [ + ("claude", mcp.SKILLS_MCP_SERVER_NAME), + ("codex", mcp.SKILLS_MCP_SERVER_NAME), + ] + + +class TestPurgeCrossWorkspaceSkillsEntry: + def test_drops_foreign_workspace_skills_entry(self, monkeypatch): + removed: list[tuple[str, str]] = [] + saved_states: list[dict] = [] + foreign = "https://other.databricks.com" + skills_entry = mcp._resolve_skills_mcp_servers(foreign, ["claude"], ["a.b"], [])[0] + # The skills URL carries a `?schema=` query; its host must still parse. + assert mcp._mcp_entry_url_host(skills_entry) == "other.databricks.com" + state = {"mcp_servers": [skills_entry]} + + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["claude"]) + monkeypatch.setattr(mcp, "load_full_state", lambda: {}) + monkeypatch.setattr( + mcp, + "remove_client_mcp_server", + lambda client, name: removed.append((client, name)) or ["user"], + ) + monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) + + mcp.purge_cross_workspace_mcp_residue(state, WS) + + assert removed == [("claude", mcp.SKILLS_MCP_SERVER_NAME)] + assert state["mcp_servers"] == [] diff --git a/tests/test_skills_download.py b/tests/test_skills_download.py new file mode 100644 index 0000000..47232d9 --- /dev/null +++ b/tests/test_skills_download.py @@ -0,0 +1,396 @@ +"""Tests for skills_download.py — the UC skill-download client, on-disk writer, +and download orchestration.""" + +from __future__ import annotations + +import pytest + +import ucode.skills_download as sd +from ucode.skills_download import skill_dir_roots, write_skill + +WS = "https://example.databricks.com" + + +class TestListSchemaSkills: + def test_keeps_finalized_skills_and_uses_bundle_name(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)) + + leaves, reason = sd.list_schema_skills(WS, "token", "main", "default") + + assert reason is None + assert leaves == ["pii-handling", "triage"] + + def test_falls_back_to_resource_name_leaf(self, monkeypatch): + payload = { + "skills": [ + { + "name": "skills/main.default.pii-handling", + "finalize_time": "2026-06-26T05:58:25Z", + } + ] + } + monkeypatch.setattr(sd, "_http_get_json", lambda url, token, timeout=30: (payload, None)) + + leaves, reason = sd.list_schema_skills(WS, "token", "main", "default") + + assert reason is None + assert leaves == ["pii-handling"] + + def test_follows_pagination(self, monkeypatch): + pages = [ + {"skills": [{"bundle_name": "a", "finalize_time": "t"}], "next_page_token": "tok"}, + {"skills": [{"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) + + leaves, reason = sd.list_schema_skills(WS, "token", "main", "default") + + assert reason is None + assert leaves == ["a", "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_walks_nested_directories_into_relative_paths(self, monkeypatch): + # The Files API returns absolute `/Volumes/...` paths. + vol = "/Volumes/main/default/triage" + listings = { + "Volumes/main/default/triage": { + "contents": [ + {"path": f"{vol}/SKILL.md", "is_directory": False}, + {"path": f"{vol}/references/", "is_directory": True}, + ] + }, + "Volumes/main/default/triage/references": { + "contents": [{"path": f"{vol}/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): + vol = "/Volumes/main/default/triage" + pages = [ + { + "contents": [{"path": f"{vol}/a.md", "is_directory": False}], + "next_page_token": "tok", + }, + {"contents": [{"path": f"{vol}/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/Volumes/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)) + assert roots == [tmp_path / ".claude/skills", tmp_path / ".agents/skills"] + + def test_defaults_to_home_when_omitted(self, tmp_path, monkeypatch): + monkeypatch.setattr(sd.Path, "home", classmethod(lambda cls: tmp_path)) + roots = skill_dir_roots(None) + assert roots == [tmp_path / ".claude/skills", tmp_path / ".agents/skills"] + + def test_relative_path_rejected(self): + with pytest.raises(ValueError, match="absolute"): + skill_dir_roots("relative/dir") + + def test_missing_directory_rejected(self, tmp_path): + with pytest.raises(ValueError, match="does not exist"): + skill_dir_roots(str(tmp_path / "nope")) + + +def _write(roots, leaf, files, *, location="main.default"): + return write_skill(roots, leaf, files, location=location) + + +class TestWriteSkill: + def test_writes_bundle_into_every_root(self, tmp_path): + roots = skill_dir_roots(str(tmp_path)) + files = {"SKILL.md": b"# skill", "scripts/run.py": b"print(1)"} + + _write(roots, "triage", files) + + for root in roots: + assert (root / "triage/SKILL.md").read_bytes() == b"# skill" + assert (root / "triage/scripts/run.py").read_bytes() == b"print(1)" + + def test_existing_skill_prompt_keep(self, tmp_path, monkeypatch): + roots = skill_dir_roots(str(tmp_path)) + _write(roots, "triage", {"SKILL.md": b"from-main"}, location="main.default") + + monkeypatch.setattr(sd, "prompt_yes_no", lambda _: False) + _write(roots, "triage", {"SKILL.md": b"from-ml"}, location="ml.prod") + + assert (roots[0] / "triage/SKILL.md").read_bytes() == b"from-main" + + def test_existing_skill_prompt_overwrite(self, tmp_path, monkeypatch): + roots = skill_dir_roots(str(tmp_path)) + _write(roots, "triage", {"SKILL.md": b"from-main"}, location="main.default") + + monkeypatch.setattr(sd, "prompt_yes_no", lambda _: True) + _write(roots, "triage", {"SKILL.md": b"from-ml"}, location="ml.prod") + + assert (roots[0] / "triage/SKILL.md").read_bytes() == b"from-ml" + + def test_invalid_leaf_is_skipped(self, tmp_path): + roots = skill_dir_roots(str(tmp_path)) + + _write(roots, "Bad_Name", {"SKILL.md": b"x"}) + + assert not (roots[0] / "Bad_Name").exists() + + def test_path_traversal_is_rejected(self, tmp_path): + roots = skill_dir_roots(str(tmp_path)) + + _write(roots, "triage", {"SKILL.md": b"ok", "../escape.md": b"nope", "/abs.md": b"nope"}) + + assert (roots[0] / "triage/SKILL.md").read_bytes() == b"ok" + assert not (tmp_path / "escape.md").exists() + + +class TestFetchBundles: + def test_empty_leaves_returns_empty_without_pool(self): + # min(workers, 0) would raise ValueError in ThreadPoolExecutor; the + # early return keeps _fetch_bundles safe regardless of caller. + assert sd._fetch_bundles(WS, "token", "main", "default", []) == {} + + +class TestDownloadSkills: + def test_fetches_and_writes_each_leaf(self, tmp_path, monkeypatch): + monkeypatch.setattr( + sd, "list_schema_skills", lambda *a, **k: (["pii-handling", "triage"], None) + ) + bundles = { + "pii-handling": {"SKILL.md": b"pii"}, + "triage": {"SKILL.md": b"triage"}, + } + monkeypatch.setattr( + sd, "fetch_skill_bundle", lambda ws, tok, c, s, leaf: (bundles[leaf], None) + ) + + sd.download_skills(WS, "token", ["main.default"], str(tmp_path)) + + assert (tmp_path / ".claude/skills/pii-handling/SKILL.md").read_bytes() == b"pii" + assert (tmp_path / ".agents/skills/triage/SKILL.md").read_bytes() == b"triage" + + def test_list_failure_skips_location(self, tmp_path, monkeypatch): + monkeypatch.setattr(sd, "list_schema_skills", lambda *a, **k: ([], "HTTP 404 Not Found")) + called = [] + monkeypatch.setattr( + sd, "fetch_skill_bundle", lambda *a, **k: called.append(1) or (None, None) + ) + + sd.download_skills(WS, "token", ["main.default"], str(tmp_path)) + + assert called == [] + + def test_bundle_failure_skips_that_skill_only(self, tmp_path, monkeypatch): + monkeypatch.setattr(sd, "list_schema_skills", lambda *a, **k: (["good", "bad"], None)) + monkeypatch.setattr( + sd, + "fetch_skill_bundle", + lambda ws, tok, c, s, leaf: ( + ({"SKILL.md": b"ok"}, None) if leaf == "good" else (None, "HTTP 500 Server Error") + ), + ) + + sd.download_skills(WS, "token", ["main.default"], str(tmp_path)) + + assert (tmp_path / ".claude/skills/good/SKILL.md").read_bytes() == b"ok" + assert not (tmp_path / ".claude/skills/bad").exists() + + def test_prints_downloaded_count_summary(self, tmp_path, monkeypatch, capsys): + monkeypatch.setattr(sd, "list_schema_skills", lambda *a, **k: (["a", "b", "c"], None)) + monkeypatch.setattr(sd, "fetch_skill_bundle", lambda *a, **k: ({"SKILL.md": b"x"}, None)) + + sd.download_skills(WS, "token", ["main.default"], str(tmp_path)) + + assert "Downloaded 3/3 skill(s) from `main.default`" in capsys.readouterr().out + + def test_summary_counts_only_written_skills(self, tmp_path, monkeypatch, capsys): + monkeypatch.setattr(sd, "list_schema_skills", lambda *a, **k: (["good", "bad"], None)) + monkeypatch.setattr( + sd, + "fetch_skill_bundle", + lambda ws, tok, c, s, leaf: ( + ({"SKILL.md": b"ok"}, None) if leaf == "good" else (None, "HTTP 500 Server Error") + ), + ) + + sd.download_skills(WS, "token", ["main.default"], str(tmp_path)) + + assert "Downloaded 1/2 skill(s) from `main.default`" in capsys.readouterr().out + + +class TestConfigureSkillsDownloadCommand: + def _stub(self, monkeypatch): + calls: dict[str, object] = {} + monkeypatch.setattr(sd, "load_state", lambda: {"state": True}) + monkeypatch.setattr( + sd, "setup_mcp_clients", lambda state, section: (WS, "profile", ["claude"]) + ) + monkeypatch.setattr(sd, "get_databricks_token", lambda ws, profile: "token") + monkeypatch.setattr( + sd, + "download_skills", + lambda ws, tok, locations, path: calls.update(download=(ws, tok, locations, path)), + ) + monkeypatch.setattr( + sd, + "register_schemaless_skills_connection", + lambda state, ws, clients: calls.update(register=(ws, clients)), + ) + return calls + + def test_downloads_then_registers_connection(self, monkeypatch): + calls = self._stub(monkeypatch) + + assert sd.configure_skills_download_command(["a.b"], path="/tmp/skills") == 0 + + assert calls["download"] == (WS, "token", ["a.b"], "/tmp/skills") + assert calls["register"] == (WS, ["claude"]) + + def test_none_path_threads_through(self, monkeypatch): + calls = self._stub(monkeypatch) + + assert sd.configure_skills_download_command(["a.b"], path=None) == 0 + + assert calls["download"] == (WS, "token", ["a.b"], None)