studio: tighten comments in the llama.cpp update confirmation flow
This commit is contained in:
danielhanchen
parent
251bec8e58
commit
c99b89b9d2
8 changed files with 66 additions and 98 deletions
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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})
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -238,8 +238,8 @@ export function LlamaUpdateBanner({
|
|||
</div>
|
||||
) : 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}
|
||||
|
|
|
|||
|
|
@ -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<boolean>;
|
||||
/** 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<LlamaApplyTarget | null>(null);
|
||||
|
|
|
|||
|
|
@ -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<LlamaApplyResult> => {
|
||||
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", {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue