diff --git a/TECHNICAL_REPORTS/1872-safe-webui-library-20260913.en.md b/TECHNICAL_REPORTS/1872-safe-webui-library-20260913.en.md new file mode 100644 index 000000000..ed66bdb75 --- /dev/null +++ b/TECHNICAL_REPORTS/1872-safe-webui-library-20260913.en.md @@ -0,0 +1,163 @@ +# Technical Report: PR #1872 — Safe WebUI model library operations + +**Date**: 2026-09-13 +**Status**: Open PR / status:review; source `84f1aeb2` passed full, strict real-model and hosted acceptance; pending central merge +**Languages**: Rust, JSON/OpenAPI fixtures +**Risk Level**: High +**Implementation snapshot**: integrated baseline `3cb4817d` plus the owned-drain revision correction described below; earlier repair checkpoints `d688ea4f` and `b7555005` + +## Executive Summary + +[PR #1872](https://github.com/lablup/mlxcel/pull/1872) implements the issue #1841 backend for WebUI model-library mutations in epic #1834. It adds authenticated WebUI download and managed-cache removal adapters on top of the existing router pool and lifecycle coordinator, while keeping startup state shared with #1838 and preserving compatibility routes. + +The highest-risk parts are disk mutation and Hub access. The implementation keeps WebUI downloads anonymous, bounded, and fd-relative until atomic publication, and confines deletion to managed cache snapshots through a descriptor-anchored private quarantine rather than deleting user model roots or preset/external paths. + +## 1. Problem Statement + +The WebUI needs to download and remove models without turning the browser layer into a second model registry or a privileged filesystem shell. Existing router compatibility endpoints already had download/remove behavior, but they were not expressed as common WebUI operations with idempotency, bounded admission, progress, cancellation, and terminal results. + +Disk deletion is irreversible and high blast radius. A model unload is not the same action as deleting a checkpoint, and deletion must not follow symlinks, race active writers/providers, or apply to models-dir/preset/external artifacts. WebUI downloads also must not silently consume ambient HuggingFace credentials just because the server host has CLI tokens. + +## 2. Technical Decisions + +### Shared coordinator authority instead of new state + +`webui::library::routes()` mounts under `/ui-api/v1` and uses `RouterServerState` directly. Download and removal actions enter the same `RouterPool` and `LifecycleCoordinator` as other WebUI operations, so idempotency, bounded active operations, cancellation state, SSE operation events, and catalog rescan behavior have one authority. + +Compatibility `POST /models` now enqueues through the same download primitive, while preserving its wire schema. The route no longer performs an unbounded synchronous metadata request before admission; duplicate and queue decisions happen first. + +### Anonymous immutable downloads + +WebUI downloads accept a public HuggingFace repo id plus optional revision, not arbitrary URLs or local paths. Anonymous WebUI metadata uses direct reqwest requests with a response byte cap instead of `hf-hub` builders that read ambient token files; CLI download behavior still uses the environment token mode. + +The worker resolves metadata once, pins transfers to the resolved SHA, bounds manifest siblings, selected files, and filename length, computes known byte totals from HEAD when available, and validates actual bytes written against known sizes before publication. The terminal operation result carries the resolved revision; the shared `RevisionRef` schema remains nullable for request compatibility. + +### Descriptor-anchored staging and cancellation boundary + +On Unix, router-managed downloads write into a unique private `.mlxcel-staging` directory opened by fd. Remote paths are created relative to directory descriptors with no symlink traversal, and publication uses atomic no-replace rename after the coordinator's `begin_publish` hook seals cancellation. Cancellation accepted before writer exit becomes terminal only after the worker observes it; a late cancel after publication linearization is refused and the completed snapshot can publish. + +The completeness policy requires `config.json`, at least one selected safetensors weight file, known-size byte agreement where metadata/headers provide a size, and LFS SHA-256 metadata for every safetensors file on the anonymous WebUI path. The stream is hashed before fd-relative rename, so same-size corrupted weights fail before publication. + +### Managed-cache deletion only + +`POST /ui-api/v1/model-removals` keeps the frozen request shape of `model_id`, `expected_revision`, and `idempotency_key`; user-facing confirmation is owned by the WebUI page work, and the backend endpoint itself is the explicit deletion intent. The pool requires a cache-sourced, current, idle model entry and marks it operation-busy before spawning deletion, which prevents rescan from rebuilding a stale entry while removal owns the catalog slot. + +Deletion opens the configured cache root, owner, target, and private quarantine with `O_NOFOLLOW`, verifies identities at the final pre-rename boundary and after quarantine rename, then recursively unlinks relative to held fds. Nested directory deletion carries expected dev/ino through open and final rmdir checks. If identity changes, deletion fails closed and leaves the quarantined snapshot for inspection; a same-UID process can still cause denial of service or a failed quarantine cleanup race, so this is not a claim of absolute hostile-same-user isolation. + +## 3. Review Refinements + +Parent pre-review found and the implementation addressed these high-risk items before publication: + +- Cancellation after the last transfer chunk but before publish is now linearized by `begin_publish_operation`, with a deterministic regression that late cancellation is refused and the operation succeeds. +- Metadata now requests `?blobs=true`; safetensors without LFS SHA-256 are rejected, streams are hashed before fd-relative rename, known metadata/HEAD sizes are checked against actual bytes, and metadata/file counts/names are bounded before selected downloads are admitted. +- Anonymous WebUI metadata no longer constructs `hf-hub` API builders that read saved token files; tests seed token files and env vars and assert no Authorization header on metadata, HEAD, or GET fake requests. +- Removal now reserves lifecycle state with `operation_busy`, rechecks after acquiring the per-entry operation guard, and rejects equality plus ancestor/descendant overlaps from configured models-dir or preset/non-cache entries so cache deletion cannot remove a user/preset-owned nested snapshot. +- Parent-swap and nested stat/open deletion races have deterministic tests and fail without deleting the wrong tree. +- New hand-written modules are below the 500-line cap after extracting `anchored_delete.rs` and adjacent test modules. + +## 4. Validation Record + +| Gate | Result at report preparation | +|---|---| +| `cargo test --profile test-fast --features metal,accelerate anchored_remove -- --nocapture` | Passed: managed-cache quarantine deletion and parent/final-unlink swap regressions | +| `cargo test --profile test-fast --features metal,accelerate anchored_publish -- --nocapture` | Passed: no-replace publish, owner/staging identity, symlink cleanup regressions | +| `cargo test --profile test-fast --features metal,accelerate anchored_delete -- --nocapture` | Passed: nested directory stat/open swap regression | +| `cargo test --lib --profile test-fast --features metal,accelerate downloader::tests:: -- --nocapture` | Passed: 71 passed, 1 ignored; fake anonymous metadata/HEAD/GET, saved-token negative, `?blobs=true`, metadata status, offline, loopback read timeout, simulated ENOSPC writer cleanup, disconnect truncation, checksum, no-weight completeness, path filtering, size mismatch, and token-mode tests | +| `cargo test --lib --profile test-fast --features metal,accelerate server::router_models::router_models_tests:: -- --nocapture` | Passed: 50 passed; shared download queue/idempotency/revision-alias/progress/cancellation, retry-after-failure, managed removal reservation, equality and descendant physical-overlap blocks, in-flight compatibility removal, and rescan guard | +| `cargo test --lib --profile test-fast --features metal,accelerate router_lifecycle_tests:: -- --nocapture` | Passed: 16 passed; whole producer comparisons for download operation/event fixtures plus lifecycle coordinator regressions | +| `cargo test --lib --profile test-fast --features metal,accelerate router_cache:: -- --nocapture` | Passed: 11 passed; descriptor-anchored publish/remove/delete race regressions | +| Focused route filters | Passed: `ui_download_route_replays_same_idempotency_key` and `ui_model_removal_route_matches_operation_accepted_fixture` | +| `cargo fmt --check` and `git diff --check` | Passed | +| `cargo clippy --lib --tests --features metal,accelerate -- -D warnings` | Passed | +| `PATH=/tmp/mlxcel-webui-contract/bin:$PATH make verify-webui-contract WEBUI_CONTRACT_PY=/tmp/mlxcel-webui-contract/bin/python` | Passed: 43 WebUI contract fixtures, DTO drift, and schema strictness | +| Root-owned real SmolLM download/load/generate/delete acceptance | Not run by this unit | +| Full workspace serial `make verify-test`, workspace all-target clippy, CUDA/GB10 checks | Not run by this unit; GB10 runner was reported down and the maintainer approved skipping only that unavailable required runner gate | + +The local evidence is deliberately CPU/fake-network scoped. It does not establish real HuggingFace transfer success, real model load/generation, CUDA behavior, or production browser UX; root owns those acceptance gates and merge decisions. + +## 5. Change Summary + +| Area | Change | +|---|---| +| WebUI API | Added authenticated `/ui-api/v1/downloads` and `/ui-api/v1/model-removals` route module using existing `RouterServerState` | +| Operation coordinator | Added bounded download admission, active download replay, cancellation registration, progress bytes, and publish sealing | +| Downloader | Added anonymous token mode, bounded anonymous metadata, fd-relative Unix download helper, immutable revision pinning, plan/progress hooks, and size validation | +| Cache source | Replaced router-managed removal with descriptor-anchored quarantine deletion for managed snapshots only | +| Router pool | Shared compatibility/WebUI download primitive, operation-busy removal reservation, rescan guard, and cache removal by model id | +| Contracts | Added full producer fixtures for running and succeeded download operations while preserving nullable `RevisionRef` compatibility | +| Tests | Added fake HTTP/token, filesystem race, lifecycle idempotency, cancellation, bounded queue, case alias, no-mkdir, and deletion race regressions | + +Statistics after repair checkpoint `d688ea4f`: the latest cycle-2 source commit changes 4 files with 437 insertions and 12 deletions on top of the prior repair/docs head. Original implementation commit: `466c6071 feat: add safe WebUI model library operations`; repair commits: `b7555005 fix: harden WebUI model library edge cases` and `d688ea4f fix: close WebUI library edge-case acceptance gaps`. + + +## 6. Fake Acceptance Matrix at Source Repair Checkpoint + +| Required fake case | Evidence in this PR | Status | +|---|---|---| +| Public valid download | `a_download_emits_the_b10621_event_sequence_and_lands_in_the_cache`; `anonymous_fd_download_sends_no_ambient_credentials_on_metadata_head_or_get` | Covered with fake downloader and fake HTTP/fd path | +| Bad repo / invalid revision / gated or private / 404 | `anonymous_fd_download_maps_metadata_status_without_publishing` covers 403 guidance and 404 missing repo/revision behavior | Covered for metadata-status handling | +| Offline | `anonymous_fd_download_offline_rejects_before_metadata_request` | Covered before metadata request | +| Timeout | `fd_stream_timeout_keeps_partial_private_and_unpublished` uses a loopback server that sends headers then stalls under a short client read timeout and verifies no final file or partial debris is left | Covered with loopback fake transport | +| Disconnect / truncation | `anonymous_fd_download_rejects_disconnect_truncation_before_publish`; `stream_file_rejects_known_size_mismatch_before_publish` | Covered | +| Disk full | `fd_stream_enospc_writer_cleans_partial_without_publishing` injects `ENOSPC` through a test-only writer factory around the production fd stream cleanup path | Covered as simulated ENOSPC; physical disk exhaustion was not run | +| Checksum and completeness | `anonymous_fd_download_rejects_same_size_checksum_mismatch_before_publish`; `anonymous_fd_download_rejects_metadata_without_weight_files_before_transfer`; `selected_safetensors_requires_lfs_sha256` | Covered for the anonymous WebUI fd path | +| Cancel/retry | `cancel_after_publish_linearization_is_refused_and_download_completes`; `remove_cancels_an_in_flight_download`; duplicate replay tests; `failed_download_can_retry_same_repo_with_new_idempotency_key` verifies a failed first operation can be retried successfully with a new operation id and one catalog entry | Covered for cancel, replay, and retry-after-failure | +| Restart / staging reconciliation | `killed_writer_and_terminal_history_reconcile_across_processes` uses three owned Rust test processes, real RouterPool/coordinator and anchored publication, and a fake local file writer; it checks killed running/queued work, session-local operation reset, explicit retry and one published catalog entry after another restart | CPU test-process coverage; root observed production CLI restart at `3cb4817d`, while the corrected complete real flow awaits rerun | +| Duplicate action | `duplicate_download_with_same_idempotency_key_replays_active_operation`; `active_download_replays_resolved_revision_alias`; `duplicate_download_idempotency_key_rejects_different_payload_before_alias_replay`; `ui_download_route_replays_same_idempotency_key` | Covered | +| Bounded queue saturation | `download_admission_rejects_queue_saturation_before_worker_network` | Covered | + +## 7. Learning Points and Follow-up + +A WebUI library action is an authority boundary, not just a button over an existing helper. Admission must be cheap and bounded, network and disk work must occur outside pool locks, and the final write/delete step must recheck the filesystem identity it is about to mutate. + +The remaining follow-up is acceptance, not hidden implementation work in this report: root should run the pinned small SmolLM real download/load/generate/delete driver, the CI-faithful full workspace gate, workspace all-target clippy, and any available CUDA/GB10 checks or document the runner outage waiver. The page layer should provide the visible deletion confirmation and distinguish unload from disk deletion without changing this backend wire contract. + +## 7. Bounded restart preparation (2026-09-14) + +The regression terminates only its owned child after observing a real running writer and queued operation. The interrupted private stage survives both restarts byte-for-byte and never becomes a catalog entry. A new server instance has no old operation history: authenticated operation GET returns the entire canonical 404 response before new admission, the same idempotency key is not replayed, and operation identity is compared as `(server_instance_id, operation_id)` rather than assuming globally unique ID strings. A new explicit retry uses fresh anchored staging, publishes once, and retains the same catalog model identity across the next process restart. Terminal history is likewise absent in that new session. No model is loaded; fake weight bytes are not a valid checkpoint and do not validate transport, checksum, production CLI startup or power-loss durability. + +The independent root CPU gate at head 42186c6a passed workspace all-target clippy, 43 contract fixtures, structural checks, formatting and diff checks. This does not replace the pending post-integration full test and real-checkpoint gates. No automatic abandoned-stage cleanup or resume behavior was added. + +Validation of this delta: the exact `server::router_cache::restart_tests::` filter passed 2 tests, and a second execution of the parent restart test passed. Scoped library/test clippy, formatting, diff checks and 43 strict contract fixtures passed. The first fixture attempts incorrectly assumed globally unique operation strings and null optional error fields; execution/review caught both assumptions and only test expectations changed. Independent bounded correctness and security reviews have no remaining HIGH/CRITICAL finding in this test/docs delta. + +## 8. Integrated acceptance and unload correction (2026-09-14) + +Integration checkpoint `3cb4817d` is based on centrally merged #1838/main `7e4577b1`. Library routes are mounted once inside the existing secured API-prefix composition. A shared pure policy drives bootstrap capabilities and mutation admission from actual cache authority, mode, offline policy and descriptor-platform support; single-model mode remains read-only and observation never creates cache directories. All 46 merged fixtures, including upstream Unicode/error cases, and 48 frontend tests passed. Independent integration reviews cleared the bounded source scope. + +Root's complete gate at `3cb4817d` passed 11,299 unique top-level tests with 361 ignored, plus two successful nested restart-child runs. The raw 11,301 aggregate includes those subprocess summaries and must not be presented as unique top-level coverage. Workspace all-target clippy, 46 strict fixtures, structural/format checks, feature-disabled compilation and both binaries' relocated empty/TTY checks passed. This is evidence for the pre-correction checkpoint, not a substitute for rerunning after the fix below. + +The first isolated real driver downloaded and hash-verified the pinned public SmolLM checkpoint, restarted the production CLI, verified changed server identity and old-operation 404, rediscovered the same complete unloaded catalog identity, loaded and generated nonempty chat. Unload then failed with `stale_revision` expected3/current4; removal was not attempted. Both owned server processes exited0, and the new temporary checkpoint was retained. This failure is preserved, not relabeled success. + +The defect already existed in merged main: unload validates the caller's revision, then its own Ready-to-Draining transition advances the revision, and its post-drain check incorrectly compares against the original token. A generic harness retry would hide this by retrying an already-draining model. The repair returns the owned drain revision atomically under the lifecycle mutex and rechecks that token/current entry after waiting. Shared revision authority means +1 cannot be assumed. RequestLease completion does not change revision. Only callers supplying a revision precondition (WebUI and explicit eviction) use the new token; legacy/shutdown callers retain identity-only cleanup, including failed-worker cleanup. The real harness is unchanged. + +The loaded fake-provider regression failed before the repair with the same expected3/current4 error as the real run. After the final compatibility refinement, five unload/token tests, 16 lifecycle tests and 50 router fake tests passed. Coverage includes a held request lease, rescan preserving the current loaded entry, observed worker exit, genuine external revision/registry replacement rejection, shared revision authority and legacy failed-worker cleanup. Independent correctness/security rechecks found no remaining HIGH/CRITICAL issue in this bounded repair. + +Final scoped library/test clippy, formatting and diff checks passed after the None-preserving refinement. Corrected full workspace and real-checkpoint acceptance remain root-owned and pending. + +## 9. UI rebase and unresolved full-gate boundary (2026-09-14) + +The branch was rebased onto merged UI PR #1863/main `b8d10fb1`, retaining the exact `@lablup/ui-common` version `0.1.0-alpha.19`, shared security/startup behavior and all 46 contract fixtures. The only rebase conflict was the generated bundle manifest. It was resolved through the canonical builder, not by hand-merging assets. Range comparison preserves all ten library commits unchanged except the generated manifest source digest; library/router/lifecycle production files are byte-identical to `857937ca`. The merged package lock, components, styles and generated JavaScript/CSS remain unchanged from main. + +Root's latest full gate at `857937ca` failed/incomplete after the core configured-depth DFlash test aborted during a confirmed Metal firmware progress-timeout recovery. The diagnostic does not establish either an external process or this WebUI repair as the cause. The historical full pass at `3cb4817d` does not replace this failed latest gate, and the GB10 outage waiver does not waive it. Root's later CPU-only workspace all-target clippy, 46 fixtures, structural/format checks and feature-disabled compilation passed at `857937ca`; the corrected actual library driver was not rerun. + +Further GPU, browser and full-suite execution is paused until the user confirms a quiet window. This rebase uses only CPU/fake tests and frontend unit/contract/build checks. Actual Safari, VoiceOver and native 200% manual checks are deferred by the user until the complete implementation is ready; they remain pending, not passed and not a per-unit blocker. + +CPU-only rebase verification passed: unload/token tests5, pure policy1, secured fake-router9, owned restart2 (plus two nested child runs), scoped clippy/format checks, all46 strict fixtures and76 frontend unit tests, type/lint checks and canonical deterministic bundle verification. No GPU, browser or full-suite execution occurred in this unit rebase. + +## 10. Latest-main integration and hosted portability repair (2026-09-14) + +The branch is now rebased onto main `e9ee5f044fa28c14db50009fffbe02283513c3ae`, including #1888 shared SwitchGLU migration, #1889 WebUI feature-disabled helper gates, #1882 f16 attention-range tests and #1890 Phixtral narrowing. All eleven pre-existing branch commits are range-diff equivalent; no merge conflict or manual asset change occurred. Upstream app/auth/router helper gates remain intact, alongside the single secured library mount, shared policy, exact ui-common alpha.19 and 46 fixtures. Historical full/GPU results do not validate this new core baseline. + +The GB10 runner subsequently executed the `71552f8c` hosted jobs: clippy failed on a same-type device-ID cast on Linux, and feature-disabled compilation failed on helpers now gated by upstream #1889. The earlier outage waiver does not apply to these executed failures. The local portability repair centralizes the existing cast in one documented helper with only a function-local unnecessary-cast allowance: Darwin's signed device representation is preserved, Linux's u64 representation remains unchanged, and neither comparison narrows identity or adds unsafe code. Temporary-directory regressions check matching identity, independent device/inode mismatch and child replacement; Darwin-specific cases check signed device values. Root's bounded read-only review cleared this delta. Actual hosted Linux clippy and feature-disabled results remain required after publication. + +The user has now confirmed GPU availability; root exclusively owns the next GPU/full and real-acceptance runs, while this unit remains CPU-only. The failed first real unload and latest Metal firmware timeout remain preserved; unchanged real download/restart/load/unload/delete acceptance must still be rerun by root. No existing user checkpoint is removed by this preparation. + +CPU preparation passed: unload5, policy1, secured router9, owned restart2 plus two child runs, anchored filesystem13 (including both new identity regressions), normal scoped library/test clippy, feature-disabled library/binary clippy, format/diff checks, 46 strict contracts, 76 frontend tests, type/lint and deterministic bundle verification. No GPU, browser or full-suite run was executed by this unit. + +## 11. Final acceptance at source 84f1aeb2 (2026-09-14) + +This section supersedes historical pending statuses above. After the user confirmed GPU availability, root's serialized chain at `84f1aeb2b4f5706bb0d63866b20ca48d74bf340a` exited0: 11,313 unique top-level tests passed, zero failed and 361 ignored across 124 summaries, plus two nested restart-child passes counted separately. Workspace all-target clippy passed with warnings denied in both default and feature-disabled configurations. All46 contracts, structural/format checks and both binaries' test-fast production relocated empty/offline/TTY checks passed. Actual hosted run [34844128193](https://github.com/lablup/mlxcel/actions/runs/34844128193) passed Linux clippy, OpenXLA feature compilation, WebUI bundle/feature-off and selected static checks; no GB10 exception was needed. Workflow-skipped CUDA sm70, MLX pin extraction and OpenXLA link are not execution claims. + +The unchanged strict fresh-cache driver downloaded public `mlx-community/SmolLM-135M-Instruct-4bit` at `642e06afe3fab57fd6cc518637c471af0a569e1e`, verifying 75,789,919 SafeTensors bytes with SHA256 `e91560ee24b13eee6ddeb14879d728a90780053d69057b15ff199ccadfcfe33b`. Production CLI restart changed server identity, returned canonical404 for the old operation and retained the complete unloaded model's stable catalog ID. Actual load, nonempty 32-token generation, unload with `worker_exit_observed=true`, API deletion of only the newly downloaded checkpoint and an empty final catalog passed. Repetitive output demonstrates execution, not generation quality. Both owned servers exited0 without forced termination. Tested server binary SHA256: `4a1d0d3c2805b9b81f5e20e9d925337776d850702e8947896cf0e7d9fe7d6fe0`. + +The first `3cb4817d` unload failure and `857937ca` Metal firmware timeout remain failures in the historical record; this later pass establishes neither an external cause nor a reboot claim. This final report update changes no source or bundle. Safari, VoiceOver and native 200% checks remain user-deferred until the whole implementation is ready, not passed. PR status remains review pending central merge. diff --git a/TECHNICAL_REPORTS/1872-safe-webui-library-20260913.ko.md b/TECHNICAL_REPORTS/1872-safe-webui-library-20260913.ko.md new file mode 100644 index 000000000..8c5640ee9 --- /dev/null +++ b/TECHNICAL_REPORTS/1872-safe-webui-library-20260913.ko.md @@ -0,0 +1,163 @@ +# 기술 리포트: PR #1872 — 안전한 WebUI 모델 라이브러리 작업 + +**작성일**: 2026-09-13 +**상태**: 열린 PR / status:review; 소스 `84f1aeb2` 전체·strict 실모델·hosted acceptance 통과, 중앙 머지 대기 +**언어**: Rust, JSON/OpenAPI fixture +**위험도**: 높음 +**구현 스냅샷**: 통합 기준 `3cb4817d`와 아래 owned-drain revision 수정; 이전 보완 체크포인트 `d688ea4f` 및 `b7555005` + +## 요약 + +[PR #1872](https://github.com/lablup/mlxcel/pull/1872)는 에픽 #1834의 이슈 #1841에 해당하는 WebUI 모델 라이브러리 변경 작업의 백엔드를 구현합니다. 인증된 WebUI 다운로드 및 관리 캐시 삭제 어댑터를 기존 router pool과 lifecycle coordinator 위에 추가하고, #1838 시작 경로와 공유할 상태 형태 및 compatibility route를 유지합니다. + +가장 위험한 부분은 디스크 변경과 Hub 접근입니다. 구현은 WebUI 다운로드를 익명·제한·fd-relative 방식으로 수행한 뒤 원자적으로 게시하고, 삭제는 descriptor-anchored private quarantine을 거치는 관리 캐시 snapshot으로 한정하여 사용자 모델 루트나 preset/external 경로를 삭제하지 않습니다. + +## 1. 문제 정의 + +WebUI는 모델을 다운로드하고 제거할 수 있어야 하지만, 브라우저 계층이 두 번째 모델 레지스트리나 권한 있는 파일시스템 shell이 되어서는 안 됩니다. 기존 router compatibility endpoint에는 다운로드·삭제 동작이 있었지만, 멱등성, 제한된 접수, 진행률, 취소, 종료 결과를 갖춘 공통 WebUI operation으로 표현되지 않았습니다. + +디스크 삭제는 되돌릴 수 없고 영향 범위가 큽니다. 모델 unload는 checkpoint 삭제와 다른 작업이며, 삭제는 symlink를 따라가거나 활성 writer/provider와 경쟁하거나 models-dir/preset/external artifact에 적용되면 안 됩니다. 또한 WebUI 다운로드는 서버 호스트에 CLI 토큰이 있다는 이유로 HuggingFace credential을 조용히 사용하면 안 됩니다. + +## 2. 기술적 선택과 그 이유 + +### 새 상태가 아니라 공유 coordinator 권한 사용 + +`webui::library::routes()`는 `/ui-api/v1` 아래에 mount되고 `RouterServerState`를 직접 사용합니다. 다운로드와 제거는 다른 WebUI operation과 같은 `RouterPool` 및 `LifecycleCoordinator`로 들어가므로 멱등성, 제한된 활성 작업 수, 취소 상태, SSE operation event, catalog rescan 동작의 권한이 하나로 유지됩니다. + +Compatibility `POST /models`도 wire schema는 유지하되 같은 다운로드 primitive로 enqueue합니다. 이제 route가 접수 전에 제한 없는 동기 metadata 요청을 수행하지 않고, 중복·queue 판단이 먼저 일어납니다. + +### 익명 immutable 다운로드 + +WebUI 다운로드는 임의 URL이나 local path가 아니라 공개 HuggingFace repo id와 선택적 revision만 받습니다. 익명 WebUI metadata는 ambient token file을 읽는 `hf-hub` builder 대신 응답 byte cap이 있는 직접 reqwest 요청을 사용합니다. CLI 다운로드 동작은 기존 environment token mode를 유지합니다. + +Worker는 metadata를 한 번 resolve하고 resolved SHA에 pin하여 전송하며, manifest sibling 수, 선택 파일 수, filename 길이를 제한하고, 가능한 경우 HEAD로 알려진 byte total을 계산한 뒤 실제 기록 byte와 비교하고 나서 게시합니다. 종료 operation result는 resolved revision을 담고, 공유 `RevisionRef` schema는 요청 호환성을 위해 nullable 그대로 둡니다. + +### Descriptor-anchored staging과 취소 경계 + +Unix에서 router 관리 다운로드는 fd로 열린 고유 private `.mlxcel-staging` 디렉터리에 기록합니다. 원격 경로는 directory descriptor 기준으로 생성되고 symlink traversal을 하지 않으며, coordinator의 `begin_publish` hook이 취소를 봉인한 뒤 atomic no-replace rename으로 게시합니다. Writer가 관찰하기 전 접수된 취소는 worker가 종료를 확인한 뒤 terminal 상태가 되고, 게시 linearization 이후의 늦은 취소는 거절되어 완성된 snapshot이 게시될 수 있습니다. + +완전성 정책은 `config.json`, 하나 이상의 선택된 safetensors weight 파일, metadata/header가 제공하는 알려진 크기와 기록 byte의 일치, 그리고 익명 WebUI 경로의 모든 safetensors에 대한 LFS SHA-256 metadata를 요구합니다. Stream은 fd-relative rename 전에 hash되므로 같은 크기의 손상된 weight도 게시 전에 실패합니다. + +### 관리 캐시 삭제만 허용 + +`POST /ui-api/v1/model-removals`는 고정된 request shape인 `model_id`, `expected_revision`, `idempotency_key`를 유지합니다. 사용자에게 보이는 확인 modal은 WebUI page 작업 범위이고, backend endpoint 자체가 명시적인 삭제 의도입니다. Pool은 cache-sourced이고 현재 revision이 맞으며 idle 상태인 model entry만 허용하고, 삭제를 spawn하기 전에 operation-busy로 표시하여 removal이 catalog slot을 소유하는 동안 rescan이 stale entry를 재생성하지 못하게 합니다. + +삭제는 설정된 cache root, owner, target, private quarantine을 `O_NOFOLLOW`로 열고, quarantine rename 직전과 rename 후 identity를 확인한 다음 held fd 기준으로 재귀 unlink를 수행합니다. 중첩 디렉터리 삭제는 예상 dev/ino를 open과 최종 rmdir 검사까지 전달합니다. Identity가 바뀌면 삭제는 fail-closed하고 quarantined snapshot을 조사용으로 남깁니다. 같은 UID 프로세스는 여전히 denial of service나 quarantine cleanup 실패 race를 유발할 수 있으므로, 이것은 hostile same-user를 절대적으로 격리한다는 주장이 아닙니다. + +## 3. 리뷰 보완 사항 + +Parent pre-review에서 지적된 다음 고위험 항목을 PR 게시 전 반영했습니다. + +- 마지막 전송 chunk 이후 publish 전 취소 race를 `begin_publish_operation`으로 linearize했고, 늦은 취소가 거절되며 operation이 성공하는 결정적 회귀 테스트를 추가했습니다. +- Metadata 요청은 `?blobs=true`를 사용하고, LFS SHA-256이 없는 safetensors는 거절하며, stream을 fd-relative rename 전에 hash하고, 알려진 metadata/HEAD 크기를 실제 byte 수와 비교하며, metadata와 선택 파일 수·파일명 길이를 제한한 뒤 다운로드를 진행합니다. +- 익명 WebUI metadata는 더 이상 저장된 token file을 읽는 `hf-hub` API builder를 생성하지 않으며, 테스트는 token file과 env var를 심고 metadata·HEAD·GET fake 요청에 Authorization header가 없음을 확인합니다. +- 제거 작업은 `operation_busy`로 lifecycle을 예약하고, per-entry operation guard 획득 후 다시 검사하며, 설정된 models-dir 또는 preset/non-cache entry와의 동일 경로 및 ancestor/descendant overlap을 거절하여 cache 삭제가 user/preset 소유 중첩 snapshot을 지우지 못하게 했습니다. +- Parent swap 및 중첩 stat/open 삭제 race는 결정적 테스트를 갖고 잘못된 tree를 삭제하지 않고 실패합니다. +- `anchored_delete.rs`와 인접 테스트 모듈로 분리하여 새 handwritten module을 500줄 이하로 유지했습니다. + +## 4. 검증 기록 + +| 게이트 | 리포트 작성 시점 결과 | +|---|---| +| `cargo test --profile test-fast --features metal,accelerate anchored_remove -- --nocapture` | 통과: 관리 캐시 quarantine 삭제 및 parent/final-unlink swap 회귀 | +| `cargo test --profile test-fast --features metal,accelerate anchored_publish -- --nocapture` | 통과: no-replace publish, owner/staging identity, symlink cleanup 회귀 | +| `cargo test --profile test-fast --features metal,accelerate anchored_delete -- --nocapture` | 통과: 중첩 디렉터리 stat/open swap 회귀 | +| `cargo test --lib --profile test-fast --features metal,accelerate downloader::tests:: -- --nocapture` | 통과: 71 passed, 1 ignored; fake anonymous metadata/HEAD/GET, saved-token 음성, `?blobs=true`, metadata status, offline, loopback read timeout, simulated ENOSPC writer cleanup, disconnect truncation, checksum, weight 없음 완전성, path filtering, size mismatch, token-mode 테스트 | +| `cargo test --lib --profile test-fast --features metal,accelerate server::router_models::router_models_tests:: -- --nocapture` | 통과: 50 passed; 공유 download queue/idempotency/revision-alias/progress/cancellation, 실패 후 재시도, 관리 removal reservation, 동일 경로 및 descendant physical-overlap 차단, in-flight compatibility removal, rescan guard | +| `cargo test --lib --profile test-fast --features metal,accelerate router_lifecycle_tests:: -- --nocapture` | 통과: 16 passed; download operation/event fixture 전체 producer 비교 및 lifecycle coordinator 회귀 | +| `cargo test --lib --profile test-fast --features metal,accelerate router_cache:: -- --nocapture` | 통과: 11 passed; descriptor-anchored publish/remove/delete race 회귀 | +| 집중 route 필터 | 통과: `ui_download_route_replays_same_idempotency_key`, `ui_model_removal_route_matches_operation_accepted_fixture` | +| `cargo fmt --check` 및 `git diff --check` | 통과 | +| `cargo clippy --lib --tests --features metal,accelerate -- -D warnings` | 통과 | +| `PATH=/tmp/mlxcel-webui-contract/bin:$PATH make verify-webui-contract WEBUI_CONTRACT_PY=/tmp/mlxcel-webui-contract/bin/python` | 통과: WebUI contract fixture 43개, DTO drift, schema strictness | +| 루트 소유 실제 SmolLM 다운로드·로드·생성·삭제 acceptance | 이 유닛은 실행하지 않음 | +| 전체 workspace 직렬 `make verify-test`, workspace all-target clippy, CUDA/GB10 검사 | 이 유닛은 실행하지 않음; GB10 runner는 down 상태였고 maintainer가 해당 unavailable required runner gate만 생략 승인 | + +로컬 증거는 의도적으로 CPU·fake-network 범위입니다. 실제 HuggingFace 전송 성공, 실제 모델 load/generation, CUDA 동작, production browser UX를 입증하지 않습니다. 해당 acceptance gate와 merge 판단은 루트가 소유합니다. + +## 5. 변경 요약 + +| 영역 | 변경 | +|---|---| +| WebUI API | 기존 `RouterServerState`를 사용하는 인증된 `/ui-api/v1/downloads` 및 `/ui-api/v1/model-removals` route 모듈 추가 | +| Operation coordinator | 제한된 download admission, active download replay, cancellation 등록, progress bytes, publish sealing 추가 | +| Downloader | anonymous token mode, 제한된 익명 metadata, fd-relative Unix download helper, immutable revision pinning, plan/progress hook, size validation 추가 | +| Cache source | Router 관리 removal을 관리 snapshot 전용 descriptor-anchored quarantine deletion으로 교체 | +| Router pool | Compatibility/WebUI 공유 download primitive, operation-busy removal reservation, rescan guard, model id 기반 cache removal 추가 | +| 계약 | nullable `RevisionRef` 호환성을 유지하면서 running/succeeded download operation 전체 producer fixture 추가 | +| 테스트 | fake HTTP/token, 파일시스템 race, lifecycle idempotency, cancellation, 제한 queue, case alias, no-mkdir, deletion race 회귀 추가 | + +보완 체크포인트 `d688ea4f` 이후 통계는 이전 보완/docs head 위에 최신 cycle-2 source commit 4개 파일 변경, 437줄 추가, 12줄 삭제입니다. 최초 구현 커밋은 `466c6071 feat: add safe WebUI model library operations`이고, 보완 커밋은 `b7555005 fix: harden WebUI model library edge cases` 및 `d688ea4f fix: close WebUI library edge-case acceptance gaps`입니다. + + +## 6. 소스 보완 체크포인트의 Fake Acceptance Matrix + +| 요구 fake case | 이 PR의 증거 | 상태 | +|---|---|---| +| 공개 valid download | `a_download_emits_the_b10621_event_sequence_and_lands_in_the_cache`; `anonymous_fd_download_sends_no_ambient_credentials_on_metadata_head_or_get` | fake downloader 및 fake HTTP/fd 경로로 커버 | +| Bad repo / invalid revision / gated·private / 404 | `anonymous_fd_download_maps_metadata_status_without_publishing`이 403 안내와 404 repo/revision 없음 동작을 커버 | metadata status 처리 커버 | +| Offline | `anonymous_fd_download_offline_rejects_before_metadata_request` | metadata 요청 전 커버 | +| Timeout | `fd_stream_timeout_keeps_partial_private_and_unpublished`가 header를 보낸 뒤 body를 지연하는 loopback 서버와 짧은 client read timeout으로 최종 파일 및 partial debris가 남지 않음을 확인합니다 | loopback fake transport로 커버 | +| Disconnect / truncation | `anonymous_fd_download_rejects_disconnect_truncation_before_publish`; `stream_file_rejects_known_size_mismatch_before_publish` | 커버 | +| Disk full | `fd_stream_enospc_writer_cleans_partial_without_publishing`가 production fd stream cleanup 경로 주위의 test-only writer factory로 `ENOSPC`를 주입합니다 | simulated ENOSPC로 커버; 실제 물리 disk exhaustion은 실행하지 않음 | +| Checksum 및 completeness | `anonymous_fd_download_rejects_same_size_checksum_mismatch_before_publish`; `anonymous_fd_download_rejects_metadata_without_weight_files_before_transfer`; `selected_safetensors_requires_lfs_sha256` | 익명 WebUI fd 경로 커버 | +| Cancel/retry | `cancel_after_publish_linearization_is_refused_and_download_completes`; `remove_cancels_an_in_flight_download`; duplicate replay 테스트; `failed_download_can_retry_same_repo_with_new_idempotency_key`가 첫 실패 후 새 operation id로 재시도 성공 및 catalog entry 1개를 확인합니다 | cancel, replay, retry-after-failure 커버 | +| Restart / staging reconciliation | `killed_writer_and_terminal_history_reconcile_across_processes`가 소유한 Rust 테스트 프로세스 3개, 실제 RouterPool/coordinator와 anchored publication, 로컬 fake writer로 실행·대기 작업 중단과 세션 이력 초기화, 명시적 재시도 및 다음 재시작의 catalog entry 1개를 검증합니다 | CPU 테스트 프로세스 검증 및 루트의 `3cb4817d` production CLI 재시작 확인; 수정본 전체 실모델 흐름 재검증 대기 | +| Duplicate action | `duplicate_download_with_same_idempotency_key_replays_active_operation`; `active_download_replays_resolved_revision_alias`; `duplicate_download_idempotency_key_rejects_different_payload_before_alias_replay`; `ui_download_route_replays_same_idempotency_key` | 커버 | +| Bounded queue saturation | `download_admission_rejects_queue_saturation_before_worker_network` | 커버 | + +## 7. 학습 포인트와 후속 조치 + +WebUI library action은 기존 helper 위의 버튼이 아니라 권한 경계입니다. 접수는 싸고 제한되어야 하며, network와 disk 작업은 pool lock 밖에서 수행되어야 하고, 마지막 write/delete 단계는 지금 변경하려는 filesystem identity를 다시 확인해야 합니다. + +남은 작업은 이 리포트 안의 숨은 구현이 아니라 acceptance입니다. 루트는 pinned small SmolLM 실제 download/load/generate/delete driver, CI-faithful 전체 workspace gate, workspace all-target clippy, 사용 가능한 CUDA/GB10 검사 또는 runner outage waiver 기록을 수행해야 합니다. Page 계층은 이 backend wire contract를 바꾸지 않으면서 눈에 보이는 삭제 확인을 제공하고 unload와 disk deletion을 명확히 구분해야 합니다. + +## 7. 제한된 재시작 검증 준비 (2026-09-14) + +회귀 테스트는 실제 writer 실행과 queued operation을 관찰한 뒤 자신이 소유한 자식 프로세스만 종료합니다. 중단된 private stage는 두 번의 재시작 동안 바이트가 유지되며 catalog에 노출되지 않습니다. 새 server instance에서는 이전 operation 조회가 신규 작업 접수 전에 전체 canonical 404 응답을 반환하며, 같은 idempotency key도 이전 작업으로 replay되지 않습니다. operation ID 문자열의 전역 유일성을 가정하지 않고 `(server_instance_id, operation_id)`를 비교합니다. 명시적 재시도는 새 anchored stage에서 한 번만 publish하고 다음 프로세스 재시작에도 동일한 catalog model identity를 유지합니다. 완료된 작업 이력도 새 세션에는 남지 않습니다. 모델은 로드하지 않으며 fake weight는 유효한 checkpoint가 아닙니다. 실제 전송·checksum·production CLI 시작·전원 장애 내구성 검증을 대신하지 않습니다. + +루트의 head 42186c6a CPU 게이트는 workspace all-target clippy, 43 contract fixture, 구조 검사, 포맷 및 diff 검사를 통과했습니다. 통합 후 전체 테스트와 실모델 게이트는 별도로 남아 있습니다. 버려진 stage의 자동 삭제나 resume 동작은 추가하지 않았습니다. + +이번 변경 검증: 정확한 `server::router_cache::restart_tests::` 필터의 테스트2개와 부모 재시작 테스트의 추가 실행을 통과했습니다. scoped library/test clippy, 포맷, diff 검사와43 strict contract fixture도 통과했습니다. 초기 fixture는 operation 문자열의 전역 유일성과 optional 오류 필드의 null 직렬화를 잘못 가정했고 실행·리뷰에서 발견하여 테스트 기대값만 수정했습니다. 이 테스트·문서 변경에 대한 독립 correctness/security 리뷰의 HIGH/CRITICAL 잔여 사항은 없습니다. + +## 8. 통합 검증과 unload 수정 (2026-09-14) + +통합 체크포인트 `3cb4817d`는 중앙에서 머지한 #1838/main `7e4577b1` 기반입니다. 기존 보안·API prefix 구성 안에 library route를 한 번만 연결하고, 순수한 공유 정책으로 실제 cache 권한·mode·offline·descriptor platform에 따른 bootstrap capability와 mutation admission을 일치시켰습니다. 단일 모델 모드는 read-only이며 관찰 중 cache 디렉터리를 생성하지 않습니다. 상위 Unicode/error 사례를 포함한 fixture 46개와 frontend 테스트 48개를 통과했고 독립 통합 리뷰도 제한된 범위에서 완료했습니다. + +루트의 `3cb4817d` 전체 게이트는 고유 top-level 테스트 11,299개 통과, 361개 ignored 및 별도 nested restart-child 실행 2개 통과입니다. 원시 합계 11,301에는 subprocess 결과가 포함되므로 고유 top-level 수로 보고하지 않습니다. Workspace all-target clippy, strict fixture 46개, 구조·포맷 검사, feature-disabled 컴파일과 두 바이너리의 relocated empty/TTY 검증도 통과했습니다. 이는 아래 수정 전 체크포인트의 증거이며 수정 후 재검증을 대신하지 않습니다. + +첫 격리 실모델 driver는 고정한 공개 SmolLM checkpoint의 실제 다운로드·hash 검증, production CLI 재시작, 새 server identity와 이전 operation 404, 동일한 complete/unloaded catalog identity, 실제 load 및 비어 있지 않은 chat 생성을 확인했습니다. 이후 unload가 `stale_revision` expected3/current4로 실패하여 삭제는 시도하지 않았습니다. 소유한 서버 두 프로세스는 모두 exit0으로 종료했고 새 임시 checkpoint는 보존했습니다. 이 실패를 성공으로 바꾸어 기록하지 않습니다. + +결함은 이미 merged main에 있었습니다. unload가 caller revision을 검사한 뒤 자신의 Ready-to-Draining 전환으로 revision을 증가시키고, 대기 후에도 이전 token을 비교하여 스스로 거부했습니다. 일반적인 harness 재시도는 이미 draining인 모델을 재시도하면서 결함을 가릴 수 있습니다. 수정은 lifecycle mutex 안에서 자신의 drain revision을 원자적으로 반환하고 대기 후 그 token과 현재 entry를 검사합니다. 공유 revision authority 때문에 +1을 가정하지 않으며 RequestLease 종료도 revision을 변경하지 않습니다. Revision precondition이 있는 WebUI·명시적 eviction caller만 새 token을 사용하고, legacy·shutdown은 failed-worker 정리를 포함한 기존 identity-only 검사를 유지합니다. 실제 harness는 변경하지 않았습니다. + +로드된 fake provider 회귀 테스트는 수정 전에 실제 실행과 동일한 expected3/current4 오류로 실패했습니다. 최종 호환성 보완 후 unload/token 테스트 5개, lifecycle 테스트 16개, router fake 테스트 50개를 통과했습니다. 활성 request lease, loaded entry를 유지하는 rescan, 실제 worker 종료 관찰, 외부 revision·registry 교체 거부, 공유 revision authority 및 legacy failed-worker 정리를 검증합니다. 독립 correctness/security 재검토에서 이 제한된 수정의 HIGH/CRITICAL 잔여 사항은 없습니다. + +None caller의 기존 동작을 보존한 최종 수정 후 scoped library/test clippy, 포맷 및 diff 검사를 통과했습니다. 수정본 전체 workspace·실모델 검증은 루트 담당으로 남아 있습니다. + +## 9. UI rebase와 미해결 전체 게이트 경계 (2026-09-14) + +머지된 UI PR #1863/main `b8d10fb1` 위로 rebase하면서 정확한 `@lablup/ui-common` 버전 `0.1.0-alpha.19`, 공유 보안·시작 동작 및 contract fixture 46개를 유지했습니다. 충돌은 생성된 bundle manifest 하나뿐이었으며 asset 수동 병합 없이 canonical builder로 해결했습니다. Range 비교에서 library 커밋 10개는 생성 manifest의 source digest 외에 동일하고, library/router/lifecycle production 파일은 `857937ca`와 바이트 단위로 동일합니다. Package lock, component, style 및 생성 JavaScript/CSS도 merged main을 유지합니다. + +루트의 최신 `857937ca` 전체 게이트는 core의 configured-depth DFlash 테스트가 확인된 Metal firmware progress-timeout recovery 중 abort되어 실패·미완료 상태입니다. 진단은 외부 프로세스나 WebUI 수정을 원인으로 확정하지 않습니다. `3cb4817d`의 과거 전체 통과는 최신 실패를 대신하지 않으며 GB10 장애 면제도 이 실패에 적용되지 않습니다. 이후 루트의 `857937ca` CPU-only workspace all-target clippy, fixture 46개, 구조·포맷 검사와 feature-disabled 컴파일은 통과했으나 수정본 실제 library driver는 재실행하지 않았습니다. + +사용자가 quiet window를 알려줄 때까지 GPU·browser·전체 suite 실행은 중단합니다. 이번 rebase에서는 CPU/fake 테스트와 frontend unit·contract·build 검사만 사용합니다. 실제 Safari·VoiceOver·native 200% 수동 검사는 사용자가 전체 구현 완료 시점으로 미뤘으므로 pending이며 통과나 개별 unit blocker로 보고하지 않습니다. + +CPU-only rebase 검증: unload/token5, 순수 policy1, 보안 fake-router9, owned restart2(별도 nested child 실행2개), scoped clippy·포맷, strict fixture46개, frontend unit76개, type·lint 및 canonical deterministic bundle 검증을 통과했습니다. 이 unit rebase에서는 GPU·browser·전체 suite를 실행하지 않았습니다. + +## 10. 최신 main 통합과 hosted 이식성 수정 (2026-09-14) + +브랜치를 main `e9ee5f044fa28c14db50009fffbe02283513c3ae` 위로 rebase했습니다. #1888 shared SwitchGLU 전환, #1889 WebUI feature-disabled helper gate, #1882 f16 attention 범위 테스트, #1890 Phixtral narrowing을 포함합니다. 기존 브랜치 커밋 11개는 모두 range-diff상 동일하며 충돌이나 수동 asset 수정이 없었습니다. Upstream app/auth/router helper gate와 단일 secured library mount, 공유 정책, 정확한 ui-common alpha.19, fixture 46개를 보존했습니다. 과거 전체·GPU 결과는 새 core 기준점의 검증 결과가 아닙니다. + +이후 GB10 runner가 `71552f8c` hosted 작업을 실제 실행했습니다. Linux에서는 device-ID의 동일 타입 cast 때문에 clippy가 실패했고, feature-disabled 컴파일은 이제 upstream #1889가 gate한 helper 때문에 실패했습니다. 앞서 승인된 장애 면제는 실제 실행된 실패에 적용되지 않습니다. 로컬 이식성 수정은 기존 cast를 설명이 있는 helper 하나로 모으고 해당 함수에만 unnecessary-cast 허용을 둡니다. Darwin의 signed device 표현과 Linux의 u64 표현을 그대로 보존하며 어느 비교도 identity를 축소하거나 unsafe 코드를 추가하지 않습니다. 임시 디렉터리 회귀는 동일 identity, 독립적인 device/inode 불일치, child 교체를 검사하고 Darwin 전용 검사는 signed device 값을 확인합니다. 루트의 한정된 읽기 전용 리뷰는 이 변경을 승인했습니다. 게시 후 실제 hosted Linux clippy 및 feature-disabled 통과가 여전히 필요합니다. + +사용자가 이제 GPU 사용 가능 상태를 확인했습니다. 다음 GPU·전체 suite 및 실제 acceptance는 루트가 독점 실행하며 이 유닛은 CPU-only를 유지합니다. 첫 실제 unload 실패와 최신 Metal firmware timeout은 그대로 기록하며, 루트가 변경하지 않은 실제 download/restart/load/unload/delete acceptance를 다시 실행해야 합니다. 이 준비 과정에서 기존 사용자 checkpoint를 삭제하지 않습니다. + +CPU 준비 검증은 unload 5개, policy 1개, secured router 9개, owned restart 2개와 child 실행 2회, anchored filesystem 13개(새 identity 회귀 2개 포함), 일반 scoped library/test clippy, feature-disabled library/binary clippy, 포맷·diff, strict contract 46개, frontend 테스트 76개, type/lint 및 결정적 bundle 검증을 통과했습니다. 이 유닛은 GPU·브라우저·전체 suite를 실행하지 않았습니다. + +## 11. 소스 84f1aeb2 최종 acceptance (2026-09-14) + +이 절은 앞선 과거의 대기 상태를 갱신합니다. 사용자가 GPU 사용 가능 상태를 확인한 뒤 루트가 `84f1aeb2b4f5706bb0d63866b20ca48d74bf340a`에서 실행한 직렬 chain은 exit0으로 종료했습니다. Summary 124개에서 고유 top-level 테스트 11,313개 통과, 실패 0개, ignored 361개이며 nested restart-child 실행 2회 통과는 별도 집계했습니다. Workspace all-target clippy는 default 및 feature-disabled 구성 모두 warnings denied로 통과했습니다. Contract 46개, 구조·포맷 검사, 두 바이너리의 test-fast production relocated empty/offline/TTY도 통과했습니다. 실제 hosted run [34844128193](https://github.com/lablup/mlxcel/actions/runs/34844128193)은 Linux clippy, OpenXLA feature compilation, WebUI bundle/feature-off 및 선택된 정적 검사를 통과했으며 GB10 예외는 필요하지 않았습니다. Workflow가 건너뛴 CUDA sm70, MLX pin extraction, OpenXLA link는 실행 주장이 아닙니다. + +변경하지 않은 strict fresh-cache driver는 공개 `mlx-community/SmolLM-135M-Instruct-4bit`의 `642e06afe3fab57fd6cc518637c471af0a569e1e`를 다운로드하고 75,789,919바이트 SafeTensors의 SHA256 `e91560ee24b13eee6ddeb14879d728a90780053d69057b15ff199ccadfcfe33b`를 검증했습니다. Production CLI 재시작 후 server identity 변경, 이전 operation의 canonical404, complete/unloaded 모델의 stable catalog ID 유지를 확인했습니다. 실제 load, 비어 있지 않은 32-token 생성, `worker_exit_observed=true`인 unload, 방금 받은 checkpoint만 API로 삭제, 최종 빈 catalog를 모두 통과했습니다. 출력은 반복적이므로 실행 검증이지 생성 품질 주장이 아닙니다. 소유한 서버 두 개는 강제 종료 없이 모두 exit0으로 종료했습니다. 검증한 서버 바이너리 SHA256: `4a1d0d3c2805b9b81f5e20e9d925337776d850702e8947896cf0e7d9fe7d6fe0`. + +첫 `3cb4817d` unload 실패와 `857937ca` Metal firmware timeout은 과거 실패 기록으로 보존합니다. 이후 통과는 외부 원인을 확정하거나 재부팅을 주장하지 않습니다. 최종 리포트 갱신은 소스와 bundle을 변경하지 않습니다. Safari·VoiceOver·native 200% 검사는 전체 구현 준비 후로 사용자가 유예했으며 통과가 아닙니다. PR은 중앙 머지 전까지 review 상태입니다. diff --git a/docs/webui/state-machines.md b/docs/webui/state-machines.md index 001bd894b..6d3a0a2d3 100644 --- a/docs/webui/state-machines.md +++ b/docs/webui/state-machines.md @@ -37,6 +37,8 @@ Deletion refusal cases are 409 for transient busy states and 422 for unsupported `idempotency_key` scope is one `server_instance_id` and expires with operation history retention unless the operation is still active. Reusing the same key with the same request returns the original operation. Reusing the same key with a different normalized request returns 409 conflict. A changed `server_instance_id` means the client must not replay POSTs automatically; it must resnapshot and ask the user when necessary. +Download operation history and idempotency records are process-local. After a server restart, clients must discard old operation references even if an operation ID string is reused in the new `server_instance_id`. Previously published managed snapshots are rediscovered from the configured store; abandoned private `.mlxcel-staging` directories are excluded from the catalog and are neither resumed nor automatically deleted by a new process. An explicit retry starts a new private stage and does not adopt another writer's partial files. This is process-restart reconciliation, not a guarantee of power-loss durability. + Every mutating model request carries `expected_revision`. If the catalog entry revision changed since the UI snapshot, return 409 `stale_revision` with the current operation/model pointer when possible. ## Snapshot to SSE fence diff --git a/src/downloader/mod.rs b/src/downloader/mod.rs index a73133f16..b0a8b7db3 100644 --- a/src/downloader/mod.rs +++ b/src/downloader/mod.rs @@ -119,6 +119,7 @@ use futures::StreamExt; use hf_hub::api::sync::{Api, ApiBuilder}; use hf_hub::{Repo, RepoType}; use percent_encoding::{AsciiSet, CONTROLS, utf8_percent_encode}; +use sha2::{Digest, Sha256}; use std::fs; use std::io::Read; use std::path::{Path, PathBuf}; @@ -174,6 +175,10 @@ const SEGMENT_ENCODE_SET: &AsciiSet = &CONTROLS /// Younger partial files are left in place to avoid racing against a concurrent /// `mlxcel` process that is mid-download in the same destination directory. const PARTIAL_TEMPFILE_STALE_AGE: Duration = Duration::from_secs(60 * 60); +const MAX_HF_SIBLINGS: usize = 16_384; +const MAX_SELECTED_DOWNLOAD_FILES: usize = 8_192; +const MAX_REPO_FILENAME_BYTES: usize = 512; +const MAX_HF_METADATA_BYTES: usize = 8 * 1024 * 1024; /// Resolved options for a download invocation. /// @@ -181,6 +186,14 @@ const PARTIAL_TEMPFILE_STALE_AGE: Duration = Duration::from_secs(60 * 60); /// shared adapter both binaries use). The struct exists so that programmatic /// callers (and unit tests) can drive [`download_repo`] without going through /// clap parsing. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TokenMode { + /// Resolve an explicit token first, then HF_TOKEN / HUGGING_FACE_HUB_TOKEN. + Environment, + /// Use anonymous HuggingFace requests; explicit and ambient tokens are ignored. + Anonymous, +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct DownloadOptions { /// HuggingFace repository identifier, e.g. `mlx-community/Qwen3-4B-4bit`. @@ -199,9 +212,13 @@ pub struct DownloadOptions { /// Repository revision (branch, tag, or commit). Defaults to `main` when /// `None`. pub revision: Option, - /// Authentication token override. When `None`, falls back to environment - /// variables (`HF_TOKEN`, then `HUGGING_FACE_HUB_TOKEN`). + /// Authentication token override. When `None` and `token_mode` is + /// [`TokenMode::Environment`], falls back to environment variables + /// (`HF_TOKEN`, then `HUGGING_FACE_HUB_TOKEN`). WebUI-managed library + /// downloads set [`TokenMode::Anonymous`] so browser actions never + /// silently consume ambient Hub credentials. pub token: Option, + pub token_mode: TokenMode, /// Optional repository-relative glob allow-list applied after the built-in /// safe file-type filter. Empty means every built-in-allowed file. pub include: Vec, @@ -219,6 +236,7 @@ impl DownloadOptions { models_dir: args.models_dir.clone(), revision: args.revision.clone(), token: args.token.clone(), + token_mode: TokenMode::Environment, include: args.include.clone(), force: args.force, } @@ -264,6 +282,13 @@ impl DownloadOptions { /// targeted error message that names the env var or flag — see /// [`validate_token`] which is invoked at HTTP-client construction time. pub fn resolve_token(explicit: Option<&str>) -> Option { + resolve_token_with_mode(explicit, TokenMode::Environment) +} + +pub fn resolve_token_with_mode(explicit: Option<&str>, mode: TokenMode) -> Option { + if mode == TokenMode::Anonymous { + return None; + } if let Some(t) = explicit { let trimmed = t.trim(); if !trimmed.is_empty() { @@ -312,6 +337,10 @@ fn build_api(token: Option) -> Result { if let Some(tok) = token { builder = builder.with_token(Some(tok)); } + finish_api_builder(builder) +} + +fn finish_api_builder(builder: ApiBuilder) -> Result { builder .build() .map_err(|err| anyhow!("Failed to initialize Hugging Face API client: {err}")) @@ -534,6 +563,153 @@ fn file_url(endpoint: &str, repo_id: &str, revision: &str, filename: &str) -> St format!("{endpoint}/{repo_enc}/resolve/{rev_enc}/{file_enc}") } +fn repo_info_url(endpoint: &str, repo_id: &str, revision: &str) -> String { + let endpoint = endpoint.trim_end_matches('/'); + let repo_enc = encode_path_segments(repo_id); + let rev_enc = encode_path_segments(revision); + format!("{endpoint}/api/models/{repo_enc}/revision/{rev_enc}?blobs=true") +} + +#[derive(Debug, Clone, serde::Deserialize)] +struct HubRepoInfo { + sha: String, + #[serde(default)] + siblings: Vec, +} + +#[derive(Debug, Clone, serde::Deserialize)] +struct HubSibling { + rfilename: String, + #[serde(default)] + size: Option, + #[serde(default)] + lfs: Option, +} + +#[derive(Debug, Clone, serde::Deserialize)] +struct HubLfs { + #[serde(default)] + sha256: Option, + #[serde(default)] + size: Option, +} + +#[derive(Debug, Clone)] +struct SelectedDownloadFile { + filename: String, + expected_size: Option, + sha256: Option, +} + +impl SelectedDownloadFile { + fn from_sibling(repo_id: &str, sibling: &HubSibling) -> Result { + validate_manifest_filename(repo_id, &sibling.rfilename)?; + let sha256 = sibling + .lfs + .as_ref() + .and_then(|lfs| lfs.sha256.as_deref()) + .map(|hash| validate_sha256_hex(&sibling.rfilename, hash)) + .transpose()?; + if sibling.rfilename.ends_with(".safetensors") && sha256.is_none() { + return Err(anyhow!( + "Repository metadata for '{}' is missing an LFS SHA-256 digest; refusing to publish unverifiable weights", + sibling.rfilename + )); + } + Ok(Self { + filename: sibling.rfilename.clone(), + expected_size: sibling + .lfs + .as_ref() + .and_then(|lfs| lfs.size) + .or(sibling.size), + sha256, + }) + } +} + +fn build_reqwest_client(token: Option<&str>, enforce_https: bool) -> Result { + let mut headers = reqwest::header::HeaderMap::new(); + if let Some(tok) = token { + validate_token(tok)?; + let auth_val = format!("Bearer {tok}"); + headers.insert( + reqwest::header::AUTHORIZATION, + reqwest::header::HeaderValue::from_str(&auth_val).with_context( + || "HF token contains invalid characters (must be ASCII, no control chars)", + )?, + ); + } + let mut builder = reqwest::Client::builder() + .connect_timeout(Duration::from_secs(10)) + .read_timeout(Duration::from_secs(30)) + .default_headers(headers); + if enforce_https { + builder = builder.https_only(true); + } + for cert in load_extra_ca_certificates()? { + builder = builder.add_root_certificate(cert); + } + builder.build().context("Failed to create HTTP client") +} + +async fn fetch_anonymous_repo_info( + client: &reqwest::Client, + endpoint: &str, + repo_id: &str, + revision: &str, +) -> Result { + let url = repo_info_url(endpoint, repo_id, revision); + let response = client + .get(&url) + .send() + .await + .with_context(|| format!("HTTP request failed for repository metadata '{repo_id}'"))?; + let status = response.status(); + if !status.is_success() { + let code = status.as_u16(); + return Err(match code { + 401 | 403 => anyhow!( + "Repository '{repo_id}' requires authentication, is gated, or is private; WebUI downloads use anonymous public Hub access only" + ), + 404 => anyhow!("Repository '{repo_id}' or revision '{revision}' was not found"), + _ => anyhow!("HTTP {code} while fetching repository metadata for '{repo_id}'"), + }); + } + + let mut body = Vec::new(); + let mut stream = response.bytes_stream(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.with_context(|| { + format!("Stream error while reading repository metadata for '{repo_id}'") + })?; + if body.len().saturating_add(chunk.len()) > MAX_HF_METADATA_BYTES { + return Err(anyhow!( + "Repository metadata for '{repo_id}' exceeds the mlxcel {MAX_HF_METADATA_BYTES}-byte safety limit" + )); + } + body.extend_from_slice(&chunk); + } + serde_json::from_slice::(&body) + .with_context(|| format!("Failed to parse HuggingFace repository metadata for '{repo_id}'")) +} + +fn fetch_anonymous_repo_info_blocking( + repo_id: &str, + revision: Option<&str>, +) -> Result { + let endpoint = hf_endpoint(); + let requested_revision = revision.unwrap_or("main"); + let rt = tokio::runtime::Runtime::new().context("Failed to create tokio runtime")?; + let client = build_reqwest_client(None, false)?; + rt.block_on(fetch_anonymous_repo_info( + &client, + &endpoint, + repo_id, + requested_revision, + )) +} + /// Download a single file via reqwest streaming, ticking the per-file and /// aggregate progress bars as each chunk arrives. /// @@ -663,9 +839,12 @@ async fn stream_to_tempfile( )); } - // Prefer the response's own Content-Length over the HEAD-pass estimate: - // it reflects the entity actually being served. - let total_size = response.content_length().unwrap_or(expected_size); + // Prefer a known metadata/HEAD size for stable UI totals, then fall back to + // the response's own Content-Length for validation when no manifest size is + // available. + let response_size = response.content_length(); + let expected_size = (expected_size > 0).then_some(expected_size); + let total_size = expected_size.or(response_size).unwrap_or(0); let mut stream = response.bytes_stream(); let mut bytes_written: u64 = 0; @@ -690,6 +869,7 @@ async fn stream_to_tempfile( .await .with_context(|| format!("Flush error for {filename}"))?; drop(out); + validate_downloaded_size(filename, expected_size.or(response_size), bytes_written)?; tokio::fs::rename(tmp, dest) .await @@ -753,6 +933,16 @@ pub fn ensure_online(what: &str) -> Result<()> { /// download when `DELETE /models` removes the model. Both needs are optional /// observations of the same download loop, so they ride along as hooks rather /// than forking the implementation. +#[derive(Debug, Clone)] +pub struct DownloadPlan { + pub repo_id: String, + pub requested_revision: String, + pub resolved_revision: String, + pub destination: PathBuf, + pub selected_files: usize, + pub total_bytes: Option, +} + #[derive(Clone, Default)] pub struct DownloadHooks { /// Called as bytes arrive for each file: `(url, downloaded, total)`. @@ -760,6 +950,15 @@ pub struct DownloadHooks { /// report one terminal call with `downloaded == total`. Called from the /// download worker thread; keep it fast and non-blocking. pub progress: Option>, + /// Called after repository metadata and selected file sizes are known but + /// before the first transfer starts. `total_bytes` is `None` when one or + /// more selected files have unknown length. + pub plan: Option>, + /// Called immediately before an already-complete staged snapshot is made + /// visible. Returning `false` aborts publication as a cooperative cancel; + /// returning `true` is the linearization point after which cancellation may + /// be refused by the coordinator. + pub begin_publish: Option bool + Send + Sync>>, /// Cooperative cancellation flag, checked between streamed chunks and /// between files. When it becomes `true`, the download aborts with an /// error wrapping [`DownloadCancelled`], leaving no partial tempfiles @@ -771,6 +970,8 @@ impl std::fmt::Debug for DownloadHooks { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("DownloadHooks") .field("progress", &self.progress.is_some()) + .field("plan", &self.plan.is_some()) + .field("begin_publish", &self.begin_publish.is_some()) .field("cancel", &self.cancel.is_some()) .finish() } @@ -806,6 +1007,80 @@ fn check_cancelled(hooks: &DownloadHooks) -> Result<()> { Ok(()) } +fn validate_manifest_filename(repo_id: &str, filename: &str) -> Result<()> { + if filename.len() > MAX_REPO_FILENAME_BYTES { + return Err(anyhow!( + "Repository '{repo_id}' contains a filename longer than {MAX_REPO_FILENAME_BYTES} bytes; refusing to download it" + )); + } + Ok(()) +} + +fn validate_manifest_bounds( + repo_id: &str, + sibling_count: usize, + selected_count: usize, +) -> Result<()> { + if sibling_count > MAX_HF_SIBLINGS { + return Err(anyhow!( + "Repository '{repo_id}' exposes {sibling_count} files, above the mlxcel limit of {MAX_HF_SIBLINGS}" + )); + } + if selected_count > MAX_SELECTED_DOWNLOAD_FILES { + return Err(anyhow!( + "Repository '{repo_id}' selected {selected_count} files for download, above the mlxcel limit of {MAX_SELECTED_DOWNLOAD_FILES}" + )); + } + Ok(()) +} + +fn validate_downloaded_size( + filename: &str, + expected_size: Option, + bytes_written: u64, +) -> Result<()> { + if let Some(expected_size) = expected_size + && bytes_written != expected_size + { + return Err(anyhow!( + "Downloaded size mismatch for '{filename}': expected {expected_size} bytes, wrote {bytes_written} bytes" + )); + } + Ok(()) +} + +fn validate_sha256_hex(filename: &str, hash: &str) -> Result { + if hash.len() != 64 || !hash.as_bytes().iter().all(|byte| byte.is_ascii_hexdigit()) { + return Err(anyhow!( + "Repository metadata for '{filename}' contains an invalid SHA-256 digest" + )); + } + Ok(hash.to_ascii_lowercase()) +} + +fn sha256_hex(digest: impl AsRef<[u8]>) -> String { + let bytes = digest.as_ref(); + let mut out = String::with_capacity(bytes.len() * 2); + for byte in bytes { + use std::fmt::Write as _; + let _ = write!(&mut out, "{byte:02x}"); + } + out +} + +fn verify_downloaded_sha256(filename: &str, expected: Option<&str>, actual: &[u8]) -> Result<()> { + let Some(expected) = expected else { + return Ok(()); + }; + let actual = sha256_hex(actual); + if actual != expected { + return Err(anyhow!( + "Downloaded checksum mismatch for '{filename}': expected SHA-256 {expected}, got {actual}" + )); + } + Ok(()) +} + /// Probe that `repo_id` names a reachable HuggingFace model repository /// (issue #1438). Resolves the ambient token, builds the API client, and /// fetches the repository manifest without downloading any file. The router @@ -813,8 +1088,20 @@ fn check_cancelled(hooks: &DownloadHooks) -> Result<()> { /// starting the background download, mirroring b10621's synchronous metadata /// fetch in its own handler. Blocking: call from a blocking-capable thread. pub fn probe_repo(repo_id: &str, revision: Option<&str>) -> Result<()> { + probe_repo_with_token_mode(repo_id, revision, TokenMode::Environment) +} + +pub fn probe_repo_with_token_mode( + repo_id: &str, + revision: Option<&str>, + token_mode: TokenMode, +) -> Result<()> { ensure_online(&format!("the repository manifest for '{repo_id}'"))?; - let token = resolve_token(None); + if token_mode == TokenMode::Anonymous { + let repo_id = normalize_repo_id(repo_id)?; + return fetch_anonymous_repo_info_blocking(&repo_id, revision).map(|_| ()); + } + let token = resolve_token_with_mode(None, token_mode); let api = build_api(token)?; let repo = build_repo_handle(repo_id, revision); api.repo(repo) @@ -867,6 +1154,347 @@ pub fn download_repo_with_hooks(opts: DownloadOptions, hooks: DownloadHooks) -> download_repo_blocking(opts, hooks) } +#[cfg(unix)] +pub fn download_repo_to_existing_dir_fd( + repo_id: &str, + revision: Option<&str>, + dir_fd: std::os::fd::RawFd, + reported_destination: PathBuf, + hooks: DownloadHooks, +) -> Result<()> { + ensure_online(&format!("the model snapshot '{repo_id}'"))?; + let repo_id = normalize_repo_id(repo_id)?; + let requested_revision = revision.unwrap_or("main").to_string(); + let endpoint = hf_endpoint(); + let rt = tokio::runtime::Runtime::new().context("Failed to create tokio runtime")?; + let client = build_reqwest_client(None, false)?; + let info = rt.block_on(fetch_anonymous_repo_info( + &client, + &endpoint, + &repo_id, + &requested_revision, + ))?; + let resolved_revision = info.sha.clone(); + if info.siblings.len() > MAX_HF_SIBLINGS { + validate_manifest_bounds(&repo_id, info.siblings.len(), 0)?; + } + let mut wanted: Vec = info + .siblings + .iter() + .map(|sibling| SelectedDownloadFile::from_sibling(&repo_id, sibling)) + .collect::>>()? + .into_iter() + .filter(|file| is_wanted_file(&file.filename)) + .collect(); + validate_manifest_bounds(&repo_id, info.siblings.len(), wanted.len())?; + if wanted.is_empty() { + return Err(anyhow!( + "Repository '{repo_id}' does not expose any supported model files. Expected config/tokenizer JSON and safetensors/weights files." + )); + } + if !wanted + .iter() + .any(|file| file.filename.ends_with(".safetensors")) + { + return Err(anyhow!( + "Repository '{repo_id}' metadata is incomplete: no safetensors weight files were selected" + )); + } + let revision = resolved_revision.as_str(); + if (hooks.progress.is_some() || hooks.plan.is_some()) + && wanted.iter().any(|file| file.expected_size.is_none()) + { + let size_map = rt.block_on(async { + let client_ref = &client; + let endpoint_ref = &endpoint; + let repo_id_ref = &repo_id; + futures::stream::iter( + wanted + .iter() + .filter(|file| file.expected_size.is_none()) + .map(|file| file.filename.clone()), + ) + .map(move |filename| async move { + let url = file_url(endpoint_ref, repo_id_ref, revision, &filename); + let size = client_ref + .head(&url) + .send() + .await + .ok() + .and_then(|r| { + r.headers() + .get(reqwest::header::CONTENT_LENGTH) + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.parse::().ok()) + }) + .filter(|size| *size > 0); + (filename, size) + }) + .buffer_unordered(8) + .collect::>>() + .await + }); + for file in &mut wanted { + if file.expected_size.is_none() + && let Some(size) = size_map.get(&file.filename).and_then(|size| *size) + { + file.expected_size = Some(size); + } + } + } + let total_bytes = wanted + .iter() + .try_fold(0u64, |acc, file| { + file.expected_size + .map(|size| acc.saturating_add(size)) + .ok_or(()) + }) + .ok(); + if let Some(plan) = &hooks.plan { + plan(DownloadPlan { + repo_id: repo_id.clone(), + requested_revision: requested_revision.clone(), + resolved_revision: resolved_revision.clone(), + destination: reported_destination, + selected_files: wanted.len(), + total_bytes, + }); + } + for file in &wanted { + check_cancelled(&hooks)?; + let url = file_url(&endpoint, &repo_id, revision, &file.filename); + rt.block_on(stream_file_to_dir_fd( + &client, + &url, + dir_fd, + &file.filename, + file.expected_size, + file.sha256.as_deref(), + &hooks, + ))?; + } + if !file_exists_nonempty_at(dir_fd, "config.json")? { + return Err(anyhow!( + "Downloaded files for '{repo_id}' are incomplete: missing config.json" + )); + } + if !wanted + .iter() + .filter(|file| file.filename.ends_with(".safetensors")) + .all(|file| file_exists_nonempty_at(dir_fd, &file.filename).unwrap_or(false)) + { + return Err(anyhow!( + "Downloaded files for '{repo_id}' are incomplete: one or more safetensors files are missing or empty" + )); + } + unsafe { + libc::fsync(dir_fd); + } + Ok(()) +} + +#[cfg(unix)] +async fn stream_file_to_dir_fd( + client: &reqwest::Client, + url: &str, + root_fd: std::os::fd::RawFd, + filename: &str, + expected_size: Option, + expected_sha256: Option<&str>, + hooks: &DownloadHooks, +) -> Result { + stream_file_to_dir_fd_with_writer_factory( + client, + url, + root_fd, + filename, + expected_size, + expected_sha256, + hooks, + tokio::fs::File::from_std, + ) + .await +} + +#[cfg(unix)] +#[allow(clippy::too_many_arguments)] +async fn stream_file_to_dir_fd_with_writer_factory( + client: &reqwest::Client, + url: &str, + root_fd: std::os::fd::RawFd, + filename: &str, + expected_size: Option, + expected_sha256: Option<&str>, + hooks: &DownloadHooks, + writer_factory: F, +) -> Result +where + W: tokio::io::AsyncWrite + Unpin, + F: FnOnce(std::fs::File) -> W, +{ + use std::os::fd::{AsRawFd, FromRawFd}; + + let (parent, leaf) = open_parent_dir_fd(root_fd, filename)?; + let tmp_name = std::ffi::CString::new(format!( + ".mlxcel-partial.{}.{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or(0), + ))?; + let tmp_fd = unsafe { + libc::openat( + parent.as_raw_fd(), + tmp_name.as_ptr(), + libc::O_WRONLY | libc::O_CREAT | libc::O_EXCL | libc::O_CLOEXEC | libc::O_NOFOLLOW, + 0o600, + ) + }; + if tmp_fd < 0 { + return Err(std::io::Error::last_os_error()) + .with_context(|| format!("Failed to create tempfile for {filename}")); + } + let std_file = unsafe { std::fs::File::from_raw_fd(tmp_fd) }; + let mut out = writer_factory(std_file); + let result = async { + let response = client + .get(url) + .send() + .await + .with_context(|| format!("HTTP request failed for {filename}"))?; + let status = response.status(); + if !status.is_success() { + let code = status.as_u16(); + return Err(anyhow!("HTTP {code} downloading '{filename}'")); + } + let response_size = response.content_length(); + let total_size = expected_size.or(response_size).unwrap_or(0); + let mut stream = response.bytes_stream(); + let mut bytes_written = 0u64; + let mut hasher = expected_sha256.map(|_| Sha256::new()); + while let Some(chunk) = stream.next().await { + check_cancelled(hooks) + .with_context(|| format!("Cancelled while downloading {filename}"))?; + let chunk = + chunk.with_context(|| format!("Stream error while downloading {filename}"))?; + out.write_all(&chunk) + .await + .with_context(|| format!("Write error while downloading {filename}"))?; + if let Some(hasher) = hasher.as_mut() { + hasher.update(&chunk); + } + bytes_written += chunk.len() as u64; + if let Some(progress) = &hooks.progress { + progress(url, bytes_written, total_size); + } + } + out.flush() + .await + .with_context(|| format!("Flush error for {filename}"))?; + drop(out); + validate_downloaded_size(filename, expected_size.or(response_size), bytes_written)?; + if let Some(hasher) = hasher { + let digest = hasher.finalize(); + verify_downloaded_sha256(filename, expected_sha256, &digest)?; + } + let rc = unsafe { + libc::renameat( + parent.as_raw_fd(), + tmp_name.as_ptr(), + parent.as_raw_fd(), + leaf.as_ptr(), + ) + }; + if rc != 0 { + return Err(std::io::Error::last_os_error()) + .with_context(|| format!("Rename error for {filename}")); + } + Ok(bytes_written) + } + .await; + if result.is_err() { + unsafe { + libc::unlinkat(parent.as_raw_fd(), tmp_name.as_ptr(), 0); + } + } + result +} + +#[cfg(unix)] +fn open_parent_dir_fd( + root_fd: std::os::fd::RawFd, + filename: &str, +) -> Result<(std::fs::File, std::ffi::CString)> { + use std::os::fd::{AsRawFd, FromRawFd}; + let mut parts = filename.split('/').collect::>(); + let leaf = parts + .pop() + .ok_or_else(|| anyhow!("invalid empty filename in repository manifest"))?; + let leaf = cstring_repo_component(leaf)?; + let dup = unsafe { libc::dup(root_fd) }; + if dup < 0 { + return Err(std::io::Error::last_os_error()).context("failed to duplicate root fd"); + } + let mut current = unsafe { std::fs::File::from_raw_fd(dup) }; + for part in parts { + let name = cstring_repo_component(part)?; + let rc = unsafe { libc::mkdirat(current.as_raw_fd(), name.as_ptr(), 0o700) }; + if rc != 0 { + let err = std::io::Error::last_os_error(); + if err.raw_os_error() != Some(libc::EEXIST) { + return Err(err) + .with_context(|| format!("failed to create directory for {filename}")); + } + } + let fd = unsafe { + libc::openat( + current.as_raw_fd(), + name.as_ptr(), + libc::O_RDONLY | libc::O_DIRECTORY | libc::O_CLOEXEC | libc::O_NOFOLLOW, + ) + }; + if fd < 0 { + return Err(std::io::Error::last_os_error()) + .with_context(|| format!("failed to open directory for {filename}")); + } + current = unsafe { std::fs::File::from_raw_fd(fd) }; + } + Ok((current, leaf)) +} + +#[cfg(unix)] +fn cstring_repo_component(component: &str) -> Result { + if component.is_empty() || component == "." || component == ".." || component.contains('\\') { + return Err(anyhow!("unsafe repository filename component")); + } + std::ffi::CString::new(component).map_err(|_| anyhow!("repository filename contains NUL byte")) +} + +#[cfg(unix)] +fn file_exists_nonempty_at(root_fd: std::os::fd::RawFd, filename: &str) -> Result { + use std::os::fd::AsRawFd; + let (parent, leaf) = open_parent_dir_fd(root_fd, filename)?; + let mut stat = std::mem::MaybeUninit::::uninit(); + let rc = unsafe { + libc::fstatat( + parent.as_raw_fd(), + leaf.as_ptr(), + stat.as_mut_ptr(), + libc::AT_SYMLINK_NOFOLLOW, + ) + }; + if rc != 0 { + let err = std::io::Error::last_os_error(); + if err.raw_os_error() == Some(libc::ENOENT) { + return Ok(false); + } + return Err(err).context("failed to stat downloaded file"); + } + let stat = unsafe { stat.assume_init() }; + Ok((stat.st_mode & libc::S_IFMT) == libc::S_IFREG && stat.st_size > 0) +} + fn download_repo_blocking(opts: DownloadOptions, hooks: DownloadHooks) -> Result<()> { // Issue #171: expand a bare, prefix-less model name (e.g. `Qwen3-4B-4bit`) // to `/` BEFORE anything is derived from `opts.repo_id` — @@ -914,7 +1542,7 @@ fn download_repo_blocking(opts: DownloadOptions, hooks: DownloadHooks) -> Result } let local_dir = opts.resolve_local_dir(); - let token = resolve_token(opts.token.as_deref()); + let token = resolve_token_with_mode(opts.token.as_deref(), opts.token_mode); let endpoint = hf_endpoint(); // M1 — refuse plaintext endpoints when a token would be sent over the @@ -938,13 +1566,24 @@ fn download_repo_blocking(opts: DownloadOptions, hooks: DownloadHooks) -> Result let info = api_repo .info() .map_err(|err| map_hf_error(err, &opts.repo_id, opts.revision.as_deref(), None))?; + let requested_revision = opts.revision.as_deref().unwrap_or("main").to_string(); + let resolved_revision = info.sha.clone(); + if info.siblings.len() > MAX_HF_SIBLINGS { + validate_manifest_bounds(&opts.repo_id, info.siblings.len(), 0)?; + } let wanted: Vec = info .siblings .iter() - .map(|s| s.rfilename.clone()) + .map(|sibling| { + validate_manifest_filename(&opts.repo_id, &sibling.rfilename)?; + Ok(sibling.rfilename.clone()) + }) + .collect::>>()? + .into_iter() .filter(|name| is_wanted_file(name)) .filter(|name| matches_include_patterns(name, &include_patterns)) .collect(); + validate_manifest_bounds(&opts.repo_id, info.siblings.len(), wanted.len())?; if wanted.is_empty() { return Err(anyhow!( @@ -1021,38 +1660,9 @@ fn download_repo_blocking(opts: DownloadOptions, hooks: DownloadHooks) -> Result // so we honor their decision here as well. let enforce_https = token.is_some() && !is_insecure_endpoint_opt_out(); let rt = tokio::runtime::Runtime::new().context("Failed to create tokio runtime")?; - let client = rt.block_on(async { - let mut headers = reqwest::header::HeaderMap::new(); - if let Some(ref tok) = token { - // L3: reject non-ASCII or control-char tokens with - // a domain-specific error instead of the panic that the prior - // `.expect("token must be ASCII")` would produce. `validate_token` - // also rejects more characters than `HeaderValue::from_str` would - // strictly need (e.g., embedded `\r\n`), making it harder to - // smuggle header-injection payloads through a malformed env var. - validate_token(tok)?; - let auth_val = format!("Bearer {tok}"); - headers.insert( - reqwest::header::AUTHORIZATION, - reqwest::header::HeaderValue::from_str(&auth_val).with_context( - || "HF token contains invalid characters (must be ASCII, no control chars)", - )?, - ); - } - let mut builder = reqwest::Client::builder() - .connect_timeout(Duration::from_secs(10)) - .read_timeout(Duration::from_secs(30)) - .default_headers(headers); - if enforce_https { - builder = builder.https_only(true); - } - for cert in load_extra_ca_certificates()? { - builder = builder.add_root_certificate(cert); - } - builder.build().context("Failed to create HTTP client") - })?; + let client = build_reqwest_client(token.as_deref(), enforce_https)?; - let revision = opts.revision.as_deref().unwrap_or("main"); + let revision = resolved_revision.as_str(); // Build per-file sizes map for accurate bar lengths. `hf-hub 0.5` does // not expose per-file sizes in the manifest (Siblings only has `rfilename`), @@ -1066,38 +1676,56 @@ fn download_repo_blocking(opts: DownloadOptions, hooks: DownloadHooks) -> Result // The HEAD-request size pass also runs when a progress hook is installed // (issue #1438): the router's SSE `download_progress` events carry // per-file totals, which come from nowhere else. - let size_map: std::collections::HashMap = if show_bars || hooks.progress.is_some() + let size_map: std::collections::HashMap = + if show_bars || hooks.progress.is_some() || hooks.plan.is_some() { + rt.block_on(async { + let client_ref = &client; + let endpoint_ref = &endpoint; + let repo_id_ref = &opts.repo_id; + futures::stream::iter(wanted.iter().cloned()) + .map(move |filename| async move { + let url = file_url(endpoint_ref, repo_id_ref, revision, &filename); + let size = client_ref + .head(&url) + .send() + .await + .ok() + .and_then(|r| { + r.headers() + .get(reqwest::header::CONTENT_LENGTH) + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.parse::().ok()) + }) + .unwrap_or(0); + (filename, size) + }) + .buffer_unordered(8) + .collect::>() + .await + }) + } else { + std::collections::HashMap::new() + }; + + let total_known_bytes: u64 = size_map.values().sum(); + let total_bytes = if wanted + .iter() + .all(|name| size_map.get(name.as_str()).copied().unwrap_or(0) > 0) { - rt.block_on(async { - let client_ref = &client; - let endpoint_ref = &endpoint; - let repo_id_ref = &opts.repo_id; - futures::stream::iter(wanted.iter().cloned()) - .map(move |filename| async move { - let url = file_url(endpoint_ref, repo_id_ref, revision, &filename); - let size = client_ref - .head(&url) - .send() - .await - .ok() - .and_then(|r| { - r.headers() - .get(reqwest::header::CONTENT_LENGTH) - .and_then(|v| v.to_str().ok()) - .and_then(|s| s.parse::().ok()) - }) - .unwrap_or(0); - (filename, size) - }) - .buffer_unordered(8) - .collect::>() - .await - }) + Some(total_known_bytes) } else { - std::collections::HashMap::new() + None }; - - let total_known_bytes: u64 = size_map.values().sum(); + if let Some(plan) = &hooks.plan { + plan(DownloadPlan { + repo_id: opts.repo_id.clone(), + requested_revision: requested_revision.clone(), + resolved_revision: resolved_revision.clone(), + destination: local_dir.clone(), + selected_files: wanted.len(), + total_bytes, + }); + } let mp = progress::create_multi_progress(); let aggregate_pb = progress::add_aggregate_bar(&mp, total_known_bytes); diff --git a/src/downloader/resolver.rs b/src/downloader/resolver.rs index 2194a4254..0e04c2812 100644 --- a/src/downloader/resolver.rs +++ b/src/downloader/resolver.rs @@ -563,6 +563,7 @@ fn download_options( // HUGGING_FACE_HUB_TOKEN fallback; `Some` is the explicit // `--hf-token` (issue #1434), which outranks both. token: token.map(str::to_string), + token_mode: super::TokenMode::Environment, include: Vec::new(), force, } diff --git a/src/downloader/store.rs b/src/downloader/store.rs index 77e58f1c2..61e212cdb 100644 --- a/src/downloader/store.rs +++ b/src/downloader/store.rs @@ -391,6 +391,9 @@ fn list_models_under(models_root: &Path) -> Vec { let Some(top_name) = top.file_name().to_str().map(str::to_owned) else { continue; }; + if top_name.starts_with('.') { + continue; + } // Case 1: a bare-id snapshot stored directly at models/. if snapshot_is_complete(&top_path) { @@ -612,6 +615,12 @@ fn remove_model_under( return Err(RemoveError::OutsideStore(target)); } + if let Ok(meta) = std::fs::symlink_metadata(&target) + && meta.file_type().is_symlink() + { + return Err(RemoveError::OutsideStore(target)); + } + if target.is_dir() { let size = dir_size(&target); std::fs::remove_dir_all(&target).map_err(|source| RemoveError::Io { diff --git a/src/downloader/tests.rs b/src/downloader/tests.rs index 1eda9151a..aca2fee44 100644 --- a/src/downloader/tests.rs +++ b/src/downloader/tests.rs @@ -385,6 +385,708 @@ fn token_resolution_explicit_wins() { restore_env("HUGGING_FACE_HUB_TOKEN", prev_alt); } +#[test] +fn token_resolution_anonymous_mode_ignores_explicit_and_env_tokens() { + let _env_guard = env_lock(); + let prev_hf = std::env::var("HF_TOKEN").ok(); + let prev_alt = std::env::var("HUGGING_FACE_HUB_TOKEN").ok(); + // SAFETY: serialized via the crate-wide ENV_LOCK acquired above. + unsafe { + std::env::set_var("HF_TOKEN", "tok-hf"); + std::env::set_var("HUGGING_FACE_HUB_TOKEN", "tok-alt"); + } + assert_eq!( + resolve_token_with_mode(Some("explicit"), TokenMode::Anonymous), + None + ); + assert_eq!(resolve_token_with_mode(None, TokenMode::Anonymous), None); + restore_env("HF_TOKEN", prev_hf); + restore_env("HUGGING_FACE_HUB_TOKEN", prev_alt); +} + +#[test] +fn anonymous_probe_clears_saved_hf_cache_token() { + let _env_guard = env_lock(); + let prev_home = std::env::var("HF_HOME").ok(); + let prev_endpoint = std::env::var("HF_ENDPOINT").ok(); + let prev_hf = std::env::var("HF_TOKEN").ok(); + let prev_alt = std::env::var("HUGGING_FACE_HUB_TOKEN").ok(); + let home = tempfile::tempdir().expect("hf home"); + std::fs::write(home.path().join("token"), "hf_saved_token").expect("token file"); + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("listener"); + let addr = listener.local_addr().expect("addr"); + let endpoint = format!("http://{addr}"); + let captured = std::sync::Arc::new(std::sync::Mutex::new(String::new())); + let captured_for_thread = captured.clone(); + let server = std::thread::spawn(move || { + use std::io::{Read, Write}; + let (mut stream, _) = listener.accept().expect("accept"); + let mut buf = [0u8; 4096]; + let n = stream.read(&mut buf).expect("read request"); + *captured_for_thread.lock().expect("captured") = + String::from_utf8_lossy(&buf[..n]).into_owned(); + let body = r#"{"siblings":[],"sha":"0123456789abcdef0123456789abcdef01234567"}"#; + write!( + stream, + "HTTP/1.1 200 OK +Content-Type: application/json +Content-Length: {} +Connection: close + +{}", + body.len(), + body + ) + .expect("write response"); + }); + + // SAFETY: serialized via the crate-wide ENV_LOCK acquired above. + unsafe { + std::env::set_var("HF_HOME", home.path()); + std::env::set_var("HF_ENDPOINT", &endpoint); + std::env::remove_var("HF_TOKEN"); + std::env::remove_var("HUGGING_FACE_HUB_TOKEN"); + } + + let result = probe_repo_with_token_mode("owner/model", None, TokenMode::Anonymous); + + restore_env("HF_HOME", prev_home); + restore_env("HF_ENDPOINT", prev_endpoint); + restore_env("HF_TOKEN", prev_hf); + restore_env("HUGGING_FACE_HUB_TOKEN", prev_alt); + server.join().expect("server thread"); + + result.expect("anonymous probe should use the local fake endpoint"); + let request = captured.lock().expect("captured").clone(); + assert!( + !request.to_ascii_lowercase().contains("authorization:"), + "anonymous WebUI probe must not send saved Hub credentials; request was {request:?}" + ); +} + +#[cfg(unix)] +#[test] +fn anonymous_fd_download_sends_no_ambient_credentials_on_metadata_head_or_get() { + use std::io::{Read, Write}; + use std::os::fd::AsRawFd; + + let _env_guard = env_lock(); + let prev_home = std::env::var("HF_HOME").ok(); + let prev_endpoint = std::env::var("HF_ENDPOINT").ok(); + let prev_hf = std::env::var("HF_TOKEN").ok(); + let prev_alt = std::env::var("HUGGING_FACE_HUB_TOKEN").ok(); + let home = tempfile::tempdir().expect("hf home"); + std::fs::write(home.path().join("token"), "hf_saved_token").expect("token file"); + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("listener"); + let addr = listener.local_addr().expect("addr"); + let endpoint = format!("http://{addr}"); + let data_hash = { + use sha2::{Digest, Sha256}; + let digest = Sha256::digest(b"data"); + digest + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::() + }; + let captured = std::sync::Arc::new(std::sync::Mutex::new(Vec::::new())); + let captured_for_thread = captured.clone(); + let server = std::thread::spawn(move || { + for _ in 0..4 { + let (mut stream, _) = listener.accept().expect("accept"); + let mut buf = [0u8; 4096]; + let n = stream.read(&mut buf).expect("read request"); + let request = String::from_utf8_lossy(&buf[..n]).into_owned(); + captured_for_thread + .lock() + .expect("captured") + .push(request.clone()); + let first = request.lines().next().unwrap_or_default(); + let mut parts = first.split_whitespace(); + let method = parts.next().unwrap_or_default(); + let path = parts.next().unwrap_or_default(); + if path.starts_with("/api/models/owner/model/revision/main") { + let body = format!( + r#"{{"siblings":[{{"rfilename":"config.json"}},{{"rfilename":"model.safetensors","lfs":{{"sha256":"{data_hash}","size":4}}}}],"sha":"0123456789abcdef0123456789abcdef01234567"}}"# + ); + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ) + .expect("metadata response"); + } else if path.ends_with("/config.json") { + let body = b"{}"; + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + body.len() + ) + .expect("config header"); + if method != "HEAD" { + stream.write_all(body).expect("config body"); + } + } else if path.ends_with("/model.safetensors") { + let body = b"data"; + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + body.len() + ) + .expect("weights header"); + if method != "HEAD" { + stream.write_all(body).expect("weights body"); + } + } else { + write!( + stream, + "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + ) + .expect("not found response"); + } + } + }); + + let root = tempfile::tempdir().expect("download root"); + let root_dir = std::fs::File::open(root.path()).expect("open root fd"); + let progress_seen = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let progress_for_hook = progress_seen.clone(); + let hooks = DownloadHooks { + progress: Some(std::sync::Arc::new(move |_, _, _| { + progress_for_hook.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + })), + ..DownloadHooks::default() + }; + + unsafe { + std::env::set_var("HF_HOME", home.path()); + std::env::set_var("HF_ENDPOINT", &endpoint); + std::env::set_var("HF_TOKEN", "hf_env_token"); + std::env::set_var("HUGGING_FACE_HUB_TOKEN", "hf_alt_token"); + } + let result = download_repo_to_existing_dir_fd( + "owner/model", + None, + root_dir.as_raw_fd(), + root.path().to_path_buf(), + hooks, + ); + + restore_env("HF_HOME", prev_home); + restore_env("HF_ENDPOINT", prev_endpoint); + restore_env("HF_TOKEN", prev_hf); + restore_env("HUGGING_FACE_HUB_TOKEN", prev_alt); + if result.is_err() { + for _ in 0..4 { + let _ = std::net::TcpStream::connect(addr); + } + } + server.join().expect("server thread"); + + result.expect("anonymous fd download should use the local fake endpoint"); + assert!(root.path().join("config.json").is_file()); + assert!(root.path().join("model.safetensors").is_file()); + assert!( + progress_seen.load(std::sync::atomic::Ordering::SeqCst) >= 2, + "progress hook should observe metadata-sized and HEAD-sized GET progress" + ); + let requests = captured.lock().expect("captured"); + assert_eq!( + requests.len(), + 4, + "expected metadata, one HEAD fallback for config.json, and two GETs" + ); + let metadata_path = requests[0] + .lines() + .next() + .and_then(|line| line.split_whitespace().nth(1)) + .expect("metadata request path"); + assert_eq!( + metadata_path, "/api/models/owner/model/revision/main?blobs=true", + "anonymous WebUI metadata request must ask HuggingFace to include LFS blob size and SHA-256 fields" + ); + for request in requests.iter() { + assert!( + !request.to_ascii_lowercase().contains("authorization:"), + "anonymous WebUI metadata/HEAD/GET must not send ambient Hub credentials; request was {request:?}" + ); + } +} + +#[cfg(unix)] +#[test] +fn anonymous_fd_download_maps_metadata_status_without_publishing() { + use std::io::{Read, Write}; + use std::os::fd::AsRawFd; + + for (status, reason, expected) in [ + (403, "Forbidden", "requires authentication"), + (404, "Not Found", "was not found"), + ] { + let _env_guard = env_lock(); + let prev_endpoint = std::env::var("HF_ENDPOINT").ok(); + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("listener"); + let addr = listener.local_addr().expect("addr"); + let endpoint = format!("http://{addr}"); + let server = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("accept"); + let mut buf = [0u8; 1024]; + let _ = stream.read(&mut buf).expect("read request"); + write!( + stream, + "HTTP/1.1 {status} {reason}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + ) + .expect("status response"); + }); + + let root = tempfile::tempdir().expect("download root"); + let root_dir = std::fs::File::open(root.path()).expect("open root fd"); + unsafe { + std::env::set_var("HF_ENDPOINT", &endpoint); + } + let result = download_repo_to_existing_dir_fd( + "owner/model", + Some("bad-revision"), + root_dir.as_raw_fd(), + root.path().to_path_buf(), + DownloadHooks::default(), + ); + restore_env("HF_ENDPOINT", prev_endpoint); + server.join().expect("server thread"); + + let err = result.expect_err("metadata status must fail"); + assert!(err.to_string().contains(expected), "{err:#}"); + assert!( + std::fs::read_dir(root.path()) + .expect("root list") + .next() + .is_none(), + "metadata errors must not publish files" + ); + } +} + +#[cfg(unix)] +#[test] +fn anonymous_fd_download_offline_rejects_before_metadata_request() { + use std::os::fd::AsRawFd; + + let root = tempfile::tempdir().expect("download root"); + let root_dir = std::fs::File::open(root.path()).expect("open root fd"); + set_offline_mode(true); + let result = download_repo_to_existing_dir_fd( + "owner/model", + None, + root_dir.as_raw_fd(), + root.path().to_path_buf(), + DownloadHooks::default(), + ); + set_offline_mode(false); + + let err = result.expect_err("offline mode must fail before metadata"); + assert!(err.to_string().contains("offline mode is on"), "{err:#}"); + assert!( + std::fs::read_dir(root.path()) + .expect("root list") + .next() + .is_none(), + "offline rejection must not write files" + ); +} + +#[cfg(unix)] +#[test] +fn anonymous_fd_download_rejects_disconnect_truncation_before_publish() { + use std::io::{Read, Write}; + use std::os::fd::AsRawFd; + + let _env_guard = env_lock(); + let prev_endpoint = std::env::var("HF_ENDPOINT").ok(); + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("listener"); + let addr = listener.local_addr().expect("addr"); + let endpoint = format!("http://{addr}"); + let data_hash = { + use sha2::{Digest, Sha256}; + let digest = Sha256::digest(b"data"); + digest + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::() + }; + let server = std::thread::spawn(move || { + for _ in 0..3 { + let (mut stream, _) = listener.accept().expect("accept"); + let mut buf = [0u8; 4096]; + let n = stream.read(&mut buf).expect("read request"); + let request = String::from_utf8_lossy(&buf[..n]).into_owned(); + let path = request + .lines() + .next() + .and_then(|line| line.split_whitespace().nth(1)) + .unwrap_or_default(); + if path.starts_with("/api/models/owner/model/revision/main") { + let body = format!( + r#"{{"siblings":[{{"rfilename":"config.json","size":2}},{{"rfilename":"model.safetensors","lfs":{{"sha256":"{data_hash}","size":4}}}}],"sha":"0123456789abcdef0123456789abcdef01234567"}}"# + ); + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ) + .expect("metadata response"); + } else if path.ends_with("/config.json") { + let body = b"{}"; + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + body.len() + ) + .expect("config header"); + stream.write_all(body).expect("config body"); + } else if path.ends_with("/model.safetensors") { + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Length: 4\r\nConnection: close\r\n\r\nda" + ) + .expect("truncated weights response"); + } else { + write!( + stream, + "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + ) + .expect("not found response"); + } + } + }); + + let root = tempfile::tempdir().expect("download root"); + let root_dir = std::fs::File::open(root.path()).expect("open root fd"); + unsafe { + std::env::set_var("HF_ENDPOINT", &endpoint); + } + let result = download_repo_to_existing_dir_fd( + "owner/model", + None, + root_dir.as_raw_fd(), + root.path().to_path_buf(), + DownloadHooks::default(), + ); + restore_env("HF_ENDPOINT", prev_endpoint); + server.join().expect("server thread"); + + let err = result.expect_err("truncated transfer must fail"); + let rendered = err.to_string(); + assert!( + rendered.contains("size mismatch") || rendered.contains("Stream error"), + "{err:#}" + ); + assert!(root.path().join("config.json").is_file()); + assert!( + !root.path().join("model.safetensors").exists(), + "truncated weights must not be published" + ); +} + +#[cfg(unix)] +#[test] +fn fd_stream_timeout_keeps_partial_private_and_unpublished() { + use std::io::{Read, Write}; + use std::os::fd::AsRawFd; + + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("listener"); + let addr = listener.local_addr().expect("addr"); + let server = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("accept"); + let mut buf = [0u8; 1024]; + let _ = stream.read(&mut buf).expect("read request"); + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Length: 4\r\nConnection: close\r\n\r\n" + ) + .expect("write headers"); + stream.flush().expect("flush headers"); + std::thread::sleep(std::time::Duration::from_millis(500)); + }); + + let root = tempfile::tempdir().expect("download root"); + let root_dir = std::fs::File::open(root.path()).expect("open root fd"); + let client = reqwest::Client::builder() + .read_timeout(std::time::Duration::from_millis(50)) + .build() + .expect("client"); + let url = format!("http://{addr}/model.safetensors"); + let err = tokio::runtime::Runtime::new() + .expect("runtime") + .block_on(stream_file_to_dir_fd( + &client, + &url, + root_dir.as_raw_fd(), + "model.safetensors", + Some(4), + None, + &DownloadHooks::default(), + )) + .expect_err("stalled body must time out"); + server.join().expect("server thread"); + + assert!( + err.chain().any(|cause| cause + .downcast_ref::() + .is_some_and(reqwest::Error::is_timeout)), + "error chain should preserve reqwest timeout category: {err:#}" + ); + assert!( + !root.path().join("model.safetensors").exists(), + "timeout must not publish final file" + ); + let leftovers: Vec<_> = std::fs::read_dir(root.path()) + .expect("root list") + .filter_map(|entry| entry.ok().map(|entry| entry.file_name())) + .collect(); + assert!( + leftovers.is_empty(), + "partial files must be cleaned: {leftovers:?}" + ); +} + +#[cfg(unix)] +#[test] +fn fd_stream_enospc_writer_cleans_partial_without_publishing() { + use std::os::fd::AsRawFd; + use std::pin::Pin; + use std::task::{Context, Poll}; + + struct EnospcWriter { + _file: std::fs::File, + } + + impl tokio::io::AsyncWrite for EnospcWriter { + fn poll_write( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + _buf: &[u8], + ) -> Poll> { + Poll::Ready(Err(std::io::Error::from_raw_os_error(libc::ENOSPC))) + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("listener"); + let addr = listener.local_addr().expect("addr"); + let server = std::thread::spawn(move || { + use std::io::{Read, Write}; + let (mut stream, _) = listener.accept().expect("accept"); + let mut buf = [0u8; 1024]; + let _ = stream.read(&mut buf).expect("read request"); + let body = b"data"; + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + body.len() + ) + .expect("write headers"); + stream.write_all(body).expect("write body"); + }); + + let root = tempfile::tempdir().expect("download root"); + let root_dir = std::fs::File::open(root.path()).expect("open root fd"); + let client = reqwest::Client::builder().build().expect("client"); + let url = format!("http://{addr}/model.safetensors"); + let err = tokio::runtime::Runtime::new() + .expect("runtime") + .block_on(stream_file_to_dir_fd_with_writer_factory( + &client, + &url, + root_dir.as_raw_fd(), + "model.safetensors", + Some(4), + None, + &DownloadHooks::default(), + |file| EnospcWriter { _file: file }, + )) + .expect_err("simulated ENOSPC must fail"); + server.join().expect("server thread"); + + assert!( + err.chain().any(|cause| cause + .downcast_ref::() + .is_some_and(|io| io.raw_os_error() == Some(libc::ENOSPC))), + "error chain should preserve simulated ENOSPC: {err:#}" + ); + assert!( + !root.path().join("model.safetensors").exists(), + "simulated ENOSPC must not publish final file" + ); + let leftovers: Vec<_> = std::fs::read_dir(root.path()) + .expect("root list") + .filter_map(|entry| entry.ok().map(|entry| entry.file_name())) + .collect(); + assert!( + leftovers.is_empty(), + "partial files must be cleaned: {leftovers:?}" + ); +} + +#[cfg(unix)] +#[test] +fn anonymous_fd_download_rejects_metadata_without_weight_files_before_transfer() { + use std::io::{Read, Write}; + use std::os::fd::AsRawFd; + + let _env_guard = env_lock(); + let prev_endpoint = std::env::var("HF_ENDPOINT").ok(); + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("listener"); + let addr = listener.local_addr().expect("addr"); + let endpoint = format!("http://{addr}"); + let captured = std::sync::Arc::new(std::sync::Mutex::new(Vec::::new())); + let captured_for_thread = captured.clone(); + let server = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("accept"); + let mut buf = [0u8; 4096]; + let n = stream.read(&mut buf).expect("read request"); + captured_for_thread + .lock() + .expect("captured") + .push(String::from_utf8_lossy(&buf[..n]).into_owned()); + let body = r#"{"siblings":[{"rfilename":"config.json","size":2}],"sha":"0123456789abcdef0123456789abcdef01234567"}"#; + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ) + .expect("metadata response"); + }); + + let root = tempfile::tempdir().expect("download root"); + let root_dir = std::fs::File::open(root.path()).expect("open root fd"); + unsafe { + std::env::set_var("HF_ENDPOINT", &endpoint); + } + let result = download_repo_to_existing_dir_fd( + "owner/model", + None, + root_dir.as_raw_fd(), + root.path().to_path_buf(), + DownloadHooks::default(), + ); + restore_env("HF_ENDPOINT", prev_endpoint); + server.join().expect("server thread"); + + let err = result.expect_err("metadata without safetensors must fail"); + assert!(err.to_string().contains("no safetensors"), "{err:#}"); + assert_eq!(captured.lock().expect("captured").len(), 1); + assert!( + !root.path().join("config.json").exists(), + "incomplete metadata must fail before publishing config-only snapshots" + ); +} + +#[cfg(unix)] +#[test] +fn anonymous_fd_download_rejects_same_size_checksum_mismatch_before_publish() { + use std::io::{Read, Write}; + use std::os::fd::AsRawFd; + + let _env_guard = env_lock(); + let prev_endpoint = std::env::var("HF_ENDPOINT").ok(); + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("listener"); + let addr = listener.local_addr().expect("addr"); + let endpoint = format!("http://{addr}"); + let good_hash = { + use sha2::{Digest, Sha256}; + let digest = Sha256::digest(b"good"); + digest + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::() + }; + let server = std::thread::spawn(move || { + for _ in 0..3 { + let (mut stream, _) = listener.accept().expect("accept"); + let mut buf = [0u8; 4096]; + let n = stream.read(&mut buf).expect("read request"); + let request = String::from_utf8_lossy(&buf[..n]).into_owned(); + let path = request + .lines() + .next() + .and_then(|line| line.split_whitespace().nth(1)) + .unwrap_or_default(); + if path.starts_with("/api/models/owner/model/revision/main") { + assert_eq!( + path, "/api/models/owner/model/revision/main?blobs=true", + "anonymous WebUI metadata request must ask for LFS blob fields" + ); + let body = format!( + r#"{{"siblings":[{{"rfilename":"config.json","size":2}},{{"rfilename":"model.safetensors","lfs":{{"sha256":"{good_hash}","size":4}}}}],"sha":"0123456789abcdef0123456789abcdef01234567"}}"# + ); + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ) + .expect("metadata response"); + } else if path.ends_with("/config.json") { + let body = b"{}"; + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + body.len() + ) + .expect("config header"); + stream.write_all(body).expect("config body"); + } else if path.ends_with("/model.safetensors") { + let body = b"evil"; + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + body.len() + ) + .expect("weights header"); + stream.write_all(body).expect("weights body"); + } else { + write!( + stream, + "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + ) + .expect("not found response"); + } + } + }); + + let root = tempfile::tempdir().expect("download root"); + let root_dir = std::fs::File::open(root.path()).expect("open root fd"); + unsafe { + std::env::set_var("HF_ENDPOINT", &endpoint); + } + let result = download_repo_to_existing_dir_fd( + "owner/model", + None, + root_dir.as_raw_fd(), + root.path().to_path_buf(), + DownloadHooks::default(), + ); + restore_env("HF_ENDPOINT", prev_endpoint); + server.join().expect("server thread"); + + let err = result.expect_err("checksum mismatch must fail"); + assert!(err.to_string().contains("checksum mismatch"), "{err:#}"); + assert!(root.path().join("config.json").is_file()); + assert!( + !root.path().join("model.safetensors").exists(), + "corrupt same-size weights must not be published" + ); +} + #[test] fn token_resolution_hf_token_env() { let _env_guard = env_lock(); @@ -482,6 +1184,28 @@ fn snapshot_complete_rejects_zero_byte_files() { assert!(!snapshot_complete(dir.path(), &wanted)); } +#[cfg(unix)] +#[test] +fn fd_relative_download_parent_refuses_symlink_component() { + use std::os::fd::AsRawFd; + use std::os::unix::fs::symlink; + + let root = tempfile::tempdir().expect("root"); + let victim = tempfile::tempdir().expect("victim"); + symlink(victim.path(), root.path().join("vision")).expect("symlink"); + let root_dir = std::fs::File::open(root.path()).expect("open root"); + + let err = open_parent_dir_fd(root_dir.as_raw_fd(), "vision/config.json") + .expect_err("fd-relative helper must not follow symlinked directories"); + + assert!( + err.to_string().contains("failed to open directory") + || err.to_string().contains("Too many levels"), + "unexpected error: {err:#}" + ); + assert!(!victim.path().join("config.json").exists()); +} + /// Verify that `stream_file` removes the partial tempfile on error. /// /// A server that immediately drops the TCP connection after accept triggers a @@ -541,6 +1265,92 @@ async fn stream_file_cleans_up_tempfile_on_error() { ); } +#[tokio::test] +async fn stream_file_rejects_known_size_mismatch_before_publish() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + + let dir = tempfile::tempdir().unwrap(); + let dest = dir.path().join("model.safetensors"); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.expect("accept"); + let mut buf = [0u8; 1024]; + let _ = stream.read(&mut buf).await.expect("request"); + stream + .write_all( + b"HTTP/1.1 200 OK +Content-Length: 2 +Connection: close + +xy", + ) + .await + .expect("response"); + }); + + let client = reqwest::Client::new(); + let url = format!("http://{addr}/x"); + let mp = indicatif::MultiProgress::with_draw_target(indicatif::ProgressDrawTarget::hidden()); + let file_pb = mp.add(indicatif::ProgressBar::hidden()); + let agg_pb = mp.add(indicatif::ProgressBar::hidden()); + + let err = stream_file( + &client, + &url, + &dest, + "model.safetensors", + &file_pb, + &agg_pb, + 4, + &DownloadHooks::default(), + ) + .await + .expect_err("known HEAD size mismatch must fail"); + + assert!(err.to_string().contains("size mismatch"), "{err:#}"); + assert!(!dest.exists(), "mismatched file must not be published"); +} + +#[test] +fn selected_safetensors_requires_lfs_sha256() { + let missing = HubSibling { + rfilename: "model.safetensors".to_string(), + size: Some(4), + lfs: None, + }; + let err = SelectedDownloadFile::from_sibling("owner/model", &missing) + .expect_err("safetensors without LFS SHA-256 must be rejected"); + assert!( + err.to_string().contains("missing an LFS SHA-256"), + "{err:#}" + ); + + let present = HubSibling { + rfilename: "model.safetensors".to_string(), + size: None, + lfs: Some(HubLfs { + sha256: Some( + "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855".to_string(), + ), + size: Some(0), + }), + }; + let selected = SelectedDownloadFile::from_sibling("owner/model", &present) + .expect("sha256-backed safetensors accepted"); + assert_eq!(selected.expected_size, Some(0)); +} + +#[test] +fn manifest_bounds_reject_oversized_names_and_file_lists() { + assert!(validate_manifest_filename("owner/model", "config.json").is_ok()); + let long_name = format!("{}.json", "x".repeat(MAX_REPO_FILENAME_BYTES)); + assert!(validate_manifest_filename("owner/model", &long_name).is_err()); + assert!(validate_manifest_bounds("owner/model", MAX_HF_SIBLINGS + 1, 0).is_err()); + assert!(validate_manifest_bounds("owner/model", 10, MAX_SELECTED_DOWNLOAD_FILES + 1).is_err()); +} + // Integration test that hits the real Hugging Face Hub. Marked `#[ignore]` // per the issue acceptance criteria so CI does not depend on network access. #[test] @@ -555,6 +1365,7 @@ fn live_download_smoke_test() { models_dir: None, revision: None, token: None, + token_mode: TokenMode::Environment, include: Vec::new(), force: true, }; @@ -1135,6 +1946,7 @@ async fn download_repo_inside_runtime_errors_instead_of_aborting() { models_dir: None, revision: None, token: None, + token_mode: TokenMode::Environment, include: Vec::new(), force: true, }; diff --git a/src/server/router_cache.rs b/src/server/router_cache.rs index c166870cb..9df866d41 100644 --- a/src/server/router_cache.rs +++ b/src/server/router_cache.rs @@ -13,40 +13,43 @@ // limitations under the License. //! Router-mode model cache source (llama-server b10621, issue #1438). -//! -//! b10621's router lists its download cache as removable model entries -//! (`source: "cache"`, `can_remove: true`), lets `POST /models` download a -//! new HuggingFace repository into that cache, and lets `DELETE /models` -//! remove a cache entry from disk. mlxcel's equivalent cache is its model -//! store (the directory `--model-store-root` / `MLXCEL_MODELS_DIR` / -//! `MLXCEL_CACHE_DIR` resolve, holding `/` MLX snapshots), so -//! the router's cache source wraps that store. -//! -//! Confinement: every path in and out of the store goes through -//! [`crate::downloader::store`]'s sanitized `/` composition, and -//! removal re-asserts containment under the store root before deleting -//! (`remove_model_with_override`), so an HTTP-supplied model name can never -//! escape the cache directory in either direction. The downloader is behind a -//! trait so tests exercise the full download flow without the network. use std::path::{Path, PathBuf}; use std::sync::Arc; -use crate::downloader::{self, DownloadHooks, DownloadOptions}; +use anyhow::Context; + +use crate::downloader::{self, DownloadHooks, TokenMode}; + +pub const STAGING_DIR: &str = ".mlxcel-staging"; + +/// Compile-time availability of the anchored atomic mutation primitives. +/// This never probes or creates the configured cache directory. +#[cfg(feature = "webui")] +pub(crate) const fn managed_mutations_supported() -> bool { + cfg!(any( + target_os = "linux", + target_os = "macos", + target_os = "ios" + )) +} /// How a router downloads a repository into the cache. The production /// implementation is [`HfRouterDownloader`]; tests substitute a fake that /// writes files locally and drives the same hooks. pub trait RouterDownloader: Send + Sync { - /// Synchronously validate that `repo_id` names a fetchable repository - /// (b10621 validates by fetching metadata before answering `POST - /// /models`). Blocking; called from a blocking-capable thread. - fn validate(&self, repo_id: &str) -> anyhow::Result<()>; + /// Synchronously validate that `repo_id` names a fetchable repository. + fn validate(&self, repo_id: &str, revision: Option<&str>) -> anyhow::Result<()>; /// Download `repo_id` into the models root at `dest_root`, reporting /// progress and honoring cancellation through `hooks`. Blocking. - fn download(&self, repo_id: &str, dest_root: &Path, hooks: DownloadHooks) - -> anyhow::Result<()>; + fn download( + &self, + repo_id: &str, + revision: Option<&str>, + dest_root: &Path, + hooks: DownloadHooks, + ) -> anyhow::Result<()>; } /// The real HuggingFace-backed downloader. @@ -54,28 +57,52 @@ pub trait RouterDownloader: Send + Sync { pub struct HfRouterDownloader; impl RouterDownloader for HfRouterDownloader { - fn validate(&self, repo_id: &str) -> anyhow::Result<()> { - downloader::probe_repo(repo_id, None) + fn validate(&self, repo_id: &str, revision: Option<&str>) -> anyhow::Result<()> { + downloader::probe_repo_with_token_mode(repo_id, revision, TokenMode::Anonymous) } fn download( &self, repo_id: &str, + revision: Option<&str>, dest_root: &Path, hooks: DownloadHooks, ) -> anyhow::Result<()> { - downloader::download_repo_with_hooks( - DownloadOptions { - repo_id: repo_id.to_string(), - local_dir: None, - models_dir: Some(dest_root.to_path_buf()), - revision: None, - token: None, - include: Vec::new(), - force: false, - }, - hooks, - ) + #[cfg(not(unix))] + { + let _ = (repo_id, revision, dest_root, hooks); + anyhow::bail!("router-managed HuggingFace downloads require Unix directory-fd safety"); + } + #[cfg(unix)] + { + let final_dir = downloader::model_dir_with_override(repo_id, Some(dest_root)) + .ok_or_else(|| { + anyhow::anyhow!("cannot resolve cache destination for '{repo_id}'") + })?; + let mut stage = AnchoredStage::create(dest_root, repo_id)?; + let begin_publish = hooks.begin_publish.clone(); + let result = downloader::download_repo_to_existing_dir_fd( + repo_id, + revision, + stage.stage_fd(), + final_dir.clone(), + hooks, + ); + if let Err(err) = result { + stage.cleanup(); + return Err(err); + } + if let Some(begin_publish) = begin_publish + && !begin_publish() + { + stage.cleanup(); + return Err(anyhow::Error::new(crate::downloader::DownloadCancelled)); + } + stage.publish().with_context(|| { + format!("failed to publish downloaded snapshot for '{repo_id}'") + })?; + Ok(()) + } } } @@ -98,14 +125,10 @@ impl CacheSource { Self { root, downloader } } - /// The models root this cache lists, downloads into, and removes from. pub fn root(&self) -> &Path { &self.root } - /// Enumerate complete snapshots as `(repo_id, path)` pairs. Uses the - /// store's own listing (a directory counts only when it holds a - /// `config.json`), so a half-finished download is not offered as a model. pub fn list(&self) -> Vec<(String, PathBuf)> { crate::downloader::list_models_with_override(Some(&self.root)) .into_iter() @@ -114,58 +137,48 @@ impl CacheSource { .collect() } - /// The snapshot directory `repo_id` would occupy (whether or not it - /// exists yet). Repo-id segments are sanitized by the store so the - /// composed path cannot escape the root. pub fn snapshot_dir(&self, repo_id: &str) -> PathBuf { crate::downloader::model_dir_with_override(repo_id, Some(&self.root)) .unwrap_or_else(|| self.root.join(repo_id)) } - /// Normalize a requested model name into the repo id used as the cache - /// entry name (bare names expand to the default organization, exactly as - /// `-m ` resolution does). pub fn normalize_name(&self, name: &str) -> anyhow::Result { downloader::normalize_repo_id(name) } - /// Validate that `repo_id` is fetchable (metadata probe, no file - /// downloads). Blocking. - pub fn validate(&self, repo_id: &str) -> anyhow::Result<()> { - self.downloader.validate(repo_id) + pub fn validate(&self, repo_id: &str, revision: Option<&str>) -> anyhow::Result<()> { + self.downloader.validate(repo_id, revision) } - /// Download `repo_id` into the cache. Blocking; run on a worker thread. - pub fn download(&self, repo_id: &str, hooks: DownloadHooks) -> anyhow::Result<()> { - self.downloader.download(repo_id, &self.root, hooks) + pub fn download( + &self, + repo_id: &str, + revision: Option<&str>, + hooks: DownloadHooks, + ) -> anyhow::Result<()> { + self.downloader + .download(repo_id, revision, &self.root, hooks) } - /// Remove `repo_id`'s snapshot from the cache. Deleting is contained to - /// the store root by `remove_model_under`'s re-assertion; a missing - /// snapshot (for example a cancelled download that never completed a - /// file) is not an error, matching b10621's best-effort - /// `common_download_remove`. pub fn remove(&self, repo_id: &str) -> anyhow::Result<()> { - use crate::downloader::{RemoveOutcome, remove_model_with_override}; - match remove_model_with_override(repo_id, None, Some(&self.root)) { - Ok(RemoveOutcome::Removed { path, size_bytes }) => { - tracing::info!( - "router: removed cache model '{repo_id}' ({size_bytes} bytes) at {}", - path.display() - ); - Ok(()) - } - Ok(RemoveOutcome::HfCacheOnly { hf_path }) => { - // The read-only HuggingFace cache is not ours to manage; the - // router's cache never lists it, so this arm is unreachable - // from HTTP. Refuse rather than pretend. - anyhow::bail!( - "model '{repo_id}' only exists in the read-only HuggingFace cache at {}", - hf_path.display() - ) + #[cfg(not(unix))] + { + let _ = repo_id; + anyhow::bail!("router-managed cache removal requires Unix directory-fd safety"); + } + #[cfg(unix)] + { + match anchored_remove::remove_managed_snapshot(&self.root, repo_id)? { + Some(removed) => { + tracing::info!( + "router: removed managed cache model '{repo_id}' ({} bytes) at {}", + removed.size_bytes, + removed.path.display() + ); + Ok(()) + } + None => Ok(()), } - Ok(RemoveOutcome::NotFound) => Ok(()), - Err(err) => Err(anyhow::anyhow!(err.to_string())), } } } @@ -175,3 +188,22 @@ fn regular_file_exists(path: &Path) -> bool { .map(|meta| meta.file_type().is_file()) .unwrap_or(false) } + +#[cfg(unix)] +#[path = "router_cache/anchored_delete.rs"] +mod anchored_delete; + +#[cfg(unix)] +#[path = "router_cache/anchored_publish.rs"] +mod anchored_publish; + +#[cfg(unix)] +#[path = "router_cache/anchored_remove.rs"] +mod anchored_remove; + +#[cfg(unix)] +use anchored_publish::AnchoredStage; + +#[cfg(all(test, unix))] +#[path = "router_cache/restart_tests.rs"] +mod restart_tests; diff --git a/src/server/router_cache/anchored_delete.rs b/src/server/router_cache/anchored_delete.rs new file mode 100644 index 000000000..7527e6428 --- /dev/null +++ b/src/server/router_cache/anchored_delete.rs @@ -0,0 +1,212 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::anchored_publish::open_child_dir; +use std::ffi::{CStr, CString}; +use std::fs::File; +use std::io; +use std::mem::MaybeUninit; +use std::os::fd::AsRawFd; +use std::os::unix::fs::MetadataExt; + +pub(super) fn delete_tree_at(parent: &File, name: &CStr) -> io::Result { + let expected = stat_child(parent, name)?; + let mut hook = None; + delete_tree_at_verified(parent, name, &expected, &mut hook) +} + +#[cfg(test)] +pub(super) enum DeleteHookEvent { + BeforeOpenDir, + BeforeRemoveDir, +} + +#[cfg(test)] +pub(super) fn delete_contents_with_hook(dir: &File, mut hook: F) -> io::Result +where + F: FnMut(DeleteHookEvent, &CStr), +{ + let mut hook: Option<&mut dyn FnMut(DeleteHookEvent, &CStr)> = Some(&mut hook); + delete_contents_internal(dir, &mut hook) +} + +pub(super) fn delete_contents(dir: &File) -> io::Result { + let mut hook = None; + delete_contents_internal(dir, &mut hook) +} + +#[cfg(not(test))] +type DeleteHook<'a> = Option<&'a mut dyn FnMut((), &CStr)>; +#[cfg(test)] +type DeleteHook<'a> = Option<&'a mut dyn FnMut(DeleteHookEvent, &CStr)>; + +fn delete_tree_at_verified( + parent: &File, + name: &CStr, + expected: &libc::stat, + hook: &mut DeleteHook<'_>, +) -> io::Result { + #[cfg(test)] + if let Some(hook) = hook.as_deref_mut() { + hook(DeleteHookEvent::BeforeOpenDir, name); + } + let child = open_child_dir(parent, name)?; + verify_file_matches_stat(&child, expected)?; + let size_bytes = delete_contents_internal(&child, hook)?; + #[cfg(test)] + if let Some(hook) = hook.as_deref_mut() { + hook(DeleteHookEvent::BeforeRemoveDir, name); + } + verify_child_matches_file(parent, name, &child)?; + let rc = unsafe { libc::unlinkat(parent.as_raw_fd(), name.as_ptr(), libc::AT_REMOVEDIR) }; + if rc == 0 { + Ok(size_bytes) + } else { + Err(io::Error::last_os_error()) + } +} + +fn delete_contents_internal(dir: &File, hook: &mut DeleteHook<'_>) -> io::Result { + let dup_fd = unsafe { libc::dup(dir.as_raw_fd()) }; + if dup_fd < 0 { + return Err(io::Error::last_os_error()); + } + let dir_ptr = unsafe { libc::fdopendir(dup_fd) }; + if dir_ptr.is_null() { + let err = io::Error::last_os_error(); + unsafe { + libc::close(dup_fd); + } + return Err(err); + } + let _guard = DirCloser(dir_ptr); + let mut total = 0u64; + loop { + errno_reset(); + let entry = unsafe { libc::readdir(dir_ptr) }; + if entry.is_null() { + let err = io::Error::last_os_error(); + if err.raw_os_error().unwrap_or(0) == 0 { + return Ok(total); + } + return Err(err); + } + let name = unsafe { CStr::from_ptr((*entry).d_name.as_ptr()) }; + if name.to_bytes() == b"." || name.to_bytes() == b".." { + continue; + } + let name = CString::new(name.to_bytes()) + .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "NUL in dir entry"))?; + let stat = match stat_child(dir, &name) { + Ok(stat) => stat, + Err(err) if err.raw_os_error() == Some(libc::ENOENT) => continue, + Err(err) => return Err(err), + }; + if is_dir_mode(stat.st_mode) { + total = total.saturating_add(delete_tree_at_verified(dir, &name, &stat, hook)?); + } else { + if is_regular_mode(stat.st_mode) && stat.st_size > 0 { + total = total.saturating_add(stat.st_size as u64); + } + let rc = unsafe { libc::unlinkat(dir.as_raw_fd(), name.as_ptr(), 0) }; + if rc != 0 { + let err = io::Error::last_os_error(); + if err.raw_os_error() != Some(libc::ENOENT) { + return Err(err); + } + } + } + } +} + +fn stat_child(parent: &File, name: &CStr) -> io::Result { + let mut stat = MaybeUninit::::uninit(); + let rc = unsafe { + libc::fstatat( + parent.as_raw_fd(), + name.as_ptr(), + stat.as_mut_ptr(), + libc::AT_SYMLINK_NOFOLLOW, + ) + }; + if rc == 0 { + Ok(unsafe { stat.assume_init() }) + } else { + Err(io::Error::last_os_error()) + } +} + +// MetadataExt::dev returns u64, while libc::dev_t is signed on Darwin and +// already u64 on Linux. Preserve Rust's signed-to-u64 representation exactly; +// narrowing the metadata value could accept a different device identity. +#[allow(clippy::unnecessary_cast)] +fn stat_device_id(stat: &libc::stat) -> u64 { + stat.st_dev as u64 +} + +fn verify_file_matches_stat(file: &File, expected: &libc::stat) -> io::Result<()> { + let actual = file.metadata()?; + if actual.dev() == stat_device_id(expected) && actual.ino() == expected.st_ino { + Ok(()) + } else { + Err(identity_mismatch()) + } +} + +fn verify_child_matches_file(parent: &File, name: &CStr, held: &File) -> io::Result<()> { + let current = stat_child(parent, name)?; + let held = held.metadata()?; + if held.dev() == stat_device_id(¤t) && held.ino() == current.st_ino { + Ok(()) + } else { + Err(identity_mismatch()) + } +} + +fn identity_mismatch() -> io::Error { + io::Error::other("directory identity mismatch") +} + +fn errno_reset() { + #[cfg(any(target_os = "macos", target_os = "ios"))] + unsafe { + *libc::__error() = 0; + } + #[cfg(target_os = "linux")] + unsafe { + *libc::__errno_location() = 0; + } +} + +fn is_dir_mode(mode: libc::mode_t) -> bool { + (mode & libc::S_IFMT) == libc::S_IFDIR +} + +fn is_regular_mode(mode: libc::mode_t) -> bool { + (mode & libc::S_IFMT) == libc::S_IFREG +} + +struct DirCloser(*mut libc::DIR); + +impl Drop for DirCloser { + fn drop(&mut self) { + unsafe { + libc::closedir(self.0); + } + } +} + +#[cfg(test)] +#[path = "anchored_delete_tests.rs"] +mod tests; diff --git a/src/server/router_cache/anchored_delete_tests.rs b/src/server/router_cache/anchored_delete_tests.rs new file mode 100644 index 000000000..352f3c0d5 --- /dev/null +++ b/src/server/router_cache/anchored_delete_tests.rs @@ -0,0 +1,91 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::super::anchored_publish::open_dir_path; +use super::*; +use std::fs; + +#[test] +fn delete_contents_detects_nested_directory_swap_before_open() { + let root = tempfile::tempdir().expect("root"); + let nested = root.path().join("nested"); + fs::create_dir(&nested).expect("nested"); + fs::write(nested.join("file"), b"original").expect("file"); + let parent = open_dir_path(root.path()).expect("open root"); + let root_path = root.path().to_path_buf(); + let mut swapped = false; + + let err = delete_contents_with_hook(&parent, move |event, name| { + if !swapped + && matches!(event, DeleteHookEvent::BeforeOpenDir) + && name.to_bytes() == b"nested" + { + swapped = true; + fs::rename(root_path.join("nested"), root_path.join("nested-old")) + .expect("move original nested"); + fs::create_dir(root_path.join("nested")).expect("replacement nested"); + fs::write(root_path.join("nested/file"), b"replacement").expect("replacement file"); + } + }) + .expect_err("nested identity swap must fail"); + + assert!( + err.to_string().contains("identity mismatch"), + "unexpected error: {err}" + ); + assert_eq!( + fs::read(root.path().join("nested-old/file")).expect("original remains"), + b"original" + ); + assert_eq!( + fs::read(root.path().join("nested/file")).expect("replacement remains"), + b"replacement" + ); +} + +#[test] +fn identity_checks_preserve_device_and_inode_without_narrowing() { + let root = tempfile::tempdir().expect("root"); + fs::create_dir(root.path().join("child")).expect("child"); + let parent = open_dir_path(root.path()).expect("parent"); + let name = CString::new("child").expect("name"); + let held = open_child_dir(&parent, &name).expect("held"); + let mut expected = stat_child(&parent, &name).expect("stat"); + assert_eq!(stat_device_id(&expected), held.metadata().unwrap().dev()); + verify_file_matches_stat(&held, &expected).expect("same identity"); + verify_child_matches_file(&parent, &name, &held).expect("same child"); + + expected.st_dev ^= 1; + assert!(verify_file_matches_stat(&held, &expected).is_err()); + expected.st_dev ^= 1; + expected.st_ino ^= 1; + assert!(verify_file_matches_stat(&held, &expected).is_err()); + + fs::rename(root.path().join("child"), root.path().join("original")).expect("retain original"); + fs::create_dir(root.path().join("child")).expect("replacement"); + assert!(verify_child_matches_file(&parent, &name, &held).is_err()); +} + +#[cfg(any(target_os = "macos", target_os = "ios"))] +#[test] +fn darwin_signed_device_id_preserves_metadata_representation() { + let root = tempfile::tempdir().expect("root"); + let parent = open_dir_path(root.path()).expect("parent"); + let name = CString::new(".").expect("name"); + let mut stat = stat_child(&parent, &name).expect("stat"); + stat.st_dev = -1; + assert_eq!(stat_device_id(&stat), u64::MAX); + stat.st_dev = i32::MIN; + assert_eq!(stat_device_id(&stat), u64::MAX - i32::MAX as u64); +} diff --git a/src/server/router_cache/anchored_publish.rs b/src/server/router_cache/anchored_publish.rs new file mode 100644 index 000000000..be1652f4a --- /dev/null +++ b/src/server/router_cache/anchored_publish.rs @@ -0,0 +1,383 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::{STAGING_DIR, anchored_delete::delete_tree_at}; +use anyhow::{Context, anyhow}; +use std::ffi::{CStr, CString}; +use std::fs::{self, File, OpenOptions}; +use std::io; +use std::mem::MaybeUninit; +use std::os::fd::{AsRawFd, FromRawFd}; +use std::os::unix::fs::{MetadataExt, OpenOptionsExt}; +use std::path::Path; +use std::sync::atomic::{AtomicU64, Ordering}; + +static STAGE_COUNTER: AtomicU64 = AtomicU64::new(0); + +pub(super) struct AnchoredStage { + root: File, + staging_parent: File, + final_parent: File, + stage: File, + owner: CString, + final_name: CString, + stage_name: CString, + published: bool, +} + +impl AnchoredStage { + pub(super) fn create(root_path: &Path, repo_id: &str) -> anyhow::Result { + let (owner, final_name) = repo_segments(repo_id)?; + fs::create_dir_all(root_path).with_context(|| { + format!("failed to create model-store root {}", root_path.display()) + })?; + let root = open_dir_path(root_path).with_context(|| { + format!( + "model-store root {} must be a real directory, not a symlink", + root_path.display() + ) + })?; + let staging_name = cstring_segment(STAGING_DIR)?; + ensure_child_dir(&root, &staging_name, 0o700) + .context("failed to create router staging directory")?; + let staging_parent = open_child_dir(&root, &staging_name) + .context("router staging path must be a real directory")?; + verify_private_staging_parent(&staging_parent) + .context("router staging directory must be private to the server user")?; + ensure_child_dir(&root, &owner, 0o755) + .with_context(|| format!("failed to create owner directory for repo '{repo_id}'"))?; + let final_parent = open_child_dir(&root, &owner).with_context(|| { + format!("cache owner directory for repo '{repo_id}' must be a real directory") + })?; + if child_exists(&final_parent, &final_name)? { + anyhow::bail!( + "cache destination for '{repo_id}' already exists; refusing to overwrite it" + ); + } + let stage_name = create_unique_stage_dir(&staging_parent, repo_id)?; + let stage = open_child_dir(&staging_parent, &stage_name) + .context("created staging directory vanished before it could be opened")?; + Ok(Self { + root, + staging_parent, + final_parent, + stage, + owner, + final_name, + stage_name, + published: false, + }) + } + + pub(super) fn stage_fd(&self) -> std::os::fd::RawFd { + use std::os::fd::AsRawFd; + self.stage.as_raw_fd() + } + + pub(super) fn publish(&mut self) -> anyhow::Result<()> { + self.verify_anchors()?; + if child_exists(&self.final_parent, &self.final_name)? { + anyhow::bail!("cache destination appeared before publish; refusing to overwrite it"); + } + rename_no_replace( + &self.staging_parent, + &self.stage_name, + &self.final_parent, + &self.final_name, + ) + .context("atomic no-replace rename failed")?; + self.published = true; + Ok(()) + } + + pub(super) fn cleanup(&mut self) { + if self.published { + return; + } + if self.verify_staging_anchor().is_ok() && self.verify_stage_identity().is_ok() { + let _ = delete_tree_at(&self.staging_parent, &self.stage_name); + } + } + + fn verify_anchors(&self) -> anyhow::Result<()> { + self.verify_staging_anchor()?; + self.verify_final_parent_anchor()?; + self.verify_stage_identity()?; + Ok(()) + } + + fn verify_staging_anchor(&self) -> anyhow::Result<()> { + verify_child_identity( + &self.root, + &cstring_segment(STAGING_DIR)?, + &self.staging_parent, + ) + .context("router staging directory identity changed during download") + } + + fn verify_final_parent_anchor(&self) -> anyhow::Result<()> { + verify_child_identity(&self.root, &self.owner, &self.final_parent) + .context("cache owner directory identity changed during download") + } + + fn verify_stage_identity(&self) -> anyhow::Result<()> { + verify_child_identity(&self.staging_parent, &self.stage_name, &self.stage) + .context("staging directory identity changed during download") + } +} + +impl Drop for AnchoredStage { + fn drop(&mut self) { + self.cleanup(); + } +} + +pub(super) fn repo_segments(repo_id: &str) -> anyhow::Result<(CString, CString)> { + let mut parts = repo_id.split('/'); + let Some(owner) = parts.next() else { + anyhow::bail!("repo id must be owner/name"); + }; + let Some(name) = parts.next() else { + anyhow::bail!("repo id must be owner/name"); + }; + if parts.next().is_some() { + anyhow::bail!("repo id must be owner/name"); + } + Ok((cstring_segment(owner)?, cstring_segment(name)?)) +} + +pub(super) fn cstring_segment(segment: &str) -> anyhow::Result { + let valid = !segment.is_empty() + && segment != "." + && segment != ".." + && segment + .bytes() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'-')); + if !valid { + anyhow::bail!("unsafe repository path segment '{segment}'"); + } + CString::new(segment).map_err(|_| anyhow!("path segment contains NUL byte")) +} + +fn safe_stage_name(repo_id: &str, seq: u64) -> CString { + let safe = repo_id + .bytes() + .map(|b| { + if b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'-') { + b as char + } else { + '_' + } + }) + .collect::(); + CString::new(format!( + "{}-{}-{}-{safe}", + std::process::id(), + chrono::Utc::now().timestamp_nanos_opt().unwrap_or_default(), + seq + )) + .expect("generated stage name has no NUL") +} + +fn create_unique_stage_dir(parent: &File, repo_id: &str) -> anyhow::Result { + for _ in 0..64 { + let seq = STAGE_COUNTER.fetch_add(1, Ordering::Relaxed); + let name = safe_stage_name(repo_id, seq); + match mkdir_child(parent, &name, 0o700) { + Ok(()) => return Ok(name), + Err(err) if err.kind() == io::ErrorKind::AlreadyExists => continue, + Err(err) => return Err(err).context("failed to create private staging directory"), + } + } + anyhow::bail!("failed to allocate a unique private staging directory") +} + +pub(super) fn open_dir_path(path: &Path) -> io::Result { + OpenOptions::new() + .read(true) + .custom_flags(libc::O_DIRECTORY | libc::O_CLOEXEC | libc::O_NOFOLLOW) + .open(path) +} + +pub(super) fn open_child_dir(parent: &File, name: &CStr) -> io::Result { + let fd = unsafe { + libc::openat( + parent.as_raw_fd(), + name.as_ptr(), + libc::O_RDONLY | libc::O_DIRECTORY | libc::O_CLOEXEC | libc::O_NOFOLLOW, + ) + }; + if fd < 0 { + Err(io::Error::last_os_error()) + } else { + Ok(unsafe { File::from_raw_fd(fd) }) + } +} + +pub(super) fn ensure_child_dir(parent: &File, name: &CStr, mode: libc::mode_t) -> io::Result<()> { + match mkdir_child(parent, name, mode) { + Ok(()) => Ok(()), + Err(err) if err.kind() == io::ErrorKind::AlreadyExists => Ok(()), + Err(err) => Err(err), + } +} + +fn mkdir_child(parent: &File, name: &CStr, mode: libc::mode_t) -> io::Result<()> { + let rc = unsafe { libc::mkdirat(parent.as_raw_fd(), name.as_ptr(), mode) }; + if rc == 0 { + Ok(()) + } else { + Err(io::Error::last_os_error()) + } +} + +pub(super) fn child_exists(parent: &File, name: &CStr) -> io::Result { + let mut stat = MaybeUninit::::uninit(); + let rc = unsafe { + libc::fstatat( + parent.as_raw_fd(), + name.as_ptr(), + stat.as_mut_ptr(), + libc::AT_SYMLINK_NOFOLLOW, + ) + }; + if rc == 0 { + Ok(true) + } else { + let err = io::Error::last_os_error(); + if err.raw_os_error() == Some(libc::ENOENT) { + Ok(false) + } else { + Err(err) + } + } +} + +pub(super) fn verify_private_staging_parent(dir: &File) -> anyhow::Result<()> { + let meta = dir.metadata()?; + let mode = meta.mode() & 0o777; + let uid = meta.uid(); + let euid = unsafe { libc::geteuid() }; + if mode == 0o700 && uid == euid { + Ok(()) + } else { + anyhow::bail!( + "staging directory must be owned by uid {euid} with mode 0700; found uid {uid} mode {mode:o}" + ) + } +} + +pub(super) fn verify_child_identity(parent: &File, name: &CStr, held: &File) -> anyhow::Result<()> { + let current = open_child_dir(parent, name)?; + verify_same_file_identity(held, ¤t) +} + +pub(super) fn verify_same_file_identity(expected: &File, actual: &File) -> anyhow::Result<()> { + let expected_meta = expected.metadata()?; + let actual_meta = actual.metadata()?; + if expected_meta.dev() == actual_meta.dev() && expected_meta.ino() == actual_meta.ino() { + Ok(()) + } else { + anyhow::bail!("directory identity mismatch") + } +} + +pub(super) fn regular_file_nonempty_at(parent: &File, name: &CStr) -> io::Result { + let mut stat = MaybeUninit::::uninit(); + let rc = unsafe { + libc::fstatat( + parent.as_raw_fd(), + name.as_ptr(), + stat.as_mut_ptr(), + libc::AT_SYMLINK_NOFOLLOW, + ) + }; + if rc != 0 { + let err = io::Error::last_os_error(); + if err.raw_os_error() == Some(libc::ENOENT) { + return Ok(false); + } + return Err(err); + } + let stat = unsafe { stat.assume_init() }; + Ok(is_regular_mode(stat.st_mode) && stat.st_size > 0) +} + +#[cfg(target_os = "linux")] +pub(super) fn rename_no_replace( + from_parent: &File, + from_name: &CStr, + to_parent: &File, + to_name: &CStr, +) -> io::Result<()> { + let rc = unsafe { + libc::syscall( + libc::SYS_renameat2, + from_parent.as_raw_fd(), + from_name.as_ptr(), + to_parent.as_raw_fd(), + to_name.as_ptr(), + libc::RENAME_NOREPLACE, + ) + }; + if rc == 0 { + Ok(()) + } else { + Err(io::Error::last_os_error()) + } +} + +#[cfg(any(target_os = "macos", target_os = "ios"))] +pub(super) fn rename_no_replace( + from_parent: &File, + from_name: &CStr, + to_parent: &File, + to_name: &CStr, +) -> io::Result<()> { + let rc = unsafe { + libc::renameatx_np( + from_parent.as_raw_fd(), + from_name.as_ptr(), + to_parent.as_raw_fd(), + to_name.as_ptr(), + libc::RENAME_EXCL, + ) + }; + if rc == 0 { + Ok(()) + } else { + Err(io::Error::last_os_error()) + } +} + +#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "ios")))] +pub(super) fn rename_no_replace( + _from_parent: &File, + _from_name: &CStr, + _to_parent: &File, + _to_name: &CStr, +) -> io::Result<()> { + Err(io::Error::new( + io::ErrorKind::Unsupported, + "atomic no-replace directory rename is not available on this platform", + )) +} + +fn is_regular_mode(mode: libc::mode_t) -> bool { + (mode & libc::S_IFMT) == libc::S_IFREG +} + +#[cfg(test)] +#[path = "anchored_publish_tests.rs"] +mod tests; diff --git a/src/server/router_cache/anchored_publish_tests.rs b/src/server/router_cache/anchored_publish_tests.rs new file mode 100644 index 000000000..10b5236e5 --- /dev/null +++ b/src/server/router_cache/anchored_publish_tests.rs @@ -0,0 +1,77 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::*; +use std::ffi::OsStr; +use std::io::Write; +use std::os::fd::FromRawFd; +use std::os::unix::ffi::OsStrExt; + +#[test] +fn publish_refuses_destination_created_after_stage() { + let root = tempfile::tempdir().expect("root"); + let mut stage = AnchoredStage::create(root.path(), "owner/model").expect("stage"); + fs::create_dir_all(root.path().join("owner/model")).expect("raced dest"); + let err = stage.publish().expect_err("destination race must fail"); + assert!(err.to_string().contains("destination appeared")); +} + +#[test] +fn publish_refuses_owner_directory_identity_swap() { + let root = tempfile::tempdir().expect("root"); + let mut stage = AnchoredStage::create(root.path(), "owner/model").expect("stage"); + fs::rename(root.path().join("owner"), root.path().join("owner-old")).expect("move owner"); + fs::create_dir(root.path().join("owner")).expect("replacement owner"); + let err = stage.publish().expect_err("identity swap must fail"); + assert!(err.to_string().contains("owner directory identity changed")); + assert!(!root.path().join("owner/model").exists()); +} + +#[test] +fn cleanup_refuses_staging_parent_identity_swap() { + let root = tempfile::tempdir().expect("root"); + let mut stage = AnchoredStage::create(root.path(), "owner/model").expect("stage"); + let moved = root.path().join("staging-old"); + fs::rename(root.path().join(STAGING_DIR), &moved).expect("move staging parent"); + fs::create_dir(root.path().join(STAGING_DIR)).expect("replacement staging parent"); + fs::create_dir(root.path().join(STAGING_DIR).join("attacker")).expect("attacker"); + let old_stage_name = OsStr::from_bytes(stage.stage_name.to_bytes()).to_owned(); + stage.cleanup(); + assert!(root.path().join(STAGING_DIR).join("attacker").exists()); + assert!(moved.join(old_stage_name).exists()); +} + +#[test] +fn cleanup_unlinks_symlink_without_deleting_target() { + let root = tempfile::tempdir().expect("root"); + let mut stage = AnchoredStage::create(root.path(), "owner/model").expect("stage"); + let victim = root.path().join("victim"); + fs::write(&victim, b"keep").expect("victim"); + let fd = unsafe { + libc::openat( + stage.stage_fd(), + c"config.json".as_ptr(), + libc::O_WRONLY | libc::O_CREAT | libc::O_EXCL | libc::O_CLOEXEC, + 0o600, + ) + }; + assert!(fd >= 0, "openat config.json failed"); + let mut file = unsafe { File::from_raw_fd(fd) }; + writeln!(file, "{{}}").expect("write"); + let victim_c = std::ffi::CString::new(victim.as_os_str().as_bytes()).expect("victim path"); + let rc = unsafe { libc::symlinkat(victim_c.as_ptr(), stage.stage_fd(), c"link".as_ptr()) }; + assert_eq!(rc, 0, "symlinkat failed"); + stage.cleanup(); + assert_eq!(fs::read(&victim).expect("victim remains"), b"keep"); +} diff --git a/src/server/router_cache/anchored_remove.rs b/src/server/router_cache/anchored_remove.rs new file mode 100644 index 000000000..271de1b45 --- /dev/null +++ b/src/server/router_cache/anchored_remove.rs @@ -0,0 +1,262 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::anchored_delete::delete_contents; +use super::anchored_publish::{ + child_exists, cstring_segment, ensure_child_dir, open_child_dir, open_dir_path, + regular_file_nonempty_at, rename_no_replace, repo_segments, verify_child_identity, + verify_private_staging_parent, verify_same_file_identity, +}; +use anyhow::Context; +use std::ffi::{CStr, CString}; +use std::fs::File; +use std::io; +use std::os::fd::AsRawFd; +use std::os::unix::ffi::OsStrExt; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; + +const QUARANTINE_DIR: &str = ".mlxcel-quarantine"; +static REMOVE_COUNTER: AtomicU64 = AtomicU64::new(0); + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) struct ManagedRemoval { + pub(super) path: PathBuf, + pub(super) size_bytes: u64, +} + +pub(super) fn remove_managed_snapshot( + root_path: &Path, + repo_id: &str, +) -> anyhow::Result> { + remove_managed_snapshot_impl(root_path, repo_id, RemovalHooks::default()) +} + +#[derive(Default)] +struct RemovalHooks { + pre_rename: Option>, + post_rename: Option>, + before_final_unlink: Option>, +} + +#[cfg(test)] +pub(super) fn remove_managed_snapshot_with_pre_rename_hook( + root_path: &Path, + repo_id: &str, + hook: F, +) -> anyhow::Result> +where + F: FnOnce() + Send + 'static, +{ + remove_managed_snapshot_impl( + root_path, + repo_id, + RemovalHooks { + pre_rename: Some(Box::new(hook)), + ..RemovalHooks::default() + }, + ) +} + +#[cfg(test)] +pub(super) fn remove_managed_snapshot_with_post_rename_hook( + root_path: &Path, + repo_id: &str, + hook: F, +) -> anyhow::Result> +where + F: FnOnce() + Send + 'static, +{ + remove_managed_snapshot_impl( + root_path, + repo_id, + RemovalHooks { + post_rename: Some(Box::new(hook)), + ..RemovalHooks::default() + }, + ) +} + +#[cfg(test)] +pub(super) fn remove_managed_snapshot_with_final_unlink_hook( + root_path: &Path, + repo_id: &str, + hook: F, +) -> anyhow::Result> +where + F: FnOnce() + Send + 'static, +{ + remove_managed_snapshot_impl( + root_path, + repo_id, + RemovalHooks { + before_final_unlink: Some(Box::new(hook)), + ..RemovalHooks::default() + }, + ) +} + +fn remove_managed_snapshot_impl( + root_path: &Path, + repo_id: &str, + mut hooks: RemovalHooks, +) -> anyhow::Result> { + let (owner, final_name) = repo_segments(repo_id)?; + let root = match open_dir_path(root_path) { + Ok(root) => root, + Err(err) if err.raw_os_error() == Some(libc::ENOENT) => return Ok(None), + Err(err) => { + return Err(err).with_context(|| { + format!( + "model-store root {} must be a real directory, not a symlink", + root_path.display() + ) + }); + } + }; + let final_parent = match open_child_dir(&root, &owner) { + Ok(owner) => owner, + Err(err) if err.raw_os_error() == Some(libc::ENOENT) => return Ok(None), + Err(err) => { + return Err(err).with_context(|| { + format!("cache owner directory for repo '{repo_id}' must be a real directory") + }); + } + }; + let target = match open_child_dir(&final_parent, &final_name) { + Ok(target) => target, + Err(err) if err.raw_os_error() == Some(libc::ENOENT) => return Ok(None), + Err(err) => { + return Err(err).with_context(|| { + format!("cache snapshot for repo '{repo_id}' must be a real directory") + }); + } + }; + if !regular_file_nonempty_at(&target, c"config.json")? { + anyhow::bail!( + "refusing to remove cache snapshot for '{repo_id}' because it is missing a non-empty regular config.json" + ); + } + + let quarantine_name = cstring_segment(QUARANTINE_DIR)?; + ensure_child_dir(&root, &quarantine_name, 0o700) + .context("failed to create router removal quarantine directory")?; + let quarantine_parent = open_child_dir(&root, &quarantine_name) + .context("router removal quarantine path must be a real directory")?; + verify_private_staging_parent(&quarantine_parent) + .context("router removal quarantine directory must be private to the server user")?; + verify_child_identity(&root, &quarantine_name, &quarantine_parent) + .context("router removal quarantine directory identity changed before removal")?; + let quarantine_child = unique_available_child_name(&quarantine_parent, repo_id)?; + + verify_child_identity(&root, &owner, &final_parent) + .context("cache owner directory identity changed before removal")?; + verify_child_identity(&final_parent, &final_name, &target) + .context("cache snapshot directory identity changed before removal")?; + if let Some(hook) = hooks.pre_rename.take() { + hook(); + } + verify_child_identity(&root, &owner, &final_parent) + .context("cache owner directory identity changed before removal")?; + rename_no_replace( + &final_parent, + &final_name, + &quarantine_parent, + &quarantine_child, + ) + .context("failed to move cache snapshot into private removal quarantine")?; + if let Some(hook) = hooks.post_rename.take() { + hook(); + } + verify_child_identity(&root, &owner, &final_parent).context( + "cache owner directory identity changed after removal quarantine; leaving quarantined snapshot for inspection", + )?; + let moved = open_child_dir(&quarantine_parent, &quarantine_child) + .context("quarantined cache snapshot vanished before identity check")?; + verify_same_file_identity(&target, &moved).context( + "quarantined cache snapshot identity changed; leaving quarantine for inspection", + )?; + + // Deletion walks the held directory fd and unlinks children relative to that + // fd, so symlinks inside a managed snapshot are unlinked rather than + // followed. The final directory unlink still names the private quarantine + // child, so a same-UID process with write access to the 0700 quarantine can + // race that last name lookup; we re-check the child identity immediately + // before unlink and leave any mismatch undeleted. Stronger protection would + // require a platform API that removes a directory by fd. + let size_bytes = delete_contents(&target)?; + verify_child_identity(&quarantine_parent, &quarantine_child, &target) + .context("quarantined cache snapshot identity changed before final unlink")?; + if let Some(hook) = hooks.before_final_unlink.take() { + hook(); + } + verify_child_identity(&quarantine_parent, &quarantine_child, &target) + .context("quarantined cache snapshot identity changed before final unlink")?; + let rc = unsafe { + libc::unlinkat( + quarantine_parent.as_raw_fd(), + quarantine_child.as_ptr(), + libc::AT_REMOVEDIR, + ) + }; + if rc != 0 { + return Err(io::Error::last_os_error()) + .context("failed to remove quarantined cache directory"); + } + + Ok(Some(ManagedRemoval { + path: managed_snapshot_path(root_path, &owner, &final_name), + size_bytes, + })) +} + +fn unique_available_child_name(parent: &File, repo_id: &str) -> anyhow::Result { + for _ in 0..64 { + let seq = REMOVE_COUNTER.fetch_add(1, Ordering::Relaxed); + let name = safe_remove_name(repo_id, seq); + if !child_exists(parent, &name)? { + return Ok(name); + } + } + anyhow::bail!("failed to allocate a unique private quarantine name") +} + +fn safe_remove_name(repo_id: &str, seq: u64) -> CString { + let safe = repo_id + .bytes() + .map(|b| { + if b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'-') { + b as char + } else { + '_' + } + }) + .collect::(); + CString::new(format!( + "{}-{}-{}-{safe}", + std::process::id(), + chrono::Utc::now().timestamp_nanos_opt().unwrap_or_default(), + seq + )) + .expect("generated quarantine name has no NUL") +} + +fn managed_snapshot_path(root: &Path, owner: &CStr, final_name: &CStr) -> PathBuf { + root.join(std::ffi::OsStr::from_bytes(owner.to_bytes())) + .join(std::ffi::OsStr::from_bytes(final_name.to_bytes())) +} + +#[cfg(test)] +#[path = "anchored_remove_tests.rs"] +mod tests; diff --git a/src/server/router_cache/anchored_remove_tests.rs b/src/server/router_cache/anchored_remove_tests.rs new file mode 100644 index 000000000..3f31803cb --- /dev/null +++ b/src/server/router_cache/anchored_remove_tests.rs @@ -0,0 +1,217 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::*; +use std::fs; + +#[test] +fn managed_removal_deletes_complete_snapshot_from_private_quarantine() { + let root = tempfile::tempdir().expect("root"); + let snapshot = root.path().join("owner/model"); + fs::create_dir_all(&snapshot).expect("snapshot"); + fs::write(snapshot.join("config.json"), b"{}").expect("config"); + fs::write(snapshot.join("model.safetensors"), b"weights").expect("weights"); + + let removed = remove_managed_snapshot(root.path(), "owner/model") + .expect("remove") + .expect("removed"); + + assert_eq!(removed.path, snapshot); + assert!( + removed.size_bytes >= 9, + "removed size should include regular files" + ); + assert!(!root.path().join("owner/model").exists()); + assert!(root.path().join(QUARANTINE_DIR).is_dir()); + assert!( + fs::read_dir(root.path().join(QUARANTINE_DIR)) + .expect("quarantine") + .next() + .is_none(), + "private quarantine should be empty after successful delete" + ); +} + +#[test] +fn managed_removal_unlinks_symlink_child_without_deleting_target() { + let root = tempfile::tempdir().expect("root"); + let snapshot = root.path().join("owner/model"); + fs::create_dir_all(&snapshot).expect("snapshot"); + fs::write(snapshot.join("config.json"), b"{}").expect("config"); + let victim = root.path().join("external-victim"); + fs::write(&victim, b"keep").expect("victim"); + std::os::unix::fs::symlink(&victim, snapshot.join("link")).expect("symlink"); + + remove_managed_snapshot(root.path(), "owner/model") + .expect("remove") + .expect("removed"); + + assert_eq!(fs::read(&victim).expect("victim remains"), b"keep"); + assert!(!snapshot.exists()); +} + +#[test] +fn managed_removal_detects_source_swap_after_verify_and_leaves_quarantine() { + let root = tempfile::tempdir().expect("root"); + let owner = root.path().join("owner"); + let snapshot = owner.join("model"); + fs::create_dir_all(&snapshot).expect("snapshot"); + fs::write(snapshot.join("config.json"), b"original").expect("config"); + let owner_for_hook = owner.clone(); + + let err = remove_managed_snapshot_with_pre_rename_hook(root.path(), "owner/model", move || { + fs::rename( + owner_for_hook.join("model"), + owner_for_hook.join("original"), + ) + .expect("move original"); + fs::create_dir(owner_for_hook.join("model")).expect("replacement"); + fs::write(owner_for_hook.join("model/config.json"), b"replacement") + .expect("replacement config"); + }) + .expect_err("identity mismatch after rename must fail"); + + assert!( + err.to_string().contains("identity changed"), + "unexpected error: {err:#}" + ); + assert_eq!( + fs::read(owner.join("original/config.json")).expect("original remains"), + b"original" + ); + let quarantine = root.path().join(QUARANTINE_DIR); + let quarantined: Vec<_> = fs::read_dir(&quarantine) + .expect("quarantine") + .filter_map(|entry| entry.ok().map(|entry| entry.path())) + .collect(); + assert_eq!( + quarantined.len(), + 1, + "wrong inode should be left quarantined" + ); + assert_eq!( + fs::read(quarantined[0].join("config.json")).expect("replacement quarantined"), + b"replacement" + ); +} + +#[test] +fn managed_removal_refuses_parent_swap_before_quarantine_rename() { + let root = tempfile::tempdir().expect("root"); + let owner = root.path().join("owner"); + let snapshot = owner.join("model"); + fs::create_dir_all(&snapshot).expect("snapshot"); + fs::write(snapshot.join("config.json"), b"original").expect("config"); + let root_for_hook = root.path().to_path_buf(); + + let err = remove_managed_snapshot_with_pre_rename_hook(root.path(), "owner/model", move || { + fs::rename(root_for_hook.join("owner"), root_for_hook.join("owner-old")) + .expect("move owner"); + fs::create_dir_all(root_for_hook.join("owner/model")).expect("replacement snapshot"); + fs::write( + root_for_hook.join("owner/model/config.json"), + b"replacement", + ) + .expect("replacement config"); + }) + .expect_err("owner identity mismatch must fail before rename"); + + assert!( + err.to_string().contains("owner directory identity changed"), + "unexpected error: {err:#}" + ); + assert_eq!( + fs::read(root.path().join("owner-old/model/config.json")).expect("old owner remains"), + b"original" + ); + assert_eq!( + fs::read(root.path().join("owner/model/config.json")).expect("replacement remains"), + b"replacement" + ); +} + +#[test] +fn managed_removal_refuses_parent_swap_after_quarantine_rename() { + let root = tempfile::tempdir().expect("root"); + let owner = root.path().join("owner"); + let snapshot = owner.join("model"); + fs::create_dir_all(&snapshot).expect("snapshot"); + fs::write(snapshot.join("config.json"), b"original").expect("config"); + let root_for_hook = root.path().to_path_buf(); + + let err = + remove_managed_snapshot_with_post_rename_hook(root.path(), "owner/model", move || { + fs::rename(root_for_hook.join("owner"), root_for_hook.join("owner-old")) + .expect("move owner"); + fs::create_dir(root_for_hook.join("owner")).expect("replacement owner"); + }) + .expect_err("owner identity mismatch after rename must fail"); + + assert!( + err.to_string().contains("owner directory identity changed"), + "unexpected error: {err:#}" + ); + let quarantined: Vec<_> = fs::read_dir(root.path().join(QUARANTINE_DIR)) + .expect("quarantine") + .filter_map(|entry| entry.ok().map(|entry| entry.path())) + .collect(); + assert_eq!(quarantined.len(), 1, "snapshot should be left quarantined"); + assert_eq!( + fs::read(quarantined[0].join("config.json")).expect("quarantined original"), + b"original" + ); +} + +#[test] +fn managed_removal_refuses_quarantine_child_swap_before_final_unlink() { + let root = tempfile::tempdir().expect("root"); + let snapshot = root.path().join("owner/model"); + fs::create_dir_all(&snapshot).expect("snapshot"); + fs::write(snapshot.join("config.json"), b"original").expect("config"); + fs::write(snapshot.join("model.safetensors"), b"weights").expect("weights"); + let root_for_hook = root.path().to_path_buf(); + + let err = + remove_managed_snapshot_with_final_unlink_hook(root.path(), "owner/model", move || { + let quarantine = root_for_hook.join(QUARANTINE_DIR); + let entry = fs::read_dir(&quarantine) + .expect("quarantine") + .next() + .expect("quarantine child") + .expect("quarantine child"); + let child_name = entry.file_name(); + fs::rename(entry.path(), quarantine.join("saved-original")) + .expect("move quarantined child"); + fs::create_dir(quarantine.join(&child_name)).expect("replacement quarantine child"); + fs::write(quarantine.join(&child_name).join("marker"), b"replacement") + .expect("replacement marker"); + }) + .expect_err("quarantine child identity mismatch must fail"); + + assert!( + err.to_string().contains("identity changed"), + "unexpected error: {err:#}" + ); + let quarantine = root.path().join(QUARANTINE_DIR); + assert!(quarantine.join("saved-original").is_dir()); + let replacement = fs::read_dir(&quarantine) + .expect("quarantine") + .filter_map(|entry| entry.ok().map(|entry| entry.path())) + .find(|path| path.join("marker").exists()) + .expect("replacement child remains"); + assert_eq!( + fs::read(replacement.join("marker")).expect("replacement marker"), + b"replacement" + ); +} diff --git a/src/server/router_cache/restart_tests.rs b/src/server/router_cache/restart_tests.rs new file mode 100644 index 000000000..250ca4e05 --- /dev/null +++ b/src/server/router_cache/restart_tests.rs @@ -0,0 +1,379 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Owned test-process restarts, not production CLI or checkpoint acceptance. +//! Abrupt termination deliberately bypasses Drop: abandoned stages are retained, +//! never promoted, resumed, or cleaned by a different process at startup. + +use super::{AnchoredStage, CacheSource, RouterDownloader, STAGING_DIR}; +use crate::downloader::{DownloadHooks, DownloadPlan}; +use crate::server::ServerStartupConfig; +use crate::server::router_lifecycle::OperationState; +use crate::server::router_models::{RouterPool, RouterSources}; +use std::ffi::CString; +use std::fs::{self, File}; +use std::io::Write; +use std::os::fd::FromRawFd; +use std::path::{Path, PathBuf}; +use std::process::{Child, Command}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +const CHILD_TEST: &str = "server::router_cache::restart_tests::restart_child"; +const REPO: &str = "test-owner/restart-model"; +const REVISION: &str = "0123456789012345678901234567890123456789"; +const LIMIT: Duration = Duration::from_secs(20); + +struct LocalDownloader { + interrupt: bool, + evidence: PathBuf, +} + +impl RouterDownloader for LocalDownloader { + fn validate(&self, _: &str, _: Option<&str>) -> anyhow::Result<()> { + Ok(()) + } + + fn download( + &self, + repo: &str, + _: Option<&str>, + root: &Path, + hooks: DownloadHooks, + ) -> anyhow::Result<()> { + let mut stage = AnchoredStage::create(root, repo)?; + if let Some(plan) = &hooks.plan { + plan(DownloadPlan { + repo_id: repo.into(), + requested_revision: REVISION.into(), + resolved_revision: REVISION.into(), + destination: root.join(repo), + selected_files: 2, + total_bytes: Some(6), + }); + } + write_at(&stage, "config.json", b"{}")?; + if self.interrupt { + write_at(&stage, "weights.partial", b"partial")?; + fs::write(self.evidence.join("writer-ready"), b"ready")?; + // A bounded escape path makes an orphaned test process harmless. + std::thread::sleep(LIMIT); + anyhow::bail!("parent did not interrupt its owned test process"); + } + write_at(&stage, "model.safetensors", b"fake")?; + if let Some(begin) = &hooks.begin_publish { + anyhow::ensure!(begin(), "unexpected cancellation"); + } + stage.publish() + } +} + +fn write_at(stage: &AnchoredStage, name: &str, bytes: &[u8]) -> anyhow::Result<()> { + let name = CString::new(name)?; + // SAFETY: the stage owns its live directory descriptor and name is a C string. + let fd = unsafe { + libc::openat( + stage.stage_fd(), + name.as_ptr(), + libc::O_WRONLY | libc::O_CREAT | libc::O_EXCL | libc::O_NOFOLLOW, + 0o600, + ) + }; + anyhow::ensure!(fd >= 0, "openat: {}", std::io::Error::last_os_error()); + // SAFETY: successful openat returned a new owned descriptor. + unsafe { File::from_raw_fd(fd) }.write_all(bytes)?; + Ok(()) +} + +fn pool(root: &Path, evidence: &Path, interrupt: bool) -> Arc { + Arc::new( + RouterPool::new( + RouterSources { + cache: Some(CacheSource::new( + root.to_path_buf(), + Arc::new(LocalDownloader { + interrupt, + evidence: evidence.into(), + }), + )), + ..Default::default() + }, + ServerStartupConfig::default(), + Default::default(), + Default::default(), + 2, + false, + ) + .expect("CPU-only pool without any model load"), + ) +} + +async fn wait_for(mut condition: impl FnMut() -> bool) { + let deadline = Instant::now() + LIMIT; + while !condition() { + assert!( + Instant::now() < deadline, + "bounded child observation timed out" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } +} + +#[test] +fn restart_child() { + let Ok(phase) = std::env::var("MLXCEL_RESTART_TEST_PHASE") else { + return; + }; + let evidence = PathBuf::from(std::env::var_os("MLXCEL_RESTART_TEST_DIR").unwrap()); + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(child_phase(&evidence, &phase)); +} + +async fn child_phase(evidence: &Path, phase: &str) { + let root = evidence.join("store"); + let pool = pool(&root, evidence, phase == "interrupt"); + let coordinator = pool.lifecycle_coordinator(); + let instance = coordinator.server_instance_id(); + if phase == "interrupt" { + assert!( + !root.exists(), + "pool construction must not create the store" + ); + let first = pool + .submit_download(REPO, Some(REVISION), Some("restart-first")) + .unwrap(); + wait_for(|| evidence.join("writer-ready").exists()).await; + let queued = pool + .submit_download("test-owner/queued", Some(REVISION), Some("restart-queued")) + .unwrap(); + assert_eq!( + coordinator + .get_operation(&first.operation_id) + .unwrap() + .state, + OperationState::Running + ); + assert_eq!( + coordinator + .get_operation(&queued.operation_id) + .unwrap() + .state, + OperationState::Queued + ); + fs::write( + evidence.join("interrupted.json"), + serde_json::to_vec(&serde_json::json!({ + "instance": instance, "running": first.operation_id, "queued": queued.operation_id + })) + .unwrap(), + ) + .unwrap(); + fs::write(evidence.join("parent-ready"), b"ready").unwrap(); + tokio::time::sleep(LIMIT).await; + panic!("parent must terminate the owned child before its deadline"); + } + let interrupted: serde_json::Value = + serde_json::from_slice(&fs::read(evidence.join("interrupted.json")).unwrap()).unwrap(); + assert_ne!(instance, interrupted["instance"].as_str().unwrap()); + for key in ["running", "queued"] { + #[cfg(feature = "webui")] + assert_missing_operation(pool.clone(), interrupted[key].as_str().unwrap()).await; + assert!( + coordinator + .get_operation(interrupted[key].as_str().unwrap()) + .is_none(), + "old session history must not report terminal success" + ); + } + let stages = fs::read_dir(root.join(STAGING_DIR)).unwrap().count(); + assert_eq!(stages, 1, "startup must retain the abandoned private stage"); + if phase == "retry" { + assert!( + pool.catalog_snapshot().is_empty(), + "config-first partial must never enter catalog" + ); + let accepted = pool + .submit_download(REPO, Some(REVISION), Some("restart-first")) + .unwrap(); + assert_ne!( + (instance, accepted.operation_id.as_str()), + ( + interrupted["instance"].as_str().unwrap(), + interrupted["running"].as_str().unwrap() + ), + "operation identity is scoped by server instance" + ); + assert!( + !accepted.idempotent_replay, + "idempotency belongs to the new session" + ); + wait_for(|| { + coordinator + .get_operation(&accepted.operation_id) + .is_some_and(|op| op.state == OperationState::Succeeded) + }) + .await; + fs::write( + evidence.join("published.json"), + serde_json::to_vec(&serde_json::json!({ + "instance": instance, "terminal": accepted.operation_id, + "model_id": pool.catalog_snapshot()[0].ui_model_id + })) + .unwrap(), + ) + .unwrap(); + } else { + assert_eq!(phase, "discover"); + let published: serde_json::Value = + serde_json::from_slice(&fs::read(evidence.join("published.json")).unwrap()).unwrap(); + assert_ne!(instance, published["instance"].as_str().unwrap()); + assert_eq!( + pool.catalog_snapshot()[0].ui_model_id, + published["model_id"].as_str().unwrap() + ); + assert!( + coordinator + .get_operation(published["terminal"].as_str().unwrap()) + .is_none() + ); + #[cfg(feature = "webui")] + assert_missing_operation(pool.clone(), published["terminal"].as_str().unwrap()).await; + } + let catalog = pool.catalog_snapshot(); + assert_eq!(catalog.len(), 1); + assert_eq!(catalog[0].name, REPO); + assert_eq!( + fs::read(root.join(REPO).join("model.safetensors")).unwrap(), + b"fake" + ); + assert_eq!( + fs::read_dir(root.join(STAGING_DIR)).unwrap().count(), + stages + ); +} + +struct OwnedChild(Child); + +impl Drop for OwnedChild { + fn drop(&mut self) { + let _ = self.0.kill(); + let _ = self.0.wait(); + } +} + +fn spawn(evidence: &Path, phase: &str) -> OwnedChild { + OwnedChild( + Command::new(std::env::current_exe().unwrap()) + .args(["--exact", CHILD_TEST, "--nocapture"]) + .env("MLXCEL_RESTART_TEST_PHASE", phase) + .env("MLXCEL_RESTART_TEST_DIR", evidence) + .spawn() + .expect("owned CPU-only child"), + ) +} + +#[test] +fn killed_writer_and_terminal_history_reconcile_across_processes() { + let evidence = tempfile::tempdir().unwrap(); + let mut child = spawn(evidence.path(), "interrupt"); + let deadline = Instant::now() + LIMIT; + while !evidence.path().join("parent-ready").exists() { + assert!( + child.0.try_wait().unwrap().is_none(), + "child exited before interruption barrier" + ); + assert!( + Instant::now() < deadline, + "parent interruption barrier timed out" + ); + std::thread::sleep(Duration::from_millis(10)); + } + child.0.kill().unwrap(); + assert!(!child.0.wait().unwrap().success()); + let abandoned = fs::read_dir(evidence.path().join("store").join(STAGING_DIR)) + .unwrap() + .next() + .unwrap() + .unwrap() + .path(); + assert_eq!( + fs::read(abandoned.join("weights.partial")).unwrap(), + b"partial" + ); + for phase in ["retry", "discover"] { + let mut child = spawn(evidence.path(), phase); + let deadline = Instant::now() + LIMIT; + loop { + if let Some(status) = child.0.try_wait().unwrap() { + assert!(status.success(), "restart phase {phase}: {status}"); + assert_eq!( + fs::read(abandoned.join("weights.partial")).unwrap(), + b"partial" + ); + break; + } + assert!(Instant::now() < deadline, "restart phase {phase} timed out"); + std::thread::sleep(Duration::from_millis(10)); + } + } +} + +#[cfg(feature = "webui")] +async fn assert_missing_operation(pool: Arc, id: &str) { + use crate::server::config::ServerConfig; + use crate::server::router_server::{ + RouterServerState, create_router_app_with_authenticated_ui, + }; + use axum::body::Body; + use axum::http::{Request, StatusCode}; + use tower::ServiceExt; + + let config = ServerConfig { + api_keys: crate::server::resolve_api_keys(&["restart-test-key".into()], &[]).unwrap(), + ..Default::default() + }; + let app = create_router_app_with_authenticated_ui(RouterServerState { + pool, + config: Arc::new(config), + startup: Arc::new(ServerStartupConfig::default()), + catalog_cache: Arc::new(crate::server::webui::catalog::CatalogProjectionCache::new()), + }); + let response = app + .oneshot( + Request::builder() + .uri(format!("/ui-api/v1/operations/{id}")) + .header("authorization", "Bearer restart-test-key") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::NOT_FOUND); + let bytes = axum::body::to_bytes(response.into_body(), 4096) + .await + .unwrap(); + let mut actual: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + assert!(actual["request_id"].as_str().unwrap().starts_with("req_")); + actual["request_id"] = serde_json::json!(""); + assert_eq!( + actual, + serde_json::json!({ + "request_id": "", + "error": {"code": "not_found", "message": "operation not found", "retryable": true} + }) + ); +} diff --git a/src/server/router_contract_test_support.rs b/src/server/router_contract_test_support.rs index 6586ad091..ad3a34e48 100644 --- a/src/server/router_contract_test_support.rs +++ b/src/server/router_contract_test_support.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +#![allow(dead_code)] + //! Whole-value producer comparisons against fixtures validated by the pinned //! JSON Schema gate. Only dynamic IDs, timestamps and revisions are normalized; //! missing fields, extra fields and all static values must match exactly. @@ -43,6 +45,39 @@ fn operation_fixture() -> Value { )) } +fn download_running_fixture() -> Value { + fixture(include_str!( + "../../tests/fixtures/webui/examples/operation.download-running.json" + )) +} + +fn download_succeeded_fixture() -> Value { + fixture(include_str!( + "../../tests/fixtures/webui/examples/operation.download-succeeded.json" + )) +} + +fn operation_accepted_fixture() -> Value { + fixture(include_str!( + "../../tests/fixtures/webui/examples/operation.accepted.json" + )) +} + +fn download_operation_dynamic(prefix: &str, include_result: bool) -> Vec<(String, Dynamic)> { + let mut fields = vec![ + ("/operation_id", Dynamic::Token), + ("/created_at", Dynamic::Timestamp), + ("/updated_at", Dynamic::Timestamp), + ]; + if include_result { + fields.push(("/result/model_id", Dynamic::ModelId)); + } + fields + .into_iter() + .map(|(path, kind)| (format!("{prefix}{path}"), kind)) + .collect() +} + fn operation_dynamic(prefix: &str) -> Vec<(String, Dynamic)> { [ ("/operation_id", Dynamic::Token), @@ -118,6 +153,44 @@ fn matches_fixture( Ok(()) } +pub(super) fn assert_operation_accepted(value: &Value) { + let expected = operation_accepted_fixture(); + let dynamic = [("/operation_id".into(), Dynamic::Token)]; + matches_fixture(value, &expected, &dynamic).unwrap_or_else(|error| panic!("{error}")); +} + +pub(super) fn assert_download_operation_running(value: &Value) { + matches_fixture( + value, + &download_running_fixture(), + &download_operation_dynamic("", false), + ) + .unwrap_or_else(|error| panic!("{error}")); +} + +pub(super) fn assert_download_operation_succeeded(value: &Value) { + matches_fixture( + value, + &download_succeeded_fixture(), + &download_operation_dynamic("", true), + ) + .unwrap_or_else(|error| panic!("{error}")); +} + +pub(super) fn assert_download_progress_event(value: &Value) { + let expected = fixture(include_str!( + "../../tests/fixtures/webui/examples/event.2.json" + )); + let dynamic = [ + ("/server_instance_id".into(), Dynamic::Token), + ("/event_id".into(), Dynamic::EventId), + ("/sequence".into(), Dynamic::Sequence), + ("/emitted_at".into(), Dynamic::Timestamp), + ("/payload/operation_id".into(), Dynamic::Token), + ]; + matches_fixture(value, &expected, &dynamic).unwrap_or_else(|error| panic!("{error}")); +} + pub(super) fn assert_operation(operation: &Value, target: &str, eviction_target: &str) { for path in ["/target/model_id", "/result/model_id"] { assert_eq!( diff --git a/src/server/router_library_integration_tests.rs b/src/server/router_library_integration_tests.rs new file mode 100644 index 000000000..b8d34e524 --- /dev/null +++ b/src/server/router_library_integration_tests.rs @@ -0,0 +1,242 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::*; +use axum::http::StatusCode; + +fn prefixed_library(root: &Path, cache: bool) -> Router { + let state = router_state_with_startup( + RouterSources { + cache: cache.then(|| { + crate::server::router_cache::CacheSource::new( + root.into(), + Arc::new(NoBootstrapDownload), + ) + }), + ..Default::default() + }, + ServerConfig { + api_prefix: "/lab".into(), + enable_settings_endpoint: true, + ..keyed_config() + }, + ServerStartupConfig { + webui_enabled: true, + model_store_root: Some(root.into()), + ..Default::default() + }, + false, + ); + let policy = + crate::server::webui::security::WebUiSecurityPolicy::with_prefixes_limits_and_rate( + vec!["127.0.0.1:18037".into()], + vec![HeaderValue::from_static("http://127.0.0.1:18037")], + "/lab/webui", + "/lab", + 32, + 16, + 120, + ) + .unwrap(); + create_router_app_with_secured_ui(state, policy) +} + +#[allow(clippy::too_many_arguments)] +async fn request( + app: Router, + method: &str, + path: &str, + host: &str, + origin: Option<&str>, + fetch: Option<&str>, + authenticated: bool, + body: &str, +) -> axum::response::Response { + let mut request = Request::builder() + .method(method) + .uri(path) + .header("host", host) + .header("content-type", "application/json"); + if authenticated { + request = request.header("authorization", format!("Bearer {ROUTER_KEY}")); + } + if let Some(origin) = origin { + request = request.header("origin", origin); + } + if let Some(fetch) = fetch { + request = request.header("sec-fetch-site", fetch); + } + app.oneshot(request.body(Body::from(body.to_string())).unwrap()) + .await + .unwrap() +} + +#[tokio::test] +async fn prefixed_library_is_secured_once_and_observation_does_not_create_cache() { + let temp = tempfile::tempdir().unwrap(); + let absent = temp.path().join("absent-store"); + let app = prefixed_library(&absent, true); + let response = request( + app.clone(), + "GET", + "/lab/ui-api/v1/bootstrap", + "127.0.0.1:18037", + None, + None, + true, + "", + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + let bytes = axum::body::to_bytes(response.into_body(), 32 * 1024) + .await + .unwrap(); + let mut actual: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + let mut expected: serde_json::Value = serde_json::from_str(include_str!( + "../../tests/fixtures/webui/examples/bootstrap.model-free.json" + )) + .unwrap(); + expected["server"]["api_base"] = serde_json::json!("/lab"); + normalize_bootstrap_fixture(&mut actual, &expected); + assert_eq!( + actual, expected, + "entire prefixed actual bootstrap producer" + ); + let response = request( + app.clone(), + "GET", + "/lab/ui-api/v1/catalog", + "127.0.0.1:18037", + None, + None, + true, + "", + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + for suffix in ["downloads", "model-removals"] { + let path = format!("/lab/ui-api/v1/{suffix}"); + for (host, origin, fetch, auth, expected) in [ + ( + "127.0.0.1:18037", + None, + None, + false, + StatusCode::UNAUTHORIZED, + ), + ("foreign.invalid", None, None, true, StatusCode::FORBIDDEN), + ( + "127.0.0.1:18037", + Some("http://foreign.invalid"), + None, + true, + StatusCode::FORBIDDEN, + ), + ( + "127.0.0.1:18037", + None, + Some("cross-site"), + true, + StatusCode::FORBIDDEN, + ), + ( + "127.0.0.1:18037", + None, + Some("same-origin"), + true, + StatusCode::BAD_REQUEST, + ), + ] { + let response = + request(app.clone(), "POST", &path, host, origin, fetch, auth, "{}").await; + assert_eq!( + response.status(), + expected, + "{path}: host={host} origin={origin:?} fetch={fetch:?}" + ); + } + for path in [ + format!("/ui-api/v1/{suffix}"), + format!("/lab/lab/ui-api/v1/{suffix}"), + ] { + let response = request( + app.clone(), + "POST", + &path, + "127.0.0.1:18037", + None, + None, + true, + "{}", + ) + .await; + // Unknown routes deliberately retain the legacy dispatcher, whose + // missing-model error differs from the mounted library JSON parser. + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let bytes = axum::body::to_bytes(response.into_body(), 4096) + .await + .unwrap(); + let body: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!( + body["error"]["message"], + "model name is missing from the request" + ); + } + } + assert!( + !absent.exists(), + "bootstrap/catalog/refused mutations must not create cache/staging or probe downloader" + ); +} + +#[tokio::test] +async fn no_cache_library_admission_is_canonical_unsupported_without_mutation() { + let temp = tempfile::tempdir().unwrap(); + let absent = temp.path().join("absent-store"); + let app = prefixed_library(&absent, false); + for (suffix, body) in [ + ( + "downloads", + serde_json::json!({"repo_id":"test-owner/model","idempotency_key":"no-cache-download"}), + ), + ( + "model-removals", + serde_json::json!({"model_id":format!("mdl_{}", "a".repeat(43)),"expected_revision":1,"idempotency_key":"no-cache-removal"}), + ), + ] { + let response = request( + app.clone(), + "POST", + &format!("/lab/ui-api/v1/{suffix}"), + "127.0.0.1:18037", + None, + None, + true, + &body.to_string(), + ) + .await; + assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY); + let bytes = axum::body::to_bytes(response.into_body(), 4096) + .await + .unwrap(); + let mut actual: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + assert!(actual["request_id"].as_str().unwrap().starts_with("req_")); + actual["request_id"] = serde_json::json!(""); + assert_eq!( + actual, + serde_json::json!({"request_id":"","error":{"code":"unsupported","message":"no managed model cache is configured","retryable":false}}) + ); + } + assert!(!absent.exists()); +} diff --git a/src/server/router_lifecycle.rs b/src/server/router_lifecycle.rs index 5ec3d69c5..f2e23c93d 100644 --- a/src/server/router_lifecycle.rs +++ b/src/server/router_lifecycle.rs @@ -88,8 +88,8 @@ pub use super::router_lifecycle_dto::{ RuntimeSnapshot, SettingsPayload, UiEvent, UiEventPayload, }; pub use super::router_lifecycle_ops::{ - EVENT_RING_LIMIT, LifecycleCoordinator, MAX_SAFE_EVENT_SEQUENCE, ReplaySubscribeError, - ResetEventKind, UiReplayCursor, + EVENT_RING_LIMIT, LifecycleCoordinator, MAX_ACTIVE_DOWNLOAD_OPERATIONS, + MAX_SAFE_EVENT_SEQUENCE, ReplaySubscribeError, ResetEventKind, UiReplayCursor, }; #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -112,6 +112,7 @@ struct LifecycleInner { generation: u64, admission_stopped: bool, worker_exit_observed: bool, + operation_busy: bool, last_error: Option, } @@ -154,6 +155,7 @@ impl ModelLifecycle { generation: revision, admission_stopped: false, worker_exit_observed: true, + operation_busy: false, last_error: None, }), notify: tokio::sync::Notify::new(), @@ -166,6 +168,12 @@ impl ModelLifecycle { self.operation_lock.lock().await } + pub fn try_operation_guard( + &self, + ) -> Result, tokio::sync::TryLockError> { + self.operation_lock.try_lock() + } + pub fn revision(&self) -> u64 { self.inner.lock().map(|g| g.revision).unwrap_or(1) } @@ -200,7 +208,8 @@ impl ModelLifecycle { | ModelLifecycleState::Unloading ) || (guard.state == ModelLifecycleState::Ready && guard.active_requests > 0) || (guard.state == ModelLifecycleState::Failed && !guard.worker_exit_observed) - || guard.download == DownloadState::Downloading; + || guard.download == DownloadState::Downloading + || guard.operation_busy; let draining_requests = if guard.admission_stopped { guard.active_requests } else { @@ -228,6 +237,7 @@ impl ModelLifecycle { | ModelLifecycleState::Draining | ModelLifecycleState::Unloading ) || (guard.state == ModelLifecycleState::Failed && !guard.worker_exit_observed) + || guard.operation_busy }) .unwrap_or(true) } @@ -253,6 +263,24 @@ impl ModelLifecycle { }); } + pub fn mark_operation_busy(&self) { + self.mutate(|g| { + if !g.operation_busy { + g.operation_busy = true; + g.revision = self.next_revision_after(g.revision); + } + }); + } + + pub fn mark_operation_idle(&self) { + self.mutate(|g| { + if g.operation_busy { + g.operation_busy = false; + g.revision = self.next_revision_after(g.revision); + } + }); + } + pub fn mark_loading(&self) { self.mutate(|g| { g.state = ModelLifecycleState::Loading; @@ -297,19 +325,29 @@ impl ModelLifecycle { } pub fn begin_drain(&self) -> bool { - let mut changed = false; + self.begin_drain_with_revision() + .is_some_and(|(changed, _)| changed) + } + + /// Capture the owned transition revision under the same lifecycle lock. + /// A caller revalidating after drain must use this token, not its pre-drain + /// revision or a guessed increment of the shared revision authority. + /// Poisoning returns no token, so mutation callers can fail closed. + pub(crate) fn begin_drain_with_revision(&self) -> Option<(bool, u64)> { + let mut transition = None; self.mutate(|g| { - if matches!( + let changed = matches!( g.state, ModelLifecycleState::Loading | ModelLifecycleState::Ready - ) { + ); + if changed { g.state = ModelLifecycleState::Draining; g.admission_stopped = true; g.revision = self.next_revision_after(g.revision); - changed = true; } + transition = Some((changed, g.revision)); }); - changed + transition } pub fn mark_unloading(&self) { diff --git a/src/server/router_lifecycle_dto.rs b/src/server/router_lifecycle_dto.rs index d70e2ffe1..4a55e65f3 100644 --- a/src/server/router_lifecycle_dto.rs +++ b/src/server/router_lifecycle_dto.rs @@ -108,7 +108,7 @@ impl OperationTarget { scope == token || model_id.as_ref().is_some_and(|id| id == token) } Self::Download { repo_id, revision } => { - repo_id == token || revision.as_ref().is_some_and(|rev| rev == token) + repo_id == token || revision.as_deref() == Some(token) } } } diff --git a/src/server/router_lifecycle_ops.rs b/src/server/router_lifecycle_ops.rs index f6cb87cd0..0ee5cff9d 100644 --- a/src/server/router_lifecycle_ops.rs +++ b/src/server/router_lifecycle_ops.rs @@ -15,7 +15,8 @@ //! Bounded WebUI operation and event store shared by router lifecycle adapters. use std::collections::{BTreeMap, BTreeSet, VecDeque}; -use std::sync::Mutex; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use super::router_lifecycle::LifecycleSnapshot; @@ -27,6 +28,7 @@ pub const MAX_SAFE_EVENT_SEQUENCE: u64 = 9_007_199_254_740_991; pub const TERMINAL_OPERATION_LIMIT: usize = 200; pub const TERMINAL_OPERATION_RETENTION: Duration = Duration::from_secs(3600); pub const MAX_ACTIVE_OPERATIONS: usize = 64; +pub const MAX_ACTIVE_DOWNLOAD_OPERATIONS: usize = 4; #[derive(Debug, Clone)] struct IdempotencyRecord { @@ -39,6 +41,8 @@ struct CoordinatorInner { active: BTreeMap, terminal: VecDeque<(Instant, Operation)>, idempotency: BTreeMap, + download_aliases: BTreeMap, + cancellations: BTreeMap>, events: VecDeque<(Instant, UiEvent)>, next_operation: u64, next_sequence: u64, @@ -185,12 +189,24 @@ impl LifecycleCoordinator { idempotency_key: Option<&str>, fingerprint: String, ) -> Result { - let now = rfc3339_now(); let mut inner = self .inner .lock() .map_err(|_| OperationError::TooManyActive)?; - prune_terminal(&mut inner); + self.begin_operation_locked(&mut inner, kind, target, idempotency_key, fingerprint, None) + } + + fn begin_operation_locked( + &self, + inner: &mut CoordinatorInner, + kind: OperationKind, + target: OperationTarget, + idempotency_key: Option<&str>, + fingerprint: String, + max_active_for_kind: Option, + ) -> Result { + let now = rfc3339_now(); + prune_terminal(inner); if let Some(key) = idempotency_key && let Some(record) = inner.idempotency.get(key) { @@ -218,6 +234,11 @@ impl LifecycleCoordinator { if inner.active.len() >= MAX_ACTIVE_OPERATIONS { return Err(OperationError::TooManyActive); } + if let Some(max_active_for_kind) = max_active_for_kind + && inner.active.values().filter(|op| op.kind == kind).count() >= max_active_for_kind + { + return Err(OperationError::TooManyActive); + } inner.next_operation += 1; let operation_id = format!("op_{}_{:06}", kind.as_str(), inner.next_operation); let operation = Operation { @@ -243,19 +264,177 @@ impl LifecycleCoordinator { }, ); } + if let OperationTarget::Download { + repo_id, + revision: Some(revision), + } = &operation.target + { + inner + .download_aliases + .insert(download_alias_key(repo_id, revision), operation_id.clone()); + } inner.active.insert(operation_id.clone(), operation.clone()); self.append_and_broadcast_locked( - &mut inner, + inner, UiEventPayload::Operation(OperationPayload { operation: operation.clone(), }), ); - let accepted = OperationAccepted { + Ok(OperationAccepted { operation_id, state: OperationState::Queued, idempotent_replay: false, + }) + } + + pub fn begin_download_operation( + &self, + target: OperationTarget, + idempotency_key: Option<&str>, + fingerprint: String, + ) -> Result { + self.begin_operation_with_kind_limit( + OperationKind::Download, + target, + idempotency_key, + fingerprint, + MAX_ACTIVE_DOWNLOAD_OPERATIONS, + ) + } + + fn begin_operation_with_kind_limit( + &self, + kind: OperationKind, + target: OperationTarget, + idempotency_key: Option<&str>, + fingerprint: String, + max_active_for_kind: usize, + ) -> Result { + let mut inner = self + .inner + .lock() + .map_err(|_| OperationError::TooManyActive)?; + self.begin_operation_locked( + &mut inner, + kind, + target, + idempotency_key, + fingerprint, + Some(max_active_for_kind), + ) + } + + pub fn idempotency_conflict(&self, key: &str, fingerprint: &str) -> Option { + let inner = self.inner.lock().ok()?; + inner.idempotency.get(key).and_then(|record| { + (record.fingerprint != fingerprint).then(|| record.operation_id.clone()) + }) + } + + pub fn find_active_download(&self, repo_id: &str, revision: &str) -> Option { + let inner = self.inner.lock().ok()?; + if let Some(operation_id) = inner + .download_aliases + .get(&download_alias_key(repo_id, revision)) + && let Some(operation) = inner.active.get(operation_id) + { + return Some(operation.clone()); + } + inner + .active + .values() + .find_map(|operation| match &operation.target { + OperationTarget::Download { + repo_id: target_repo, + revision: target_revision, + } if target_repo == repo_id && target_revision.as_deref() == Some(revision) => { + Some(operation.clone()) + } + _ => None, + }) + } + + pub fn register_download_revision_alias( + &self, + operation_id: &str, + repo_id: &str, + revision: &str, + ) -> bool { + let Ok(mut inner) = self.inner.lock() else { + return false; + }; + if !inner + .active + .get(operation_id) + .is_some_and(|operation| operation.kind == OperationKind::Download) + { + return false; + } + inner.download_aliases.insert( + download_alias_key(repo_id, revision), + operation_id.to_string(), + ); + true + } + + pub fn register_cancellation(&self, operation_id: &str, cancel: Arc) { + let Ok(mut inner) = self.inner.lock() else { + return; + }; + inner.cancellations.insert(operation_id.to_string(), cancel); + if let Some(operation) = inner.active.get_mut(operation_id) { + operation.cancellable = true; + operation.cancel_reason = None; + operation.updated_at = rfc3339_now(); + let cloned = operation.clone(); + self.append_and_broadcast_locked( + &mut inner, + UiEventPayload::Operation(OperationPayload { operation: cloned }), + ); + } + } + + pub fn begin_publish_operation(&self, operation_id: &str) -> bool { + let Ok(mut inner) = self.inner.lock() else { + return false; }; - Ok(accepted) + let Some(operation) = inner.active.get_mut(operation_id) else { + return false; + }; + if operation.state == OperationState::Cancelling { + return false; + } + operation.cancellable = false; + operation.cancel_reason = + Some("download is being published and can no longer be cancelled".to_string()); + operation.updated_at = rfc3339_now(); + let cloned = operation.clone(); + inner.cancellations.remove(operation_id); + self.append_and_broadcast_locked( + &mut inner, + UiEventPayload::Operation(OperationPayload { operation: cloned }), + ); + true + } + + pub fn update_operation_progress( + &self, + operation_id: &str, + progress: ProgressBytes, + ) -> Option { + let mut inner = self.inner.lock().ok()?; + let operation = inner.active.get_mut(operation_id)?; + operation.progress = progress.clone(); + operation.updated_at = rfc3339_now(); + let cloned = operation.clone(); + self.append_and_broadcast_locked( + &mut inner, + UiEventPayload::DownloadProgress(DownloadProgressPayload { + operation_id: operation_id.to_string(), + progress, + }), + ); + Some(cloned) } pub fn update_operation( @@ -277,6 +456,8 @@ impl LifecycleCoordinator { ) { operation.cancellable = false; operation.cancel_reason = None; + inner.cancellations.remove(operation_id); + inner.download_aliases.retain(|_, id| id != operation_id); inner .terminal .push_back((Instant::now(), operation.clone())); @@ -349,11 +530,41 @@ impl LifecycleCoordinator { } pub fn cancel_operation(&self, operation_id: &str) -> Result { - let operation = self - .get_operation(operation_id) + let mut inner = self.inner.lock().map_err(|_| CancelError::NotFound)?; + let Some(cancel) = inner.cancellations.get(operation_id).cloned() else { + let operation = inner + .active + .get(operation_id) + .or_else(|| { + inner + .terminal + .iter() + .find_map(|(_, op)| (op.operation_id == operation_id).then_some(op)) + }) + .ok_or(CancelError::NotFound)?; + return Err(CancelError::Unsupported { + operation_id: operation.operation_id.clone(), + }); + }; + cancel.store(true, Ordering::Relaxed); + let operation = inner + .active + .get_mut(operation_id) .ok_or(CancelError::NotFound)?; - Err(CancelError::Unsupported { - operation_id: operation.operation_id, + operation.state = OperationState::Cancelling; + operation.updated_at = rfc3339_now(); + operation.cancellable = false; + operation.cancel_reason = + Some("cancellation accepted; waiting for the writer to exit".to_string()); + let cloned = operation.clone(); + self.append_and_broadcast_locked( + &mut inner, + UiEventPayload::Operation(OperationPayload { operation: cloned }), + ); + Ok(OperationAccepted { + operation_id: operation_id.to_string(), + state: OperationState::Cancelling, + idempotent_replay: false, }) } @@ -516,6 +727,13 @@ fn prune_idempotency(inner: &mut CoordinatorInner) { inner .idempotency .retain(|_, record| retained_operations.contains(&record.operation_id)); + inner + .download_aliases + .retain(|_, operation_id| inner.active.contains_key(operation_id)); +} + +fn download_alias_key(repo_id: &str, revision: &str) -> String { + format!("{repo_id}::{revision}") } fn prune_events(inner: &mut CoordinatorInner) { diff --git a/src/server/router_lifecycle_tests.rs b/src/server/router_lifecycle_tests.rs index 14fc0253a..d4331a2f4 100644 --- a/src/server/router_lifecycle_tests.rs +++ b/src/server/router_lifecycle_tests.rs @@ -15,8 +15,13 @@ use std::sync::Arc; use std::time::Duration; +use super::super::router_lifecycle_dto::DownloadProgressPayload; use super::*; +#[allow(clippy::duplicate_mod)] +#[path = "router_contract_test_support.rs"] +mod contract; + #[derive(serde::Deserialize)] struct IdentityVectors { identity_vectors: Vec, @@ -222,6 +227,8 @@ fn lifecycle_response_dtos_round_trip_contract_fixtures() { for path in [ "tests/fixtures/webui/examples/operation.running.json", "tests/fixtures/webui/examples/operation.succeeded.json", + "tests/fixtures/webui/examples/operation.download-running.json", + "tests/fixtures/webui/examples/operation.download-succeeded.json", ] { let value = fixture_value(path); let dto: Operation = serde_json::from_value(value.clone()).unwrap_or_else(|err| { @@ -272,6 +279,88 @@ fn lifecycle_response_dtos_round_trip_contract_fixtures() { } } +#[test] +fn lifecycle_download_operation_producers_match_contract_fixtures() { + const REPO_ID: &str = "mlx-community/SmolLM-135M-Instruct-4bit"; + const RESOLVED_REVISION: &str = "642e06afe3fab57fd6cc518637c471af0a569e1e"; + const TOTAL_BYTES: u64 = 75_789_919; + + let coordinator = LifecycleCoordinator::new(); + let accepted = coordinator + .begin_download_operation( + OperationTarget::Download { + repo_id: REPO_ID.to_string(), + revision: Some("main".to_string()), + }, + Some("download-fixture-key"), + format!("download:{REPO_ID}:main"), + ) + .expect("download operation accepted"); + coordinator.register_cancellation( + &accepted.operation_id, + Arc::new(std::sync::atomic::AtomicBool::new(false)), + ); + coordinator.update_operation_progress( + &accepted.operation_id, + ProgressBytes { + completed_bytes: 1_048_576, + total_bytes: Some(TOTAL_BYTES), + indeterminate: false, + }, + ); + coordinator.update_operation(&accepted.operation_id, OperationState::Running, None, None); + let running = serde_json::to_value( + coordinator + .get_operation(&accepted.operation_id) + .expect("running operation present"), + ) + .expect("operation json"); + contract::assert_download_operation_running(&running); + + coordinator.update_operation_progress( + &accepted.operation_id, + ProgressBytes { + completed_bytes: TOTAL_BYTES, + total_bytes: Some(TOTAL_BYTES), + indeterminate: false, + }, + ); + coordinator.update_operation( + &accepted.operation_id, + OperationState::Succeeded, + Some(OperationResult::Download { + repo_id: REPO_ID.to_string(), + revision: Some(RESOLVED_REVISION.to_string()), + model_id: Some("mdl_lR1nHQwFUguxLqHbEzH2DJLdZYiDFJ0S3FzxIzY5MUU".to_string()), + download: DownloadState::Complete, + }), + None, + ); + let succeeded = serde_json::to_value( + coordinator + .get_operation(&accepted.operation_id) + .expect("terminal operation present"), + ) + .expect("operation json"); + contract::assert_download_operation_succeeded(&succeeded); +} + +#[test] +fn lifecycle_download_progress_event_matches_contract_fixture() { + let coordinator = LifecycleCoordinator::new(); + let event = + coordinator.publish_payload(UiEventPayload::DownloadProgress(DownloadProgressPayload { + operation_id: "op_dl_001".to_string(), + progress: ProgressBytes { + completed_bytes: 1_048_576, + total_bytes: None, + indeterminate: true, + }, + })); + let value = serde_json::to_value(event).expect("event json"); + contract::assert_download_progress_event(&value); +} + #[test] fn lifecycle_coordinator_lists_gets_and_reports_cancel_unsupported() { let coordinator = LifecycleCoordinator::new(); diff --git a/src/server/router_models.rs b/src/server/router_models.rs index 6880f89c2..205bbfae2 100644 --- a/src/server/router_models.rs +++ b/src/server/router_models.rs @@ -43,6 +43,7 @@ use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering}; use std::sync::{Arc, Mutex, RwLock}; +use anyhow::Context; use axum::body::Body; use futures::StreamExt; @@ -51,7 +52,7 @@ use super::router_cache::CacheSource; use super::router_lifecycle::{ DownloadState, ErrorBody, LifecycleCoordinator, LifecycleSnapshot, ModelEvictionOutcome, ModelEvictionReport, ModelLifecycle, ModelLifecycleState, OperationError, OperationKind, - OperationResult, OperationState, OperationTarget, stable_model_identity, + OperationResult, OperationState, OperationTarget, ProgressBytes, stable_model_identity, }; use super::router_presets::{PresetCliOverrides, PresetSection, RouterPresets}; use super::{AppState, ChatTemplateProcessor, ModelProvider, ServerStartupConfig}; @@ -148,11 +149,58 @@ struct EntryState { struct DownloadInFlight { cancel: Arc, + operation_id: String, /// b10621 `loaded_info` during a download: /// `{"progress": {url: {"done": n, "total": n}}}`. progress: serde_json::Value, } +#[derive(Debug)] +struct DownloadOperationProgressState { + resolved_revision: String, + total_bytes: Option, + completed_by_url: BTreeMap, + observed_total_by_url: BTreeMap, +} + +impl DownloadOperationProgressState { + fn new(requested_revision: String) -> Self { + Self { + resolved_revision: requested_revision, + total_bytes: None, + completed_by_url: BTreeMap::new(), + observed_total_by_url: BTreeMap::new(), + } + } + + fn progress_bytes(&self) -> ProgressBytes { + let completed_bytes = self + .completed_by_url + .values() + .fold(0u64, |acc, value| acc.saturating_add(*value)); + if let Some(total_bytes) = self.total_bytes { + return ProgressBytes { + completed_bytes, + total_bytes: Some(total_bytes), + indeterminate: false, + }; + } + let all_known = !self.observed_total_by_url.is_empty() + && self.observed_total_by_url.len() == self.completed_by_url.len() + && self.observed_total_by_url.values().all(|total| *total > 0); + let total_bytes = all_known.then(|| { + self.observed_total_by_url + .values() + .fold(0u64, |acc, value| acc.saturating_add(*value)) + }); + ProgressBytes { + completed_bytes, + total_bytes, + indeterminate: total_bytes.is_none(), + } + } +} + #[derive(Clone)] struct LoadedApp { state: AppState, @@ -276,6 +324,7 @@ pub struct RouterCatalogModel { pub ui_model_id: String, pub source_key_hash: String, pub hidden: bool, + pub removal_blocked_reason: Option, pub lifecycle: LifecycleSnapshot, pub revision: u64, pub generation: u64, @@ -331,6 +380,9 @@ pub struct RouterPool { /// Serializes model loads so two concurrent autoloads cannot race the /// capacity check or contend the accelerator during weight upload. load_lock: tokio::sync::Mutex<()>, + /// Serializes cache writers; additional accepted downloads remain queued + /// in the bounded operation coordinator until this lock is available. + download_lock: tokio::sync::Mutex<()>, #[cfg(test)] rescan_after_snapshot_hook: Mutex>, #[cfg(test)] @@ -388,6 +440,46 @@ struct LoadEntryExpectation { revision: u64, } +#[derive(Debug, Clone, PartialEq, Eq)] +struct SnapshotPhysicalIdentity { + canonical_path: PathBuf, + #[cfg(unix)] + dev: u64, + #[cfg(unix)] + ino: u64, +} + +fn snapshot_physical_identity(path: &Path) -> anyhow::Result { + let canonical_path = path + .canonicalize() + .with_context(|| format!("failed to canonicalize model snapshot {}", path.display()))?; + let metadata = std::fs::metadata(&canonical_path) + .with_context(|| format!("failed to stat model snapshot {}", canonical_path.display()))?; + if !metadata.is_dir() { + anyhow::bail!( + "model snapshot is not a directory: {}", + canonical_path.display() + ); + } + Ok(SnapshotPhysicalIdentity { + canonical_path, + #[cfg(unix)] + dev: { + use std::os::unix::fs::MetadataExt; + metadata.dev() + }, + #[cfg(unix)] + ino: { + use std::os::unix::fs::MetadataExt; + metadata.ino() + }, + }) +} + +fn path_tree_overlap(left: &Path, right: &Path) -> bool { + left.starts_with(right) || right.starts_with(left) +} + #[derive(Clone)] struct EvictionTarget { entry: Arc, @@ -423,6 +515,7 @@ impl RouterPool { revision_authority: Arc::new(AtomicU64::new(1)), catalog_epoch: AtomicU64::new(0), load_lock: tokio::sync::Mutex::new(()), + download_lock: tokio::sync::Mutex::new(()), #[cfg(test)] rescan_after_snapshot_hook: Mutex::new(None), #[cfg(test)] @@ -752,6 +845,26 @@ impl RouterPool { .map(|entry| entry.config.clone()) } + fn has_case_alias_entry(&self, name: &str) -> bool { + self.entries + .read() + .map(|entries| { + entries + .keys() + .any(|existing| existing.eq_ignore_ascii_case(name)) + }) + .unwrap_or(true) + } + + fn ensure_no_case_alias_in_registry( + entries: &BTreeMap>, + name: &str, + ) -> bool { + !entries + .keys() + .any(|existing| existing.eq_ignore_ascii_case(name)) + } + fn ensure_entry_is_current( &self, entry: &Arc, @@ -788,6 +901,171 @@ impl RouterPool { Ok(()) } + fn cache_removal_config_block_reason( + &self, + entries: &BTreeMap>, + entry: &Arc, + ) -> Option { + if entry.source != RouterModelSource::Cache { + return Some("only managed cache entries can be removed".to_string()); + } + if entry.hidden { + return Some("managed cache snapshot is also exposed by a preset alias".to_string()); + } + if let Some(models_dir) = &self.sources.models_dir + && path_tree_overlap(&entry.path, models_dir) + { + return Some( + "managed cache snapshot overlaps the configured --models-dir root".to_string(), + ); + } + entries.values().find_map(|other| { + if Arc::ptr_eq(other, entry) || !path_tree_overlap(&entry.path, &other.path) { + return None; + } + if other.source != RouterModelSource::Cache { + return Some(format!( + "managed cache snapshot is also exposed by an overlapping {} entry", + other.source.as_str() + )); + } + if other.path != entry.path { + return Some("managed cache snapshot overlaps another cache entry".to_string()); + } + (other.reserves_capacity() || other.is_downloading()) + .then(|| "managed cache snapshot is in use by another cache alias".to_string()) + }) + } + + fn entry_idle_for_webui_removal(entry: &RouterModelEntry) -> bool { + let snapshot = entry.lifecycle_snapshot(); + snapshot.state == ModelLifecycleState::Unloaded + && snapshot.download != DownloadState::Downloading + && snapshot.active_requests == 0 + && snapshot.draining_requests == 0 + && !entry.is_downloading() + && !entry.is_running() + } + + fn removal_conflict(operation_id: &str, message: impl Into) -> ErrorBody { + ErrorBody { + code: "conflict".to_string(), + message: message.into(), + retryable: true, + field_errors: None, + operation_id: Some(operation_id.to_string()), + } + } + + fn removal_unsupported(operation_id: &str, message: impl Into) -> ErrorBody { + ErrorBody { + code: "unsupported".to_string(), + message: message.into(), + retryable: false, + field_errors: None, + operation_id: Some(operation_id.to_string()), + } + } + + fn ensure_cache_removal_physical_owner( + &self, + entry: &Arc, + operation_id: &str, + ) -> Result<(), ErrorBody> { + if entry.hidden { + return Err(Self::removal_unsupported( + operation_id, + "managed cache snapshot is also exposed by a preset alias", + )); + } + let target_identity = snapshot_physical_identity(&entry.path).map_err(|_| { + Self::removal_conflict( + operation_id, + "model removal could not verify the managed cache snapshot; refresh and retry", + ) + })?; + if let Some(models_dir) = &self.sources.models_dir { + let models_dir_identity = models_dir.canonicalize().ok(); + match models_dir_identity { + Some(root) if path_tree_overlap(&target_identity.canonical_path, &root) => { + return Err(Self::removal_unsupported( + operation_id, + "managed cache snapshot overlaps the configured --models-dir root", + )); + } + Some(_) => {} + None => { + return Err(Self::removal_conflict( + operation_id, + "model removal could not verify --models-dir does not overlap the cache snapshot", + )); + } + } + } + let entries: Vec> = self + .entries + .read() + .map_err(|_| { + Self::removal_conflict(operation_id, "router pool is unavailable; retry removal") + })? + .values() + .cloned() + .collect(); + for other in entries { + if Arc::ptr_eq(&other, entry) { + continue; + } + let same_literal_path = other.path == entry.path; + let overlap_identity = if same_literal_path { + Some(target_identity.clone()) + } else { + match snapshot_physical_identity(&other.path) { + Ok(identity) => (identity == target_identity + || path_tree_overlap( + &target_identity.canonical_path, + &identity.canonical_path, + )) + .then_some(identity), + Err(_) if other.source != RouterModelSource::Cache => { + return Err(Self::removal_conflict( + operation_id, + "model removal could not prove a configured model source is disjoint from the cache snapshot", + )); + } + Err(_) => None, + } + }; + let Some(overlap_identity) = overlap_identity else { + continue; + }; + if other.source != RouterModelSource::Cache { + return Err(Self::removal_unsupported( + operation_id, + format!( + "managed cache snapshot is also exposed by an overlapping {} entry", + other.source.as_str() + ), + )); + } + if overlap_identity != target_identity { + return Err(Self::removal_conflict( + operation_id, + "managed cache snapshot overlaps another cache entry", + )); + } + if other.reserves_capacity() + || other.is_downloading() + || !Self::entry_idle_for_webui_removal(&other) + { + return Err(Self::removal_conflict( + operation_id, + "managed cache snapshot is in use by another model entry", + )); + } + } + Ok(()) + } + fn mark_loading_if_current( &self, entry: &Arc, @@ -905,6 +1183,7 @@ impl RouterPool { ui_model_id: entry.ui_model_id.clone(), source_key_hash: entry.source_key_hash.clone(), hidden: entry.hidden, + removal_blocked_reason: self.cache_removal_config_block_reason(&entries, entry), lifecycle: entry.lifecycle_snapshot(), revision: entry.lifecycle_revision(), generation: entry.lifecycle.generation(), @@ -1612,7 +1891,13 @@ impl RouterPool { let Some(app) = app else { return Err(RouterPoolError::NotLoaded); }; - entry.lifecycle.begin_drain(); + let (_, drain_revision) = entry.lifecycle.begin_drain_with_revision().ok_or_else(|| { + RouterPoolError::LoadFailed("model lifecycle lock poisoned during drain".into()) + })?; + let drain_expectation = expectation.map(|expected| LoadEntryExpectation { + model_id: expected.model_id.clone(), + revision: drain_revision, + }); self.notify_lifecycle(&entry); self.notify_status(&entry); if !entry @@ -1632,7 +1917,12 @@ impl RouterPool { ))); } - self.ensure_entry_is_current(&entry, expectation)?; + // The caller's revision was checked before our own drain transition. + // Recheck the owned token and registry identity after the wait: request + // lease completion preserves it, external lifecycle changes do not. + // Legacy/shutdown callers supplied no revision precondition: preserve + // their identity-only cleanup semantics for failed workers. + self.ensure_entry_is_current(&entry, drain_expectation.as_ref())?; entry.lifecycle.mark_unloading(); self.notify_lifecycle(&entry); let observer = app.state.model_provider.worker_exit_observer(); @@ -1812,6 +2102,28 @@ fn reject_stale_revision( Ok(()) } +fn operation_begin_error(err: OperationError, operation_id: Option) -> ErrorBody { + match err { + OperationError::Conflict { + operation_id: conflicting, + } => ErrorBody { + code: "conflict".to_string(), + message: "operation idempotency key conflicts with another active operation" + .to_string(), + retryable: true, + field_errors: None, + operation_id: operation_id.or(conflicting), + }, + OperationError::TooManyActive => ErrorBody { + code: "rate_limited".to_string(), + message: "too many active operations".to_string(), + retryable: true, + field_errors: None, + operation_id, + }, + } +} + fn operation_error_for_router_error(err: &RouterPoolError, operation_id: String) -> ErrorBody { match err { RouterPoolError::OperationRejected(error) => { @@ -1944,38 +2256,80 @@ impl RouterPool { /// Synchronously validate that `repo_id` is fetchable (b10621's metadata /// probe before it starts a download). Blocking. - pub fn validate_cache_repo(&self, repo_id: &str) -> anyhow::Result<()> { + pub fn validate_cache_repo(&self, repo_id: &str, revision: Option<&str>) -> anyhow::Result<()> { let cache = self .sources .cache .as_ref() .ok_or_else(|| anyhow::anyhow!("no model cache is configured"))?; - cache.validate(repo_id) + cache.validate(repo_id, revision) } - /// Start downloading `name` into the cache (`POST /models`). The name - /// must already be validated (non-empty, normalized, not present, - /// metadata probe passed). Registers a transient `downloading` entry, - /// then fetches on a background task, forwarding progress into the SSE - /// stream and rescanning when the snapshot lands. + /// Start downloading `name` into the cache (`POST /models`). The legacy + /// compatibility facade keeps its wire schema but now uses the same + /// bounded operation primitive as the WebUI adapter. pub fn start_download(self: &Arc, name: &str) -> Result<(), RouterPoolError> { + self.submit_download(name, None, None).map(|_| ()) + } + + pub fn submit_download( + self: &Arc, + repo_id: &str, + revision: Option<&str>, + idempotency_key: Option<&str>, + ) -> Result { let Some(cache) = &self.sources.cache else { return Err(RouterPoolError::LoadFailed( - "no model cache is configured (set --model-store-root, MLXCEL_MODELS_DIR, or \ - MLXCEL_CACHE_DIR)" + "no model cache is configured (set --model-store-root, MLXCEL_MODELS_DIR, or MLXCEL_CACHE_DIR)" .to_string(), )); }; - if self.lookup(name).is_some() { - return Err(RouterPoolError::AlreadyExists(name.to_string())); + let requested_revision = revision.unwrap_or("main").to_string(); + let fingerprint = format!("download:{repo_id}:{requested_revision}"); + if let Some(idempotency_key) = idempotency_key + && let Some(operation_id) = self + .lifecycle + .idempotency_conflict(idempotency_key, &fingerprint) + { + return Err(RouterPoolError::OperationRejected(operation_begin_error( + OperationError::Conflict { + operation_id: Some(operation_id), + }, + None, + ))); + } + if let Some(operation) = self + .lifecycle + .find_active_download(repo_id, &requested_revision) + { + return Ok(super::router_lifecycle::OperationAccepted { + operation_id: operation.operation_id, + state: operation.state, + idempotent_replay: true, + }); + } + if self.has_case_alias_entry(repo_id) { + return Err(RouterPoolError::AlreadyExists(repo_id.to_string())); + } + let target = OperationTarget::Download { + repo_id: repo_id.to_string(), + revision: Some(requested_revision.clone()), + }; + let accepted = self + .lifecycle + .begin_download_operation(target, idempotency_key, fingerprint) + .map_err(|err| RouterPoolError::OperationRejected(operation_begin_error(err, None)))?; + if accepted.idempotent_replay { + return Ok(accepted); } let cancel = Arc::new(AtomicBool::new(false)); - let path = cache.snapshot_dir(name); + let path = cache.snapshot_dir(repo_id); let section = PresetSection::default(); - let config = self.build_entry_config(name, &path, §ion); - let (ui_model_id, source_key_hash) = self.ui_identity_for(name, RouterModelSource::Cache); + let config = self.build_entry_config(repo_id, &path, §ion); + let (ui_model_id, source_key_hash) = + self.ui_identity_for(repo_id, RouterModelSource::Cache); let entry = Arc::new(RouterModelEntry { - name: name.to_string(), + name: repo_id.to_string(), path, source: RouterModelSource::Cache, aliases: Vec::new(), @@ -1990,6 +2344,7 @@ impl RouterPool { failed: false, download: Some(DownloadInFlight { cancel: cancel.clone(), + operation_id: accepted.operation_id.clone(), progress: serde_json::json!({ "progress": {} }), }), }), @@ -2001,16 +2356,58 @@ impl RouterPool { .entries .write() .map_err(|_| RouterPoolError::LoadFailed("router pool poisoned".into()))?; + if !Self::ensure_no_case_alias_in_registry(&entries, repo_id) { + self.lifecycle.update_operation( + &accepted.operation_id, + OperationState::Failed, + None, + Some(operation_begin_error( + OperationError::Conflict { operation_id: None }, + Some(accepted.operation_id.clone()), + )), + ); + return Err(RouterPoolError::AlreadyExists(repo_id.to_string())); + } entries.insert(entry.name.clone(), entry.clone()); } + self.lifecycle + .register_cancellation(&accepted.operation_id, cancel.clone()); self.notify_status(&entry); let pool = self.clone(); let task_entry = entry; - let repo = name.to_string(); + let repo = repo_id.to_string(); + let requested_revision_for_task = requested_revision.clone(); + let revision_for_task = revision.map(str::to_string); + let operation_id = accepted.operation_id.clone(); tokio::spawn(async move { + let _download_permit = pool.download_lock.lock().await; + if cancel.load(Ordering::Relaxed) { + pool.finish_download_task( + &task_entry, + &repo, + &operation_id, + &requested_revision_for_task, + Err(anyhow::Error::new(crate::downloader::DownloadCancelled)), + requested_revision_for_task.clone(), + ); + return; + } + pool.lifecycle + .update_operation(&operation_id, OperationState::Running, None, None); + let plan_state = Arc::new(Mutex::new(DownloadOperationProgressState::new( + requested_revision_for_task.clone(), + ))); let hooks = crate::downloader::DownloadHooks { - progress: Some(pool.download_progress_hook(&task_entry)), + progress: Some(pool.download_progress_hook( + &task_entry, + Some(operation_id.clone()), + Some(plan_state.clone()), + )), + plan: Some(pool.download_plan_hook(&operation_id, plan_state.clone())), + begin_publish: Some( + pool.download_begin_publish_hook(&operation_id, cancel.clone()), + ), cancel: Some(cancel), }; let blocking_pool = pool.clone(); @@ -2019,52 +2416,152 @@ impl RouterPool { let Some(cache) = &blocking_pool.sources.cache else { return Err(anyhow::anyhow!("no model cache configured")); }; - cache.download(&blocking_repo, hooks) + cache.download(&blocking_repo, revision_for_task.as_deref(), hooks) }) .await .unwrap_or_else(|join_err| Err(anyhow::anyhow!(join_err.to_string()))); + let resolved_revision = plan_state + .lock() + .ok() + .map(|guard| guard.resolved_revision.clone()) + .unwrap_or_else(|| requested_revision_for_task.clone()); + pool.finish_download_task( + &task_entry, + &repo, + &operation_id, + &requested_revision_for_task, + result, + resolved_revision, + ); + }); + Ok(accepted) + } - let ok = result.is_ok(); - if let Ok(mut guard) = task_entry.state.lock() { - guard.download = None; - guard.failed = !ok; + fn download_plan_hook( + self: &Arc, + operation_id: &str, + state: Arc>, + ) -> Arc { + let pool = self.clone(); + let operation_id = operation_id.to_string(); + Arc::new(move |plan| { + let progress = if let Ok(mut guard) = state.lock() { + guard.resolved_revision = plan.resolved_revision.clone(); + guard.total_bytes = plan.total_bytes; + guard.progress_bytes() + } else { + ProgressBytes { + completed_bytes: 0, + total_bytes: plan.total_bytes, + indeterminate: plan.total_bytes.is_none(), + } + }; + pool.lifecycle.register_download_revision_alias( + &operation_id, + &plan.repo_id, + &plan.resolved_revision, + ); + pool.lifecycle + .update_operation_progress(&operation_id, progress); + }) + } + + fn download_begin_publish_hook( + self: &Arc, + operation_id: &str, + cancel: Arc, + ) -> Arc bool + Send + Sync> { + let pool = self.clone(); + let operation_id = operation_id.to_string(); + Arc::new(move || { + if cancel.load(Ordering::Relaxed) { + return false; } - task_entry.lifecycle.mark_download_terminal(if ok { + pool.lifecycle.begin_publish_operation(&operation_id) + }) + } + + fn finish_download_task( + self: &Arc, + entry: &Arc, + repo: &str, + operation_id: &str, + requested_revision: &str, + result: anyhow::Result<()>, + resolved_revision: String, + ) { + let ok = result.is_ok(); + if let Ok(mut guard) = entry.state.lock() { + guard.download = None; + guard.failed = !ok; + } + let cancelled = result + .as_ref() + .err() + .is_some_and(crate::downloader::is_download_cancelled); + entry.lifecycle.mark_download_terminal(if ok { + DownloadState::Complete + } else if cancelled { + DownloadState::Incomplete + } else { + DownloadState::Failed + }); + self.notify_lifecycle(entry); + let terminal_state = if ok { + OperationState::Succeeded + } else if cancelled { + OperationState::Cancelled + } else { + OperationState::Failed + }; + let error = if ok || cancelled { + None + } else { + Some(ErrorBody { + code: "conflict".to_string(), + message: "download failed; see server logs".to_string(), + retryable: true, + field_errors: None, + operation_id: Some(operation_id.to_string()), + }) + }; + let result_body = Some(OperationResult::Download { + repo_id: repo.to_string(), + revision: Some(if ok { + resolved_revision + } else { + requested_revision.to_string() + }), + model_id: ok.then(|| entry.ui_model_id.clone()), + download: if ok { DownloadState::Complete - } else if result - .as_ref() - .err() - .is_some_and(crate::downloader::is_download_cancelled) - { + } else if cancelled { DownloadState::Incomplete } else { DownloadState::Failed - }); - pool.notify_lifecycle(&task_entry); - match &result { - Ok(()) => tracing::info!("router: download of '{repo}' finished"), - Err(err) if crate::downloader::is_download_cancelled(err) => { - tracing::info!("router: download of '{repo}' cancelled"); - } - Err(err) => tracing::warn!("router: download of '{repo}' failed: {err:#}"), - } - // b10621 order: the terminal download event first, then the - // reload that reconciles the entry (a finished snapshot becomes a - // regular cache entry; a failed one drops out of the list). - pool.notify( - if ok { - "download_finished" - } else { - "download_failed" - }, - &repo, - serde_json::Value::Null, - ); - if let Err(err) = pool.rescan() { - tracing::warn!("router: rescan after download of '{repo}' failed: {err:#}"); - } + }, }); - Ok(()) + self.lifecycle + .update_operation(operation_id, terminal_state, result_body, error); + match &result { + Ok(()) => tracing::info!("router: download of '{repo}' finished"), + Err(err) if crate::downloader::is_download_cancelled(err) => { + tracing::info!("router: download of '{repo}' cancelled"); + } + Err(err) => tracing::warn!("router: download of '{repo}' failed: {err:#}"), + } + self.notify( + if ok { + "download_finished" + } else { + "download_failed" + }, + repo, + serde_json::Value::Null, + ); + if let Err(err) = self.rescan() { + tracing::warn!("router: rescan after download of '{repo}' failed: {err:#}"); + } } /// The per-chunk progress hook: updates the entry's progress block and @@ -2073,10 +2570,13 @@ impl RouterPool { fn download_progress_hook( self: &Arc, entry: &Arc, + operation_id: Option, + operation_progress: Option>>, ) -> Arc { let pool = self.clone(); let entry = entry.clone(); - let last_emit: Mutex> = Mutex::new(None); + let last_legacy_emit: Mutex> = + Mutex::new(Some(std::time::Instant::now())); Arc::new(move |url: &str, done: u64, total: u64| { let snapshot = { let Ok(mut guard) = entry.state.lock() else { @@ -2091,7 +2591,7 @@ impl RouterPool { }; let terminal = total > 0 && done >= total; let due = { - let Ok(mut last) = last_emit.lock() else { + let Ok(mut last) = last_legacy_emit.lock() else { return; }; let now = std::time::Instant::now(); @@ -2104,12 +2604,265 @@ impl RouterPool { } due }; + if let (Some(operation_id), Some(operation_progress)) = + (operation_id.as_deref(), operation_progress.as_ref()) + && due + && let Ok(mut guard) = operation_progress.lock() + { + guard.completed_by_url.insert(url.to_string(), done); + if total > 0 { + guard.observed_total_by_url.insert(url.to_string(), total); + } + let progress_bytes = guard.progress_bytes(); + drop(guard); + pool.lifecycle + .update_operation_progress(operation_id, progress_bytes); + } if due { pool.notify("download_progress", &entry.name, snapshot); } }) } + /// Submit a WebUI-safe cache removal. Unlike the compatibility + /// `DELETE /models` path, this never unloads or cancels implicitly: the + /// catalog revision must match and the model must be idle, cache-sourced, + /// and not downloading before deletion starts. + pub fn submit_cache_removal_by_model_id( + self: &Arc, + model_id: &str, + expected_revision: u64, + idempotency_key: &str, + ) -> Result { + let target = OperationTarget::Model { + model_id: model_id.to_string(), + requested_revision: Some(expected_revision), + eviction_target_id: None, + }; + let fingerprint = format!("model_removal:{model_id}:{expected_revision}"); + let accepted = self + .lifecycle + .begin_operation( + OperationKind::ModelRemoval, + target, + Some(idempotency_key), + fingerprint, + ) + .map_err(|err| RouterPoolError::OperationRejected(operation_begin_error(err, None)))?; + if accepted.idempotent_replay { + return Ok(accepted); + } + let Some(entry) = self.get_by_model_id(model_id) else { + let error = ErrorBody { + code: "not_found".to_string(), + message: "model was not found; refresh the catalog before retrying".to_string(), + retryable: true, + field_errors: None, + operation_id: Some(accepted.operation_id.clone()), + }; + self.lifecycle.update_operation( + &accepted.operation_id, + OperationState::Failed, + None, + Some(error.clone()), + ); + return Err(RouterPoolError::OperationRejected(error)); + }; + let expectation = LoadEntryExpectation { + model_id: model_id.to_string(), + revision: expected_revision, + }; + let entry_permit = match entry.lifecycle.try_operation_guard() { + Ok(guard) => guard, + Err(_) => { + let error = Self::removal_conflict( + &accepted.operation_id, + "model is busy; unload or wait for current lifecycle work before removing it", + ); + self.lifecycle.update_operation( + &accepted.operation_id, + OperationState::Failed, + None, + Some(error.clone()), + ); + return Err(RouterPoolError::OperationRejected(error)); + } + }; + if let Err(err) = self.ensure_entry_is_current(&entry, Some(&expectation)) { + let error = operation_error_for_router_error(&err, accepted.operation_id.clone()); + self.lifecycle.update_operation( + &accepted.operation_id, + OperationState::Failed, + None, + Some(error.clone()), + ); + return Err(RouterPoolError::OperationRejected(error)); + } + if entry.source != RouterModelSource::Cache { + let error = Self::removal_unsupported( + &accepted.operation_id, + "only cache-sourced models can be removed from the WebUI", + ); + self.lifecycle.update_operation( + &accepted.operation_id, + OperationState::Failed, + None, + Some(error.clone()), + ); + return Err(RouterPoolError::OperationRejected(error)); + } + if !Self::entry_idle_for_webui_removal(&entry) { + let error = Self::removal_conflict( + &accepted.operation_id, + "model is busy; unload or wait for current lifecycle work before removing it", + ); + self.lifecycle.update_operation( + &accepted.operation_id, + OperationState::Failed, + None, + Some(error.clone()), + ); + return Err(RouterPoolError::OperationRejected(error)); + } + if let Err(error) = self.ensure_cache_removal_physical_owner(&entry, &accepted.operation_id) + { + self.lifecycle.update_operation( + &accepted.operation_id, + OperationState::Failed, + None, + Some(error.clone()), + ); + return Err(RouterPoolError::OperationRejected(error)); + } + entry.lifecycle.mark_operation_busy(); + self.notify_lifecycle(&entry); + let operation_revision = entry.lifecycle_revision(); + drop(entry_permit); + + let pool = self.clone(); + let operation_id = accepted.operation_id.clone(); + let model_id = model_id.to_string(); + let removal_entry = entry.clone(); + tokio::spawn(async move { + pool.lifecycle + .update_operation(&operation_id, OperationState::Running, None, None); + let expectation = LoadEntryExpectation { + model_id: model_id.clone(), + revision: operation_revision, + }; + let entry = removal_entry; + let _guard = entry.lifecycle.operation_guard().await; + if let Err(err) = pool.ensure_entry_is_current(&entry, Some(&expectation)) { + let error = operation_error_for_router_error(&err, operation_id.clone()); + entry.lifecycle.mark_operation_idle(); + pool.notify_lifecycle(&entry); + pool.lifecycle.update_operation( + &operation_id, + OperationState::Failed, + None, + Some(error), + ); + return; + } + if entry.source != RouterModelSource::Cache { + let error = Self::removal_unsupported( + &operation_id, + "only cache-sourced models can be removed from the WebUI", + ); + entry.lifecycle.mark_operation_idle(); + pool.notify_lifecycle(&entry); + pool.lifecycle.update_operation( + &operation_id, + OperationState::Failed, + None, + Some(error), + ); + return; + } + if !Self::entry_idle_for_webui_removal(&entry) { + let error = Self::removal_conflict( + &operation_id, + "model is busy; unload or wait for current lifecycle work before removing it", + ); + entry.lifecycle.mark_operation_idle(); + pool.notify_lifecycle(&entry); + pool.lifecycle.update_operation( + &operation_id, + OperationState::Failed, + None, + Some(error), + ); + return; + } + if let Err(error) = pool.ensure_cache_removal_physical_owner(&entry, &operation_id) { + entry.lifecycle.mark_operation_idle(); + pool.notify_lifecycle(&entry); + pool.lifecycle.update_operation( + &operation_id, + OperationState::Failed, + None, + Some(error), + ); + return; + } + let snapshot = entry.lifecycle_snapshot(); + let remove_name = entry.name.clone(); + let remove_pool = pool.clone(); + let remove_name_for_blocking = remove_name.clone(); + let removed = tokio::task::spawn_blocking(move || { + let Some(cache) = &remove_pool.sources.cache else { + anyhow::bail!("no model cache is configured"); + }; + cache.remove(&remove_name_for_blocking) + }) + .await + .unwrap_or_else(|join_err| Err(anyhow::anyhow!(join_err.to_string()))); + match removed { + Ok(()) => { + if let Ok(mut entries) = pool.entries.write() + && entries + .get(&remove_name) + .is_some_and(|current| Arc::ptr_eq(current, &entry)) + { + entries.remove(&remove_name); + } + pool.catalog_epoch.fetch_add(1, Ordering::SeqCst); + pool.notify("model_remove", &remove_name, serde_json::Value::Null); + pool.lifecycle.update_operation( + &operation_id, + OperationState::Succeeded, + Some(OperationResult::ModelRemoval { + model_id: model_id.clone(), + revision: entry.lifecycle_revision(), + lifecycle: snapshot, + eviction: None, + }), + None, + ); + } + Err(err) => { + tracing::warn!(model_id = %model_id, error = %format_args!("{err:#}"), "router: model removal failed"); + entry.lifecycle.mark_operation_idle(); + pool.notify_lifecycle(&entry); + let error = ErrorBody { + code: "conflict".to_string(), + message: "model removal failed; see server logs".to_string(), + retryable: true, + field_errors: None, + operation_id: Some(operation_id.clone()), + }; + pool.lifecycle.update_operation( + &operation_id, + OperationState::Failed, + None, + Some(error), + ); + } + } + }); + Ok(accepted) + } + /// Remove `name` from the cache (`DELETE /models`, b10621 /// `server_models::remove`): cancel an in-flight download or stop a /// running instance, delete the snapshot from disk (containment-checked @@ -2126,11 +2879,15 @@ impl RouterPool { }; self.ensure_entry_is_current(&entry, None)?; - if let Ok(guard) = entry.state.lock() - && let Some(download) = &guard.download - { + let download_operation_id = entry.state.lock().ok().and_then(|guard| { + guard + .download + .as_ref() + .map(|download| download.operation_id.clone()) + }); + if let Some(operation_id) = download_operation_id { tracing::info!("router: cancelling download for model '{name}'"); - download.cancel.store(true, Ordering::Relaxed); + let _ = self.lifecycle.cancel_operation(&operation_id); } // Wait for the download worker to acknowledge the cancel (it clears // the download state in its completion block). The method does not @@ -2160,12 +2917,24 @@ impl RouterPool { None => false, } }; + if !still_registered && !entry.path.exists() { + self.notify("model_remove", name, serde_json::Value::Null); + return Ok(()); + } + if still_registered && entry.reserves_capacity() { tracing::info!("router: stopping model instance '{name}' before removal"); self.unload_entry(&entry).await?; } let _entry_permit = entry.lifecycle.operation_guard().await; + if let Err(error) = self.ensure_cache_removal_physical_owner(&entry, "op_compat_remove") { + return if error.code == "unsupported" { + Err(RouterPoolError::NotRemovable(name.to_string())) + } else { + Err(RouterPoolError::LoadFailed(error.message)) + }; + } { let mut entries = self .entries @@ -2293,3 +3062,7 @@ mod router_models_tests; #[cfg(test)] #[path = "router_models_discovery_tests.rs"] mod router_models_discovery_tests; + +#[cfg(test)] +#[path = "router_unload_revision_tests.rs"] +mod router_unload_revision_tests; diff --git a/src/server/router_models_discovery_tests.rs b/src/server/router_models_discovery_tests.rs index 093bea410..9454c8ba2 100644 --- a/src/server/router_models_discovery_tests.rs +++ b/src/server/router_models_discovery_tests.rs @@ -61,13 +61,14 @@ fn router_sources_reject_symlinked_config_evidence_consistently() { struct NoopDownloader; impl RouterDownloader for NoopDownloader { - fn validate(&self, _repo_id: &str) -> anyhow::Result<()> { + fn validate(&self, _repo_id: &str, _revision: Option<&str>) -> anyhow::Result<()> { Ok(()) } fn download( &self, _repo_id: &str, + _revision: Option<&str>, _dest_root: &std::path::Path, _hooks: DownloadHooks, ) -> anyhow::Result<()> { diff --git a/src/server/router_models_tests.rs b/src/server/router_models_tests.rs index 73e54de8a..03eed76ae 100644 --- a/src/server/router_models_tests.rs +++ b/src/server/router_models_tests.rs @@ -49,6 +49,19 @@ fn add_fake_model(root: &std::path::Path, name: &str) { std::fs::write(dir.join("config.json"), "{}").expect("config.json"); } +#[test] +fn cache_source_list_does_not_create_absent_store_root() { + let temp = tempfile::tempdir().expect("tempdir"); + let absent = temp.path().join("missing-store"); + let cache = CacheSource::new(absent.clone(), FakeDownloader::ok()); + + assert!(cache.list().is_empty()); + assert!( + !absent.exists(), + "catalog construction/listing must not create cache or staging directories" + ); +} + /// A downloader that materializes a fake snapshot locally, driving the same /// hooks the HuggingFace downloader drives. `delay_until_cancel` makes the /// download hang until the cancel flag flips, for cancellation tests. @@ -85,13 +98,14 @@ impl FakeDownloader { } impl RouterDownloader for FakeDownloader { - fn validate(&self, _repo_id: &str) -> anyhow::Result<()> { + fn validate(&self, _repo_id: &str, _revision: Option<&str>) -> anyhow::Result<()> { Ok(()) } fn download( &self, repo_id: &str, + _revision: Option<&str>, dest_root: &Path, hooks: DownloadHooks, ) -> anyhow::Result<()> { @@ -124,6 +138,190 @@ impl RouterDownloader for FakeDownloader { } } +struct PublishingDownloader { + entered_publish: std::sync::Mutex>>, +} + +impl PublishingDownloader { + fn new(sender: std::sync::mpsc::Sender<()>) -> Arc { + Arc::new(Self { + entered_publish: std::sync::Mutex::new(Some(sender)), + }) + } +} + +impl RouterDownloader for PublishingDownloader { + fn validate(&self, _repo_id: &str, _revision: Option<&str>) -> anyhow::Result<()> { + Ok(()) + } + + fn download( + &self, + repo_id: &str, + _revision: Option<&str>, + dest_root: &Path, + hooks: DownloadHooks, + ) -> anyhow::Result<()> { + let url = format!("https://example.invalid/{repo_id}/config.json"); + if let Some(progress) = &hooks.progress { + progress(&url, 1, 1); + } + if let Some(begin_publish) = &hooks.begin_publish + && !begin_publish() + { + return Err(anyhow::Error::new(crate::downloader::DownloadCancelled)); + } + if let Some(sender) = self + .entered_publish + .lock() + .ok() + .and_then(|mut guard| guard.take()) + { + let _ = sender.send(()); + } + std::thread::sleep(std::time::Duration::from_millis(100)); + let dest = dest_root.join(repo_id); + std::fs::create_dir_all(&dest)?; + std::fs::write(dest.join("config.json"), "{}")?; + Ok(()) + } +} + +struct ResolvingHangingDownloader { + plan_sent: std::sync::Mutex>>, + resolved_revision: String, +} + +impl ResolvingHangingDownloader { + fn new(sender: std::sync::mpsc::Sender<()>, resolved_revision: &str) -> Arc { + Arc::new(Self { + plan_sent: std::sync::Mutex::new(Some(sender)), + resolved_revision: resolved_revision.to_string(), + }) + } +} + +impl RouterDownloader for ResolvingHangingDownloader { + fn validate(&self, _repo_id: &str, _revision: Option<&str>) -> anyhow::Result<()> { + Ok(()) + } + + fn download( + &self, + repo_id: &str, + revision: Option<&str>, + _dest_root: &Path, + hooks: DownloadHooks, + ) -> anyhow::Result<()> { + if let Some(plan) = &hooks.plan { + plan(crate::downloader::DownloadPlan { + repo_id: repo_id.to_string(), + requested_revision: revision.unwrap_or("main").to_string(), + resolved_revision: self.resolved_revision.clone(), + destination: PathBuf::from("/redacted/cache"), + selected_files: 2, + total_bytes: Some(100), + }); + } + if let Some(sender) = self + .plan_sent + .lock() + .ok() + .and_then(|mut guard| guard.take()) + { + let _ = sender.send(()); + } + let cancel = hooks.cancel.clone().expect("cancel flag"); + while !cancel.load(Ordering::Relaxed) { + std::thread::sleep(std::time::Duration::from_millis(10)); + } + Err(anyhow::Error::new(crate::downloader::DownloadCancelled)) + } +} + +struct BurstProgressDownloader; + +impl RouterDownloader for BurstProgressDownloader { + fn validate(&self, _repo_id: &str, _revision: Option<&str>) -> anyhow::Result<()> { + Ok(()) + } + + fn download( + &self, + repo_id: &str, + revision: Option<&str>, + dest_root: &Path, + hooks: DownloadHooks, + ) -> anyhow::Result<()> { + if let Some(plan) = &hooks.plan { + plan(crate::downloader::DownloadPlan { + repo_id: repo_id.to_string(), + requested_revision: revision.unwrap_or("main").to_string(), + resolved_revision: "abc123".to_string(), + destination: dest_root.join(repo_id), + selected_files: 2, + total_bytes: Some(100), + }); + } + let url = format!("https://example.invalid/{repo_id}/model.safetensors"); + if let Some(progress) = &hooks.progress { + for done in 1..20 { + progress(&url, done, 20); + } + progress(&url, 20, 20); + } + let dest = dest_root.join(repo_id); + std::fs::create_dir_all(&dest)?; + std::fs::write(dest.join("config.json"), "{}")?; + Ok(()) + } +} + +struct FailOnceDownloader { + attempts: AtomicUsize, +} + +impl FailOnceDownloader { + fn new() -> Arc { + Arc::new(Self { + attempts: AtomicUsize::new(0), + }) + } +} + +impl RouterDownloader for FailOnceDownloader { + fn validate(&self, _repo_id: &str, _revision: Option<&str>) -> anyhow::Result<()> { + Ok(()) + } + + fn download( + &self, + repo_id: &str, + revision: Option<&str>, + dest_root: &Path, + hooks: DownloadHooks, + ) -> anyhow::Result<()> { + if self.attempts.fetch_add(1, Ordering::SeqCst) == 0 { + anyhow::bail!("simulated first download failure"); + } + if let Some(plan) = &hooks.plan { + plan(crate::downloader::DownloadPlan { + repo_id: repo_id.to_string(), + requested_revision: revision.unwrap_or("main").to_string(), + resolved_revision: "retry-sha".to_string(), + destination: dest_root.join(repo_id), + selected_files: 2, + total_bytes: Some(2), + }); + } + let dest = dest_root.join(repo_id); + std::fs::create_dir_all(&dest)?; + std::fs::write(dest.join("config.json"), "{}")?; + std::fs::write(dest.join("model.safetensors"), b"ok")?; + Ok(()) + } +} + fn sources_dir_only(root: PathBuf) -> RouterSources { RouterSources { models_dir: Some(root), @@ -394,6 +592,417 @@ async fn a_download_emits_the_b10621_event_sequence_and_lands_in_the_cache() { assert_eq!(downloader.downloads.load(Ordering::SeqCst), 1); } +#[tokio::test] +async fn duplicate_download_with_same_idempotency_key_replays_active_operation() { + let cache = tempfile::tempdir().unwrap(); + let downloader = FakeDownloader::hanging(); + let pool = Arc::new(pool_from( + RouterSources { + cache: Some(CacheSource::new( + cache.path().to_path_buf(), + downloader.clone(), + )), + ..RouterSources::default() + }, + 2, + false, + )); + let mut events = pool.subscribe(); + let first = pool + .submit_download( + "mlx-community/replay-model", + None, + Some("idem_download_replay"), + ) + .expect("first download"); + let second = pool + .submit_download( + "mlx-community/replay-model", + None, + Some("idem_download_replay"), + ) + .expect("replay download"); + assert_eq!(second.operation_id, first.operation_id); + assert!(second.idempotent_replay); + wait_for_operation_state( + &pool, + &first.operation_id, + &[crate::server::router_lifecycle::OperationState::Running], + ) + .await; + pool.lifecycle_coordinator() + .cancel_operation(&first.operation_id) + .expect("cancel replayed op"); + wait_for_event(&mut events, "download_failed").await; + assert_eq!(downloader.downloads.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn duplicate_download_idempotency_key_rejects_different_payload_before_alias_replay() { + let cache = tempfile::tempdir().unwrap(); + let downloader = FakeDownloader::hanging(); + let pool = Arc::new(pool_from( + RouterSources { + cache: Some(CacheSource::new(cache.path().to_path_buf(), downloader)), + ..RouterSources::default() + }, + 2, + false, + )); + let mut events = pool.subscribe(); + let first = pool + .submit_download( + "mlx-community/idempotency-conflict", + Some("main"), + Some("idem_download_conflict"), + ) + .expect("first download"); + + let err = pool + .submit_download( + "mlx-community/idempotency-conflict", + Some("dev"), + Some("idem_download_conflict"), + ) + .expect_err("same idempotency key with different payload must fail"); + match err { + RouterPoolError::OperationRejected(error) => { + assert_eq!(error.code, "conflict"); + assert_eq!( + error.operation_id.as_deref(), + Some(first.operation_id.as_str()) + ); + } + other => panic!("unexpected error: {other:?}"), + } + + pool.lifecycle_coordinator() + .cancel_operation(&first.operation_id) + .expect("cancel first op"); + wait_for_event(&mut events, "download_failed").await; +} + +#[tokio::test] +async fn active_download_replays_resolved_revision_alias() { + let cache = tempfile::tempdir().unwrap(); + let (plan_tx, plan_rx) = std::sync::mpsc::channel(); + let downloader = ResolvingHangingDownloader::new(plan_tx, "abc123"); + let pool = Arc::new(pool_from( + RouterSources { + cache: Some(CacheSource::new(cache.path().to_path_buf(), downloader)), + ..RouterSources::default() + }, + 2, + false, + )); + let mut events = pool.subscribe(); + let first = pool + .submit_download( + "mlx-community/revision-alias", + Some("main"), + Some("idem_revision_alias_main"), + ) + .expect("first download"); + tokio::task::spawn_blocking(move || { + plan_rx + .recv_timeout(std::time::Duration::from_secs(10)) + .expect("plan hook") + }) + .await + .expect("plan waiter"); + + let second = pool + .submit_download( + "mlx-community/revision-alias", + Some("abc123"), + Some("idem_revision_alias_sha"), + ) + .expect("resolved revision replay"); + assert_eq!(second.operation_id, first.operation_id); + assert!(second.idempotent_replay); + + pool.lifecycle_coordinator() + .cancel_operation(&first.operation_id) + .expect("cancel first op"); + wait_for_event(&mut events, "download_failed").await; +} + +#[tokio::test] +async fn download_progress_keeps_plan_total_and_throttles_coordinator_events() { + let cache = tempfile::tempdir().unwrap(); + let pool = Arc::new(pool_from( + RouterSources { + cache: Some(CacheSource::new( + cache.path().to_path_buf(), + Arc::new(BurstProgressDownloader), + )), + ..RouterSources::default() + }, + 2, + false, + )); + let mut lifecycle_events = pool.lifecycle_coordinator().subscribe(); + let accepted = pool + .submit_download( + "mlx-community/progress-total", + None, + Some("idem_progress_total"), + ) + .expect("download accepted"); + wait_for_operation_state( + &pool, + &accepted.operation_id, + &[crate::server::router_lifecycle::OperationState::Succeeded], + ) + .await; + + let mut progress_events = Vec::new(); + while let Ok(event) = lifecycle_events.try_recv() { + if let crate::server::router_lifecycle::UiEventPayload::DownloadProgress(payload) = + event.payload + && payload.operation_id == accepted.operation_id + { + progress_events.push(payload.progress); + } + } + assert!( + progress_events.len() <= 2, + "coordinator progress should be coalesced, got {progress_events:?}" + ); + assert!( + progress_events + .iter() + .all(|progress| progress.total_bytes == Some(100)), + "known plan total must not regress to per-file totals: {progress_events:?}" + ); +} + +#[tokio::test] +async fn failed_download_can_retry_same_repo_with_new_idempotency_key() { + let cache = tempfile::tempdir().unwrap(); + let downloader = FailOnceDownloader::new(); + let pool = Arc::new(pool_from( + RouterSources { + cache: Some(CacheSource::new( + cache.path().to_path_buf(), + downloader.clone(), + )), + ..RouterSources::default() + }, + 2, + false, + )); + let mut events = pool.subscribe(); + let first = pool + .submit_download( + "mlx-community/retry-once", + None, + Some("idem_retry_once_first"), + ) + .expect("first download accepted"); + let failed = wait_for_operation_state( + &pool, + &first.operation_id, + &[crate::server::router_lifecycle::OperationState::Failed], + ) + .await; + assert_eq!( + failed.state, + crate::server::router_lifecycle::OperationState::Failed + ); + wait_for_event(&mut events, "download_failed").await; + assert!( + pool.get("mlx-community/retry-once").is_none(), + "failed transient entry must be dropped before retry" + ); + assert!( + !cache.path().join("mlx-community/retry-once").exists(), + "failed fake transfer must leave no cache snapshot" + ); + + let second = pool + .submit_download( + "mlx-community/retry-once", + None, + Some("idem_retry_once_second"), + ) + .expect("retry download accepted"); + assert_ne!( + first.operation_id, second.operation_id, + "new idempotency key should create a distinct terminal operation" + ); + let succeeded = wait_for_operation_state( + &pool, + &second.operation_id, + &[crate::server::router_lifecycle::OperationState::Succeeded], + ) + .await; + assert_eq!( + succeeded.state, + crate::server::router_lifecycle::OperationState::Succeeded + ); + wait_for_event(&mut events, "download_finished").await; + + assert_eq!(downloader.attempts.load(Ordering::SeqCst), 2); + assert!( + !cache.path().join(".mlxcel-staging").exists(), + "retry fake transfer must not leave staging debris" + ); + let cache_entries = pool + .catalog_snapshot() + .into_iter() + .filter(|model| { + model.name == "mlx-community/retry-once" && model.source == RouterModelSource::Cache + }) + .count(); + assert_eq!(cache_entries, 1, "retry should publish exactly one entry"); +} + +#[tokio::test] +async fn download_rejects_case_only_cache_alias() { + let cache = tempfile::tempdir().unwrap(); + let downloader = FakeDownloader::hanging(); + let pool = Arc::new(pool_from( + RouterSources { + cache: Some(CacheSource::new(cache.path().to_path_buf(), downloader)), + ..RouterSources::default() + }, + 2, + false, + )); + let mut events = pool.subscribe(); + let first = pool + .submit_download("mlx-community/CaseModel", None, Some("idem_case_one")) + .expect("first download"); + assert_eq!( + pool.submit_download("mlx-community/casemodel", None, Some("idem_case_two")), + Err(RouterPoolError::AlreadyExists( + "mlx-community/casemodel".to_string() + )) + ); + pool.lifecycle_coordinator() + .cancel_operation(&first.operation_id) + .expect("cancel first op"); + wait_for_event(&mut events, "download_failed").await; +} + +#[tokio::test] +async fn download_admission_rejects_queue_saturation_before_worker_network() { + let cache = tempfile::tempdir().unwrap(); + let downloader = FakeDownloader::hanging(); + let pool = Arc::new(pool_from( + RouterSources { + cache: Some(CacheSource::new( + cache.path().to_path_buf(), + downloader.clone(), + )), + ..RouterSources::default() + }, + 8, + false, + )); + let mut events = pool.subscribe(); + let mut accepted = Vec::new(); + for idx in 0..crate::server::router_lifecycle::MAX_ACTIVE_DOWNLOAD_OPERATIONS { + accepted.push( + pool.submit_download( + &format!("mlx-community/saturated-{idx}"), + None, + Some(&format!("idem_saturation_{idx}")), + ) + .expect("download admission"), + ); + } + + let err = pool + .submit_download( + "mlx-community/saturated-extra", + None, + Some("idem_saturation_extra"), + ) + .expect_err("fifth active download must be rejected before network"); + match err { + RouterPoolError::OperationRejected(error) => { + assert_eq!(error.code, "rate_limited"); + assert!( + downloader.downloads.load(Ordering::SeqCst) <= accepted.len(), + "rejected request must not start an extra worker" + ); + } + other => panic!("unexpected error: {other:?}"), + } + + for op in &accepted { + pool.lifecycle_coordinator() + .cancel_operation(&op.operation_id) + .expect("cancel accepted op"); + } + let mut failed = 0usize; + tokio::time::timeout(std::time::Duration::from_secs(10), async { + while failed < accepted.len() { + let event = events.recv().await.expect("events open"); + if event["event"] == "download_failed" { + failed += 1; + } + } + }) + .await + .expect("all accepted downloads should cancel"); +} + +#[tokio::test] +async fn cancel_after_publish_linearization_is_refused_and_download_completes() { + let cache = tempfile::tempdir().unwrap(); + let (tx, rx) = std::sync::mpsc::channel(); + let pool = Arc::new(pool_from( + RouterSources { + cache: Some(CacheSource::new( + cache.path().to_path_buf(), + PublishingDownloader::new(tx), + )), + ..RouterSources::default() + }, + 2, + false, + )); + let mut events = pool.subscribe(); + let accepted = pool + .submit_download( + "mlx-community/publish-race", + None, + Some("idem_publish_race"), + ) + .expect("download accepted"); + + tokio::task::spawn_blocking(move || { + rx.recv_timeout(std::time::Duration::from_secs(10)) + .expect("publish linearization") + }) + .await + .expect("publish waiter"); + + assert!(matches!( + pool.lifecycle_coordinator() + .cancel_operation(&accepted.operation_id), + Err(crate::server::router_lifecycle::CancelError::Unsupported { .. }) + )); + wait_for_event(&mut events, "download_finished").await; + let operation = pool + .lifecycle_coordinator() + .get_operation(&accepted.operation_id) + .expect("operation"); + assert_eq!( + operation.state, + crate::server::router_lifecycle::OperationState::Succeeded + ); + assert!( + cache + .path() + .join("mlx-community/publish-race/config.json") + .is_file() + ); +} + #[tokio::test] async fn download_progress_events_carry_per_url_done_and_total() { let cache_root = temp_models_dir("dl-progress"); @@ -832,6 +1441,326 @@ async fn rescan_and_remove_preserve_reserved_entries_until_release() { ); } +#[test] +fn rescan_preserves_webui_removal_busy_entry_until_operation_finishes() { + let cache_root = temp_models_dir("removal-busy-rescan"); + add_fake_model(&cache_root.join("mlx-community"), "remove-busy"); + let pool = pool_from( + RouterSources { + models_dir: None, + cache: Some(CacheSource::new(cache_root.clone(), FakeDownloader::ok())), + presets: Default::default(), + }, + 4, + true, + ); + let entry = pool.get("mlx-community/remove-busy").expect("entry"); + entry.lifecycle.mark_operation_busy(); + std::fs::remove_dir_all(cache_root.join("mlx-community/remove-busy")).unwrap(); + + pool.rescan().expect("rescan"); + + let current = pool + .get("mlx-community/remove-busy") + .expect("entry after rescan"); + assert!( + Arc::ptr_eq(&entry, ¤t), + "rescan must not rebuild an entry while WebUI removal owns its busy reservation" + ); + assert!(current.lifecycle_snapshot().busy); + entry.lifecycle.mark_operation_idle(); +} + +#[tokio::test] +async fn webui_cache_removal_deletes_managed_snapshot_and_drops_entry() { + let cache_root = temp_models_dir("webui-remove-ok"); + add_fake_model(&cache_root.join("mlx-community"), "remove-me"); + let pool = Arc::new(pool_from( + RouterSources { + models_dir: None, + cache: Some(CacheSource::new(cache_root.clone(), FakeDownloader::ok())), + presets: Default::default(), + }, + 4, + true, + )); + let mut events = pool.subscribe(); + let entry = pool.get("mlx-community/remove-me").expect("entry"); + let accepted = pool + .submit_cache_removal_by_model_id( + &entry.ui_model_id, + entry.lifecycle_revision(), + "idem_remove_model", + ) + .expect("remove accepted"); + + let terminal = wait_for_operation_state( + &pool, + &accepted.operation_id, + &[crate::server::router_lifecycle::OperationState::Succeeded], + ) + .await; + assert_eq!( + terminal.kind, + crate::server::router_lifecycle::OperationKind::ModelRemoval + ); + wait_for_event(&mut events, "model_remove").await; + assert!(!cache_root.join("mlx-community/remove-me").exists()); + assert!(pool.get("mlx-community/remove-me").is_none()); +} + +#[tokio::test] +async fn webui_cache_removal_rejects_when_operation_guard_is_held() { + let cache_root = temp_models_dir("webui-remove-guard-held"); + add_fake_model(&cache_root.join("mlx-community"), "guarded"); + let pool = Arc::new(pool_from( + RouterSources { + models_dir: None, + cache: Some(CacheSource::new(cache_root.clone(), FakeDownloader::ok())), + presets: Default::default(), + }, + 4, + true, + )); + let entry = pool.get("mlx-community/guarded").expect("entry"); + let _guard = entry.lifecycle.operation_guard().await; + let err = pool + .submit_cache_removal_by_model_id( + &entry.ui_model_id, + entry.lifecycle_revision(), + "idem_remove_guarded", + ) + .expect_err("held operation guard must reject removal"); + match err { + RouterPoolError::OperationRejected(error) => assert_eq!(error.code, "conflict"), + other => panic!("unexpected error: {other:?}"), + } + assert!(cache_root.join("mlx-community/guarded").exists()); +} + +#[tokio::test] +async fn webui_cache_removal_rejects_models_dir_physical_overlap() { + let cache_root = temp_models_dir("webui-remove-models-dir-overlap"); + add_fake_model(&cache_root.join("mlx-community"), "shared"); + let pool = Arc::new(pool_from( + RouterSources { + models_dir: Some(cache_root.join("mlx-community")), + cache: Some(CacheSource::new(cache_root.clone(), FakeDownloader::ok())), + presets: Default::default(), + }, + 4, + true, + )); + let entry = pool.get("mlx-community/shared").expect("cache entry"); + let err = pool + .submit_cache_removal_by_model_id( + &entry.ui_model_id, + entry.lifecycle_revision(), + "idem_remove_models_dir_overlap", + ) + .expect_err("models-dir alias must block cache removal"); + match err { + RouterPoolError::OperationRejected(error) => assert_eq!(error.code, "unsupported"), + other => panic!("unexpected error: {other:?}"), + } + let catalog = pool.catalog_snapshot(); + let cache_entry = catalog + .iter() + .find(|model| model.name == "mlx-community/shared") + .expect("catalog cache entry"); + assert!(cache_entry.removal_blocked_reason.is_some()); + assert!( + cache_root + .join("mlx-community/shared/config.json") + .is_file() + ); +} + +#[tokio::test] +async fn webui_cache_removal_rejects_descendant_models_dir_overlap() { + let cache_root = temp_models_dir("webui-remove-models-dir-descendant"); + add_fake_model(&cache_root.join("mlx-community"), "parent"); + let nested_models_dir = cache_root.join("mlx-community/parent/nested-models"); + add_fake_model(&nested_models_dir, "user-owned"); + let pool = Arc::new(pool_from( + RouterSources { + models_dir: Some(nested_models_dir.clone()), + cache: Some(CacheSource::new(cache_root.clone(), FakeDownloader::ok())), + presets: Default::default(), + }, + 4, + true, + )); + assert!( + pool.get("user-owned").is_some(), + "descendant models-dir entry should be discovered" + ); + let entry = pool.get("mlx-community/parent").expect("cache entry"); + let err = pool + .submit_cache_removal_by_model_id( + &entry.ui_model_id, + entry.lifecycle_revision(), + "idem_remove_models_dir_descendant", + ) + .expect_err("descendant models-dir root must block parent cache removal"); + match err { + RouterPoolError::OperationRejected(error) => assert_eq!(error.code, "unsupported"), + other => panic!("unexpected error: {other:?}"), + } + let catalog = pool.catalog_snapshot(); + let cache_entry = catalog + .iter() + .find(|model| model.name == "mlx-community/parent") + .expect("catalog cache entry"); + assert!(cache_entry.removal_blocked_reason.is_some()); + assert!( + cache_root + .join("mlx-community/parent/config.json") + .is_file(), + "managed parent snapshot must remain after rejected removal" + ); + assert!( + nested_models_dir.join("user-owned/config.json").is_file(), + "nested user-owned models-dir files must remain after rejected removal" + ); +} + +#[tokio::test] +async fn webui_cache_removal_rejects_active_preset_physical_alias_without_hidden_cache_row() { + let cache_root = temp_models_dir("webui-remove-active-preset-overlap"); + add_fake_model(&cache_root.join("mlx-community"), "active-twin"); + let ini = "[served-active] +hf-repo = mlx-community/active-twin +"; + let presets = parse_preset_text(ini).expect("parse presets"); + let pool = Arc::new(pool_from( + RouterSources { + models_dir: None, + cache: Some(CacheSource::new(cache_root.clone(), FakeDownloader::ok())), + presets, + }, + 4, + true, + )); + let cache_entry = pool.get("mlx-community/active-twin").expect("cache entry"); + assert!(!cache_entry.hidden); + let served = pool.get("served-active").expect("preset alias"); + served.lifecycle.mark_ready(); + + let err = pool + .submit_cache_removal_by_model_id( + &cache_entry.ui_model_id, + cache_entry.lifecycle_revision(), + "idem_remove_active_preset_overlap", + ) + .expect_err("active preset alias must block cache removal"); + match err { + RouterPoolError::OperationRejected(error) => assert_eq!(error.code, "unsupported"), + other => panic!("unexpected error: {other:?}"), + } + let catalog = pool.catalog_snapshot(); + let cache_catalog = catalog + .iter() + .find(|model| model.name == "mlx-community/active-twin") + .expect("catalog cache entry"); + assert!(cache_catalog.removal_blocked_reason.is_some()); + assert!( + cache_root + .join("mlx-community/active-twin/config.json") + .is_file() + ); +} + +#[tokio::test] +async fn webui_cache_removal_rejects_preset_physical_alias_even_when_hidden() { + let cache_root = temp_models_dir("webui-remove-preset-overlap"); + add_fake_model(&cache_root.join("mlx-community"), "twin"); + let ini = "[served]\nhf-repo = mlx-community/twin\ndedup-cache-models = 1\n"; + let presets = parse_preset_text(ini).expect("parse presets"); + let pool = Arc::new(pool_from( + RouterSources { + models_dir: None, + cache: Some(CacheSource::new(cache_root.clone(), FakeDownloader::ok())), + presets, + }, + 4, + true, + )); + let cache_entry = pool.get("mlx-community/twin").expect("hidden cache entry"); + assert!(cache_entry.hidden); + let served = pool.get("served").expect("preset entry"); + served.lifecycle.mark_ready(); + + let err = pool + .submit_cache_removal_by_model_id( + &cache_entry.ui_model_id, + cache_entry.lifecycle_revision(), + "idem_remove_hidden_preset_overlap", + ) + .expect_err("preset alias must block cache removal"); + match err { + RouterPoolError::OperationRejected(error) => assert_eq!(error.code, "unsupported"), + other => panic!("unexpected error: {other:?}"), + } + assert!(cache_root.join("mlx-community/twin/config.json").is_file()); +} + +#[tokio::test] +async fn webui_cache_removal_rejects_loaded_preset_descendant_overlap() { + let cache_root = temp_models_dir("webui-remove-preset-descendant"); + add_fake_model(&cache_root.join("mlx-community"), "parent"); + let preset_dir = cache_root.join("mlx-community/parent/preset-owned"); + std::fs::create_dir_all(&preset_dir).expect("preset dir"); + std::fs::write(preset_dir.join("config.json"), "{}").expect("preset config"); + let ini = format!("[served-descendant]\nmodel = {}\n", preset_dir.display()); + let presets = parse_preset_text(&ini).expect("parse presets"); + let pool = Arc::new(pool_from( + RouterSources { + models_dir: None, + cache: Some(CacheSource::new(cache_root.clone(), FakeDownloader::ok())), + presets, + }, + 4, + true, + )); + let cache_entry = pool.get("mlx-community/parent").expect("cache entry"); + let served = pool.get("served-descendant").expect("preset entry"); + assert_eq!(served.source, RouterModelSource::Preset); + served.lifecycle.mark_ready(); + + let err = pool + .submit_cache_removal_by_model_id( + &cache_entry.ui_model_id, + cache_entry.lifecycle_revision(), + "idem_remove_preset_descendant", + ) + .expect_err("loaded descendant preset must block parent cache removal"); + match err { + RouterPoolError::OperationRejected(error) => assert_eq!(error.code, "unsupported"), + other => panic!("unexpected error: {other:?}"), + } + let catalog = pool.catalog_snapshot(); + let cache_catalog = catalog + .iter() + .find(|model| model.name == "mlx-community/parent") + .expect("catalog cache entry"); + assert!(cache_catalog.removal_blocked_reason.is_some()); + assert!( + cache_root + .join("mlx-community/parent/config.json") + .is_file(), + "managed parent snapshot must remain after rejected removal" + ); + assert!( + preset_dir.join("config.json").is_file(), + "nested preset snapshot must remain after rejected removal" + ); + assert_eq!( + served.lifecycle.state(), + crate::server::router_lifecycle::ModelLifecycleState::Ready + ); +} + #[test] fn rescan_preserves_entry_that_becomes_reserved_after_snapshot_clone() { let cache_root = temp_models_dir("reserved-after-clone"); diff --git a/src/server/router_server.rs b/src/server/router_server.rs index d9131a701..e37d742d3 100644 --- a/src/server/router_server.rs +++ b/src/server/router_server.rs @@ -1234,10 +1234,10 @@ async fn ui_events( ) } -/// POST /models (b10621 `post_router_models`): validate the name as a -/// fetchable HuggingFace repository, then download it into the model cache +/// POST /models (b10621 `post_router_models`): enqueue a cache download /// in the background, reporting progress through `GET /models/sse` -/// (issue #1438). +/// (issue #1438). Network validation runs inside the bounded download +/// operation so duplicate/queue admission happens before any Hub request. async fn router_models_add( State(state): State, body: axum::body::Bytes, @@ -1266,17 +1266,6 @@ async fn router_models_add( if repo_id != name && state.pool.lookup(&repo_id).is_some() { return llama_invalid_request(&format!("model '{repo_id}' already exists")); } - // b10621 validates by fetching repository metadata before answering; a - // failed probe is a 500 from its handler wrapper. The probe blocks on - // the network, so it runs on the blocking pool. - let probe_pool = state.pool.clone(); - let probe_repo = repo_id.clone(); - let probed = tokio::task::spawn_blocking(move || probe_pool.validate_cache_repo(&probe_repo)) - .await - .unwrap_or_else(|join_err| Err(anyhow::anyhow!(join_err.to_string()))); - if let Err(err) = probed { - return llama_server_error(&format!("model validation failed: {err:#}")); - } match state.pool.start_download(&repo_id) { Ok(()) => Json(serde_json::json!({ "success": true })).into_response(), Err(err) => pool_error_response(err), @@ -1501,6 +1490,7 @@ fn router_ui_routes(state: RouterServerState) -> axum::Router post(ui_operation_cancel), ) .route("/ui-api/v1/events", get(ui_events)) + .nest("/ui-api/v1", super::webui::library::routes()) .layer(middleware::from_fn_with_state( state, router_ui_api_key_auth, diff --git a/src/server/router_server_security_support_tests.rs b/src/server/router_server_security_support_tests.rs index a3e10c4ad..482a9f20b 100644 --- a/src/server/router_server_security_support_tests.rs +++ b/src/server/router_server_security_support_tests.rs @@ -179,13 +179,14 @@ fn normalize_bootstrap_fixture(value: &mut serde_json::Value, expected: &serde_j struct NoBootstrapDownload; impl crate::server::router_cache::RouterDownloader for NoBootstrapDownload { - fn validate(&self, _: &str) -> anyhow::Result<()> { + fn validate(&self, _: &str, _: Option<&str>) -> anyhow::Result<()> { panic!("bootstrap must not probe the network"); } fn download( &self, _: &str, + _: Option<&str>, _: &Path, _: crate::downloader::DownloadHooks, ) -> anyhow::Result<()> { @@ -245,8 +246,18 @@ async fn assert_mounted_bootstrap_fixture(with_cache: bool) { // compare every producer field without normalizing away this state. expected["roots"] = serde_json::json!([]); for action in ["download", "cache_delete"] { - expected["actions"][action]["reason"] = - serde_json::json!("no writable managed cache route is available in this mode"); + expected["actions"][action] = serde_json::json!({ + "state": "disabled", "reason": "no managed model cache is configured", + "instructions": "Configure --model-store-root and restart to enable library mutations" + }); + for pointer in ["/features", "/server/build/features"] { + expected + .pointer_mut(pointer) + .unwrap() + .as_array_mut() + .unwrap() + .retain(|feature| feature != action); + } } } normalize_bootstrap_fixture(&mut actual, &expected); @@ -368,7 +379,7 @@ async fn secured_router_mounts_static_bootstrap_catalog_runtime_and_events() { let bootstrap_json: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); assert_eq!(bootstrap_json["schema_version"], "webui.ui-api.v1"); assert_eq!(bootstrap_json["server"]["mode"], "router_pool"); - assert_eq!(bootstrap_json["actions"]["download"]["state"], "read_only"); + assert_eq!(bootstrap_json["actions"]["download"]["state"], "disabled"); let catalog = secured_request( secured_router_app_with_limits(32, 16), @@ -696,3 +707,6 @@ async fn assembled_ui_router_body_boundary_reaches_actual_action_handler() { async fn assembled_ui_off_router_body_boundary_preserves_legacy_load_handler() { assembled_router_boundary(false).await; } + +#[path = "router_library_integration_tests.rs"] +mod library_integration_tests; diff --git a/src/server/router_server_tests.rs b/src/server/router_server_tests.rs index 60e9b8221..8878f8465 100644 --- a/src/server/router_server_tests.rs +++ b/src/server/router_server_tests.rs @@ -76,13 +76,14 @@ fn add_catalog_model(root: &std::path::Path, name: &str, model_type: &str) { struct InstantDownloader; impl RouterDownloader for InstantDownloader { - fn validate(&self, _repo_id: &str) -> anyhow::Result<()> { + fn validate(&self, _repo_id: &str, _revision: Option<&str>) -> anyhow::Result<()> { Ok(()) } fn download( &self, repo_id: &str, + _revision: Option<&str>, dest_root: &Path, hooks: DownloadHooks, ) -> anyhow::Result<()> { @@ -239,20 +240,10 @@ fn assert_bootstrap_reports_cache_authority(body: &serde_json::Value) { assert_eq!(body["server"]["mode"], "model_free", "{body}"); assert_eq!(body["actions"]["load"]["state"], "enabled", "{body}"); assert_eq!(body["actions"]["unload"]["state"], "enabled", "{body}"); - assert_eq!(body["actions"]["download"]["state"], "read_only", "{body}"); - assert_eq!( - body["actions"]["download"]["reason"], "download adapter is not mounted in this build", - "{body}" - ); - assert_eq!( - body["actions"]["cache_delete"]["state"], "read_only", - "{body}" - ); - assert_eq!( - body["actions"]["cache_delete"]["reason"], - "cache removal adapter is not mounted in this build", - "{body}" - ); + for name in ["download", "cache_delete"] { + assert_eq!(body["actions"][name]["state"], "enabled", "{body}"); + assert!(body["actions"][name]["reason"].is_null(), "{body}"); + } let roots = body["roots"].as_array().expect("roots array"); assert!( roots.iter().any(|root| root["kind"] == "cache" @@ -262,6 +253,7 @@ fn assert_bootstrap_reports_cache_authority(body: &serde_json::Value) { ); } +#[allow(clippy::duplicate_mod)] #[path = "router_contract_test_support.rs"] mod contract; #[cfg(feature = "webui")] @@ -1114,6 +1106,87 @@ async fn ui_operations_routes_list_get_and_report_cancel_unsupported() { assert_eq!(body["error"]["operation_id"], accepted.operation_id); } +#[tokio::test] +async fn ui_download_route_replays_same_idempotency_key() { + let root = temp_models_dir("ui-download-route"); + let cache_root = temp_models_dir("ui-download-route-cache"); + let state = router_state_from( + RouterSources { + models_dir: Some(root), + cache: Some(CacheSource::new(cache_root, Arc::new(InstantDownloader))), + presets: Default::default(), + }, + keyed_config(), + true, + ); + let app = create_router_app_with_authenticated_ui(state); + let body = serde_json::json!({ + "repo_id": "mlx-community/replay-http", + "idempotency_key": "download-route-replay-0001" + }) + .to_string(); + + let (first_status, first) = send( + app.clone(), + Method::POST, + "/ui-api/v1/downloads", + &body, + Some(ROUTER_KEY), + ) + .await; + assert_eq!(first_status, StatusCode::ACCEPTED, "{first}"); + contract::assert_operation_accepted(&first); + assert_eq!(first["idempotent_replay"], false); + + let (second_status, second) = send( + app, + Method::POST, + "/ui-api/v1/downloads", + &body, + Some(ROUTER_KEY), + ) + .await; + assert_eq!(second_status, StatusCode::ACCEPTED, "{second}"); + assert_eq!(second["operation_id"], first["operation_id"]); + assert_eq!(second["idempotent_replay"], true); +} + +#[tokio::test] +async fn ui_model_removal_route_matches_operation_accepted_fixture() { + let root = temp_models_dir("ui-removal-route"); + let cache_root = temp_models_dir("ui-removal-route-cache"); + add_fake_model(&cache_root.join("mlx-community"), "remove-route"); + let state = router_state_from( + RouterSources { + models_dir: Some(root), + cache: Some(CacheSource::new(cache_root, Arc::new(InstantDownloader))), + presets: Default::default(), + }, + keyed_config(), + true, + ); + let entry = state + .pool + .get("mlx-community/remove-route") + .expect("cache entry"); + let body = serde_json::json!({ + "model_id": entry.ui_model_id, + "expected_revision": entry.lifecycle_revision(), + "idempotency_key": "removal-route-0001" + }); + let app = create_router_app_with_authenticated_ui(state); + let (status, response) = send( + app, + Method::POST, + "/ui-api/v1/model-removals", + &body.to_string(), + Some(ROUTER_KEY), + ) + .await; + assert_eq!(status, StatusCode::ACCEPTED, "{response}"); + contract::assert_operation_accepted(&response); +} + #[tokio::test] async fn ui_model_action_route_validates_profile_fields_and_idempotency() { let root = temp_models_dir("ui-action-profile"); diff --git a/src/server/router_unload_revision_tests.rs b/src/server/router_unload_revision_tests.rs new file mode 100644 index 000000000..97bf24b5b --- /dev/null +++ b/src/server/router_unload_revision_tests.rs @@ -0,0 +1,254 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Real recording-provider worker threads and lifecycle leases, without MLX. + +use super::*; +use std::time::{Duration, Instant}; + +struct Fixture { + _root: tempfile::TempDir, + pool: Arc, + entry: Arc, + provider: Arc, +} + +impl Fixture { + fn new() -> Self { + let root = tempfile::tempdir().unwrap(); + let path = root.path().join("fake"); + std::fs::create_dir(&path).unwrap(); + std::fs::write(path.join("config.json"), "{}").unwrap(); + let pool = Arc::new( + RouterPool::new( + RouterSources { + models_dir: Some(root.path().into()), + ..Default::default() + }, + ServerStartupConfig::default(), + Default::default(), + PresetCliOverrides::default(), + 2, + false, + ) + .unwrap(), + ); + let entry = pool.get("fake").unwrap(); + let (tx, _rx) = std::sync::mpsc::channel(); + let provider = Arc::new(ModelProvider::recording_for_route_tests(tx)); + let state = AppState::new( + provider.clone(), + entry.config.clone(), + ChatTemplateProcessor::with_template("ok".into()), + crate::tokenizer::MlxcelTokenizer::stub(), + path, + provider.batch_metrics().clone(), + ); + entry.state.lock().unwrap().app = Some(LoadedApp { + state, + router: axum::Router::new(), + }); + entry.lifecycle.mark_loading(); + entry.lifecycle.mark_ready(); + Self { + _root: root, + pool, + entry, + provider, + } + } + + fn unload(&self) -> super::super::router_lifecycle::OperationAccepted { + self.pool + .submit_model_action( + &self.entry.ui_model_id, + RouterModelAction::Unload, + self.entry.lifecycle_revision(), + "unload-revision-test", + None, + ) + .unwrap() + } +} + +impl Drop for Fixture { + fn drop(&mut self) { + self.provider.shutdown_worker(); + assert!( + self.provider + .worker_exit_observer() + .wait_timeout(Duration::from_secs(2)) + ); + } +} + +async fn wait_for(mut predicate: impl FnMut() -> bool) { + let deadline = Instant::now() + Duration::from_secs(5); + while !predicate() { + assert!( + Instant::now() < deadline, + "bounded lifecycle observation timed out" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } +} + +async fn terminal(fixture: &Fixture, id: &str) -> super::super::router_lifecycle::Operation { + let coordinator = fixture.pool.lifecycle_coordinator(); + wait_for(|| { + coordinator.get_operation(id).is_some_and(|op| { + matches!( + op.state, + OperationState::Succeeded | OperationState::Failed | OperationState::Cancelled + ) + }) + }) + .await; + coordinator.get_operation(id).unwrap() +} + +#[tokio::test] +async fn valid_revision_unload_drains_request_and_observes_worker_exit() { + let fixture = Fixture::new(); + let caller_revision = fixture.entry.lifecycle_revision(); + let lease = fixture.entry.lifecycle.clone().try_request_lease().unwrap(); + let accepted = fixture.unload(); + wait_for(|| fixture.entry.lifecycle.state() == ModelLifecycleState::Draining).await; + let drain_revision = fixture.entry.lifecycle_revision(); + assert!(drain_revision > caller_revision); + assert!(!fixture.provider.worker_exit_observed()); + assert!(fixture.entry.state.lock().unwrap().app.is_some()); + assert_eq!( + fixture + .pool + .lifecycle_coordinator() + .get_operation(&accepted.operation_id) + .unwrap() + .state, + OperationState::Running + ); + // Legitimate discovery cannot replace a loaded/draining entry. + fixture.pool.rescan().unwrap(); + assert!(Arc::ptr_eq( + &fixture.pool.get("fake").unwrap(), + &fixture.entry + )); + drop(lease); + let operation = terminal(&fixture, &accepted.operation_id).await; + assert_eq!(operation.state, OperationState::Succeeded, "{operation:?}"); + assert!(fixture.provider.worker_exit_observed()); + assert!(fixture.entry.state.lock().unwrap().app.is_none()); + assert_eq!( + fixture.entry.lifecycle.state(), + ModelLifecycleState::Unloaded + ); + assert!(fixture.entry.lifecycle_revision() > drain_revision); + assert!( + matches!(operation.result, Some(OperationResult::ModelUnload { lifecycle, .. }) if lifecycle.worker_exit_observed) + ); +} + +#[tokio::test] +async fn unload_rejects_external_revision_change_during_drain_without_stopping_worker() { + let fixture = Fixture::new(); + let lease = fixture.entry.lifecycle.clone().try_request_lease().unwrap(); + let accepted = fixture.unload(); + wait_for(|| fixture.entry.lifecycle.state() == ModelLifecycleState::Draining).await; + let drain_revision = fixture.entry.lifecycle_revision(); + fixture.entry.lifecycle.mark_operation_busy(); + assert!(fixture.entry.lifecycle_revision() > drain_revision); + drop(lease); + let operation = terminal(&fixture, &accepted.operation_id).await; + assert_eq!(operation.state, OperationState::Failed); + assert_eq!(operation.error.unwrap().code, "stale_revision"); + assert!(!fixture.provider.worker_exit_observed()); + assert!(fixture.entry.state.lock().unwrap().app.is_some()); +} + +#[tokio::test] +async fn unload_rejects_replaced_entry_during_drain_without_stopping_either_worker() { + let fixture = Fixture::new(); + let replacement = Fixture::new(); + let lease = fixture.entry.lifecycle.clone().try_request_lease().unwrap(); + let accepted = fixture.unload(); + wait_for(|| fixture.entry.lifecycle.state() == ModelLifecycleState::Draining).await; + // Test-only adversarial registry replacement; production rescan preserves + // the draining entry, as the success regression above verifies. + fixture + .pool + .entries + .write() + .unwrap() + .insert("fake".into(), replacement.entry.clone()); + drop(lease); + let operation = terminal(&fixture, &accepted.operation_id).await; + assert_eq!(operation.state, OperationState::Failed); + assert_eq!(operation.error.unwrap().code, "stale_revision"); + assert!(!fixture.provider.worker_exit_observed()); + assert!(!replacement.provider.worker_exit_observed()); +} + +#[test] +fn owned_drain_token_uses_shared_authority_and_survives_lease_completion() { + let authority = Arc::new(AtomicU64::new(10)); + let lifecycle = Arc::new(ModelLifecycle::new_with_revision_authority( + DownloadState::Complete, + authority.clone(), + )); + lifecycle.mark_loading(); + lifecycle.mark_ready(); + let before = lifecycle.revision(); + let generation = lifecycle.generation(); + let lease = lifecycle.clone().try_request_lease().unwrap(); + let other = ModelLifecycle::new_with_revision_authority(DownloadState::Complete, authority); + assert!(other.revision() > before); + let (changed, token) = lifecycle.begin_drain_with_revision().unwrap(); + assert!(changed); + assert!( + token > before + 1, + "another model consumed a shared revision" + ); + drop(lease); + assert_eq!(lifecycle.revision(), token); + assert_eq!(lifecycle.generation(), generation); + assert_eq!(lifecycle.snapshot().active_requests, 0); + assert_eq!(lifecycle.begin_drain_with_revision(), Some((false, token))); +} + +#[tokio::test] +async fn legacy_unload_cleans_failed_worker_without_new_revision_precondition() { + let fixture = Fixture::new(); + let lease = fixture.entry.lifecycle.clone().try_request_lease().unwrap(); + let pool = fixture.pool.clone(); + let task = tokio::spawn(async move { pool.unload("fake").await }); + wait_for(|| fixture.entry.lifecycle.state() == ModelLifecycleState::Draining).await; + fixture.provider.shutdown_worker(); + wait_for(|| fixture.provider.worker_exit_observed()).await; + fixture + .entry + .lifecycle + .mark_failed("recording worker exited during drain", true); + drop(lease); + tokio::time::timeout(Duration::from_secs(5), task) + .await + .unwrap() + .unwrap() + .unwrap(); + assert!(fixture.entry.state.lock().unwrap().app.is_none()); + assert_eq!( + fixture.entry.lifecycle.state(), + ModelLifecycleState::Unloaded + ); + assert!(fixture.entry.lifecycle_snapshot().worker_exit_observed); +} diff --git a/src/server/webui/api.rs b/src/server/webui/api.rs index e4c04f3d3..55180913c 100644 --- a/src/server/webui/api.rs +++ b/src/server/webui/api.rs @@ -175,32 +175,13 @@ pub(crate) fn bootstrap_response( ); actions.insert( "download", - action( - WebUiActionState::ReadOnly, - Some( - if cache_available && !matches!(mode, WebUiServerMode::SingleModel) { - "download adapter is not mounted in this build" - } else { - "no writable managed cache route is available in this mode" - }, - ), - Some("Downloads are handled by the WebUI download adapter in issue #1841"), - ), + super::library_policy::availability(mode, cache_available, true), ); actions.insert( "cache_delete", - action( - WebUiActionState::ReadOnly, - Some( - if cache_available && !matches!(mode, WebUiServerMode::SingleModel) { - "cache removal adapter is not mounted in this build" - } else { - "no writable managed cache route is available in this mode" - }, - ), - Some("Cache deletion is handled by the WebUI removal adapter in issue #1841"), - ), + super::library_policy::availability(mode, cache_available, false), ); + let features = feature_flags(config, &actions); BootstrapResponse { schema_version: SCHEMA_VERSION.to_string(), server: BackendIdentity { @@ -208,9 +189,9 @@ pub(crate) fn bootstrap_response( mode: mode.as_str(), api_base: config.api_prefix.clone(), auth_required: !config.api_keys.is_empty(), - build: build_info(config, cache_available), + build: build_info(features.clone()), }, - features: feature_flags(config, cache_available), + features, actions, roots: root_summaries(startup, cache_available, mode), limits: limit_summary(), @@ -229,21 +210,31 @@ fn action( } } -fn feature_flags(config: &ServerConfig, cache_available: bool) -> Vec<&'static str> { +fn feature_flags( + config: &ServerConfig, + actions: &BTreeMap<&'static str, ActionAvailability>, +) -> Vec<&'static str> { let mut features = vec!["webui", "catalog", "load", "unload", "chat", "runtime"]; - let _ = cache_available; + for name in ["download", "cache_delete"] { + if actions + .get(name) + .is_some_and(|action| action.state == "enabled") + { + features.push(name); + } + } if config.enable_settings_endpoint { features.push("settings"); } features } -fn build_info(config: &ServerConfig, cache_available: bool) -> BuildInfo { +fn build_info(features: Vec<&'static str>) -> BuildInfo { BuildInfo { version: env!("CARGO_PKG_VERSION"), git_commit: option_env!("MLXCEL_GIT_COMMIT"), target: format!("{}-{}", std::env::consts::ARCH, std::env::consts::OS), - features: feature_flags(config, cache_available), + features, } } @@ -258,7 +249,11 @@ fn root_summaries( kind: "cache", display_name: "mlxcel managed cache".to_string(), redacted: true, - writable: Some(true), + writable: Some( + cache_available + && mode != WebUiServerMode::SingleModel + && crate::server::router_cache::managed_mutations_supported(), + ), error: None, }); } diff --git a/src/server/webui/catalog_metadata.rs b/src/server/webui/catalog_metadata.rs index 7b3772cb0..a1e3384b5 100644 --- a/src/server/webui/catalog_metadata.rs +++ b/src/server/webui/catalog_metadata.rs @@ -45,7 +45,11 @@ pub(super) fn catalog_entry(model: RouterCatalogModel) -> CatalogEntry { let supported = metadata.support.architecturally_supported && metadata.support.runnable_on_backend && complete; - let removal = removal_status(model.source, &model.lifecycle); + let removal = removal_status( + model.source, + &model.lifecycle, + model.removal_blocked_reason.as_deref(), + ); let mut entry = CatalogEntry { identity: ModelIdentity { id: model.ui_model_id, @@ -77,7 +81,11 @@ pub(super) fn apply_runtime_fields(entry: &mut CatalogEntry, model: &RouterCatal entry.identity.revision = model.revision; entry.identity.generation = model.generation; entry.lifecycle = model.lifecycle.clone(); - let removal = removal_status(model.source, &model.lifecycle); + let removal = removal_status( + model.source, + &model.lifecycle, + model.removal_blocked_reason.as_deref(), + ); entry.removable = removal.eligible; entry.removal = removal; apply_provider_confirmed_capabilities(entry, model.provider_capabilities); @@ -450,7 +458,11 @@ fn support_reason(architectural: bool, runnable: bool, complete_reason: &str) -> } } -fn removal_status(source: RouterModelSource, lifecycle: &LifecycleSnapshot) -> RemovalStatus { +fn removal_status( + source: RouterModelSource, + lifecycle: &LifecycleSnapshot, + blocked_reason: Option<&str>, +) -> RemovalStatus { if source != RouterModelSource::Cache { return RemovalStatus { eligible: false, @@ -461,6 +473,16 @@ fn removal_status(source: RouterModelSource, lifecycle: &LifecycleSnapshot) -> R ), }; } + if let Some(reason) = blocked_reason { + return RemovalStatus { + eligible: false, + reason: Some(reason.to_string()), + instructions: Some( + "Remove the non-cache alias or reconfigure overlapping roots before deleting the managed cache snapshot." + .to_string(), + ), + }; + } if lifecycle.busy { return RemovalStatus { eligible: false, diff --git a/src/server/webui/catalog_tests.rs b/src/server/webui/catalog_tests.rs index 55257629d..bc84aa685 100644 --- a/src/server/webui/catalog_tests.rs +++ b/src/server/webui/catalog_tests.rs @@ -96,6 +96,7 @@ fn model_with_lifecycle( ui_model_id, source_key_hash, hidden: false, + removal_blocked_reason: None, lifecycle, revision, generation: 1, diff --git a/src/server/webui/library.rs b/src/server/webui/library.rs new file mode 100644 index 000000000..dc1c657cd --- /dev/null +++ b/src/server/webui/library.rs @@ -0,0 +1,416 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! WebUI library management routes. +//! +//! These routes expose safe model-library operations to the browser without +//! changing router startup state: they operate only on the existing +//! [`RouterServerState`] and delegate all authority to [`RouterPool`]. + +use axum::Router; +use axum::extract::State; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Json, Response}; +use axum::routing::post; + +use super::super::router_lifecycle::{ErrorBody, ErrorEnvelope, FieldError}; +use super::super::router_models::RouterPoolError; +use super::super::router_server::RouterServerState; + +const MODEL_ID_PREFIX: &str = "mdl_"; +const MODEL_ID_SUFFIX_LEN: usize = 43; +const IDEMPOTENCY_KEY_MIN: usize = 8; +const IDEMPOTENCY_KEY_MAX: usize = 128; +const REVISION_REF_MAX: usize = 128; + +#[derive(Debug, serde::Deserialize)] +#[serde(deny_unknown_fields)] +struct DownloadRequest { + repo_id: String, + revision: Option, + idempotency_key: String, +} + +#[derive(Debug, serde::Deserialize)] +#[serde(deny_unknown_fields)] +struct RemovalRequest { + model_id: String, + expected_revision: u64, + idempotency_key: String, +} + +/// Routes mounted by the startup/WebUI owner under `/ui-api/v1`. +pub(crate) fn routes() -> Router { + Router::new() + .route("/downloads", post(ui_downloads)) + .route("/model-removals", post(ui_model_removals)) +} + +async fn ui_downloads(State(state): State, body: axum::body::Bytes) -> Response { + let request: DownloadRequest = match serde_json::from_slice(&body) { + Ok(request) => request, + Err(err) => { + return webui_error( + StatusCode::BAD_REQUEST, + "invalid_request", + format!("invalid download request: {err}"), + true, + ); + } + }; + if let Some(response) = validate_repo_id(&request.repo_id) { + return response; + } + if let Some(revision) = request.revision.as_deref() + && let Some(response) = validate_revision_ref(revision) + { + return response; + } + if let Some(response) = validate_idempotency_key(&request.idempotency_key) { + return response; + } + if let Some(response) = mutation_unavailable(&state, true) { + return response; + } + match state.pool.submit_download( + &request.repo_id, + request.revision.as_deref(), + Some(&request.idempotency_key), + ) { + Ok(accepted) => (StatusCode::ACCEPTED, Json(accepted)).into_response(), + Err(err) => webui_pool_error_response(err), + } +} + +async fn ui_model_removals( + State(state): State, + body: axum::body::Bytes, +) -> Response { + let request: RemovalRequest = match serde_json::from_slice(&body) { + Ok(request) => request, + Err(err) => { + return webui_error( + StatusCode::BAD_REQUEST, + "invalid_request", + format!("invalid model removal request: {err}"), + true, + ); + } + }; + if let Some(response) = validate_model_id(&request.model_id, "model_id") { + return response; + } + if request.expected_revision == 0 { + return invalid_webui_field( + "expected_revision", + "out_of_range", + "expected_revision must be at least 1", + ); + } + if let Some(response) = validate_idempotency_key(&request.idempotency_key) { + return response; + } + if let Some(response) = mutation_unavailable(&state, false) { + return response; + } + match state.pool.submit_cache_removal_by_model_id( + &request.model_id, + request.expected_revision, + &request.idempotency_key, + ) { + Ok(accepted) => (StatusCode::ACCEPTED, Json(accepted)).into_response(), + Err(err) => webui_pool_error_response(err), + } +} + +fn mutation_unavailable(state: &RouterServerState, download: bool) -> Option { + let availability = super::library_policy::availability( + super::api::WebUiServerMode::RouterPool, + state.pool.has_cache(), + download, + ); + (availability.state != "enabled").then(|| { + webui_error( + StatusCode::UNPROCESSABLE_ENTITY, + "unsupported", + availability + .reason + .unwrap_or("library mutation is unavailable"), + false, + ) + }) +} + +fn request_id() -> String { + format!("req_{}", chrono::Utc::now().timestamp_micros()) +} + +fn webui_error( + status: StatusCode, + code: &str, + message: impl Into, + retryable: bool, +) -> Response { + ( + status, + Json(ErrorEnvelope { + error: ErrorBody { + code: code.to_string(), + message: message.into(), + retryable, + field_errors: None, + operation_id: None, + }, + request_id: request_id(), + }), + ) + .into_response() +} + +fn webui_field_error( + status: StatusCode, + code: &str, + message: impl Into, + field: &str, + field_code: &str, + field_message: impl Into, +) -> Response { + ( + status, + Json(ErrorEnvelope { + error: ErrorBody { + code: code.to_string(), + message: message.into(), + retryable: true, + field_errors: Some(vec![FieldError { + field: field.to_string(), + code: field_code.to_string(), + message: field_message.into(), + }]), + operation_id: None, + }, + request_id: request_id(), + }), + ) + .into_response() +} + +fn invalid_webui_field( + field: &str, + field_code: &str, + field_message: impl Into, +) -> Response { + webui_field_error( + StatusCode::BAD_REQUEST, + "invalid_request", + "request does not match the WebUI contract", + field, + field_code, + field_message, + ) +} + +fn webui_pool_error_response(err: RouterPoolError) -> Response { + match err { + RouterPoolError::OperationRejected(error) => operation_error_response(error), + RouterPoolError::NotFound(_) => webui_error( + StatusCode::NOT_FOUND, + "not_found", + "model was not found; refresh the catalog before retrying", + true, + ), + RouterPoolError::NotRemovable(_) => webui_error( + StatusCode::UNPROCESSABLE_ENTITY, + "unsupported", + "model is not removable", + false, + ), + RouterPoolError::AlreadyExists(_) => webui_error( + StatusCode::CONFLICT, + "conflict", + "model already exists", + true, + ), + RouterPoolError::Capacity(_) => webui_error( + StatusCode::CONFLICT, + "conflict", + "model lifecycle capacity is unavailable; refresh and retry", + true, + ), + RouterPoolError::MissingName => webui_error( + StatusCode::BAD_REQUEST, + "invalid_request", + "model name is required", + true, + ), + RouterPoolError::NotLoaded => webui_error( + StatusCode::CONFLICT, + "conflict", + "model is not loaded", + true, + ), + RouterPoolError::LoadFailed(_) | RouterPoolError::LoadFailedWithEviction { .. } => { + webui_error( + StatusCode::CONFLICT, + "conflict", + "model load failed; see server logs", + true, + ) + } + } +} + +fn operation_error_response(error: ErrorBody) -> Response { + let status = match error.code.as_str() { + "not_found" => StatusCode::NOT_FOUND, + "unsupported" | "auth_required" => StatusCode::UNPROCESSABLE_ENTITY, + "rate_limited" => StatusCode::TOO_MANY_REQUESTS, + "invalid_request" => StatusCode::BAD_REQUEST, + _ => StatusCode::CONFLICT, + }; + ( + status, + Json(ErrorEnvelope { + error, + request_id: request_id(), + }), + ) + .into_response() +} + +fn validate_idempotency_key(value: &str) -> Option { + if value.len() < IDEMPOTENCY_KEY_MIN { + return Some(invalid_webui_field( + "idempotency_key", + "too_short", + "idempotency_key must contain at least 8 characters", + )); + } + if value.len() > IDEMPOTENCY_KEY_MAX { + return Some(invalid_webui_field( + "idempotency_key", + "too_long", + "idempotency_key must contain at most 128 characters", + )); + } + if !is_webui_token(value, IDEMPOTENCY_KEY_MIN, IDEMPOTENCY_KEY_MAX) { + return Some(invalid_webui_field( + "idempotency_key", + "invalid_format", + "idempotency_key may only contain A-Z, a-z, 0-9, dot, underscore, tilde and dash", + )); + } + None +} + +fn validate_model_id(value: &str, field: &str) -> Option { + let Some(suffix) = value.strip_prefix(MODEL_ID_PREFIX) else { + return Some(invalid_webui_field( + field, + "invalid_format", + "model_id must match ^mdl_[A-Za-z0-9_-]{43}$", + )); + }; + if suffix.len() != MODEL_ID_SUFFIX_LEN + || !suffix + .as_bytes() + .iter() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-')) + { + return Some(invalid_webui_field( + field, + "invalid_format", + "model_id must match ^mdl_[A-Za-z0-9_-]{43}$", + )); + } + None +} + +fn validate_repo_id(value: &str) -> Option { + let mut parts = value.split('/'); + let owner = parts.next().unwrap_or_default(); + let name = parts.next().unwrap_or_default(); + if parts.next().is_some() || !valid_hf_repo_segment(owner) || !valid_hf_repo_segment(name) { + return Some(invalid_webui_field( + "repo_id", + "invalid_format", + "repo_id must be a public HuggingFace owner/name id using A-Z, a-z, 0-9, dot, underscore and dash", + )); + } + None +} + +fn validate_revision_ref(value: &str) -> Option { + if value.len() > REVISION_REF_MAX || value.is_empty() || !valid_revision_ref(value) { + return Some(invalid_webui_field( + "revision", + "invalid_format", + "revision must be 1-128 characters and may only contain A-Z, a-z, 0-9, dot, underscore, tilde, plus, slash and dash; dot-only path segments are rejected", + )); + } + None +} + +fn is_webui_token_byte(byte: u8) -> bool { + byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'~' | b'-') +} + +fn is_webui_token(value: &str, min: usize, max: usize) -> bool { + let bytes = value.as_bytes(); + (min..=max).contains(&bytes.len()) && bytes.iter().copied().all(is_webui_token_byte) +} + +fn valid_hf_repo_segment(segment: &str) -> bool { + let bytes = segment.as_bytes(); + (1..=96).contains(&bytes.len()) + && bytes[0].is_ascii_alphanumeric() + && segment != "." + && segment != ".." + && bytes + .iter() + .copied() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'-')) +} + +fn valid_revision_ref(value: &str) -> bool { + value.split('/').all(|segment| { + !segment.is_empty() + && segment != "." + && segment != ".." + && segment + .bytes() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'~' | b'+' | b'-')) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn repo_id_validation_rejects_path_and_url_shapes() { + assert!(validate_repo_id("mlx-community/Qwen3-4B-4bit").is_none()); + assert!(validate_repo_id("mlx-community/../secret").is_some()); + assert!(validate_repo_id("https://huggingface.co/mlx-community/Qwen3").is_some()); + assert!(validate_repo_id("mlx-community").is_some()); + } + + #[test] + fn revision_validation_rejects_dot_only_segments() { + assert!(validate_revision_ref("main").is_none()); + assert!(validate_revision_ref("refs/pr/1").is_none()); + assert!(validate_revision_ref("refs/../main").is_some()); + assert!(validate_revision_ref("").is_some()); + } +} diff --git a/src/server/webui/library_policy.rs b/src/server/webui/library_policy.rs new file mode 100644 index 000000000..d8284ce55 --- /dev/null +++ b/src/server/webui/library_policy.rs @@ -0,0 +1,77 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Observation-only policy shared by bootstrap and mutation admission. +//! Cache presence means configured authority, not a filesystem permission probe. + +use super::api::{ActionAvailability, WebUiServerMode}; + +pub(crate) fn availability( + mode: WebUiServerMode, + cache_available: bool, + download: bool, +) -> ActionAvailability { + project( + mode, + cache_available, + download, + crate::downloader::offline_mode(), + crate::server::router_cache::managed_mutations_supported(), + ) +} + +fn project( + mode: WebUiServerMode, + cache_available: bool, + download: bool, + offline: bool, + platform_supported: bool, +) -> ActionAvailability { + let (state, reason, instructions) = if mode == WebUiServerMode::SingleModel { + ( + "read_only", + Some("single-model mode does not own the router managed cache"), + Some("Restart without -m/--model to manage cached models"), + ) + } else if !cache_available { + ( + "disabled", + Some("no managed model cache is configured"), + Some("Configure --model-store-root and restart to enable library mutations"), + ) + } else if !platform_supported { + ( + "disabled", + Some("managed cache mutations require supported directory-descriptor operations"), + Some("Use a supported macOS or Linux host"), + ) + } else if download && offline { + ( + "disabled", + Some("offline mode is enabled"), + Some("Restart without --offline / LLAMA_ARG_OFFLINE to download public models"), + ) + } else { + ("enabled", None, None) + }; + ActionAvailability { + state, + reason, + instructions, + } +} + +#[cfg(test)] +#[path = "library_policy_tests.rs"] +mod tests; diff --git a/src/server/webui/library_policy_tests.rs b/src/server/webui/library_policy_tests.rs new file mode 100644 index 000000000..22c42254b --- /dev/null +++ b/src/server/webui/library_policy_tests.rs @@ -0,0 +1,44 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::{WebUiServerMode, project}; + +#[test] +fn capability_matrix_is_pure_and_fail_closed() { + for mode in [ + WebUiServerMode::SingleModel, + WebUiServerMode::ModelFree, + WebUiServerMode::RouterPool, + ] { + for cache in [false, true] { + for offline in [false, true] { + for platform in [false, true] { + for download in [false, true] { + let value = project(mode, cache, download, offline, platform); + let enabled = mode != WebUiServerMode::SingleModel + && cache + && platform + && !(download && offline); + assert_eq!(value.state == "enabled", enabled); + assert_eq!(value.reason.is_none(), enabled); + assert_eq!(value.instructions.is_none(), enabled); + if mode == WebUiServerMode::SingleModel { + assert_eq!(value.state, "read_only"); + } + } + } + } + } + } +} diff --git a/src/server/webui/mod.rs b/src/server/webui/mod.rs index f6d44f69e..286b33349 100644 --- a/src/server/webui/mod.rs +++ b/src/server/webui/mod.rs @@ -22,6 +22,8 @@ pub(crate) mod api; pub(crate) mod assets; pub(crate) mod catalog; pub(crate) mod events; +pub(crate) mod library; +pub(crate) mod library_policy; pub(crate) mod security; pub use assets::{WEBUI_PREFIX, manifest_json, router}; diff --git a/src/webui/assets/mlxcel-webui-manifest.json b/src/webui/assets/mlxcel-webui-manifest.json index c0bbc8d7b..6fd595123 100644 --- a/src/webui/assets/mlxcel-webui-manifest.json +++ b/src/webui/assets/mlxcel-webui-manifest.json @@ -55,5 +55,5 @@ }, "pnpm_version": "11.18.0", "schema_version": 1, - "source_digest_sha256": "cfbcaada7d5e9734b2d305eacf534e003d8b8ea3916ce163c5a53f294c6bbd78" + "source_digest_sha256": "299e9615a7e6730bdba3e505522e9498231fdcd33b9fff3acbf4c52558f537cf" } diff --git a/tests/fixtures/webui/examples/bootstrap.model-free.json b/tests/fixtures/webui/examples/bootstrap.model-free.json index 920a65926..d0bfb94d0 100644 --- a/tests/fixtures/webui/examples/bootstrap.model-free.json +++ b/tests/fixtures/webui/examples/bootstrap.model-free.json @@ -16,6 +16,8 @@ "unload", "chat", "runtime", + "download", + "cache_delete", "settings" ] } @@ -27,6 +29,8 @@ "unload", "chat", "runtime", + "download", + "cache_delete", "settings" ], "actions": { @@ -36,14 +40,14 @@ "instructions": null }, "download": { - "state": "read_only", - "reason": "download adapter is not mounted in this build", - "instructions": "Downloads are handled by the WebUI download adapter in issue #1841" + "state": "enabled", + "reason": null, + "instructions": null }, "cache_delete": { - "state": "read_only", - "reason": "cache removal adapter is not mounted in this build", - "instructions": "Cache deletion is handled by the WebUI removal adapter in issue #1841" + "state": "enabled", + "reason": null, + "instructions": null }, "unload": { "state": "enabled", diff --git a/tests/fixtures/webui/examples/operation.download-running.json b/tests/fixtures/webui/examples/operation.download-running.json new file mode 100644 index 000000000..052867f31 --- /dev/null +++ b/tests/fixtures/webui/examples/operation.download-running.json @@ -0,0 +1,22 @@ +{ + "operation_id": "op_download_003", + "kind": "download", + "state": "running", + "created_at": "2026-09-12T03:06:00Z", + "updated_at": "2026-09-12T03:06:05Z", + "idempotency_scope": "server_instance", + "target": { + "target_kind": "download", + "repo_id": "mlx-community/SmolLM-135M-Instruct-4bit", + "revision": "main" + }, + "progress": { + "completed_bytes": 1048576, + "total_bytes": 75789919, + "indeterminate": false + }, + "result": null, + "error": null, + "cancellable": true, + "cancel_reason": null +} diff --git a/tests/fixtures/webui/examples/operation.download-succeeded.json b/tests/fixtures/webui/examples/operation.download-succeeded.json new file mode 100644 index 000000000..5731f490d --- /dev/null +++ b/tests/fixtures/webui/examples/operation.download-succeeded.json @@ -0,0 +1,28 @@ +{ + "operation_id": "op_download_003", + "kind": "download", + "state": "succeeded", + "created_at": "2026-09-12T03:06:00Z", + "updated_at": "2026-09-12T03:06:20Z", + "idempotency_scope": "server_instance", + "target": { + "target_kind": "download", + "repo_id": "mlx-community/SmolLM-135M-Instruct-4bit", + "revision": "main" + }, + "progress": { + "completed_bytes": 75789919, + "total_bytes": 75789919, + "indeterminate": false + }, + "result": { + "result_kind": "download", + "repo_id": "mlx-community/SmolLM-135M-Instruct-4bit", + "revision": "642e06afe3fab57fd6cc518637c471af0a569e1e", + "model_id": "mdl_lR1nHQwFUguxLqHbEzH2DJLdZYiDFJ0S3FzxIzY5MUU", + "download": "complete" + }, + "error": null, + "cancellable": false, + "cancel_reason": null +} diff --git a/webui/src/api/contract.test.ts b/webui/src/api/contract.test.ts index 18b29c1fb..1ff387332 100644 --- a/webui/src/api/contract.test.ts +++ b/webui/src/api/contract.test.ts @@ -18,6 +18,8 @@ const exampleSchemas = new Map([ ['event.3.json', 'UiEvent'], ['event.gap.json', 'UiEvent'], ['operation.accepted.json', 'OperationAccepted'], + ['operation.download-running.json', 'Operation'], + ['operation.download-succeeded.json', 'Operation'], ['operation.running.json', 'Operation'], ['operation.succeeded.json', 'Operation'], ['operations.list.json', 'OperationsListResponse'], @@ -61,7 +63,7 @@ function failSchema(path: string): never { describe('canonical WebUI contract fixtures', () => { it('validates every shared fixture at the JavaScript runtime boundary', () => { const entries = Object.entries(fixtures).sort(([left], [right]) => left.localeCompare(right)); - expect(entries).toHaveLength(44); + expect(entries).toHaveLength(46); for (const [path] of entries) validateAgainstSchema(schemaFor(path), cloneFixture(path), path); });