Skip to content

Commit c468eec

Browse files
committed
skills: prompt before fetching so declined skills aren't downloaded
Download mode fetched every skill's bytes up front, then prompted to overwrite existing dirs at write time — so declining a skill threw away an already-completed download. Move the overwrite prompt and invalid-name check ahead of the fetch: split write_skill into should_download_skill (the disk-only decision, extracting existing_skill_on_disk) and a pure write_skill, and filter each schema's leaves through the decision before _fetch_bundles runs. The per-schema parallel fetch and the sequential location loop are unchanged, so cross-location same-leaf overwrite prompting still works.
1 parent 41ec3f9 commit c468eec

2 files changed

Lines changed: 80 additions & 37 deletions

File tree

src/ucode/skills_download.py

Lines changed: 30 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -197,27 +197,36 @@ def _write_bundle(skill_dir: Path, leaf: str, files: dict[str, bytes]) -> None:
197197
destination.write_bytes(content)
198198

199199

200-
def write_skill(roots: list[Path], leaf: str, files: dict[str, bytes], *, location: str) -> bool:
201-
"""Write ``leaf``'s bundle (``{relpath: bytes}``) into every root.
200+
def existing_skill_on_disk(roots: list[Path], leaf: str) -> bool:
201+
"""Whether ``leaf`` already has a skill directory under any root."""
202+
return any((root / leaf).exists() for root in roots)
202203

203-
Prompts before overwriting an existing skill dir. ``location`` is the source
204-
``<catalog>.<schema>``, shown in that prompt. Returns True if the skill was
205-
written, False if it was skipped or kept.
204+
205+
def should_download_skill(roots: list[Path], leaf: str, *, location: str) -> bool:
206+
"""Whether ``leaf`` should be fetched and written into ``roots``.
207+
208+
Applies the disk-only checks that need no bundle bytes, so a declined or
209+
invalid skill is never downloaded: skips invalid leaf names, and prompts
210+
before overwriting a skill already on disk (``location`` is the source
211+
``<catalog>.<schema>`` shown in that prompt).
206212
"""
207213
if not _is_valid_leaf(leaf):
208214
print_warning(f"Skipping `{leaf}`: not a valid skill name (lowercase a-z, 0-9, -).")
209215
return False
210216

211-
already_on_disk = any((root / leaf).exists() for root in roots)
212-
if already_on_disk and not prompt_yes_no(
217+
if existing_skill_on_disk(roots, leaf) and not prompt_yes_no(
213218
f"A skill named `{leaf}` already exists. Overwrite it with `{location}.{leaf}`?"
214219
):
215220
print_note(f"Kept existing `{leaf}`.")
216221
return False
217222

223+
return True
224+
225+
226+
def write_skill(roots: list[Path], leaf: str, files: dict[str, bytes]) -> None:
227+
"""Write ``leaf``'s bundle (``{relpath: bytes}``) into every root."""
218228
for root in roots:
219229
_write_bundle(root / leaf, leaf, files)
220-
return True
221230

222231

223232
# --- Orchestration ---------------------------------------------------------
@@ -250,9 +259,12 @@ def _fetch_bundles(
250259
def download_skills(workspace: str, token: str, locations: list[str], path: str | None) -> None:
251260
"""Download every skill in each ``<catalog>.<schema>`` location to disk.
252261
253-
Bundles are fetched concurrently (with a progress bar) per schema, then
254-
written sequentially so overwrite prompts don't interleave. A failure on one
255-
skill warns and skips it without aborting the batch.
262+
Locations are processed sequentially. Within each, overwrite prompts run
263+
first so declined skills are never fetched, then the survivors' bundles are
264+
fetched concurrently (with a progress bar) and written. Processing one
265+
location fully before the next lets a same-named skill from a later location
266+
see the earlier one on disk and prompt. A failure on one skill warns and
267+
skips it without aborting the batch.
256268
"""
257269
roots = skill_dir_roots(path)
258270
roots_display = " and ".join(str(root) for root in roots)
@@ -266,15 +278,18 @@ def download_skills(workspace: str, token: str, locations: list[str], path: str
266278
print_note(f"No skills found in `{location}`.")
267279
continue
268280

269-
bundles = _fetch_bundles(workspace, token, catalog, schema, leaves)
281+
to_download = [
282+
leaf for leaf in leaves if should_download_skill(roots, leaf, location=location)
283+
]
284+
bundles = _fetch_bundles(workspace, token, catalog, schema, to_download)
270285
written = 0
271-
for leaf in leaves:
286+
for leaf in to_download:
272287
files, reason = bundles[leaf]
273288
if reason or files is None:
274289
print_warning(f"Skipping `{location}.{leaf}`: {reason}.")
275290
continue
276-
if write_skill(roots, leaf, files, location=location):
277-
written += 1
291+
write_skill(roots, leaf, files)
292+
written += 1
278293
console.print()
279294
print_success(
280295
f"Downloaded {written}/{len(leaves)} skill(s) from `{location}` in {roots_display}."

tests/test_skills_download.py

Lines changed: 50 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,12 @@
66
import pytest
77

88
import ucode.skills_download as sd
9-
from ucode.skills_download import skill_dir_roots, write_skill
9+
from ucode.skills_download import (
10+
existing_skill_on_disk,
11+
should_download_skill,
12+
skill_dir_roots,
13+
write_skill,
14+
)
1015

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

@@ -237,50 +242,56 @@ def test_missing_directory_rejected(self, tmp_path):
237242
skill_dir_roots(str(tmp_path / "nope"))
238243

239244

240-
def _write(roots, leaf, files, *, location="main.default"):
241-
return write_skill(roots, leaf, files, location=location)
242-
243-
244-
class TestWriteSkill:
245-
def test_writes_bundle_into_every_root(self, tmp_path):
245+
class TestShouldDownloadSkill:
246+
def test_new_skill_is_downloaded(self, tmp_path):
246247
roots = skill_dir_roots(str(tmp_path))
247-
files = {"SKILL.md": b"# skill", "scripts/run.py": b"print(1)"}
248-
249-
_write(roots, "triage", files)
250248

251-
for root in roots:
252-
assert (root / "triage/SKILL.md").read_bytes() == b"# skill"
253-
assert (root / "triage/scripts/run.py").read_bytes() == b"print(1)"
249+
assert should_download_skill(roots, "triage", location="main.default")
254250

255251
def test_existing_skill_prompt_keep(self, tmp_path, monkeypatch):
256252
roots = skill_dir_roots(str(tmp_path))
257-
_write(roots, "triage", {"SKILL.md": b"from-main"}, location="main.default")
253+
write_skill(roots, "triage", {"SKILL.md": b"from-main"})
258254

259255
monkeypatch.setattr(sd, "prompt_yes_no", lambda _: False)
260-
_write(roots, "triage", {"SKILL.md": b"from-ml"}, location="ml.prod")
261256

262-
assert (roots[0] / "triage/SKILL.md").read_bytes() == b"from-main"
257+
assert not should_download_skill(roots, "triage", location="ml.prod")
263258

264259
def test_existing_skill_prompt_overwrite(self, tmp_path, monkeypatch):
265260
roots = skill_dir_roots(str(tmp_path))
266-
_write(roots, "triage", {"SKILL.md": b"from-main"}, location="main.default")
261+
write_skill(roots, "triage", {"SKILL.md": b"from-main"})
267262

268263
monkeypatch.setattr(sd, "prompt_yes_no", lambda _: True)
269-
_write(roots, "triage", {"SKILL.md": b"from-ml"}, location="ml.prod")
270264

271-
assert (roots[0] / "triage/SKILL.md").read_bytes() == b"from-ml"
265+
assert should_download_skill(roots, "triage", location="ml.prod")
272266

273267
def test_invalid_leaf_is_skipped(self, tmp_path):
274268
roots = skill_dir_roots(str(tmp_path))
275269

276-
_write(roots, "Bad_Name", {"SKILL.md": b"x"})
270+
assert not should_download_skill(roots, "Bad_Name", location="main.default")
271+
272+
def test_existing_skill_on_disk_checks_every_root(self, tmp_path):
273+
roots = skill_dir_roots(str(tmp_path))
274+
assert not existing_skill_on_disk(roots, "triage")
275+
276+
(roots[1] / "triage").mkdir(parents=True)
277+
assert existing_skill_on_disk(roots, "triage")
278+
279+
280+
class TestWriteSkill:
281+
def test_writes_bundle_into_every_root(self, tmp_path):
282+
roots = skill_dir_roots(str(tmp_path))
283+
files = {"SKILL.md": b"# skill", "scripts/run.py": b"print(1)"}
284+
285+
write_skill(roots, "triage", files)
277286

278-
assert not (roots[0] / "Bad_Name").exists()
287+
for root in roots:
288+
assert (root / "triage/SKILL.md").read_bytes() == b"# skill"
289+
assert (root / "triage/scripts/run.py").read_bytes() == b"print(1)"
279290

280291
def test_path_traversal_is_rejected(self, tmp_path):
281292
roots = skill_dir_roots(str(tmp_path))
282293

283-
_write(roots, "triage", {"SKILL.md": b"ok", "../escape.md": b"nope", "/abs.md": b"nope"})
294+
write_skill(roots, "triage", {"SKILL.md": b"ok", "../escape.md": b"nope", "/abs.md": b"nope"})
284295

285296
assert (roots[0] / "triage/SKILL.md").read_bytes() == b"ok"
286297
assert not (tmp_path / "escape.md").exists()
@@ -322,6 +333,23 @@ def test_list_failure_skips_location(self, tmp_path, monkeypatch):
322333

323334
assert called == []
324335

336+
def test_declined_skill_is_not_fetched(self, tmp_path, monkeypatch):
337+
roots = skill_dir_roots(str(tmp_path))
338+
write_skill(roots, "triage", {"SKILL.md": b"kept"})
339+
monkeypatch.setattr(sd, "list_schema_skills", lambda *a, **k: (["triage"], None))
340+
monkeypatch.setattr(sd, "prompt_yes_no", lambda _: False)
341+
fetched = []
342+
monkeypatch.setattr(
343+
sd,
344+
"fetch_skill_bundle",
345+
lambda ws, tok, c, s, leaf: fetched.append(leaf) or ({"SKILL.md": b"new"}, None),
346+
)
347+
348+
sd.download_skills(WS, "token", ["main.default"], str(tmp_path))
349+
350+
assert fetched == []
351+
assert (roots[0] / "triage/SKILL.md").read_bytes() == b"kept"
352+
325353
def test_bundle_failure_skips_that_skill_only(self, tmp_path, monkeypatch):
326354
monkeypatch.setattr(sd, "list_schema_skills", lambda *a, **k: (["good", "bad"], None))
327355
monkeypatch.setattr(

0 commit comments

Comments
 (0)