906 lines
35 KiB
Rust
906 lines
35 KiB
Rust
//! 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 setup
|
|
//! 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.
|
|
|
|
use std::path::{Path, PathBuf};
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicBool, Ordering};
|
|
|
|
use anyhow::{Context, Result, bail};
|
|
use serde::{Deserialize, Serialize};
|
|
use serde_json::json;
|
|
|
|
use super::driver::{AttachmentRef, Driver, Event, EventSink, SessionStatus};
|
|
use super::process;
|
|
use super::transport::{Launch, Streams, Transport};
|
|
use crate::config::{ProviderConfig, SessionConfig};
|
|
|
|
/// 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.
|
|
#[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<String, serde_json::Value>,
|
|
/// Set by [`Driver::interrupt`]; the streaming loop checks it between
|
|
/// chunks and stops, leaving what was generated in the transcript.
|
|
cancel: Arc<AtomicBool>,
|
|
/// Where this session's process record lives, so [`Driver::stop`] can find
|
|
/// the server it has to end.
|
|
session_dir: PathBuf,
|
|
}
|
|
|
|
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<super::subagent::Subagents>,
|
|
) -> Result<Self> {
|
|
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_on(transport, models_dir, model)?;
|
|
|
|
// 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.
|
|
if let Some(process::Record {
|
|
detail: process::Detail::Http { port },
|
|
pid,
|
|
..
|
|
}) = process::live(session_dir)
|
|
{
|
|
tracing::info!(
|
|
"session {} reattaching to the llama-server it left loaded (pid {pid}, port {port})",
|
|
meta.id
|
|
);
|
|
return Ok(Self::attached(
|
|
format!("http://127.0.0.1:{port}"),
|
|
meta,
|
|
model,
|
|
transcript,
|
|
session_dir,
|
|
sink,
|
|
));
|
|
}
|
|
|
|
// 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<String> = vec![
|
|
"-m".into(),
|
|
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(),
|
|
];
|
|
// 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.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(Self::attached(
|
|
format!("http://127.0.0.1:{}", forward.here),
|
|
meta,
|
|
model,
|
|
transcript,
|
|
session_dir,
|
|
sink,
|
|
))
|
|
}
|
|
|
|
/// The driver for a `llama-server` at `endpoint`, however it got there.
|
|
///
|
|
/// Shared by starting one and adopting one, because everything after "there
|
|
/// is a server at this address" is identical -- including waiting for it to
|
|
/// answer, which an adopted one still owes: a recorded pid says a process
|
|
/// exists, not that its model is loaded.
|
|
fn attached(
|
|
endpoint: String,
|
|
meta: &SessionConfig,
|
|
model: &str,
|
|
transcript: &Path,
|
|
session_dir: &Path,
|
|
sink: EventSink,
|
|
) -> Self {
|
|
// Loading is slow enough to be worth saying so: the session shows as
|
|
// running until the model is in memory, 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();
|
|
let session_dir = session_dir.to_path_buf();
|
|
std::thread::spawn(move || match wait_until_ready(&endpoint, &session_dir) {
|
|
Ok(()) => {
|
|
tracing::info!("{model} loaded and answering at {endpoint}");
|
|
let _ = sink.send(Event::Status {
|
|
state: SessionStatus::Idle,
|
|
});
|
|
watch(session_dir, sink);
|
|
}
|
|
Err(err) => {
|
|
let _ = sink.send(Event::Error {
|
|
message: format!("{model} never became ready: {err:#}"),
|
|
});
|
|
let _ = sink.send(Event::Status {
|
|
state: SessionStatus::Exited,
|
|
});
|
|
process::clear(&session_dir);
|
|
}
|
|
});
|
|
}
|
|
|
|
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::<f64>()
|
|
{
|
|
sampling.insert(field.to_string(), json!(number));
|
|
}
|
|
}
|
|
|
|
Self {
|
|
sink,
|
|
endpoint,
|
|
transcript: transcript.to_path_buf(),
|
|
sampling,
|
|
cancel: Arc::new(AtomicBool::new(false)),
|
|
session_dir: session_dir.to_path_buf(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// 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<std::fs::File> {
|
|
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)) => {
|
|
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 Driver for LlamaDriver {
|
|
fn send_user_message(&self, text: String, attachments: Vec<AttachmentRef>) {
|
|
if !attachments.is_empty() {
|
|
let _ = self.sink.send(Event::Error {
|
|
message: "this model can't be sent attachments or files".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 || {
|
|
// Nothing is ever held back here -- there is no queue to wait in --
|
|
// so the message is taken the moment it arrives. Said anyway,
|
|
// because this is what records it: see `MessageTaken`.
|
|
let _ = sink.send(Event::MessageTaken {
|
|
id: None,
|
|
text: text.clone(),
|
|
// Never any: this driver refuses attachments above.
|
|
attachments: Vec::new(),
|
|
});
|
|
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 message's own transcript entry is still on its
|
|
// way 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, _answers: &[String]) {
|
|
// Nothing here asks questions: this driver has no tools.
|
|
}
|
|
|
|
fn interrupt(&self) {
|
|
self.cancel.store(true, Ordering::Relaxed);
|
|
}
|
|
|
|
// 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) {}
|
|
|
|
fn set_permission_mode(&self, _mode: &str) {
|
|
let _ = self.sink.send(Event::Error {
|
|
message: "a llama.cpp session runs no tools, so there is nothing for a permission \
|
|
mode to govern."
|
|
.to_string(),
|
|
});
|
|
}
|
|
|
|
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 run_command(&self, text: &str) {
|
|
let _ = self.sink.send(Event::Error {
|
|
message: format!(
|
|
"a llama.cpp session has no commands of its own, so {text} means nothing to it."
|
|
),
|
|
});
|
|
}
|
|
|
|
fn compact(&self) {
|
|
let _ = self.sink.send(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.
|
|
let _ = self.sink.send(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.cancel.store(true, Ordering::Relaxed);
|
|
}
|
|
|
|
fn stop(&self) {
|
|
self.cancel.store(true, Ordering::Relaxed);
|
|
if let Some(record) = process::live(&self.session_dir) {
|
|
process::stop(&record, process::STOP_GRACE);
|
|
}
|
|
process::clear(&self.session_dir);
|
|
}
|
|
}
|
|
|
|
/// 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.
|
|
///
|
|
/// 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<Message> {
|
|
let Ok(events) = crate::session::transcript::read_after(path, 0) else {
|
|
return Vec::new();
|
|
};
|
|
let mut messages: Vec<Message> = Vec::new();
|
|
let mut pending = String::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[..],
|
|
};
|
|
for event in events.iter().cloned() {
|
|
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),
|
|
Event::AssistantTextFinal { text } => {
|
|
pending = text;
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
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<PathBuf> {
|
|
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)
|
|
}
|
|
|
|
/// The model file's path **on the machine that will serve it**, confirmed to be
|
|
/// there.
|
|
///
|
|
/// 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. The alternative is a path built here from a `~`
|
|
/// this machine cannot expand.
|
|
fn model_on(transport: &Transport, models_dir: &Path, key: &str) -> Result<String> {
|
|
let Transport::Ssh { name, .. } = transport else {
|
|
return Ok(model_path(models_dir, key)?.to_string_lossy().into_owned());
|
|
};
|
|
// 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 = "p=$1; case $p in \"~\") p=$HOME;; \"~/\"*) p=$HOME/${p#\"~/\"};; esac; \
|
|
[ -f \"$p\" ] && printf 'at\\t%s\\n' \"$p\" || printf 'missing\\n'"
|
|
.to_string();
|
|
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"))?;
|
|
match answer.trim().split_once('\t') {
|
|
Some(("at", resolved)) => Ok(resolved.to_string()),
|
|
_ => 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."
|
|
),
|
|
}
|
|
}
|
|
|
|
/// 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::<Vec<_>>().join(" / ")
|
|
)
|
|
}
|
|
|
|
/// How much of that log to carry into a message somebody reads on a phone.
|
|
const LOG_TAIL_LINES: usize = 6;
|
|
|
|
/// One streamed completion: posts the conversation, emits each delta as it
|
|
/// arrives. Emits rather than returns, because the transcript those events land
|
|
/// in is what the next turn reads back.
|
|
fn generate(
|
|
endpoint: &str,
|
|
messages: &[Message],
|
|
sampling: &serde_json::Map<String, serde_json::Value>,
|
|
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;
|
|
// The prompt side only, which is what the model is holding -- the same
|
|
// definition the other dialects report, so one word on the phone means one
|
|
// thing whichever kind of session it is.
|
|
let mut context = None;
|
|
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.
|
|
let Some(payload) = line.strip_prefix("data: ") else {
|
|
continue;
|
|
};
|
|
if payload.trim() == "[DONE]" {
|
|
break;
|
|
}
|
|
let Ok(chunk) = serde_json::from_str::<serde_json::Value>(payload) else {
|
|
continue;
|
|
};
|
|
if let Some(usage) = chunk.get("usage") {
|
|
if let Some(total) = usage
|
|
.get("total_tokens")
|
|
.and_then(serde_json::Value::as_u64)
|
|
{
|
|
tokens = total;
|
|
}
|
|
if let Some(prompt) = usage
|
|
.get("prompt_tokens")
|
|
.and_then(serde_json::Value::as_u64)
|
|
{
|
|
context = Some(prompt);
|
|
}
|
|
}
|
|
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, context });
|
|
}
|
|
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 {
|
|
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::<Vec<_>>(),
|
|
[
|
|
("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]
|
|
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",
|
|
);
|
|
}
|
|
}
|
|
}
|