From c99b89b9d242284098454a66f7ed1d8268732c19 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 19 Jul 2026 18:52:59 +0000 Subject: [PATCH] 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", {