From d554b77b9dd51108739d92eb91679edd5b26d133 Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Fri, 24 Jul 2026 21:12:48 +0200 Subject: [PATCH] feat: first cut at openai compatible server --- Cargo.lock | 3 + Cargo.toml | 4 +- PLAN.md | 150 ++- README.md | 21 +- .../20260725000000_add_endpoint_port/down.sql | 1 + .../20260725000000_add_endpoint_port/up.sql | 2 + src/app.rs | 332 ++--- src/app/view.rs | 629 ++++++---- src/database.rs | 15 +- src/engine.rs | 314 ++++- src/main.rs | 6 + src/model.rs | 7 + src/native_edit.rs | 315 +++++ src/native_menu.rs | 76 +- src/runtime.rs | 151 +++ src/schema.rs | 1 + src/server.rs | 1116 +++++++++++++++++ src/settings.rs | 4 + 18 files changed, 2612 insertions(+), 535 deletions(-) create mode 100644 migrations/20260725000000_add_endpoint_port/down.sql create mode 100644 migrations/20260725000000_add_endpoint_port/up.sql create mode 100644 src/native_edit.rs create mode 100644 src/runtime.rs create mode 100644 src/server.rs diff --git a/Cargo.lock b/Cargo.lock index 5f4740c..a008d2a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1150,6 +1150,8 @@ dependencies = [ "muda", "png", "rfd", + "serde", + "serde_json", "sha2", "ureq", ] @@ -4037,6 +4039,7 @@ version = "1.0.151" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" dependencies = [ + "indexmap", "itoa", "memchr", "serde", diff --git a/Cargo.toml b/Cargo.toml index 6fb84ba..cd3a49c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,10 +14,12 @@ cc = "1.3.0" [dependencies] diesel = { version = "2.3.11", features = ["sqlite", "returning_clauses_for_sqlite_3_35", "64-column-tables"] } diesel_migrations = "2.3.2" -iced = { version = "0.13.1", features = ["highlighter", "markdown", "svg", "tokio"] } +iced = { version = "0.13.1", features = ["advanced", "highlighter", "markdown", "svg", "tokio"] } memmap2 = "0.9.11" png = "0.17.16" rfd = "0.15.4" +serde = { version = "1.0.228", features = ["derive"] } +serde_json = { version = "1.0.149", features = ["preserve_order"] } sha2 = "0.11.0" ureq = { version = "3.3.0", default-features = false, features = ["rustls"] } diff --git a/PLAN.md b/PLAN.md index 22d528c..d0acfd1 100644 --- a/PLAN.md +++ b/PLAN.md @@ -13,11 +13,11 @@ sessions, tools, and turn orchestration. | Area | Status | Available now | Still open | | --- | --- | --- | --- | | App shell | Implemented | Native multi-window Iced app, Codex-inspired layout, native Application/Edit/Window menus, Preferences panel, Model Manager, and streaming composer | Native app release packaging | -| Projects and sessions | Implemented with transcripts | Native folder picker, project-name dialog, create/select/delete projects and sessions, and durable structured chat history | KV payloads, rename/archive, and model binding | -| Persistence | Implemented for metadata, preferences, and chat | SQLite in Application Support, Diesel ORM, embedded migrations, constrained defaults, streamed message updates, and reopen tests | Downloaded-artifact metadata and session-runtime migrations | -| Model runtime | DeepSeek Flash vertical implemented | Rust-owned mmap/model/session graph, reused Metal kernels, 2K compressed-attention generation, DS4 sampling, lazy background loading, cancellation, and idle unloading | Ratio-4 sparse indexer beyond 2K, DSpark/SSD/steering, GLM/Pro, KV checkpoint persistence, and shared acquisition for HTTP | -| Local endpoint | Not started | Nothing listening | OpenAI-compatible HTTP and streaming surfaces | -| Agent | Basic local chat implemented | Multi-turn role rendering, structured reasoning disclosure, streamed DeepSeek Flash text, composer send/stop, and durable transcript rehydration | Tools, approvals, compaction, statistics, and KV-backed resume | +| Projects and sessions | Implemented with transcripts and KV resume | Native folder picker, project-name dialog, create/select/delete projects and sessions, durable structured chat history, and per-session KV checkpoints | Rename/archive and model binding | +| Persistence | Implemented for local chat | SQLite in Application Support, Diesel ORM, embedded migrations, constrained defaults, streamed message updates, reopen tests, context/speed state, and local KV rehydration | Downloaded-artifact metadata and session-runtime migrations; HTTP chats deliberately never enter SQLite | +| Model runtime | DeepSeek Flash vertical implemented | Rust-owned mmap/model/session graph, local Metal kernels, full configured sparse context, DS4 sampling, one process-wide UI/HTTP owner, cancellation, idle unloading, prefix continuation, and KV checkpoint save/load | DSpark/SSD/steering and GLM/Pro execution | +| Local endpoint | Chat Completions vertical implemented | Automatic configurable localhost listener (port 4000 by default), verified-on-disk model discovery, transient requests, non-streaming and SSE Chat Completions, reasoning, usage, sampling, stops, cancellation, tools, and exact in-process tool replay | GUI status and enable/disable controls, Responses, legacy Completions, Anthropic Messages, full `ds4_server.c` fixture parity, and cross-restart exact tool replay | +| Agent | Solid local chat implemented | Multi-turn role rendering, structured reasoning disclosure, Markdown, automatic scrolling, context/speed display, durable transcript rehydration, and KV-backed resume | Tools, approvals, compaction, and richer statistics | | A2UI | Future extension | bDS2 provides the reference structured render-tool contract | Native inline surfaces for local chat after core chat and tools are stable | | 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 | @@ -28,8 +28,9 @@ without blocking the UI. Local chat now lazily acquires the Rust DeepSeek Flash executor, streams generated text, supports cancellation and multi-turn role replay, persists reasoning and answers while they stream, rehydrates selected sessions, and unloads model resources after the configured idle timeout. The -current Rust graph supports DS4 compressed attention through 2K context; the sparse ratio-4 indexer, -KV checkpoints, HTTP service, and agent tools remain open. +current Rust graph supports the full configured context with ratio-4 sparse +attention and durable local KV checkpoints. The shared HTTP Chat Completions +service is available; the remaining protocol routes and agent tools remain open. ## Phase 0 — project and session shell @@ -49,16 +50,16 @@ Exit criterion: restart the app and see the same project/session tree. ## Phase 1 — model library and on-demand loading Status: **partially implemented**. The typed preference/download/intake path and -a DeepSeek Flash Rust/Metal vertical are implemented, including compressed KV -through 2K context, DS4 sampling, a single background owner, cancellation, and -idle unload. Sparse indexed attention, persistent KV/session sync, other model -families, and optional execution modes remain open. +a DeepSeek Flash Rust/Metal vertical are implemented, including full configured +context, ratio-4 sparse indexed attention, prefix continuation, KV checkpoint +save/load, DS4 sampling, cancellation, idle unload, and shared HTTP acquisition. +Other model families and optional execution modes remain open. 1. **Implemented:** the typed model/runtime preference contract described below is complete before chat, the HTTP endpoint, or the engine lifecycle. -2. **Partially implemented:** mmap-backed loading, tokenizer/prompt rendering, - Rust-owned Flash session creation/evaluation/sampling, and model ownership - are complete. Prefix sync plus KV save/load remain. +2. **Implemented for DeepSeek Flash:** mmap-backed loading, tokenizer/prompt + rendering, Rust-owned session creation/evaluation/sampling, prefix sync, and + KV save/load are complete. 3. **Implemented:** retain the tested model restrictions from DwarfStar. Main and DSpark artifacts reject invalid GGUF versions, catalog mismatches, and incompatible metadata, tensor shapes, offsets, or quantization before @@ -66,8 +67,11 @@ families, and optional execution modes remain open. 4. **Partially implemented:** reuse the existing `.metal` kernels and preserve mmap-backed resident loading plus the explicit SSD-streaming path. Keep Objective-C only at the Metal interop boundary. -5. **Implemented for local chat:** keep one engine resident at most, matching DwarfStar's instance-lock and - memory assumptions. +5. Keep exactly one model resident process-wide. Local chat and every endpoint + request must use the same owner and queue; model selection may replace the + resident model only after active work reaches a safe boundary. Multiple + simultaneously loaded models are out of scope because DS4Server targets + models large enough to consume most available memory. 6. **Implemented for local chat:** move engine work off the UI thread and support cooperative cancellation at the safe session boundaries already defined by `ds4_session_sync`. @@ -92,21 +96,25 @@ to Rust while preserving its semantics and on-disk compatibility: 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. +- Use one process-wide Rust model owner and KV checkpoint implementation for + local chat and the HTTP server. Model access is serialized so two callers + never mutate the single resident inference graph concurrently. +- Application projects and sessions belong exclusively to local chat. HTTP + requests are transient and must never create projects, sessions, messages, + transcript rows, request-identity mappings, or any other durable chat record. + The outside client owns conversation identity and resends the history needed + for each request, exactly as `ds4_server.c` expects. +- Transient HTTP requests may read and write engine-owned KV cache files for + compatible-prefix reuse. Those opaque cache files are the only server-side + persistence permitted for external conversations and do not establish an + application session or recoverable transcript. +- Store engine-owned checkpoints under Application Support in `kv-cache/`. + Local chats use their application session ID; HTTP uses content-addressed + entries under `kv-cache/http/`. The large opaque payload is never stored in + SQLite, and HTTP cache entries do not create application sessions. +- Cache sharing does not merge conversations or expose transcript text. It + shares only the DS4 implementation and compatible opaque prefix checkpoints; + the request body remains the source of truth for external conversation 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 @@ -140,6 +148,8 @@ to Rust while preserving its semantics and on-disk compatibility: - 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. +- Let the user configure the localhost endpoint port; default to 4000 and apply + a changed port when preferences are saved. - Persist preferences in the application SQLite database. Changing preferences never loads a model by itself. @@ -258,43 +268,58 @@ engine lifecycle. ## Phase 2 — OpenAI-compatible local endpoint -Status: **open; not started**. +Status: **partially implemented**. The shared single-model runtime, +verified-model discovery, configurable localhost port, and the first +`ds4_server.c`-compatible Chat Completions vertical are implemented. The +remaining protocol routes plus GUI status and enable/disable controls remain +open. -1. Add a localhost-only HTTP service whose lifecycle is independent of the +1. **Partially implemented:** add a localhost-only HTTP service whose lifecycle is independent of the engine. It can listen while the model is unloaded and acquires the configured - model only for inference. Start and stop it from the GUI; never expose a LAN - listener without an explicit setting. -2. Implement the DwarfStar compatibility surface in this order: + model only for inference. It listens on configurable port 4000 by default; + start/stop controls remain open. Never expose a LAN listener without an + explicit setting. +2. Implement the DwarfStar compatibility surface from `ds4_server.c` in this + order, preserving its request parsing, prompt formatting, reasoning fields, + streaming events, errors, usage accounting, stop behavior, tool replay, and + cancellation semantics rather than inventing an app-specific protocol: - `GET /v1/models` - `POST /v1/chat/completions` - `POST /v1/responses` - `POST /v1/completions` - `GET /v1/models` reports the configured model without forcing it to load. + `GET /v1/models` scans managed artifacts without forcing a model to load and + reports only main model files that are fully downloaded and have a matching + on-disk verification marker. Missing, partial, unverified, or checksum-invalid + artifacts are never advertised. 3. Support streaming SSE, reasoning output, usage accounting, cancellation, sampling parameters, tool schemas, and tool choice. 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. 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. +5. Route API work through the single process-wide model owner and engine-owned + KV store defined in Phase 1. External requests are always stateless at the + application layer: never write their messages to SQLite and never bind a + Responses `conversation`, request identifier, or client-supplied session key + to a local project/session. The client must replay its protocol history; + DS4-compatible prefix lookup may accelerate that replay using opaque KV cache + files. Add resident batching only if one-model serialized execution is later + proven insufficient. 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. 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. +multi-turn tool call. Verified installed models are listed without loading one; +missing/unverified models are absent; local and HTTP generations cannot load two +models concurrently; external calls leave no project, session, message, or +transcript rows behind; and compatible KV cache files remain reusable. ## Phase 3 — durable agent sessions -Status: **partially implemented**. Project/session metadata and structured -transcripts are durable; model binding and KV payload persistence remain open. +Status: **partially implemented for local chat only**. Project/session metadata, +structured transcripts, context statistics, and per-session KV checkpoints are +durable. Model binding and archive behavior remain open. This phase never +applies to external endpoint conversations. Extend each metadata-only session with transcript and shared-checkpoint metadata: @@ -306,14 +331,14 @@ Application Support/DS4Server.rfc1437.de/ main.gguf dspark.gguf # only when enabled kv-cache/ - .kv + .bin + http/.bin ``` -- **Partially implemented:** migrated message tables persist user text, - reasoning state, and assistant output while generation streams. Add the - 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. +- **Implemented:** migrated message tables persist local user text, reasoning + state, assistant output, context usage, and generation speed while generation + streams. Local session KV payloads remain engine-owned files written with + atomic replacement; SQLite never stores opaque tensor data. - Record model identity, context size, rendered token history, title, timestamps, and working directory. - Restore compatible KV immediately. If KV is absent or incompatible, rebuild @@ -365,6 +390,11 @@ ephemeral presentation preferences are the only exception. ### Native macOS menus (required with chat) +Status: **partially implemented**. Application, Edit, and Window menus are +installed; Cut, Copy, Paste, and Select All now bridge into the focused Iced +text field. File/View/Help menus, dynamic enabled state, the remaining Edit +actions, and selectable transcript text remain open. + The application must install and maintain a complete native macOS menu bar; the application-name menu alone is not sufficient. Use platform menu roles, ordering, selectors, and standard key equivalents rather than drawing menus @@ -493,11 +523,13 @@ tests, and release packaging remain open. 1. Finish the Flash session core as one larger slice: port the ratio-4 indexer and indexed attention for full configured context, then add prefix sync and - compatible KV checkpoint save/load. + compatible KV checkpoint save/load. **Implemented.** 2. Put the OpenAI-compatible endpoint on the same engine lifecycle so local - chats and HTTP requests cannot load duplicate engines. -3. Add transcript/KV persistence, then connect the existing conversation UI and - finally port the agent tools. + chats and HTTP requests cannot load duplicate models. Keep HTTP conversation + state transient except for opaque KV cache files. **Chat Completions vertical + implemented; remaining reference routes are open.** +3. Local transcript/KV persistence and the conversation UI are implemented; + port the agent tools after the shared endpoint lifecycle is stable. 4. Add the opt-in Dev Brain tool after local file/search boundaries and agent approvals are stable. 5. Add bDS2-style A2UI render tools and native inline surfaces after the core diff --git a/README.md b/README.md index 9017838..8fea35a 100644 --- a/README.md +++ b/README.md @@ -4,11 +4,15 @@ DS4Server is a native macOS coding-agent application that rewrites the DwarfStar (`ds4`) inference engine in Rust. It uses Rust and Iced and will combine local model loading, an OpenAI-compatible localhost endpoint, and project-scoped agent chat in one app. +DS4Server vendors and adapts the Metal kernels and Objective-C Metal glue from +[DwarfStar (`ds4`)](https://github.com/antirez/ds4). Their copyright and license +notices are retained in [`native/metal/LICENSE`](native/metal/LICENSE). + 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 configure model, generation, -runtime, and idle-unload settings. The separate Model Manager +runtime, local endpoint, and idle-unload settings. The separate Model Manager (`Shift-Command-M`) lists local main and DSpark artifacts, their on-disk sizes and state, and lets you download, resume, validate, or delete them. Rust-native background work @@ -21,9 +25,17 @@ The selected DeepSeek V4 Flash model can run directly from a project session. The model, KV/compressor state, 43-layer graph, sampling, and lifecycle are owned by Rust; a fixed snapshot of the Objective-C Metal boundary and unchanged Metal kernels is vendored and built inside this repository. Tokens stream into the chat UI, Stop cancels generation, -follow-up turns replay their role history, and the model unloads after the idle -timeout. The current graph supports 2K context while the sparse long-context -indexer and persistent transcript/KV restore remain the next engine slice. +follow-up turns reuse durable transcript and KV state, and the model unloads +after the idle timeout. The graph uses the full configured context with the +ratio-4 sparse indexer. + +The app also listens on `127.0.0.1:4000` by default for `GET /v1/models` and +`POST /v1/chat/completions`; the port is configurable in Preferences. The +endpoint and local chat share the single model owner. External conversations +are client-managed and never enter the project, session, message, or transcript +database; only opaque content-addressed KV cache files are retained. Model +discovery advertises only supported main artifacts that are fully downloaded +and verified on disk. ```sh cargo install cargo-packager --locked --version 0.11.8 @@ -35,6 +47,7 @@ State is stored at: ```text ~/Library/Application Support/DS4Server.rfc1437.de/data.sqlite3 +~/Library/Application Support/DS4Server.rfc1437.de/kv-cache/ ``` Deleting a project or session removes only DS4Server metadata. It never deletes diff --git a/migrations/20260725000000_add_endpoint_port/down.sql b/migrations/20260725000000_add_endpoint_port/down.sql new file mode 100644 index 0000000..175dc6a --- /dev/null +++ b/migrations/20260725000000_add_endpoint_port/down.sql @@ -0,0 +1 @@ +ALTER TABLE preferences DROP COLUMN endpoint_port; diff --git a/migrations/20260725000000_add_endpoint_port/up.sql b/migrations/20260725000000_add_endpoint_port/up.sql new file mode 100644 index 0000000..0e840b8 --- /dev/null +++ b/migrations/20260725000000_add_endpoint_port/up.sql @@ -0,0 +1,2 @@ +ALTER TABLE preferences ADD COLUMN endpoint_port INTEGER NOT NULL DEFAULT 4000 + CHECK (endpoint_port BETWEEN 1 AND 65535); diff --git a/src/app.rs b/src/app.rs index 0da67d4..349907e 100644 --- a/src/app.rs +++ b/src/app.rs @@ -4,8 +4,10 @@ pub(crate) use view::app_theme; use crate::database::{AppPreferences, Database, ProjectWithSessions, StoredMessage}; #[cfg(target_os = "macos")] -use crate::engine::{ChatTurn, Generator}; +use crate::engine::ChatTurn; use crate::model::{self, DownloadOutcome, DownloadProgress, ManagedArtifactId, ModelChoice}; +#[cfg(target_os = "macos")] +use crate::runtime::{ActiveGeneration, CheckpointTarget, GenerationEvent, GenerationService}; use crate::settings::{ DiagnosticPreferences, ExecutionPreferences, GIB, GenerationPreferences, ReasoningMode, RuntimePreferences, SpeculativePreferences, SsdPreferences, SteeringPreferences, @@ -16,9 +18,9 @@ use iced::{Size, Subscription, Task, keyboard, window}; use rfd::AsyncFileDialog; use std::fs; use std::path::PathBuf; -use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::mpsc::{self, TryRecvError}; +use std::sync::{Arc, RwLock}; use std::thread; use std::time::{Duration, Instant}; @@ -29,6 +31,7 @@ struct PreferenceDraft { model: ModelChoice, dspark_enabled: bool, idle_timeout_minutes: String, + endpoint_port: String, context_tokens: String, max_generated_tokens: String, system_prompt: String, @@ -73,6 +76,7 @@ impl PreferenceDraft { model, dspark_enabled: speculative.dspark_enabled, idle_timeout_minutes: preferences.idle_timeout_minutes.to_string(), + endpoint_port: preferences.endpoint_port.to_string(), context_tokens: generation.context_tokens.to_string(), max_generated_tokens: generation.max_generated_tokens.to_string(), system_prompt: generation.system_prompt, @@ -168,6 +172,7 @@ impl PreferenceDraft { self.model = ModelChoice::default(); self.dspark_enabled = false; self.idle_timeout_minutes = "10".into(); + self.endpoint_port = "4000".into(); self.context_tokens = defaults.context_tokens.to_string(); self.max_generated_tokens = defaults.max_generated_tokens.to_string(); self.system_prompt = defaults.system_prompt; @@ -389,6 +394,8 @@ pub(crate) struct App { pub(super) pending_model_delete: Option, #[cfg(target_os = "macos")] _native_menu: Option, + #[cfg(target_os = "macos")] + pub(super) native_edit_commands: crate::native_edit::EditCommandQueue, database: Option, projects: Vec, preferences: AppPreferences, @@ -408,7 +415,13 @@ pub(crate) struct App { pub(super) context_limit: u32, pub(super) tokens_per_second: Option, #[cfg(target_os = "macos")] - generation_worker: Option, + generation_service: Option, + #[cfg(target_os = "macos")] + active_generation: Option, + #[cfg(target_os = "macos")] + runtime_preferences: Arc>, + #[cfg(target_os = "macos")] + _endpoint: Option, error: Option, } @@ -461,40 +474,6 @@ impl From for ChatMessage { } } -#[cfg(target_os = "macos")] -struct GenerationWorker { - commands: mpsc::Sender, - events: mpsc::Receiver, - cancel: Option>, -} - -#[cfg(target_os = "macos")] -enum GenerationCommand { - Generate { - engine: crate::settings::EngineSettings, - turn: crate::settings::TurnSettings, - messages: Vec, - checkpoint: PathBuf, - idle_timeout: Duration, - cancel: Arc, - }, -} - -#[cfg(target_os = "macos")] -enum GenerationEvent { - Loading, - Chunk { - reasoning: bool, - content: String, - }, - Context { - used: u32, - limit: u32, - tokens_per_second: Option, - }, - Finished(Result<(), String>), -} - #[derive(Debug)] pub(super) enum ModelDownload { Idle, @@ -526,6 +505,8 @@ pub(super) struct ActiveDownload { #[derive(Debug, Clone)] pub(crate) enum Message { Noop, + #[cfg(target_os = "macos")] + NativeEdit(crate::native_edit::EditCommand), OpenPreferences, OpenModelManager, ModelManagerOpened(window::Id), @@ -535,6 +516,7 @@ pub(crate) enum Message { PreferenceModelChanged(ModelChoice), PreferenceDsparkChanged(bool), PreferenceTimeoutChanged(String), + PreferenceEndpointPortChanged(String), PreferenceContextChanged(String), PreferenceMaxTokensChanged(String), PreferenceSystemPromptChanged(String), @@ -602,12 +584,17 @@ impl App { Err(error) => return Self::failed(error, main_window), }; let context_limit = preferences.context_tokens.max(0) as u32; + #[cfg(target_os = "macos")] + let (runtime_preferences, generation_service, endpoint, service_error) = + spawn_services(&preferences); Self { main_window, model_manager_window: None, pending_model_delete: None, #[cfg(target_os = "macos")] _native_menu: None, + #[cfg(target_os = "macos")] + native_edit_commands: crate::native_edit::command_queue(), database: Some(database), projects, preferences, @@ -627,8 +614,23 @@ impl App { context_limit, tokens_per_second: None, #[cfg(target_os = "macos")] - generation_worker: None, - error: None, + generation_service, + #[cfg(target_os = "macos")] + active_generation: None, + #[cfg(target_os = "macos")] + runtime_preferences, + #[cfg(target_os = "macos")] + _endpoint: endpoint, + error: { + #[cfg(target_os = "macos")] + { + service_error + } + #[cfg(not(target_os = "macos"))] + { + None + } + }, } } (Err(error), _) | (_, Err(error)) => Self::failed(error, main_window), @@ -642,12 +644,21 @@ impl App { let preference_draft = PreferenceDraft::from_saved(&preferences) .expect("default preferences must use a supported model"); let context_limit = preferences.context_tokens.max(0) as u32; + #[cfg(target_os = "macos")] + let (runtime_preferences, generation_service, endpoint, service_error) = + spawn_services(&preferences); + #[cfg(target_os = "macos")] + let startup_error = service_error; + #[cfg(not(target_os = "macos"))] + let startup_error = None::; Self { main_window, model_manager_window: None, pending_model_delete: None, #[cfg(target_os = "macos")] _native_menu: None, + #[cfg(target_os = "macos")] + native_edit_commands: crate::native_edit::command_queue(), database: None, projects: Vec::new(), preferences, @@ -667,14 +678,29 @@ impl App { context_limit, tokens_per_second: None, #[cfg(target_os = "macos")] - generation_worker: None, - error: Some(format!("Could not open the project database: {error}")), + generation_service, + #[cfg(target_os = "macos")] + active_generation: None, + #[cfg(target_os = "macos")] + runtime_preferences, + #[cfg(target_os = "macos")] + _endpoint: endpoint, + error: Some(match startup_error { + Some(service_error) => { + format!("Could not open the project database: {error}. {service_error}") + } + None => format!("Could not open the project database: {error}"), + }), } } pub(crate) fn update(&mut self, message: Message) -> Task { match message { Message::Noop => {} + #[cfg(target_os = "macos")] + Message::NativeEdit(command) => { + crate::native_edit::queue_command(&self.native_edit_commands, command) + } Message::OpenPreferences => self.open_preferences(), Message::OpenModelManager => return self.open_model_manager(), Message::ModelManagerOpened(id) => { @@ -749,6 +775,10 @@ impl App { self.preference_draft.idle_timeout_minutes = value; self.preference_error = None; } + Message::PreferenceEndpointPortChanged(value) => { + self.preference_draft.endpoint_port = value; + self.preference_error = None; + } Message::PreferenceContextChanged(value) => { self.preference_draft.context_tokens = value; self.preference_error = None; @@ -953,14 +983,11 @@ impl App { self.start_generation(); return scroll_chat_to_end(); } - Message::StopGeneration => { + Message::StopGeneration => + { #[cfg(target_os = "macos")] - if let Some(cancel) = self - .generation_worker - .as_ref() - .and_then(|worker| worker.cancel.as_ref()) - { - cancel.store(true, Ordering::Relaxed); + if let Some(active) = &self.active_generation { + active.cancel.store(true, Ordering::Relaxed); } } Message::GenerationTick => { @@ -1133,6 +1160,9 @@ impl App { Some(crate::native_menu::NativeMenuEvent::ModelManager) => { Message::OpenModelManager } + Some(crate::native_menu::NativeMenuEvent::Edit(command)) => { + Message::NativeEdit(command) + } None => Message::Noop, } })); @@ -1255,6 +1285,14 @@ impl App { self.preference_error = Some("Idle timeout must be between 1 and 1440 minutes.".into()); return; } + let Ok(endpoint_port) = self.preference_draft.endpoint_port.trim().parse::() else { + self.preference_error = Some("Endpoint port must be a whole number.".into()); + return; + }; + if endpoint_port == 0 { + self.preference_error = Some("Endpoint port must be between 1 and 65535.".into()); + return; + } let generation = match self.preference_draft.generation() { Ok(generation) => generation, Err(error) => { @@ -1275,12 +1313,50 @@ impl App { self.preference_error = Some(error); return; } + #[cfg(target_os = "macos")] + let pending_endpoint = if self.preferences.endpoint_port != i32::from(endpoint_port) + || self._endpoint.is_none() + { + let Some(generation) = &self.generation_service else { + self.preference_error = Some("The model runtime is unavailable.".into()); + return; + }; + match crate::server::ServerHandle::spawn( + generation.clone(), + Arc::clone(&self.runtime_preferences), + models_path(), + application_support_path().join("kv-cache").join("http"), + endpoint_port, + ) { + Ok(endpoint) => Some(endpoint), + Err(error) => { + self.preference_error = Some(error); + return; + } + } + } else { + None + }; let Some(database) = &mut self.database else { return; }; - match database.update_preferences(model.id(), idle_timeout_minutes, &generation, &runtime) { + match database.update_preferences( + model.id(), + idle_timeout_minutes, + i32::from(endpoint_port), + &generation, + &runtime, + ) { Ok(preferences) => { self.preferences = preferences; + #[cfg(target_os = "macos")] + if let Ok(mut runtime_preferences) = self.runtime_preferences.write() { + *runtime_preferences = self.preferences.clone(); + } + #[cfg(target_os = "macos")] + if let Some(endpoint) = pending_endpoint { + self._endpoint = Some(endpoint); + } self.preference_draft = PreferenceDraft::from_saved(&self.preferences) .expect("the saved model was selected from the supported catalog"); self.preferences_open = false; @@ -1508,15 +1584,10 @@ impl App { #[cfg(target_os = "macos")] { - if self.generation_worker.is_none() { - match spawn_generation_worker() { - Ok(worker) => self.generation_worker = Some(worker), - Err(error) => { - self.error = Some(error); - return; - } - } - } + let Some(service) = &self.generation_service else { + self.error = Some("The model runtime is unavailable.".into()); + return; + }; let Some(database) = &mut self.database else { return; }; @@ -1527,24 +1598,22 @@ impl App { return; } }; - let cancel = Arc::new(AtomicBool::new(false)); let idle_timeout = Duration::from_secs(self.preferences.idle_timeout_minutes.max(1) as u64 * 60); - let command = GenerationCommand::Generate { - engine: effective.engine, - turn: effective.turn, + self.active_generation = match service.generate( + effective.engine, + effective.turn, messages, - checkpoint: session_checkpoint_path(session_id), + CheckpointTarget::Local(session_checkpoint_path(session_id)), idle_timeout, - cancel: Arc::clone(&cancel), + ) { + Ok(active) => Some(active), + Err(error) => { + self.generation_service = None; + self.error = Some(error); + return; + } }; - let worker = self.generation_worker.as_mut().expect("worker was created"); - if worker.commands.send(command).is_err() { - self.generation_worker = None; - self.error = Some("The local generation worker stopped unexpectedly.".into()); - return; - } - worker.cancel = Some(cancel); let user = ChatMessage::from(saved.0); let mut assistant = ChatMessage::from(saved.1); assistant.reasoning_open = assistant_reasoning; @@ -1565,7 +1634,7 @@ impl App { fn poll_generation(&mut self) -> bool { #[cfg(target_os = "macos")] - let Some(worker) = &mut self.generation_worker else { + let Some(active) = &mut self.active_generation else { self.generating = false; return false; }; @@ -1575,7 +1644,7 @@ impl App { let mut context_changed = false; #[cfg(target_os = "macos")] loop { - match worker.events.try_recv() { + match active.events.try_recv() { Ok(GenerationEvent::Loading) => {} Ok(GenerationEvent::Chunk { reasoning, content }) => { if let Some(message) = self.conversation.last_mut() @@ -1597,17 +1666,17 @@ impl App { } Ok(GenerationEvent::Finished(result)) => { self.generating = false; - worker.cancel = None; if let Err(error) = result { self.error = Some(error); } + self.active_generation = None; break; } Err(TryRecvError::Empty) => break, Err(TryRecvError::Disconnected) => { self.generating = false; - self.generation_worker = None; - self.error = Some("The local generation worker stopped unexpectedly.".into()); + self.active_generation = None; + self.error = Some("The model runtime stopped unexpectedly.".into()); break; } } @@ -1627,12 +1696,8 @@ impl App { &message.content, ) { - if let Some(cancel) = self - .generation_worker - .as_ref() - .and_then(|worker| worker.cancel.as_ref()) - { - cancel.store(true, Ordering::Relaxed); + if let Some(active) = &self.active_generation { + active.cancel.store(true, Ordering::Relaxed); } self.error = Some(format!("Could not save generated chat text: {error}")); } @@ -1667,77 +1732,30 @@ impl App { } #[cfg(target_os = "macos")] -fn spawn_generation_worker() -> Result { - let (command_sender, command_receiver) = mpsc::channel(); - let (event_sender, event_receiver) = mpsc::channel(); - thread::Builder::new() - .name("local-generation".into()) - .spawn(move || { - let mut loaded = None::<(crate::settings::EngineSettings, Generator)>; - let mut last_used = Instant::now(); - let mut idle_timeout = Duration::from_secs(15 * 60); - loop { - match command_receiver.recv_timeout(Duration::from_secs(1)) { - Ok(GenerationCommand::Generate { - engine, - turn, - messages, - checkpoint, - idle_timeout: requested_timeout, - cancel, - }) => { - idle_timeout = requested_timeout; - if loaded - .as_ref() - .is_none_or(|(current, _)| current != &engine) - { - let _ = event_sender.send(GenerationEvent::Loading); - loaded = match Generator::open(&engine) { - Ok(generator) => Some((engine.clone(), generator)), - Err(error) => { - let _ = - event_sender.send(GenerationEvent::Finished(Err(error))); - None - } - }; - } - if let Some((_, generator)) = &mut loaded { - let result = generator.generate( - &checkpoint, - &messages, - &turn, - &cancel, - |reasoning, content| { - let _ = event_sender - .send(GenerationEvent::Chunk { reasoning, content }); - }, - |used, limit, tokens_per_second| { - let _ = event_sender.send(GenerationEvent::Context { - used, - limit, - tokens_per_second, - }); - }, - ); - let _ = event_sender.send(GenerationEvent::Finished(result)); - last_used = Instant::now(); - } - } - Err(mpsc::RecvTimeoutError::Timeout) => { - if loaded.is_some() && last_used.elapsed() >= idle_timeout { - loaded = None; - } - } - Err(mpsc::RecvTimeoutError::Disconnected) => break, - } - } - }) - .map_err(|error| format!("Could not start local generation: {error}"))?; - Ok(GenerationWorker { - commands: command_sender, - events: event_receiver, - cancel: None, - }) +fn spawn_services( + preferences: &AppPreferences, +) -> ( + Arc>, + Option, + Option, + Option, +) { + let runtime_preferences = Arc::new(RwLock::new(preferences.clone())); + let generation = match GenerationService::spawn() { + Ok(generation) => generation, + Err(error) => return (runtime_preferences, None, None, Some(error)), + }; + let endpoint = crate::server::ServerHandle::spawn( + generation.clone(), + Arc::clone(&runtime_preferences), + models_path(), + application_support_path().join("kv-cache").join("http"), + u16::try_from(preferences.endpoint_port).unwrap_or(4000), + ); + match endpoint { + Ok(endpoint) => (runtime_preferences, Some(generation), Some(endpoint), None), + Err(error) => (runtime_preferences, Some(generation), None, Some(error)), + } } impl Drop for App { @@ -1746,12 +1764,8 @@ impl Drop for App { download.cancel.store(true, Ordering::Relaxed); } #[cfg(target_os = "macos")] - if let Some(cancel) = self - .generation_worker - .as_ref() - .and_then(|worker| worker.cancel.as_ref()) - { - cancel.store(true, Ordering::Relaxed); + if let Some(active) = &self.active_generation { + active.cancel.store(true, Ordering::Relaxed); } } } diff --git a/src/app/view.rs b/src/app/view.rs index ec5074e..a3bea63 100644 --- a/src/app/view.rs +++ b/src/app/view.rs @@ -31,7 +31,12 @@ impl App { if self.model_manager_window == Some(id) { self.model_manager() } else { - self.main_view() + let content = self.main_view(); + #[cfg(target_os = "macos")] + return crate::native_edit::native_edit(content, self.native_edit_commands.clone()) + .into(); + #[cfg(not(target_os = "macos"))] + content } } @@ -457,286 +462,331 @@ impl App { dspark_confidence.on_input(Message::PreferenceDsparkConfidenceChanged); } - let mut fields = column![ - text("MODEL").size(11), - pick_list( - &MODEL_CHOICES[..], - Some(self.preference_draft.model), - Message::PreferenceModelChanged, - ) - .width(Length::Fill), - text(format!( - "Main: {}{}", - engine.map_or_else( - || "Invalid settings".to_owned(), - |engine| engine.artifacts.model.display().to_string(), - ), - engine - .and_then(|engine| engine.artifacts.mtp.as_ref()) - .map_or_else(String::new, |path| format!(" • support: {}", path.display())), - )) - .size(12), - Space::with_height(8), - text("EXECUTION").size(11), - preference_input_row( - "CPU helper threads", - text_input("Automatic", &self.preference_draft.cpu_threads) - .on_input(Message::PreferenceCpuThreadsChanged), - ), - preference_input_row("GPU power percent", power), - preference_input_row("Prefill chunk", prefill), - checkbox("Prefer exact quality kernels", self.preference_draft.quality) - .on_toggle(Message::PreferenceQualityChanged), - checkbox("Warm mapped weights at load time", self.preference_draft.warm_weights) - .on_toggle(Message::PreferenceWarmWeightsChanged), - text(if self.preference_draft.model == ModelChoice::Glm52 { - "GLM 5.2 uses full GPU power and selects prefill chunks automatically." - } else { - "Blank numeric values preserve DS4's automatic engine behavior." - }) - .size(12), - text(engine.as_ref().map_or_else( - || "Effective execution settings will appear after valid values are entered." - .to_owned(), - |engine| { - let settings = engine.execution; - format!( - "Metal engine: threads {} • power {}% • prefill {} • quality {} • warm weights {}", - if settings.cpu_threads == 0 { "auto".to_owned() } else { settings.cpu_threads.to_string() }, - if settings.power_percent == 0 { 100 } else { settings.power_percent }, - if settings.prefill_chunk == 0 { "auto".to_owned() } else { settings.prefill_chunk.to_string() }, - if settings.quality { "on" } else { "off" }, - if settings.warm_weights { "on" } else { "off" }, - )}, - )) - .size(12), - Space::with_height(8), - text("SPECULATIVE DECODING").size(11), - preference_input_row( - "MTP draft tokens", - text_input("1", &self.preference_draft.mtp_draft_tokens) - .on_input(Message::PreferenceMtpDraftChanged), - ), - preference_input_row( - "MTP verifier margin", - text_input("3", &self.preference_draft.mtp_margin) - .on_input(Message::PreferenceMtpMarginChanged), - ), - checkbox("Enable integrated GLM MTP", self.preference_draft.glm_mtp) - .on_toggle_maybe(glm_mtp_toggle), - checkbox( - "Log GLM MTP timing counters", - self.preference_draft.glm_mtp_timing, - ) - .on_toggle_maybe(glm_mtp_timing_toggle), - dspark, - preference_input_row("DSpark confidence threshold", dspark_confidence), - checkbox( - "DSpark target-only decode", - self.preference_draft.dspark_strict, - ) - .on_toggle_maybe(dspark_strict_toggle), - text(if self.preference_draft.model.supports_dspark() { - "DSpark uses the managed support artifact; entering a threshold or enabling strict mode also enables DSpark." - } else if self.preference_draft.model == ModelChoice::Glm52 { - "GLM MTP is integrated; DSpark is unavailable for this model." - } else { - "No managed MTP support artifact is available for this model." - }) - .size(12), - text(engine.as_ref().map_or_else( - || "Effective speculative settings will appear after valid values are entered." - .to_owned(), - |engine| { - let settings = engine.speculative; - format!( - "Engine: MTP draft {} • margin {} • GLM MTP {} • timing {} • DSpark {} • confidence {}{} • target-only {}", - settings.mtp_draft_tokens, - settings.mtp_margin, - if settings.glm_mtp { "on" } else { "off" }, - if settings.glm_mtp_timing { "on" } else { "off" }, - if settings.dspark { "on" } else { "off" }, - settings.dspark_confidence_threshold, - if settings.dspark_confidence_threshold_set { " explicit" } else { " default" }, - if settings.dspark_strict { "on" } else { "off" }, - )}, - )) - .size(12), - Space::with_height(8), - text("SSD STREAMING").size(11), - checkbox("Enable SSD-backed model streaming", self.preference_draft.ssd_streaming) - .on_toggle(Message::PreferenceSsdChanged), - checkbox("Skip automatic expert preload", self.preference_draft.ssd_streaming_cold) - .on_toggle(Message::PreferenceSsdColdChanged), - preference_input_row( - "Expert cache count or GiB", - text_input("Automatic, 128, or 64GB", &self.preference_draft.ssd_cache) - .on_input(Message::PreferenceSsdCacheChanged), - ), - preference_input_row("Fully resident GLM layers", ssd_full_layers), - preference_input_row( - "Explicit expert preload count", - text_input("Automatic", &self.preference_draft.ssd_preload_experts) - .on_input(Message::PreferenceSsdPreloadChanged), - ), - text("A blank full-layer value is automatic; an explicit 0 disables fully resident GLM layers. SSD streaming and DSpark are mutually exclusive.") - .size(12), - text(engine.as_ref().map_or_else( - || "Effective SSD settings will appear after valid values are entered." - .to_owned(), - |engine| { - let settings = engine.ssd; - let cache = if settings.cache_bytes > 0 { - format!("{} GiB", settings.cache_bytes / GIB) - } else if settings.cache_experts > 0 { - format!("{} experts", settings.cache_experts) - } else { - "auto".to_owned() - }; - format!( - "Engine: streaming {} • cold {} • cache {} • full layers {}{} • preload {}", - if settings.enabled { "on" } else { "off" }, - if settings.cold { "on" } else { "off" }, - cache, - settings.full_layers, - if settings.full_layers_set { " explicit" } else { " auto" }, - if settings.preload_experts == 0 { "auto".to_owned() } else { settings.preload_experts.to_string() }, - ) - }, - )) - .size(12), - Space::with_height(8), - text("DIRECTIONAL STEERING").size(11), - text("Direction-vector file").size(13), - steering_file.padding(9), - preference_input_row("FFN scale", steering_ffn), - preference_input_row("Attention scale", steering_attn), - text(if self.preference_draft.model == ModelChoice::Glm52 { - "Directional steering is not supported for GLM 5.2." - } else { - "With a file and no explicit scale, DS4 defaults the FFN scale to 1. Scales accept -100 through 100." - }) - .size(12), - text(engine.as_ref().map_or_else( - || "Effective steering settings will appear after valid values are entered." - .to_owned(), - |engine| format!( - "Engine: file {} • FFN scale {} • attention scale {}", - if engine.steering.file.is_some() { "set" } else { "off" }, - engine.steering.ffn_scale, - engine.steering.attention_scale, - ), - )) - .size(12), - Space::with_height(8), - text("CAPACITY").size(11), - preference_input_row( - "Context tokens", - text_input("32768", &self.preference_draft.context_tokens) - .on_input(Message::PreferenceContextChanged), - ), - preference_input_row( - "Maximum generated tokens", - text_input("50000", &self.preference_draft.max_generated_tokens) - .on_input(Message::PreferenceMaxTokensChanged), - ), - text("System prompt").size(13), - text_input( - "You are a helpful assistant", - &self.preference_draft.system_prompt, - ) - .on_input(Message::PreferenceSystemPromptChanged) - .padding(9), - Space::with_height(8), - text("SAMPLING AND REASONING").size(11), - preference_input_row( - "Temperature", - text_input("DS4 default", &self.preference_draft.temperature) - .on_input(Message::PreferenceTemperatureChanged), - ), - preference_input_row( - "Top-p", - text_input("DS4 default", &self.preference_draft.top_p) - .on_input(Message::PreferenceTopPChanged), - ), - preference_input_row( - "Min-p", - text_input("DS4 default", &self.preference_draft.min_p) - .on_input(Message::PreferenceMinPChanged), - ), - preference_input_row( - "Seed", - text_input("Random", &self.preference_draft.seed) - .on_input(Message::PreferenceSeedChanged), - ), - row![ - text("Reasoning").size(13).width(Length::Fill), + let model_group = preference_group( + "MODEL & LIFECYCLE", + column![ pick_list( - &REASONING_MODES[..], - Some(self.preference_draft.reasoning_mode), - Message::PreferenceReasoningChanged, + &MODEL_CHOICES[..], + Some(self.preference_draft.model), + Message::PreferenceModelChanged, ) - .width(240), - ] - .spacing(12) - .align_y(Alignment::Center), - text("Blank sampling values retain DS4's model-family defaults. Think Max needs at least 393216 context tokens.") + .width(Length::Fill), + text(format!( + "Main: {}{}", + engine.map_or_else( + || "Invalid settings".to_owned(), + |engine| engine.artifacts.model.display().to_string(), + ), + engine + .and_then(|engine| engine.artifacts.mtp.as_ref()) + .map_or_else(String::new, |path| format!( + " • support: {}", + path.display() + )), + )) .size(12), - text(turn.map_or_else( - || "Effective settings will appear after valid values are entered.".to_owned(), - |settings| format!( - "Effective: {} context • {} max • temp {} • top-p {} • min-p {} • seed {} • {} • system prompt {}", - settings.context_tokens, - settings.max_generated_tokens, - settings.temperature, - settings.top_p, - settings.min_p, - settings.seed.map_or_else(|| "random".to_owned(), |seed| seed.to_string()), - settings.reasoning_mode, - if settings.system_prompt.is_empty() { "off" } else { "on" }, - ), - )) - .size(12), - Space::with_height(8), - text("ADVANCED DIAGNOSTICS").size(11), - preference_input_row( - "Simulated used memory (GiB)", - text_input("Disabled", &self.preference_draft.simulated_used_memory_gib) - .on_input(Message::PreferenceSimulatedMemoryChanged), - ), - text("Routed expert profile output").size(13), - text_input("Output file path", &self.preference_draft.expert_profile_path) - .on_input(Message::PreferenceExpertProfileChanged) - .padding(9), - text(engine.as_ref().map_or_else( - || "Effective diagnostic settings will appear after valid values are entered." - .to_owned(), - |engine| format!( - "{} load: simulated memory {} • expert profile {}", - engine.model, - if engine.diagnostics.simulated_used_memory_bytes == 0 { - "off".to_owned() - } else { - format!("{} GiB", engine.diagnostics.simulated_used_memory_bytes / GIB) - }, - if engine.diagnostics.expert_profile_path.is_some() { "set" } else { "off" }, - ), - )) - .size(12), - Space::with_height(8), - text("INACTIVITY").size(11), - row![ - text_input("10", &self.preference_draft.idle_timeout_minutes) - .on_input(Message::PreferenceTimeoutChanged) - .width(90) - .padding(9), - text("minutes before unloading the model").size(13), + row![ + text_input("10", &self.preference_draft.idle_timeout_minutes) + .on_input(Message::PreferenceTimeoutChanged) + .width(90) + .padding(9), + text("minutes before unloading the model").size(13), + ] + .spacing(10) + .align_y(Alignment::Center), + text("Enter a whole number from 1 to 1440.").size(12), ] - .spacing(10) - .align_y(Alignment::Center), - text("Enter a whole number from 1 to 1440.").size(12), + .spacing(10), + ); + let endpoint_group = preference_group( + "LOCAL ENDPOINT", + column![ + preference_input_row( + "Port", + text_input("4000", &self.preference_draft.endpoint_port) + .on_input(Message::PreferenceEndpointPortChanged), + ), + text("Listens on 127.0.0.1. Saving a changed port restarts the local endpoint.") + .size(12), + ] + .spacing(10), + ); + let generation_group = preference_group( + "GENERATION", + column![ + preference_input_row( + "Context tokens", + text_input("32768", &self.preference_draft.context_tokens) + .on_input(Message::PreferenceContextChanged), + ), + preference_input_row( + "Maximum generated tokens", + text_input("50000", &self.preference_draft.max_generated_tokens) + .on_input(Message::PreferenceMaxTokensChanged), + ), + text("System prompt").size(13), + text_input( + "You are a helpful assistant", + &self.preference_draft.system_prompt, + ) + .on_input(Message::PreferenceSystemPromptChanged) + .padding(9), + Space::with_height(4), + text("SAMPLING & REASONING").size(11).color(muted_text()), + preference_input_row( + "Temperature", + text_input("DS4 default", &self.preference_draft.temperature) + .on_input(Message::PreferenceTemperatureChanged), + ), + preference_input_row( + "Top-p", + text_input("DS4 default", &self.preference_draft.top_p) + .on_input(Message::PreferenceTopPChanged), + ), + preference_input_row( + "Min-p", + text_input("DS4 default", &self.preference_draft.min_p) + .on_input(Message::PreferenceMinPChanged), + ), + preference_input_row( + "Seed", + text_input("Random", &self.preference_draft.seed) + .on_input(Message::PreferenceSeedChanged), + ), + row![ + text("Reasoning").size(13).width(Length::Fill), + pick_list( + &REASONING_MODES[..], + Some(self.preference_draft.reasoning_mode), + Message::PreferenceReasoningChanged, + ) + .width(240), + ] + .spacing(12) + .align_y(Alignment::Center), + text("Blank sampling values retain DS4's model-family defaults. Think Max needs at least 393216 context tokens.") + .size(12), + text(turn.map_or_else( + || "Effective settings will appear after valid values are entered.".to_owned(), + |settings| format!( + "Effective: {} context • {} max • temp {} • top-p {} • min-p {} • seed {} • {} • system prompt {}", + settings.context_tokens, + settings.max_generated_tokens, + settings.temperature, + settings.top_p, + settings.min_p, + settings.seed.map_or_else(|| "random".to_owned(), |seed| seed.to_string()), + settings.reasoning_mode, + if settings.system_prompt.is_empty() { "off" } else { "on" }, + ), + )) + .size(12), + ] + .spacing(10), + ); + let execution_group = preference_group( + "EXECUTION", + column![ + preference_input_row( + "CPU helper threads", + text_input("Automatic", &self.preference_draft.cpu_threads) + .on_input(Message::PreferenceCpuThreadsChanged), + ), + preference_input_row("GPU power percent", power), + preference_input_row("Prefill chunk", prefill), + checkbox("Prefer exact quality kernels", self.preference_draft.quality) + .on_toggle(Message::PreferenceQualityChanged), + checkbox("Warm mapped weights at load time", self.preference_draft.warm_weights) + .on_toggle(Message::PreferenceWarmWeightsChanged), + text(if self.preference_draft.model == ModelChoice::Glm52 { + "GLM 5.2 uses full GPU power and selects prefill chunks automatically." + } else { + "Blank numeric values preserve DS4's automatic engine behavior." + }) + .size(12), + text(engine.as_ref().map_or_else( + || "Effective execution settings will appear after valid values are entered." + .to_owned(), + |engine| { + let settings = engine.execution; + format!( + "Metal engine: threads {} • power {}% • prefill {} • quality {} • warm weights {}", + if settings.cpu_threads == 0 { "auto".to_owned() } else { settings.cpu_threads.to_string() }, + if settings.power_percent == 0 { 100 } else { settings.power_percent }, + if settings.prefill_chunk == 0 { "auto".to_owned() } else { settings.prefill_chunk.to_string() }, + if settings.quality { "on" } else { "off" }, + if settings.warm_weights { "on" } else { "off" }, + ) + }, + )) + .size(12), + ] + .spacing(10), + ); + let acceleration_group = preference_group( + "ACCELERATION & MEMORY", + column![ + text("SPECULATIVE DECODING").size(11).color(muted_text()), + preference_input_row( + "MTP draft tokens", + text_input("1", &self.preference_draft.mtp_draft_tokens) + .on_input(Message::PreferenceMtpDraftChanged), + ), + preference_input_row( + "MTP verifier margin", + text_input("3", &self.preference_draft.mtp_margin) + .on_input(Message::PreferenceMtpMarginChanged), + ), + checkbox("Enable integrated GLM MTP", self.preference_draft.glm_mtp) + .on_toggle_maybe(glm_mtp_toggle), + checkbox( + "Log GLM MTP timing counters", + self.preference_draft.glm_mtp_timing, + ) + .on_toggle_maybe(glm_mtp_timing_toggle), + dspark, + preference_input_row("DSpark confidence threshold", dspark_confidence), + checkbox( + "DSpark target-only decode", + self.preference_draft.dspark_strict, + ) + .on_toggle_maybe(dspark_strict_toggle), + text(if self.preference_draft.model.supports_dspark() { + "DSpark uses the managed support artifact; entering a threshold or enabling strict mode also enables DSpark." + } else if self.preference_draft.model == ModelChoice::Glm52 { + "GLM MTP is integrated; DSpark is unavailable for this model." + } else { + "No managed MTP support artifact is available for this model." + }) + .size(12), + text(engine.as_ref().map_or_else( + || "Effective speculative settings will appear after valid values are entered." + .to_owned(), + |engine| { + let settings = engine.speculative; + format!( + "Engine: MTP draft {} • margin {} • GLM MTP {} • timing {} • DSpark {} • confidence {}{} • target-only {}", + settings.mtp_draft_tokens, + settings.mtp_margin, + if settings.glm_mtp { "on" } else { "off" }, + if settings.glm_mtp_timing { "on" } else { "off" }, + if settings.dspark { "on" } else { "off" }, + settings.dspark_confidence_threshold, + if settings.dspark_confidence_threshold_set { " explicit" } else { " default" }, + if settings.dspark_strict { "on" } else { "off" }, + ) + }, + )) + .size(12), + Space::with_height(6), + text("SSD STREAMING").size(11).color(muted_text()), + checkbox("Enable SSD-backed model streaming", self.preference_draft.ssd_streaming) + .on_toggle(Message::PreferenceSsdChanged), + checkbox("Skip automatic expert preload", self.preference_draft.ssd_streaming_cold) + .on_toggle(Message::PreferenceSsdColdChanged), + preference_input_row( + "Expert cache count or GiB", + text_input("Automatic, 128, or 64GB", &self.preference_draft.ssd_cache) + .on_input(Message::PreferenceSsdCacheChanged), + ), + preference_input_row("Fully resident GLM layers", ssd_full_layers), + preference_input_row( + "Explicit expert preload count", + text_input("Automatic", &self.preference_draft.ssd_preload_experts) + .on_input(Message::PreferenceSsdPreloadChanged), + ), + text("A blank full-layer value is automatic; an explicit 0 disables fully resident GLM layers. SSD streaming and DSpark are mutually exclusive.") + .size(12), + text(engine.as_ref().map_or_else( + || "Effective SSD settings will appear after valid values are entered." + .to_owned(), + |engine| { + let settings = engine.ssd; + let cache = if settings.cache_bytes > 0 { + format!("{} GiB", settings.cache_bytes / GIB) + } else if settings.cache_experts > 0 { + format!("{} experts", settings.cache_experts) + } else { + "auto".to_owned() + }; + format!( + "Engine: streaming {} • cold {} • cache {} • full layers {}{} • preload {}", + if settings.enabled { "on" } else { "off" }, + if settings.cold { "on" } else { "off" }, + cache, + settings.full_layers, + if settings.full_layers_set { " explicit" } else { " auto" }, + if settings.preload_experts == 0 { "auto".to_owned() } else { settings.preload_experts.to_string() }, + ) + }, + )) + .size(12), + ] + .spacing(10), + ); + let steering_group = preference_group( + "STEERING & DIAGNOSTICS", + column![ + text("DIRECTIONAL STEERING").size(11).color(muted_text()), + text("Direction-vector file").size(13), + steering_file.padding(9), + preference_input_row("FFN scale", steering_ffn), + preference_input_row("Attention scale", steering_attn), + text(if self.preference_draft.model == ModelChoice::Glm52 { + "Directional steering is not supported for GLM 5.2." + } else { + "With a file and no explicit scale, DS4 defaults the FFN scale to 1. Scales accept -100 through 100." + }) + .size(12), + text(engine.as_ref().map_or_else( + || "Effective steering settings will appear after valid values are entered." + .to_owned(), + |engine| format!( + "Engine: file {} • FFN scale {} • attention scale {}", + if engine.steering.file.is_some() { "set" } else { "off" }, + engine.steering.ffn_scale, + engine.steering.attention_scale, + ), + )) + .size(12), + Space::with_height(6), + text("ADVANCED DIAGNOSTICS").size(11).color(muted_text()), + preference_input_row( + "Simulated used memory (GiB)", + text_input("Disabled", &self.preference_draft.simulated_used_memory_gib) + .on_input(Message::PreferenceSimulatedMemoryChanged), + ), + text("Routed expert profile output").size(13), + text_input("Output file path", &self.preference_draft.expert_profile_path) + .on_input(Message::PreferenceExpertProfileChanged) + .padding(9), + text(engine.as_ref().map_or_else( + || "Effective diagnostic settings will appear after valid values are entered." + .to_owned(), + |engine| format!( + "{} load: simulated memory {} • expert profile {}", + engine.model, + if engine.diagnostics.simulated_used_memory_bytes == 0 { + "off".to_owned() + } else { + format!("{} GiB", engine.diagnostics.simulated_used_memory_bytes / GIB) + }, + if engine.diagnostics.expert_profile_path.is_some() { "set" } else { "off" }, + ), + )) + .size(12), + ] + .spacing(10), + ); + let mut fields = column![ + model_group, + endpoint_group, + generation_group, + execution_group, + acceleration_group, + steering_group, ] - .spacing(10); + .spacing(12); if let Some(error) = &self.preference_error { fields = fields.push(text(error).style(iced::widget::text::danger)); @@ -971,6 +1021,17 @@ fn preference_input_row<'a>( .into() } +fn preference_group<'a>( + title: &'a str, + content: impl Into>, +) -> Element<'a, Message> { + container(column![text(title).size(11).color(muted_text()), content.into(),].spacing(10)) + .width(Length::Fill) + .padding(14) + .style(preference_group_style) + .into() +} + fn download_status_bar(download: &ActiveDownload) -> Element<'_, Message> { let progress = &download.progress; let percent = progress.fraction() * 100.0; @@ -1201,6 +1262,18 @@ fn overview_style(_: &Theme) -> container::Style { } } +fn preference_group_style(_: &Theme) -> container::Style { + container::Style { + background: Some(Background::Color(Color::from_rgb8(38, 38, 40))), + border: Border { + color: Color::from_rgb8(58, 58, 61), + width: 1.0, + radius: 14.0.into(), + }, + ..container::Style::default() + } +} + fn chat_message_style(theme: &Theme, user: bool) -> container::Style { if !user { return overview_style(theme); @@ -1247,6 +1320,10 @@ mod tests { chat_message_style(&theme, true).background, chat_message_style(&theme, false).background ); + assert_eq!( + preference_group_style(&theme).background, + Some(Background::Color(Color::from_rgb8(38, 38, 40))) + ); } #[test] diff --git a/src/database.rs b/src/database.rs index 12458c7..89c60de 100644 --- a/src/database.rs +++ b/src/database.rs @@ -51,6 +51,7 @@ pub struct AppPreferences { pub directional_steering_attn: Option, pub simulated_used_memory_gib: Option, pub expert_profile_path: Option, + pub endpoint_port: i32, } impl Default for AppPreferences { @@ -90,6 +91,7 @@ impl Default for AppPreferences { directional_steering_attn: None, simulated_used_memory_gib: None, expert_profile_path: None, + endpoint_port: 4000, } } } @@ -234,6 +236,7 @@ struct PreferenceChanges<'a> { directional_steering_attn: Option, simulated_used_memory_gib: Option, expert_profile_path: Option<&'a str>, + endpoint_port: i32, } #[derive(Clone, Debug, Identifiable, Queryable, Selectable)] @@ -360,6 +363,7 @@ impl Database { &mut self, selected_model: &str, idle_timeout_minutes: i32, + endpoint_port: i32, generation: &GenerationPreferences, runtime: &RuntimePreferences, ) -> Result { @@ -367,6 +371,9 @@ impl Database { let model = crate::model::ModelChoice::from_id(selected_model) .ok_or_else(|| format!("Unsupported model: {selected_model}"))?; runtime.validate(model)?; + if !(1..=65_535).contains(&endpoint_port) { + return Err("Endpoint port must be between 1 and 65535.".into()); + } let execution = &runtime.execution; let speculative = &runtime.speculative; let (ssd_cache_experts, ssd_cache_gib) = match runtime.ssd.cache { @@ -416,6 +423,7 @@ impl Database { .simulated_used_memory_gib .map(|value| value as i64), expert_profile_path: runtime.diagnostics.expert_profile_path.as_deref(), + endpoint_port, }) .returning(AppPreferences::as_returning()) .get_result(&mut self.connection) @@ -572,6 +580,7 @@ mod tests { assert_eq!(preferences.selected_model, "deepseek-v4-flash"); assert!(!preferences.dspark_enabled); assert_eq!(preferences.idle_timeout_minutes, 10); + assert_eq!(preferences.endpoint_port, 4000); let generation = GenerationPreferences::default(); let runtime = RuntimePreferences::default(); assert!( @@ -579,6 +588,7 @@ mod tests { .update_preferences( "glm-5.2", 30, + 4000, &generation, &RuntimePreferences { speculative: SpeculativePreferences { @@ -592,7 +602,7 @@ mod tests { ); assert!( database - .update_preferences("deepseek-v4-flash", 0, &generation, &runtime,) + .update_preferences("deepseek-v4-flash", 0, 4000, &generation, &runtime,) .is_err() ); let generation = GenerationPreferences { @@ -630,7 +640,7 @@ mod tests { ..RuntimePreferences::default() }; database - .update_preferences("glm-5.2", 30, &generation, &runtime) + .update_preferences("glm-5.2", 30, 4567, &generation, &runtime) .unwrap(); let project = database.create_project("DS4", "/tmp/ds4").unwrap(); @@ -652,6 +662,7 @@ mod tests { let preferences = reopened.load_preferences().unwrap(); assert_eq!(preferences.selected_model, "glm-5.2"); assert_eq!(preferences.idle_timeout_minutes, 30); + assert_eq!(preferences.endpoint_port, 4567); assert_eq!(preferences.generation().unwrap(), generation); assert_eq!(preferences.runtime().unwrap(), runtime); drop(reopened); diff --git a/src/engine.rs b/src/engine.rs index fd59fd2..973aaed 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -305,6 +305,14 @@ pub(crate) struct ChatTurn { pub(crate) content: String, } +pub(crate) struct GenerationOutput { + pub(crate) message: ChatTurn, + pub(crate) prompt_tokens: u32, + pub(crate) cached_tokens: u32, + pub(crate) completion_tokens: u32, + pub(crate) finish_reason: &'static str, +} + #[cfg(target_os = "macos")] impl Generator { pub(crate) fn open(settings: &EngineSettings) -> Result { @@ -337,12 +345,12 @@ impl Generator { cancelled: &AtomicBool, mut emit: impl FnMut(bool, String), mut progress: impl FnMut(u32, u32, Option), - ) -> Result<(), String> { + ) -> Result { self.select_checkpoint(checkpoint)?; - let (generated, prompt_complete) = + let (output, prompt_complete) = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?; let mut completed = messages.to_vec(); - completed.push(generated); + completed.push(output.message.clone()); self.executor.save_checkpoint( checkpoint, if prompt_complete { @@ -350,7 +358,57 @@ impl Generator { } else { [0; 32] }, - ) + )?; + Ok(output) + } + + pub(crate) fn generate_transient( + &mut self, + directory: &Path, + messages: &[ChatTurn], + settings: &TurnSettings, + cancelled: &AtomicBool, + mut emit: impl FnMut(bool, String), + mut progress: impl FnMut(u32, u32, Option), + ) -> Result { + let history = messages + .split_last() + .map_or(messages, |(_, history)| history); + let history_tag = + conversation_tag(&settings.system_prompt, settings.reasoning_mode, history); + let checkpoint = directory.join(format!("{}.bin", hex_tag(history_tag))); + if self.executor.checkpoint_tag() != history_tag { + self.select_checkpoint(&checkpoint)?; + if self.executor.checkpoint_tag() != history_tag { + self.executor.reset()?; + self.checkpoint = None; + let _ = std::fs::remove_file(&checkpoint); + } + } + + let (output, prompt_complete) = + self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?; + let mut completed = messages.to_vec(); + completed.push(output.message.clone()); + let completed_tag = + conversation_tag(&settings.system_prompt, settings.reasoning_mode, &completed); + std::fs::create_dir_all(directory).map_err(|error| { + format!( + "Could not create transient KV cache directory {}: {error}", + directory.display() + ) + })?; + let completed_checkpoint = directory.join(format!("{}.bin", hex_tag(completed_tag))); + self.executor.save_checkpoint( + &completed_checkpoint, + if prompt_complete { + completed_tag + } else { + [0; 32] + }, + )?; + self.checkpoint = Some(completed_checkpoint); + Ok(output) } fn select_checkpoint(&mut self, checkpoint: &Path) -> Result<(), String> { @@ -373,7 +431,7 @@ impl Generator { cancelled: &AtomicBool, emit: &mut impl FnMut(bool, String), progress: &mut impl FnMut(u32, u32, Option), - ) -> Result<(ChatTurn, bool), String> { + ) -> Result<(GenerationOutput, bool), String> { let tokens = match messages.split_last() { Some((latest, history)) if latest.user @@ -418,10 +476,21 @@ impl Generator { reasoning_complete: !reasoning, content: String::new(), }; + let mut emitted_reasoning = 0; + let mut emitted_content = 0; let prompt_tokens = tokens.len(); for (index, token) in tokens.into_iter().enumerate().skip(reused) { if cancelled.load(Ordering::Relaxed) { - return Ok((generated, false)); + return Ok(( + GenerationOutput { + message: generated, + prompt_tokens: prompt_tokens as u32, + cached_tokens: reused as u32, + completion_tokens: 0, + finish_reason: "stop", + }, + false, + )); } self.executor.eval(token)?; if (index + 1).is_multiple_of(16) || index + 1 == prompt_tokens { @@ -436,13 +505,30 @@ impl Generator { .min((max_context - self.executor.position() as usize) as i32) { if cancelled.load(Ordering::Relaxed) { - return Ok((generated, true)); + flush_generated( + &mut generated, + &mut emitted_reasoning, + &mut emitted_content, + &settings.stops, + emit, + ); + return Ok(( + GenerationOutput { + message: generated, + prompt_tokens: prompt_tokens as u32, + cached_tokens: reused as u32, + completion_tokens: generated_tokens, + finish_reason: "stop", + }, + true, + )); } let token = sample( self.executor.logits(), settings.temperature, settings.top_p, settings.min_p, + settings.top_k, &mut rng, ); if self @@ -450,27 +536,89 @@ impl Generator { .model() .is_stop_token_for_reasoning(token, settings.reasoning_mode) { - return Ok((generated, true)); + flush_generated( + &mut generated, + &mut emitted_reasoning, + &mut emitted_content, + &settings.stops, + emit, + ); + return Ok(( + GenerationOutput { + message: generated, + prompt_tokens: prompt_tokens as u32, + cached_tokens: reused as u32, + completion_tokens: generated_tokens, + finish_reason: "stop", + }, + true, + )); } if self.executor.model().is_think_start_token(token) { reasoning = true; generated.reasoning.get_or_insert_default(); } else if self.executor.model().is_think_end_token(token) { + if let Some(reasoning_text) = &mut generated.reasoning + && emit_safe_text( + reasoning_text, + &mut emitted_reasoning, + &settings.stops, + true, + true, + emit, + ) + { + return Ok(( + GenerationOutput { + message: generated, + prompt_tokens: prompt_tokens as u32, + cached_tokens: reused as u32, + completion_tokens: generated_tokens + 1, + finish_reason: "stop", + }, + false, + )); + } reasoning = false; generated.reasoning_complete = true; emit(false, String::new()); } else if let Some(bytes) = self.executor.model().token_bytes(token) { let content = String::from_utf8_lossy(&bytes).into_owned(); - if reasoning { - generated - .reasoning - .get_or_insert_default() - .push_str(&content); + let stopped = if reasoning { + let text = generated.reasoning.get_or_insert_default(); + text.push_str(&content); + emit_safe_text( + text, + &mut emitted_reasoning, + &settings.stops, + false, + true, + emit, + ) } else { generated.reasoning_complete = true; generated.content.push_str(&content); + emit_safe_text( + &mut generated.content, + &mut emitted_content, + &settings.stops, + false, + false, + emit, + ) + }; + if stopped { + return Ok(( + GenerationOutput { + message: generated, + prompt_tokens: prompt_tokens as u32, + cached_tokens: reused as u32, + completion_tokens: generated_tokens + 1, + finish_reason: "stop", + }, + false, + )); } - emit(reasoning, content); } self.executor.eval(token)?; generated_tokens += 1; @@ -483,10 +631,103 @@ impl Generator { ), ); } - Ok((generated, true)) + flush_generated( + &mut generated, + &mut emitted_reasoning, + &mut emitted_content, + &settings.stops, + emit, + ); + Ok(( + GenerationOutput { + message: generated, + prompt_tokens: prompt_tokens as u32, + cached_tokens: reused as u32, + completion_tokens: generated_tokens, + finish_reason: "length", + }, + true, + )) } } +#[cfg(target_os = "macos")] +fn flush_generated( + generated: &mut ChatTurn, + emitted_reasoning: &mut usize, + emitted_content: &mut usize, + stops: &[String], + emit: &mut impl FnMut(bool, String), +) { + if let Some(reasoning) = &mut generated.reasoning { + let _ = emit_safe_text(reasoning, emitted_reasoning, stops, true, true, emit); + } + let _ = emit_safe_text( + &mut generated.content, + emitted_content, + stops, + true, + false, + emit, + ); +} + +#[cfg(target_os = "macos")] +fn emit_safe_text( + text: &mut String, + emitted: &mut usize, + stops: &[String], + final_flush: bool, + reasoning: bool, + emit: &mut impl FnMut(bool, String), +) -> bool { + let stop = stops + .iter() + .filter_map(|stop| { + text[*emitted..] + .find(stop) + .map(|position| *emitted + position) + }) + .min(); + if let Some(stop) = stop { + if stop > *emitted { + emit(reasoning, text[*emitted..stop].to_owned()); + } + text.truncate(stop); + *emitted = stop; + return true; + } + let hold = if final_flush { + 0 + } else { + stops + .iter() + .map(|stop| stop.len().saturating_sub(1)) + .max() + .unwrap_or(0) + }; + let mut safe = text.len().saturating_sub(hold); + while safe > *emitted && !text.is_char_boundary(safe) { + safe -= 1; + } + if safe > *emitted { + emit(reasoning, text[*emitted..safe].to_owned()); + *emitted = safe; + } + false +} + +#[cfg(target_os = "macos")] +fn hex_tag(tag: [u8; 32]) -> String { + const HEX: &[u8; 16] = b"0123456789abcdef"; + let mut output = String::with_capacity(64); + for byte in tag { + output.push(HEX[(byte >> 4) as usize] as char); + output.push(HEX[(byte & 0x0f) as usize] as char); + } + output +} + #[cfg(target_os = "macos")] fn conversation_tag(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn]) -> [u8; 32] { fn text(hasher: &mut Sha256, value: &str) { @@ -518,7 +759,14 @@ fn conversation_tag(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn } #[cfg(target_os = "macos")] -fn sample(logits: &[f32], temperature: f32, top_p: f32, min_p: f32, rng: &mut Rng) -> i32 { +fn sample( + logits: &[f32], + temperature: f32, + top_p: f32, + min_p: f32, + top_k: i32, + rng: &mut Rng, +) -> i32 { if temperature <= 0.0 { return logits .iter() @@ -554,12 +802,16 @@ fn sample(logits: &[f32], temperature: f32, top_p: f32, min_p: f32, rng: &mut Rn .max_by(|a, b| a.1.total_cmp(b.1)) .map_or(0, |(index, _)| index as i32); } - if top_p < 1.0 { + if top_p < 1.0 || top_k > 0 { probabilities.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0))); - let total: f32 = logits + if top_k > 0 { + probabilities.truncate(probabilities.len().min(top_k as usize)); + } + } + if top_p < 1.0 { + let total: f32 = probabilities .iter() - .filter(|logit| logit.is_finite()) - .map(|logit| ((*logit - maximum) / temperature).exp()) + .map(|(_, probability)| probability) .sum(); let mut kept = 0.0; let count = probabilities @@ -612,7 +864,27 @@ mod sampling_tests { #[test] fn zero_temperature_is_greedy() { let mut rng = Rng::new(1); - assert_eq!(sample(&[1.0, 4.0, 2.0], 0.0, 1.0, 0.0, &mut rng), 1); + assert_eq!(sample(&[1.0, 4.0, 2.0], 0.0, 1.0, 0.0, 0, &mut rng), 1); + } + + #[test] + fn top_k_and_stream_stops_are_applied_before_output() { + let mut rng = Rng::new(1); + assert_eq!(sample(&[1.0, 4.0, 2.0], 1.0, 1.0, 0.0, 1, &mut rng), 1); + + let mut text = "hello STOP hidden".to_owned(); + let mut emitted = 0; + let mut chunks = Vec::new(); + assert!(emit_safe_text( + &mut text, + &mut emitted, + &["STOP".into()], + false, + false, + &mut |_, chunk| chunks.push(chunk), + )); + assert_eq!(text, "hello "); + assert_eq!(chunks, ["hello "]); } #[test] diff --git a/src/main.rs b/src/main.rs index 0190446..d9c2a02 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,8 +3,14 @@ mod database; mod engine; mod model; #[cfg(target_os = "macos")] +mod native_edit; +#[cfg(target_os = "macos")] mod native_menu; +#[cfg(target_os = "macos")] +mod runtime; mod schema; +#[cfg(target_os = "macos")] +mod server; mod settings; use app::{App, Message, app_icon, app_theme}; diff --git a/src/model.rs b/src/model.rs index edd3a69..8264e16 100644 --- a/src/model.rs +++ b/src/model.rs @@ -327,6 +327,13 @@ pub(crate) fn managed_artifacts(models_path: &Path) -> Vec { .collect() } +pub(crate) fn installed_models(models_path: &Path) -> Vec { + MODEL_CHOICES + .into_iter() + .filter(|model| model.main_artifact().is_installed(*model, models_path)) + .collect() +} + pub(crate) fn artifact_download_progress( id: ManagedArtifactId, models_path: &Path, diff --git a/src/native_edit.rs b/src/native_edit.rs new file mode 100644 index 0000000..7fb4ba9 --- /dev/null +++ b/src/native_edit.rs @@ -0,0 +1,315 @@ +use std::collections::VecDeque; +use std::sync::{Arc, Mutex}; + +use iced::advanced::layout; +use iced::advanced::overlay; +use iced::advanced::renderer; +use iced::advanced::widget::{Operation, Tree, tree}; +use iced::advanced::{Clipboard, Layout, Shell, Widget}; +use iced::event; +use iced::keyboard::{self, Key, Location, Modifiers, key}; +use iced::mouse; +use iced::{Element, Event, Length, Rectangle, Size, Vector}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum EditCommand { + Undo, + Redo, + Cut, + Copy, + Paste, + SelectAll, +} + +pub(crate) type EditCommandQueue = Arc>>; + +pub(crate) fn command_queue() -> EditCommandQueue { + Arc::new(Mutex::new(VecDeque::new())) +} + +pub(crate) fn queue_command(queue: &EditCommandQueue, command: EditCommand) { + queue + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push_back(command); +} + +fn pop_command(queue: &EditCommandQueue) -> Option { + queue + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .pop_front() +} + +/// Replays native Edit menu actions through the focused Iced widget. +pub(crate) struct NativeEdit<'a, Message, Theme = iced::Theme, Renderer = iced::Renderer> { + content: Element<'a, Message, Theme, Renderer>, + commands: EditCommandQueue, +} + +impl<'a, Message, Theme, Renderer> NativeEdit<'a, Message, Theme, Renderer> { + fn new( + content: impl Into>, + commands: EditCommandQueue, + ) -> Self { + Self { + content: content.into(), + commands, + } + } +} + +#[derive(Default)] +struct State; + +fn modifier_sync_event(event: &Event) -> Option { + let Event::Keyboard(keyboard::Event::KeyPressed { modifiers, .. }) = event else { + return None; + }; + Some(Event::Keyboard(keyboard::Event::ModifiersChanged( + *modifiers, + ))) +} + +fn command_events(command: EditCommand) -> [Event; 4] { + let (character, modified_character, physical_key, modifiers) = match command { + EditCommand::Undo => ("z", "z", key::Code::KeyZ, Modifiers::COMMAND), + EditCommand::Redo => ( + "z", + "Z", + key::Code::KeyZ, + Modifiers::COMMAND | Modifiers::SHIFT, + ), + EditCommand::Cut => ("x", "x", key::Code::KeyX, Modifiers::COMMAND), + EditCommand::Copy => ("c", "c", key::Code::KeyC, Modifiers::COMMAND), + EditCommand::Paste => ("v", "v", key::Code::KeyV, Modifiers::COMMAND), + EditCommand::SelectAll => ("a", "a", key::Code::KeyA, Modifiers::COMMAND), + }; + let key = Key::Character(character.into()); + [ + Event::Keyboard(keyboard::Event::ModifiersChanged(modifiers)), + Event::Keyboard(keyboard::Event::KeyPressed { + key: key.clone(), + modified_key: Key::Character(modified_character.into()), + physical_key: key::Physical::Code(physical_key), + location: Location::Standard, + modifiers, + text: None, + }), + Event::Keyboard(keyboard::Event::KeyReleased { + key, + location: Location::Standard, + modifiers, + }), + Event::Keyboard(keyboard::Event::ModifiersChanged(Modifiers::default())), + ] +} + +impl<'a, Message, Theme, Renderer> Widget + for NativeEdit<'a, Message, Theme, Renderer> +where + Renderer: renderer::Renderer, +{ + fn tag(&self) -> tree::Tag { + tree::Tag::of::() + } + + fn state(&self) -> tree::State { + tree::State::new(State) + } + + fn children(&self) -> Vec { + vec![Tree::new(&self.content)] + } + + fn diff(&self, tree: &mut Tree) { + tree.diff_children(std::slice::from_ref(&self.content)); + } + + fn size(&self) -> Size { + self.content.as_widget().size() + } + + fn layout( + &self, + tree: &mut Tree, + renderer: &Renderer, + limits: &layout::Limits, + ) -> layout::Node { + self.content + .as_widget() + .layout(&mut tree.children[0], renderer, limits) + } + + fn operate( + &self, + tree: &mut Tree, + layout: Layout<'_>, + renderer: &Renderer, + operation: &mut dyn Operation, + ) { + self.content + .as_widget() + .operate(&mut tree.children[0], layout, renderer, operation); + } + + fn on_event( + &mut self, + tree: &mut Tree, + event: Event, + layout: Layout<'_>, + cursor: mouse::Cursor, + renderer: &Renderer, + clipboard: &mut dyn Clipboard, + shell: &mut Shell<'_, Message>, + viewport: &Rectangle, + ) -> event::Status { + let mut status = event::Status::Ignored; + if matches!( + event, + Event::Window(iced::window::Event::RedrawRequested(_)) + ) { + while let Some(command) = pop_command(&self.commands) { + for command_event in command_events(command) { + if self.content.as_widget_mut().on_event( + &mut tree.children[0], + command_event, + layout, + cursor, + renderer, + clipboard, + shell, + viewport, + ) == event::Status::Captured + { + status = event::Status::Captured; + } + } + } + } + + if let Some(sync_event) = modifier_sync_event(&event) { + let _ = self.content.as_widget_mut().on_event( + &mut tree.children[0], + sync_event, + layout, + cursor, + renderer, + clipboard, + shell, + viewport, + ); + } + + if self.content.as_widget_mut().on_event( + &mut tree.children[0], + event, + layout, + cursor, + renderer, + clipboard, + shell, + viewport, + ) == event::Status::Captured + { + event::Status::Captured + } else { + status + } + } + + fn mouse_interaction( + &self, + tree: &Tree, + layout: Layout<'_>, + cursor: mouse::Cursor, + viewport: &Rectangle, + renderer: &Renderer, + ) -> mouse::Interaction { + self.content.as_widget().mouse_interaction( + &tree.children[0], + layout, + cursor, + viewport, + renderer, + ) + } + + fn draw( + &self, + tree: &Tree, + renderer: &mut Renderer, + theme: &Theme, + style: &renderer::Style, + layout: Layout<'_>, + cursor: mouse::Cursor, + viewport: &Rectangle, + ) { + self.content.as_widget().draw( + &tree.children[0], + renderer, + theme, + style, + layout, + cursor, + viewport, + ); + } + + fn overlay<'b>( + &'b mut self, + tree: &'b mut Tree, + layout: Layout<'_>, + renderer: &Renderer, + translation: Vector, + ) -> Option> { + self.content + .as_widget_mut() + .overlay(&mut tree.children[0], layout, renderer, translation) + } +} + +impl<'a, Message, Theme, Renderer> From> + for Element<'a, Message, Theme, Renderer> +where + Message: 'a, + Theme: 'a, + Renderer: 'a + renderer::Renderer, +{ + fn from(bridge: NativeEdit<'a, Message, Theme, Renderer>) -> Self { + Element::new(bridge) + } +} + +pub(crate) fn native_edit<'a, Message, Theme, Renderer>( + content: impl Into>, + commands: EditCommandQueue, +) -> NativeEdit<'a, Message, Theme, Renderer> { + NativeEdit::new(content, commands) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn paste_replays_a_complete_command_shortcut() { + let events = command_events(EditCommand::Paste); + assert!(matches!( + events[0], + Event::Keyboard(keyboard::Event::ModifiersChanged(Modifiers::COMMAND)) + )); + assert!(matches!( + events[1], + Event::Keyboard(keyboard::Event::KeyPressed { + key: Key::Character(ref key), + modifiers: Modifiers::COMMAND, + .. + }) if key == "v" + )); + assert!(matches!( + events[3], + Event::Keyboard(keyboard::Event::ModifiersChanged(modifiers)) if modifiers.is_empty() + )); + } +} diff --git a/src/native_menu.rs b/src/native_menu.rs index 866a5be..e123f71 100644 --- a/src/native_menu.rs +++ b/src/native_menu.rs @@ -1,8 +1,16 @@ -use muda::accelerator::{Accelerator, Code, Modifiers}; +use muda::accelerator::{Accelerator, CMD_OR_CTRL, Code, Modifiers}; use muda::{Menu, MenuEvent, MenuItem, PredefinedMenuItem, Submenu}; +use crate::native_edit::EditCommand; + const PREFERENCES: &str = "preferences"; const MODEL_MANAGER: &str = "model-manager"; +const UNDO: &str = "undo"; +const REDO: &str = "redo"; +const CUT: &str = "cut"; +const COPY: &str = "copy"; +const PASTE: &str = "paste"; +const SELECT_ALL: &str = "select-all"; pub(crate) struct NativeMenu { _menu: Menu, @@ -12,6 +20,7 @@ pub(crate) struct NativeMenu { pub(crate) enum NativeMenuEvent { Preferences, ModelManager, + Edit(EditCommand), } pub(crate) fn install() -> Result { @@ -34,6 +43,12 @@ pub(crate) fn install() -> Result { Code::KeyM, )), ); + let undo = edit_item(UNDO, "Undo", Code::KeyZ, None); + let redo = edit_item(REDO, "Redo", Code::KeyZ, Some(Modifiers::SHIFT)); + let cut = edit_item(CUT, "Cut", Code::KeyX, None); + let copy = edit_item(COPY, "Copy", Code::KeyC, None); + let paste = edit_item(PASTE, "Paste", Code::KeyV, None); + let select_all = edit_item(SELECT_ALL, "Select All", Code::KeyA, None); application .append_items(&[ @@ -51,13 +66,13 @@ pub(crate) fn install() -> Result { ]) .map_err(|error| error.to_string())?; edit.append_items(&[ - &PredefinedMenuItem::undo(None), - &PredefinedMenuItem::redo(None), + &undo, + &redo, &PredefinedMenuItem::separator(), - &PredefinedMenuItem::cut(None), - &PredefinedMenuItem::copy(None), - &PredefinedMenuItem::paste(None), - &PredefinedMenuItem::select_all(None), + &cut, + ©, + &paste, + &select_all, ]) .map_err(|error| error.to_string())?; window @@ -79,12 +94,47 @@ pub(crate) fn install() -> Result { pub(crate) fn next_event() -> Option { while let Ok(event) = MenuEvent::receiver().try_recv() { - if event.id == PREFERENCES { - return Some(NativeMenuEvent::Preferences); - } - if event.id == MODEL_MANAGER { - return Some(NativeMenuEvent::ModelManager); - } + let event = match event.id.0.as_str() { + PREFERENCES => NativeMenuEvent::Preferences, + MODEL_MANAGER => NativeMenuEvent::ModelManager, + UNDO => NativeMenuEvent::Edit(EditCommand::Undo), + REDO => NativeMenuEvent::Edit(EditCommand::Redo), + CUT => NativeMenuEvent::Edit(EditCommand::Cut), + COPY => NativeMenuEvent::Edit(EditCommand::Copy), + PASTE => NativeMenuEvent::Edit(EditCommand::Paste), + SELECT_ALL => NativeMenuEvent::Edit(EditCommand::SelectAll), + _ => continue, + }; + return Some(event); } None } + +fn edit_item( + id: &'static str, + label: &'static str, + code: Code, + extra_modifier: Option, +) -> MenuItem { + MenuItem::with_id( + id, + label, + true, + Some(Accelerator::new( + Some(CMD_OR_CTRL | extra_modifier.unwrap_or_else(Modifiers::empty)), + code, + )), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn edit_items_are_custom_menu_events() { + let paste = edit_item(PASTE, "Paste", Code::KeyV, None); + assert_eq!(paste.id().0, PASTE); + assert!(paste.is_enabled()); + } +} diff --git a/src/runtime.rs b/src/runtime.rs new file mode 100644 index 0000000..9a64fb6 --- /dev/null +++ b/src/runtime.rs @@ -0,0 +1,151 @@ +use crate::engine::{ChatTurn, GenerationOutput, Generator}; +use crate::settings::{EngineSettings, TurnSettings}; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::mpsc::{self, Receiver, Sender}; +use std::thread; +use std::time::{Duration, Instant}; + +#[derive(Clone)] +pub(crate) struct GenerationService { + commands: Sender, +} + +pub(crate) struct ActiveGeneration { + pub(crate) events: Receiver, + pub(crate) cancel: Arc, +} + +pub(crate) enum CheckpointTarget { + Local(PathBuf), + Transient(PathBuf), +} + +pub(crate) enum GenerationEvent { + Loading, + Chunk { + reasoning: bool, + content: String, + }, + Context { + used: u32, + limit: u32, + tokens_per_second: Option, + }, + Finished(Result), +} + +struct Command { + engine: EngineSettings, + turn: TurnSettings, + messages: Vec, + checkpoint: CheckpointTarget, + idle_timeout: Duration, + cancel: Arc, + events: Sender, +} + +impl GenerationService { + pub(crate) fn spawn() -> Result { + let (commands, receiver) = mpsc::channel::(); + thread::Builder::new() + .name("model-runtime".into()) + .spawn(move || run(receiver)) + .map_err(|error| format!("Could not start the model runtime: {error}"))?; + Ok(Self { commands }) + } + + pub(crate) fn generate( + &self, + engine: EngineSettings, + turn: TurnSettings, + messages: Vec, + checkpoint: CheckpointTarget, + idle_timeout: Duration, + ) -> Result { + let cancel = Arc::new(AtomicBool::new(false)); + let (events, receiver) = mpsc::channel(); + self.commands + .send(Command { + engine, + turn, + messages, + checkpoint, + idle_timeout, + cancel: Arc::clone(&cancel), + events, + }) + .map_err(|_| "The model runtime stopped unexpectedly.".to_owned())?; + Ok(ActiveGeneration { + events: receiver, + cancel, + }) + } +} + +fn run(commands: Receiver) { + let mut loaded = None::<(EngineSettings, Generator)>; + let mut last_used = Instant::now(); + let mut idle_timeout = Duration::from_secs(15 * 60); + loop { + match commands.recv_timeout(Duration::from_secs(1)) { + Ok(command) => { + idle_timeout = command.idle_timeout; + if loaded + .as_ref() + .is_none_or(|(settings, _)| settings != &command.engine) + { + let _ = command.events.send(GenerationEvent::Loading); + loaded = match Generator::open(&command.engine) { + Ok(generator) => Some((command.engine.clone(), generator)), + Err(error) => { + let _ = command.events.send(GenerationEvent::Finished(Err(error))); + None + } + }; + } + if let Some((_, generator)) = &mut loaded { + let mut emit = |reasoning, content| { + let _ = command + .events + .send(GenerationEvent::Chunk { reasoning, content }); + }; + let mut progress = |used, limit, tokens_per_second| { + let _ = command.events.send(GenerationEvent::Context { + used, + limit, + tokens_per_second, + }); + }; + let result = match command.checkpoint { + CheckpointTarget::Local(path) => generator.generate( + &path, + &command.messages, + &command.turn, + &command.cancel, + &mut emit, + &mut progress, + ), + CheckpointTarget::Transient(directory) => generator.generate_transient( + &directory, + &command.messages, + &command.turn, + &command.cancel, + &mut emit, + &mut progress, + ), + }; + let _ = command.events.send(GenerationEvent::Finished(result)); + last_used = Instant::now(); + } + } + Err(mpsc::RecvTimeoutError::Timeout) => { + if loaded.is_some() && last_used.elapsed() >= idle_timeout { + loaded = None; + } + } + Err(mpsc::RecvTimeoutError::Disconnected) => break, + } + } +} diff --git a/src/schema.rs b/src/schema.rs index ec3b8cf..972254f 100644 --- a/src/schema.rs +++ b/src/schema.rs @@ -53,6 +53,7 @@ diesel::table! { directional_steering_attn -> Nullable, simulated_used_memory_gib -> Nullable, expert_profile_path -> Nullable, + endpoint_port -> Integer, } } diff --git a/src/server.rs b/src/server.rs new file mode 100644 index 0000000..01c65a0 --- /dev/null +++ b/src/server.rs @@ -0,0 +1,1116 @@ +use crate::database::AppPreferences; +use crate::engine::ChatTurn; +use crate::model::{self, ModelChoice}; +use crate::runtime::{CheckpointTarget, GenerationEvent, GenerationService}; +use crate::settings::{ReasoningMode, effective_settings}; +use serde::Deserialize; +use serde_json::{Map, Value, json}; +use std::collections::HashMap; +use std::io::{Read, Write}; +use std::net::{TcpListener, TcpStream}; +use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, RwLock}; +use std::thread::{self, JoinHandle}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +const MAX_HEADER_BYTES: usize = 64 * 1024; +const MAX_BODY_BYTES: usize = 64 * 1024 * 1024; +const TOOLS_PROMPT: &str = "## Tools\n\n\ +You have access to a set of tools to help answer the user question. You can invoke tools by writing a \"<｜DSML｜tool_calls>\" block like the following:\n\n\ +<｜DSML｜tool_calls>\n\ +<｜DSML｜invoke name=\"$TOOL_NAME\">\n\ +<｜DSML｜parameter name=\"$PARAMETER_NAME\" string=\"true|false\">$PARAMETER_VALUE\n\ +...\n\ +\n\ +<｜DSML｜invoke name=\"$TOOL_NAME2\">\n\ +...\n\ +\n\ +\n\n\ +String parameters should be specified as raw text and set `string=\"true\"`. Preserve characters such as `>`, `&`, and `&&` exactly; never replace normal string characters with XML or HTML entity escapes. Only if a string value itself contains the exact closing parameter tag ``, write that tag as `</｜DSML｜parameter>` inside the value. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string=\"false\"`.\n\n\ +If thinking_mode is enabled (triggered by ), you MUST output your complete reasoning inside ... BEFORE any tool calls or final response.\n\n\ +Otherwise, output directly after with tool calls or final response.\n\n\ +### Available Tool Schemas\n\n"; + +pub(crate) struct ServerHandle { + stop: Arc, + thread: Option>, +} + +struct State { + generation: GenerationService, + preferences: Arc>, + models_path: PathBuf, + cache_path: PathBuf, + sequence: AtomicU64, + tool_memory: Mutex>, +} + +#[derive(Deserialize)] +struct ChatRequest { + #[serde(default)] + model: Option, + #[serde(default)] + messages: Vec, + #[serde(default)] + tools: Vec, + #[serde(default)] + tool_choice: Option, + #[serde(default)] + max_tokens: Option, + #[serde(default)] + max_completion_tokens: Option, + #[serde(default)] + temperature: Option, + #[serde(default)] + top_p: Option, + #[serde(default)] + min_p: Option, + #[serde(default)] + top_k: Option, + #[serde(default)] + seed: Option, + #[serde(default)] + stream: bool, + #[serde(default)] + stream_options: Option, + #[serde(default)] + thinking: Option, + #[serde(default)] + think: Option, + #[serde(default)] + reasoning_effort: Option, + #[serde(default)] + stop: Option, +} + +#[derive(Deserialize)] +struct ApiMessage { + #[serde(default = "default_role")] + role: String, + #[serde(default)] + content: Value, + #[serde(default)] + reasoning_content: Value, + #[serde(default)] + tool_calls: Vec, +} + +#[derive(Clone, Deserialize)] +struct ApiToolCall { + #[serde(default)] + id: String, + function: ApiFunction, +} + +#[derive(Clone, Deserialize)] +struct ApiFunction { + name: String, + #[serde(default = "empty_arguments")] + arguments: String, +} + +#[derive(Deserialize)] +struct StreamOptions { + #[serde(default)] + include_usage: bool, +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum OneOrMany { + One(String), + Many(Vec), +} + +struct ParsedRequest { + model_id: String, + messages: Vec, + turn: crate::settings::TurnSettings, + engine: crate::settings::EngineSettings, + idle_timeout: Duration, + stream: bool, + include_usage: bool, + has_tools: bool, +} + +struct ResponseOptions { + model_id: String, + stream: bool, + include_usage: bool, + has_tools: bool, +} + +struct HttpRequest { + method: String, + path: String, + body: Vec, +} + +impl ServerHandle { + pub(crate) fn spawn( + generation: GenerationService, + preferences: Arc>, + models_path: PathBuf, + cache_path: PathBuf, + port: u16, + ) -> Result { + let address = format!("127.0.0.1:{port}"); + let listener = TcpListener::bind(("127.0.0.1", port)) + .map_err(|error| format!("Could not listen on http://{address}: {error}"))?; + listener + .set_nonblocking(true) + .map_err(|error| format!("Could not configure http://{address}: {error}"))?; + let stop = Arc::new(AtomicBool::new(false)); + let worker_stop = Arc::clone(&stop); + let state = Arc::new(State { + generation, + preferences, + models_path, + cache_path, + sequence: AtomicU64::new(0), + tool_memory: Mutex::new(HashMap::new()), + }); + let thread = thread::Builder::new() + .name("local-http".into()) + .spawn(move || serve(listener, state, worker_stop)) + .map_err(|error| format!("Could not start the local HTTP service: {error}"))?; + Ok(Self { + stop, + thread: Some(thread), + }) + } +} + +impl Drop for ServerHandle { + fn drop(&mut self) { + self.stop.store(true, Ordering::Relaxed); + if let Some(thread) = self.thread.take() { + let _ = thread.join(); + } + } +} + +fn serve(listener: TcpListener, state: Arc, stop: Arc) { + while !stop.load(Ordering::Relaxed) { + match listener.accept() { + Ok((stream, _)) => { + let state = Arc::clone(&state); + let _ = thread::Builder::new() + .name("http-request".into()) + .spawn(move || handle(stream, &state)); + } + Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { + thread::sleep(Duration::from_millis(50)); + } + Err(error) => { + eprintln!("DS4Server: endpoint accept failed: {error}"); + thread::sleep(Duration::from_millis(100)); + } + } + } +} + +fn handle(mut stream: TcpStream, state: &State) { + let _ = stream.set_read_timeout(Some(Duration::from_secs(30))); + let request = match read_request(&mut stream) { + Ok(request) => request, + Err(error) => { + let _ = send_error(&mut stream, 400, &error); + return; + } + }; + match (request.method.as_str(), request.path.as_str()) { + ("OPTIONS", _) => { + let _ = send_response(&mut stream, 204, None, &[]); + } + ("GET", "/v1/models") => { + let body = models_json(state); + let _ = send_json(&mut stream, 200, &body); + } + ("POST", "/v1/chat/completions") => { + if let Err(error) = chat_completion(&mut stream, state, &request.body) { + let _ = send_error(&mut stream, error.0, &error.1); + } + } + ("GET", path) if path.starts_with("/v1/models/") => { + let id = &path[11..]; + let model = installed_endpoint_models(&state.models_path) + .into_iter() + .find(|model| model.id() == id); + if let Some(model) = model { + let context = state + .preferences + .read() + .map_or(32_768, |preferences| preferences.context_tokens); + let _ = send_json(&mut stream, 200, &model_json(model, context)); + } else { + let _ = send_error(&mut stream, 404, "unknown model"); + } + } + _ => { + let _ = send_error(&mut stream, 404, "unknown endpoint"); + } + } +} + +fn chat_completion( + stream: &mut TcpStream, + state: &State, + body: &[u8], +) -> Result<(), (u16, String)> { + let request: ChatRequest = + serde_json::from_slice(body).map_err(|_| (400, "invalid JSON request".to_owned()))?; + let parsed = parse_chat_request(state, request)?; + let response = ResponseOptions { + model_id: parsed.model_id, + stream: parsed.stream, + include_usage: parsed.include_usage, + has_tools: parsed.has_tools, + }; + let active = state + .generation + .generate( + parsed.engine, + parsed.turn, + parsed.messages, + CheckpointTarget::Transient(state.cache_path.clone()), + parsed.idle_timeout, + ) + .map_err(|error| (500, error))?; + let sequence = state.sequence.fetch_add(1, Ordering::Relaxed) + 1; + let id = format!("chatcmpl-{sequence}"); + if response.stream { + stream_response(stream, state, response, active, &id) + } else { + final_response(stream, state, response, active, &id) + } +} + +fn parse_chat_request(state: &State, request: ChatRequest) -> Result { + if request.messages.is_empty() { + return Err((400, "missing messages".into())); + } + let requested_id = request.model.as_deref().unwrap_or_default(); + let preferences = state + .preferences + .read() + .map_err(|_| (500, "preferences are unavailable".to_owned()))? + .clone(); + let model = if requested_id.is_empty() { + ModelChoice::from_id(&preferences.selected_model) + } else { + model_alias(requested_id) + } + .ok_or_else(|| (400, format!("unknown model: {requested_id}")))?; + if !installed_endpoint_models(&state.models_path).contains(&model) { + return Err((400, format!("model is not installed and verified: {model}"))); + } + + let mut generation = preferences.generation().map_err(|error| (500, error))?; + generation.system_prompt.clear(); + generation.max_generated_tokens = request + .max_completion_tokens + .or(request.max_tokens) + .unwrap_or(generation.max_generated_tokens); + if generation.max_generated_tokens <= 0 { + return Err((400, "max_tokens must be positive".into())); + } + generation.temperature = request.temperature.or(Some(1.0)); + generation.top_p = request.top_p.or(Some(1.0)); + generation.min_p = request.min_p.or(Some(0.05)); + generation.seed = request.seed.filter(|seed| *seed > 0); + generation.reasoning_mode = request_reasoning(&request, requested_id)?; + + let tools_enabled = !request.tools.is_empty() + && request.tool_choice.as_ref().and_then(Value::as_str) != Some("none"); + let (system, messages) = + render_messages(state, &request.messages, &request.tools, tools_enabled)?; + generation.system_prompt = system; + let runtime = preferences.runtime().map_err(|error| (500, error))?; + let mut effective = effective_settings(model, &generation, &runtime, &state.models_path) + .map_err(|error| (400, error))?; + effective.turn.top_k = request.top_k.unwrap_or(0); + if effective.turn.top_k < 0 { + return Err((400, "top_k must not be negative".into())); + } + effective.turn.stops = match request.stop { + Some(OneOrMany::One(stop)) => vec![stop], + Some(OneOrMany::Many(stops)) => stops, + None => Vec::new(), + }; + effective.turn.stops.retain(|stop| !stop.is_empty()); + + Ok(ParsedRequest { + model_id: if requested_id.is_empty() { + model.id().to_owned() + } else { + requested_id.to_owned() + }, + messages, + turn: effective.turn, + engine: effective.engine, + idle_timeout: Duration::from_secs(preferences.idle_timeout_minutes.max(1) as u64 * 60), + stream: request.stream, + include_usage: request + .stream_options + .is_some_and(|options| options.include_usage), + has_tools: tools_enabled, + }) +} + +fn request_reasoning( + request: &ChatRequest, + model_id: &str, +) -> Result { + let explicit_thinking = request + .think + .or_else(|| request.thinking.as_ref().and_then(thinking_enabled)); + let mut reasoning = match request.reasoning_effort.as_deref() { + Some("max") => ReasoningMode::Max, + Some("none") => ReasoningMode::Direct, + Some("xhigh" | "high" | "medium" | "low" | "minimal") | None => ReasoningMode::High, + Some(value) => return Err((400, format!("unsupported reasoning_effort: {value}"))), + }; + if explicit_thinking == Some(false) + || (explicit_thinking.is_none() + && matches!( + model_id, + "deepseek-chat" | "glm-5.2-chat" | "glm-5.2-no-think" | "glm-5.2-nothink" + )) + { + reasoning = ReasoningMode::Direct; + } + Ok(reasoning) +} + +fn thinking_enabled(value: &Value) -> Option { + value.as_bool().or_else(|| { + value.as_str().map(|value| value != "disabled").or_else(|| { + value + .get("type") + .and_then(Value::as_str) + .map(|value| value != "disabled") + }) + }) +} + +fn render_messages( + state: &State, + messages: &[ApiMessage], + tools: &[Value], + tools_enabled: bool, +) -> Result<(String, Vec), (u16, String)> { + let preserve_reasoning = tools_enabled + || messages.iter().any(|message| { + matches!(message.role.as_str(), "tool" | "function") || !message.tool_calls.is_empty() + }); + let mut system = String::new(); + if tools_enabled { + system.push_str(TOOLS_PROMPT); + for tool in tools { + let schema = tool.get("function").unwrap_or(tool); + if !system.ends_with("\n\n") { + system.push('\n'); + } + system.push_str( + &serde_json::to_string(schema) + .map_err(|error| (400, format!("invalid tool schema: {error}")))?, + ); + system.push('\n'); + } + system.push_str( + "\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls. Use the exact parameter names from the schemas.", + ); + } + + let mut turns = Vec::::new(); + for message in messages { + let content = content_text(&message.content); + match message.role.as_str() { + "system" | "developer" => { + if !system.is_empty() { + system.push_str("\n\n"); + } + system.push_str(&content); + } + "user" => turns.push(ChatTurn { + user: true, + reasoning: None, + reasoning_complete: true, + content, + }), + "tool" | "function" => { + let wrapped = format!( + "{}", + escape_tool_result(&content) + ); + if let Some(previous) = turns.last_mut() + && previous.user + && previous.content.starts_with("") + { + previous.content.push_str(&wrapped); + } else { + turns.push(ChatTurn { + user: true, + reasoning: None, + reasoning_complete: true, + content: wrapped, + }); + } + } + "assistant" => { + let mut content = content; + if !message.tool_calls.is_empty() { + content.push_str(&replayed_or_canonical_tools(state, &message.tool_calls)); + } + let reasoning = content_text(&message.reasoning_content); + turns.push(ChatTurn { + user: false, + reasoning: (preserve_reasoning && !reasoning.is_empty()).then_some(reasoning), + reasoning_complete: true, + content, + }); + } + role => return Err((400, format!("unsupported message role: {role}"))), + } + } + Ok((system, turns)) +} + +fn replayed_or_canonical_tools(state: &State, calls: &[ApiToolCall]) -> String { + if let Ok(memory) = state.tool_memory.lock() + && let Some(raw) = calls.iter().find_map(|call| { + (!call.id.is_empty()) + .then(|| memory.get(&call.id)) + .flatten() + }) + { + return raw.clone(); + } + canonical_tools(calls) +} + +fn canonical_tools(calls: &[ApiToolCall]) -> String { + let mut output = String::from("\n\n<｜DSML｜tool_calls>\n"); + for call in calls { + output.push_str("<｜DSML｜invoke name=\""); + output.push_str(&escape_attribute(&call.function.name)); + output.push_str("\">\n"); + match serde_json::from_str::(&call.function.arguments) { + Ok(Value::Object(arguments)) => { + for (name, value) in arguments { + output.push_str("<｜DSML｜parameter name=\""); + output.push_str(&escape_attribute(&name)); + let string = value.as_str(); + output.push_str(if string.is_some() { + "\" string=\"true\">" + } else { + "\" string=\"false\">" + }); + if let Some(value) = string { + output.push_str(&escape_parameter(value)); + } else { + output.push_str(&escape_json_parameter(&value.to_string())); + } + output.push_str("\n"); + } + } + _ => { + output.push_str("<｜DSML｜parameter name=\"arguments\" string=\"true\">"); + output.push_str(&escape_parameter(&call.function.arguments)); + output.push_str("\n"); + } + } + output.push_str("\n"); + } + output.push_str(""); + output +} + +fn final_response( + stream: &mut TcpStream, + state: &State, + request: ResponseOptions, + active: crate::runtime::ActiveGeneration, + id: &str, +) -> Result<(), (u16, String)> { + let output = wait_for_output(active)?; + let (content, calls) = parse_generated_tools(state, &output.message.content); + let finish = if calls.is_empty() { + output.finish_reason + } else { + "tool_calls" + }; + let mut message = json!({"role": "assistant", "content": content}); + if let Some(reasoning) = output.message.reasoning.filter(|value| !value.is_empty()) { + message["reasoning_content"] = Value::String(reasoning); + } + if !calls.is_empty() { + message["tool_calls"] = tool_calls_json(&calls); + } + let body = json!({ + "id": id, + "object": "chat.completion", + "created": unix_time(), + "model": request.model_id, + "choices": [{"index": 0, "message": message, "finish_reason": finish}], + "usage": usage_json(output.prompt_tokens, output.cached_tokens, output.completion_tokens), + }); + send_json(stream, 200, &body).map_err(|error| (500, error)) +} + +fn stream_response( + stream: &mut TcpStream, + state: &State, + request: ResponseOptions, + active: crate::runtime::ActiveGeneration, + id: &str, +) -> Result<(), (u16, String)> { + send_sse_headers(stream).map_err(|error| (500, error))?; + let role = chunk_json(id, &request.model_id, json!({"role": "assistant"}), None); + send_sse(stream, &role).map_err(|error| (500, error))?; + let mut buffered_content = String::new(); + let mut buffered_reasoning = String::new(); + let mut output = None; + while let Ok(event) = active.events.recv() { + match event { + GenerationEvent::Loading | GenerationEvent::Context { .. } => {} + GenerationEvent::Chunk { reasoning, content } => { + if reasoning { + buffered_reasoning.push_str(&content); + } else { + buffered_content.push_str(&content); + } + if request.has_tools && buffered_content.contains("") { + active.cancel.store(true, Ordering::Relaxed); + } + if !request.has_tools && !content.is_empty() { + let field = if reasoning { + "reasoning_content" + } else { + "content" + }; + let chunk = chunk_json(id, &request.model_id, json!({field: content}), None); + if send_sse(stream, &chunk).is_err() { + active.cancel.store(true, Ordering::Relaxed); + return Ok(()); + } + } + } + GenerationEvent::Finished(result) => { + match result { + Ok(result) => output = Some(result), + Err(error) => { + let event = json!({"error": {"message": error, "type": "server_error"}}); + let _ = send_sse(stream, &event); + let _ = stream.write_all(b"data: [DONE]\n\n"); + return Ok(()); + } + } + break; + } + } + } + let output = output.ok_or_else(|| (500, "The model runtime stopped unexpectedly.".into()))?; + let (content, calls) = parse_generated_tools(state, &output.message.content); + if request.has_tools { + if let Some(reasoning) = output + .message + .reasoning + .as_deref() + .filter(|value| !value.is_empty()) + { + send_sse( + stream, + &chunk_json( + id, + &request.model_id, + json!({"reasoning_content": reasoning}), + None, + ), + ) + .map_err(|error| (500, error))?; + } + if !content.is_empty() { + send_sse( + stream, + &chunk_json(id, &request.model_id, json!({"content": content}), None), + ) + .map_err(|error| (500, error))?; + } + if !calls.is_empty() { + send_sse( + stream, + &chunk_json( + id, + &request.model_id, + json!({"tool_calls": tool_calls_json(&calls)}), + None, + ), + ) + .map_err(|error| (500, error))?; + } + } + let finish = if calls.is_empty() { + output.finish_reason + } else { + "tool_calls" + }; + send_sse( + stream, + &chunk_json(id, &request.model_id, json!({}), Some(finish)), + ) + .map_err(|error| (500, error))?; + if request.include_usage { + let usage = json!({ + "id": id, + "object": "chat.completion.chunk", + "created": unix_time(), + "model": request.model_id, + "choices": [], + "usage": usage_json(output.prompt_tokens, output.cached_tokens, output.completion_tokens), + }); + send_sse(stream, &usage).map_err(|error| (500, error))?; + } + stream + .write_all(b"data: [DONE]\n\n") + .map_err(|error| (500, error.to_string())) +} + +fn wait_for_output( + active: crate::runtime::ActiveGeneration, +) -> Result { + let mut content = String::new(); + while let Ok(event) = active.events.recv() { + match event { + GenerationEvent::Chunk { + reasoning: false, + content: chunk, + } => { + content.push_str(&chunk); + if content.contains("") { + active.cancel.store(true, Ordering::Relaxed); + } + } + GenerationEvent::Finished(result) => { + return result.map_err(|error| (500, error)); + } + _ => {} + } + } + Err((500, "The model runtime stopped unexpectedly.".into())) +} + +fn parse_generated_tools(state: &State, text: &str) -> (String, Vec) { + let Some(start) = text.find("<｜DSML｜tool_calls>") else { + return (text.to_owned(), Vec::new()); + }; + let Some(relative_end) = text[start..].find("") else { + return (text.to_owned(), Vec::new()); + }; + let end = start + relative_end + "".len(); + let raw = &text[start..end]; + let mut calls = Vec::new(); + let mut cursor = 0; + while let Some(relative) = raw[cursor..].find("<｜DSML｜invoke name=\"") { + let name_start = cursor + relative + "<｜DSML｜invoke name=\"".len(); + let Some(name_end_relative) = raw[name_start..].find("\">") else { + return (text.to_owned(), Vec::new()); + }; + let name_end = name_start + name_end_relative; + let body_start = name_end + 2; + let Some(body_end_relative) = raw[body_start..].find("") else { + return (text.to_owned(), Vec::new()); + }; + let body_end = body_start + body_end_relative; + let arguments = parse_dsml_arguments(&raw[body_start..body_end]); + calls.push(ApiToolCall { + id: String::new(), + function: ApiFunction { + name: unescape_dsml(&raw[name_start..name_end]), + arguments: Value::Object(arguments).to_string(), + }, + }); + cursor = body_end + "".len(); + } + if calls.is_empty() { + return (text.to_owned(), Vec::new()); + } + let sequence = state + .sequence + .fetch_add(calls.len() as u64, Ordering::Relaxed); + for (index, call) in calls.iter_mut().enumerate() { + call.id = format!("call_{:032x}", sequence + index as u64 + 1); + } + if let Ok(mut memory) = state.tool_memory.lock() { + // ponytail: one process-local replay table; add LRU eviction if 100k live tool ids is measured insufficient. + if memory.len() >= 100_000 { + memory.clear(); + } + for call in &calls { + memory.insert(call.id.clone(), raw.to_owned()); + } + } + (text[..start].trim_end().to_owned(), calls) +} + +fn parse_dsml_arguments(body: &str) -> Map { + let mut arguments = Map::new(); + let mut cursor = 0; + while let Some(relative) = body[cursor..].find("<｜DSML｜parameter name=\"") { + let name_start = cursor + relative + "<｜DSML｜parameter name=\"".len(); + let Some(name_end_relative) = body[name_start..].find("\" string=\"") else { + break; + }; + let name_end = name_start + name_end_relative; + let flag_start = name_end + "\" string=\"".len(); + let Some(flag_end_relative) = body[flag_start..].find("\">") else { + break; + }; + let flag_end = flag_start + flag_end_relative; + let value_start = flag_end + 2; + let Some(value_end_relative) = body[value_start..].find("") else { + break; + }; + let value_end = value_start + value_end_relative; + let raw = unescape_dsml(&body[value_start..value_end]); + let value = if &body[flag_start..flag_end] == "true" { + Value::String(raw) + } else { + serde_json::from_str(&raw).unwrap_or(Value::String(raw)) + }; + arguments.insert(unescape_dsml(&body[name_start..name_end]), value); + cursor = value_end + "".len(); + } + arguments +} + +fn tool_calls_json(calls: &[ApiToolCall]) -> Value { + Value::Array( + calls + .iter() + .map(|call| { + json!({ + "id": call.id, + "type": "function", + "function": { + "name": call.function.name, + "arguments": call.function.arguments, + } + }) + }) + .collect(), + ) +} + +fn installed_endpoint_models(models_path: &std::path::Path) -> Vec { + model::installed_models(models_path) + .into_iter() + .filter(|model| *model == ModelChoice::DeepSeekV4Flash) + .collect() +} + +fn models_json(state: &State) -> Value { + let context = state + .preferences + .read() + .map_or(32_768, |preferences| preferences.context_tokens); + json!({ + "object": "list", + "data": installed_endpoint_models(&state.models_path) + .into_iter() + .map(|model| model_json(model, context)) + .collect::>() + }) +} + +fn model_json(model: ModelChoice, context: i32) -> Value { + json!({ + "id": model.id(), + "object": "model", + "created": 1767225600_i64, + "owned_by": "ds4.c", + "name": model.to_string(), + "context_length": context, + "top_provider": { + "context_length": context, + "max_completion_tokens": context, + "is_moderated": false + }, + "supported_parameters": [ + "tools", "tool_choice", "max_tokens", "temperature", "top_p", + "top_k", "min_p", "stop", "seed", "stream", "reasoning_effort" + ] + }) +} + +fn model_alias(id: &str) -> Option { + match id { + "deepseek-chat" | "deepseek-reasoner" => Some(ModelChoice::DeepSeekV4Flash), + "glm-5.2-chat" + | "glm-5.2-no-think" + | "glm-5.2-nothink" + | "glm-5.2-reasoner" + | "zai/glm-5.2" + | "zai/glm-5.2-chat" + | "zai/glm-5.2-reasoner" => Some(ModelChoice::Glm52), + _ => ModelChoice::from_id(id), + } +} + +fn content_text(value: &Value) -> String { + match value { + Value::String(text) => text.clone(), + Value::Array(parts) => parts + .iter() + .filter_map(|part| match part { + Value::String(text) => Some(text.as_str()), + Value::Object(object) => object.get("text").and_then(Value::as_str), + _ => None, + }) + .collect(), + _ => String::new(), + } +} + +fn escape_attribute(text: &str) -> String { + text.replace('&', "&") + .replace('<', "<") + .replace('>', ">") + .replace('"', """) +} + +fn escape_parameter(text: &str) -> String { + text.replace("", "</｜DSML｜parameter>") +} + +fn escape_json_parameter(text: &str) -> String { + text.replace("", "\\u003c/｜DSML｜parameter>") +} + +fn escape_tool_result(text: &str) -> String { + text.replace("", "</tool_result>") +} + +fn unescape_dsml(text: &str) -> String { + text.replace(""", "\"") + .replace(">", ">") + .replace("<", "<") + .replace("&", "&") +} + +fn usage_json(prompt: u32, cached: u32, completion: u32) -> Value { + let cached = cached.min(prompt); + json!({ + "prompt_tokens": prompt, + "completion_tokens": completion, + "total_tokens": prompt + completion, + "prompt_tokens_details": { + "cached_tokens": cached, + "cache_write_tokens": prompt - cached + } + }) +} + +fn chunk_json(id: &str, model: &str, delta: Value, finish: Option<&str>) -> Value { + json!({ + "id": id, + "object": "chat.completion.chunk", + "created": unix_time(), + "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}] + }) +} + +fn send_sse_headers(stream: &mut TcpStream) -> Result<(), String> { + stream + .write_all( + b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nCache-Control: no-cache\r\nConnection: close\r\n\r\n", + ) + .map_err(|error| error.to_string()) +} + +fn send_sse(stream: &mut TcpStream, value: &Value) -> Result<(), String> { + let body = serde_json::to_vec(value).map_err(|error| error.to_string())?; + stream + .write_all(b"data: ") + .map_err(|error| error.to_string())?; + stream.write_all(&body).map_err(|error| error.to_string())?; + stream.write_all(b"\n\n").map_err(|error| error.to_string()) +} + +fn send_json(stream: &mut TcpStream, code: u16, value: &Value) -> Result<(), String> { + let mut body = serde_json::to_vec(value).map_err(|error| error.to_string())?; + body.push(b'\n'); + send_response(stream, code, Some("application/json"), &body) +} + +fn send_error(stream: &mut TcpStream, code: u16, message: &str) -> Result<(), String> { + send_json( + stream, + code, + &json!({"error": {"message": message, "type": "invalid_request_error"}}), + ) +} + +fn send_response( + stream: &mut TcpStream, + code: u16, + content_type: Option<&str>, + body: &[u8], +) -> Result<(), String> { + let reason = match code { + 200 => "OK", + 204 => "No Content", + 400 => "Bad Request", + 404 => "Not Found", + 409 => "Conflict", + 500 => "Internal Server Error", + _ => "Error", + }; + let mut header = format!( + "HTTP/1.1 {code} {reason}\r\nContent-Length: {}\r\n", + body.len() + ); + if let Some(content_type) = content_type { + header.push_str("Content-Type: "); + header.push_str(content_type); + header.push_str("\r\n"); + } + header.push_str("Connection: close\r\n\r\n"); + stream + .write_all(header.as_bytes()) + .and_then(|()| stream.write_all(body)) + .map_err(|error| error.to_string()) +} + +fn read_request(stream: &mut TcpStream) -> Result { + let mut bytes = Vec::new(); + let header_end = loop { + if bytes.len() >= MAX_HEADER_BYTES { + return Err("HTTP headers are too large".into()); + } + let mut chunk = [0_u8; 4096]; + let read = stream.read(&mut chunk).map_err(|error| error.to_string())?; + if read == 0 { + return Err("bad HTTP request".into()); + } + bytes.extend_from_slice(&chunk[..read]); + if let Some(end) = find_header_end(&bytes) { + break end; + } + }; + let header = std::str::from_utf8(&bytes[..header_end]) + .map_err(|_| "HTTP headers are not UTF-8".to_owned())?; + let mut lines = header.lines(); + let request_line = lines.next().ok_or_else(|| "bad HTTP request".to_owned())?; + let mut parts = request_line.split_whitespace(); + let method = parts + .next() + .ok_or_else(|| "bad HTTP request".to_owned())? + .to_owned(); + let path = parts + .next() + .ok_or_else(|| "bad HTTP request".to_owned())? + .to_owned(); + let length = lines + .find_map(|line| { + line.split_once(':').and_then(|(name, value)| { + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + }) + .unwrap_or(0); + if length > MAX_BODY_BYTES { + return Err("HTTP body is too large".into()); + } + while bytes.len() < header_end + length { + let mut chunk = [0_u8; 8192]; + let read = stream.read(&mut chunk).map_err(|error| error.to_string())?; + if read == 0 { + return Err("incomplete HTTP body".into()); + } + bytes.extend_from_slice(&chunk[..read]); + } + Ok(HttpRequest { + method, + path: path.split('?').next().unwrap_or(&path).to_owned(), + body: bytes[header_end..header_end + length].to_vec(), + }) +} + +fn find_header_end(bytes: &[u8]) -> Option { + bytes + .windows(4) + .position(|window| window == b"\r\n\r\n") + .map(|position| position + 4) + .or_else(|| { + bytes + .windows(2) + .position(|window| window == b"\n\n") + .map(|position| position + 2) + }) +} + +fn unix_time() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +fn default_role() -> String { + "user".into() +} + +fn empty_arguments() -> String { + "{}".into() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn canonical_tool_replay_preserves_argument_types() { + let calls = vec![ApiToolCall { + id: "call_1".into(), + function: ApiFunction { + name: "shell".into(), + arguments: r#"{"command":"printf hi","timeout":30,"login":false}"#.into(), + }, + }]; + let rendered = canonical_tools(&calls); + assert!(rendered.contains("name=\"command\" string=\"true\">printf hi")); + assert!(rendered.contains("name=\"timeout\" string=\"false\">30")); + assert!(rendered.contains("name=\"login\" string=\"false\">false")); + } + + #[test] + fn generated_dsml_becomes_openai_tool_calls() { + let state = State { + generation: GenerationService::spawn().unwrap(), + preferences: Arc::new(RwLock::new(AppPreferences::default())), + models_path: PathBuf::new(), + cache_path: PathBuf::new(), + sequence: AtomicU64::new(0), + tool_memory: Mutex::new(HashMap::new()), + }; + let text = "done\n\n<｜DSML｜tool_calls>\n<｜DSML｜invoke name=\"shell\">\n<｜DSML｜parameter name=\"command\" string=\"true\">pwd\n\n"; + let (content, calls) = parse_generated_tools(&state, text); + assert_eq!(content, "done"); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].function.name, "shell"); + assert_eq!(calls[0].function.arguments, r#"{"command":"pwd"}"#); + } + + #[test] + fn content_arrays_match_ds4_text_projection() { + assert_eq!( + content_text(&json!([{"type": "text", "text": "one"}, " two"])), + "one two" + ); + } +} diff --git a/src/settings.rs b/src/settings.rs index 48a49f6..82eb644 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -443,6 +443,8 @@ impl GenerationPreferences { temperature: self.temperature.unwrap_or(1.0), top_p: self.top_p.unwrap_or(if glm { 0.95 } else { 1.0 }), min_p: self.min_p.unwrap_or(if glm { 0.0 } else { 0.05 }), + top_k: 0, + stops: Vec::new(), seed: self.seed, reasoning_mode: if self.reasoning_mode == ReasoningMode::Max && self.context_tokens < THINK_MAX_MIN_CONTEXT @@ -489,6 +491,8 @@ pub(crate) struct TurnSettings { pub(crate) temperature: f32, pub(crate) top_p: f32, pub(crate) min_p: f32, + pub(crate) top_k: i32, + pub(crate) stops: Vec, pub(crate) seed: Option, pub(crate) reasoning_mode: ReasoningMode, }