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()