//! 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. //! //! **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 second half of what a transport is -- "run this" //! plus "reach this port" -- and it is what lets a session run on another //! machine: [`Transport::reserve_port`] hands back a port the server binds //! *there* and one that reaches it *here*, and the ssh connection carrying the //! command carries the tunnel between them. The far `llama-server` binds //! loopback only, so a model is never served to that machine's network. //! //! **The model file is the far machine's, not this one's.** A remote machine //! names its own models directory (`SshConfig::models_dir`, defaulting to where //! this backend keeps its downloads), and the file is looked for *there* -- so //! a session naming a model that machine does not have says so, instead of //! starting a server that will never load one. Downloading to another machine //! is not built; the model gets there however anything else does. //! //! **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. //! //! 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. Resolve any inconsistency in this direction. //! //! **It runs the agent loop itself.** `llama-server` hosts the tools and runs //! them (`--tools`, `GET`/`POST /tools`) but does not drive the conversation: //! a completion comes back with tool calls in it and stops. So the loop -- //! call, ask, run, feed the result back, ask again -- is here, which is also //! what puts the permission gate on this side, where a phone can answer it. //! [`tools`] is the catalog and the running; [`mcp`] is the half of it that //! this backend reaches rather than the model's machine. //! //! **Loading is a state, not a fast bit of starting.** A multi-gigabyte model //! takes a while to reach memory, and for that while the server refuses //! everything. It is [`SessionStatus::Loading`] on screen and a message sent //! into it waits rather than failing -- see [`Serving`]. use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Condvar, Mutex}; use anyhow::{Context, Result, bail}; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; use super::driver::{ AttachmentRef, Driver, Event, EventSink, QuestionOption, SessionStatus, Unqueued, }; use super::process; use super::transport::{Launch, Streams, Transport}; use crate::config::{ProviderConfig, SessionConfig}; mod mcp; mod tools; pub use tools::{DEFAULT_MODE as DEFAULT_PERMISSION_MODE, MODES as PERMISSION_MODES}; use mcp::McpServer; use tools::Tools; /// The sampling half of a session's settings, in the wire's own names. /// /// One function because two callers need the identical mapping: the spawn, and /// a later change through [`Driver::set_params`]. A value that will not parse /// as a number is left out rather than passed through -- `llama-server` would /// refuse the whole request for it, which would look like the session breaking /// rather than like one field being wrong. fn sampling_from( params: &std::collections::BTreeMap, ) -> serde_json::Map { 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) = params.get(key) && let Ok(number) = raw.parse::() { sampling.insert(field.to_string(), json!(number)); } } sampling } /// Exa's own hosted MCP server, which is what a llama session searches the web /// with. The address llama.cpp's web UI offers under "Exa" in its recommended /// servers, so a session here reaches the same thing that UI does. pub const EXA_MCP_URL: &str = "https://mcp.exa.ai/mcp"; /// The spawn parameter that turns speculative decoding off for a session whose /// model would otherwise use it. `"off"` and nothing else, because there is /// only one thing to say: the model either has a head or it does not, and this /// is the escape for a machine where drafting turns out not to pay. const SPECULATIVE: &str = "speculative"; /// The spawn parameter naming which built-in tools a session gets, as /// `llama-server`'s own comma-separated list. Absent is all of them, and /// `"none"` is the way to ask for a session that only talks. const TOOLS: &str = "tools"; /// How many times one message may go round the call-a-tool loop. /// /// A bound rather than a budget: a small model that has decided to read the /// same file for ever costs nothing but will never stop on its own, and the /// transcript fills with it. Reached, it is said in the transcript rather than /// hidden, because a turn that stopped for this reason looks exactly like one /// that finished. const MAX_STEPS: usize = 32; /// How long to wait for a model to load before giving up. 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. 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. /// /// The OpenAI chat shape, which is what `llama-server` renders through the /// model's own chat template. `content` is always present and sometimes empty /// rather than absent: a template reads it unconditionally, and a missing key /// renders the word "None" into the prompt on the ones that do. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] struct Message { role: String, content: String, /// What the assistant asked to run, in the wire's own shape so it goes /// back exactly as it came. #[serde(default, skip_serializing_if = "Vec::is_empty")] tool_calls: Vec, /// Which call this message is the result of. Only on `role: "tool"`. #[serde(default, skip_serializing_if = "Option::is_none")] tool_call_id: Option, } impl Message { fn new(role: &str, content: impl Into) -> Self { Self { role: role.to_string(), content: content.into(), tool_calls: Vec::new(), tool_call_id: None, } } fn result_of(call: &str, content: impl Into) -> Self { Self { tool_call_id: Some(call.to_string()), ..Self::new("tool", content) } } } /// One call the model asked for, assembled out of however many chunks it /// arrived in. #[derive(Debug, Clone, Default)] struct Call { id: String, name: String, /// JSON text, because that is what the wire carries and what has to go /// back unchanged; parsed once, at the point of running it. arguments: String, } impl Call { /// The call as it goes back into the conversation. /// /// Empty arguments are written as the empty object, not as the empty /// string the wire sent: a tool that takes none streams `""`, and a chat /// template rendering that back produces a call whose arguments are not /// JSON. fn wire(&self) -> Value { let arguments = if self.arguments.trim().is_empty() { "{}" } else { &self.arguments }; json!({ "id": self.id, "type": "function", "function": {"name": self.name, "arguments": arguments}, }) } /// The arguments as an object, or an empty one for a model that sent /// something that is not. Never an error: the tool is about to say what it /// makes of them, and it says it better than this could. fn arguments(&self) -> Value { serde_json::from_str(&self.arguments).unwrap_or_else(|_| json!({})) } } /// What one completion produced. struct Reply { text: String, calls: Vec, } /// Where this session's server has got to. /// /// Three states and not two: a model that will never load is neither loading /// nor ready, and a message sent to it has to be told something. Everything /// that needs the server goes through [`Shared::serving`], so there is one /// place that knows which of the three it is. enum Serving { /// Started, not answering yet. Anything sent now waits here. Loading, Ready { endpoint: String, tools: Arc, }, /// It exited, or never came up. Carries what to tell somebody, because by /// the time a message arrives the log that explained it is long gone from /// the screen. Failed(String), } /// Whether a turn is running, and the messages written during it -- each with /// the id of the `MessageQueued` that announced it, so the `UserMessage` can /// say which bubble it resolves. #[derive(Default)] struct Turns { running: bool, waiting: std::collections::VecDeque<(String, String)>, } /// What it takes to put this session on a different model. /// /// Kept whole rather than reduced to the two fields that change, because /// starting a server is one function and giving it a second set of inputs is /// how the two come to disagree about, say, which flags a session gets. struct Respawn { meta: SessionConfig, provider: ProviderConfig, transport: Transport, models_dir: PathBuf, } /// Everything the driver's own threads need, which is nearly all of it: a turn /// runs on a thread of its own and outlives any borrow of the driver. struct Shared { sink: EventSink, /// Where the conversation is read back from, one line per event. transcript: PathBuf, /// Where this session's process record lives, so [`Driver::stop`] can find /// the server it has to end. session_dir: PathBuf, /// Sampling settings, sent with every request. Behind a lock because they /// are changeable while the session runs: they ride on the next request, /// so unlike the server's own flags there is nothing to reload. sampling: Mutex>, /// The session's working directory, which is where its tools act. `None` /// leaves that to `llama-server`, which is the honest answer rather than a /// guess at one. cwd: Option, /// Set by [`Driver::interrupt`]; the streaming loop and the tool loop both /// check it, leaving what was produced in the transcript. cancel: AtomicBool, /// Whether a turn is running, and what is waiting behind it. /// /// One lock over both, because the two decide each other: with a flag and /// a queue apart, a turn ending can find the queue empty and a message /// arriving can find the flag set, in that order -- and the message is /// then in the queue with nothing running and nothing that will look at /// it again. A lost message, once in a while, with no sign of why. turns: Mutex, serving: Mutex, /// Woken whenever `serving` changes, which is what a waiting message /// waits on. settled: Condvar, /// MCP servers, connected once and kept across model changes: they are /// nothing to do with which model is loaded, and reconnecting on every /// switch would spend a round trip to learn the same list. mcp: Vec>>, /// How much this session asks before acting. Live, so a reader who has /// seen enough can stop being asked without restarting anything. mode: Mutex, /// Tools this session has been told to stop asking about, folded from the /// transcript at launch and added to as they are answered. In memory as /// well as on disk because an allowance given in this turn has to hold for /// the next call in it, and the transcript is written behind us. allowed: Mutex>, /// Questions a turn is blocked on, by question id. asked: Mutex>>>, } pub struct LlamaDriver { shared: Arc, respawn: Respawn, } impl LlamaDriver { /// Takes charge of this session's `llama-server`: the one already loaded if /// there is one, otherwise a new one. /// /// One entry point, for the reason `ClaudeDriver::launch` gives, expensive /// in a different currency: two servers holding the same model is twice the /// memory, and the second would bind a different port while the phone kept /// talking to the first. #[allow(clippy::too_many_arguments)] pub fn launch( meta: &SessionConfig, provider: &ProviderConfig, transport: &Transport, models_dir: &Path, transcript: &Path, session_dir: &Path, sink: EventSink, // llama.cpp has no notion of a Task call, so this is accepted only // to keep one shape across every driver's launch -- see // `SUBAGENTS.md`'s "Server layout". _subagents: Arc, ) -> Result { let model = meta.model.as_deref().context( "a llama.cpp session needs a model -- one of the downloaded ones, by its key", )?; let sampling = sampling_from(&meta.params); let driver = Self { shared: Arc::new(Shared { sink, transcript: transcript.to_path_buf(), session_dir: session_dir.to_path_buf(), sampling: Mutex::new(sampling), cwd: meta .cwd .as_ref() .map(|path| path.to_string_lossy().into_owned()), cancel: AtomicBool::new(false), turns: Mutex::new(Turns::default()), serving: Mutex::new(Serving::Loading), settled: Condvar::new(), // Connected before anything is started, and on this thread: // a spawn is already paying for a round trip to find the model, // and a session whose tools arrive after its first message is // one that answers that message with fewer tools than it has. mcp: connect_mcp(provider), mode: Mutex::new( meta.permission_mode .clone() .unwrap_or_else(|| tools::DEFAULT_MODE.to_string()), ), allowed: Mutex::new(allowances(transcript)), asked: Mutex::new(HashMap::new()), }), respawn: Respawn { meta: meta.clone(), provider: provider.clone(), transport: transport.clone(), models_dir: models_dir.to_path_buf(), }, }; driver.start(model)?; Ok(driver) } /// Puts a `llama-server` behind this session and starts watching for it to /// be ready -- adopting the one already there, or running a new one. /// /// Also the model-change path, which is what makes it take the model /// rather than read `respawn.meta`: a switch stops the old server and /// calls this, so the two ways a session comes to have a server are one /// piece of code and cannot drift. fn start(&self, model: &str) -> Result<()> { let shared = &self.shared; let Respawn { meta, provider, transport, models_dir, } = &self.respawn; let found = model_on(transport, models_dir, model)?; // Loading is slow enough to be worth its own state: the session shows // as loading until the model is in memory, rather than looking ready // and refusing the first message. *shared.serving.lock().unwrap() = Serving::Loading; let _ = shared.sink.send(Event::Status { state: SessionStatus::Loading, }); // Already loaded and still running: keep talking to it. The health poll // below confirms it is really answering, so adopting a pid whose server // has wedged still reports as a failure rather than as a session that // silently never replies. let endpoint = if let Some(process::Record { detail: process::Detail::Http { port }, pid, .. }) = process::live(&shared.session_dir) { tracing::info!( "session {} reattaching to the llama-server it left loaded (pid {pid}, port {port})", meta.id ); format!("http://127.0.0.1:{port}") } else { spawn_server( meta, provider, transport, &found, model, &shared.session_dir, )? }; let shared = Arc::clone(shared); let model = model.to_string(); std::thread::spawn(move || { let settled = match wait_until_ready(&endpoint, &shared.session_dir) .and_then(|()| Tools::discover(&endpoint, shared.mcp.clone())) { Ok(tools) => { tracing::info!( "{model} loaded and answering at {endpoint} with {} tools", tools.offered().map_or(0, |offered| offered.len()), ); // Asked now rather than carried from the spawn flags: a // session that named no context size gets the model's // own, which only the loaded server knows, and one that // named an impossible one gets whatever it settled for. // Either way this is the measurement rather than the // request. if let Some(window) = context_window(&endpoint) { shared.emit(Event::ContextWindow { tokens: window }); } Serving::Ready { endpoint: endpoint.clone(), tools: Arc::new(tools), } } Err(err) => { let why = format!("{model} never became ready: {err:#}"); let _ = shared.sink.send(Event::Error { message: why.clone(), }); let _ = shared.sink.send(Event::Status { state: SessionStatus::Exited, }); process::clear(&shared.session_dir); Serving::Failed(why) } }; let ready = matches!(settled, Serving::Ready { .. }); shared.settle(settled); if ready { let _ = shared.sink.send(Event::Status { state: SessionStatus::Idle, }); watch(shared.session_dir.clone(), shared.sink.clone()); } }); Ok(()) } } impl Shared { /// Blocks until the server can be spoken to, and says what to talk to. /// /// The whole of what [`SessionStatus::Loading`] means in practice: a /// message that arrives during a load is held here rather than refused, /// which is the thing a phone could not otherwise do anything about -- /// the reader cannot see that the model is still coming off disk, and /// retrying until it works is not an interface. fn await_ready(&self) -> Result<(String, Arc)> { let serving = self .serving .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); let serving = self .settled .wait_while(serving, |serving| matches!(serving, Serving::Loading)) .unwrap_or_else(std::sync::PoisonError::into_inner); match &*serving { Serving::Ready { endpoint, tools } => Ok((endpoint.clone(), Arc::clone(tools))), Serving::Failed(why) => bail!("{why}"), // `wait_while` does not return while this holds. Serving::Loading => unreachable!("waited out of Loading"), } } /// Leaves [`Serving::Loading`] for whatever it turned out to be, waking /// everything waiting. /// /// The write and the notification are one hold of the lock. Released /// between them, a waiter could read `Loading` and go to sleep in the gap /// -- and no second notification is coming, because the whole point of /// this state is that it is left once. fn settle(&self, settled: Serving) { let mut serving = self .serving .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); *serving = settled; self.settled.notify_all(); } /// Releases every turn blocked on a permission, as a refusal. /// /// Three callers, all of them the session being taken away from under a /// question: stopped, interrupted, or detached because this server is /// going away. Without it the turn's thread waits on a channel nobody will /// ever send to, and the question card stays on screen on every device -- /// only an `Answered` closes one, which is why this records one rather /// than quietly dropping the question. Answered as a refusal because that /// is what happened: the call did not run. fn abandon_questions(&self) { for (id, answer) in std::mem::take( &mut *self .asked .lock() .unwrap_or_else(std::sync::PoisonError::into_inner), ) { let _ = answer.send(vec![tools::REFUSE.to_string()]); self.emit(Event::Answered { id, answers: vec![tools::REFUSE.to_string()], }); } } fn emit(&self, event: Event) { let _ = self.sink.send(event); } } /// Runs a `llama-server` for this session and records it, returning where it /// is reached from here. /// /// Split out of [`LlamaDriver::start`] because adopting one and starting one /// share everything after "there is a server at this address" and nothing /// before it. fn spawn_server( meta: &SessionConfig, provider: &ProviderConfig, transport: &Transport, found: &Model, model: &str, session_dir: &Path, ) -> Result { // Where it listens on its own machine, and where that is reached // from here -- the same number when that machine is this one. let forward = transport .reserve_port() .context("finding a port for llama-server")?; let mut args: Vec = vec![ "-m".into(), found.path.clone(), // Loopback there, whichever machine there is: what reaches it // from outside that machine is the ssh tunnel and nothing // else. "--host".into(), "127.0.0.1".into(), "--port".into(), forward.there.to_string(), // The built-in agent tools -- read, search, edit, shell. All of them // unless the session says otherwise, because whether a particular call // should happen is the permission gate's question rather than a flag's. // They run on the machine serving the model, which is the machine the // files are on. // // One slot, not the four `llama-server` picks on its own. A session is // one conversation making one request at a time -- the driver holds a // second message until the turn ends -- so the other three are context // this session could have been given and was not. // // It is also what decides whether the MTP head below is worth having. // Measured 2026-09-19 on the 27B here: 41.5 tok/s plain at any slot // count, **61.4** with the head at one slot, and **28** with the head // at four. Speculation against a split KV cache is slower than not // speculating at all, which is a much bigger effect than the head // itself and reads exactly like the head being broken. "-np".into(), "1".into(), ]; // Settable because it is not free: the definitions of all seven are around // 2,000 tokens of every prompt -- measured at 2,191 against 1,322 for two // of them -- which on a small context window is a quarter of it spent // before anything is said. // // "none" omits the flag rather than passing it on: `--tools none` is // `tools setup failed: unknown tool "none"` and a server that exits, since // the argument is a list of tool names and no-tools is what having no flag // means. match meta.params.get(TOOLS).map(|chosen| chosen.trim()) { Some("none") => {} chosen => { args.push("--tools".into()); args.push( chosen .filter(|c| !c.is_empty()) .unwrap_or("all") .to_string(), ); } } // A model that carries a multi-token-prediction head drafts with it, which // is most of a 50% speed-up for free -- the tensors are in the file // whether or not they are used, and without the flag `llama-server` says // "unused tensor blk.N.nextn.* -- ignoring" and leaves them there. // // Conditional because it cannot be otherwise: asked for on a model without // one, `llama-server` **exits** ("context type MTP requested but model // doesn't contain MTP layers"), which is a session that never starts. The // answer comes from the file itself -- see `Model::mtp`. if found.mtp && meta.params.get(SPECULATIVE).map(String::as_str) != Some("off") { args.push("--spec-type".into()); args.push("draft-mtp".into()); } // 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"), // How far ahead the draft head guesses. Not defaulted here: 2 measured // 7% faster than llama.cpp's 3 on this machine's GPU, once, which is // a reason to make the knob reachable and not a reason to move it for // everybody. ("specDraftNMax", "--spec-draft-n-max"), ] { if let Some(value) = meta.params.get(key) { args.push(flag.to_string()); args.push(value.clone()); } } let program = provider.program(); let launch = Launch::new(program, args, meta.cwd.as_deref()).reaching(forward); // Its output goes to files, not pipes. Not only so the process can // outlive this server: nothing ever read those pipes, so a chatty // llama-server filled the 64 KB buffer and blocked mid-load with no sign // of why. let child = transport.spawn( &launch, Streams::Detached { stdin: std::process::Stdio::null(), stdout: log_file(&session_dir.join(SERVER_LOG))?.into(), stderr: log_file(&session_dir.join(SERVER_LOG))?.into(), }, )?; let pid = child .id() .context("llama-server exited before it could be recorded")?; tracing::info!( "session {} running {program} for {model} {} on 127.0.0.1:{} there, \ reached at 127.0.0.1:{} here, as pid {pid}", meta.id, transport.describe(), forward.there, forward.here, ); // Reaped so it does not become a zombie while this server is still its // parent; the health poll and the record are what say whether the // session is alive, because after a restart there is no `Child` to ask. tokio::spawn(async move { let mut child = child; let _ = child.wait().await; }); // The *near* port, because that is the one anything reaching this // server has to dial -- including a later run of this backend, // which adopts the record without knowing which machine the server // is on. For a remote session the recorded pid is the ssh // client's, which is the process this machine owns and which holds // the tunnel open for exactly as long as the far server lives. let record = process::Record::of(pid, process::Detail::Http { port: forward.here }) .context("llama-server was gone before its start time could be read")?; process::write(session_dir, &record); Ok(format!("http://127.0.0.1:{}", forward.here)) } /// Connects to every MCP server this provider names, dropping the ones that /// would not answer. /// /// A failure here is logged and survived rather than failing the spawn: a /// machine with no route out should still get a session with the tools that /// do not need one, and the alternative -- refusing to start -- makes a web /// search Exa happens to be down for into a session that cannot be created. /// What must not happen is silence, so it is said in the transcript too: /// the reader is about to be given a session whose web search is missing, and /// nothing else on the screen would say why. fn connect_mcp(provider: &ProviderConfig) -> Vec>> { provider .mcp_servers .iter() .filter_map( |configured| match McpServer::connect(&configured.name, &configured.url) { Ok(server) => { tracing::info!( "MCP server {} at {} offers {}", configured.name, configured.url, server .tools() .iter() .map(|tool| tool.qualified.as_str()) .collect::>() .join(", "), ); Some(Arc::new(Mutex::new(server))) } Err(err) => { tracing::warn!("MCP server {} is not available: {err:#}", configured.name); None } }, ) .collect() } /// Tools this session has already been told to stop asking about. /// /// Folded out of the transcript rather than stored beside it, for the reason /// the conversation is: this driver keeps nothing that a second device or a /// restarted backend could not see. What makes it foldable is that the answer /// carries the tool's name -- see [`tools::ALWAYS_PREFIX`] -- so one pass over /// the answers is the whole set, with no need to pair each one back to the /// question and the call it was about. fn allowances(transcript: &Path) -> std::collections::HashSet { let Ok(events) = crate::session::transcript::read_after(transcript, 0) else { return std::collections::HashSet::new(); }; events .into_iter() .filter_map(|entry| match entry.event { Event::Answered { answers, .. } => Some(answers), _ => None, }) .flatten() .filter_map(|answer| { answer .strip_prefix(tools::ALWAYS_PREFIX) .map(str::to_string) }) .collect() } /// Where llama-server's own output goes. One file for both streams: it is /// diagnostics nobody parses, and interleaving them is how it reads in a /// terminal anyway. const SERVER_LOG: &str = "llama-server.log"; /// How often a loaded server is checked for still being there. Slower than the /// Claude driver's stdout poll because nothing is waiting on it: this only has /// to notice a server that has gone. const WATCH_INTERVAL: std::time::Duration = std::time::Duration::from_secs(2); /// An owner-only log opened for appending, so the two streams pointed at /// it do not overwrite each other and a reattach keeps what came before. fn log_file(path: &Path) -> Result { use std::os::unix::fs::OpenOptionsExt; std::fs::OpenOptions::new() .create(true) .append(true) .mode(0o600) .open(path) .with_context(|| format!("opening {}", path.display())) } /// Reports the server going away, for as long as the session is there to report /// it to. /// /// Polled rather than waited on, for the reason the Claude driver gives: after a /// restart this server is not the process's parent, so liveness has to be a /// question asked of the record -- and asking it two different ways is how the /// two answers come to disagree. fn watch(session_dir: PathBuf, sink: EventSink) { std::thread::spawn(move || { loop { std::thread::sleep(WATCH_INTERVAL); match process::recorded(&session_dir) { Some((_, process::Liveness::Alive)) => {} // Nothing recorded means the session was stopped or deleted // deliberately, and whoever did that has already said so. None => return, Some((_, process::Liveness::Dead)) => { if !process::stopping(&session_dir) { let _ = sink.send(Event::Error { message: "llama-server exited".to_string(), }); } let _ = sink.send(Event::Status { state: SessionStatus::Exited, }); process::clear(&session_dir); return; } Some((_, process::Liveness::Unknown)) => { let _ = sink.send(Event::Status { state: SessionStatus::Unknown, }); } } if sink.is_closed() { return; } } }); } impl LlamaDriver { /// Takes the next waiting message, if any, and runs it. /// /// Called at the end of every turn as well as at the start of one, which /// is what drains the queue: a turn that ends with something waiting /// starts the next immediately rather than leaving the session idle with /// a bubble on screen. fn take_next(shared: &Arc) { let next = { let mut turns = shared.turns.lock().unwrap(); match turns.waiting.pop_front() { Some(next) => Some(next), None => { turns.running = false; None } } }; if let Some((id, text)) = next { Self::run_turn(shared, Some(id), text); } } /// One message, from being read to the session going idle. /// /// Runs on a thread of its own: the request blocks for as long as the /// model takes to generate, and a tool call inside it blocks for as long /// as the tool takes, which for a shell command is unbounded by anything /// this server knows. fn run_turn(shared: &Arc, queued: Option, text: String) { let shared = Arc::clone(shared); std::thread::spawn(move || { shared.cancel.store(false, Ordering::SeqCst); // The message goes into the transcript here, at the moment it is // read -- see `MessageTaken`. `queued` names the bubble this // resolves, and is `None` for one that never waited. shared.emit(Event::MessageTaken { id: queued, text: text.clone(), // Never any: this driver refuses them where they arrive, so // nothing is carried this far. attachments: Vec::new(), }); // Waits out a model still coming off disk rather than failing. // The status stays `Loading` while it does, which is the whole // difference from a session that is thinking. match shared.await_ready() { Ok((endpoint, tools)) => { shared.emit(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 message's own // transcript entry is still on its way when this runs. let mut messages = conversation(&shared.transcript); messages.push(Message::new("user", text)); if let Err(err) = converse(&shared, &endpoint, &tools, messages) { shared.emit(Event::Error { message: format!("{err:#}"), }); } shared.emit(Event::Status { state: SessionStatus::Idle, }); } // No status here, either side of the error. The server this // was waiting for is gone, and whoever established that has // already said `exited` -- saying `idle` over it would take // away the Start button and claim the session was waiting for // a person. Err(err) => shared.emit(Event::Error { message: format!("{err:#}"), }), } Self::take_next(&shared); }); } } /// The loop: generate, run what was asked for, generate again. /// /// `messages` is carried rather than re-read from the transcript each time /// round, which is the one place this driver's "fold it out of the record" /// rule has to bend. The record is written behind us -- the pump appends what /// the sink was sent -- so re-reading mid-turn would ask the model to act on a /// call whose result had not landed in the file yet. The events emitted here /// are exactly what `conversation` folds back, so the next turn reads the same /// thing this one built. fn converse( shared: &Arc, endpoint: &str, tools: &Tools, mut messages: Vec, ) -> Result<()> { for _ in 0..MAX_STEPS { if shared.cancel.load(Ordering::SeqCst) { return Ok(()); } // Read per call rather than per turn, so a sampling change made while // a long turn is running reaches the rest of it. let sampling = shared.sampling.lock().unwrap().clone(); let reply = generate(endpoint, &messages, tools, &sampling, shared)?; let calls = reply.calls; messages.push(Message { tool_calls: calls.iter().map(Call::wire).collect(), ..Message::new("assistant", reply.text) }); if calls.is_empty() { return Ok(()); } for call in &calls { let output = run_call(shared, tools, call); messages.push(Message::result_of(&call.id, output)); } } // Said rather than left as a turn that simply stopped: a reply that ends // here and one that ends because the model was finished look identical on // screen, and only one of them is worth sending the same message again // about. shared.emit(Event::Error { message: format!( "stopped after {MAX_STEPS} tool calls in one turn. Send another message to carry on, \ or ask for something narrower." ), }); Ok(()) } /// One tool call: announced, asked about if it has to be, run, and reported. /// /// Always returns something for the model to read, including when it was /// refused or interrupted. A call with no result is a conversation a chat /// template cannot render -- see [`tools::UNFINISHED`] -- so there is no path /// out of here that leaves one. fn run_call(shared: &Arc, tools: &Tools, call: &Call) -> String { let arguments = call.arguments(); shared.emit(Event::ToolStart { id: call.id.clone(), tool: call.name.clone(), input: arguments.clone(), }); let output = if !tools.knows(&call.name) { // The model invented one. Told plainly, because the alternative is a // silent empty result it reads as the tool having done nothing. format!( "There is no tool called {}. Use one of the tools you were given.", call.name ) } else if shared.cancel.load(Ordering::SeqCst) { tools::UNFINISHED.to_string() } else if !permitted(shared, call) { tools::REFUSED.to_string() } else { match tools.execute(&call.name, &arguments, shared.cwd.as_deref()) { Ok(output) => output, // Reaching the tool failed, which is this server's problem and // not the model's work going wrong -- but the model is still what // has to carry on, so it is told in the result rather than only // in the log. Err(err) => format!("This tool could not be run: {err:#}"), } }; shared.emit(Event::ToolEnd { id: call.id.clone(), output: output.clone(), }); output } /// Whether this call may go ahead, asking whoever is reading if it has to. /// /// Blocks the turn while the question is out, which is what the question is /// for. `AwaitingInput` while it waits, so every screen showing this session /// says it wants something. fn permitted(shared: &Arc, call: &Call) -> bool { if shared.mode.lock().unwrap().as_str() == "bypassPermissions" || shared.allowed.lock().unwrap().contains(&call.name) { return true; } let id = super::random_hex(); let (answer, wait) = std::sync::mpsc::channel(); shared.asked.lock().unwrap().insert(id.clone(), answer); let always = format!("{}{}", tools::ALWAYS_PREFIX, call.name); shared.emit(Event::Question { id: id.clone(), prompt: format!("Run {}?", call.name), // No header: this is about the call drawn directly above it, and // naming the tool twice reads as two different things. header: None, options: vec![ QuestionOption::plain(tools::ALLOW_ONCE), QuestionOption { label: always.clone(), description: Some(format!( "Stop asking about {} for the rest of this session.", call.name )), preview: None, }, QuestionOption::plain(tools::REFUSE), ], multi_select: false, // The call it is permission for, so a phone draws the ask on the // tool's own row rather than as a card repeating its input. about: Some(call.id.clone()), }); shared.emit(Event::Status { state: SessionStatus::AwaitingInput, }); let answers = wait.recv().unwrap_or_default(); shared.asked.lock().unwrap().remove(&id); shared.emit(Event::Status { state: SessionStatus::Running, }); if answers.contains(&always) { shared.allowed.lock().unwrap().insert(call.name.clone()); return true; } answers.iter().any(|answer| answer == tools::ALLOW_ONCE) } impl Driver for LlamaDriver { fn send_user_message(&self, text: String, attachments: Vec) { if !attachments.is_empty() { self.shared.emit(Event::Error { message: "this model can't be sent attachments or files".to_string(), }); } // A message written during a turn waits for it, rather than starting a // second conversation against the same server. Nothing was queued here // until tools arrived and turns grew long enough for it to matter -- // two turns interleaving their deltas into one transcript is what that // looked like. { let mut turns = self.shared.turns.lock().unwrap(); if turns.running { let id = super::random_hex(); turns.waiting.push_back((id.clone(), text.clone())); drop(turns); self.shared.emit(Event::MessageQueued { id, text, // Refused above, so there are none to wait with it. Said // as an empty list rather than the caller's, which would // draw a thumbnail on a bubble whose message will arrive // without it. attachments: Vec::new(), }); return; } turns.running = true; } Self::run_turn(&self.shared, None, text); } fn unqueue(&self, id: &str) -> Unqueued { let mut turns = self.shared.turns.lock().unwrap(); let Some(at) = turns.waiting.iter().position(|(waiting, ..)| waiting == id) else { return Unqueued::Unknown; }; turns.waiting.remove(at); drop(turns); self.shared .emit(Event::MessageDropped { id: id.to_string() }); Unqueued::Dropped } fn between_turns(&self) -> bool { !self.shared.turns.lock().unwrap().running } fn answer_question(&self, id: &str, answers: &[String]) { let waiting = self.shared.asked.lock().unwrap().remove(id); match waiting { Some(answer) => { let _ = answer.send(answers.to_vec()); } None => self.shared.emit(Event::Error { message: format!("no question {id} is awaiting an answer"), }), } } fn interrupt(&self) { self.shared.cancel.store(true, Ordering::SeqCst); self.shared.abandon_questions(); } // Nothing to forward: this process has no notion of what the conversation // is called, and the rename has already happened where the name lives. fn set_title(&self, _title: &str) {} /// Takes new settings: the sampling half now, and says so about the rest. /// /// The split is what [`crate::config::ParamSpec::restart`] describes, and /// it is said out loud rather than left to the screen, because the screen /// can only say what a setting *usually* does -- this is the one place /// that knows whether this session's server was started with the old /// value. A session already stopped needs no such note: its next start /// will read all of them. fn set_params(&self, params: &std::collections::BTreeMap) { *self.shared.sampling.lock().unwrap() = sampling_from(params); // Only the settings that actually differ from what this session's // server was started with. Listing every restart-only one on every // save would be a wall of text about nothing having changed. let waiting: Vec<&str> = crate::config::DriverKind::LlamaCpp .params() .iter() .filter(|spec| { spec.restart && params.get(spec.key) != self.respawn.meta.params.get(spec.key) }) .map(|spec| spec.label) .collect(); // Nothing to say to a session with no server: its next start reads all // of them, which is what the note would have been asking for. let running = matches!(&*self.shared.serving.lock().unwrap(), Serving::Ready { .. }); if !waiting.is_empty() && running { let one = waiting.len() == 1; self.shared.emit(Event::Error { message: format!( "{} {} saved. {} when this session's server next starts -- stop and start \ the session, or change its model, to load {} now.", waiting.join(", "), if one { "is" } else { "are" }, if one { "It takes effect" } else { "They take effect" }, if one { "it" } else { "them" }, ), }); } } fn set_permission_mode(&self, mode: &str) { if !tools::MODES.contains(&mode) { self.shared.emit(Event::Error { message: format!( "a llama.cpp session has no \"{mode}\" mode -- it is one of {}.", tools::MODES.join(" or "), ), }); return; } *self.shared.mode.lock().unwrap() = mode.to_string(); // What it is *set to*, which is the only thing a phone acts on. See // `Event::Settings`. self.shared.emit(Event::Settings { model: None, permission_mode: Some(mode.to_string()), }); } /// Puts this session on a different model, by loading one. /// /// A `llama-server` holds exactly one model, so this stops the one it has /// and starts another -- which costs a load and nothing else. The /// conversation survives it because the conversation was never in the /// server: it is folded out of the transcript on the next message, and the /// new model is given the same history the old one had. /// /// What is lost is the prompt cache, so the next turn reprocesses the whole /// conversation. That is exactly what the phone warns about before /// switching, and it is the same cost the other drivers pay for the same /// thing. fn set_model(&self, model: &str) { if self.shared.turns.lock().unwrap().running { self.shared.emit(Event::Error { message: "this session is mid-turn -- stop it first, then change the model." .to_string(), }); return; } self.shared.cancel.store(true, Ordering::SeqCst); if let Some(record) = process::live(&self.shared.session_dir) { process::stop(&record, process::STOP_GRACE); } process::clear(&self.shared.session_dir); self.shared.cancel.store(false, Ordering::SeqCst); match self.start(model) { // Reported when it is true and not before: `start` has put the // session into `Loading`, and the model it is loading is this one. Ok(()) => self.shared.emit(Event::Settings { model: Some(model.to_string()), permission_mode: None, }), Err(err) => { let why = format!("couldn't load {model}: {err:#}"); self.shared.emit(Event::Error { message: why.clone(), }); self.shared.emit(Event::Status { state: SessionStatus::Exited, }); self.shared.settle(Serving::Failed(why)); } } } fn run_command(&self, text: &str) { self.shared.emit(Event::Error { message: format!( "a llama.cpp session has no commands of its own, so {text} means nothing to it." ), }); } fn compact(&self) { self.shared.emit(Event::Error { message: "llama.cpp has no compaction. Clear the session instead, which costs nothing." .to_string(), }); } fn clear(&self) { // All of it. `conversation` folds from the last of these, so recording // the marker *is* the reset -- there is no driver state to keep in step // with it, which is the same property that makes a second device see the // same conversation this one does. self.shared.emit(Event::Cleared); } /// Stops generating and leaves the server loaded. /// /// Worth being deliberate about, because the cost points the other way from /// the Claude driver's: a `llama-server` holds its whole model in memory, so /// a leaked one is gigabytes nobody is using. It is left anyway, because the /// alternative is unloading and reloading that model on every backend /// restart -- minutes of disk, for a session somebody is in the middle of. /// The record is what keeps it from being *nobody's*. fn detach(&self) { self.shared.cancel.store(true, Ordering::SeqCst); self.shared.abandon_questions(); } fn stop(&self) { self.shared.cancel.store(true, Ordering::SeqCst); self.shared.abandon_questions(); if let Some(record) = process::live(&self.shared.session_dir) { process::stop(&record, process::STOP_GRACE); } process::clear(&self.shared.session_dir); } } /// The conversation so far, folded out of the transcript. /// /// Consecutive `AssistantText` deltas are one assistant turn, closed by the /// next user message or by the first tool call after them -- which is also /// what makes an interrupted reply come back as the partial text the phone /// actually saw. /// /// **A tool call and its result are part of the conversation, not decoration /// on it.** They go back as the assistant message that asked and the `tool` /// message that answered, which is the shape a chat template renders, and a /// call whose result never arrived is given [`tools::UNFINISHED`] rather than /// dropped: a template that finds a call with no answer either refuses the /// request or renders a conversation where the model asked for something and /// nothing came back, and the second is worse than the first. /// /// 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(); }; // Everything before the last clear is still in the transcript and is // deliberately not in the conversation. Folding from zero would put it back, // which is the whole of what clearing had to undo. let events = match events.iter().rposition(|e| e.event == Event::Cleared) { Some(at) => &events[at + 1..], None => &events[..], }; let mut fold = Fold::default(); for event in events.iter().cloned() { match event.event { Event::UserMessage { text, .. } => { fold.close(); fold.messages.push(Message::new("user", text)); } Event::AssistantText { delta } => { // Text after a call belongs to the reply that follows it, not // to the one that asked for it. if !fold.calls.is_empty() { fold.close(); } fold.text.push_str(&delta); } Event::AssistantTextFinal { text } => { if !fold.calls.is_empty() { fold.close(); } fold.text = text; } Event::ToolStart { id, tool, input } => fold.calls.push(Call { id, name: tool, arguments: input.to_string(), }), Event::ToolEnd { id, output } => { fold.results.insert(id, output); } _ => {} } } fold.close(); fold.messages } /// The running state of [`conversation`]: the assistant turn being assembled, /// and what is known about the calls in it. #[derive(Default)] struct Fold { messages: Vec, text: String, calls: Vec, /// Outputs by call id. A map rather than a field on `Call` because the /// result arrives as its own event and, when two calls were made at once, /// not necessarily in the order they were asked for. results: HashMap, } impl Fold { /// Closes the assistant turn being assembled and its results, if there is /// one. Nothing is emitted for a turn with neither text nor calls, which /// is what the start of a conversation looks like. fn close(&mut self) { let text = std::mem::take(&mut self.text); let calls = std::mem::take(&mut self.calls); if text.is_empty() && calls.is_empty() { return; } self.messages.push(Message { tool_calls: calls.iter().map(Call::wire).collect(), ..Message::new("assistant", text) }); for call in calls { let output = self .results .remove(&call.id) .unwrap_or_else(|| tools::UNFINISHED.to_string()); self.messages.push(Message::result_of(&call.id, output)); } } } /// 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) } /// A model file on the machine that will serve it: where it is, and what its /// own metadata says about how to load it. #[derive(Debug, Clone, PartialEq)] struct Model { /// Absolute, on that machine. path: String, /// Whether it carries a multi-token-prediction head, which decides one /// flag -- see [`crate::gguf::has_mtp_head`], and note that guessing wrong /// in the "yes" direction is a server that exits rather than one that runs /// slightly differently. mtp: bool, } /// The model file **on the machine that will serve it**, confirmed to be /// there and read far enough to know how to load it. /// /// One function rather than a local check and hope for the other case: the same /// question has to be asked of two filesystems. The remote answer is measured /// for the reason the local one is -- a missing file otherwise becomes a /// `llama-server` that starts, fails to load, and reports as a session that /// never became ready, which reads as the machine being slow. /// /// One blocking round trip on a remote spawn, which is what the spawn is /// already paying to start ssh, and it carries the head of the file as well as /// its path -- for the same reason it carries the path: the file is over /// there, and a second round trip to read a few hundred bytes of it would be /// a second way for the two answers to disagree. fn model_on(transport: &Transport, models_dir: &Path, key: &str) -> Result { let Transport::Ssh { name, .. } = transport else { let path = model_path(models_dir, key)?; let mtp = std::fs::File::open(&path) .map(|mut file| crate::gguf::has_mtp_head(&mut file)) .unwrap_or(false); return Ok(Model { path: path.to_string_lossy().into_owned(), mtp, }); }; // The same directory the spawn screen listed for this machine, and one // function for the same reason: a list from one place and a load from // another is a model that appears and then fails. let dir = crate::models::dir_on(transport, models_dir); // Checked here rather than in the script: `..` in a key would walk out of // the models directory on a machine this server can start processes on, // and the phone is where the key comes from. for part in key.split('/') { if part.is_empty() || part == "." || part == ".." { bail!("\"{key}\" is not a model key this can resolve"); } } let path = format!("{}/{key}", dir.trim_end_matches('/')); // `$HOME` on the far side, which is the only machine that knows what it is, // and the resolved path printed back so the launch hands `llama-server` // something absolute. "Not there" is answered rather than failed, because a // machine that could not be asked at all has to say so in its own words -- // it would otherwise arrive as this same sentence about a missing model. let script = format!( "p=$1; case $p in \"~\") p=$HOME;; \"~/\"*) p=$HOME/${{p#\"~/\"}};; esac; \ [ -f \"$p\" ] || {{ printf 'missing\\n'; exit 0; }}; \ printf 'at\\t%s\\t%s\\n' \"$(head -c {prefix} \"$p\" | base64 | tr -d '\\n')\" \"$p\"", prefix = crate::gguf::PREFIX_BYTES, ); let launch = Launch::new( "sh", vec!["-c".to_string(), script, "sh".to_string(), path.clone()], None, ); let answer = transport .capture_blocking(&launch) .with_context(|| format!("couldn't ask {name} where its models are"))?; // The path last, so a `\t` in it survives; the middle field is base64, // whose alphabet has no tab. let found = answer .trim() .strip_prefix("at\t") .and_then(|rest| rest.split_once('\t')); match found { Some((head, resolved)) => Ok(Model { path: resolved.to_string(), mtp: mtp_in_prefix(head), }), None => bail!( "{name} has no model at {path}. A llama.cpp session serves the file from the \ machine it runs on, so the model has to be on {name} -- what this backend has \ downloaded is somewhere else." ), } } /// Whether a base64 prefix of a model file shows a multi-token-prediction /// head. `false` for a prefix that did not survive the trip, which is the /// safe direction: the flag is what makes a server exit. fn mtp_in_prefix(head: &str) -> bool { use base64::Engine as _; base64::engine::general_purpose::STANDARD .decode(head.trim()) .is_ok_and(|bytes| crate::gguf::has_mtp_head(&mut bytes.as_slice())) } /// Polls until the server says it is ready, or gives up. /// /// Watches the process as well as the port, because the two failures need /// different words and one of them is common: a model that will not load, /// a port already taken on the far machine, a `llama-server` too old for /// a flag. All of those exit within a second and none of them will ever /// answer `/health`, so waiting out the timeout turns a server that said /// exactly what was wrong into "gave up after 300s". fn wait_until_ready(endpoint: &str, session_dir: &Path) -> 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(()); } // `None` is the session having been stopped or deleted while this // waited, which is nobody's fault and still not worth waiting on. match process::recorded(session_dir) { Some((_, process::Liveness::Alive | process::Liveness::Unknown)) => {} Some((_, process::Liveness::Dead)) | None => { bail!("it exited before it answered.{}", log_tail(session_dir)); } } if std::time::Instant::now() > deadline { bail!( "gave up after {}s.{}", READY_TIMEOUT.as_secs(), log_tail(session_dir) ); } std::thread::sleep(std::time::Duration::from_millis(250)); } } /// The end of `llama-server`'s own log, for a failure message. /// /// Its account of what went wrong is the useful half -- "failed to load /// model", "bind: Address already in use" -- and on a remote session it /// is the only half, since nobody reading the phone can open a file on /// that machine. Bounded, because this ends up in an event a phone draws. fn log_tail(session_dir: &Path) -> String { let Ok(text) = std::fs::read_to_string(session_dir.join(SERVER_LOG)) else { return String::new(); }; let tail: Vec<&str> = text.lines().rev().take(LOG_TAIL_LINES).collect(); if tail.is_empty() { return String::new(); } format!( " It last said: {}", tail.into_iter().rev().collect::>().join(" / ") ) } /// How much of that log to carry into a message somebody reads on a phone. const LOG_TAIL_LINES: usize = 6; /// How many tokens this server can hold, from the server itself. /// /// `/props` reports the per-slot context, which is the whole of it because /// this driver always starts one slot -- see the `-np` argument. `None` for a /// server that would not answer, which draws as no ceiling rather than as a /// guessed one. fn context_window(endpoint: &str) -> Option { ureq::get(format!("{endpoint}/props")) .call() .ok()? .body_mut() .read_json::() .ok()? .pointer("/default_generation_settings/n_ctx") .and_then(Value::as_u64) } /// One streamed completion: posts the conversation, emits each text delta as /// it arrives, and assembles whatever tool calls came with it. /// /// Emits *and* returns, and the two carry different things on purpose. Text is /// emitted, because the transcript those events land in is what the reader /// watches and what the next turn reads back. Tool calls are returned, because /// what happens to them next -- asking, running, reporting -- is the caller's, /// and a call is not in the transcript until it has actually been made. /// /// `reasoning_content` is dropped, which is what the Claude driver does with /// thinking deltas. A transcript is what was said, and this app does not draw /// a model's working. fn generate( endpoint: &str, messages: &[Message], tools: &Tools, sampling: &serde_json::Map, shared: &Shared, ) -> 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"); if let Some(offered) = tools.offered() { map.insert("tools".to_string(), json!(offered)); } for (key, value) in sampling { map.insert(key.clone(), value.clone()); } let mut response = ureq::post(format!("{endpoint}/v1/chat/completions")) .config() // A turn can be long: a slow model on a long prompt, and the whole // reply arrives down this one response. Without a ceiling at all a // wedged server holds the turn for ever; this is the generous version // of one. .timeout_global(Some(GENERATE_TIMEOUT)) // The refusal is read rather than thrown away -- see `refusal`. .http_status_as_error(false) .build() .header("Content-Type", "application/json") .send_json(&body) .context("asking llama-server to generate")?; if let Some(why) = refusal(&mut response) { bail!("{why}"); } let reader = std::io::BufReader::new(response.body_mut().as_reader()); let mut text = String::new(); // Calls in the order the stream numbered them. A model asking for several // at once interleaves their fragments, each tagged with its index. let mut calls: Vec = Vec::new(); let mut tokens = 0u64; // What the model is holding when this call ends: the prompt it was given // plus the reply it produced, which is exactly what the next call's prompt // begins with. // // The prompt alone was wrong, and measurably: a five-word reply reported // 2,042 against a slot holding 2,355 (2026-09-19, checked against // `stop processing: n_tokens` in the server's own log). The gap is the // reply, so it grows with how much the model just said -- which is the // worst direction for a figure somebody is watching to see how much room // is left. let mut context = None; for line in std::io::BufRead::lines(reader) { if shared.cancel.load(Ordering::SeqCst) { break; } let line = line.context("reading the generation stream")?; // Server-sent events: the payload lines are the ones that matter. let Some(payload) = line.strip_prefix("data: ") else { continue; }; if payload.trim() == "[DONE]" { break; } let Ok(chunk) = serde_json::from_str::(payload) else { continue; }; // One figure answering both questions, which for this dialect it // genuinely does: what the call was charged for and what the model is // left holding are the same tokens, because nothing here is billed and // the whole conversation is resent every time. if let Some(total) = chunk.pointer("/usage/total_tokens").and_then(Value::as_u64) { tokens = total; context = Some(total); } let Some(delta) = chunk.pointer("/choices/0/delta") else { continue; }; if let Some(fragment) = delta.get("content").and_then(Value::as_str) && !fragment.is_empty() { text.push_str(fragment); shared.emit(Event::AssistantText { delta: fragment.to_string(), }); } for fragment in delta .get("tool_calls") .and_then(Value::as_array) .into_iter() .flatten() { absorb(&mut calls, fragment); } } if tokens > 0 { shared.emit(Event::UsageDelta { tokens, context }); } // A call whose name never arrived is not a call. It happens when a stream // is cut mid-fragment, and running it would mean inventing what was asked // for. calls.retain(|call| !call.name.is_empty()); Ok(Reply { text, calls }) } /// How long one completion may take before the turn is abandoned. const GENERATE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(1800); /// What `llama-server` said when it refused a request, or `None` when it did /// not refuse. /// /// Read out of the body rather than left as the status code, because the body /// is the half that says what to do: a turn whose tool results have outgrown /// the context comes back as *"request (9960 tokens) exceeds the available /// context size (8192 tokens), try increasing it"*, and the reader of a phone /// cannot open the server's log to find that out. As a bare status it was /// "asking llama-server to generate: http status: 400", which names neither /// the cause nor the fix. /// /// Common rather than exotic now that sessions have tools: one web search is /// thousands of tokens of result, and a default context is a few thousand. fn refusal(response: &mut ureq::http::Response) -> Option { let status = response.status(); if status.is_success() { return None; } let body = response.body_mut().read_to_string().unwrap_or_default(); // The dialect's own shape first, then whatever it sent, then the code // alone -- which is all there is for a server that refused with no body. let message = serde_json::from_str::(&body) .ok() .and_then(|body| { ["/error/message", "/message"] .iter() .find_map(|at| body.pointer(at).and_then(Value::as_str).map(str::to_string)) }) .unwrap_or_else(|| body.trim().to_string()); Some(if message.is_empty() { format!("llama-server refused the request ({status})") } else { format!("llama-server refused the request: {message}") }) } /// Folds one `tool_calls` fragment into the calls assembled so far. /// /// The wire sends a call in pieces, each carrying the index it belongs to: the /// id and name once, then the arguments a few characters at a time. A server /// that sends the whole call in one fragment -- which `llama-server` does /// today -- goes through the same path and simply arrives complete. /// /// `index` is trusted only as a slot number, and a fragment without one is /// taken as the newest call: that is what a stream sending exactly one call /// and omitting the field means, and appending a fresh call for each of its /// fragments instead would produce a call per character of arguments. fn absorb(calls: &mut Vec, fragment: &Value) { let at = match fragment.get("index").and_then(Value::as_u64) { Some(index) => index as usize, None => calls.len().saturating_sub(1), }; if calls.len() <= at { calls.resize_with(at + 1, Call::default); } let call = &mut calls[at]; if let Some(id) = fragment.get("id").and_then(Value::as_str) && !id.is_empty() { call.id = id.to_string(); } if let Some(name) = fragment .pointer("/function/name") .and_then(Value::as_str) .filter(|name| !name.is_empty()) { call.name = name.to_string(); } if let Some(arguments) = fragment .pointer("/function/arguments") .and_then(Value::as_str) { call.arguments.push_str(arguments); } // A call the server gave no id is still a call, and every later message // about it is addressed by that id -- so one is minted rather than left // empty, which would pair every result with every call. if call.id.is_empty() && !call.name.is_empty() { call.id = super::random_hex(); } } #[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 { id: None, text: "hello".into(), attachments: Vec::new(), }, Event::AssistantText { delta: "hi ".into(), }, Event::AssistantText { delta: "there".into(), }, Event::Status { state: SessionStatus::Idle, }, Event::UserMessage { id: None, text: "again".into(), attachments: Vec::new(), }, 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. fn an_interrupted_reply_stays_in_the_conversation() { let (_dir, path) = transcript_with(&[ Event::UserMessage { id: None, text: "count".into(), attachments: Vec::new(), }, 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 { id: None, text: "hello".into(), attachments: Vec::new(), }, Event::Error { message: "something went wrong".into(), }, Event::AssistantText { delta: "still here".into(), }, Event::UsageDelta { tokens: 12, context: Some(12), }, ]); let messages = conversation(&path); assert_eq!(messages.len(), 2); assert_eq!(messages[0].content, "hello"); assert_eq!(messages[1].content, "still here"); } #[test] /// Clearing decides what the *model* is given, not just what the phone /// draws. Everything above the marker stays in the transcript and none of it /// is sent. fn the_conversation_starts_after_the_last_clear() { let (_dir, path) = transcript_with(&[ Event::UserMessage { id: None, text: "the long expensive conversation".into(), attachments: Vec::new(), }, Event::AssistantText { delta: "at length".into(), }, Event::Cleared, Event::UserMessage { id: None, text: "a fresh start".into(), attachments: Vec::new(), }, Event::AssistantText { delta: "cheaply".into(), }, ]); let messages = conversation(&path); assert_eq!(messages.len(), 2); assert_eq!(messages[0].content, "a fresh start"); assert_eq!(messages[1].content, "cheaply"); } #[test] /// The *last* one, so clearing twice does not resurrect what the /// first clear dropped. fn only_the_newest_clear_counts() { let (_dir, path) = transcript_with(&[ Event::UserMessage { id: None, text: "one".into(), attachments: Vec::new(), }, Event::Cleared, Event::UserMessage { id: None, text: "two".into(), attachments: Vec::new(), }, Event::Cleared, Event::UserMessage { id: None, text: "three".into(), attachments: Vec::new(), }, ]); let messages = conversation(&path); assert_eq!(messages.len(), 1); assert_eq!(messages[0].content, "three"); } #[test] /// The shape a chat template renders a tool call in: the assistant message /// that asked, then a `tool` message per call, in the order asked. What /// this is really testing is that a turn's events survive the trip out to /// the transcript and back, because that round trip is the only memory /// this driver has. fn a_call_and_its_result_come_back_as_two_messages() { let (_dir, path) = transcript_with(&[ Event::UserMessage { id: None, text: "what is in notes.txt".into(), attachments: Vec::new(), }, Event::AssistantText { delta: "Let me look.".into(), }, Event::ToolStart { id: "call-1".into(), tool: "read_file".into(), input: json!({"path": "notes.txt"}), }, Event::ToolEnd { id: "call-1".into(), output: "alpha\nbeta\n".into(), }, Event::AssistantText { delta: "It says alpha and beta.".into(), }, ]); let messages = conversation(&path); let shape: Vec<(&str, &str)> = messages .iter() .map(|m| (m.role.as_str(), m.content.as_str())) .collect(); assert_eq!( shape, [ ("user", "what is in notes.txt"), ("assistant", "Let me look."), ("tool", "alpha\nbeta\n"), ("assistant", "It says alpha and beta."), ], ); assert_eq!(messages[1].tool_calls.len(), 1); assert_eq!( messages[1].tool_calls[0].pointer("/function/name"), Some(&json!("read_file")), ); assert_eq!(messages[2].tool_call_id.as_deref(), Some("call-1")); } #[test] /// Two calls in one turn, whose results arrive out of order -- which they /// do, because the driver runs them one after another and the transcript /// records whatever finished. Each result has to find its own call. fn results_are_paired_by_id_rather_than_by_position() { let (_dir, path) = transcript_with(&[ Event::UserMessage { id: None, text: "look at both".into(), attachments: Vec::new(), }, Event::ToolStart { id: "a".into(), tool: "read_file".into(), input: json!({"path": "one"}), }, Event::ToolStart { id: "b".into(), tool: "read_file".into(), input: json!({"path": "two"}), }, Event::ToolEnd { id: "b".into(), output: "second".into(), }, Event::ToolEnd { id: "a".into(), output: "first".into(), }, ]); let messages = conversation(&path); assert_eq!(messages[1].tool_calls.len(), 2); assert_eq!( messages[2..] .iter() .map(|m| (m.tool_call_id.as_deref(), m.content.as_str())) .collect::>(), [(Some("a"), "first"), (Some("b"), "second")], ); } #[test] /// A turn stopped between the call and its result. Every call owes an /// answer or the conversation will not render, so the gap is filled with /// a sentence saying what happened rather than an invented outcome -- and /// the call itself is kept, because dropping it would tell the model it /// never asked. fn a_call_with_no_result_is_answered_as_unfinished() { let (_dir, path) = transcript_with(&[ Event::UserMessage { id: None, text: "run it".into(), attachments: Vec::new(), }, Event::ToolStart { id: "call-1".into(), tool: "exec_shell_command".into(), input: json!({"command": "sleep 100"}), }, Event::Status { state: SessionStatus::Idle, }, ]); let messages = conversation(&path); assert_eq!(messages[1].tool_calls.len(), 1); assert_eq!(messages[2].role, "tool"); assert_eq!(messages[2].content, tools::UNFINISHED); } #[test] /// Text after a call belongs to the reply that follows it. Folded into the /// message that asked, it would go back as an assistant turn that both /// requested a tool and reported its answer -- before the answer existed. fn text_after_a_call_opens_a_new_assistant_message() { let (_dir, path) = transcript_with(&[ Event::UserMessage { id: None, text: "go".into(), attachments: Vec::new(), }, Event::AssistantText { delta: "checking".into(), }, Event::ToolStart { id: "c".into(), tool: "get_info".into(), input: json!({}), }, Event::ToolEnd { id: "c".into(), output: "linux".into(), }, Event::AssistantText { delta: "it is linux".into(), }, ]); let messages = conversation(&path); assert_eq!(messages[1].content, "checking"); assert_eq!(messages[1].tool_calls.len(), 1); assert_eq!(messages[3].content, "it is linux"); assert!(messages[3].tool_calls.is_empty()); } #[test] /// The streaming shape: a call arrives as an id and a name once, then its /// arguments a few characters at a time. fn a_call_streamed_in_fragments_is_assembled() { let mut calls = Vec::new(); for fragment in [ json!({"index": 0, "id": "call-9", "function": {"name": "read_file", "arguments": ""}}), json!({"index": 0, "function": {"arguments": "{\"pa"}}), json!({"index": 0, "function": {"arguments": "th\": \"a.txt\"}"}}), ] { absorb(&mut calls, &fragment); } assert_eq!(calls.len(), 1); assert_eq!(calls[0].id, "call-9"); assert_eq!(calls[0].name, "read_file"); assert_eq!(calls[0].arguments(), json!({"path": "a.txt"})); } #[test] /// Two at once, interleaved. The index is the only thing that says which /// fragment belongs to which call. fn interleaved_fragments_are_kept_apart_by_index() { let mut calls = Vec::new(); for fragment in [ json!({"index": 0, "id": "a", "function": {"name": "read_file"}}), json!({"index": 1, "id": "b", "function": {"name": "grep_search"}}), json!({"index": 1, "function": {"arguments": "{\"q\":1}"}}), json!({"index": 0, "function": {"arguments": "{\"q\":0}"}}), ] { absorb(&mut calls, &fragment); } assert_eq!( calls .iter() .map(|call| (call.id.as_str(), call.arguments())) .collect::>(), [("a", json!({"q": 0})), ("b", json!({"q": 1}))], ); } #[test] /// A server that sends the whole call at once and numbers nothing. Each /// fragment appended as its own call would make one call per chunk, and /// the arguments would arrive as a row of empty ones. fn fragments_without_an_index_extend_the_newest_call() { let mut calls = Vec::new(); absorb( &mut calls, &json!({"id": "a", "function": {"name": "read_file", "arguments": "{\"p\":"}}), ); absorb(&mut calls, &json!({"function": {"arguments": "1}"}})); assert_eq!(calls.len(), 1); assert_eq!(calls[0].arguments(), json!({"p": 1})); } #[test] /// The allowances a session has been given, folded out of the answers /// alone -- which is what makes them survive a backend restart without /// being stored anywhere of their own. fn always_allow_answers_fold_into_the_allowed_set() { let (_dir, path) = transcript_with(&[ Event::Answered { id: "q1".into(), answers: vec![tools::ALLOW_ONCE.to_string()], }, Event::Answered { id: "q2".into(), answers: vec![format!("{}read_file", tools::ALWAYS_PREFIX)], }, Event::Answered { id: "q3".into(), answers: vec![tools::REFUSE.to_string()], }, ]); let allowed = allowances(&path); assert_eq!( allowed, std::collections::HashSet::from(["read_file".to_string()]), ); } #[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", ); } } }