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
7 changes: 5 additions & 2 deletions src/ucode/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,7 @@
configure_selected_skills_download_command,
configure_skills_download_picker_command,
reconcile_managed_skills,
refresh_downloaded_skills_on_launch,
remove_downloaded_skills_command,
)
from ucode.skills_list import configured_skill_counts_by_agent, list_configured_skills_command
Expand Down Expand Up @@ -2791,6 +2792,8 @@ def _launch_tool(
# Claude re-adds an out-of-catalog saved model to /model even when built-ins are
# replaced. Keep the managed catalog launch-scoped and leave the user's settings alone.
state["_claude_launch_picker_models"] = picker_catalog.model_ids
if not skip_preflight:
refresh_downloaded_skills_on_launch(state)
# Relayed = a Claude subscription: forward the model to Claude Code's own flag, like `-- --model X`.
should_forward_relayed_model = (
tool == "claude"
Expand All @@ -2817,8 +2820,8 @@ def _launch_tool(
)
if recommendation is not None:
_print_budget_panel(recommendation, tool, managed)
# The managed config's MCP servers and skills are both applied at `ug configure`, not here,
# so the launch hot path makes no per-launch discovery calls for them.
# The managed config's MCP servers and skills are applied at `ug configure`, not here.
# Downloaded skills get a rate-limited refresh above (refresh_downloaded_skills_on_launch).
if tool == "claude":
if provider:
state["_claude_launch_provider"] = provider
Expand Down
21 changes: 4 additions & 17 deletions src/ucode/managed_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
fetch_model_recommendation,
get_databricks_token,
)
from ucode.string_utils import parse_update_time
from ucode.ui import console, print_warning

MANAGED_CONFIG_PATH = config_io.APP_DIR / "managed-config.json"
Expand Down Expand Up @@ -447,29 +448,15 @@ def managed_update_time(managed: dict | None) -> str | None:
return _str(_as_dict(managed).get("update_time"))


def _parse_update_time(value: str | None) -> datetime | None:
if not value:
return None
try:
dt = datetime.fromisoformat(value.replace("Z", "+00:00"))
except ValueError:
return None
# An offset-less timestamp (e.g. a stub value) parses tz-naive; pin it to UTC so it can be
# compared against the tz-aware persisted watermark without raising.
if dt.tzinfo is None:
dt = dt.replace(tzinfo=UTC)
return dt


def managed_config_is_newer(fetched: dict | None, applied_update_time: str | None) -> bool:
"""True when ``fetched`` is a newer version than the last one applied locally.

A fetched config whose ``update_time`` is missing or unparseable is treated as newer, so a launch
re-applies it rather than trusting possibly-stale local settings; no previously-applied watermark
also counts as newer (the first apply).
"""
fetched_ut = _parse_update_time(managed_update_time(fetched))
applied_ut = _parse_update_time(applied_update_time)
fetched_ut = parse_update_time(managed_update_time(fetched))
applied_ut = parse_update_time(applied_update_time)
if fetched_ut is None or applied_ut is None:
return True
return fetched_ut > applied_ut
Expand Down Expand Up @@ -731,7 +718,7 @@ def _cached_result_if_fresh(workspace: str) -> ManagedConfigResult | None:
if data.get("workspace") != workspace:
return None
# Reuses the RFC-3339 parser the update-time watermark uses; None (missing/unparseable) is stale.
retrieved_at = _parse_update_time(_str(data.get("retrieved_at")))
retrieved_at = parse_update_time(_str(data.get("retrieved_at")))
if retrieved_at is None:
return None
age = _utcnow() - retrieved_at
Expand Down
113 changes: 113 additions & 0 deletions src/ucode/skills_download.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,11 @@

from __future__ import annotations

import os
import shutil
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import UTC, datetime
from pathlib import Path

import questionary
Expand All @@ -21,15 +23,20 @@
list_schema_skills,
)
from ucode.skills_state import (
SKILL_UPDATE_CHECK_INTERVAL,
SkillInstall,
forget,
last_update_check,
list_downloaded,
record_downloads,
records_for_fqns,
records_for_schema,
records_for_scope,
remove_downloads,
set_last_update_check,
)
from ucode.state import load_state
from ucode.string_utils import parse_update_time
from ucode.ui import (
console,
picker_style,
Expand Down Expand Up @@ -423,6 +430,106 @@ def reconcile_managed_skills(managed: dict) -> tuple[list[str], list[str]]:
return [ref.bundle_name for ref in installed], removed


# --- Launch-time refresh ---------------------------------------------------


def _eligible_launch_refresh_records(records: list[dict], workspace: str) -> list[dict]:
"""The current workspace's own (non-managed) downloads, which a launch may refresh.

Managed skills are left to ``ug configure``, and other workspaces' downloads are skipped
because the launch token authenticates only this workspace.
"""
return [
record
for record in records
if record.get("scope") != "managed" and record.get("workspace") == workspace
]


def _get_updated_refs(
workspace: str, token: str, records: list[dict]
) -> list[tuple[dict, SkillRef]]:
"""Pair each record to re-download with its current skill: one whose UC source is newer than
its download, or one whose on-disk copy is only partly present and needs restoring to mirror UC.

Resolves every record concurrently; one that no longer resolves (deleted, unfinalized,
unauthorized) is skipped, leaving its on-disk copy alone. Times are parsed before comparing
so the two RFC-3339 forms UC emits sort chronologically. A record with no parseable recorded
``uc_update_time`` predates attribution, so it is refreshed once to backfill the field.
"""
if not records:
return []
pairs: list[tuple[dict, SkillRef]] = []
with ThreadPoolExecutor(max_workers=min(_MAX_FETCH_WORKERS, len(records))) as pool:
futures = {pool.submit(get_skill, workspace, token, r["fqn"]): r for r in records}
for future in as_completed(futures):
ref = future.result()
if ref is None:
continue
record = futures[future]
stored = parse_update_time(record.get("uc_update_time"))
current = parse_update_time(ref.uc_update_time)
is_newer = stored is None or (current is not None and current > stored)
if is_newer or _record_dirs_missing(record):
pairs.append((record, ref))
return pairs


def _update_stale_skills(workspace: str, token: str, pairs: list[tuple[dict, SkillRef]]) -> int:
"""Re-download each stale skill into its own base and refresh its manifest record.

Overwrites in place with no prompt, since the developer already chose to download these,
and only manifest-attributed directories are touched, so a user-authored skill of the same
name is never overwritten. Returns how many skills were rewritten.
"""
home = os.path.normpath(str(Path.home()))
refs_by_base: dict[str, list[SkillRef]] = {}
for record, ref in pairs:
refs_by_base.setdefault(os.path.normpath(record["base"]), []).append(ref)

updated = 0
for base, refs in refs_by_base.items():
path = None if base == home else base
roots = skill_dir_roots(path)
written = _fetch_bundles_and_write(workspace, token, refs, roots, label="Updating skills")
record_downloads(_skill_installs(written, roots, path, workspace))
updated += len(written)
return updated


def refresh_downloaded_skills_on_launch(state: dict) -> None:
"""Update downloaded skills whose UC source changed, before an agent launches.

Rate-limited to once per ``SKILL_UPDATE_CHECK_INTERVAL`` via the manifest's
``last_update_check`` stamp, so back-to-back launches make no network calls. A record whose
directories the user deleted entirely is forgotten; one only partly deleted is re-downloaded
to restore the mirror. Best-effort: any failure is reported and the launch proceeds on
whatever is already on disk.
"""
try:
now = datetime.now(UTC)
last = last_update_check()
if last is not None and now - last < SKILL_UPDATE_CHECK_INTERVAL:
return
workspace = state.get("workspace")
if not workspace:
return
set_last_update_check(now)
deleted, present = [], []
for record in _eligible_launch_refresh_records(list_downloaded(), workspace):
(deleted if _record_dirs_all_missing(record) else present).append(record)
forget(deleted)
if present:
print_note("Checking Unity Catalog for downloaded skill updates...")
token = get_databricks_token(workspace, state.get("profile"))
pairs = _get_updated_refs(workspace, token, present)
updated = _update_stale_skills(workspace, token, pairs)
if updated:
print_success(f"Updated {updated} downloaded skill(s) from Unity Catalog.")
except Exception as exc: # noqa: BLE001 - a skill refresh must never block a launch
print_note(f"Skipped checking for skill updates: {exc}")


def configure_location_skills_download_command(locations: list[str], *, path: str | None) -> int:
"""Download every skill in each schema to disk and register the skills connection.

Expand Down Expand Up @@ -537,6 +644,12 @@ def _record_dirs_missing(record: dict) -> bool:
return any(not Path(directory).exists() for directory in record.get("dirs") or [])


def _record_dirs_all_missing(record: dict) -> bool:
"""Whether every one of a record's on-disk directories no longer exists."""
dirs = record.get("dirs") or []
return bool(dirs) and all(not Path(directory).exists() for directory in dirs)


def _download_label(record: dict) -> str:
label = f"{record.get('fqn')} ({record.get('scope')}: {record.get('base')})"
return f"{label} (missing)" if _record_dirs_missing(record) else label
Expand Down
44 changes: 35 additions & 9 deletions src/ucode/skills_state.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,18 @@
import shutil
import time
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from pathlib import Path

from ucode import config_io
from ucode.string_utils import parse_update_time
from ucode.ui import print_warning

SKILLS_STATE_VERSION = 1

# Matches Isaac's plugin marketplace staleness window (see plugin-marketplace/CLAUDE.md).
SKILL_UPDATE_CHECK_INTERVAL = timedelta(hours=24)


@dataclass(frozen=True)
class SkillInstall:
Expand Down Expand Up @@ -57,8 +62,8 @@ def _quarantine_corrupt(path: Path) -> None:
pass


def _load() -> list[dict]:
"""Records in the manifest; ``[]`` if it is absent, unreadable, or an unrecognized version.
def _load_manifest() -> dict:
"""The whole manifest; ``{}`` if it is absent, unreadable, or an unrecognized version.

A file that fails to parse is quarantined (see ``_quarantine_corrupt``) rather than read as
empty, so a single bad byte doesn't let the next write silently erase every tracked skill.
Expand All @@ -67,22 +72,43 @@ def _load() -> list[dict]:
try:
text = path.read_text(encoding="utf-8")
except OSError:
return []
return {}
try:
data = json.loads(text)
except json.JSONDecodeError:
_quarantine_corrupt(path)
return []
return {}
if not isinstance(data, dict) or data.get("version") != SKILLS_STATE_VERSION:
return []
downloads = data.get("skill_downloads")
return {}
return data


def _load() -> list[dict]:
"""The download records in the manifest, or ``[]`` when there are none."""
downloads = _load_manifest().get("skill_downloads")
return [r for r in downloads if isinstance(r, dict)] if isinstance(downloads, list) else []


def _save(downloads: list[dict]) -> None:
config_io.atomic_write_json(
_skills_state_path(), {"version": SKILLS_STATE_VERSION, "skill_downloads": downloads}
)
"""Write the download records, preserving other manifest keys (e.g. ``last_update_check``)."""
manifest = _load_manifest()
manifest["version"] = SKILLS_STATE_VERSION
manifest["skill_downloads"] = downloads
config_io.atomic_write_json(_skills_state_path(), manifest)


def last_update_check() -> datetime | None:
"""When the launch-time update sweep last ran, or None if it never has."""
raw = _load_manifest().get("last_update_check")
return parse_update_time(raw) if isinstance(raw, str) else None


def set_last_update_check(when: datetime) -> None:
"""Record when the launch-time update sweep last ran."""
manifest = _load_manifest()
manifest["version"] = SKILLS_STATE_VERSION
manifest["last_update_check"] = when.astimezone(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
config_io.atomic_write_json(_skills_state_path(), manifest)


def _norm(path: str) -> str:
Expand Down
20 changes: 19 additions & 1 deletion src/ucode/string_utils.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,25 @@
"""Shared string validation helpers."""
"""Shared string helpers."""

from __future__ import annotations

from datetime import UTC, datetime


def parse_update_time(value: str | None) -> datetime | None:
"""Parse an RFC-3339 ``update_time`` into an aware UTC datetime, or None if absent/unparseable.

UC serializes fractional seconds only when non-zero (protobuf JSON), so ``...:25Z`` and
``...:25.400Z`` both occur; parsing before comparing avoids the wrong lexicographic ordering of
those two forms. An offset-less value is pinned to UTC.
"""
if not value:
return None
try:
parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
except ValueError:
return None
return parsed if parsed.tzinfo is not None else parsed.replace(tzinfo=UTC)


def is_valid_catalog_schema(value: str) -> bool:
"""Return whether value is a safe ``<catalog>.<schema>`` reference."""
Expand Down
Loading
Loading