Studio: stop the llama.cpp update banner flickering and show the download size

studio/backend/routes/llama.py

@@ -55,6 +55,9 @@ class LlamaUpdateStatusResponse(BaseModel):
     source_build: bool = Field(
         False, description = "True when there is no marker (source build) but a prebuilt is offered."
     )
+    update_size_bytes: Optional[int] = Field(
+        None, description = "Download size of the prebuilt Update would fetch, in bytes."
+    )
     job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob)
 
 

studio/backend/tests/test_llama_cpp_freshness.py

@@ -520,3 +520,114 @@ def test_in_memory_only_reset_replays_stale_same_base_mix(monkeypatch, tmp_path)
     assert info["latest_tag"] == "b9596-mix-aaa"
     assert info["behind"] is True
     assert info["stale"] is True
+
+
+# update_download_size_bytes (banner download-size lookup).
+
+
+def _patch_assets(monkeypatch, mapping):
+    """Stub latest_release_assets with a per-repo {asset_name: size} lookup."""
+    monkeypatch.setattr(
+        fr,
+        "latest_release_assets",
+        lambda repo, *, force_refresh = False: mapping.get(repo),
+    )
+
+
+def test_update_size_unsloth_prebuilt_exact_match(monkeypatch):
+    # The unsloth fork's own bundle (app-<tag>-<platform>): the want= exact match
+    # on app-<latest>-<suffix> wins.
+    marker = {
+        "asset": "app-b9190-linux-x64-cuda13-newer.tar.gz",
+        "published_repo": "unslothai/llama.cpp",
+    }
+    _patch_assets(
+        monkeypatch,
+        {
+            "unslothai/llama.cpp": {
+                "app-b9300-linux-x64-cuda13-newer.tar.gz": 123_456_789,
+                "app-b9300-windows-x64-cuda13-newer.zip": 999,
+            }
+        },
+    )
+    assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 123_456_789
+
+
+def test_update_size_macos_fork_asset_suffix_fallback(monkeypatch):
+    # macOS bundles use the upstream-style llama-<tag>-bin-macos-*, matched via the
+    # endswith fallback in the publish repo.
+    marker = {
+        "asset": "llama-b9190-bin-macos-arm64.tar.gz",
+        "published_repo": "unslothai/llama.cpp",
+    }
+    _patch_assets(
+        monkeypatch,
+        {"unslothai/llama.cpp": {"llama-b9300-bin-macos-arm64.tar.gz": 55_000_000}},
+    )
+    assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 55_000_000
+
+
+def test_update_size_upstream_ubuntu_uses_binary_repo(monkeypatch):
+    # #6338 P2: ggml-org ubuntu-* prebuilt lives in binary_repo, not the fork
+    # publish repo. The size must still resolve.
+    marker = {
+        "asset": "llama-b9190-bin-ubuntu-x64.tar.gz",
+        "published_repo": "unslothai/llama.cpp",
+        "binary_repo": "ggml-org/llama.cpp",
+    }
+    _patch_assets(
+        monkeypatch,
+        {
+            "unslothai/llama.cpp": {"app-b9300-linux-x64-cuda13-newer.tar.gz": 1},
+            "ggml-org/llama.cpp": {
+                "llama-b9673-bin-ubuntu-x64.tar.gz": 42_000_000,
+                "llama-b9673-bin-ubuntu-vulkan-x64.tar.gz": 7,
+            },
+        },
+    )
+    assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 42_000_000
+
+
+def test_update_size_upstream_windows_uses_binary_repo(monkeypatch):
+    # Regression (#6338 P2): the Windows upstream CPU prebuilt uses a win-* token.
+    marker = {
+        "asset": "llama-b9190-bin-win-cpu-x64.zip",
+        "published_repo": "unslothai/llama.cpp",
+        "binary_repo": "ggml-org/llama.cpp",
+    }
+    _patch_assets(
+        monkeypatch,
+        {"ggml-org/llama.cpp": {"llama-b9673-bin-win-cpu-x64.zip": 33_000_000}},
+    )
+    assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") == 33_000_000
+
+
+def test_update_size_no_matching_asset_fails_open(monkeypatch):
+    # A ROCm version drift (installed 6.4 vs latest 7.2) leaves no suffix match;
+    # the helper fails open to None rather than guessing a wrong artifact.
+    marker = {
+        "asset": "llama-b9190-bin-ubuntu-rocm-6.4-x64.tar.gz",
+        "published_repo": "unslothai/llama.cpp",
+        "binary_repo": "ggml-org/llama.cpp",
+    }
+    _patch_assets(
+        monkeypatch,
+        {"ggml-org/llama.cpp": {"llama-b9673-bin-ubuntu-rocm-7.2-x64.tar.gz": 9}},
+    )
+    assert fr.update_download_size_bytes(marker, "b9300", "unslothai/llama.cpp") is None
+
+
+def test_update_size_missing_inputs_fail_open(monkeypatch):
+    _patch_assets(
+        monkeypatch,
+        {"unslothai/llama.cpp": {"app-b9300-linux-x64-cpu.tar.gz": 5}},
+    )
+    # No marker, no latest tag, or no asset string -> None (never raise).
+    assert fr.update_download_size_bytes(None, "b9300", "unslothai/llama.cpp") is None
+    assert (
+        fr.update_download_size_bytes(
+            {"asset": "app-b9190-linux-x64-cpu.tar.gz"}, None, "unslothai/llama.cpp"
+        )
+        is None
+    )
+    assert fr.update_download_size_bytes({"asset": None}, "b9300", "unslothai/llama.cpp") is None

studio/backend/tests/test_llama_cpp_update.py

@@ -901,3 +901,49 @@ def test_start_update_marked_refuses_when_not_behind(monkeypatch, tmp_path):
     res = upd.start_update()
     assert res["started"] is False
     assert res["reason"] == "up_to_date"
+
+
+def test_status_update_available_includes_size(monkeypatch, tmp_path):
+    # Marker (prebuilt) update path attaches the download size of the asset the
+    # banner would fetch.
+    binary = _write_install(tmp_path, "b9493", asset = "app-b9493-linux-x64-cuda13-newer.tar.gz")
+    monkeypatch.setattr(upd, "_find_binary", lambda: binary)
+    monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
+    monkeypatch.setattr(
+        freshness,
+        "latest_release_assets",
+        lambda repo, *, force_refresh = False: {
+            "app-b9518-linux-x64-cuda13-newer.tar.gz": 88_000_000
+        },
+    )
+    st = upd.get_update_status(force_refresh = True)
+    assert st["update_available"] is True
+    assert st["update_size_bytes"] == 88_000_000
+
+
+def test_status_source_build_includes_update_size(monkeypatch, tmp_path):
+    # #6338 P3: a source build offered a prebuilt must carry the asset size too.
+    binary = tmp_path / "llama.cpp" / "build" / "bin" / "llama-server"
+    binary.parent.mkdir(parents = True)
+    binary.write_text("stub")  # no marker -> source build
+    monkeypatch.setattr(upd, "_find_binary", lambda: str(binary))
+    _prebuilt(
+        monkeypatch,
+        repo = "unslothai/llama.cpp",
+        release_tag = "b9585",
+        asset = "app-b9585-linux-x64-cpu.tar.gz",
+    )
+    monkeypatch.setattr(upd, "_installed_build_number", lambda b: None)
+    monkeypatch.setattr(
+        upd,
+        "latest_release_assets",
+        lambda repo, *, force_refresh = False: (
+            {"app-b9585-linux-x64-cpu.tar.gz": 77_000_000}
+            if repo == "unslothai/llama.cpp"
+            else None
+        ),
+    )
+    st = upd.get_update_status()
+    assert st["source_build"] is True
+    assert st["update_available"] is True
+    assert st["update_size_bytes"] == 77_000_000

studio/backend/tests/test_llama_route.py

@@ -78,6 +78,27 @@ def test_status_response_exposes_source_build():
     rl.LlamaUpdateStatusResponse(**{**payload, "unexpected": 1})
 
 
+def test_status_response_exposes_update_size_bytes():
+    payload = {
+        "supported": True,
+        "update_available": True,
+        "stale": False,
+        "installed_tag": "b9493",
+        "latest_tag": "b9518",
+        "published_repo": "unslothai/llama.cpp",
+        "installed_at_utc": None,
+        "age_days": None,
+        "source_build": False,
+        "update_size_bytes": 123_456_789,
+        "job": {"state": "idle"},
+    }
+    model = rl.LlamaUpdateStatusResponse(**payload)
+    assert model.model_dump()["update_size_bytes"] == 123_456_789
+    # Omitted -> defaults to None (the offline / no-matching-asset case).
+    without = {k: v for k, v in payload.items() if k != "update_size_bytes"}
+    assert rl.LlamaUpdateStatusResponse(**without).model_dump()["update_size_bytes"] is None
+
+
 def test_status_handler_runs_off_event_loop(monkeypatch):
     seen = {}
 

studio/backend/utils/llama_cpp_freshness.py

@@ -33,6 +33,8 @@
 
 _marker_cache: dict[str, Optional[dict]] = {}
 _release_memo: dict[str, tuple[float, Optional[str]]] = {}
+# Newest-release asset sizes (name -> bytes), memoized like the tag (24h TTL).
+_assets_memo: dict[str, tuple[float, dict[str, int]]] = {}
 
 
 def _cache_dir() -> Path:
@@ -180,6 +182,113 @@ def latest_published_release(repo: str, *, force_refresh: bool = False) -> Optio
     return latest
 
 
+def _fetch_latest_release_assets(repo: str, timeout: float = 5.0) -> Optional[dict[str, int]]:
+    """Asset name -> size (bytes) for the newest published release of `repo`,
+    selected exactly like _fetch_latest_release_tag. None on any failure."""
+    import urllib.error
+    import urllib.request
+
+    url = f"https://api.github.com/repos/{repo}/releases?per_page=30"
+    headers = {
+        "Accept": "application/vnd.github+json",
+        "User-Agent": "unsloth-studio-freshness-check",
+    }
+    token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN")
+    if token:
+        headers["Authorization"] = f"Bearer {token}"
+    req = urllib.request.Request(url, headers = headers)
+    try:
+        with urllib.request.urlopen(req, timeout = timeout) as resp:
+            data = json.loads(resp.read().decode("utf-8"))
+    except (
+        urllib.error.URLError,
+        urllib.error.HTTPError,
+        OSError,
+        json.JSONDecodeError,
+    ) as exc:
+        logger.debug("freshness asset fetch failed", repo = repo, error = str(exc))
+        return None
+    if not isinstance(data, list):
+        return None
+    published = [
+        r
+        for r in data
+        if isinstance(r, dict)
+        and not r.get("draft")
+        and not r.get("prerelease")
+        and isinstance(r.get("tag_name"), str)
+        and r.get("tag_name")
+    ]
+    if not published:
+        return None
+    newest = max(published, key = lambda r: r.get("published_at") or "")
+    assets: dict[str, int] = {}
+    for a in newest.get("assets") or []:
+        name, size = a.get("name"), a.get("size")
+        if isinstance(name, str) and isinstance(size, int):
+            assets[name] = size
+    return assets
+
+
+def latest_release_assets(repo: str, *, force_refresh: bool = False) -> Optional[dict[str, int]]:
+    """Newest-release asset sizes for `repo`, memoized (24h TTL). None when
+    offline and never fetched. In-memory only -- a restart simply re-fetches."""
+    if not repo:
+        return None
+    now = time.time()
+    if not force_refresh:
+        memo = _assets_memo.get(repo)
+        if memo and now - memo[0] < _RELEASE_CACHE_TTL_SECONDS:
+            return memo[1]
+    assets = _fetch_latest_release_assets(repo)
+    if assets is None:
+        memo = _assets_memo.get(repo)
+        return memo[1] if memo else None
+    _assets_memo[repo] = (now, assets)
+    return assets
+
+
+def update_download_size_bytes(
+    marker: Optional[dict],
+    latest_tag: Optional[str],
+    repo: Optional[str],
+    *,
+    force_refresh: bool = False,
+) -> Optional[int]:
+    """Download size of the latest-release asset matching this host's installed
+    bundle (same platform/arch/runtime suffix as the installed asset). None when
+    there is no marker asset, the latest assets can't be read, or no match."""
+    if not marker or not latest_tag or not repo:
+        return None
+    installed_asset = marker.get("asset")
+    if not isinstance(installed_asset, str):
+        return None
+    # Tag-independent platform suffix: accept the fork's "app-*" bundles and the
+    # upstream ggml-org "ubuntu-*"/"win-*" prebuilts ("windows" before "win").
+    m = re.search(r"-((?:linux|ubuntu|windows|win|macos|darwin)-.*)$", installed_asset)
+    if not m:
+        return None
+    suffix = m.group(1)
+    # Upstream ubuntu/win assets live in the marker's binary_repo, not the fork
+    # publish repo; try the publish repo first, then it.
+    repos = [repo]
+    binary_repo = marker.get("binary_repo")
+    if isinstance(binary_repo, str) and binary_repo and binary_repo != repo:
+        repos.append(binary_repo)
+    want = f"app-{latest_tag}-{suffix}"
+    for r in repos:
+        assets = latest_release_assets(r, force_refresh = force_refresh)
+        if not assets:
+            continue
+        if want in assets:
+            return assets[want]
+        # Tag formatting can vary (mix suffixes); fall back to the platform suffix.
+        for name, size in assets.items():
+            if name.endswith(suffix):
+                return size
+    return None
+
+
 def _parse_installed_at(value: object) -> Optional[datetime]:
     if not isinstance(value, str) or not value:
         return None
@@ -313,6 +422,7 @@ def reset_caches(*, drop_disk: bool = False) -> None:
     open (off) instead of pointing at the just-replaced build."""
     _marker_cache.clear()
     _release_memo.clear()
+    _assets_memo.clear()
     if drop_disk:
         import shutil