diff --git a/server/Cargo.lock b/server/Cargo.lock index bfd919a..58fea9e 100644 --- a/server/Cargo.lock +++ b/server/Cargo.lock @@ -395,6 +395,35 @@ version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" +[[package]] +name = "cookie" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a373e3602691c3cdea496d2f0ee5935151e6168fe87739483c463db1b2f2f87" +dependencies = [ + "percent-encoding", + "time", + "version_check", +] + +[[package]] +name = "cookie_store" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15b2c103cf610ec6cae3da84a766285b42fd16aad564758459e6ecf128c75206" +dependencies = [ + "cookie", + "document-features", + "idna", + "indexmap", + "log", + "serde", + "serde_derive", + "serde_json", + "time", + "url", +] + [[package]] name = "cpufeatures" version = "0.3.0" @@ -470,6 +499,15 @@ dependencies = [ "syn 3.0.4", ] +[[package]] +name = "document-features" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4b8a88685455ed29a21542a33abd9cb6510b6b129abadabdcef0f4c55bc8f61" +dependencies = [ + "litrs", +] + [[package]] name = "dunce" version = "1.0.5" @@ -743,6 +781,110 @@ dependencies = [ "tower-service", ] +[[package]] +name = "icu_collections" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa68d21081c4a05d5a901a1c62add574c77048b6a1c67be3b50ce0b60d4ca513" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d56e28588da92eee5c3201a6eff33fabdd49b62269c8938d4ff050ce4d900deb" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "12f9cf5f235641ed274641dd81c3f28d870e276763d0797aeeab72317b1c646f" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1563da1ed3e0b3bf3d74c9b85917ac9c56464d2f57242270c09c9e752f8021a0" + +[[package]] +name = "icu_properties" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e7ca276ad3145661a65914e6daf131ca5120cd3dcee8f8f3214b8875184a148" +dependencies = [ + "displaydoc", + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e590f038c1464a96894fd6d10127e90a8be4509f56ff7ecef851b15cee0b7caa" + +[[package]] +name = "icu_provider" +version = "2.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d27bbb9d3abbefac45d55f647c9de1d44aafcd1186eb91879afef17c396c3e73" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + [[package]] name = "if-addrs" version = "0.15.0" @@ -803,6 +945,18 @@ version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" +[[package]] +name = "litemap" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47d9d19d1d6efa0109d2f65ff4c85cddd50bd572e5a00127ab10987290bcefae" + +[[package]] +name = "litrs" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" + [[package]] name = "log" version = "0.4.34" @@ -982,6 +1136,15 @@ version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548" +[[package]] +name = "potential_utf" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d83eb9bc6d8e5cf568e7a1101d60ee05e81ed50ea106026f3d18deeb046d7661" +dependencies = [ + "zerovec", +] + [[package]] name = "powerfmt" version = "0.2.0" @@ -1301,6 +1464,12 @@ version = "0.9.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3763264f6b73151db08c50ff20d7d8a0b8796e021cdea7ceedad07b80155fa0e" +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + [[package]] name = "strsim" version = "0.11.1" @@ -1424,6 +1593,16 @@ dependencies = [ "time-core", ] +[[package]] +name = "tinystr" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1e27c91459209c2986af3dcf603a5a74a4368754ce37414f59acc971167f643" +dependencies = [ + "displaydoc", + "zerovec", +] + [[package]] name = "tokio" version = "1.53.1" @@ -1607,11 +1786,14 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "972d7902c8735f2695410b8aed7df6ed12a47394aa1c8d7af49f0497b731a94d" dependencies = [ "base64 0.23.1", + "cookie_store", "flate2", "log", "percent-encoding", "rustls", "rustls-pki-types", + "serde", + "serde_json", "ureq-proto", "utf8-zero", "webpki-roots", @@ -1629,12 +1811,30 @@ dependencies = [ "log", ] +[[package]] +name = "url" +version = "2.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", + "serde", +] + [[package]] name = "utf8-zero" version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8c0a043c9540bae7c578c88f91dda8bd82e59ae27c21baca69c8b191aaf5a6e" +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + [[package]] name = "utf8parse" version = "0.2.2" @@ -1756,6 +1956,12 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "writeable" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ad82d2a33cdc9674dc7465672f271e096168fcdbe0f799d9e6db8c5892679dc" + [[package]] name = "x509-parser" version = "0.18.1" @@ -1784,12 +1990,89 @@ dependencies = [ "time", ] +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", + "synstructure", +] + +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", + "synstructure", +] + [[package]] name = "zeroize" version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" +[[package]] +name = "zerotrie" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ea269c3bd32f0a32c321907a2ae912ba6f4649bb0fc764a15627e99a7095a3f" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb0464e17806c1d976d5cba29399c7f08e516e279e2ba493f63123b5fca67dd8" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34df6fc39dbd26ddc9c10e6a2984476e13acce22e64e4487636ef494369225da" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.4", +] + [[package]] name = "zmij" version = "1.0.23" diff --git a/server/Cargo.toml b/server/Cargo.toml index d96a9ba..c72c265 100644 --- a/server/Cargo.toml +++ b/server/Cargo.toml @@ -10,7 +10,7 @@ path = "src/main.rs" [dependencies] axum = { version = "0.8", features = ["json", "multipart"] } axum-server = { version = "0.8", features = ["tls-rustls"] } -tokio = { version = "1", features = ["rt-multi-thread", "macros", "net", "sync", "time", "process", "io-util"] } +tokio = { version = "1", features = ["rt-multi-thread", "macros", "net", "sync", "time", "process", "io-util", "signal"] } tokio-stream = "0.1" tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } @@ -46,7 +46,7 @@ rcgen = { version = "0.14", features = ["pem", "x509-parser"] } # Outbound HTTPS for the usage endpoint. A small blocking client fits an # every-few-minutes poll better than pulling in reqwest's tower stack; # rustls-backed like the rest of the TLS here. -ureq = "3" +ureq = { version = "3", features = ["json"] } # Direct dependency only to pick the process-level CryptoProvider in main: # ureq pulls rustls-with-ring, axum-server rustls-with-aws-lc-rs, and with # both in the graph rustls refuses to auto-select one. diff --git a/server/src/auth.rs b/server/src/auth.rs index d6ddc45..758eeb0 100644 --- a/server/src/auth.rs +++ b/server/src/auth.rs @@ -108,7 +108,12 @@ mod tests { fn manager_with_token(dir: &std::path::Path, token: &str) -> Arc { let manager = Arc::new( - SessionManager::new(dir.join("config.ron"), dir.join("sessions")).expect("manager"), + SessionManager::new( + dir.join("config.ron"), + dir.join("sessions"), + dir.join("models"), + ) + .expect("manager"), ); manager .set_tokens(vec![TokenEntry { diff --git a/server/src/config.rs b/server/src/config.rs index beb1a15..fdc4be1 100644 --- a/server/src/config.rs +++ b/server/src/config.rs @@ -15,6 +15,7 @@ //! JSONL file in its own directory (see `session::transcript`); this file //! holds only the metadata needed to list and respawn sessions. +use std::collections::BTreeMap; use std::path::{Path, PathBuf}; use anyhow::{Context, Result}; @@ -156,6 +157,10 @@ pub enum DriverKind { /// involved, and stays useful as a connectivity check that costs no /// tokens. Always available as a built-in provider. Echo, + /// A GGUF model served by llama.cpp's `llama-server` (see + /// `session::llama`). The model itself is one this machine has + /// downloaded; the provider's command is the server binary. + LlamaCpp, /// The Claude Code CLI over stream-json (see `session::claude`). /// Named for the CLI specifically: bare "claude" would suggest the /// credit-billed API, which this is not. @@ -200,6 +205,17 @@ pub struct SessionConfig { /// needs no change on this side. #[serde(skip_serializing_if = "Option::is_none")] pub permission_mode: Option, + /// Settings the driver interprets, chosen at spawn. + /// + /// Deliberately untyped here: what a temperature or a context size + /// means is the driver's business, and giving this schema a field per + /// driver is how a shared model starts carrying one dialect's + /// vocabulary. `permission_mode` above predates this and should fold + /// into it. A map rather than a list so the phone can send exactly + /// what a person changed, and BTreeMap so the file's order is stable + /// across writes. + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] + pub params: BTreeMap, /// Epoch seconds when the session was spawned. pub created: f64, } @@ -338,6 +354,7 @@ mod tests { model: None, cwd: None, permission_mode: None, + params: BTreeMap::new(), created: 1234.5, }], }; diff --git a/server/src/main.rs b/server/src/main.rs index 05fe403..2ae73ab 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -30,6 +30,7 @@ use std::sync::Arc; use anyhow::{Context, Result}; use clap::Parser; +use tokio::signal::unix::{SignalKind, signal}; use config::TokenEntry; use session::SessionManager; @@ -193,7 +194,7 @@ async fn main() -> Result<()> { .unwrap_or_else(|| data_home().join("models")); let models = Arc::new(models::ModelStore::new(models_dir.clone())); let manager = Arc::new( - SessionManager::new(config_path.clone(), data_dir) + SessionManager::new(config_path.clone(), data_dir, models_dir.clone()) .with_context(|| format!("failed to load {}", config_path.display()))?, ); tracing::info!("config: {}", config_path.display()); @@ -283,10 +284,21 @@ async fn main() -> Result<()> { let addr = SocketAddr::new(bind_ip, args.port); tracing::info!("serving https://{addr}"); - axum_server::bind_rustls(addr, tls_config) - .serve(app.into_make_service_with_connect_info::()) - .await - .context("TLS listener failed")?; + + // Stop the sessions' processes on the way out. Without this a signal + // kills this process and leaves its children running -- which for the + // Claude CLI is untidy and for a `llama-server` holding a model is + // gigabytes of memory belonging to nobody. Both signals, because + // systemd and OpenRC send TERM while a terminal sends INT. + let serving = axum_server::bind_rustls(addr, tls_config) + .serve(app.into_make_service_with_connect_info::()); + let mut terminate = signal(SignalKind::terminate()).context("listening for SIGTERM")?; + tokio::select! { + served = serving => served.context("TLS listener failed")?, + _ = terminate.recv() => tracing::info!("SIGTERM -- stopping sessions"), + _ = tokio::signal::ctrl_c() => tracing::info!("interrupted -- stopping sessions"), + } + manager.shutdown_all(); Ok(()) } diff --git a/server/src/routes.rs b/server/src/routes.rs index 0470277..55ca9f8 100644 --- a/server/src/routes.rs +++ b/server/src/routes.rs @@ -179,6 +179,11 @@ struct SpawnRequest { cwd: Option, #[serde(default)] permission_mode: Option, + /// Whatever the chosen driver understands -- llama.cpp's context size + /// and sampling, for instance. Opaque here on purpose: see + /// `SessionConfig::params`. + #[serde(default)] + params: std::collections::BTreeMap, } async fn spawn_session( @@ -193,6 +198,7 @@ async fn spawn_session( model: body.model, cwd: body.cwd, permission_mode: body.permission_mode, + params: body.params, }) .map_err(bad_request)?; tracing::info!( diff --git a/server/src/session/llama.rs b/server/src/session/llama.rs new file mode 100644 index 0000000..f0d8694 --- /dev/null +++ b/server/src/session/llama.rs @@ -0,0 +1,542 @@ +//! The llama.cpp driver: a `llama-server` process per session, spoken to +//! over its OpenAI-compatible HTTP API and translated into the common +//! event model. +//! +//! Two things make this shaped differently from the Claude driver, and +//! both are worth knowing before changing anything here. +//! +//! **It is spawned but not spoken to over stdio.** The process is started +//! through the same [`Transport`] as any other, and then reached over +//! HTTP on a loopback port. That is the case the transport's doc comment +//! flags: a remote llama-server would need its port forwarded as well as +//! its command wrapped, which is not built, so a session on an ssh host +//! is refused rather than silently talking to the wrong machine. +//! +//! **The server is stateless between requests**, so the whole +//! conversation goes with every one. It is rebuilt from the session's +//! transcript rather than kept in this struct, which is not tidiness: a +//! copy in driver memory is invisible to a second device and gone when +//! this process restarts, and the app is meant to work across devices. +//! The transcript is already the source of truth for everything else, and +//! this makes it the source of truth for the prompt too. +//! +//! That leaves the Claude driver as the odd one out rather than this one: +//! the CLI's own memory of a conversation is a cache in front of the same +//! transcript, not a second truth. Anyone tempted to "fix" the +//! inconsistency should resolve it in this direction. + +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; + +use anyhow::{Context, Result, bail}; +use serde::{Deserialize, Serialize}; +use serde_json::json; + +use super::driver::{Driver, Event, EventSink, ImageRef, SessionStatus}; +use super::transport::{Launch, Transport}; +use crate::config::{ProviderConfig, SessionConfig}; + +/// How long to wait for a model to load before giving up on it. Loading +/// is mostly disk, and a large quantised model on a cold cache is +/// genuinely slow, so this is generous -- the failure it exists for is a +/// server that will never answer, not one that is taking its time. +const READY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300); + +/// One turn in the conversation this driver keeps on the server's behalf. +#[derive(Debug, Clone, Serialize, Deserialize)] +struct Message { + role: String, + content: String, +} + +pub struct LlamaDriver { + sink: EventSink, + /// Where this session's own llama-server answers. + endpoint: String, + /// Where the conversation is read back from, one line per event. + transcript: PathBuf, + /// Sampling settings chosen at spawn, sent with every request. + sampling: serde_json::Map, + /// Set by [`Driver::interrupt`]; the streaming loop checks it between + /// chunks and stops, leaving what was generated in the transcript. + cancel: Arc, + /// Taken by shutdown to stop the server. + kill: Mutex>>, +} + +impl LlamaDriver { + pub fn spawn( + meta: &SessionConfig, + provider: &ProviderConfig, + transport: &Transport, + models_dir: &Path, + transcript: &Path, + sink: EventSink, + ) -> Result { + if !matches!(transport, Transport::Here) { + bail!( + "llama.cpp sessions can only run on this machine for now: the model is served \ + over HTTP, and forwarding that port to another host isn't built yet." + ); + } + let model = meta.model.as_deref().context( + "a llama.cpp session needs a model -- one of the downloaded ones, by its key", + )?; + let path = model_path(models_dir, model)?; + + let port = free_port().context("finding a port for llama-server")?; + let mut args: Vec = vec![ + "-m".into(), + path.to_string_lossy().into_owned(), + "--host".into(), + "127.0.0.1".into(), + "--port".into(), + port.to_string(), + ]; + // Settings that belong to the server because they decide how the + // model is loaded; the sampling ones ride on each request instead, + // so changing them later needn't reload anything. + for (key, flag) in [ + ("contextSize", "-c"), + ("gpuLayers", "-ngl"), + ("threads", "-t"), + ] { + if let Some(value) = meta.params.get(key) { + args.push(flag.to_string()); + args.push(value.clone()); + } + } + + let program = provider.command.as_deref().unwrap_or("llama-server"); + let launch = Launch::new(program, args, meta.cwd.as_deref()); + let mut child = transport.spawn(&launch)?; + tracing::info!( + "session {} running {program} for {model} on 127.0.0.1:{port}", + meta.id + ); + + let endpoint = format!("http://127.0.0.1:{port}"); + let (kill_tx, kill_rx) = tokio::sync::oneshot::channel::<()>(); + { + let sink = sink.clone(); + let label = format!("{} ({model})", provider.name); + tokio::spawn(async move { + let status = tokio::select! { + status = child.wait() => status.ok(), + _ = kill_rx => { + let _ = child.kill().await; + None + } + }; + if let Some(status) = status + && !status.success() + { + let _ = sink.send(Event::Error { + message: format!("{label} exited: {status}"), + }); + } + let _ = sink.send(Event::Status { + state: SessionStatus::Exited, + }); + }); + } + + // Loading is slow enough to be worth saying so: the session shows + // as running until the model is in memory, then goes idle, rather + // than looking ready and refusing the first message. + let _ = sink.send(Event::Status { + state: SessionStatus::Running, + }); + { + let sink = sink.clone(); + let endpoint = endpoint.clone(); + let model = model.to_string(); + std::thread::spawn(move || match wait_until_ready(&endpoint) { + Ok(()) => { + tracing::info!("{model} loaded and answering at {endpoint}"); + let _ = sink.send(Event::Status { + state: SessionStatus::Idle, + }); + } + Err(err) => { + let _ = sink.send(Event::Error { + message: format!("{model} never became ready: {err:#}"), + }); + let _ = sink.send(Event::Status { + state: SessionStatus::Exited, + }); + } + }); + } + + let mut sampling = serde_json::Map::new(); + for (key, field) in [ + ("temperature", "temperature"), + ("topP", "top_p"), + ("topK", "top_k"), + ("maxTokens", "max_tokens"), + ] { + if let Some(raw) = meta.params.get(key) + && let Ok(number) = raw.parse::() + { + sampling.insert(field.to_string(), json!(number)); + } + } + + Ok(Self { + sink, + endpoint, + transcript: transcript.to_path_buf(), + sampling, + cancel: Arc::new(AtomicBool::new(false)), + kill: Mutex::new(Some(kill_tx)), + }) + } +} + +impl Driver for LlamaDriver { + fn send_user_message(&self, text: String, images: Vec) { + if !images.is_empty() { + let _ = self.sink.send(Event::Error { + message: "this model can't be sent images".to_string(), + }); + } + let sink = self.sink.clone(); + let endpoint = self.endpoint.clone(); + let transcript = self.transcript.clone(); + let sampling = self.sampling.clone(); + let cancel = Arc::clone(&self.cancel); + cancel.store(false, Ordering::Relaxed); + + // Its own thread: the request blocks for as long as the model + // takes to generate, which is the whole point of streaming it. + std::thread::spawn(move || { + let _ = sink.send(Event::Status { + state: SessionStatus::Running, + }); + // Everything before this message, plus this message. Read + // rather than remembered, and `text` is appended here rather + // than waited for, because the manager's UserMessage event is + // still on its way to the transcript when this runs. + let mut messages = conversation(&transcript); + messages.push(Message { + role: "user".into(), + content: text, + }); + // The reply is not stored: the deltas below are the durable + // record, so the next turn reads back exactly what the phone + // was shown -- including a partial one that was interrupted. + if let Err(err) = generate(&endpoint, &messages, &sampling, &cancel, &sink) { + let _ = sink.send(Event::Error { + message: format!("{err:#}"), + }); + } + let _ = sink.send(Event::Status { + state: SessionStatus::Idle, + }); + }); + } + + fn answer_question(&self, _id: &str, _answer: &str) { + // Nothing here asks questions: this driver has no tools, so no + // permission prompts and no AskUserQuestion. + } + + fn interrupt(&self) { + self.cancel.store(true, Ordering::Relaxed); + } + + fn set_model(&self, _model: &str) { + let _ = self.sink.send(Event::Error { + message: "a llama.cpp session's model is fixed when it starts, because the server \ + loads one model into memory. Spawn another session to use a different one." + .to_string(), + }); + } + + fn compact(&self) { + let _ = self.sink.send(Event::Error { + message: "llama.cpp has no compaction. When the context fills, start a new session." + .to_string(), + }); + } + + fn shutdown(&self) { + self.cancel.store(true, Ordering::Relaxed); + if let Some(kill) = self.kill.lock().unwrap().take() { + let _ = kill.send(()); + } + } +} + +/// The conversation so far, folded out of the transcript. +/// +/// Consecutive `AssistantText` deltas are one assistant turn, closed by +/// the next user message -- which is also what makes an interrupted reply +/// come back as the partial text the phone actually saw, rather than +/// vanishing or being invented. +/// +/// This must stay a pure function of the transcript and must never +/// re-render earlier turns. llama.cpp caches the prompt prefix, so a +/// growing conversation reprocesses almost nothing -- but only while +/// every turn is byte-identical to last time. Changing how an old turn is +/// rendered silently reprocesses the whole history on every message. +fn conversation(path: &Path) -> Vec { + let Ok(events) = crate::session::transcript::read_after(path, 0) else { + return Vec::new(); + }; + let mut messages: Vec = Vec::new(); + let mut pending = String::new(); + for event in events { + match event.event { + Event::UserMessage { text } => { + if !pending.is_empty() { + messages.push(Message { + role: "assistant".into(), + content: std::mem::take(&mut pending), + }); + } + messages.push(Message { + role: "user".into(), + content: text, + }); + } + Event::AssistantText { delta } => pending.push_str(&delta), + _ => {} + } + } + if !pending.is_empty() { + messages.push(Message { + role: "assistant".into(), + content: pending, + }); + } + messages +} + +/// Where a model key resolves to on disk, refusing anything that climbs +/// out of the models directory -- the key arrives from a phone. +fn model_path(models_dir: &Path, key: &str) -> Result { + let mut path = models_dir.to_path_buf(); + for part in key.split('/') { + if part.is_empty() || part == "." || part == ".." { + bail!("\"{key}\" is not a model key this can resolve"); + } + path.push(part); + } + if !path.is_file() { + bail!("no downloaded model called \"{key}\" -- download it first"); + } + Ok(path) +} + +/// An unused loopback port, by asking the OS for one and letting it go. +/// +/// Racy in principle: something else could take it between here and +/// llama-server binding. In practice nothing on this machine is hunting +/// for ports, and the alternative -- parsing the port back out of the +/// server's log -- couples us to its output format for no real gain. +fn free_port() -> Result { + let listener = std::net::TcpListener::bind("127.0.0.1:0")?; + Ok(listener.local_addr()?.port()) +} + +/// Polls until the server says it is ready, or gives up. +fn wait_until_ready(endpoint: &str) -> Result<()> { + let deadline = std::time::Instant::now() + READY_TIMEOUT; + let url = format!("{endpoint}/health"); + loop { + if let Ok(response) = ureq::get(&url).call() + && response.status() == 200 + { + return Ok(()); + } + if std::time::Instant::now() > deadline { + bail!("gave up after {}s", READY_TIMEOUT.as_secs()); + } + std::thread::sleep(std::time::Duration::from_millis(250)); + } +} + +/// One streamed completion: posts the conversation, emits each delta as it +/// arrives. Emits rather than returns: the transcript those events land +/// in is what the next turn reads back, so there is nothing to hand up. +fn generate( + endpoint: &str, + messages: &[Message], + sampling: &serde_json::Map, + cancel: &AtomicBool, + sink: &EventSink, +) -> Result<()> { + let mut body = json!({ + "messages": messages, + "stream": true, + "stream_options": {"include_usage": true}, + }); + let map = body.as_object_mut().expect("built as an object"); + for (key, value) in sampling { + map.insert(key.clone(), value.clone()); + } + + let mut response = ureq::post(format!("{endpoint}/v1/chat/completions")) + .header("Content-Type", "application/json") + .send_json(&body) + .context("asking llama-server to generate")?; + + let reader = std::io::BufReader::new(response.body_mut().as_reader()); + let mut tokens = 0u64; + for line in std::io::BufRead::lines(reader) { + if cancel.load(Ordering::Relaxed) { + break; + } + let line = line.context("reading the generation stream")?; + // Server-sent events: the payload lines are the ones that matter, + // and blank lines separate events. + let Some(payload) = line.strip_prefix("data: ") else { + continue; + }; + if payload.trim() == "[DONE]" { + break; + } + let Ok(chunk) = serde_json::from_str::(payload) else { + continue; + }; + if let Some(usage) = chunk.get("usage").and_then(|u| u.get("total_tokens")) + && let Some(total) = usage.as_u64() + { + tokens = total; + } + let delta = chunk + .get("choices") + .and_then(|c| c.get(0)) + .and_then(|c| c.get("delta")) + .and_then(|d| d.get("content")) + .and_then(serde_json::Value::as_str) + .unwrap_or_default(); + if !delta.is_empty() { + let _ = sink.send(Event::AssistantText { + delta: delta.to_string(), + }); + } + } + if tokens > 0 { + let _ = sink.send(Event::UsageDelta { tokens }); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::session::transcript::Transcript; + + /// Writes a transcript the way the pump does, so the fold is tested + /// against the real file format rather than a hand-built vector. + fn transcript_with(events: &[Event]) -> (tempfile::TempDir, PathBuf) { + let dir = tempfile::tempdir().expect("tempdir"); + let path = dir.path().join("transcript.jsonl"); + let mut transcript = Transcript::open(&path).expect("open"); + for event in events { + transcript.append(event.clone(), 0.0).expect("append"); + } + (dir, path) + } + + #[test] + fn deltas_between_user_messages_are_one_assistant_turn() { + let (_dir, path) = transcript_with(&[ + Event::UserMessage { + text: "hello".into(), + }, + Event::AssistantText { + delta: "hi ".into(), + }, + Event::AssistantText { + delta: "there".into(), + }, + Event::Status { + state: SessionStatus::Idle, + }, + Event::UserMessage { + text: "again".into(), + }, + Event::AssistantText { + delta: "yes".into(), + }, + ]); + let messages = conversation(&path); + assert_eq!( + messages + .iter() + .map(|m| (m.role.as_str(), m.content.as_str())) + .collect::>(), + [ + ("user", "hello"), + ("assistant", "hi there"), + ("user", "again"), + ("assistant", "yes") + ], + ); + } + + #[test] + /// The interrupted case, which decides what a resumed conversation is + /// built from: whatever the phone was shown. The deltas that arrived + /// before the stop are in the transcript, so they are in the prompt -- + /// the model is never told it said something the user did not see, and + /// never has a turn silently dropped from under it. + fn an_interrupted_reply_stays_in_the_conversation() { + let (_dir, path) = transcript_with(&[ + Event::UserMessage { + text: "count".into(), + }, + Event::AssistantText { + delta: "one two".into(), + }, + Event::Status { + state: SessionStatus::Idle, + }, + ]); + let messages = conversation(&path); + assert_eq!(messages.len(), 2); + assert_eq!(messages[1].content, "one two"); + } + + #[test] + /// Events this driver does not produce must not disturb the fold: a + /// transcript can carry errors and status changes from a session that + /// was, say, relaunched. + fn other_events_are_not_part_of_the_conversation() { + let (_dir, path) = transcript_with(&[ + Event::Status { + state: SessionStatus::Running, + }, + Event::UserMessage { + text: "hello".into(), + }, + Event::Error { + message: "something went wrong".into(), + }, + Event::AssistantText { + delta: "still here".into(), + }, + Event::UsageDelta { tokens: 12 }, + ]); + let messages = conversation(&path); + assert_eq!(messages.len(), 2); + assert_eq!(messages[0].content, "hello"); + assert_eq!(messages[1].content, "still here"); + } + + #[test] + fn a_model_key_cannot_climb_out_of_the_models_directory() { + let dir = tempfile::tempdir().expect("tempdir"); + for attempt in ["../../etc/passwd", "unsloth/../../escape.gguf", ""] { + assert!( + model_path(dir.path(), attempt).is_err(), + "{attempt:?} should have been refused", + ); + } + } +} diff --git a/server/src/session/mod.rs b/server/src/session/mod.rs index 7ae466e..5774e1e 100644 --- a/server/src/session/mod.rs +++ b/server/src/session/mod.rs @@ -12,6 +12,7 @@ pub mod claude; pub mod driver; pub mod echo; +pub mod llama; pub mod transcript; pub mod transport; @@ -28,6 +29,7 @@ use crate::config::{Config, DriverKind, HostConfig, ProviderConfig, SessionConfi use claude::ClaudeDriver; use driver::{Driver, Event, ImageRef, SessionStatus}; use echo::EchoDriver; +use llama::LlamaDriver; use transcript::{SeqEvent, Transcript}; use transport::Transport; @@ -52,6 +54,8 @@ pub struct SpawnSpec { pub model: Option, pub cwd: Option, pub permission_mode: Option, + /// Driver-interpreted settings; see `SessionConfig::params`. + pub params: std::collections::BTreeMap, } /// One row of `GET /sessions`. @@ -124,6 +128,11 @@ impl LiveSession { self.driver.interrupt(); } + /// Stops this session's process without deleting anything. + pub fn shutdown(&self) { + self.driver.shutdown(); + } + pub fn compact(&self) { self.driver.compact(); } @@ -184,6 +193,10 @@ pub struct SessionManager { /// Per-session directories (transcript, attachments, produced images) /// live under here, each named by session id. data_dir: PathBuf, + /// Downloaded GGUF models, shared by every session that names one -- + /// which is why they live beside the session directories rather than + /// inside one. + models_dir: PathBuf, inner: RwLock, } @@ -193,7 +206,7 @@ impl SessionManager { /// crash-recovery story; the echo driver just starts fresh over the /// same transcript. Must be called inside a tokio runtime (each /// session spawns its event pump). - pub fn new(config_path: PathBuf, data_dir: PathBuf) -> Result { + pub fn new(config_path: PathBuf, data_dir: PathBuf, models_dir: PathBuf) -> Result { let config = Config::load(&config_path)?; crate::private::create_dir(&data_dir)?; @@ -204,7 +217,13 @@ impl SessionManager { // shows as exited rather than taking the whole server down // with it, and can still be deleted from the phone. match resolve(&config, meta).and_then(|(provider, host)| { - launch(meta.clone(), &provider, host.as_ref(), &data_dir) + launch( + meta.clone(), + &provider, + host.as_ref(), + &data_dir, + &models_dir, + ) }) { Ok(session) => { live.insert(meta.id.clone(), session); @@ -217,6 +236,7 @@ impl SessionManager { let manager = Self { config_path, data_dir, + models_dir, inner: RwLock::new(Inner { config, live }), }; manager.seed_providers()?; @@ -265,6 +285,22 @@ impl SessionManager { /// Every session, in config order, with live status joined in. A /// session that failed to relaunch reports as exited. + /// Stops every session's process, for a server that is going away. + /// + /// Drivers set `kill_on_drop`, which covers a session being deleted + /// while the server keeps running -- but not the server itself being + /// signalled, because nothing drops on the way out of a SIGTERM. That + /// leaves the children orphaned, which for a `llama-server` holding a + /// model means gigabytes of memory nobody owns any more. So exiting + /// asks them all to stop first. + pub fn shutdown_all(&self) { + let inner = self.inner.read().unwrap(); + for session in inner.live.values() { + session.shutdown(); + } + tracing::info!("stopped {} session process(es)", inner.live.len()); + } + pub fn sessions(&self) -> Vec { let inner = self.inner.read().unwrap(); inner @@ -346,10 +382,17 @@ impl SessionManager { model: spec.model.or_else(|| provider.models.first().cloned()), cwd: spec.cwd, permission_mode: spec.permission_mode, + params: spec.params, created: now(), }; - let session = launch(meta.clone(), &provider, host.as_ref(), &self.data_dir)?; + let session = launch( + meta.clone(), + &provider, + host.as_ref(), + &self.data_dir, + &self.models_dir, + )?; let mut candidate = inner.config.clone(); candidate.sessions.push(meta); if let Err(err) = candidate.save(&self.config_path) { @@ -455,6 +498,7 @@ fn launch( provider: &ProviderConfig, host: Option<&HostConfig>, data_dir: &Path, + models_dir: &Path, ) -> Result> { let dir = data_dir.join(&meta.id); crate::private::create_dir(&dir)?; @@ -471,6 +515,14 @@ fn launch( let driver: Box = match provider.kind { DriverKind::Echo => Box::new(EchoDriver::new(sink.clone())), + DriverKind::LlamaCpp => Box::new(LlamaDriver::spawn( + &meta, + provider, + &Transport::for_host(host), + models_dir, + &transcript_path, + sink.clone(), + )?), DriverKind::ClaudeCli => Box::new(ClaudeDriver::spawn( &meta, provider, @@ -534,6 +586,7 @@ mod tests { fn echo_spec() -> SpawnSpec { SpawnSpec { + params: Default::default(), provider: crate::config::ECHO_PROVIDER.to_string(), host: None, title: None, @@ -591,7 +644,12 @@ mod tests { let dir = tempfile::tempdir().expect("tempdir"); let config_path = dir.path().join("config.ron"); let data_dir = dir.path().join("sessions"); - let manager = SessionManager::new(config_path.clone(), data_dir.clone()).expect("manager"); + let manager = SessionManager::new( + config_path.clone(), + data_dir.clone(), + data_dir.join("models"), + ) + .expect("manager"); let info = manager.spawn_session(echo_spec()).expect("spawn"); // Untitled sessions are named after the provider that runs them. @@ -646,9 +704,12 @@ mod tests { #[tokio::test] async fn questions_round_trip_through_answer() { let dir = tempfile::tempdir().expect("tempdir"); - let manager = - SessionManager::new(dir.path().join("config.ron"), dir.path().join("sessions")) - .expect("manager"); + let manager = SessionManager::new( + dir.path().join("config.ron"), + dir.path().join("sessions"), + dir.path().join("models"), + ) + .expect("manager"); let info = manager.spawn_session(echo_spec()).expect("spawn"); let session = manager.session(&info.id).expect("live session"); @@ -685,7 +746,12 @@ mod tests { let config_path = dir.path().join("config.ron"); let data_dir = dir.path().join("sessions"); - let manager = SessionManager::new(config_path.clone(), data_dir.clone()).expect("manager"); + let manager = SessionManager::new( + config_path.clone(), + data_dir.clone(), + data_dir.join("models"), + ) + .expect("manager"); let info = manager.spawn_session(echo_spec()).expect("spawn"); let session = manager.session(&info.id).expect("live session"); let mut rx = session.subscribe(); @@ -699,7 +765,8 @@ mod tests { // A new manager over the same state: the session is back, and new // events continue the sequence rather than restarting it -- which // is what makes a phone's cursor survive a backend restart. - let manager = SessionManager::new(config_path, data_dir).expect("manager restart"); + let manager = SessionManager::new(config_path, data_dir.clone(), data_dir.join("models")) + .expect("manager restart"); let listed = manager.sessions(); assert_eq!(listed.len(), 1); assert_eq!(listed[0].id, info.id);