diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py index 540647e3bc..1f3f0cc039 100644 --- a/studio/backend/routes/llama.py +++ b/studio/backend/routes/llama.py @@ -14,19 +14,61 @@ never blocks on a missing marker / offline GitHub. from __future__ import annotations import asyncio +import socket +import sys import threading from typing import Optional -from fastapi import APIRouter, Depends, Query +from fastapi import APIRouter, Body, Depends, Query from pydantic import BaseModel, Field from auth.authentication import get_current_subject from loggers import get_logger from utils.llama_cpp_update import get_update_status, start_update +from utils.update_confirm import ( + CONFIRM_TOKEN_TTL_SECONDS, + consume_confirm_token, + mint_confirm_token, +) logger = get_logger(__name__) router = APIRouter() +# Messages for a refused apply, keyed by outcome. A refusal never runs the swap. +_REFUSAL_MESSAGES = { + "confirmation_required": ( + "Confirm the llama.cpp update before it runs. It will download and swap " + "the binary on the machine running Studio." + ), + "invalid_token": ( + "The confirmation is missing or unrecognized. Re-check for the update and " + "confirm again before it runs." + ), + "expired_token": ("The confirmation expired. Re-check for the update and confirm again."), + "stale_target": ( + "The available build changed since you confirmed. Re-check for the update " + "and confirm the new build before it runs." + ), +} + + +class UpdateMachine(BaseModel): + """The host a swap targets, so a remote operator sees which machine changes.""" + + hostname: str = Field("", description = "Hostname of the machine running Studio.") + platform: str = Field( + "", description = "Host OS tag (sys.platform), e.g. 'linux', 'darwin', 'win32'." + ) + + +def _current_machine() -> UpdateMachine: + """Best-effort host identity; cross-OS and never branched on.""" + try: + hostname = socket.gethostname() or "unknown" + except Exception: # pragma: no cover + hostname = "unknown" + return UpdateMachine(hostname = hostname, platform = sys.platform) + class LlamaUpdateJob(BaseModel): state: str = Field("idle", description = "idle | running | success | error") @@ -62,13 +104,58 @@ class LlamaUpdateStatusResponse(BaseModel): update_size_bytes: Optional[int] = Field( None, description = "Download size of the prebuilt Update would fetch, in bytes." ) + machine: UpdateMachine = Field( + default_factory = UpdateMachine, + description = "The host this status describes; the machine an Update would change.", + ) job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob) +class LlamaUpdateRequest(BaseModel): + """Body for POST /update. Empty body = no confirmation, so the swap is refused.""" + + confirm_token: Optional[str] = Field( + None, + description = "Single-use token from POST /update/confirm, bound to the exact offered build.", + ) + confirmed: bool = Field( + False, + description = "Explicit confirmation for non-interactive/CLI callers that do not use a token.", + ) + + +class LlamaUpdateConfirmResponse(BaseModel): + """Pending swap details plus, when appliable, a confirm token for the prompt.""" + + update_available: bool = False + appliable: bool = Field( + False, description = "True when an update exists and can actually be applied here." + ) + reason: Optional[str] = Field( + None, description = "Why not appliable: up_to_date | local_link | ..." + ) + machine: UpdateMachine = Field(default_factory = UpdateMachine) + installed_tag: Optional[str] = None + latest_tag: Optional[str] = None + update_size_bytes: Optional[int] = None + confirm_token: Optional[str] = Field( + None, description = "Echo back to POST /update to run the swap. Absent when not appliable." + ) + confirm_expires_at: Optional[str] = None + confirm_ttl_seconds: Optional[int] = None + + class LlamaUpdateActionResponse(BaseModel): started: bool reason: Optional[str] = None message: Optional[str] = None + machine: UpdateMachine = Field( + default_factory = UpdateMachine, + description = "The host the swap ran (or would run) on; names the result's machine.", + ) + installed_tag: Optional[str] = None + latest_tag: Optional[str] = None + update_size_bytes: Optional[int] = None job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob) @@ -77,8 +164,7 @@ _last_llama_update_step = -1 def _log_llama_update_progress(job: LlamaUpdateJob) -> None: - """One llama_update_progress line per 10% step so a prebuilt update reports - progress without a line per poll. Resyncs when a new update starts.""" + """Log one progress line per 10% step, not per poll; resyncs on a new update.""" global _last_llama_update_step if job.state != "running" or job.progress is None: return @@ -102,14 +188,116 @@ async def llama_update_status( ) -> LlamaUpdateStatusResponse: # Off the event loop: detection may probe the host and read GitHub. status = await asyncio.to_thread(get_update_status, force_refresh = force_refresh) - resp = LlamaUpdateStatusResponse(**status) + resp = LlamaUpdateStatusResponse(machine = _current_machine(), **status) _log_llama_update_progress(resp.job) return resp +@router.post("/update/confirm", response_model = LlamaUpdateConfirmResponse) +async def llama_update_confirm( + current_subject: str = Depends(get_current_subject), +) -> LlamaUpdateConfirmResponse: + """Step one of the two-step apply: describe the pending swap and, when + appliable, mint a single-use token bound to the offered build. The gate is + confirmation, not caller location.""" + # Force-refresh so the token binds the current build, not a stale cached tag. + status = await asyncio.to_thread(get_update_status, force_refresh = True) + machine = _current_machine() + installed_tag = status.get("installed_tag") + latest_tag = status.get("latest_tag") + size = status.get("update_size_bytes") + + # A --with-llama-cpp-dir local link reports update_available=False, so this must + # run before the up_to_date branch or the local_link reason gets masked. + if status.get("local_link"): + return LlamaUpdateConfirmResponse( + update_available = True, + appliable = False, + reason = "local_link", + machine = machine, + installed_tag = installed_tag, + latest_tag = latest_tag, + update_size_bytes = size, + ) + if not status.get("update_available"): + return LlamaUpdateConfirmResponse( + update_available = False, + appliable = False, + reason = status.get("reason") or "up_to_date", + machine = machine, + installed_tag = installed_tag, + latest_tag = latest_tag, + update_size_bytes = size, + ) + + token, expires_at = mint_confirm_token(latest_tag or "") + return LlamaUpdateConfirmResponse( + update_available = True, + appliable = True, + reason = None, + machine = machine, + installed_tag = installed_tag, + latest_tag = latest_tag, + update_size_bytes = size, + confirm_token = token, + confirm_expires_at = expires_at, + confirm_ttl_seconds = CONFIRM_TOKEN_TTL_SECONDS, + ) + + @router.post("/update", response_model = LlamaUpdateActionResponse) async def llama_update( + request: Optional[LlamaUpdateRequest] = Body(default = None), current_subject: str = Depends(get_current_subject), ) -> LlamaUpdateActionResponse: - action = await asyncio.to_thread(start_update) - return LlamaUpdateActionResponse(**action) + """Apply the swap, but only with an explicit, fresh confirmation. + + The installer replaces the host binary, so a caller must either echo the + single-use ``confirm_token`` from POST /update/confirm (replay-safe, build-bound) + or send ``confirmed=true`` (non-interactive callers); with neither, the swap is + refused untouched. The gate is confirmation, not the caller's location.""" + req = request or LlamaUpdateRequest() + machine = _current_machine() + # Force-refresh so the token is validated against the same build start_update will + # resolve; a stale cache could accept an old-tag token then install a newer build. + status = await asyncio.to_thread(get_update_status, force_refresh = True) + installed_tag = status.get("installed_tag") + target_tag = status.get("latest_tag") + size = status.get("update_size_bytes") + + confirmed = False + refuse_reason = "confirmation_required" + if req.confirm_token: + ok, why = consume_confirm_token(req.confirm_token, target_tag or "") + if ok: + confirmed = True + else: + refuse_reason = why or "invalid_token" + elif req.confirmed: + confirmed = True + + if not confirmed: + # Safe default: no confirmation, no swap; return a visible, actionable refusal. + return LlamaUpdateActionResponse( + started = False, + reason = refuse_reason, + message = _REFUSAL_MESSAGES.get( + refuse_reason, _REFUSAL_MESSAGES["confirmation_required"] + ), + machine = machine, + installed_tag = installed_tag, + latest_tag = target_tag, + update_size_bytes = size, + job = LlamaUpdateJob(**status.get("job", {})), + ) + + # Pass the confirmed target so the updater installs exactly that build and + # aborts if latest moved between this refresh and its own (see start_update). + action = await asyncio.to_thread(start_update, target_tag) + return LlamaUpdateActionResponse( + machine = machine, + installed_tag = installed_tag, + latest_tag = target_tag, + update_size_bytes = size, + **action, + ) diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index f12384231f..478f06ea10 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -396,9 +396,62 @@ def test_start_update_source_build_installs_prebuilt(monkeypatch, tmp_path): assert "--llama-tag" in cmd and "latest" in cmd assert cmd[cmd.index("--rocm-gfx") + 1] == "gfx110x" assert "--simple-policy" not in cmd and "--cpu-fallback" not in cmd - # No pin: source-build detection and the unpinned apply share the same - # "latest" resolver, so they already agree. - assert "--published-release-tag" not in cmd + # Pin to the release the resolver picked, so one published before the installer's + # own re-resolve can't swap in an unconfirmed build (matches the marker path). + assert "--published-release-tag" in cmd + assert cmd[cmd.index("--published-release-tag") + 1] == "b9585" + + +def test_start_update_source_build_pins_resolver_release_tag(monkeypatch, tmp_path): + # The source-build apply pins the installer to the resolver's release_tag, not the + # display tag. For a fork wrapper they differ ("v1.0" vs "b9457"); only the real + # release tag is a valid --published-release-tag. Confirming the displayed tag + # must still proceed and pin the real release. + install_dir = tmp_path / "llama.cpp" + binary = install_dir / "build" / "bin" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("stub") # no marker -> source-build path + monkeypatch.delenv("UNSLOTH_LLAMA_CPP_PATH", raising = False) + monkeypatch.setattr(upd, "_find_binary", lambda: str(binary)) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(upd, "_installed_build_number", lambda b: 9000) # behind b9457 + _prebuilt( + monkeypatch, + repo = "unslothai/llama.cpp", + release_tag = "v1.0", + llama_tag = "b9457", + asset = "llama-b9457-bin-linux-x64.tar.gz", + ) + + captured = {} + + class _Proc: + returncode = 0 + stdout = "installed" + stderr = "" + + def _fake_run(cmd, **kwargs): + return _Proc() + + def _on_start(cmd): + captured["cmd"] = cmd + # Installer writes the marker for the pinned release (real tag v1.0). + _write_install(install_dir, "b9457", release_tag = "v1.0") + + monkeypatch.setattr(upd.subprocess, "run", _fake_run) + _patch_installer_popen(monkeypatch, on_start = _on_start) + + res = upd.start_update(expected_tag = "b9457") # the displayed/confirmed tag + assert res["started"] is True, res + deadline = time.time() + 10 + while time.time() < deadline: + if upd.get_update_status()["job"]["state"] in ("success", "error"): + break + time.sleep(0.05) + cmd = captured["cmd"] + assert "--published-release-tag" in cmd + assert cmd[cmd.index("--published-release-tag") + 1] == "v1.0" + assert upd.get_update_status()["job"]["state"] == "success" def test_start_update_happy_path(monkeypatch, tmp_path): diff --git a/studio/backend/tests/test_llama_route.py b/studio/backend/tests/test_llama_route.py index 0ecfeee018..a3e3a76c7a 100644 --- a/studio/backend/tests/test_llama_route.py +++ b/studio/backend/tests/test_llama_route.py @@ -123,11 +123,14 @@ def test_status_handler_runs_off_event_loop(monkeypatch): def test_update_handler_runs_off_event_loop(monkeypatch): seen = {} - def fake_start(): + def fake_start(expected_tag = None): seen["thread"] = threading.current_thread() return {"started": True, "reason": None, "job": {"state": "running"}} monkeypatch.setattr(rl, "start_update", fake_start) - out = asyncio.run(rl.llama_update(current_subject = "t")) + # The swap now needs explicit confirmation; confirm here so the off-loop path runs. + out = asyncio.run( + rl.llama_update(request = rl.LlamaUpdateRequest(confirmed = True), current_subject = "t") + ) assert out.started is True assert seen["thread"] is not threading.main_thread() diff --git a/studio/backend/tests/test_update_contract.py b/studio/backend/tests/test_update_contract.py new file mode 100644 index 0000000000..aaa2795aaf --- /dev/null +++ b/studio/backend/tests/test_update_contract.py @@ -0,0 +1,527 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Contract tests for the confirmation + visible-result llama.cpp update (#7003). + +Properties covered: + - remote authenticated + confirmed -> update proceeds, result names the host + - remote authenticated, no confirm -> refused with a clear message, NO swap + - local / desktop -> identical behavior (no same-machine axis) + - unauthenticated -> refused (401), NO swap + - stale / expired / replayed token -> refused, NO swap + +routes/llama.py loads standalone with stubbed auth/loggers/llama_cpp_update and the +real utils.update_confirm. Exercised via handler calls and a FastAPI TestClient. +""" + +from __future__ import annotations + +import asyncio +import importlib.util +import sys +import types +from pathlib import Path +from typing import Optional + +import pytest + +pytest.importorskip("fastapi") + +_HERE = Path(__file__).resolve().parent +_BACKEND = _HERE.parent + +# Stub modules _install_stubs writes into the shared sys.modules, tracked so they +# are restored after the route loads and a regression test can assert none leak +# into the rest of the one-process backend suite. +_INSTALLED_STUBS: dict[str, types.ModuleType] = {} +_STUBBED_KEYS = ( + "auth", + "auth.authentication", + "loggers", + "utils", + "utils.llama_cpp_update", + "utils.update_confirm", +) + + +def _install_stubs(): + """Stub packages so routes/llama.py imports cleanly, plus the real update_confirm.""" + # auth.get_current_subject: a real FastAPI dep that 401s without a valid bearer. + from fastapi import Header, HTTPException + + def get_current_subject(authorization: Optional[str] = Header(default = None)) -> str: + if authorization == "Bearer good": + return "operator" + raise HTTPException(status_code = 401, detail = "unauthorized") + + auth_pkg = types.ModuleType("auth") + auth_pkg.__path__ = [] + auth_mod = types.ModuleType("auth.authentication") + auth_mod.get_current_subject = get_current_subject + + # loggers.get_logger -> no-op structured logger. + loggers_mod = types.ModuleType("loggers") + + class _NopLogger: + def info(self, *a, **k): + pass + + def debug(self, *a, **k): + pass + + def warning(self, *a, **k): + pass + + loggers_mod.get_logger = lambda *_a, **_k: _NopLogger() + + # utils.llama_cpp_update -> monkeypatchable get_update_status / start_update. + utils_pkg = sys.modules.get("utils") + if utils_pkg is None: + utils_pkg = types.ModuleType("utils") + utils_pkg.__path__ = [] + sys.modules["utils"] = utils_pkg + _INSTALLED_STUBS["utils"] = utils_pkg + + lcu_mod = types.ModuleType("utils.llama_cpp_update") + + def _default_status(force_refresh: bool = False): + return { + "supported": True, + "update_available": True, + "stale": False, + "installed_tag": "b9860", + "latest_tag": "b9909", + "published_repo": "unslothai/llama.cpp", + "installed_at_utc": None, + "age_days": None, + "source_build": False, + "update_size_bytes": 42_000_000, + "job": {"state": "idle"}, + } + + def _default_start(expected_tag = None): + return {"started": True, "reason": None, "job": {"state": "running"}} + + lcu_mod.get_update_status = _default_status + lcu_mod.start_update = _default_start + + # The real confirmation-token module, under its production import name. + uc_spec = importlib.util.spec_from_file_location( + "utils.update_confirm", str(_BACKEND / "utils" / "update_confirm.py") + ) + uc_mod = importlib.util.module_from_spec(uc_spec) + uc_spec.loader.exec_module(uc_mod) + + for name, mod in ( + ("auth", auth_pkg), + ("auth.authentication", auth_mod), + ("loggers", loggers_mod), + ("utils.llama_cpp_update", lcu_mod), + ("utils.update_confirm", uc_mod), + ): + _INSTALLED_STUBS[name] = mod + sys.modules[name] = mod + return uc_mod + + +# Snapshot the real modules before installing the collection-time stubs, then +# restore them once the route has bound its stubbed imports. Leaving the stubs in +# sys.modules poisons the shared process: a bare ``auth`` package and a NopLogger +# break every later test that imports the real auth/loggers/utils modules. +_saved_modules = {name: sys.modules.get(name) for name in _STUBBED_KEYS} +_uc = _install_stubs() + + +def _load_route(): + spec = importlib.util.spec_from_file_location( + "llama_route_under_test", str(_BACKEND / "routes" / "llama.py") + ) + mod = importlib.util.module_from_spec(spec) + sys.modules["llama_route_under_test"] = mod # pydantic forward-ref resolution + spec.loader.exec_module(mod) + return mod + + +rl = _load_route() + +# rl now holds its stub references; put the real modules back so nothing leaks. +for _name, _orig in _saved_modules.items(): + if _orig is None: + sys.modules.pop(_name, None) + else: + sys.modules[_name] = _orig + + +@pytest.fixture(autouse = True) +def _reset(monkeypatch): + _uc.reset_tokens_for_tests() + # Restore default status/start each test; individual tests override as needed. + monkeypatch.setattr( + rl, + "get_update_status", + lambda force_refresh = False: { + "supported": True, + "update_available": True, + "installed_tag": "b9860", + "latest_tag": "b9909", + "update_size_bytes": 42_000_000, + "job": {"state": "idle"}, + }, + ) + yield + + +def _track_start(monkeypatch): + calls = {"n": 0, "expected_tags": []} + + def _start(expected_tag = None): + calls["n"] += 1 + calls["expected_tags"].append(expected_tag) + return {"started": True, "reason": None, "job": {"state": "running"}} + + monkeypatch.setattr(rl, "start_update", _start) + return calls + + +# update_confirm token unit tests + + +def test_token_roundtrip_single_use(): + tok, exp = _uc.mint_confirm_token("b9909") + assert isinstance(tok, str) and tok + ok, reason = _uc.consume_confirm_token(tok, "b9909") + assert ok and reason is None + # Single use: a replay is refused. + ok2, reason2 = _uc.consume_confirm_token(tok, "b9909") + assert ok2 is False and reason2 == "invalid_token" + + +def test_token_stale_target_refused(): + tok, _ = _uc.mint_confirm_token("b9909") + ok, reason = _uc.consume_confirm_token(tok, "b9910") # offered build changed + assert ok is False and reason == "stale_target" + + +def test_token_expired_refused(): + tok, _ = _uc.mint_confirm_token("b9909", ttl_seconds = 0) + ok, reason = _uc.consume_confirm_token(tok, "b9909") + assert ok is False and reason == "expired_token" + + +def test_token_missing_refused(): + ok, reason = _uc.consume_confirm_token(None, "b9909") + assert ok is False and reason == "invalid_token" + + +# Handler-level: the swap only runs with an explicit confirmation + + +def test_apply_without_confirmation_is_refused_and_never_swaps(monkeypatch): + calls = _track_start(monkeypatch) + out = asyncio.run(rl.llama_update(request = None, current_subject = "operator")) + assert out.started is False + assert out.reason == "confirmation_required" + assert "swap" in out.message.lower() or "confirm" in out.message.lower() + # The result still names the host + build so the refusal is actionable. + assert out.machine.hostname + assert out.latest_tag == "b9909" + assert calls["n"] == 0 # installer NEVER started + + +def test_apply_with_confirmed_true_proceeds(monkeypatch): + calls = _track_start(monkeypatch) + body = rl.LlamaUpdateRequest(confirmed = True) + out = asyncio.run(rl.llama_update(request = body, current_subject = "operator")) + assert out.started is True + assert out.machine.hostname # which machine + assert out.latest_tag == "b9909" # which version + assert calls["n"] == 1 + # The confirmed target is threaded into start_update so it installs exactly + # that build (and aborts if latest moved since this refresh). + assert calls["expected_tags"] == ["b9909"] + + +def test_apply_with_valid_token_proceeds(monkeypatch): + calls = _track_start(monkeypatch) + tok, _ = _uc.mint_confirm_token("b9909") + body = rl.LlamaUpdateRequest(confirm_token = tok) + out = asyncio.run(rl.llama_update(request = body, current_subject = "operator")) + assert out.started is True + assert calls["n"] == 1 + # Token is single-use: replaying the same body is refused, no second swap. + out2 = asyncio.run(rl.llama_update(request = body, current_subject = "operator")) + assert out2.started is False + assert out2.reason == "invalid_token" + assert calls["n"] == 1 + + +def test_apply_with_stale_token_is_refused(monkeypatch): + calls = _track_start(monkeypatch) + tok, _ = _uc.mint_confirm_token("b9900") # bound to an older offered build + body = rl.LlamaUpdateRequest(confirm_token = tok) + out = asyncio.run(rl.llama_update(request = body, current_subject = "operator")) + assert out.started is False + assert out.reason == "stale_target" + assert calls["n"] == 0 + + +def test_confirm_endpoint_describes_and_mints(monkeypatch): + out = asyncio.run(rl.llama_update_confirm(current_subject = "operator")) + assert out.update_available is True + assert out.appliable is True + assert out.installed_tag == "b9860" + assert out.latest_tag == "b9909" + assert out.machine.hostname + assert out.confirm_token + # The freshly minted token applies the update end to end. + calls = _track_start(monkeypatch) + body = rl.LlamaUpdateRequest(confirm_token = out.confirm_token) + applied = asyncio.run(rl.llama_update(request = body, current_subject = "operator")) + assert applied.started is True + assert calls["n"] == 1 + + +def test_confirm_endpoint_up_to_date_offers_no_token(monkeypatch): + monkeypatch.setattr( + rl, + "get_update_status", + lambda force_refresh = False: { + "update_available": False, + "installed_tag": "b9909", + "latest_tag": "b9909", + "job": {"state": "idle"}, + }, + ) + out = asyncio.run(rl.llama_update_confirm(current_subject = "operator")) + assert out.update_available is False + assert out.appliable is False + assert out.confirm_token is None + + +def test_confirm_endpoint_reports_local_link_not_up_to_date(monkeypatch): + # A --with-llama-cpp-dir tree reports local_link=True with update_available=False; + # the confirm endpoint must surface reason="local_link", not the generic up_to_date. + monkeypatch.setattr( + rl, + "get_update_status", + lambda force_refresh = False: { + "update_available": False, + "local_link": True, + "installed_tag": "b9909", + "latest_tag": None, + "job": {"state": "idle"}, + }, + ) + out = asyncio.run(rl.llama_update_confirm(current_subject = "operator")) + assert out.reason == "local_link" + assert out.appliable is False + assert out.confirm_token is None + + +def test_apply_revalidates_token_against_refreshed_target(monkeypatch): + # Token confirmed for b9909; b9910 publishes before apply. The apply must + # re-resolve and refuse the now-stale token, not install an unconfirmed build. + def _status(force_refresh: bool = False): + return { + "supported": True, + "update_available": True, + "installed_tag": "b9860", + "latest_tag": "b9910" if force_refresh else "b9909", + "update_size_bytes": 42_000_000, + "job": {"state": "idle"}, + } + + monkeypatch.setattr(rl, "get_update_status", _status) + calls = _track_start(monkeypatch) + tok, _ = _uc.mint_confirm_token("b9909") + body = rl.LlamaUpdateRequest(confirm_token = tok) + out = asyncio.run(rl.llama_update(request = body, current_subject = "operator")) + assert out.started is False + assert out.reason == "stale_target" + assert calls["n"] == 0 # never installed the unconfirmed b9910 + + +def test_status_reports_machine(): + out = asyncio.run(rl.llama_update_status(force_refresh = False, current_subject = "operator")) + assert out.machine.hostname + assert out.machine.platform + assert out.latest_tag == "b9909" + + +# start_update pins the confirmed target: a release publishing after confirmation +# must not be installed in place of the confirmed build. + + +def _load_real_llama_cpp_update(): + """Load the real utils.llama_cpp_update under a private name (leaving the + module-level stub for the handler tests) plus its two real deps, to exercise + start_update's confirmed-target guard directly.""" + + def _load(mod_name: str, filename: str): + spec = importlib.util.spec_from_file_location(mod_name, str(_BACKEND / "utils" / filename)) + mod = importlib.util.module_from_spec(spec) + sys.modules[mod_name] = mod + spec.loader.exec_module(mod) + return mod + + for mod_name, filename in ( + ("utils.llama_cpp_freshness", "llama_cpp_freshness.py"), + ("utils.process_lifetime", "process_lifetime.py"), + ): + if mod_name not in sys.modules: + _load(mod_name, filename) + return _load("real_llama_cpp_update", "llama_cpp_update.py") + + +def _stub_marker_update(monkeypatch, lcu, *, resolved_latest: str): + """Route start_update down the marker branch with a fixed freshly-resolved + latest tag; no network, no real install-root probing.""" + monkeypatch.setattr(lcu, "_find_binary", lambda: "/llama.cpp/build/bin/llama-server") + monkeypatch.setattr(lcu, "_active_install_is_local_link", lambda binary: False) + monkeypatch.setattr( + lcu, + "read_install_marker", + lambda binary: { + "tag": "b9860", + "published_repo": "unslothai/llama.cpp", + "asset": "cuda-x64.zip", + }, + ) + monkeypatch.setattr(lcu, "_installer_script", lambda: Path("/fake/install_llama_prebuilt.py")) + monkeypatch.setattr(lcu, "_install_dir_for", lambda binary: Path("/llama.cpp/build")) + monkeypatch.setattr( + lcu, + "get_update_status", + lambda force_refresh = False: { + "update_available": True, + "installed_tag": "b9860", + "latest_tag": resolved_latest, + "job": {"state": "idle"}, + }, + ) + + +def test_start_update_aborts_when_resolved_latest_differs_from_confirmed(monkeypatch): + # Operator confirmed b9909, but a newer b9910 publishes before the updater + # re-resolves latest. The confirmed target must win: abort, never install. + lcu = _load_real_llama_cpp_update() + lcu._reset_job_for_tests() + _stub_marker_update(monkeypatch, lcu, resolved_latest = "b9910") + installs = {"n": 0} + monkeypatch.setattr( + lcu, "_run_update", lambda *a, **k: installs.__setitem__("n", installs["n"] + 1) + ) + out = lcu.start_update("b9909") + assert out["started"] is False + assert out["reason"] == "stale_target" + assert installs["n"] == 0 # never swapped to the unconfirmed b9910 + + +def test_start_update_proceeds_when_resolved_latest_matches_confirmed(monkeypatch): + # Confirmed target still matches the freshly-resolved latest: install it and + # pin the installer to exactly that build. + lcu = _load_real_llama_cpp_update() + lcu._reset_job_for_tests() + _stub_marker_update(monkeypatch, lcu, resolved_latest = "b9909") + spawned: dict = {} + + class _CapturingThread: + def __init__( + self, + *, + target = None, + args = (), + name = None, + daemon = None, + ): + spawned["args"] = args + + def start(self): + spawned["started"] = True + + monkeypatch.setattr(lcu.threading, "Thread", _CapturingThread) + out = lcu.start_update("b9909") + assert out["started"] is True + assert out["reason"] is None + assert spawned.get("started") is True + # args = (install_dir, repo, asset, script, pin_release_tag); the installer + # is pinned to exactly the confirmed build (a pin is disabled on macOS). + if sys.platform != "darwin": + assert spawned["args"][4] == "b9909" + lcu._reset_job_for_tests() + + +# HTTP-level: auth gate + wiring + backwards-compatible bodyless POST + + +def _client(): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + app = FastAPI() + app.include_router(rl.router, prefix = "/api/llama") + return TestClient(app, raise_server_exceptions = True) + + +def test_http_unauthenticated_update_is_refused(monkeypatch): + calls = _track_start(monkeypatch) + client = _client() + # No Authorization header at all -> 401 from the auth dependency. + r = client.post("/api/llama/update", json = {"confirmed": True}) + assert r.status_code == 401 + assert calls["n"] == 0 # installer never reached + + +def test_http_authenticated_bodyless_post_refused_no_swap(monkeypatch): + # Backwards compat: an OLD frontend posts /update with no body -> refused + # (confirmation_required), never a silent swap. + calls = _track_start(monkeypatch) + client = _client() + r = client.post("/api/llama/update", headers = {"Authorization": "Bearer good"}) + assert r.status_code == 200 + body = r.json() + assert body["started"] is False + assert body["reason"] == "confirmation_required" + assert body["machine"]["hostname"] + assert calls["n"] == 0 + + +def test_http_two_step_confirm_then_apply(monkeypatch): + calls = _track_start(monkeypatch) + client = _client() + h = {"Authorization": "Bearer good"} + + step1 = client.post("/api/llama/update/confirm", headers = h) + assert step1.status_code == 200 + tok = step1.json()["confirm_token"] + assert tok + + step2 = client.post("/api/llama/update", headers = h, json = {"confirm_token": tok}) + assert step2.status_code == 200 + assert step2.json()["started"] is True + assert step2.json()["machine"]["hostname"] + assert calls["n"] == 1 + + +def test_http_local_and_remote_behave_identically(monkeypatch): + # No same-machine axis: the contract depends only on auth + confirm, so local + # and remote authenticated callers get identical outcomes. + calls = _track_start(monkeypatch) + client = _client() + h = {"Authorization": "Bearer good"} + # Proxy/forwarding headers a "not host" gate would key on change nothing here. + remote_h = {**h, "X-Forwarded-For": "203.0.113.7", "CF-Connecting-IP": "203.0.113.7"} + + local = client.post("/api/llama/update", headers = h, json = {"confirmed": True}) + remote = client.post("/api/llama/update", headers = remote_h, json = {"confirmed": True}) + assert local.json()["started"] == remote.json()["started"] == True + assert calls["n"] == 2 # both proceeded; location was never the gate + + +def test_collection_stubs_do_not_leak_into_sys_modules(): + # Regression: the collection-time stubs (bare auth package, NopLogger loggers, + # default update fns) must be removed from sys.modules after the route loads. + # Leaving them poisons the one-process backend suite -- every later test that + # imports the real auth/loggers/utils modules fails with "unknown location". + for name, stub in _INSTALLED_STUBS.items(): + assert sys.modules.get(name) is not stub, f"{name} stub leaked into sys.modules" diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 67733bde35..f5032bb102 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -629,9 +629,14 @@ def _run_update( pass -def start_update() -> dict: +def start_update(expected_tag: Optional[str] = None) -> dict: """Kick off a background update. Idempotent: a second call while one is - running returns the in-flight job rather than starting another.""" + running returns the in-flight job rather than starting another. + + ``expected_tag`` is the build the caller confirmed. Since the updater re-resolves + "latest", a release published after confirmation could move the target; when the + freshly-resolved latest differs from ``expected_tag`` this aborts rather than swap + an unconfirmed build. None -> no target pinning (prior behaviour).""" binary = _find_binary() # Refuse to update a --with-llama-cpp-dir local link: installing a prebuilt # here would write through the link into the user's own checkout (or fail) @@ -686,6 +691,7 @@ def start_update() -> dict: # (skipping too-new prebuilts); elsewhere an unusable latest now fails # the job loudly (retryable) instead of walking back. pin_release_tag = None if sys.platform == "darwin" else status.get("latest_tag") + resolved_tag = status.get("latest_tag") else: # Source build / custom path: only proceed when the same detection logic # would offer the update (prebuilt exists, install is behind, root is @@ -713,11 +719,29 @@ def start_update() -> dict: repo = (res or {}).get("repo") or DEFAULT_PUBLISHED_REPO from_tag = None asset = (res or {}).get("asset") + # Pin the installer to the release the host-aware resolver already picked (res + # release_tag: the real GitHub tag with any macOS walk-back applied), so a + # release published before the installer's own re-resolve can't swap in an + # unconfirmed build. Host-compatible by construction, so pinning disables no + # needed walk-back (unlike the marker path); also arms _run_update's post- + # install tag check. Unpinned only if the resolver reported no release tag. + pin_release_tag = (res or {}).get("release_tag") or None + resolved_tag = src.get("latest_tag") # Source builds carry no forced-CPU marker, so nothing to preserve here. force_cpu = False - # No pin: source-build detection resolves via --resolve-prebuilt latest, - # the same resolver the unpinned apply uses, so the two already agree. - pin_release_tag = None + + # Install exactly the build the caller confirmed: a release published since + # confirmation moves latest above, so abort rather than swap an unconfirmed build. + if expected_tag is not None and resolved_tag != expected_tag: + return { + "started": False, + "reason": "stale_target", + "message": ( + "The available llama.cpp build changed since it was confirmed. " + "Re-check for the update and confirm the new build before it runs." + ), + "job": get_update_status()["job"], + } if install_dir is None: return { diff --git a/studio/backend/utils/update_confirm.py b/studio/backend/utils/update_confirm.py new file mode 100644 index 0000000000..3bb7493945 --- /dev/null +++ b/studio/backend/utils/update_confirm.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Single-use, short-lived confirmation tokens for the llama.cpp host-binary swap. + +The update runs an OS installer that replaces the binary on the machine running +Studio, so it must not fire from an unconfirmed click, a stale banner, or a replay. +Confirmation is required uniformly for every caller (not gated on "same machine", +since a headless SSH server has no host-local session). A token binds its offered +build (``target_tag``), is single-use with a short TTL, and lives in-process +(matching the single-process backend; use an HMAC stateless token for multi-worker). +""" + +from __future__ import annotations + +import secrets +import threading +import time +from datetime import datetime, timezone +from typing import Optional, Tuple + +# How long a freshly minted confirmation token stays valid. +CONFIRM_TOKEN_TTL_SECONDS = 300 +# Cap the store so unapplied confirm calls can't grow memory unbounded; oldest first. +_MAX_TOKENS = 64 + +_lock = threading.Lock() +# token -> (target_tag, expires_at_monotonic) +_tokens: "dict[str, Tuple[str, float]]" = {} + + +def _iso(ts_epoch: float) -> str: + return ( + datetime.fromtimestamp(ts_epoch, tz = timezone.utc) + .replace(microsecond = 0) + .isoformat() + .replace("+00:00", "Z") + ) + + +def _purge_expired_locked(now: float) -> None: + expired = [tok for tok, (_tag, exp) in _tokens.items() if exp <= now] + for tok in expired: + _tokens.pop(tok, None) + + +def mint_confirm_token( + target_tag: str, *, ttl_seconds: int = CONFIRM_TOKEN_TTL_SECONDS +) -> Tuple[str, str]: + """Mint a single-use token bound to ``target_tag``. + + Returns ``(token, expires_at_iso)``. ``expires_at_iso`` is wall-clock UTC for + display; validity itself is tracked on a monotonic clock so a system time + change cannot extend or shorten it. + """ + now_mono = time.monotonic() + now_wall = time.time() + token = secrets.token_urlsafe(32) + with _lock: + _purge_expired_locked(now_mono) + if len(_tokens) >= _MAX_TOKENS: + oldest = min(_tokens, key = lambda t: _tokens[t][1]) # evict nearest expiry + _tokens.pop(oldest, None) + _tokens[token] = (target_tag, now_mono + ttl_seconds) + return token, _iso(now_wall + ttl_seconds) + + +def consume_confirm_token( + token: Optional[str], current_target_tag: str +) -> Tuple[bool, Optional[str]]: + """Validate and consume a token for a swap to ``current_target_tag``. + + Returns ``(ok, reason)``. On success ``(True, None)`` and the token is burned + (single use). On failure ``ok`` is False and ``reason`` is one of: + ``invalid_token`` (missing/unknown), ``expired_token``, ``stale_target`` + (bound to a different build than the one about to install). + """ + if not token: + return False, "invalid_token" + now_mono = time.monotonic() + with _lock: + # Pop first so any outcome consumes the token, and an expired hit still + # reports "expired" rather than collapsing to "invalid". + entry = _tokens.pop(token, None) + _purge_expired_locked(now_mono) + if entry is None: + return False, "invalid_token" + target_tag, expires_at = entry + if expires_at <= now_mono: + return False, "expired_token" + if target_tag != (current_target_tag or ""): + return False, "stale_target" + return True, None + + +def reset_tokens_for_tests() -> None: + """Test-only: clear the token store.""" + with _lock: + _tokens.clear() diff --git a/studio/frontend/src/components/llama-update-banner.tsx b/studio/frontend/src/components/llama-update-banner.tsx index 3db15ffe30..75e5c38efb 100644 --- a/studio/frontend/src/components/llama-update-banner.tsx +++ b/studio/frontend/src/components/llama-update-banner.tsx @@ -2,6 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { Button } from "@/components/ui/button"; +import { useLlamaUpdateConfirmGate } from "@/components/llama-update-confirm-dialog"; import { resyncInferenceStatusAfterServerModelChange } from "@/features/chat"; import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check"; import { useShowLlamaUpdateBanner } from "@/hooks/use-llama-update-pref"; @@ -91,16 +92,23 @@ export function LlamaUpdateBanner({ enabled, onReloadRequired: resyncInferenceStatusAfterServerModelChange, }); + // Explicit accept/cancel prompt naming the exact build + host before the swap. + const { requestConfirm, dialog } = useLlamaUpdateConfirmGate(); async function handleUpdate() { - const result = await apply(); + const result = await apply(requestConfirm); if (result?.ok) { - const updatedTag = result.tag ?? status?.latest_tag ?? "the latest build"; + const host = result.machine?.hostname; + const where = host ? ` on ${host}` : ""; + const fromTag = result.fromTag ?? status?.installed_tag ?? "unknown"; + const toTag = 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}`); - } else if (result) { + toast.success( + `llama.cpp${where} updated ${fromTag} to ${toTag}.${reloadHint}`, + ); + } else if (result && result.error !== "canceled") { toast.error( `llama.cpp update failed: ${result.error ?? "unknown error"}`, ); @@ -127,7 +135,7 @@ export function LlamaUpdateBanner({ ); // Avoid opacity/transform transitions; GPU layer churn can flash. - return show ? ( + const banner = show ? (
+ {target?.fromTag ?? "unknown"} →{" "} + + {target?.toTag ?? "the latest build"} + +
+ {target?.machine?.hostname && ( ++ {target.machine.hostname} + {target.machine.platform ? ` ยท ${target.machine.platform}` : ""} +
+ )} +