From 6ffa4bfe90ee34f999e7750c8012853778beeabc Mon Sep 17 00:00:00 2001 From: wasimysaid Date: Sun, 21 Jun 2026 12:09:39 +0200 Subject: [PATCH] Fix/adjust llama update toast for PR #6493 --- studio/backend/routes/llama.py | 1 + studio/backend/tests/test_llama_cpp_update.py | 31 +++++++++++++++++++ studio/backend/tests/test_llama_route.py | 3 +- studio/backend/utils/llama_cpp_update.py | 6 +++- .../src/components/llama-update-banner.tsx | 12 ++----- .../src/features/chat/chat-settings-sheet.tsx | 5 ++- .../src/hooks/use-llama-update-check.ts | 10 +++++- 7 files changed, 54 insertions(+), 14 deletions(-) diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py index 5559cc0404..123349126d 100644 --- a/studio/backend/routes/llama.py +++ b/studio/backend/routes/llama.py @@ -30,6 +30,7 @@ class LlamaUpdateJob(BaseModel): message: str = "" from_tag: Optional[str] = None to_tag: Optional[str] = None + reload_required: Optional[bool] = None error: Optional[str] = None progress: Optional[float] = Field(None, description = "0..1 while running, 1 on success.") started_at: Optional[str] = None diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index 828d439e0a..1b509fea6c 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -445,6 +445,7 @@ def test_start_update_happy_path(monkeypatch, tmp_path): time.sleep(0.05) assert job["state"] == "success", job assert job["to_tag"] == "b9518" + assert job["reload_required"] is False # Installer was invoked with the resolved install dir + latest + repo. assert "--install-dir" in captured["cmd"] assert str(install_dir) in captured["cmd"] @@ -456,6 +457,35 @@ def test_start_update_happy_path(monkeypatch, tmp_path): assert popen_kwargs["env"]["UNSLOTH_PROGRESS_PERCENT_STEP"] == "5" +def test_start_update_reports_full_release_tag(monkeypatch, tmp_path): + install_dir = tmp_path / "llama.cpp" + binary = _write_install(install_dir, "b9595") + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr( + freshness, + "_fetch_latest_release_tag", + lambda repo, timeout = 5.0: "b9596-mix-e6f2453", + ) + + def _on_start(cmd): + _write_install(install_dir, "b9596", release_tag = "b9596-mix-e6f2453") + + _patch_installer_popen(monkeypatch, on_start = _on_start) + + res = upd.start_update() + assert res["started"] is True + deadline = time.time() + 10 + while time.time() < deadline: + job = upd.get_update_status()["job"] + if job["state"] in ("success", "error"): + break + time.sleep(0.05) + assert job["state"] == "success", job + assert job["to_tag"] == "b9596-mix-e6f2453" + assert "Updated llama.cpp to b9596-mix-e6f2453." in job["message"] + + def test_start_update_installer_failure_reports_error(monkeypatch, tmp_path): install_dir = tmp_path / "llama.cpp" binary = _write_install(install_dir, "b9493") @@ -674,6 +704,7 @@ def test_update_sets_maintenance_flag_and_unloads(monkeypatch, tmp_path): time.sleep(0.05) assert backend.unloaded is True + assert upd.get_update_status()["job"]["reload_required"] is True assert seen.get("flag_during_install") is True # Cleared in the finally so model loads work again after the swap. assert backend._llama_update_in_progress is False diff --git a/studio/backend/tests/test_llama_route.py b/studio/backend/tests/test_llama_route.py index 333f710b44..0ecfeee018 100644 --- a/studio/backend/tests/test_llama_route.py +++ b/studio/backend/tests/test_llama_route.py @@ -70,10 +70,11 @@ def test_status_response_exposes_source_build(): "installed_at_utc": None, "age_days": None, "source_build": True, - "job": {"state": "idle"}, + "job": {"state": "idle", "reload_required": False}, } model = rl.LlamaUpdateStatusResponse(**payload) assert model.model_dump()["source_build"] is True + assert model.model_dump()["job"]["reload_required"] is False # Extra/unknown keys must not crash the response model. rl.LlamaUpdateStatusResponse(**{**payload, "unexpected": 1}) diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 965f1a970a..4e15fbe074 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -62,6 +62,7 @@ _job: dict = { "message": "", "from_tag": None, "to_tag": None, + "reload_required": None, "error": None, "progress": None, "started_at": None, @@ -505,7 +506,7 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path except Exception as exc: # pragma: no cover - network defensive logger.debug("llama update: post-install freshness refresh failed", error = str(exc)) new_marker = read_install_marker(_find_binary()) - new_tag = (new_marker or {}).get("tag") or (new_marker or {}).get("release_tag") + new_tag = (new_marker or {}).get("release_tag") or (new_marker or {}).get("tag") with _job_lock: _job.update( @@ -515,6 +516,7 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path + (" Reload your model to use it." if model_was_active else "") ), to_tag = new_tag, + reload_required = model_was_active, error = None, progress = 1.0, finished_at = _utcnow(), @@ -618,6 +620,7 @@ def start_update() -> dict: message = "Downloading and installing the latest llama.cpp prebuilt...", from_tag = from_tag, to_tag = None, + reload_required = None, error = None, progress = 0.0, started_at = _utcnow(), @@ -643,6 +646,7 @@ def _reset_job_for_tests() -> None: message = "", from_tag = None, to_tag = None, + reload_required = None, error = None, progress = None, started_at = None, diff --git a/studio/frontend/src/components/llama-update-banner.tsx b/studio/frontend/src/components/llama-update-banner.tsx index 0daade512b..f53e6e0401 100644 --- a/studio/frontend/src/components/llama-update-banner.tsx +++ b/studio/frontend/src/components/llama-update-banner.tsx @@ -2,7 +2,6 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { Button } from "@/components/ui/button"; -import { isExternalModelId, useChatRuntimeStore } from "@/features/chat"; import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check"; import { useShowLlamaUpdateBanner } from "@/hooks/use-llama-update-pref"; import { toast } from "@/lib/toast"; @@ -100,15 +99,8 @@ export function LlamaUpdateBanner({ async function handleUpdate() { const result = await apply(); if (result?.ok) { - // Prefer the full release tag (e.g. b9726-mix-); the job's to_tag and - // installed_tag are the bare bNNNN build number. - const updatedTag = status?.latest_tag ?? result.tag ?? "the latest build"; - // Only a loaded local model can be reloaded to pick up the new binary; - // external-provider models do not use llama.cpp at all. - const checkpoint = useChatRuntimeStore.getState().params.checkpoint; - const hasLocalModel = - Boolean(checkpoint) && !isExternalModelId(checkpoint); - const reloadHint = hasLocalModel + const updatedTag = result.tag ?? status?.latest_tag ?? "the latest build"; + const reloadHint = result.reloadRequired ? " Reload your model to use it." : ""; toast.success(`llama.cpp updated to ${updatedTag}.${reloadHint}`); diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 935f7e1bc6..8a7f09062a 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -518,8 +518,11 @@ export function ChatSettingsPanel({ const handleMtpUpdate = useCallback(async () => { const result = await applyLlamaUpdate(); if (result.ok) { + const reloadHint = result.reloadRequired + ? " Reload your model to enable MTP." + : ""; toast.success( - `llama.cpp updated to ${result.tag ?? "the latest build"}. Reload your model to enable MTP.`, + `llama.cpp updated to ${result.tag ?? "the latest build"}.${reloadHint}`, ); } else { toast.error(`llama.cpp update failed: ${result.error ?? "unknown error"}`); diff --git a/studio/frontend/src/hooks/use-llama-update-check.ts b/studio/frontend/src/hooks/use-llama-update-check.ts index ccfbb0d377..3b891cddd0 100644 --- a/studio/frontend/src/hooks/use-llama-update-check.ts +++ b/studio/frontend/src/hooks/use-llama-update-check.ts @@ -21,6 +21,7 @@ export interface LlamaUpdateJob { message: string; from_tag: string | null; to_tag: string | null; + reload_required: boolean | null; error: string | null; // Download fraction (0..1) while running, 1 on success, null when unknown. progress: number | null; @@ -52,6 +53,8 @@ function parseStatus(value: unknown): LlamaUpdateStatus | null { message: typeof job.message === "string" ? job.message : "", from_tag: typeof job.from_tag === "string" ? job.from_tag : null, to_tag: typeof job.to_tag === "string" ? job.to_tag : null, + reload_required: + typeof job.reload_required === "boolean" ? job.reload_required : null, error: typeof job.error === "string" ? job.error : null, progress: typeof job.progress === "number" ? job.progress : null, }, @@ -85,6 +88,7 @@ interface UseLlamaUpdateCheckOptions { export interface LlamaApplyResult { ok: boolean; tag?: string | null; + reloadRequired?: boolean | null; error?: string | null; } @@ -126,7 +130,11 @@ export function useLlamaUpdateCheck({ if (s.job.state === "success") { setVisible(false); void refreshHardwareInfo(); - onDone?.({ ok: true, tag: s.job.to_tag }); + onDone?.({ + ok: true, + tag: s.job.to_tag, + reloadRequired: s.job.reload_required, + }); } else if (s.job.state === "error") { // Leave the banner up so the user can retry; clearing applying drops // the "Updating..." state.