iris-aiandClaude Opus 5
ac476ab0c9
A llama session was a chat box: no tools, a fixed model, no permission
mode, and a model name drawn as the path the file sits at. It now runs the
agent loop itself, which is what the pieces below all hang off.
Tools are `llama-server`'s own (`--tools all`), which that server both
publishes and runs -- `GET /tools` for the definitions, `POST /tools` to
call one. Web search is Exa's MCP server, reached from this backend rather
than from the machine serving the model: that is what llama.cpp's own web
UI does, and it puts the search on the machine with a route out instead of
the one with the GPU. `llama-server`'s `--mcp-servers-json` can only spawn
local commands, so using it would have meant a Node bridge on every
machine that serves a model.
Driving the loop is what makes the permission gate ours. Two modes,
`manual` and `bypassPermissions`, which is what the mechanism has: the web
UI asks before every call and remembers the tools you say "always" to. The
allowances fold back out of the transcript's own answers, so they survive
a restart and a model change without being stored anywhere else.
Also here, because tools made each of them matter:
- **Loading is a state.** A 12 GB model takes twenty seconds to reach
memory and refuses everything until it has; the session used to report
`running` for that whole time, and a message sent meanwhile came back as
an error. It is `loading` now, and the message waits.
- **The model can be changed.** A `llama-server` holds one model, so this
stops it and starts another. The conversation survives because it was
never in the server.
- **Models are named, not pathed.** `general.name` read out of the file
itself -- over ssh too, in the round trip the spawn was already making.
Where two models share a name the file name breaks the tie.
- **`-np 1`, and the MTP draft head where the file has one.** Measured on
the 27B here: 41.5 tok/s plain, 61.4 with `--spec-type draft-mtp` at one
slot, and 28 with it at four -- speculating against a split KV cache is
worse than not speculating. The flag is conditional because asking for a
head that is not there makes `llama-server` exit.
- **A refusal says what to do.** Tool results are thousands of tokens, so
an overrun context is now ordinary; it was "http status: 400" and is now
the server's own "exceeds the available context size, try increasing it".
`GET /machines/{id}/models` is gone: the provider models route answers the
same question, and two answers to one question is how a picker comes to
offer a model the spawn screen does not.
Verified end to end against real models: a tool call asked and allowed, an
Exa search, a shell command, a 27B loaded while a message waited on it, a
model switch mid-session, a second message queued behind a running turn,
and the whole of it again on a session running over ssh.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2053 lines
81 KiB
Rust
2053 lines
81 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 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;
|
|
|
|
/// 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";
|
|
|
|
/// 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<Value>,
|
|
/// Which call this message is the result of. Only on `role: "tool"`.
|
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
|
tool_call_id: Option<String>,
|
|
}
|
|
|
|
impl Message {
|
|
fn new(role: &str, content: impl Into<String>) -> 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<String>) -> 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<Call>,
|
|
}
|
|
|
|
/// 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<Tools>,
|
|
},
|
|
/// 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 chosen at spawn, sent with every request.
|
|
sampling: serde_json::Map<String, Value>,
|
|
/// 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<String>,
|
|
/// 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<Turns>,
|
|
serving: Mutex<Serving>,
|
|
/// 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<Arc<Mutex<McpServer>>>,
|
|
/// How much this session asks before acting. Live, so a reader who has
|
|
/// seen enough can stop being asked without restarting anything.
|
|
mode: Mutex<String>,
|
|
/// 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<std::collections::HashSet<String>>,
|
|
/// Questions a turn is blocked on, by question id.
|
|
asked: Mutex<HashMap<String, std::sync::mpsc::Sender<Vec<String>>>>,
|
|
}
|
|
|
|
pub struct LlamaDriver {
|
|
shared: Arc<Shared>,
|
|
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<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 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));
|
|
}
|
|
}
|
|
|
|
let driver = Self {
|
|
shared: Arc::new(Shared {
|
|
sink,
|
|
transcript: transcript.to_path_buf(),
|
|
session_dir: session_dir.to_path_buf(),
|
|
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()),
|
|
);
|
|
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<Tools>)> {
|
|
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<String> {
|
|
// 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(),
|
|
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. Every one of
|
|
// them, because a session offered a subset is a session that says "I
|
|
// can't do that" about something it was installed to do, and the
|
|
// question of whether a particular call should happen is the
|
|
// permission gate's rather than a flag's. They run on the machine
|
|
// serving the model, which is the machine the files are on.
|
|
"--tools".into(),
|
|
"all".into(),
|
|
// 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(),
|
|
];
|
|
// 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"),
|
|
] {
|
|
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<Arc<Mutex<McpServer>>> {
|
|
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::<Vec<_>>()
|
|
.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<String> {
|
|
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<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)) => {
|
|
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<Shared>) {
|
|
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<Shared>, queued: Option<String>, 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<Shared>,
|
|
endpoint: &str,
|
|
tools: &Tools,
|
|
mut messages: Vec<Message>,
|
|
) -> Result<()> {
|
|
for _ in 0..MAX_STEPS {
|
|
if shared.cancel.load(Ordering::SeqCst) {
|
|
return Ok(());
|
|
}
|
|
let reply = generate(endpoint, &messages, tools, &shared.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<Shared>, 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<Shared>, 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<AttachmentRef>) {
|
|
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) {}
|
|
|
|
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<Message> {
|
|
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<Message>,
|
|
text: String,
|
|
calls: Vec<Call>,
|
|
/// 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<String, String>,
|
|
}
|
|
|
|
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<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)
|
|
}
|
|
|
|
/// 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<Model> {
|
|
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::<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 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<String, Value>,
|
|
shared: &Shared,
|
|
) -> Result<Reply> {
|
|
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<Call> = Vec::new();
|
|
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 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::<Value>(payload) else {
|
|
continue;
|
|
};
|
|
if let Some(usage) = chunk.get("usage") {
|
|
if let Some(total) = usage.get("total_tokens").and_then(Value::as_u64) {
|
|
tokens = total;
|
|
}
|
|
if let Some(prompt) = usage.get("prompt_tokens").and_then(Value::as_u64) {
|
|
context = Some(prompt);
|
|
}
|
|
}
|
|
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<ureq::Body>) -> Option<String> {
|
|
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::<Value>(&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<Call>, 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::<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]
|
|
/// 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::<Vec<_>>(),
|
|
[(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::<Vec<_>>(),
|
|
[("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",
|
|
);
|
|
}
|
|
}
|
|
}
|