diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 16367e9..602ae1a 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -90,7 +90,10 @@ purge_cross_workspace_mcp_residue, revert_mcp_configs, ) -from ucode.skills_download import configure_skills_download_command +from ucode.skills_download import ( + configure_skills_download_command, + download_managed_skills_on_launch, +) from ucode.smart_routing import claude_routing, codex_routing from ucode.state import ( STATE_PATH, @@ -1456,13 +1459,47 @@ def _register_managed_mcp_servers(managed: dict, tool: str, state: dict) -> None print_note(f"Registered workspace MCP server(s) for {TOOL_SPECS[tool]['display']}: {names}") +def _managed_skill_locations(managed: dict) -> list[str]: + """The ``.`` skill locations the admin published, or ``[]``.""" + return [ + loc + for loc in ((managed.get("skills") or {}).get("names") or []) + if isinstance(loc, str) and loc + ] + + +def _download_managed_skills(managed: dict, state: dict) -> None: + """Download the admin-published skill schemas to disk (user scope). + + Registering the skills MCP connection (see :func:`_apply_managed_skills`) exposes the skill + *tools* over the gateway, but the agent's ``/skills`` picker reads skill bundles from + ``~/.claude/skills`` / ``~/.agents/skills`` on disk. Without this download those directories stay + empty, so a workspace-published skill never shows up in ``/skills``. Skills already on disk are + left untouched, so a steady-state launch only lists each schema and writes nothing. Best-effort: + a failure here never blocks the launch. + """ + locations = _managed_skill_locations(managed) + if not locations: + return + try: + token = get_databricks_token(state["workspace"], state.get("profile")) + written = download_managed_skills_on_launch(state["workspace"], token, locations) + except RuntimeError as exc: + print_warning(f"Could not download your workspace's skills: {exc}") + return + if written: + print_note(f"Downloaded workspace skill(s) to disk: {', '.join(written)}") + + def _apply_managed_skills(managed: dict, tool: str, state: dict) -> None: - """Register the managed config's skill schemas on ``tool``'s skills MCP connection. + """Register the managed config's skill schemas on ``tool``'s skills MCP connection and disk. Sibling of :func:`_register_managed_mcp_servers` for the skills registry: the managed config lists the skill schemas the admin published, and nothing else on the launch path routes them to the agent. ``apply_managed_skills`` persists the connection (and the applied set, for diffing a - later removal) into ``state`` itself. A failure here never blocks the launch. + later removal) into ``state`` itself, then ``_download_managed_skills`` writes the skill bundles + to disk so the agent's ``/skills`` picker lists them. A failure in either step never blocks the + launch. """ try: applied = apply_managed_skills( @@ -1475,12 +1512,13 @@ def _apply_managed_skills(managed: dict, tool: str, state: dict) -> None: ) except RuntimeError as exc: print_warning(f"Could not register your workspace's skills: {exc}") - return - if applied: - names = ", ".join(applied) - print_note( - f"Registered workspace skill schema(s) for {TOOL_SPECS[tool]['display']}: {names}" - ) + else: + if applied: + names = ", ".join(applied) + print_note( + f"Registered workspace skill schema(s) for {TOOL_SPECS[tool]['display']}: {names}" + ) + _download_managed_skills(managed, state) def _launch_tool( diff --git a/src/ucode/skills_download.py b/src/ucode/skills_download.py index e6aa9d9..2168b52 100644 --- a/src/ucode/skills_download.py +++ b/src/ucode/skills_download.py @@ -396,6 +396,43 @@ def download_skills( ) +def download_managed_skills_on_launch( + workspace: str, token: str, locations: list[str], path: str | None = None +) -> list[str]: + """Download admin-published skills to disk so the agent's ``/skills`` lists them. + + Runs on the managed launch path: the config only registers the skills MCP + connection, so nothing else writes the bundles that ``/skills`` reads. Writes + only skills not already on disk -- no overwrite prompt, so the launch never + blocks on input and a developer's own same-named skill is never clobbered. + Best-effort and never raises, so it can't block the launch. Returns the bundle + names newly written. + """ + roots = skill_dir_roots(path) + written: list[str] = [] + for location in locations: + if location.count(".") != 1: + continue + catalog, schema = location.split(".") + refs, reason = list_schema_skills(workspace, token, catalog, schema) + if reason: + print_warning(f"Could not list workspace skills in `{location}`: {reason}.") + continue + refs = _reject_bundle_name_collisions(refs, location=location) + missing = [ref for ref in refs if not existing_skill_on_disk(roots, ref.bundle_name)] + if not missing: + continue + bundles = _fetch_bundles(workspace, token, catalog, schema, missing) + for ref in missing: + files, reason = bundles[ref.securable_name] + if reason or files is None: + print_warning(f"Skipping `{location}.{ref.securable_name}`: {reason}.") + continue + write_skill(roots, ref, files) + written.append(ref.bundle_name) + return written + + def configure_skills_download_command( locations: list[str], *, path: str | None, skills: set[str] | None = None ) -> int: diff --git a/tests/test_cli.py b/tests/test_cli.py index 7125528..7daa4bf 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -681,6 +681,76 @@ def test_path_without_location_exit_1(self): mock_download.assert_not_called() +class TestApplyManagedSkills: + """The launch path both registers the skills MCP connection and downloads bundles to disk.""" + + def _state(self): + return {"workspace": "https://example.databricks.com", "profile": "prod"} + + def test_downloads_managed_skill_schemas_to_disk(self): + managed = {"skills": {"names": ["main.default", "ml.prod"]}} + with ( + patch("ucode.cli.apply_managed_skills", return_value=["main.default"]) as mock_apply, + patch("ucode.cli.get_databricks_token", return_value="tok") as mock_token, + patch( + "ucode.cli.download_managed_skills_on_launch", return_value=["triage"] + ) as mock_dl, + ): + from ucode import cli + + cli._apply_managed_skills(managed, "claude", self._state()) + + mock_apply.assert_called_once() + mock_token.assert_called_once_with("https://example.databricks.com", "prod") + mock_dl.assert_called_once_with( + "https://example.databricks.com", "tok", ["main.default", "ml.prod"] + ) + + def test_no_managed_skills_skips_the_download(self): + with ( + patch("ucode.cli.apply_managed_skills", return_value=[]), + patch("ucode.cli.get_databricks_token") as mock_token, + patch("ucode.cli.download_managed_skills_on_launch") as mock_dl, + ): + from ucode import cli + + cli._apply_managed_skills({}, "claude", self._state()) + + mock_token.assert_not_called() + mock_dl.assert_not_called() + + def test_download_still_runs_when_mcp_registration_fails(self): + # A failure registering the MCP connection must not stop the disk download — the two are + # independent ways skills reach the agent, and /skills depends only on the disk write. + with ( + patch("ucode.cli.apply_managed_skills", side_effect=RuntimeError("boom")), + patch("ucode.cli.get_databricks_token", return_value="tok"), + patch("ucode.cli.download_managed_skills_on_launch", return_value=[]) as mock_dl, + ): + from ucode import cli + + cli._apply_managed_skills( + {"skills": {"names": ["main.default"]}}, "claude", self._state() + ) + + mock_dl.assert_called_once() + + def test_download_failure_never_blocks_launch(self): + with ( + patch("ucode.cli.apply_managed_skills", return_value=[]), + patch("ucode.cli.get_databricks_token", side_effect=RuntimeError("no auth")), + patch("ucode.cli.download_managed_skills_on_launch") as mock_dl, + ): + from ucode import cli + + # Must not raise. + cli._apply_managed_skills( + {"skills": {"names": ["main.default"]}}, "claude", self._state() + ) + + mock_dl.assert_not_called() + + class TestStatusSkillsSection: def _run(self, state): with patch("ucode.cli.load_state", return_value=state): diff --git a/tests/test_skills_download.py b/tests/test_skills_download.py index a7e7924..c792701 100644 --- a/tests/test_skills_download.py +++ b/tests/test_skills_download.py @@ -610,6 +610,95 @@ def test_none_skill_filter_downloads_everything(self, tmp_path, monkeypatch): assert (tmp_path / ".claude/skills/b/SKILL.md").exists() +class TestDownloadManagedSkillsOnLaunch: + def test_writes_missing_skills_and_returns_their_bundle_names(self, tmp_path, monkeypatch): + monkeypatch.setattr( + sd, "list_schema_skills", lambda *a, **k: ([ref("triage"), ref("pii")], None) + ) + monkeypatch.setattr( + sd, + "fetch_skill_bundle", + lambda ws, tok, c, s, leaf: ({"SKILL.md": leaf.encode()}, None), + ) + + written = sd.download_managed_skills_on_launch(WS, "token", ["main.default"], str(tmp_path)) + + assert sorted(written) == ["pii", "triage"] + assert (tmp_path / ".claude/skills/triage/SKILL.md").read_bytes() == b"triage" + assert (tmp_path / ".agents/skills/pii/SKILL.md").read_bytes() == b"pii" + + def test_skips_already_downloaded_skills_without_prompting(self, tmp_path, monkeypatch): + roots = skill_dir_roots(str(tmp_path)) + write_skill(roots, ref("triage"), {"SKILL.md": b"kept"}) + monkeypatch.setattr( + sd, "list_schema_skills", lambda *a, **k: ([ref("triage"), ref("pii")], None) + ) + fetched = [] + monkeypatch.setattr( + sd, + "fetch_skill_bundle", + lambda ws, tok, c, s, leaf: fetched.append(leaf) or ({"SKILL.md": b"new"}, None), + ) + monkeypatch.setattr(sd, "prompt_yes_no", lambda msg: pytest.fail(f"prompted: {msg}")) + + written = sd.download_managed_skills_on_launch(WS, "token", ["main.default"], str(tmp_path)) + + # Only the missing one is fetched; the existing skill is left untouched. + assert fetched == ["pii"] + assert written == ["pii"] + assert (roots[0] / "triage/SKILL.md").read_bytes() == b"kept" + + def test_nothing_missing_fetches_nothing(self, tmp_path, monkeypatch): + roots = skill_dir_roots(str(tmp_path)) + write_skill(roots, ref("triage"), {"SKILL.md": b"kept"}) + monkeypatch.setattr(sd, "list_schema_skills", lambda *a, **k: ([ref("triage")], None)) + monkeypatch.setattr( + sd, "fetch_skill_bundle", lambda *a, **k: pytest.fail("should not fetch") + ) + + assert ( + sd.download_managed_skills_on_launch(WS, "token", ["main.default"], str(tmp_path)) == [] + ) + + def test_list_failure_warns_and_skips_location(self, tmp_path, monkeypatch, capsys): + monkeypatch.setattr(sd, "list_schema_skills", lambda *a, **k: ([], "HTTP 404 Not Found")) + monkeypatch.setattr( + sd, "fetch_skill_bundle", lambda *a, **k: pytest.fail("should not fetch") + ) + + assert ( + sd.download_managed_skills_on_launch(WS, "token", ["main.default"], str(tmp_path)) == [] + ) + assert "Could not list workspace skills in `main.default`" in capsys.readouterr().out + + def test_bundle_failure_skips_that_skill_only(self, tmp_path, monkeypatch): + monkeypatch.setattr( + sd, "list_schema_skills", lambda *a, **k: ([ref("good"), ref("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") + ), + ) + + written = sd.download_managed_skills_on_launch(WS, "token", ["main.default"], str(tmp_path)) + + assert written == ["good"] + assert (tmp_path / ".claude/skills/good/SKILL.md").read_bytes() == b"ok" + assert not (tmp_path / ".claude/skills/bad").exists() + + def test_malformed_location_is_skipped(self, tmp_path, monkeypatch): + monkeypatch.setattr( + sd, "list_schema_skills", lambda *a, **k: pytest.fail("should not list a bad location") + ) + + assert ( + sd.download_managed_skills_on_launch(WS, "token", ["not-a-schema"], str(tmp_path)) == [] + ) + + class TestConfigureSkillsDownloadCommand: def _stub(self, monkeypatch): calls: dict[str, object] = {}