From 952fb79edc09b889410dbb9a588fb862d15cedda Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Fri, 24 Jul 2026 13:10:04 +0200 Subject: [PATCH] Add resumable background model downloads --- Cargo.lock | 206 ++++++- Cargo.toml | 2 + PLAN.md | 138 ++++- README.md | 9 +- .../down.sql | 1 + .../up.sql | 1 + src/app.rs | 227 +++++-- src/app/view.rs | 214 ++++++- src/database.rs | 12 +- src/main.rs | 1 + src/model.rs | 567 ++++++++++++++++++ src/schema.rs | 2 +- 12 files changed, 1265 insertions(+), 115 deletions(-) create mode 100644 migrations/20260724160000_rename_dflash_to_dspark/down.sql create mode 100644 migrations/20260724160000_rename_dflash_to_dspark/up.sql create mode 100644 src/model.rs diff --git a/Cargo.lock b/Cargo.lock index 3ce9334..debc301 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -359,6 +359,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "block-buffer" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" +dependencies = [ + "hybrid-array", +] + [[package]] name = "block2" version = "0.5.1" @@ -614,6 +623,12 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "const-oid" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" + [[package]] name = "core-foundation" version = "0.9.4" @@ -686,6 +701,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + [[package]] name = "crc32fast" version = "1.5.0" @@ -736,6 +760,15 @@ dependencies = [ "typenum", ] +[[package]] +name = "crypto-common" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" +dependencies = [ + "hybrid-array", +] + [[package]] name = "ctor" version = "0.10.1" @@ -889,8 +922,19 @@ version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ - "block-buffer", - "crypto-common", + "block-buffer 0.10.4", + "crypto-common 0.1.7", +] + +[[package]] +name = "digest" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" +dependencies = [ + "block-buffer 0.12.1", + "const-oid", + "crypto-common 0.2.2", ] [[package]] @@ -1023,6 +1067,8 @@ dependencies = [ "diesel_migrations", "iced", "rfd", + "sha2", + "ureq", ] [[package]] @@ -1626,6 +1672,31 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dfa686283ad6dd069f105e5ab091b04c62850d3e4cf5d67debad1933f55023df" +[[package]] +name = "http" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "hybrid-array" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "818356c5132c1fede50f837ca96afbe78ff42413047f4abb886217845e1b6c8c" +dependencies = [ + "typenum", +] + [[package]] name = "iced" version = "0.13.1" @@ -1940,6 +2011,12 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + [[package]] name = "jni" version = "0.22.4" @@ -3215,6 +3292,20 @@ dependencies = [ "bytemuck", ] +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + [[package]] name = "roxmltree" version = "0.20.0" @@ -3288,6 +3379,41 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "rustls" +version = "0.23.42" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" +dependencies = [ + "log", + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-pki-types" +version = "1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96" +dependencies = [ + "zeroize", +] + +[[package]] +name = "rustls-webpki" +version = "0.103.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + [[package]] name = "rustversion" version = "1.0.23" @@ -3414,8 +3540,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a978451301f4db1d02937a4ab3ccce137717b81826e79b7d49ffe3244a13c3b8" dependencies = [ "cfg-if", - "cpufeatures", - "digest", + "cpufeatures 0.2.17", + "digest 0.10.7", +] + +[[package]] +name = "sha2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.3", ] [[package]] @@ -3654,6 +3791,12 @@ version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + [[package]] name = "svg_fmt" version = "0.4.5" @@ -4075,6 +4218,40 @@ version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + +[[package]] +name = "ureq" +version = "3.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dea7109cdcd5864d4eeb1b58a1648dc9bf520360d7af16ec26d0a9354bafcfc0" +dependencies = [ + "base64", + "log", + "percent-encoding", + "rustls", + "rustls-pki-types", + "ureq-proto", + "utf8-zero", + "webpki-roots", +] + +[[package]] +name = "ureq-proto" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e994ba84b0bd1b1b0cf92878b7ef898a5c1760108fe7b6010327e274917a808c" +dependencies = [ + "base64", + "http", + "httparse", + "log", +] + [[package]] name = "url" version = "2.5.8" @@ -4121,6 +4298,12 @@ dependencies = [ "xmlwriter", ] +[[package]] +name = "utf8-zero" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8c0a043c9540bae7c578c88f91dda8bd82e59ae27c21baca69c8b191aaf5a6e" + [[package]] name = "utf8_iter" version = "1.0.4" @@ -4400,6 +4583,15 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "webpki-roots" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "weezl" version = "0.1.12" @@ -5046,6 +5238,12 @@ dependencies = [ "synstructure", ] +[[package]] +name = "zeroize" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" + [[package]] name = "zerotrie" version = "0.2.4" diff --git a/Cargo.toml b/Cargo.toml index c563ebe..a8653fd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,8 @@ diesel = { version = "2.3.11", features = ["sqlite", "returning_clauses_for_sqli diesel_migrations = "2.3.2" iced = { version = "0.13.1", features = ["svg"] } rfd = "0.15.4" +sha2 = "0.11.0" +ureq = { version = "3.3.0", default-features = false, features = ["rustls"] } [package.metadata.packager] product-name = "DS4Server" diff --git a/PLAN.md b/PLAN.md index 0a5d10a..9338357 100644 --- a/PLAN.md +++ b/PLAN.md @@ -15,15 +15,16 @@ sessions, tools, and turn orchestration. | App shell | Implemented | Native Iced window, Codex-inspired two-pane layout, dark theme, embedded SVG icons, `Command-,` Preferences panel | Functional chat content and native app packaging | | Projects and sessions | Implemented for metadata | Native folder picker, project-name dialog, create/select/delete projects and sessions | Transcripts, KV payloads, rename/archive, and model binding | | Persistence | Implemented for metadata and preferences | SQLite in Application Support, Diesel ORM, embedded migrations, constrained defaults and CRUD tests | Messages, downloaded-artifact metadata, and session-runtime migrations | -| Model runtime | Preferences only | Persisted model choice, DFlash toggle, and idle timeout | Downloads, loading, inference, idle unloading, and Metal integration | +| Model runtime | Download workflow implemented | Persisted model choice, DSpark toggle, idle timeout, and Rust-native background main/DSpark downloads with progress, ETA, stop, restart-resume, size checks, and SHA-256 verification | Loading, inference, idle unloading, and Metal integration | | Local endpoint | Not started | Nothing listening | OpenAI-compatible HTTP and streaming surfaces | | Agent | UI shell only | Empty conversation pane and composer placeholder | Messages, generation, tools, approvals, compaction, and cancellation | | Dev brain | Not started | No vault access | Obsidian vault selection, durable access, and agent knowledge/memory tools | | Release | Not started | Development binary builds and tests | `.app` packaging, signing, entitlements, and notarization | The application currently manages durable project/session metadata and model -preferences. It does not yet download or load a model, serve HTTP, save chat -messages, or run an agent. +preferences, and can download the selected model's main and optional DSpark +GGUF into managed Application Support storage without blocking the UI. It does +not yet load a model, serve HTTP, save chat messages, or run an agent. ## Phase 0 — project and session shell @@ -42,8 +43,9 @@ Exit criterion: restart the app and see the same project/session tree. ## Phase 1 — model library and on-demand loading -Status: **partially implemented**. Preferences and their schema are complete; -model downloads, loading, inference, and idle unloading remain open. +Status: **partially implemented**. Preferences and the complete background +download workflow are implemented; loading, inference, and idle unloading +remain open. 1. Port the narrow GGUF loader, tokenizer, prompt renderer, and Metal session API behind a small Rust `Engine` surface modeled on `ds4.h`: @@ -58,6 +60,48 @@ model downloads, loading, inference, and idle unloading remain open. 5. Move engine work off the UI thread and support cooperative cancellation at the safe session boundaries already defined by `ds4_session_sync`. +### DS4 KV-cache port and shared session core + +KV-cache behavior is part of the DS4 engine contract, not a new application +cache to redesign. Port the existing `ds4_session` and `ds4_kvstore` behavior +to Rust while preserving its semantics and on-disk compatibility: + +- A live engine session owns one mutable inference timeline: exact rendered + tokens, logits, and the DS4-specific KV graph state. Callers provide the full + token prefix and the Rust port of `ds4_session_sync` must retain the longest + compatible prefix, evaluate only its suffix, or safely rewind, invalidate, + and rebuild exactly when DS4 does. +- Carry over DS4's safe cancellation/checkpoint boundaries, common-prefix + matching, canonical rewrite and replay, snapshots, payload staging, + save/load, strip-to-transcript behavior, and atomic replacement. The engine + continues to own the opaque KV payload format; UI, HTTP, and persistence code + must not reinterpret tensor state. +- Port `ds4_kvstore` policy rather than replacing it with a generic cache: + exact token and rendered-text prefix lookup, cold and continued aligned + checkpoints, model/quantization/context/payload-ABI compatibility checks, + tool-call replay metadata, disk-budget enforcement, hit tracking, and + eviction behavior. +- Use one process-wide Rust session registry and KV checkpoint store for both + local chat and the HTTP server. A remote request bound to an application + session resolves to the same logical DS4 session as the GUI and mutations are + serialized so two callers never write one live timeline concurrently. + Stateless HTTP requests still participate in the same compatible-prefix + checkpoint pool, so a prefix produced locally can be reused remotely and + vice versa. +- Store that shared checkpoint pool under Application Support in + `kv-cache/`, using the DS4 key/header/payload rules. Session rows hold only an + optional reference to their current checkpoint; the large opaque payload is + not copied into separate GUI and HTTP caches. +- Sharing does not merge unrelated conversations or expose one session's + transcript to another. It shares the DS4 implementation, checkpoint pool, and + explicitly bound logical session; each unrelated live conversation keeps its + own mutable session state. +- Engine unload must first leave every reusable live timeline in a valid + persisted checkpoint. Reload restores a compatible checkpoint immediately; + if model identity, quantization, context size, payload ABI, or rendered + history is incompatible, rebuild from the transcript and replace the stale + checkpoint only after the new state is valid. + ### Lazy lifecycle - Start the application and the local HTTP endpoint with no model loaded. @@ -82,24 +126,45 @@ model downloads, loading, inference, and idle unloading remain open. - Provide a Preferences panel from the sidebar and the standard macOS `Command-,` shortcut. - Let the user choose the active model. Default to `deepseek-v4-flash`. -- Let the user enable or disable DFlash for the selected model. Disable the - control for models that do not support DFlash. +- Let the user enable or disable DSpark for the selected model. Disable the + control for models that do not support DSpark. - Let the user configure the inactivity timeout; default to 10 minutes. - Persist preferences in the application SQLite database. Changing preferences never loads a model by itself. ### Targeted downloads and storage +- The Preferences panel can download the selected model's main GGUF and the + DSpark support GGUF when enabled. A dedicated worker thread performs HTTP and + verification work so Preferences and the rest of the application stay + interactive. +- Use an in-process Rust HTTP client with Rustls rather than invoking a system + downloader. Stream response chunks directly into a `.part` file and use a + Rust SHA-256 implementation for verification. +- Report aggregate bytes present, total bytes, bytes missing, percentage, + smoothed transfer rate, current artifact/verification phase, and estimated + time remaining. Keep the detailed display in Preferences and a persistent + status bar with a Stop action while the transfer runs. +- Stop is non-destructive: close the active transfer and retain its `.part` + file. Starting the download again in the same process or after quitting and + relaunching must derive progress from that file and issue an HTTP range + request from its exact byte length. If the server ignores the range request, + safely restart that artifact from byte zero rather than appending corrupt + data. +- Promote a `.part` file only after its exact size and SHA-256 are verified. A + stopped verification keeps the complete partial file so the next run can + verify and promote it without downloading it again. - Store managed model artifacts under `~/Library/Application Support/DS4Server.rfc1437.de/models//` and mmap/load them from there. Do not place managed models in the app bundle or a project directory. -- Download only the selected model's main artifact. Download its DFlash - artifact only when DFlash is enabled; never prefetch other catalog models or +- Download only the selected model's main artifact. Download its DSpark + artifact only when DSpark is enabled; never prefetch other catalog models or optional artifacts. -- Verify size/checksum before making a download usable, resume partial - downloads, and remove failed temporary files without touching a valid model. -- When the selected model changes or DFlash is disabled, offer removal of the +- Remove checksum-invalid partial files without touching a valid model. Keep + partial files after cancellation, application exit, network failure, or an + interrupted verification so they remain restart-resumable. +- When the selected model changes or DSpark is disabled, offer removal of the now-unused artifact; do not delete large existing downloads without explicit confirmation. @@ -128,41 +193,52 @@ Status: **open; not started**. 4. Port exact sampled tool-call replay and deterministic canonicalization from `ds4_server.c` so a client's normalized JSON does not destroy KV-prefix reuse. -5. Bind API conversations to engine sessions and persist/restore KV payloads - using the engine-owned serialization format. Add resident batching only - after single-session correctness is established. +5. Route API conversations through the shared session registry and KV store + defined in Phase 1. Responses `conversation` bindings and any other + persistent request identity map to the same application session used by the + GUI; stateless requests use shared DS4 prefix lookup. Persist/restore only + through the engine-owned serialization format. Add resident batching only + after single-session and cross-surface correctness are established. 6. Display endpoint address, model, active requests, token rates, and errors in the GUI. Exit criterion: the upstream DwarfStar server tests pass against the app for non-streaming and streaming Chat Completions and Responses calls, including a -multi-turn tool call. +multi-turn tool call. A compatible prefix created by local chat must register a +remote cache hit, a remote checkpoint must be reusable by local chat, and a +GUI/HTTP race on one bound session must serialize without corrupting its token +or KV frontier. ## Phase 3 — durable agent sessions Status: **partially open**. Project/session metadata rows exist; transcripts, model binding, and KV payload persistence are not implemented. -Replace each metadata-only session with a directory: +Extend each metadata-only session with transcript and shared-checkpoint +metadata: ```text Application Support/DS4Server.rfc1437.de/ data.sqlite3 models// main.gguf - dflash.gguf # only when enabled - sessions// - kv.bin + dspark.gguf # only when enabled + kv-cache/ + .kv ``` -- Add migrated message/transcript tables to SQLite. Keep only the engine-owned - large KV payload in per-session files, written with atomic replacement. +- Add migrated message/transcript tables and an optional current-checkpoint key + to SQLite. Keep large engine-owned KV payloads only in the process-wide Phase + 1 store, written with atomic replacement; do not introduce separate + agent-only or HTTP-only cache formats. - Record model identity, context size, rendered token history, title, timestamps, and working directory. - Restore compatible KV immediately. If KV is absent or incompatible, rebuild from the transcript and show prefill progress. -- Add strip/archive behavior equivalent to the native agent: retain transcript - while removing the large KV payload. +- Add strip/archive behavior equivalent to the native agent: retain the + transcript and clear its checkpoint reference. Remove the cache file only + when it is not reusable by another live/session prefix; otherwise normal DS4 + disk-budget eviction reclaims it. - Version formats before the first release and add migrations only when a real format change exists. @@ -176,7 +252,8 @@ visual placeholders with no message or agent runtime behind them. 1. Build the classic chat surface: transcript, reasoning disclosure, tool-call cards, composer, queued input, stop button, prefill progress, and generation - statistics. + statistics. Local turns must call the same session registry and KV path used + by the HTTP endpoint. 2. Port the native agent turn loop from `ds4_agent.c`, including model-specific system prompts and DeepSeek DSML/GLM tool-call parsing. 3. Implement the local tool set with the project directory as its boundary: @@ -231,6 +308,7 @@ and the preferences shortcut have unit coverage. Model fixtures, integration tests, and release packaging remain open. - Keep unit coverage narrow: persistence round trips, format compatibility, + DS4 KV prefix/rewrite/store/load parity, cross-surface session binding, request parsing, and tool-boundary checks. - Reuse `../ds4` fixtures for prompt rendering, sampling, server streaming, and KV correctness; run live Metal tests only when a supported GGUF is available. @@ -241,13 +319,11 @@ tests, and release packaging remain open. ## Recommended next milestones -1. Add the targeted download manager and model catalog, storing only selected - artifacts under Application Support with resumable, verified downloads. -2. Implement the single-engine lifecycle and local generation path: lazy load, +1. Implement the single-engine lifecycle and local generation path: lazy load, shared concurrent acquisition, activity tracking, and idle unload. -3. Put the OpenAI-compatible endpoint on that same engine lifecycle so local +2. Put the OpenAI-compatible endpoint on that same engine lifecycle so local chats and HTTP requests cannot load duplicate engines. -4. Add transcript/KV persistence, then connect the existing conversation UI and +3. Add transcript/KV persistence, then connect the existing conversation UI and finally port the agent tools. -5. Add the opt-in Dev Brain tool after local file/search boundaries and agent +4. Add the opt-in Dev Brain tool after local file/search boundaries and agent approvals are stable. diff --git a/README.md b/README.md index 63763ba..109f602 100644 --- a/README.md +++ b/README.md @@ -7,8 +7,13 @@ OpenAI-compatible localhost endpoint, and project-scoped agent chat in one app. The current milestone provides a Codex-inspired project/session layout. A native macOS folder picker selects each workspace, then the app asks for its display name. Projects, sessions, and model preferences are persisted through Diesel in -SQLite. Open Preferences with `Command-,` to select the model, toggle DFlash, -and configure the model idle timeout. +SQLite. Open Preferences with `Command-,` to select the model, toggle DSpark, +configure the model idle timeout, and download the selected main GGUF with +its optional DSpark support. Rust-native background downloads show live byte +progress, speed, and ETA in Preferences and the app status bar. Stopping or +quitting keeps the partial file; the next Download/Resume action continues from +that exact byte after relaunch. Exact size and SHA-256 verification happen +before an artifact becomes usable. ```sh cargo install cargo-packager --locked --version 0.11.8 diff --git a/migrations/20260724160000_rename_dflash_to_dspark/down.sql b/migrations/20260724160000_rename_dflash_to_dspark/down.sql new file mode 100644 index 0000000..15feb6e --- /dev/null +++ b/migrations/20260724160000_rename_dflash_to_dspark/down.sql @@ -0,0 +1 @@ +ALTER TABLE preferences RENAME COLUMN dspark_enabled TO dflash_enabled; diff --git a/migrations/20260724160000_rename_dflash_to_dspark/up.sql b/migrations/20260724160000_rename_dflash_to_dspark/up.sql new file mode 100644 index 0000000..906bd25 --- /dev/null +++ b/migrations/20260724160000_rename_dflash_to_dspark/up.sql @@ -0,0 +1 @@ +ALTER TABLE preferences RENAME COLUMN dflash_enabled TO dspark_enabled; diff --git a/src/app.rs b/src/app.rs index 0345ebe..63cdb6f 100644 --- a/src/app.rs +++ b/src/app.rs @@ -3,59 +3,24 @@ mod view; pub(crate) use view::app_theme; use crate::database::{AppPreferences, Database, ProjectWithSessions}; +use crate::model::{self, DownloadOutcome, DownloadProgress, ModelChoice}; +use iced::futures::{SinkExt, Stream}; use iced::{Subscription, Task, keyboard}; use rfd::AsyncFileDialog; -use std::fmt; use std::fs; use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::mpsc::{self, TryRecvError}; +use std::thread; +use std::time::{Duration, Instant}; const APP_ID: &str = "DS4Server.rfc1437.de"; -const MODEL_CHOICES: [ModelChoice; 3] = [ - ModelChoice::DeepSeekV4Flash, - ModelChoice::DeepSeekV4Pro, - ModelChoice::Glm52, -]; - -#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] -pub(crate) enum ModelChoice { - #[default] - DeepSeekV4Flash, - DeepSeekV4Pro, - Glm52, -} - -impl ModelChoice { - fn id(self) -> &'static str { - match self { - Self::DeepSeekV4Flash => "deepseek-v4-flash", - Self::DeepSeekV4Pro => "deepseek-v4-pro", - Self::Glm52 => "glm-5.2", - } - } - - fn from_id(id: &str) -> Option { - MODEL_CHOICES.into_iter().find(|model| model.id() == id) - } - - fn supports_dflash(self) -> bool { - self == Self::DeepSeekV4Flash - } -} - -impl fmt::Display for ModelChoice { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.write_str(match self { - Self::DeepSeekV4Flash => "DeepSeek V4 Flash", - Self::DeepSeekV4Pro => "DeepSeek V4 Pro", - Self::Glm52 => "GLM 5.2", - }) - } -} #[derive(Clone)] struct PreferenceDraft { model: ModelChoice, - dflash_enabled: bool, + dspark_enabled: bool, idle_timeout_minutes: String, } @@ -65,7 +30,7 @@ impl PreferenceDraft { .ok_or_else(|| format!("Unsupported model: {}", preferences.selected_model))?; Ok(Self { model, - dflash_enabled: model.supports_dflash() && preferences.dflash_enabled, + dspark_enabled: model.supports_dspark() && preferences.dspark_enabled, idle_timeout_minutes: preferences.idle_timeout_minutes.to_string(), }) } @@ -83,17 +48,42 @@ pub(crate) struct App { choosing_folder: bool, pending_project_path: Option, project_name_input: String, + model_download: ModelDownload, error: Option, } +#[derive(Debug)] +pub(super) enum ModelDownload { + Idle, + Active(ActiveDownload), + Complete(ModelChoice, DownloadProgress), + Failed(ModelChoice, String, DownloadProgress), +} + +#[derive(Debug)] +pub(super) struct ActiveDownload { + model: ModelChoice, + dspark_enabled: bool, + progress: DownloadProgress, + sampled_at: Instant, + sampled_bytes: u64, + bytes_per_second: f64, + cancel: Arc, + result: mpsc::Receiver>, + stopping: bool, +} + #[derive(Debug, Clone)] pub(crate) enum Message { OpenPreferences, DismissPanel, PreferenceModelChanged(ModelChoice), - PreferenceDflashChanged(bool), + PreferenceDsparkChanged(bool), PreferenceTimeoutChanged(String), SavePreferences, + DownloadSelectedModel, + StopModelDownload, + DownloadProgressTick, ChooseProjectFolder, ProjectFolderPicked(Option), ProjectNameChanged(String), @@ -128,6 +118,7 @@ impl App { choosing_folder: false, pending_project_path: None, project_name_input: String::new(), + model_download: ModelDownload::Idle, error: None, } } @@ -153,6 +144,7 @@ impl App { choosing_folder: false, pending_project_path: None, project_name_input: String::new(), + model_download: ModelDownload::Idle, error: Some(format!("Could not open the project database: {error}")), } } @@ -172,14 +164,14 @@ impl App { } Message::PreferenceModelChanged(model) => { self.preference_draft.model = model; - if !model.supports_dflash() { - self.preference_draft.dflash_enabled = false; + if !model.supports_dspark() { + self.preference_draft.dspark_enabled = false; } self.preference_error = None; } - Message::PreferenceDflashChanged(enabled) => { - self.preference_draft.dflash_enabled = - self.preference_draft.model.supports_dflash() && enabled; + Message::PreferenceDsparkChanged(enabled) => { + self.preference_draft.dspark_enabled = + self.preference_draft.model.supports_dspark() && enabled; self.preference_error = None; } Message::PreferenceTimeoutChanged(value) => { @@ -187,6 +179,48 @@ impl App { self.preference_error = None; } Message::SavePreferences => self.save_preferences(), + Message::DownloadSelectedModel => { + if matches!(self.model_download, ModelDownload::Active(_)) { + return Task::none(); + } + let model = self.preference_draft.model; + let dspark_enabled = + model.supports_dspark() && self.preference_draft.dspark_enabled; + let models_path = models_path(); + let progress = model::download_progress(model, dspark_enabled, &models_path); + let cancel = Arc::new(AtomicBool::new(false)); + let worker_cancel = Arc::clone(&cancel); + let (result_sender, result_receiver) = mpsc::channel(); + if let Err(error) = thread::Builder::new() + .name("model-download".to_owned()) + .spawn(move || { + let result = + model::download(model, dspark_enabled, &models_path, &worker_cancel); + let _ = result_sender.send(result); + }) + { + self.error = Some(format!("Could not start model download: {error}")); + return Task::none(); + } + self.model_download = ModelDownload::Active(ActiveDownload { + model, + dspark_enabled, + sampled_at: Instant::now(), + sampled_bytes: progress.downloaded, + progress, + bytes_per_second: 0.0, + cancel: Arc::clone(&cancel), + result: result_receiver, + stopping: false, + }); + } + Message::StopModelDownload => { + if let ModelDownload::Active(download) = &mut self.model_download { + download.stopping = true; + download.cancel.store(true, Ordering::Relaxed); + } + } + Message::DownloadProgressTick => self.update_download_progress(), Message::ChooseProjectFolder => { self.choosing_folder = true; return Task::perform( @@ -256,7 +290,15 @@ impl App { } pub(crate) fn subscription(&self) -> Subscription { - keyboard::on_key_press(shortcut) + let shortcuts = keyboard::on_key_press(shortcut); + if matches!(self.model_download, ModelDownload::Active(_)) { + Subscription::batch([ + shortcuts, + Subscription::run(download_ticks).map(|_| Message::DownloadProgressTick), + ]) + } else { + shortcuts + } } fn open_preferences(&mut self) { @@ -289,11 +331,11 @@ impl App { } let model = self.preference_draft.model; - let dflash_enabled = model.supports_dflash() && self.preference_draft.dflash_enabled; + let dspark_enabled = model.supports_dspark() && self.preference_draft.dspark_enabled; let Some(database) = &mut self.database else { return; }; - match database.update_preferences(model.id(), dflash_enabled, idle_timeout_minutes) { + match database.update_preferences(model.id(), dspark_enabled, idle_timeout_minutes) { Ok(preferences) => { self.preferences = preferences; self.preference_draft = PreferenceDraft::from_saved(&self.preferences) @@ -395,6 +437,73 @@ impl App { } } } + + fn update_download_progress(&mut self) { + let ModelDownload::Active(download) = &mut self.model_download else { + return; + }; + let now = Instant::now(); + let progress = + model::download_progress(download.model, download.dspark_enabled, &models_path()); + let elapsed = now.duration_since(download.sampled_at).as_secs_f64(); + let transferred = progress.downloaded.saturating_sub(download.sampled_bytes); + if transferred > 0 && elapsed > 0.0 { + let current = transferred as f64 / elapsed; + download.bytes_per_second = if download.bytes_per_second == 0.0 { + current + } else { + download.bytes_per_second * 0.75 + current * 0.25 + }; + } + download.progress = progress; + download.sampled_at = now; + download.sampled_bytes = download.progress.downloaded; + + let result = match download.result.try_recv() { + Ok(result) => Some(result), + Err(TryRecvError::Empty) => None, + Err(TryRecvError::Disconnected) => { + Some(Err("Download worker stopped unexpectedly.".into())) + } + }; + let model = download.model; + let progress = download.progress.clone(); + if let Some(result) = result { + match result { + Ok(DownloadOutcome::Complete) => { + self.model_download = ModelDownload::Complete(model, progress); + self.error = None; + } + Ok(DownloadOutcome::Stopped) => { + self.model_download = ModelDownload::Idle; + self.error = None; + } + Err(error) => { + self.error = Some(error.clone()); + self.model_download = ModelDownload::Failed(model, error, progress); + } + } + } + } +} + +impl Drop for App { + fn drop(&mut self) { + if let ModelDownload::Active(download) = &self.model_download { + download.cancel.store(true, Ordering::Relaxed); + } + } +} + +fn download_ticks() -> impl Stream { + iced::stream::channel(1, |mut output| async move { + loop { + std::thread::sleep(Duration::from_secs(1)); + if output.send(()).await.is_err() { + break; + } + } + }) } fn shortcut(key: keyboard::Key, modifiers: keyboard::Modifiers) -> Option { @@ -414,19 +523,23 @@ fn application_support_path() -> PathBuf { .join(APP_ID) } +fn models_path() -> PathBuf { + application_support_path().join("models") +} + #[cfg(test)] mod tests { use super::*; #[test] - fn preferences_shortcut_and_dflash_support_are_explicit() { + fn preferences_shortcut_and_dspark_support_are_explicit() { let message = shortcut( keyboard::Key::Character(",".into()), keyboard::Modifiers::COMMAND, ); assert!(matches!(message, Some(Message::OpenPreferences))); - assert!(ModelChoice::DeepSeekV4Flash.supports_dflash()); - assert!(!ModelChoice::DeepSeekV4Pro.supports_dflash()); - assert!(!ModelChoice::Glm52.supports_dflash()); + assert!(ModelChoice::DeepSeekV4Flash.supports_dspark()); + assert!(!ModelChoice::DeepSeekV4Pro.supports_dspark()); + assert!(!ModelChoice::Glm52.supports_dspark()); } } diff --git a/src/app/view.rs b/src/app/view.rs index d2a4c2b..87000f5 100644 --- a/src/app/view.rs +++ b/src/app/view.rs @@ -1,9 +1,10 @@ -use super::{App, MODEL_CHOICES, Message, ModelChoice}; +use super::{ActiveDownload, App, Message, ModelDownload, models_path}; use crate::database::{ProjectWithSessions, Session}; +use crate::model::{self, DownloadPhase, DownloadProgress, MODEL_CHOICES, ModelChoice}; use iced::theme::{Palette, palette}; use iced::widget::{ - Space, Svg, button, checkbox, column, container, opaque, pick_list, row, scrollable, stack, - svg, text, text_input, + Space, Svg, button, checkbox, column, container, opaque, pick_list, progress_bar, row, + scrollable, stack, svg, text, text_input, }; use iced::{Alignment, Color, Element, Length, Theme}; use std::path::Path; @@ -27,6 +28,9 @@ impl App { .width(Length::Fill) .height(Length::Fill) ]; + if let ModelDownload::Active(download) = &self.model_download { + shell = shell.push(download_status_bar(download)); + } if let Some(error) = &self.error { shell = shell.push( container(text(error).style(iced::widget::text::danger)) @@ -277,16 +281,43 @@ impl App { } fn preferences_panel(&self) -> Element<'_, Message> { - let dflash_toggle: Option Message> = self + let dspark_toggle: Option Message> = self .preference_draft .model - .supports_dflash() - .then_some(Message::PreferenceDflashChanged); - let dflash = checkbox( - "Enable DFlash for this model", - self.preference_draft.dflash_enabled, + .supports_dspark() + .then_some(Message::PreferenceDsparkChanged); + let dspark = checkbox( + "Enable DSpark for this model", + self.preference_draft.dspark_enabled, ) - .on_toggle_maybe(dflash_toggle); + .on_toggle_maybe(dspark_toggle); + let model = self.preference_draft.model; + let dspark_enabled = model.supports_dspark() && self.preference_draft.dspark_enabled; + let selected_progress = model::download_progress(model, dspark_enabled, &models_path()); + let installed = selected_progress.phase == DownloadPhase::Complete; + let download = if let ModelDownload::Active(download) = &self.model_download { + if download.stopping { + button("Stopping…") + } else { + button("Stop download") + .on_press(Message::StopModelDownload) + .style(button::danger) + } + } else if installed { + button("Downloaded") + } else { + let action = if selected_progress.downloaded > 0 { + "Resume" + } else { + "Download" + }; + button(text(format!( + "{action} {} ({})", + model, + format_bytes(selected_progress.total) + ))) + .on_press(Message::DownloadSelectedModel) + }; let mut content = column![ row![ @@ -305,13 +336,15 @@ impl App { Message::PreferenceModelChanged, ) .width(Length::Fill), - dflash, - text(if self.preference_draft.model.supports_dflash() { - "DFlash is downloaded and used only when enabled." + dspark, + text(if self.preference_draft.model.supports_dspark() { + "DSpark support is downloaded and used only when enabled." } else { - "DFlash is not available for the selected model." + "DSpark is not available for the selected model." }) .size(12), + download.style(button::secondary), + model_download_status(&self.model_download, &selected_progress), Space::with_height(8), text("INACTIVITY").size(11), row![ @@ -408,6 +441,152 @@ impl App { } } +fn download_status_bar(download: &ActiveDownload) -> Element<'_, Message> { + let progress = &download.progress; + let percent = progress.fraction() * 100.0; + let transfer = if download.bytes_per_second > 0.0 && progress.remaining() > 0 { + format!( + "{}/s • about {} remaining", + format_bytes(download.bytes_per_second as u64), + format_duration(progress.remaining() as f64 / download.bytes_per_second), + ) + } else { + "Calculating time remaining…".to_owned() + }; + let stop = if download.stopping { + button("Stopping…") + } else { + button("Stop") + .on_press(Message::StopModelDownload) + .style(button::danger) + }; + + container( + row![ + text(phase_text(progress.phase)).size(12), + progress_bar(0.0..=1.0, progress.fraction()) + .width(180) + .height(7), + text(format!( + "{} of {} ({percent:.1}%)", + format_bytes(progress.downloaded), + format_bytes(progress.total), + )) + .size(12), + text(transfer).size(12), + Space::with_width(Length::Fill), + stop, + ] + .spacing(12) + .align_y(Alignment::Center), + ) + .width(Length::Fill) + .padding([7, 14]) + .style(sidebar_style) + .into() +} + +fn model_download_status<'a>( + download: &ModelDownload, + selected: &DownloadProgress, +) -> Element<'a, Message> { + let (progress, heading, speed, failed) = match download { + ModelDownload::Idle => ( + selected, + if selected.downloaded > 0 { + "Paused — the next download will resume from this point.".to_owned() + } else { + phase_text(selected.phase) + }, + 0.0, + false, + ), + ModelDownload::Active(active) => ( + &active.progress, + if active.stopping { + format!("Stopping {}…", active.model) + } else { + phase_text(active.progress.phase) + }, + active.bytes_per_second, + false, + ), + ModelDownload::Complete(model, progress) => ( + progress, + format!("{model} is downloaded and verified."), + 0.0, + false, + ), + ModelDownload::Failed(model, error, progress) => ( + progress, + format!("{model} download failed: {error}"), + 0.0, + true, + ), + }; + + let percent = progress.fraction() * 100.0; + let mut status = column![ + if failed { + text(heading).size(12).style(iced::widget::text::danger) + } else { + text(heading).size(12) + }, + progress_bar(0.0..=1.0, progress.fraction()).height(8), + text(format!( + "{} of {} ({percent:.1}%) • {} remaining", + format_bytes(progress.downloaded), + format_bytes(progress.total), + format_bytes(progress.remaining()), + )) + .size(12), + ] + .spacing(6); + if speed > 0.0 && progress.remaining() > 0 { + status = status.push( + text(format!( + "{}/s • about {} remaining", + format_bytes(speed as u64), + format_duration(progress.remaining() as f64 / speed), + )) + .size(12), + ); + } + status.into() +} + +fn phase_text(phase: DownloadPhase) -> String { + match phase { + DownloadPhase::Pending(artifact) => format!("Ready to download {artifact}."), + DownloadPhase::Downloading(artifact) => format!("Downloading {artifact}…"), + DownloadPhase::Verifying(artifact) => format!("Verifying {artifact}…"), + DownloadPhase::Complete => "Download complete.".to_owned(), + } +} + +fn format_bytes(bytes: u64) -> String { + const GB: f64 = 1_000_000_000.0; + const MB: f64 = 1_000_000.0; + if bytes >= 1_000_000_000 { + format!("{:.1} GB", bytes as f64 / GB) + } else { + format!("{:.1} MB", bytes as f64 / MB) + } +} + +fn format_duration(seconds: f64) -> String { + let seconds = seconds.max(0.0).round() as u64; + let hours = seconds / 3600; + let minutes = seconds % 3600 / 60; + if hours > 0 { + format!("{hours}h {minutes}m") + } else if minutes > 0 { + format!("{minutes}m {}s", seconds % 60) + } else { + format!("{seconds}s") + } +} + pub(crate) fn app_theme() -> Theme { let palette = Palette { background: Color::from_rgb8(29, 29, 31), @@ -451,4 +630,11 @@ mod tests { assert_eq!(palette.background.weak.color, Color::from_rgb8(38, 38, 40)); assert_eq!(palette.secondary.base.color, Color::from_rgb8(47, 47, 50)); } + + #[test] + fn download_measurements_are_readable() { + assert_eq!(format_bytes(1_500_000_000), "1.5 GB"); + assert_eq!(format_duration(7_384.0), "2h 3m"); + assert_eq!(format_duration(125.0), "2m 5s"); + } } diff --git a/src/database.rs b/src/database.rs index bf1b541..2c8ca60 100644 --- a/src/database.rs +++ b/src/database.rs @@ -14,7 +14,7 @@ pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations"); pub struct AppPreferences { pub id: i32, pub selected_model: String, - pub dflash_enabled: bool, + pub dspark_enabled: bool, pub idle_timeout_minutes: i32, } @@ -23,7 +23,7 @@ impl Default for AppPreferences { Self { id: 1, selected_model: "deepseek-v4-flash".into(), - dflash_enabled: false, + dspark_enabled: false, idle_timeout_minutes: 10, } } @@ -33,7 +33,7 @@ impl Default for AppPreferences { #[diesel(table_name = preferences)] struct PreferenceChanges<'a> { selected_model: &'a str, - dflash_enabled: bool, + dspark_enabled: bool, idle_timeout_minutes: i32, } @@ -135,13 +135,13 @@ impl Database { pub fn update_preferences( &mut self, selected_model: &str, - dflash_enabled: bool, + dspark_enabled: bool, idle_timeout_minutes: i32, ) -> Result { diesel::update(preferences::table.find(1)) .set(PreferenceChanges { selected_model, - dflash_enabled, + dspark_enabled, idle_timeout_minutes, }) .returning(AppPreferences::as_returning()) @@ -200,7 +200,7 @@ mod tests { let preferences = database.load_preferences().unwrap(); assert_eq!(preferences.selected_model, "deepseek-v4-flash"); - assert!(!preferences.dflash_enabled); + assert!(!preferences.dspark_enabled); assert_eq!(preferences.idle_timeout_minutes, 10); assert!(database.update_preferences("glm-5.2", true, 30).is_err()); assert!( diff --git a/src/main.rs b/src/main.rs index 7e969c0..91619bd 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,6 @@ mod app; mod database; +mod model; mod schema; use app::{App, app_theme}; diff --git a/src/model.rs b/src/model.rs new file mode 100644 index 0000000..85491e7 --- /dev/null +++ b/src/model.rs @@ -0,0 +1,567 @@ +use std::fmt; +use std::fs::{self, File, OpenOptions}; +use std::io::{Read, Write}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; + +use sha2::{Digest, Sha256}; + +pub(crate) const MODEL_CHOICES: [ModelChoice; 3] = [ + ModelChoice::DeepSeekV4Flash, + ModelChoice::DeepSeekV4Pro, + ModelChoice::Glm52, +]; + +const DEEPSEEK_REPOSITORY: &str = "antirez/deepseek-v4-gguf"; +const GLM_REPOSITORY: &str = "antirez/glm-5.2-gguf"; + +const FLASH: Artifact = Artifact { + label: "DeepSeek V4 Flash model", + file_name: "DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf", + repository: DEEPSEEK_REPOSITORY, + size: 86_720_111_488, + sha256: "efc7ed607ff27076e3e501fc3fefefa33c0ed8cf1eff483a2b7fdc0c2e616668", +}; +const FLASH_DSPARK: Artifact = Artifact { + label: "DSpark support", + file_name: "DeepSeek-V4-Flash-DSpark-support.gguf", + repository: DEEPSEEK_REPOSITORY, + size: 5_989_114_272, + sha256: "8b3adf5942bec22ae2ea867cd7079cf13530ba83ffcffaf00f5de48664a1a34e", +}; +const PRO: Artifact = Artifact { + label: "DeepSeek V4 Pro model", + file_name: "DeepSeek-V4-Pro-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-Instruct-imatrix.gguf", + repository: DEEPSEEK_REPOSITORY, + size: 464_627_334_560, + sha256: "a0314d9c0e16122cd60071079124a2d17185d317c55a8f95ecb3ed3506278a96", +}; +const GLM: Artifact = Artifact { + label: "GLM 5.2 model", + file_name: "GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf", + repository: GLM_REPOSITORY, + size: 211_075_856_448, + sha256: "a49de64c5020432bdae23de36a423a9660a5621bc0db8d12b66bd8814b07fea0", +}; + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub(crate) enum ModelChoice { + #[default] + DeepSeekV4Flash, + DeepSeekV4Pro, + Glm52, +} + +impl ModelChoice { + pub(crate) fn id(self) -> &'static str { + match self { + Self::DeepSeekV4Flash => "deepseek-v4-flash", + Self::DeepSeekV4Pro => "deepseek-v4-pro", + Self::Glm52 => "glm-5.2", + } + } + + pub(crate) fn from_id(id: &str) -> Option { + MODEL_CHOICES.into_iter().find(|model| model.id() == id) + } + + pub(crate) fn supports_dspark(self) -> bool { + self == Self::DeepSeekV4Flash + } + + fn main_artifact(self) -> &'static Artifact { + match self { + Self::DeepSeekV4Flash => &FLASH, + Self::DeepSeekV4Pro => &PRO, + Self::Glm52 => &GLM, + } + } + + fn artifacts(self, dspark_enabled: bool) -> impl Iterator { + [ + Some(self.main_artifact()), + (self.supports_dspark() && dspark_enabled).then_some(&FLASH_DSPARK), + ] + .into_iter() + .flatten() + } + + pub(crate) fn download_size(self, dspark_enabled: bool) -> u64 { + self.artifacts(dspark_enabled) + .map(|artifact| artifact.size) + .sum() + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum DownloadPhase { + Pending(&'static str), + Downloading(&'static str), + Verifying(&'static str), + Complete, +} + +#[derive(Clone, Debug)] +pub(crate) struct DownloadProgress { + pub(crate) downloaded: u64, + pub(crate) total: u64, + pub(crate) phase: DownloadPhase, +} + +impl DownloadProgress { + pub(crate) fn remaining(&self) -> u64 { + self.total.saturating_sub(self.downloaded) + } + + pub(crate) fn fraction(&self) -> f32 { + if self.total == 0 { + 0.0 + } else { + self.downloaded as f32 / self.total as f32 + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum DownloadOutcome { + Complete, + Stopped, +} + +impl fmt::Display for ModelChoice { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(match self { + Self::DeepSeekV4Flash => "DeepSeek V4 Flash", + Self::DeepSeekV4Pro => "DeepSeek V4 Pro", + Self::Glm52 => "GLM 5.2", + }) + } +} + +struct Artifact { + label: &'static str, + file_name: &'static str, + repository: &'static str, + size: u64, + sha256: &'static str, +} + +impl Artifact { + fn path(&self, model: ModelChoice, models_path: &Path) -> PathBuf { + models_path.join(model.id()).join(self.file_name) + } + + fn url(&self) -> String { + format!( + "https://huggingface.co/{}/resolve/main/{}", + self.repository, self.file_name + ) + } + + fn verification_path(&self, model: ModelChoice, models_path: &Path) -> PathBuf { + self.path(model, models_path).with_extension("gguf.sha256") + } + + fn partial_path(&self, model: ModelChoice, models_path: &Path) -> PathBuf { + self.path(model, models_path).with_extension("gguf.part") + } + + fn is_installed(&self, model: ModelChoice, models_path: &Path) -> bool { + self.path(model, models_path) + .metadata() + .is_ok_and(|metadata| metadata.len() == self.size) + && fs::read_to_string(self.verification_path(model, models_path)) + .is_ok_and(|sha256| sha256.trim() == self.sha256) + } + + fn downloaded_bytes(&self, model: ModelChoice, models_path: &Path) -> u64 { + if self.is_installed(model, models_path) { + return self.size; + } + self.path(model, models_path) + .metadata() + .or_else(|_| self.partial_path(model, models_path).metadata()) + .map_or(0, |metadata| metadata.len().min(self.size)) + } +} + +pub(crate) fn download_progress( + model: ModelChoice, + dspark_enabled: bool, + models_path: &Path, +) -> DownloadProgress { + let mut downloaded = 0; + let mut phase = DownloadPhase::Complete; + for artifact in model.artifacts(dspark_enabled) { + let bytes = artifact.downloaded_bytes(model, models_path); + downloaded += bytes; + if phase == DownloadPhase::Complete && !artifact.is_installed(model, models_path) { + phase = if bytes >= artifact.size { + DownloadPhase::Verifying(artifact.label) + } else if bytes > 0 { + DownloadPhase::Downloading(artifact.label) + } else { + DownloadPhase::Pending(artifact.label) + }; + } + } + DownloadProgress { + downloaded, + total: model.download_size(dspark_enabled), + phase, + } +} + +pub(crate) fn download( + model: ModelChoice, + dspark_enabled: bool, + models_path: &Path, + cancel: &AtomicBool, +) -> Result { + for artifact in model.artifacts(dspark_enabled) { + if download_artifact_with_cancel(model, artifact, models_path, cancel)? + == DownloadOutcome::Stopped + { + return Ok(DownloadOutcome::Stopped); + } + } + Ok(DownloadOutcome::Complete) +} + +#[cfg(test)] +fn download_artifact( + model: ModelChoice, + artifact: &Artifact, + models_path: &Path, +) -> Result { + download_artifact_with_cancel(model, artifact, models_path, &AtomicBool::new(false)) +} + +fn download_artifact_with_cancel( + model: ModelChoice, + artifact: &Artifact, + models_path: &Path, + cancel: &AtomicBool, +) -> Result { + if cancel.load(Ordering::Relaxed) { + return Ok(DownloadOutcome::Stopped); + } + let destination = artifact.path(model, models_path); + if artifact.is_installed(model, models_path) { + return Ok(DownloadOutcome::Complete); + } + if destination.exists() { + if verify(&destination, artifact, cancel)? == DownloadOutcome::Stopped { + return Ok(DownloadOutcome::Stopped); + } + mark_verified(model, artifact, models_path)?; + return Ok(DownloadOutcome::Complete); + } + + let directory = destination + .parent() + .ok_or_else(|| "model artifact path has no parent directory".to_owned())?; + fs::create_dir_all(directory).map_err(|error| error.to_string())?; + let partial = artifact.partial_path(model, models_path); + let partial_size = partial.metadata().map_or(0, |metadata| metadata.len()); + if partial_size > artifact.size { + File::create(&partial).map_err(|error| error.to_string())?; + } + if partial.metadata().map_or(0, |metadata| metadata.len()) != artifact.size + && download_to_partial(artifact, &partial, cancel)? == DownloadOutcome::Stopped + { + return Ok(DownloadOutcome::Stopped); + } + + if cancel.load(Ordering::Relaxed) { + return Ok(DownloadOutcome::Stopped); + } + match verify(&partial, artifact, cancel) { + Ok(DownloadOutcome::Stopped) => return Ok(DownloadOutcome::Stopped), + Ok(DownloadOutcome::Complete) => {} + Err(error) => { + fs::remove_file(&partial).map_err(|remove_error| { + format!("{error}; could not remove partial file: {remove_error}") + })?; + return Err(error); + } + } + fs::rename(partial, destination).map_err(|error| error.to_string())?; + mark_verified(model, artifact, models_path)?; + Ok(DownloadOutcome::Complete) +} + +fn mark_verified( + model: ModelChoice, + artifact: &Artifact, + models_path: &Path, +) -> Result<(), String> { + fs::write( + artifact.verification_path(model, models_path), + artifact.sha256, + ) + .map_err(|error| error.to_string()) +} + +fn verify( + path: &Path, + artifact: &Artifact, + cancel: &AtomicBool, +) -> Result { + let size = path.metadata().map_err(|error| error.to_string())?.len(); + if size != artifact.size { + return Err(format!( + "{} has size {size}, expected {}", + path.display(), + artifact.size + )); + } + + let mut file = File::open(path).map_err(|error| error.to_string())?; + let mut hasher = Sha256::new(); + let mut buffer = vec![0; 1024 * 1024]; + loop { + if cancel.load(Ordering::Relaxed) { + return Ok(DownloadOutcome::Stopped); + } + let count = file.read(&mut buffer).map_err(|error| error.to_string())?; + if count == 0 { + break; + } + hasher.update(&buffer[..count]); + } + let actual = hex(&hasher.finalize()); + if actual != artifact.sha256 { + return Err(format!( + "Checksum verification failed for {}", + path.display() + )); + } + Ok(DownloadOutcome::Complete) +} + +fn hex(bytes: &[u8]) -> String { + const DIGITS: &[u8; 16] = b"0123456789abcdef"; + let mut result = String::with_capacity(bytes.len() * 2); + for byte in bytes { + result.push(DIGITS[(byte >> 4) as usize] as char); + result.push(DIGITS[(byte & 0x0f) as usize] as char); + } + result +} + +fn download_to_partial( + artifact: &Artifact, + partial: &Path, + cancel: &AtomicBool, +) -> Result { + download_url_to_partial(&artifact.url(), partial, cancel) +} + +fn download_url_to_partial( + url: &str, + partial: &Path, + cancel: &AtomicBool, +) -> Result { + let offset = partial.metadata().map_or(0, |metadata| metadata.len()); + let agent: ureq::Agent = ureq::Agent::config_builder() + .https_only(url.starts_with("https://")) + .build() + .into(); + let mut request = agent.get(url); + if offset > 0 { + request = request.header("Range", format!("bytes={offset}-")); + } + let mut response = request + .call() + .map_err(|error| format!("Model download failed: {error}"))?; + let status = response.status().as_u16(); + let append = offset > 0 && status == 206; + if offset > 0 && status != 200 && status != 206 { + return Err(format!( + "Model server returned HTTP {status} while resuming at byte {offset}" + )); + } + if append { + validate_content_range(&response, offset)?; + } + + let mut output = OpenOptions::new() + .create(true) + .write(true) + .append(append) + .truncate(!append) + .open(partial) + .map_err(|error| error.to_string())?; + let mut body = response.body_mut().as_reader(); + let mut buffer = vec![0; 1024 * 1024]; + loop { + if cancel.load(Ordering::Relaxed) { + return Ok(DownloadOutcome::Stopped); + } + let count = body + .read(&mut buffer) + .map_err(|error| format!("Model download failed: {error}"))?; + if count == 0 { + break; + } + output + .write_all(&buffer[..count]) + .map_err(|error| error.to_string())?; + } + Ok(DownloadOutcome::Complete) +} + +fn validate_content_range( + response: &ureq::http::Response, + offset: u64, +) -> Result<(), String> { + let expected = format!("bytes {offset}-"); + let content_range = response + .headers() + .get("content-range") + .and_then(|value| value.to_str().ok()) + .unwrap_or_default(); + if content_range.starts_with(&expected) { + Ok(()) + } else { + Err(format!( + "Model server returned an invalid Content-Range while resuming at byte {offset}" + )) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::net::TcpListener; + use std::thread; + use std::time::{SystemTime, UNIX_EPOCH}; + + #[test] + fn catalog_and_checksum_verification_are_explicit() { + assert_eq!(ModelChoice::from_id("glm-5.2"), Some(ModelChoice::Glm52)); + assert!(ModelChoice::from_id("unknown").is_none()); + assert_eq!( + ModelChoice::DeepSeekV4Flash.main_artifact().size, + 86_720_111_488 + ); + assert_eq!(ModelChoice::Glm52.main_artifact().size, 211_075_856_448); + assert_eq!(ModelChoice::DeepSeekV4Flash.artifacts(true).count(), 2); + assert_eq!(ModelChoice::Glm52.artifacts(true).count(), 1); + + let id = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + let models_path = std::env::temp_dir().join(format!("ds4-server-models-{id}")); + let empty = Artifact { + label: "empty", + file_name: "empty", + repository: "", + size: 0, + sha256: "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + }; + let partial = empty.partial_path(ModelChoice::DeepSeekV4Flash, &models_path); + fs::create_dir_all(partial.parent().unwrap()).unwrap(); + fs::write(&partial, []).unwrap(); + download_artifact(ModelChoice::DeepSeekV4Flash, &empty, &models_path).unwrap(); + assert!(empty.is_installed(ModelChoice::DeepSeekV4Flash, &models_path)); + assert!(!partial.exists()); + assert_eq!( + fs::read_to_string(empty.verification_path(ModelChoice::DeepSeekV4Flash, &models_path)) + .unwrap(), + empty.sha256 + ); + fs::remove_dir_all(models_path).unwrap(); + } + + #[test] + fn restart_resumes_at_the_existing_partial_byte() { + let content = b"restart-resume works"; + let offset = 8; + let directory = std::env::temp_dir().join(format!( + "ds4-server-resume-{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + fs::create_dir_all(&directory).unwrap(); + let partial = directory.join("model.gguf.part"); + fs::write(&partial, &content[..offset]).unwrap(); + + let listener = match TcpListener::bind("127.0.0.1:0") { + Ok(listener) => listener, + Err(error) if error.kind() == std::io::ErrorKind::PermissionDenied => { + fs::remove_dir_all(directory).unwrap(); + return; + } + Err(error) => panic!("could not start test server: {error}"), + }; + let address = listener.local_addr().unwrap(); + let server = thread::spawn(move || { + let (mut connection, _) = listener.accept().unwrap(); + let mut request = [0; 2048]; + let count = connection.read(&mut request).unwrap(); + let request = String::from_utf8_lossy(&request[..count]).to_ascii_lowercase(); + assert!(request.contains("range: bytes=8-")); + let remaining = &content[offset..]; + write!( + connection, + "HTTP/1.1 206 Partial Content\r\nContent-Length: {}\r\nContent-Range: bytes {offset}-{}/{}\r\nConnection: close\r\n\r\n", + remaining.len(), + content.len() - 1, + content.len(), + ) + .unwrap(); + connection.write_all(remaining).unwrap(); + }); + + let outcome = download_url_to_partial( + &format!("http://{address}/model.gguf"), + &partial, + &AtomicBool::new(false), + ) + .unwrap(); + server.join().unwrap(); + assert_eq!(outcome, DownloadOutcome::Complete); + assert_eq!(fs::read(&partial).unwrap(), content); + fs::remove_dir_all(directory).unwrap(); + } + + #[test] + fn cancellation_keeps_the_partial_file_for_the_next_run() { + let directory = std::env::temp_dir().join(format!( + "ds4-server-cancel-{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + let artifact = Artifact { + label: "test model", + file_name: "model.gguf", + repository: "unused", + size: 10, + sha256: "unused", + }; + let partial = artifact.partial_path(ModelChoice::DeepSeekV4Flash, &directory); + fs::create_dir_all(partial.parent().unwrap()).unwrap(); + fs::write(&partial, b"part").unwrap(); + let cancel = AtomicBool::new(true); + + assert_eq!( + download_artifact_with_cancel( + ModelChoice::DeepSeekV4Flash, + &artifact, + &directory, + &cancel, + ) + .unwrap(), + DownloadOutcome::Stopped + ); + assert_eq!(fs::read(&partial).unwrap(), b"part"); + fs::remove_dir_all(directory).unwrap(); + } +} diff --git a/src/schema.rs b/src/schema.rs index 7600ed2..737dbcd 100644 --- a/src/schema.rs +++ b/src/schema.rs @@ -10,7 +10,7 @@ diesel::table! { preferences (id) { id -> Integer, selected_model -> Text, - dflash_enabled -> Bool, + dspark_enabled -> Bool, idle_timeout_minutes -> Integer, } }