Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 39 additions & 18 deletions src/ucode/skills_download.py
Original file line number Diff line number Diff line change
Expand Up @@ -197,27 +197,36 @@ def _write_bundle(skill_dir: Path, leaf: str, files: dict[str, bytes]) -> None:
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.
def existing_skill_on_disk(roots: list[Path], leaf: str) -> bool:
"""Whether ``leaf`` already has a skill directory under any root."""
return any((root / leaf).exists() for root in roots)

Prompts before overwriting an existing skill dir. ``location`` is the source
``<catalog>.<schema>``, shown in that prompt. Returns True if the skill was
written, False if it was skipped or kept.

def should_download_skill(roots: list[Path], leaf: str, *, location: str) -> bool:
"""Whether ``leaf`` should be fetched and written into ``roots``.
Applies the disk-only checks that need no bundle bytes, so a declined or
invalid skill is never downloaded: skips invalid leaf names, and prompts
before overwriting a skill already on disk (``location`` is the source
``<catalog>.<schema>`` shown in that prompt).
"""
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(
if existing_skill_on_disk(roots, leaf) 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

return True


def write_skill(roots: list[Path], leaf: str, files: dict[str, bytes]) -> None:
"""Write ``leaf``'s bundle (``{relpath: bytes}``) into every root."""
for root in roots:
_write_bundle(root / leaf, leaf, files)
return True


# --- Orchestration ---------------------------------------------------------
Expand Down Expand Up @@ -256,12 +265,21 @@ def download_skills(
) -> None:
"""Download every skill in each ``<catalog>.<schema>`` 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.
When ``skills`` is given, only those leaf names are downloaded; names absent
from a schema warn and are skipped. ``None`` downloads the whole schema.
Locations are processed one at a time, and each runs three stages:
1. **List** the schema's skill leaves. When ``skills`` is given, restrict to
those leaf names; names absent from the schema warn and are skipped, and
``None`` keeps the whole schema.
2. **Decide** which to download via ``should_download_skill`` (skips invalid
names and prompts before overwriting a skill already on disk), so a
declined skill is never fetched.
3. **Fetch** the survivors' bundles concurrently (with a progress bar) and
**write** them.
Finishing one location before starting the next means a skill written for an
earlier location is already on disk when a same-named skill in a later
location reaches its decide stage, so the overwrite prompt still fires. A
failure on one skill warns and skips it without aborting the batch.
"""
roots = skill_dir_roots(path)
roots_display = " and ".join(str(root) for root in roots)
Expand All @@ -286,15 +304,18 @@ def download_skills(
print_note(f"No skills found in `{location}`.")
continue

bundles = _fetch_bundles(workspace, token, catalog, schema, leaves)
to_download = [
leaf for leaf in leaves if should_download_skill(roots, leaf, location=location)
]
bundles = _fetch_bundles(workspace, token, catalog, schema, to_download)
written = 0
for leaf in leaves:
for leaf in to_download:
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
write_skill(roots, leaf, files)
written += 1
console.print()
print_success(
f"Downloaded {written}/{len(leaves)} skill(s) from `{location}` in {roots_display}."
Expand Down
74 changes: 52 additions & 22 deletions tests/test_skills_download.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,12 @@
import pytest

import ucode.skills_download as sd
from ucode.skills_download import skill_dir_roots, write_skill
from ucode.skills_download import (
existing_skill_on_disk,
should_download_skill,
skill_dir_roots,
write_skill,
)

WS = "https://example.databricks.com"

Expand Down Expand Up @@ -237,50 +242,58 @@ def test_missing_directory_rejected(self, tmp_path):
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):
class TestShouldDownloadSkill:
def test_new_skill_is_downloaded(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)"
assert should_download_skill(roots, "triage", location="main.default")

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")
write_skill(roots, "triage", {"SKILL.md": b"from-main"})

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"
assert not should_download_skill(roots, "triage", location="ml.prod")

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")
write_skill(roots, "triage", {"SKILL.md": b"from-main"})

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"
assert should_download_skill(roots, "triage", location="ml.prod")

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 should_download_skill(roots, "Bad_Name", location="main.default")

def test_existing_skill_on_disk_checks_every_root(self, tmp_path):
roots = skill_dir_roots(str(tmp_path))
assert not existing_skill_on_disk(roots, "triage")

(roots[1] / "triage").mkdir(parents=True)
assert existing_skill_on_disk(roots, "triage")


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_skill(roots, "triage", files)

assert not (roots[0] / "Bad_Name").exists()
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_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"})
write_skill(
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()
Expand Down Expand Up @@ -322,6 +335,23 @@ def test_list_failure_skips_location(self, tmp_path, monkeypatch):

assert called == []

def test_declined_skill_is_not_fetched(self, tmp_path, monkeypatch):
roots = skill_dir_roots(str(tmp_path))
write_skill(roots, "triage", {"SKILL.md": b"kept"})
monkeypatch.setattr(sd, "list_schema_skills", lambda *a, **k: (["triage"], None))
monkeypatch.setattr(sd, "prompt_yes_no", lambda _: False)
fetched = []
monkeypatch.setattr(
sd,
"fetch_skill_bundle",
lambda ws, tok, c, s, leaf: fetched.append(leaf) or ({"SKILL.md": b"new"}, None),
)

sd.download_skills(WS, "token", ["main.default"], str(tmp_path))

assert fetched == []
assert (roots[0] / "triage/SKILL.md").read_bytes() == b"kept"

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(
Expand Down
Loading