From 3b446b6d24e9b89cb181204253b0326c198afd97 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 19 Jul 2026 10:36:43 +0000 Subject: [PATCH 1/8] Studio: require confirmation and report the target machine for llama.cpp updates A remote authenticated browser could swap the host llama.cpp binary with no confirmation and no visible result. POST /api/llama/update/confirm now mints a single-use, build-bound token, and POST /api/llama/update refuses a bodyless or unconfirmed call instead of running the swap. Update responses include the host name and platform so the result is attributable. Confirmation is required uniformly, so a headless remote server is not gated out; the desktop app, local browser, and CLI update path are unchanged. --- studio/backend/routes/llama.py | 193 +++++++++- studio/backend/tests/test_llama_route.py | 5 +- studio/backend/tests/test_update_contract.py | 350 +++++++++++++++++++ studio/backend/utils/update_confirm.py | 104 ++++++ 4 files changed, 648 insertions(+), 4 deletions(-) create mode 100644 studio/backend/tests/test_update_contract.py create mode 100644 studio/backend/utils/update_confirm.py diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py index 540647e3bc..6d4b2415dc 100644 --- a/studio/backend/routes/llama.py +++ b/studio/backend/routes/llama.py @@ -14,19 +14,63 @@ 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): + """Host a swap targets, so a remote operator sees which machine would change.""" + + 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 +106,59 @@ 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 means no confirmation, so the swap is + refused (safe default) and a stale banner or replay cannot swap the binary.""" + + 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) @@ -102,14 +192,111 @@ 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 it can + be applied, mint a single-use token bound to the offered build. Available to + any authenticated operator; confirmation, not location, is the gate.""" + # 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") + + 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, + ) + # Update exists but the tree is a --with-llama-cpp-dir local link: describe, do not apply. + 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, + ) + + 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: + """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 (preferred: replay-safe, + bound to the build) or send ``confirmed=true`` (non-interactive callers). With + neither, the swap is refused and the binary is left untouched. The gate is the + confirmation, not the caller's location, so a headless SSH server confirms like a local one.""" + req = request or LlamaUpdateRequest() + machine = _current_machine() + # Cached read: names the host and build for the prompt and binds the token. + status = await asyncio.to_thread(get_update_status) + 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", {})), + ) + action = await asyncio.to_thread(start_update) - return LlamaUpdateActionResponse(**action) + 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_route.py b/studio/backend/tests/test_llama_route.py index 0ecfeee018..ac3ec01797 100644 --- a/studio/backend/tests/test_llama_route.py +++ b/studio/backend/tests/test_llama_route.py @@ -128,6 +128,9 @@ def test_update_handler_runs_off_event_loop(monkeypatch): 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..cf84df143c --- /dev/null +++ b/studio/backend/tests/test_update_contract.py @@ -0,0 +1,350 @@ +# 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 is loaded standalone with stubbed auth / loggers / llama_cpp_update +and the real utils.update_confirm, so no heavy backend deps are needed. Both handler +calls and a real FastAPI TestClient (HTTP + auth gate) are exercised. +""" + +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 + + +def _install_stubs(): + """Register stub packages so routes/llama.py imports cleanly, plus the real + update_confirm module under its production name.""" + # auth.authentication.get_current_subject -> a real FastAPI dependency that + # 401s without a valid bearer, so the HTTP tests can prove the auth gate. + 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 + + 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(): + 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), + ): + sys.modules[name] = mod + return uc_mod + + +_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() + + +@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} + + def _start(): + calls["n"] += 1 + 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 + + +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_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" + + +# --------------------------------------------------------------------------- # +# 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. That must be + # safe -- 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): + # There is no same-machine axis: the contract depends only on auth + confirm, + # so a "local" and a "remote" authenticated caller get identical outcomes. + calls = _track_start(monkeypatch) + client = _client() + h = {"Authorization": "Bearer good"} + # Simulate a remote caller by adding proxy/forwarding headers that the closed + # PR would have treated as "not host" -- here they change nothing. + 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 diff --git a/studio/backend/utils/update_confirm.py b/studio/backend/utils/update_confirm.py new file mode 100644 index 0000000000..1735ba98c0 --- /dev/null +++ b/studio/backend/utils/update_confirm.py @@ -0,0 +1,104 @@ +# 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 (local, desktop, or remote over +SSH/Cloudflare); we do not gate on "same machine" since a headless SSH server never +has a host-local session. A token binds the build it was offered for (``target_tag``) +and is single-use with a short TTL. The store is in-process, matching Studio's +single-process backend; for a multi-worker deploy swap it for an HMAC stateless token. +""" + +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 a burst of confirm calls that are never applied cannot grow +# memory without bound; oldest entries are evicted 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() From cc3259d425402bb7f9b124bcffd922fda778ce5e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 19 Jul 2026 10:40:47 +0000 Subject: [PATCH 2/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/llama.py | 8 ++++---- studio/backend/tests/test_update_contract.py | 11 ++++++----- studio/backend/utils/update_confirm.py | 7 ++----- 3 files changed, 12 insertions(+), 14 deletions(-) diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py index 6d4b2415dc..d3d247eb2b 100644 --- a/studio/backend/routes/llama.py +++ b/studio/backend/routes/llama.py @@ -44,9 +44,7 @@ _REFUSAL_MESSAGES = { "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." - ), + "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." @@ -284,7 +282,9 @@ async def llama_update( return LlamaUpdateActionResponse( started = False, reason = refuse_reason, - message = _REFUSAL_MESSAGES.get(refuse_reason, _REFUSAL_MESSAGES["confirmation_required"]), + message = _REFUSAL_MESSAGES.get( + refuse_reason, _REFUSAL_MESSAGES["confirmation_required"] + ), machine = machine, installed_tag = installed_tag, latest_tag = target_tag, diff --git a/studio/backend/tests/test_update_contract.py b/studio/backend/tests/test_update_contract.py index cf84df143c..81386dc503 100644 --- a/studio/backend/tests/test_update_contract.py +++ b/studio/backend/tests/test_update_contract.py @@ -162,6 +162,7 @@ def _track_start(monkeypatch): # update_confirm token unit tests # --------------------------------------------------------------------------- # + def test_token_roundtrip_single_use(): tok, exp = _uc.mint_confirm_token("b9909") assert isinstance(tok, str) and tok @@ -193,6 +194,7 @@ def test_token_missing_refused(): # 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")) @@ -210,8 +212,8 @@ def test_apply_with_confirmed_true_proceeds(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 out.machine.hostname # which machine + assert out.latest_tag == "b9909" # which version assert calls["n"] == 1 @@ -273,9 +275,7 @@ def test_confirm_endpoint_up_to_date_offers_no_token(monkeypatch): def test_status_reports_machine(): - out = asyncio.run( - rl.llama_update_status(force_refresh = False, current_subject = "operator") - ) + 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" @@ -285,6 +285,7 @@ def test_status_reports_machine(): # HTTP-level: auth gate + wiring + backwards-compatible bodyless POST # --------------------------------------------------------------------------- # + def _client(): from fastapi import FastAPI from fastapi.testclient import TestClient diff --git a/studio/backend/utils/update_confirm.py b/studio/backend/utils/update_confirm.py index 1735ba98c0..342ebb81ac 100644 --- a/studio/backend/utils/update_confirm.py +++ b/studio/backend/utils/update_confirm.py @@ -47,9 +47,7 @@ def _purge_expired_locked(now: float) -> None: def mint_confirm_token( - target_tag: str, - *, - ttl_seconds: int = CONFIRM_TOKEN_TTL_SECONDS, + target_tag: str, *, ttl_seconds: int = CONFIRM_TOKEN_TTL_SECONDS ) -> Tuple[str, str]: """Mint a single-use token bound to ``target_tag``. @@ -70,8 +68,7 @@ def mint_confirm_token( def consume_confirm_token( - token: Optional[str], - current_target_tag: str, + token: Optional[str], current_target_tag: str ) -> Tuple[bool, Optional[str]]: """Validate and consume a token for a swap to ``current_target_tag``. From 3d7abda4343b7522793b89fd2b3860d2035cfa57 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 19 Jul 2026 11:26:40 +0000 Subject: [PATCH 3/8] Studio: fix local-link confirm reason, token build binding, and wire the client Three fixes to the update-confirmation flow. The confirm endpoint checked update_available before local_link, so a local-link tree was reported as up_to_date instead of local_link. The apply route validated the confirm token against a cached update status while start_update re-resolved the latest tag, so a token minted for an older build could apply a newer one; force-refresh so both see the same target. And the frontend still posted the update with no token, which the new gate refuses, so the Update button did nothing; it now confirms to mint a single-use token and applies with it. --- studio/backend/routes/llama.py | 30 +++++++------ studio/backend/tests/test_update_contract.py | 45 +++++++++++++++++++ .../src/hooks/use-llama-update-check.ts | 29 +++++++++++- 3 files changed, 90 insertions(+), 14 deletions(-) diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py index d3d247eb2b..10e89a27f1 100644 --- a/studio/backend/routes/llama.py +++ b/studio/backend/routes/llama.py @@ -209,17 +209,9 @@ async def llama_update_confirm( latest_tag = status.get("latest_tag") size = status.get("update_size_bytes") - 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, - ) - # Update exists but the tree is a --with-llama-cpp-dir local link: describe, do not apply. + # A --with-llama-cpp-dir local link reports update_available=False, so this + # must run before the up_to_date branch below or the local_link reason is + # masked and callers are wrongly told there is no update. if status.get("local_link"): return LlamaUpdateConfirmResponse( update_available = True, @@ -230,6 +222,16 @@ async def llama_update_confirm( 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( @@ -260,8 +262,10 @@ async def llama_update( confirmation, not the caller's location, so a headless SSH server confirms like a local one.""" req = request or LlamaUpdateRequest() machine = _current_machine() - # Cached read: names the host and build for the prompt and binds the token. - status = await asyncio.to_thread(get_update_status) + # Force-refresh so the token is validated against the same build start_update + # will resolve: a stale cache would accept a token minted for an older tag and + # then install a newer one, bypassing the exact-build binding. + 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") diff --git a/studio/backend/tests/test_update_contract.py b/studio/backend/tests/test_update_contract.py index 81386dc503..650ef3ef58 100644 --- a/studio/backend/tests/test_update_contract.py +++ b/studio/backend/tests/test_update_contract.py @@ -274,6 +274,51 @@ def test_confirm_endpoint_up_to_date_offers_no_token(monkeypatch): 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 mask it behind the + # generic up_to_date refusal. + 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): + # A token confirmed for b9909; a newer build b9910 publishes before apply. + # The apply must re-resolve the target and refuse the now-stale token rather + # than install a build the operator never confirmed. + 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 diff --git a/studio/frontend/src/hooks/use-llama-update-check.ts b/studio/frontend/src/hooks/use-llama-update-check.ts index e3a735ffde..ba8a110e1c 100644 --- a/studio/frontend/src/hooks/use-llama-update-check.ts +++ b/studio/frontend/src/hooks/use-llama-update-check.ts @@ -311,8 +311,35 @@ export function useLlamaUpdateCheck({ message?: string | null; job?: unknown; } | null = null; + // Two-step apply: confirm to mint a single-use token bound to the offered + // build, then apply with it. A bare POST is refused with confirmation_required. + let confirmToken: string; try { - const res = await authFetch("/api/llama/update", { method: "POST" }); + const cres = await authFetch("/api/llama/update/confirm", { method: "POST" }); + if (!cres.ok) { + setApplying(false); + return { ok: false, error: `HTTP ${cres.status}` }; + } + const confirm = (await cres.json().catch(() => null)) as { + appliable?: boolean; + reason?: string | null; + confirm_token?: string | null; + } | null; + if (!confirm?.appliable || !confirm.confirm_token) { + setApplying(false); + return { ok: false, error: confirm?.reason ?? "update is not applicable" }; + } + confirmToken = confirm.confirm_token; + } catch (e) { + setApplying(false); + return { ok: false, error: String(e) }; + } + try { + const res = await authFetch("/api/llama/update", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ confirm_token: confirmToken }), + }); if (!res.ok) { setApplying(false); return { ok: false, error: `HTTP ${res.status}` }; From ec9ba169a54f58ceb424c1230b6cd8951b340c8c Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 19 Jul 2026 12:04:08 +0000 Subject: [PATCH 4/8] Studio: require an explicit confirmation before applying a llama.cpp update Complete the confirmation flow. The apply route now passes the confirmed target into start_update, which aborts with stale_target if the freshly resolved latest moved since confirmation, so it never installs a build the operator did not see. The client no longer applies on a single click: it fetches the target, shows a confirmation dialog naming the host and the from and to builds, and only applies on explicit acceptance. The result now carries the machine and versions so the toast names which host was updated. --- studio/backend/routes/llama.py | 4 +- studio/backend/tests/test_llama_route.py | 2 +- studio/backend/tests/test_update_contract.py | 114 +++++++++- studio/backend/utils/llama_cpp_update.py | 26 ++- .../src/components/llama-update-banner.tsx | 27 ++- .../llama-update-confirm-dialog.tsx | 99 +++++++++ .../src/features/chat/chat-settings-sheet.tsx | 19 +- .../src/hooks/use-llama-update-check.ts | 208 ++++++++++++------ 8 files changed, 413 insertions(+), 86 deletions(-) create mode 100644 studio/frontend/src/components/llama-update-confirm-dialog.tsx diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py index 10e89a27f1..72f369ba39 100644 --- a/studio/backend/routes/llama.py +++ b/studio/backend/routes/llama.py @@ -296,7 +296,9 @@ async def llama_update( job = LlamaUpdateJob(**status.get("job", {})), ) - action = await asyncio.to_thread(start_update) + # 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, diff --git a/studio/backend/tests/test_llama_route.py b/studio/backend/tests/test_llama_route.py index ac3ec01797..a3e3a76c7a 100644 --- a/studio/backend/tests/test_llama_route.py +++ b/studio/backend/tests/test_llama_route.py @@ -123,7 +123,7 @@ 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"}} diff --git a/studio/backend/tests/test_update_contract.py b/studio/backend/tests/test_update_contract.py index 650ef3ef58..6b5009fb79 100644 --- a/studio/backend/tests/test_update_contract.py +++ b/studio/backend/tests/test_update_contract.py @@ -88,7 +88,7 @@ def _install_stubs(): "job": {"state": "idle"}, } - def _default_start(): + def _default_start(expected_tag = None): return {"started": True, "reason": None, "job": {"state": "running"}} lcu_mod.get_update_status = _default_status @@ -148,10 +148,11 @@ def _reset(monkeypatch): def _track_start(monkeypatch): - calls = {"n": 0} + calls = {"n": 0, "expected_tags": []} - def _start(): + 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) @@ -215,6 +216,9 @@ def test_apply_with_confirmed_true_proceeds(monkeypatch): 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): @@ -326,6 +330,110 @@ def test_status_reports_machine(): assert out.latest_tag == "b9909" +# --------------------------------------------------------------------------- # +# start_update pins the confirmed target: a release that publishes in the gap +# 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 to exercise start_update's + confirmed-target guard directly. The module-level stubs shadow only + utils.llama_cpp_update, so this loads it under a private name (leaving that + stub in place for the handler tests) and its two real deps under their + production names.""" + + 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 # --------------------------------------------------------------------------- # diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 31dbda63ea..fe13f16da5 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -622,9 +622,15 @@ 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 already confirmed. The updater + re-resolves "latest" itself, so a release published in the gap after + confirmation would move the target; when the freshly-resolved latest differs + from ``expected_tag`` this aborts rather than swapping to a build the caller + never confirmed. Left None, it behaves as before (no target pinning).""" 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) @@ -678,6 +684,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 @@ -708,6 +715,21 @@ def start_update() -> dict: # 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 + resolved_tag = src.get("latest_tag") + + # Install exactly the build the caller confirmed. A release published in the + # gap since confirmation moves the freshly-resolved latest above, so abort + # rather than swap to a build the caller never saw or confirmed. + 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/frontend/src/components/llama-update-banner.tsx b/studio/frontend/src/components/llama-update-banner.tsx index 3db15ffe30..05c0f94cd8 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 ? (
) : null; + + // The confirm dialog renders unconditionally (inert until Update is clicked) + // so it survives the banner hiding once the swap begins. + return ( + <> + {dialog} + {banner} + + ); } diff --git a/studio/frontend/src/components/llama-update-confirm-dialog.tsx b/studio/frontend/src/components/llama-update-confirm-dialog.tsx new file mode 100644 index 0000000000..caee3f1529 --- /dev/null +++ b/studio/frontend/src/components/llama-update-confirm-dialog.tsx @@ -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 + +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; +import type { LlamaApplyTarget } from "@/hooks/use-llama-update-check"; +import { useCallback, useRef, useState, type ReactElement } from "react"; + +interface LlamaUpdateConfirmGate { + /** + * Open the confirmation prompt for `target` and resolve true only on an + * explicit accept, false on cancel/dismiss. Pass this as the `apply()` gate so + * the destructive host-binary swap never runs without a visible, accepted + * build + target host. + */ + requestConfirm: (target: LlamaApplyTarget) => Promise; + /** Render once in the tree; inert until requestConfirm opens it. */ + dialog: ReactElement; +} + +/** + * Reusable accept/cancel gate for the llama.cpp host-binary swap. Shows the + * exact build (from -> to) and the machine the swap targets, and only resolves + * true when the user explicitly accepts. + */ +export function useLlamaUpdateConfirmGate(): LlamaUpdateConfirmGate { + const [target, setTarget] = useState(null); + const resolveRef = useRef<((accepted: boolean) => void) | null>(null); + + const decide = useCallback((accepted: boolean) => { + const resolve = resolveRef.current; + resolveRef.current = null; + setTarget(null); + resolve?.(accepted); + }, []); + + const requestConfirm = useCallback((next: LlamaApplyTarget) => { + // A still-open prior prompt (e.g. a double click) resolves as declined. + resolveRef.current?.(false); + setTarget(next); + return new Promise((resolve) => { + resolveRef.current = resolve; + }); + }, []); + + const machineLabel = target?.machine?.hostname?.trim() || "this machine"; + + const dialog = ( + { + if (!open) decide(false); + }} + > + + + Update llama.cpp? + + This downloads and swaps the llama.cpp binary on {machineLabel}. It + replaces the running build, so only continue if you started this + update. + + +
+

+ {target?.fromTag ?? "unknown"} →{" "} + + {target?.toTag ?? "the latest build"} + +

+ {target?.machine?.hostname && ( +

+ {target.machine.hostname} + {target.machine.platform ? ` ยท ${target.machine.platform}` : ""} +

+ )} +
+ + decide(false)}> + Cancel + + decide(true)}> + Update + + + + + ); + + return { requestConfirm, dialog }; +} diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index cedd298ecf..f63d4595d2 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -57,6 +57,7 @@ import { InfoHint } from "@/components/ui/info-hint"; import { Tooltip, TooltipContent } from "@/components/ui/tooltip"; import { useIsMobile } from "@/hooks/use-mobile"; import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check"; +import { useLlamaUpdateConfirmGate } from "@/components/llama-update-confirm-dialog"; import { cn } from "@/lib/utils"; import { ArrowTurnBackwardIcon, @@ -570,19 +571,28 @@ export function ChatSettingsPanel({ enabled: mtpUpdatable, onReloadRequired: resyncInferenceStatusAfterServerModelChange, }); + // Explicit accept/cancel prompt naming the exact build + host before the swap. + const { + requestConfirm: requestLlamaUpdateConfirm, + dialog: llamaUpdateConfirmDialog, + } = useLlamaUpdateConfirmGate(); const handleMtpUpdate = useCallback(async () => { - const result = await applyLlamaUpdate(); + const result = await applyLlamaUpdate(requestLlamaUpdateConfirm); if (result.ok) { + const host = result.machine?.hostname; + const where = host ? ` on ${host}` : ""; + const fromTag = result.fromTag ?? llamaUpdateStatus?.installed_tag ?? "unknown"; + const toTag = result.tag ?? llamaUpdateStatus?.latest_tag ?? "the latest build"; const reloadHint = result.reloadRequired ? " Reload your model to enable MTP." : ""; toast.success( - `llama.cpp updated to ${result.tag ?? "the latest build"}.${reloadHint}`, + `llama.cpp${where} updated ${fromTag} to ${toTag}.${reloadHint}`, ); - } else { + } else if (result.error !== "canceled") { toast.error(`llama.cpp update failed: ${result.error ?? "unknown error"}`); } - }, [applyLlamaUpdate]); + }, [applyLlamaUpdate, requestLlamaUpdateConfirm, llamaUpdateStatus]); const specDraftNMax = useChatRuntimeStore((s) => s.specDraftNMax); const setSpecDraftNMax = useChatRuntimeStore((s) => s.setSpecDraftNMax); const loadedSpecDraftNMax = useChatRuntimeStore( @@ -1887,6 +1897,7 @@ export function ChatSettingsPanel({ + {llamaUpdateConfirmDialog} ); diff --git a/studio/frontend/src/hooks/use-llama-update-check.ts b/studio/frontend/src/hooks/use-llama-update-check.ts index ba8a110e1c..dc41a05621 100644 --- a/studio/frontend/src/hooks/use-llama-update-check.ts +++ b/studio/frontend/src/hooks/use-llama-update-check.ts @@ -120,13 +120,41 @@ interface UseLlamaUpdateCheckOptions { onReloadRequired?: () => void; } +export interface LlamaMachine { + hostname: string; + platform: string; +} + +/** The pending swap a confirmation prompt describes: the exact build (from -> + * to), the host it targets, and the single-use token that applies it. */ +export interface LlamaApplyTarget { + token: string; + machine: LlamaMachine | null; + fromTag: string | null; + toTag: string | null; +} + export interface LlamaApplyResult { ok: boolean; + // Post-swap build tag (the "to" version). tag?: string | null; + // Build the swap replaced (the "from" version). + fromTag?: string | null; + // Host the swap ran on, so the result can name which machine changed. + machine?: LlamaMachine | null; reloadRequired?: boolean | null; error?: string | null; } +function parseMachine(value: unknown): LlamaMachine | null { + if (!value || typeof value !== "object") return null; + const m = value as Record; + return { + hostname: typeof m.hostname === "string" ? m.hostname : "", + platform: typeof m.platform === "string" ? m.platform : "", + }; +} + /** Tracks llama.cpp update visibility and apply progress. */ export function useLlamaUpdateCheck({ enabled = true, @@ -301,81 +329,121 @@ export function useLlamaUpdateCheck({ }, SNOOZE_DELAY_MS); }, [surfaceIfAvailable]); - const apply = useCallback(async (): Promise => { - if (applying) return { ok: false, error: "already running" }; - setApplying(true); - setVisible(true); - let action: { - started?: boolean; - reason?: string | null; - message?: string | null; - job?: unknown; - } | null = null; - // Two-step apply: confirm to mint a single-use token bound to the offered - // build, then apply with it. A bare POST is refused with confirmation_required. - let confirmToken: string; - try { - const cres = await authFetch("/api/llama/update/confirm", { method: "POST" }); - if (!cres.ok) { - setApplying(false); - return { ok: false, error: `HTTP ${cres.status}` }; - } - const confirm = (await cres.json().catch(() => null)) as { - appliable?: boolean; - reason?: string | null; - confirm_token?: string | null; - } | null; - if (!confirm?.appliable || !confirm.confirm_token) { - setApplying(false); - return { ok: false, error: confirm?.reason ?? "update is not applicable" }; - } - confirmToken = confirm.confirm_token; - } catch (e) { - setApplying(false); - return { ok: false, error: String(e) }; - } - try { - const res = await authFetch("/api/llama/update", { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ confirm_token: confirmToken }), - }); - if (!res.ok) { - setApplying(false); - return { ok: false, error: `HTTP ${res.status}` }; - } + const apply = useCallback( + async ( + confirm: (target: LlamaApplyTarget) => boolean | Promise, + ): Promise => { + if (applying) return { ok: false, error: "already running" }; + + // Step 1: describe the pending swap. /confirm force-refreshes and mints a + // single-use token bound to the offered build, and reports the host + the + // from/to tags. No install starts here, so `applying` stays false while the + // confirmation prompt is open. + let target: LlamaApplyTarget; try { - action = await res.json(); - } catch { - action = null; + const cres = await authFetch("/api/llama/update/confirm", { + method: "POST", + }); + if (!cres.ok) return { ok: false, error: `HTTP ${cres.status}` }; + const confirmResp = (await cres.json().catch(() => null)) as { + appliable?: boolean; + reason?: string | null; + confirm_token?: string | null; + machine?: unknown; + installed_tag?: string | null; + latest_tag?: string | null; + } | null; + if (!confirmResp?.appliable || !confirmResp.confirm_token) { + return { + ok: false, + error: confirmResp?.reason ?? "update is not applicable", + }; + } + target = { + token: confirmResp.confirm_token, + machine: parseMachine(confirmResp.machine), + fromTag: + typeof confirmResp.installed_tag === "string" + ? confirmResp.installed_tag + : null, + toTag: + typeof confirmResp.latest_tag === "string" + ? confirmResp.latest_tag + : null, + }; + } catch (e) { + return { ok: false, error: String(e) }; } - } catch (e) { - setApplying(false); - return { ok: false, error: String(e) }; - } - // Non-started jobs stay idle; already_running is tracked below. - if ( - action && - action.started === false && - action.reason !== "already_running" - ) { - // A stale banner's click can land after another tab already applied the - // update (e.g. "up_to_date"): the response still carries that tab's - // completed job, so process reload_required here too, not just from the - // poll path -- otherwise this rejection silently drops it. - notifyReloadIfNeeded(parseJob(action.job)); - setApplying(false); + // Step 2: require an explicit user confirmation of the exact build + host + // before the destructive binary swap. Declining aborts, untouched. + const accepted = await confirm(target); + if (!accepted) return { ok: false, error: "canceled" }; + + // Step 3: apply with the confirmed single-use token. + setApplying(true); + setVisible(true); + let action: { + started?: boolean; + reason?: string | null; + message?: string | null; + job?: unknown; + } | null = null; + try { + const res = await authFetch("/api/llama/update", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ confirm_token: target.token }), + }); + if (!res.ok) { + setApplying(false); + return { ok: false, error: `HTTP ${res.status}` }; + } + try { + action = await res.json(); + } catch { + action = null; + } + } catch (e) { + setApplying(false); + return { ok: false, error: String(e) }; + } + + // Non-started jobs stay idle; already_running is tracked below. + if ( + action && + action.started === false && + action.reason !== "already_running" + ) { + // A stale banner's click can land after another tab already applied the + // update (e.g. "up_to_date"): the response still carries that tab's + // completed job, so process reload_required here too, not just from the + // poll path -- otherwise this rejection silently drops it. + notifyReloadIfNeeded(parseJob(action.job)); + setApplying(false); + return { + ok: false, + error: action.message ?? action.reason ?? "update was not started", + machine: target.machine, + fromTag: target.fromTag, + tag: target.toTag, + }; + } + + // Surface the confirmed host + from/to build alongside the job's post-swap + // tag, so the result can report what changed and where. + const polled = await new Promise((resolve) => + startJobPoll(resolve), + ); return { - ok: false, - error: action.message ?? action.reason ?? "update was not started", + ...polled, + machine: target.machine, + fromTag: target.fromTag, + tag: polled.tag ?? target.toTag, }; - } - - return await new Promise((resolve) => - startJobPoll(resolve), - ); - }, [applying, startJobPoll, notifyReloadIfNeeded]); + }, + [applying, startJobPoll, notifyReloadIfNeeded], + ); return { status: enabled ? status : null, From ff074472fde2779a458c405c562f1dc52bfa8969 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 19 Jul 2026 12:05:53 +0000 Subject: [PATCH 5/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_update_contract.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/studio/backend/tests/test_update_contract.py b/studio/backend/tests/test_update_contract.py index 6b5009fb79..d6fd0ffe2c 100644 --- a/studio/backend/tests/test_update_contract.py +++ b/studio/backend/tests/test_update_contract.py @@ -344,9 +344,7 @@ def _load_real_llama_cpp_update(): production names.""" def _load(mod_name: str, filename: str): - spec = importlib.util.spec_from_file_location( - mod_name, str(_BACKEND / "utils" / filename) - ) + 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) @@ -375,9 +373,7 @@ def _stub_marker_update(monkeypatch, lcu, *, resolved_latest: str): "asset": "cuda-x64.zip", }, ) - monkeypatch.setattr( - lcu, "_installer_script", lambda: Path("/fake/install_llama_prebuilt.py") - ) + 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, @@ -416,7 +412,14 @@ def test_start_update_proceeds_when_resolved_latest_matches_confirmed(monkeypatc spawned: dict = {} class _CapturingThread: - def __init__(self, *, target = None, args = (), name = None, daemon = None): + def __init__( + self, + *, + target = None, + args = (), + name = None, + daemon = None, + ): spawned["args"] = args def start(self): From 251bec8e583ed6d35ec665f84fbfc1d0680ec194 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 19 Jul 2026 12:50:27 +0000 Subject: [PATCH 6/8] Pin the source-build llama.cpp install to the resolved release tag The source-build path left the installer unpinned, so a release published between the resolve here and the installer's own latest re-resolve could install a build the user never confirmed. Pin to the host-aware resolver's release_tag (any macOS walk-back already applied), which cannot disable a needed walk-back and arms the post-install tag check. --- studio/backend/tests/test_llama_cpp_update.py | 61 ++++++++++++++++++- studio/backend/utils/llama_cpp_update.py | 12 +++- 2 files changed, 67 insertions(+), 6 deletions(-) diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index 83ea07a066..86614606a1 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -393,9 +393,64 @@ 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 host-aware resolver already picked, so a release + # published between resolve and the installer's own "latest" re-resolve + # cannot swap in an unconfirmed build (matches the marker path's pin). + 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 must pin the installer to the release the host-aware + # resolver picked (res release_tag), not the display tag. For a fork-wrapper + # release the two differ ("v1.0" release vs "b9457" display); only the real + # release tag is a valid --published-release-tag and post-install anchor. + # 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/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index fe13f16da5..dd80d8ef7e 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -712,9 +712,15 @@ def start_update(expected_tag: Optional[str] = None) -> dict: repo = (res or {}).get("repo") or DEFAULT_PUBLISHED_REPO from_tag = None asset = (res or {}).get("asset") - # 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 + # Pin the installer to the exact release the host-aware resolver already + # picked (res release_tag, which is the real GitHub release tag and has + # any macOS walk-back already applied), so a release published between + # this resolve and the installer's own "latest" re-resolve cannot swap in + # an unconfirmed build. The pinned tag is host-compatible by construction, + # so pinning does not disable a needed walk-back (unlike the marker path); + # it also arms the post-install tag check in _run_update. Falls back to + # 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") # Install exactly the build the caller confirmed. A release published in the From c99b89b9d242284098454a66f7ed1d8268732c19 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 19 Jul 2026 18:52:59 +0000 Subject: [PATCH 7/8] studio: tighten comments in the llama.cpp update confirmation flow --- studio/backend/routes/llama.py | 31 +++++------ studio/backend/tests/test_llama_cpp_update.py | 14 +++-- studio/backend/tests/test_update_contract.py | 52 +++++++------------ studio/backend/utils/llama_cpp_update.py | 28 +++++----- studio/backend/utils/update_confirm.py | 12 ++--- .../src/components/llama-update-banner.tsx | 4 +- .../llama-update-confirm-dialog.tsx | 12 ++--- .../src/hooks/use-llama-update-check.ts | 11 ++-- 8 files changed, 66 insertions(+), 98 deletions(-) diff --git a/studio/backend/routes/llama.py b/studio/backend/routes/llama.py index 72f369ba39..1f3f0cc039 100644 --- a/studio/backend/routes/llama.py +++ b/studio/backend/routes/llama.py @@ -53,7 +53,7 @@ _REFUSAL_MESSAGES = { class UpdateMachine(BaseModel): - """Host a swap targets, so a remote operator sees which machine would change.""" + """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( @@ -112,8 +112,7 @@ class LlamaUpdateStatusResponse(BaseModel): class LlamaUpdateRequest(BaseModel): - """Body for POST /update. Empty body means no confirmation, so the swap is - refused (safe default) and a stale banner or replay cannot swap the binary.""" + """Body for POST /update. Empty body = no confirmation, so the swap is refused.""" confirm_token: Optional[str] = Field( None, @@ -165,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 @@ -199,9 +197,9 @@ async def llama_update_status( 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 it can - be applied, mint a single-use token bound to the offered build. Available to - any authenticated operator; confirmation, not location, is the gate.""" + """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() @@ -209,9 +207,8 @@ async def llama_update_confirm( 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 below or the local_link reason is - # masked and callers are wrongly told there is no update. + # 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, @@ -256,15 +253,13 @@ async def llama_update( """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 (preferred: replay-safe, - bound to the build) or send ``confirmed=true`` (non-interactive callers). With - neither, the swap is refused and the binary is left untouched. The gate is the - confirmation, not the caller's location, so a headless SSH server confirms like a local one.""" + 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 would accept a token minted for an older tag and - # then install a newer one, bypassing the exact-build binding. + # 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") diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index 86614606a1..43a4eb5f29 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -393,19 +393,17 @@ 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 - # Pin to the release the host-aware resolver already picked, so a release - # published between resolve and the installer's own "latest" re-resolve - # cannot swap in an unconfirmed build (matches the marker path's pin). + # 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 must pin the installer to the release the host-aware - # resolver picked (res release_tag), not the display tag. For a fork-wrapper - # release the two differ ("v1.0" release vs "b9457" display); only the real - # release tag is a valid --published-release-tag and post-install anchor. - # Confirming the displayed tag must still proceed and pin the real release. + # 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) diff --git a/studio/backend/tests/test_update_contract.py b/studio/backend/tests/test_update_contract.py index d6fd0ffe2c..e9562bb5f6 100644 --- a/studio/backend/tests/test_update_contract.py +++ b/studio/backend/tests/test_update_contract.py @@ -10,9 +10,8 @@ Properties covered: - unauthenticated -> refused (401), NO swap - stale / expired / replayed token -> refused, NO swap -routes/llama.py is loaded standalone with stubbed auth / loggers / llama_cpp_update -and the real utils.update_confirm, so no heavy backend deps are needed. Both handler -calls and a real FastAPI TestClient (HTTP + auth gate) are exercised. +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 @@ -33,10 +32,8 @@ _BACKEND = _HERE.parent def _install_stubs(): - """Register stub packages so routes/llama.py imports cleanly, plus the real - update_confirm module under its production name.""" - # auth.authentication.get_current_subject -> a real FastAPI dependency that - # 401s without a valid bearer, so the HTTP tests can prove the auth gate. + """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: @@ -159,9 +156,7 @@ def _track_start(monkeypatch): return calls -# --------------------------------------------------------------------------- # # update_confirm token unit tests -# --------------------------------------------------------------------------- # def test_token_roundtrip_single_use(): @@ -191,9 +186,7 @@ def test_token_missing_refused(): 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): @@ -279,9 +272,8 @@ def test_confirm_endpoint_up_to_date_offers_no_token(monkeypatch): 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 mask it behind the - # generic up_to_date refusal. + # 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", @@ -300,9 +292,8 @@ def test_confirm_endpoint_reports_local_link_not_up_to_date(monkeypatch): def test_apply_revalidates_token_against_refreshed_target(monkeypatch): - # A token confirmed for b9909; a newer build b9910 publishes before apply. - # The apply must re-resolve the target and refuse the now-stale token rather - # than install a build the operator never confirmed. + # 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, @@ -330,18 +321,14 @@ def test_status_reports_machine(): assert out.latest_tag == "b9909" -# --------------------------------------------------------------------------- # -# start_update pins the confirmed target: a release that publishes in the gap -# after confirmation must not be installed in place of the confirmed build. -# --------------------------------------------------------------------------- # +# 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 to exercise start_update's - confirmed-target guard directly. The module-level stubs shadow only - utils.llama_cpp_update, so this loads it under a private name (leaving that - stub in place for the handler tests) and its two real deps under their - production names.""" + """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)) @@ -437,9 +424,7 @@ def test_start_update_proceeds_when_resolved_latest_matches_confirmed(monkeypatc lcu._reset_job_for_tests() -# --------------------------------------------------------------------------- # # HTTP-level: auth gate + wiring + backwards-compatible bodyless POST -# --------------------------------------------------------------------------- # def _client(): @@ -461,8 +446,8 @@ def test_http_unauthenticated_update_is_refused(monkeypatch): def test_http_authenticated_bodyless_post_refused_no_swap(monkeypatch): - # Backwards compat: an OLD frontend posts /update with no body. That must be - # safe -- refused (confirmation_required), NEVER a silent swap. + # 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"}) @@ -492,13 +477,12 @@ def test_http_two_step_confirm_then_apply(monkeypatch): def test_http_local_and_remote_behave_identically(monkeypatch): - # There is no same-machine axis: the contract depends only on auth + confirm, - # so a "local" and a "remote" authenticated caller get identical outcomes. + # 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"} - # Simulate a remote caller by adding proxy/forwarding headers that the closed - # PR would have treated as "not host" -- here they change nothing. + # 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}) diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index dd80d8ef7e..ee1d415272 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -626,11 +626,10 @@ 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. - ``expected_tag`` is the build the caller already confirmed. The updater - re-resolves "latest" itself, so a release published in the gap after - confirmation would move the target; when the freshly-resolved latest differs - from ``expected_tag`` this aborts rather than swapping to a build the caller - never confirmed. Left None, it behaves as before (no target pinning).""" + ``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) @@ -712,20 +711,17 @@ def start_update(expected_tag: Optional[str] = None) -> dict: repo = (res or {}).get("repo") or DEFAULT_PUBLISHED_REPO from_tag = None asset = (res or {}).get("asset") - # Pin the installer to the exact release the host-aware resolver already - # picked (res release_tag, which is the real GitHub release tag and has - # any macOS walk-back already applied), so a release published between - # this resolve and the installer's own "latest" re-resolve cannot swap in - # an unconfirmed build. The pinned tag is host-compatible by construction, - # so pinning does not disable a needed walk-back (unlike the marker path); - # it also arms the post-install tag check in _run_update. Falls back to - # unpinned only if the resolver reported no release tag. + # 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") - # Install exactly the build the caller confirmed. A release published in the - # gap since confirmation moves the freshly-resolved latest above, so abort - # rather than swap to a build the caller never saw or confirmed. + # 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, diff --git a/studio/backend/utils/update_confirm.py b/studio/backend/utils/update_confirm.py index 342ebb81ac..3bb7493945 100644 --- a/studio/backend/utils/update_confirm.py +++ b/studio/backend/utils/update_confirm.py @@ -5,11 +5,10 @@ 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 (local, desktop, or remote over -SSH/Cloudflare); we do not gate on "same machine" since a headless SSH server never -has a host-local session. A token binds the build it was offered for (``target_tag``) -and is single-use with a short TTL. The store is in-process, matching Studio's -single-process backend; for a multi-worker deploy swap it for an HMAC stateless token. +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 @@ -22,8 +21,7 @@ from typing import Optional, Tuple # How long a freshly minted confirmation token stays valid. CONFIRM_TOKEN_TTL_SECONDS = 300 -# Cap the store so a burst of confirm calls that are never applied cannot grow -# memory without bound; oldest entries are evicted first. +# Cap the store so unapplied confirm calls can't grow memory unbounded; oldest first. _MAX_TOKENS = 64 _lock = threading.Lock() diff --git a/studio/frontend/src/components/llama-update-banner.tsx b/studio/frontend/src/components/llama-update-banner.tsx index 05c0f94cd8..75e5c38efb 100644 --- a/studio/frontend/src/components/llama-update-banner.tsx +++ b/studio/frontend/src/components/llama-update-banner.tsx @@ -238,8 +238,8 @@ export function LlamaUpdateBanner({ ) : null; - // The confirm dialog renders unconditionally (inert until Update is clicked) - // so it survives the banner hiding once the swap begins. + // Render the confirm dialog unconditionally (inert until Update) so it survives + // the banner hiding once the swap begins. return ( <> {dialog} diff --git a/studio/frontend/src/components/llama-update-confirm-dialog.tsx b/studio/frontend/src/components/llama-update-confirm-dialog.tsx index caee3f1529..ffbf6dd416 100644 --- a/studio/frontend/src/components/llama-update-confirm-dialog.tsx +++ b/studio/frontend/src/components/llama-update-confirm-dialog.tsx @@ -16,10 +16,9 @@ import { useCallback, useRef, useState, type ReactElement } from "react"; interface LlamaUpdateConfirmGate { /** - * Open the confirmation prompt for `target` and resolve true only on an - * explicit accept, false on cancel/dismiss. Pass this as the `apply()` gate so - * the destructive host-binary swap never runs without a visible, accepted - * build + target host. + * Open the prompt for `target`; resolve true only on explicit accept, false on + * cancel/dismiss. Pass as the `apply()` gate so the destructive swap never runs + * without a visible, accepted build + target host. */ requestConfirm: (target: LlamaApplyTarget) => Promise; /** Render once in the tree; inert until requestConfirm opens it. */ @@ -27,9 +26,8 @@ interface LlamaUpdateConfirmGate { } /** - * Reusable accept/cancel gate for the llama.cpp host-binary swap. Shows the - * exact build (from -> to) and the machine the swap targets, and only resolves - * true when the user explicitly accepts. + * Accept/cancel gate for the llama.cpp host-binary swap: shows the exact build + * (from -> to) and target machine, resolving true only on explicit accept. */ export function useLlamaUpdateConfirmGate(): LlamaUpdateConfirmGate { const [target, setTarget] = useState(null); diff --git a/studio/frontend/src/hooks/use-llama-update-check.ts b/studio/frontend/src/hooks/use-llama-update-check.ts index dc41a05621..b9bffb68c1 100644 --- a/studio/frontend/src/hooks/use-llama-update-check.ts +++ b/studio/frontend/src/hooks/use-llama-update-check.ts @@ -125,8 +125,8 @@ export interface LlamaMachine { platform: string; } -/** The pending swap a confirmation prompt describes: the exact build (from -> - * to), the host it targets, and the single-use token that applies it. */ +/** The pending swap a confirm prompt describes: build (from -> to), target host, + * and the single-use token that applies it. */ export interface LlamaApplyTarget { token: string; machine: LlamaMachine | null; @@ -335,10 +335,9 @@ export function useLlamaUpdateCheck({ ): Promise => { if (applying) return { ok: false, error: "already running" }; - // Step 1: describe the pending swap. /confirm force-refreshes and mints a - // single-use token bound to the offered build, and reports the host + the - // from/to tags. No install starts here, so `applying` stays false while the - // confirmation prompt is open. + // Step 1: describe the pending swap. /confirm force-refreshes, mints a + // single-use token bound to the offered build, and reports host + from/to tags. + // No install starts here, so `applying` stays false while the prompt is open. let target: LlamaApplyTarget; try { const cres = await authFetch("/api/llama/update/confirm", { From 4b6a6c774816c6b0454b38136bc6e1e0b7dd31d4 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 07:55:46 +0000 Subject: [PATCH 8/8] Stop the update-contract test stubs from leaking into the shared suite test_update_contract installs stub auth/loggers/utils modules into sys.modules at collection time so routes/llama.py loads standalone, but it never removed them. In the one-process backend suite that poisoned every later test: a bare auth package broke 'from auth import storage' and 'from auth.authentication import create_access_token', and the NopLogger stub (no error method) broke code paths that log errors. Snapshot the affected sys.modules entries before installing the stubs and restore them right after the route module has bound its imports; the loaded route keeps the stub references it captured, so the contract tests are unchanged. Add a regression test asserting no installed stub is left in sys.modules. Mirrors the pop-after-load hygiene already in test_llama_route. --- studio/backend/tests/test_update_contract.py | 36 ++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/studio/backend/tests/test_update_contract.py b/studio/backend/tests/test_update_contract.py index e9562bb5f6..aaa2795aaf 100644 --- a/studio/backend/tests/test_update_contract.py +++ b/studio/backend/tests/test_update_contract.py @@ -30,6 +30,19 @@ 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.""" @@ -67,6 +80,7 @@ def _install_stubs(): 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") @@ -105,10 +119,16 @@ def _install_stubs(): ("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() @@ -124,6 +144,13 @@ def _load_route(): 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): @@ -489,3 +516,12 @@ def test_http_local_and_remote_behave_identically(monkeypatch): 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"