iris/ai-app
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Serve a machine's models from one shared llama-server

Browse Source 
A llama.cpp session had its own `llama-server`: two sessions on one model
held two copies of it in memory, a model change bought a load only that
session benefited from, and the process was a session's to end. A machine's
models are now served by one `llama-server` in **router mode** -- no `-m`,
a preset file naming models and their flags, a child server per model asked
for, and each request routed by its `model` field. So one server per model
with that model's own settings is what a machine runs, while this backend
has one process, one port and one record per machine to keep track of.

The record is the mechanism every other driver already uses, so a restart
adopts it; a session records the same pid in its own directory as
`Detail::Shared`, and `process::signal` refuses to signal one of those --
which is what keeps stopping, deleting or cleaning up after one session
from unloading a model every other session is using. Nothing stops a router
on its own. That is deliberate (a loaded model is minutes of disk) and it is
why the machines tab now has a card per provider that opens its own screen:
how each model is loaded, how many stay in memory, Unload, and Stop.

How a model is *loaded* therefore belongs to the model on its machine rather
than to a session -- context size, GPU layers, threads, slots, speculative
decoding -- written into the preset as llama-server's own argument names.
Saving them re-reads that file, which unloads the model; that is the change
taking effect, and the dialog says so before you save. What stays a
session's is everything that rides on a request, including which tools it
offers: the router hosts one set for the machine and the choice is a filter
applied here, so it costs no reload (2,181 tokens of prompt with all seven,
698 with none).

Verified end to end against the scratch backend and the emulator: two
sessions sharing one loaded model with one child process, a second session
joining it with a 26ms prefill, a backend restart adopting the router and
answering with the prompt cache intact, the same over ssh to this VM, a
model's settings reaching the running server, Unload, and Stop leaving every
session `exited` with no error line.
This commit is contained in:
iris-ai committed 2026-09-19 17:37:31 -04:00
1 parent 74cda485e5
commit 8c323fc7a9
19 files changed
+2601 -511
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/config.rs
+109 -44
View File
@@ -102,6 +102,20 @@ pub struct ProviderConfig {
/// and this would be a second, quieter answer to the same question.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub mcp_servers: Vec<McpServerConfig>,
/// How each model this provider serves is loaded, by model key -- the
/// settings in [`LLAMA_MODEL_PARAMS`].
///
/// On the model rather than on the session because one loaded model is
/// what several sessions talk to: a machine's `llama-server` holds it
/// once, and a context size or a layer count that two sessions disagreed
/// about would be one of them being ignored. See `session::llama::router`.
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
pub model_settings: BTreeMap<String, BTreeMap<String, String>>,
/// How many models this provider's server keeps loaded at once before it
/// evicts the least recently used. `None` is one, which is the right
/// answer for a machine with one GPU -- see `router`'s `DEFAULT_MAX_LOADED`.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_loaded: Option<u32>,
}
/// An MCP server reached over HTTP.
@@ -334,6 +348,17 @@ impl DriverKind {
}
}
/// What a *model* this kind serves takes, which is nothing for a kind
/// that does not load models of its own. Separate from [`params`] because
/// the two have different owners, not different shapes: a session's ride
/// on its requests, a model's decide how the machine loads it.
pub fn model_params(self) -> &'static [ParamSpec] {
match self {
Self::LlamaCpp => LLAMA_MODEL_PARAMS,
Self::Echo | Self::ClaudeCli | Self::CodexCli => &[],
}
}
/// The mode used when a new-session form first selects this kind.
pub fn default_permission_mode(self) -> Option<&'static str> {
match self {
@@ -390,58 +415,25 @@ pub enum ParamKind {
},
}
/// What a llama.cpp session takes.
/// What a llama.cpp **session** takes: everything that rides on a request.
///
/// The server flags first, in the order they matter, then the sampling ones --
/// which is also the order of how disruptive changing one is.
/// Nothing here waits for a restart, and that is a property of the split
/// rather than a coincidence. A machine's `llama-server` holds one copy of a
/// model for every session using it, so how that model is *loaded* cannot be
/// one session's to decide -- those settings are [`LLAMA_MODEL_PARAMS`],
/// against the model on its machine.
const LLAMA_PARAMS: &[ParamSpec] = &[
ParamSpec {
key: "contextSize",
label: "Context size",
unset: "the model's own trained context",
kind: ParamKind::Integer,
restart: true,
},
ParamSpec {
key: "tools",
label: "Tools",
// Worth a control rather than a constant because of what it costs:
// the definitions of all seven are ~2,000 tokens of the context,
// every turn, before anything is said. On a small window that is the
// difference between a usable session and one that overruns.
// a prompt measured 2,181 tokens with all seven and 698 with none,
// every turn, before anything is said. A filter this backend applies
// to what the machine's server offers, so unlike the flag it replaced
// it takes effect on the next message.
unset: "all of them -- or a comma-separated list, or \"none\"",
kind: ParamKind::Text,
restart: true,
},
ParamSpec {
key: "gpuLayers",
label: "Layers on the GPU",
unset: "as many as fit",
kind: ParamKind::Integer,
restart: true,
},
ParamSpec {
key: "threads",
label: "Threads",
unset: "one per core",
kind: ParamKind::Integer,
restart: true,
},
ParamSpec {
key: "speculative",
label: "Speculative decoding",
unset: "on, for a model whose file carries a draft head",
kind: ParamKind::Choice {
options: &["auto", "off"],
},
restart: true,
},
ParamSpec {
key: "specDraftNMax",
label: "Tokens drafted ahead",
unset: "llama.cpp's own default",
kind: ParamKind::Integer,
restart: true,
restart: false,
},
ParamSpec {
key: "thinking",
@@ -490,6 +482,71 @@ const LLAMA_PARAMS: &[ParamSpec] = &[
},
];
/// What a llama.cpp **model** takes: everything that decides how it is loaded.
///
/// Per model on its machine rather than per session, because one loaded copy
/// is what every session on that model is talking to. Changing one reloads
/// that model -- for everybody using it, which is the honest consequence of
/// sharing it and is what the machine's provider view says before saving.
///
/// Every key here is written into the router's preset file as
/// `llama-server`'s own argument name, so adding a setting is a row here and a
/// row in `session::llama::router`'s `section`.
pub const LLAMA_MODEL_PARAMS: &[ParamSpec] = &[
ParamSpec {
key: "contextSize",
label: "Context size",
unset: "the model's own trained context",
kind: ParamKind::Integer,
// Every one of these does, which is what makes them the model's: see
// the doc comment above.
restart: true,
},
ParamSpec {
key: "gpuLayers",
label: "Layers on the GPU",
unset: "as many as fit",
kind: ParamKind::Integer,
restart: true,
},
ParamSpec {
key: "threads",
label: "Threads",
unset: "one per core",
kind: ParamKind::Integer,
restart: true,
},
ParamSpec {
key: "slots",
label: "Sessions answered at once",
// One, so that a second session's turn waits rather than splitting the
// model's cache. Measured 2026-09-19 on the 27B here: 41.5 tok/s
// plain, 61.4 with the draft head at one slot, and 28 with the head at
// four -- speculating against a split cache is slower than not
// speculating at all. Worth raising on a machine where several
// sessions really are used together and drafting does not pay.
unset: "one -- a second session's turn waits for the first",
kind: ParamKind::Integer,
restart: true,
},
ParamSpec {
key: "speculative",
label: "Speculative decoding",
unset: "on, for a model whose file carries a draft head",
kind: ParamKind::Choice {
options: &["auto", "off"],
},
restart: true,
},
ParamSpec {
key: "specDraftNMax",
label: "Tokens drafted ahead",
unset: "llama.cpp's own default",
kind: ParamKind::Integer,
restart: true,
},
];
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct TokenEntry {
@@ -698,6 +755,8 @@ impl Config {
command: None,
models: Vec::new(),
mcp_servers: Vec::new(),
model_settings: BTreeMap::new(),
max_loaded: None,
}
}
@@ -779,6 +838,8 @@ mod tests {
command: Some("/usr/bin/claude".to_string()),
models: Vec::new(),
mcp_servers: Vec::new(),
model_settings: BTreeMap::new(),
max_loaded: None,
},
]),
MachineConfig {
@@ -798,6 +859,8 @@ mod tests {
command: None,
models: vec!["haiku".to_string()],
mcp_servers: Vec::new(),
model_settings: BTreeMap::new(),
max_loaded: None,
}],
},
],
@@ -936,6 +999,8 @@ sessions: [(
command: Some("/usr/bin/claude".to_string()),
models: Vec::new(),
mcp_servers: Vec::new(),
model_settings: BTreeMap::new(),
max_loaded: None,
},
]);
assert_eq!(
server/src/machines.rs
+5
View File
@@ -80,6 +80,11 @@ pub async fn discover(transport: &Transport) -> Result<Vec<ProviderConfig>> {
_ => Vec::new(),
},
mcp_servers: mcp_defaults(*kind),
// What a probe cannot know: how this machine's models are loaded
// is configured after the fact, and a re-probe keeps it -- see
// `SessionManager::update_machine`.
model_settings: Default::default(),
max_loaded: None,
});
}
Ok(providers)
server/src/models.rs
+4 -3
View File
@@ -797,9 +797,10 @@ fn get_json(url: &str) -> Result<serde_json::Value> {
}
/// Percent-encodes a query string. Deliberately minimal -- this escapes what a
/// model search actually contains rather than implementing the whole rule set,
/// and anything unexpected becomes `%XX` rather than being passed through.
fn urlencode(value: &str) -> String {
/// model search and a model key actually contain rather than implementing the
/// whole rule set, and anything unexpected becomes `%XX` rather than being
/// passed through.
pub fn urlencode(value: &str) -> String {
value
.bytes()
.map(|b| match b {
server/src/provider_auth.rs
+2
View File
@@ -437,6 +437,8 @@ mod tests {
command: Some(cli.display().to_string()),
models: Vec::new(),
mcp_servers: Vec::new(),
model_settings: Default::default(),
max_loaded: None,
};
let started = logins.start(machine, provider);
server/src/routes.rs
+286 -3
View File
@@ -7,6 +7,12 @@
//! POST /machines add {name, ssh?} -- providers are discovered
//! POST /machines/probe dry run {ssh?}: what would be found there
//! GET /machines/{id} one machine, for refetching after a change
//! GET /machines/{id}/providers/{provider} what it is, its models and their
//! settings, and what its server is doing
//! POST /machines/{id}/providers/{provider}/settings {maxLoaded} -- null for the default
//! POST /machines/{id}/providers/{provider}/model-settings {model, params} -- how it loads
//! POST /machines/{id}/providers/{provider}/stop end the machine's shared server
//! POST /machines/{id}/providers/{provider}/unload {model} -- out of memory, server stays
//! GET /machines/{id}/providers/{provider}/models models that provider offers
//! POST /machines/{id}/providers/{provider}/auth begin provider sign-in
//! GET /machines/{id}/providers/{provider}/auth/{attempt} sign-in state
@@ -136,6 +142,26 @@ pub fn router(manager: Arc<SessionManager>) -> Router {
"/machines/{id}/providers/{provider}/models",
get(provider_models),
)
// What one provider on one machine is and how it is set up. Under the
// machine because that is what a provider belongs to: the same program
// on two machines is two of these, with their own models loaded.
.route("/machines/{id}/providers/{provider}", get(provider_view))
.route(
"/machines/{id}/providers/{provider}/settings",
post(set_provider_settings),
)
.route(
"/machines/{id}/providers/{provider}/model-settings",
post(set_model_settings),
)
.route(
"/machines/{id}/providers/{provider}/stop",
post(stop_provider_server),
)
.route(
"/machines/{id}/providers/{provider}/unload",
post(unload_provider_model),
)
// The filesystem of a configured machine. Under the machine
// rather than under a session because a filesystem is a property of
// a machine; a session only says where to start looking.
@@ -558,6 +584,18 @@ fn files_on(
))
}
/// The named provider on a machine, or the 404 saying it is not there. One
/// function because four routes ask the same question and the wording of the
/// answer is part of the surface.
fn provider_on<'a>(
machine: &'a crate::config::MachineConfig,
name: &str,
) -> Result<&'a crate::config::ProviderConfig, ApiError> {
machine
.provider(name)
.ok_or_else(|| ApiError::NotFound(format!("no provider {name} on {}", machine.name)))
}
/// A failure from one of the scripts is the *machine's* message, written to
/// be read where it happened, which is the phone. So it comes back as a 400
/// with those words rather than a 500 and a log line only the backend sees.
@@ -581,9 +619,7 @@ async fn provider_models(
UrlPath((id, provider_name)): UrlPath<(String, String)>,
) -> Result<axum::Json<Vec<crate::machines::OfferedModel>>, ApiError> {
let machine = machine_by_id(&manager, &id)?;
let provider = machine.provider(&provider_name).ok_or_else(|| {
ApiError::NotFound(format!("no provider {provider_name} on {}", machine.name))
})?;
let provider = provider_on(&machine, &provider_name)?;
let transport = crate::session::transport::Transport::for_machine(&machine);
crate::machines::provider_models(&transport, provider, manager.models_dir())
.await
@@ -591,6 +627,253 @@ async fn provider_models(
.map_err(from_machine)
}
/// One provider on one machine: what it is, every model it offers with the
/// settings that decide how that model is loaded, and what its shared server
/// is currently holding.
///
/// One answer rather than three requests, because this is one screen and the
/// phone reaching it is on the far end of a tunnel.
#[derive(serde::Serialize)]
#[serde(rename_all = "camelCase")]
struct ProviderView {
machine: String,
name: String,
kind: crate::config::DriverKind,
/// The program that was found, which is the honest answer to "what is
/// this" and is not something the phone may change -- see `machines`.
#[serde(skip_serializing_if = "Option::is_none")]
command: Option<String>,
/// What each of this provider's models takes, for drawing the controls.
/// Empty for a provider that loads no models, which draws no settings at
/// all rather than an empty form.
#[serde(skip_serializing_if = "<[_]>::is_empty")]
model_params: &'static [crate::config::ParamSpec],
/// How many models its server keeps loaded at once; absent is one.
#[serde(skip_serializing_if = "Option::is_none")]
max_loaded: Option<u32>,
models: Vec<ProviderModel>,
/// The tools this provider's sessions have over and above its own.
#[serde(skip_serializing_if = "Vec::is_empty")]
mcp_servers: Vec<String>,
/// Its shared server, or `None` for a provider that has none. The
/// difference matters on screen: "nothing is loaded" and "there is nothing
/// here to load" are not the same sentence.
#[serde(skip_serializing_if = "Option::is_none")]
server: Option<ServerView>,
}
#[derive(serde::Serialize)]
#[serde(rename_all = "camelCase")]
struct ProviderModel {
id: String,
label: String,
#[serde(skip_serializing_if = "std::collections::BTreeMap::is_empty")]
settings: std::collections::BTreeMap<String, String>,
/// What the server is doing with it -- `loaded`, `loading`, `unloaded`,
/// in llama.cpp's own words. Absent where nothing was asked, which is a
/// provider with no server or a server that is not running: the phone
/// draws nothing rather than guessing "unloaded".
#[serde(skip_serializing_if = "Option::is_none")]
status: Option<String>,
}
#[derive(serde::Serialize)]
#[serde(rename_all = "camelCase")]
struct ServerView {
running: bool,
/// Where this backend reaches it, which is the local end of the tunnel for
/// a remote machine. Shown because it is the one fact that makes a running
/// server checkable from outside the app.
#[serde(skip_serializing_if = "Option::is_none")]
port: Option<u16>,
}
async fn provider_view(
State(manager): State<Arc<SessionManager>>,
UrlPath((id, provider_name)): UrlPath<(String, String)>,
) -> Result<axum::Json<ProviderView>, ApiError> {
let machine = machine_by_id(&manager, &id)?;
let provider = provider_on(&machine, &provider_name)?.clone();
let transport = crate::session::transport::Transport::for_machine(&machine);
let offered = crate::machines::provider_models(&transport, &provider, manager.models_dir())
.await
.map_err(from_machine)?;
let router = manager.router_for(&machine, &provider);
// Asked of the server only when there is one running: a router that is
// down has no opinion about which models are loaded, and inventing
// "unloaded" for each would be an answer nobody checked. On a blocking
// thread because asking is an HTTP request, which for a remote machine
// goes down the tunnel.
let held: std::collections::HashMap<String, String> = match router.clone() {
Some(router) if router.endpoint().is_some() => tokio::task::spawn_blocking(move || {
router
.loaded()
.into_iter()
.map(|model| (model.model, model.status))
.collect()
})
.await
.map_err(|err| ApiError::BadRequest(err.to_string()))?,
_ => std::collections::HashMap::new(),
};
let models = offered
.into_iter()
.map(|model| ProviderModel {
settings: provider
.model_settings
.get(&model.id)
.cloned()
.unwrap_or_default(),
status: held.get(&model.id).cloned(),
id: model.id,
label: model.label,
})
.collect();
Ok(axum::Json(ProviderView {
machine: machine.name.clone(),
name: provider.name.clone(),
kind: provider.kind,
command: provider.command.clone(),
model_params: provider.kind.model_params(),
max_loaded: provider.max_loaded,
models,
mcp_servers: provider
.mcp_servers
.iter()
.map(|server| server.name.clone())
.collect(),
server: router.map(|router| ServerView {
running: router.endpoint().is_some(),
port: router.record().and_then(|record| match record.detail {
crate::session::process::Detail::Http { port }
| crate::session::process::Detail::Shared { port } => Some(port),
crate::session::process::Detail::Stdio { .. } => None,
}),
}),
}))
}
/// What the provider itself takes, as opposed to what one of its models does.
#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct ProviderSettingsRequest {
/// Absent or null is the default -- one model loaded at a time, which is
/// the right answer for a machine with one GPU.
#[serde(default)]
max_loaded: Option<u32>,
}
async fn set_provider_settings(
State(manager): State<Arc<SessionManager>>,
UrlPath((id, provider_name)): UrlPath<(String, String)>,
axum::Json(body): axum::Json<ProviderSettingsRequest>,
) -> Result<axum::Json<MachineInfo>, ApiError> {
let machine = manager
.set_provider_settings(&id, &provider_name, Some(body.max_loaded), None)
.map_err(bad_request)?;
Ok(axum::Json(info_for(machine)))
}
/// How one model is loaded, which is the machine's business rather than any
/// session's: one copy of it in memory serves every session that names it.
#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct ModelSettingsRequest {
model: String,
/// The settings whole, like a session's params: what is absent is unset
/// rather than unchanged, so a save cannot leave a value nobody can see.
params: std::collections::BTreeMap<String, String>,
}
async fn set_model_settings(
State(manager): State<Arc<SessionManager>>,
UrlPath((id, provider_name)): UrlPath<(String, String)>,
axum::Json(body): axum::Json<ModelSettingsRequest>,
) -> Result<axum::Json<MachineInfo>, ApiError> {
let machine = manager
.set_provider_settings(
&id,
&provider_name,
None,
Some((body.model.clone(), body.params)),
)
.map_err(bad_request)?;
// Saved first, then told to the machine: a save that reached the config is
// one a phone can rely on having made, and the server may be unreachable
// for reasons that have nothing to do with it. What the machine is told is
// read back from what was saved, so the two cannot describe the model
// differently.
let provider = provider_on(&machine, &provider_name)?.clone();
if let Some(router) = manager.router_for(&machine, &provider) {
let transport = crate::session::transport::Transport::for_machine(&machine);
let models_dir = manager.models_dir().to_path_buf();
let settings = provider
.model_settings
.get(&body.model)
.cloned()
.unwrap_or_default();
// On a blocking thread: it asks the serving machine where the model
// is, which over ssh is a round trip.
tokio::task::spawn_blocking(move || {
crate::session::llama::apply_model_settings(
&router,
&transport,
&models_dir,
&body.model,
&settings,
)
})
.await
.map_err(|err| ApiError::BadRequest(err.to_string()))?
.map_err(from_machine)?;
}
Ok(axum::Json(info_for(machine)))
}
/// Ends the machine's shared server, and with it every model it was holding.
///
/// Offered rather than done automatically, and said plainly on the button:
/// sessions using it will report as exited and their next message will load
/// their model again. It is the only thing that frees the memory, which is why
/// hiding it would just move the decision somewhere with no warning attached.
async fn stop_provider_server(
State(manager): State<Arc<SessionManager>>,
UrlPath((id, provider_name)): UrlPath<(String, String)>,
) -> Result<StatusCode, ApiError> {
let machine = machine_by_id(&manager, &id)?;
let provider = provider_on(&machine, &provider_name)?.clone();
let router = manager
.router_for(&machine, &provider)
.ok_or_else(|| bad_request(anyhow::anyhow!("{provider_name} runs no server of its own")))?;
router.stop().map_err(bad_request)?;
Ok(StatusCode::NO_CONTENT)
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct UnloadRequest {
model: String,
}
/// Takes one model out of memory, leaving the server and every other model
/// alone -- the cheap half of the button above.
async fn unload_provider_model(
State(manager): State<Arc<SessionManager>>,
UrlPath((id, provider_name)): UrlPath<(String, String)>,
axum::Json(body): axum::Json<UnloadRequest>,
) -> Result<StatusCode, ApiError> {
let machine = machine_by_id(&manager, &id)?;
let provider = provider_on(&machine, &provider_name)?.clone();
let router = manager
.router_for(&machine, &provider)
.ok_or_else(|| bad_request(anyhow::anyhow!("{provider_name} loads no models")))?;
tokio::task::spawn_blocking(move || router.unload(&body.model))
.await
.map_err(|err| ApiError::BadRequest(err.to_string()))?
.map_err(from_machine)?;
Ok(StatusCode::NO_CONTENT)
}
/// What is in a directory, and what that directory resolved to.
async fn list_dir(
State(manager): State<Arc<SessionManager>>,
server/src/session/llama/mod.rs
+250 -388
View File
@@ -1,16 +1,26 @@
//! 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.
//! The llama.cpp driver: a session that talks to a model its machine is
//! serving, over an OpenAI-compatible HTTP API, 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.
//! **It has no process of its own.** The machine's models are served by one
//! `llama-server` in router mode, shared by every session on that machine and
//! outliving this backend -- [`router`] is all of that. A session asks it to
//! load a model and then addresses that model by name; several sessions on one
//! model are several conversations against one copy of it in memory. What a
//! session records in its own directory is the router's pid as a
//! [`process::Detail::Shared`], so that "is the thing I am talking to still
//! there?" has the same answer here as for every other driver, while ending
//! this session ends nothing anybody else is using.
//!
//! The router is reached over the same [`Transport`] as any other process --
//! "run this" plus "reach this port" -- which 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
@@ -39,7 +49,9 @@
//! **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`].
//! into it waits rather than failing -- see [`Serving`]. A session joining a
//! model another session already loaded passes through it in an instant,
//! which is the whole benefit of sharing one.
use std::collections::HashMap;
use std::path::{Path, PathBuf};
@@ -54,16 +66,18 @@ use super::driver::{
AttachmentRef, Driver, Event, EventSink, QuestionOption, SessionStatus, Unqueued,
};
use super::process;
use super::transport::{Launch, Streams, Transport};
use super::transport::{Launch, Transport};
use crate::config::{ProviderConfig, SessionConfig};
mod mcp;
pub mod router;
mod tools;
pub use tools::{DEFAULT_MODE as DEFAULT_PERMISSION_MODE, MODES as PERMISSION_MODES};
use mcp::McpServer;
use tools::Tools;
use router::Router;
use tools::{Chosen, Tools};
/// The sampling half of a session's settings, in the wire's own names.
///
@@ -109,12 +123,6 @@ fn chosen_thinking(params: &std::collections::BTreeMap<String, String>) -> Optio
/// servers, so a session here reaches the same thing that UI does.
pub const EXA_MCP_URL: &str = "https://mcp.exa.ai/mcp";
/// The spawn parameter that turns speculative decoding off for a session whose
/// model would otherwise use it. `"off"` and nothing else, because there is
/// only one thing to say: the model either has a head or it does not, and this
/// is the escape for a machine where drafting turns out not to pay.
const SPECULATIVE: &str = "speculative";
/// The parameter naming how hard the model should think, which is a chat
/// template argument rather than a server flag -- so unlike the flags it takes
/// effect on the next request. `"off"` asks the template for no thinking at
@@ -129,9 +137,15 @@ const THINKING_LEVELS: &[&str] = &["low", "medium", "high", "xhigh", "max"];
/// `enable_thinking: false`, a different argument to the template.
const THINKING_OFF: &str = "off";
/// The spawn parameter naming which built-in tools a session gets, as
/// `llama-server`'s own comma-separated list. Absent is all of them, and
/// `"none"` is the way to ask for a session that only talks.
/// The parameter naming which built-in tools a session offers its model, as a
/// comma-separated list of `llama-server`'s own names. Absent is all of them,
/// and `"none"` is the way to ask for a session that only talks.
///
/// A filter applied here rather than a flag passed to the server: the router
/// hosts one set of tools for the machine, and what goes into a request is
/// what this session chose. So it costs no reload -- and it is worth choosing,
/// because the definitions of all seven are most of a prompt on a small
/// window: 2,181 tokens against 698 with none, measured 2026-09-19.
const TOOLS: &str = "tools";
/// How many times one message may go round the call-a-tool loop.
@@ -143,11 +157,6 @@ const TOOLS: &str = "tools";
/// 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
@@ -240,7 +249,7 @@ enum Serving {
/// Started, not answering yet. Anything sent now waits here.
Loading,
Ready {
endpoint: String,
serves: Serves,
tools: Arc<Tools>,
},
/// It exited, or never came up. Carries what to tell somebody, because by
@@ -249,6 +258,35 @@ enum Serving {
Failed(String),
}
/// Where a session's model is, which is a router and a name rather than an
/// address on its own: one `llama-server` serves every model its machine has
/// loaded, and which one a request means is the `model` field in it.
///
/// The two travel together everywhere, because either without the other is a
/// request to the wrong model -- and on a shared server "the wrong model" is
/// another session's conversation.
#[derive(Debug, Clone)]
struct Serves {
endpoint: String,
model: String,
}
impl Serves {
fn url(&self, path: &str) -> String {
format!("{}{path}", self.endpoint)
}
/// The same address as a query, for the `GET` endpoints that take it
/// there rather than in a body.
fn query(&self, path: &str) -> String {
format!(
"{}{path}?model={}",
self.endpoint,
crate::models::urlencode(&self.model)
)
}
}
/// 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.
@@ -268,6 +306,10 @@ struct Respawn {
provider: ProviderConfig,
transport: Transport,
models_dir: PathBuf,
/// The machine's shared `llama-server`. Held rather than looked up each
/// time, so a session keeps talking to the one it started against even
/// while the machine is being edited.
router: Arc<Router>,
}
/// Everything the driver's own threads need, which is nearly all of it: a turn
@@ -316,6 +358,14 @@ struct Shared {
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>>>>,
/// Which of the machine's tools this session offers its model. Live like
/// the sampling settings and for the same reason -- it is applied to the
/// next request rather than to anything that was started.
tools_wanted: Mutex<Chosen>,
/// Whether a thread is already watching the shared server for this
/// session. One is enough, and a model change would otherwise add another
/// every time -- each reporting the same exit to the same transcript.
watching: AtomicBool,
/// How hard this session asks the model to think: a level, `"off"`, or
/// `None` for the model's own default. Live like the sampling settings,
/// and for the same reason -- it rides on the next request.
@@ -338,19 +388,21 @@ pub struct LlamaDriver {
}
impl LlamaDriver {
/// Takes charge of this session's `llama-server`: the one already loaded if
/// there is one, otherwise a new one.
/// Puts this session behind its machine's `llama-server`, asking for its
/// model to be loaded there.
///
/// 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.
/// One entry point, for the reason `ClaudeDriver::launch` gives. What it
/// costs is different here and mostly somebody else's: the model may
/// already be in memory because another session asked for it, in which
/// case this is a round trip, and the expensive case is the first session
/// to want a model nobody has loaded.
#[allow(clippy::too_many_arguments)]
pub fn launch(
meta: &SessionConfig,
provider: &ProviderConfig,
transport: &Transport,
models_dir: &Path,
router: Arc<Router>,
transcript: &Path,
session_dir: &Path,
sink: EventSink,
@@ -391,6 +443,8 @@ impl LlamaDriver {
),
allowed: Mutex::new(allowances(transcript)),
asked: Mutex::new(HashMap::new()),
tools_wanted: Mutex::new(Chosen::from(meta.params.get(TOOLS).map(String::as_str))),
watching: AtomicBool::new(false),
thinking: Mutex::new(chosen_thinking(&meta.params)),
thinking_options: Mutex::new(None),
}),
@@ -399,19 +453,24 @@ impl LlamaDriver {
provider: provider.clone(),
transport: transport.clone(),
models_dir: models_dir.to_path_buf(),
router,
},
};
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.
/// Asks this session's machine to have its model loaded, and starts
/// watching for that to have happened.
///
/// 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.
/// rather than read `respawn.meta`: a switch calls this, so the two ways a
/// session comes to have a model are one piece of code and cannot drift.
///
/// The slow half runs on a thread of its own because it is genuinely slow
/// -- a model coming off disk -- and because a *second* session naming a
/// model that is already loaded must not wait behind the first one's load
/// to find that out.
fn start(&self, model: &str) -> Result<()> {
let shared = &self.shared;
let Respawn {
@@ -419,7 +478,11 @@ impl LlamaDriver {
provider,
transport,
models_dir,
router,
} = &self.respawn;
// Asked of the machine that will serve it, before anything is started:
// a model that is not there says so rather than becoming a server that
// will not load one.
let found = model_on(transport, models_dir, model)?;
// Loading is slow enough to be worth its own state: the session shows
@@ -433,42 +496,43 @@ impl LlamaDriver {
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,
)?
};
// How this model is loaded belongs to the model rather than to the
// session, because one loaded copy is what several sessions share.
let settings = provider
.model_settings
.get(model)
.cloned()
.unwrap_or_default();
let router = Arc::clone(router);
let session = meta.id.clone();
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) => {
let ready = router.load(&model, &found, &settings).and_then(|endpoint| {
// What this session is reaching its model through, recorded
// where every other driver records its process -- as a
// `Shared`, which is what keeps this session's end from
// ending every other session's model. Written before the
// tools are asked for, so a phone that looks during the round
// trip sees a session with something behind it.
if let Some(record) = router.shared_record() {
process::write(&shared.session_dir, &record);
}
let serves = Serves {
endpoint,
model: model.clone(),
};
let tools = Tools::discover(&serves.endpoint, shared.mcp.clone())?;
Ok((serves, tools))
});
let settled = match ready {
Ok((serves, tools)) => {
tracing::info!(
"{model} loaded and answering at {endpoint} with {} tools",
tools.offered().map_or(0, |offered| offered.len()),
"session {session} talking to {model} at {} with {} tools",
serves.endpoint,
tools
.offered(&shared.tools_wanted.lock().unwrap())
.map_or(0, |offered| offered.len()),
);
// Asked now rather than carried from the spawn flags: a
// session that named no context size gets the model's
@@ -476,7 +540,7 @@ impl LlamaDriver {
// named an impossible one gets whatever it settled for.
// Either way this is the measurement rather than the
// request.
if let Some(window) = context_window(&endpoint) {
if let Some(window) = context_window(&serves) {
shared.emit(Event::ContextWindow { tokens: window });
}
// Asked of the model that is now loaded, for the same
@@ -484,10 +548,10 @@ impl LlamaDriver {
// is carrying that this model cannot take is said here
// rather than at the next turn, which is where it would
// otherwise surface as the model simply not doing it.
*shared.thinking_options.lock().unwrap() = Some(thinking_options(&endpoint));
*shared.thinking_options.lock().unwrap() = Some(thinking_options(&serves));
shared.note_unusable_thinking();
Serving::Ready {
endpoint: endpoint.clone(),
serves,
tools: Arc::new(tools),
}
}
@@ -503,13 +567,13 @@ impl LlamaDriver {
Serving::Failed(why)
}
};
let ready = matches!(settled, Serving::Ready { .. });
let serving = matches!(settled, Serving::Ready { .. });
shared.settle(settled);
if ready {
if serving {
let _ = shared.sink.send(Event::Status {
state: SessionStatus::Idle,
});
watch(shared.session_dir.clone(), shared.sink.clone());
watch(Arc::clone(&shared), router);
}
});
Ok(())
@@ -558,7 +622,7 @@ impl Shared {
/// 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>)> {
fn await_ready(&self) -> Result<(Serves, Arc<Tools>)> {
let serving = self
.serving
.lock()
@@ -568,7 +632,7 @@ impl Shared {
.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::Ready { serves, tools } => Ok((serves.clone(), Arc::clone(tools))),
Serving::Failed(why) => bail!("{why}"),
// `wait_while` does not return while this holds.
Serving::Loading => unreachable!("waited out of Loading"),
@@ -619,151 +683,32 @@ impl Shared {
let _ = self.sink.send(event);
}
}
/// Runs a `llama-server` for this session and records it, returning where it
/// is reached from here.
/// Puts a model's settings in front of the machine already serving it.
///
/// 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,
/// The half of a settings change that is not storage: the machine's server
/// reads how to load a model out of its preset file, so a change nobody wrote
/// there is one that takes effect at some unpredictable later point -- the
/// next time a session happened to start. Written now, the model is unloaded
/// and the sessions using it load it again, with the new settings, on their
/// next message.
///
/// Blocking: it asks the serving machine where the model is, which for a
/// remote one is an ssh round trip.
pub fn apply_model_settings(
router: &Router,
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. All of them
// unless the session says otherwise, because whether a particular call
// should happen is the permission gate's question rather than a flag's.
// They run on the machine serving the model, which is the machine the
// files are on.
//
// One slot, not the four `llama-server` picks on its own. A session is
// one conversation making one request at a time -- the driver holds a
// second message until the turn ends -- so the other three are context
// this session could have been given and was not.
//
// It is also what decides whether the MTP head below is worth having.
// Measured 2026-09-19 on the 27B here: 41.5 tok/s plain at any slot
// count, **61.4** with the head at one slot, and **28** with the head
// at four. Speculation against a split KV cache is slower than not
// speculating at all, which is a much bigger effect than the head
// itself and reads exactly like the head being broken.
"-np".into(),
"1".into(),
];
// Settable because it is not free: the definitions of all seven are around
// 2,000 tokens of every prompt -- measured at 2,191 against 1,322 for two
// of them -- which on a small context window is a quarter of it spent
// before anything is said.
//
// "none" omits the flag rather than passing it on: `--tools none` is
// `tools setup failed: unknown tool "none"` and a server that exits, since
// the argument is a list of tool names and no-tools is what having no flag
// means.
match meta.params.get(TOOLS).map(|chosen| chosen.trim()) {
Some("none") => {}
chosen => {
args.push("--tools".into());
args.push(
chosen
.filter(|c| !c.is_empty())
.unwrap_or("all")
.to_string(),
);
}
models_dir: &Path,
key: &str,
settings: &std::collections::BTreeMap<String, String>,
) -> Result<()> {
// Nothing to tell a server that is not running, and nothing to look up:
// it reads the file when it starts, and the file is written by the first
// session to ask for this model.
if router.endpoint().is_none() {
return Ok(());
}
// A model that carries a multi-token-prediction head drafts with it, which
// is most of a 50% speed-up for free -- the tensors are in the file
// whether or not they are used, and without the flag `llama-server` says
// "unused tensor blk.N.nextn.* -- ignoring" and leaves them there.
//
// Conditional because it cannot be otherwise: asked for on a model without
// one, `llama-server` **exits** ("context type MTP requested but model
// doesn't contain MTP layers"), which is a session that never starts. The
// answer comes from the file itself -- see `Model::mtp`.
if found.mtp && meta.params.get(SPECULATIVE).map(String::as_str) != Some("off") {
args.push("--spec-type".into());
args.push("draft-mtp".into());
}
// Settings that belong to the server because they decide how the model
// is loaded; the sampling ones ride on each request instead, so changing
// them later needn't reload anything.
for (key, flag) in [
("contextSize", "-c"),
("gpuLayers", "-ngl"),
("threads", "-t"),
// How far ahead the draft head guesses. Not defaulted here: 2 measured
// 7% faster than llama.cpp's 3 on this machine's GPU, once, which is
// a reason to make the knob reachable and not a reason to move it for
// everybody.
("specDraftNMax", "--spec-draft-n-max"),
] {
if let Some(value) = meta.params.get(key) {
args.push(flag.to_string());
args.push(value.clone());
}
}
let program = provider.program();
let launch = Launch::new(program, args, meta.cwd.as_deref()).reaching(forward);
// Its output goes to files, not pipes. Not only so the process can
// outlive this server: nothing ever read those pipes, so a chatty
// llama-server filled the 64 KB buffer and blocked mid-load with no sign
// of why.
let child = transport.spawn(
&launch,
Streams::Detached {
stdin: std::process::Stdio::null(),
stdout: log_file(&session_dir.join(SERVER_LOG))?.into(),
stderr: log_file(&session_dir.join(SERVER_LOG))?.into(),
},
)?;
let pid = child
.id()
.context("llama-server exited before it could be recorded")?;
tracing::info!(
"session {} running {program} for {model} {} on 127.0.0.1:{} there, \
reached at 127.0.0.1:{} here, as pid {pid}",
meta.id,
transport.describe(),
forward.there,
forward.here,
);
// Reaped so it does not become a zombie while this server is still its
// parent; the health poll and the record are what say whether the
// session is alive, because after a restart there is no `Child` to ask.
tokio::spawn(async move {
let mut child = child;
let _ = child.wait().await;
});
// The *near* port, because that is the one anything reaching this
// server has to dial -- including a later run of this backend,
// which adopts the record without knowing which machine the server
// is on. For a remote session the recorded pid is the ssh
// client's, which is the process this machine owns and which holds
// the tunnel open for exactly as long as the far server lives.
let record = process::Record::of(pid, process::Detail::Http { port: forward.here })
.context("llama-server was gone before its start time could be read")?;
process::write(session_dir, &record);
Ok(format!("http://127.0.0.1:{}", forward.here))
let found = model_on(transport, models_dir, key)?;
router.describe_model(key, &found, settings)
}
/// Connects to every MCP server this provider names, dropping the ones that
@@ -832,46 +777,39 @@ fn allowances(transcript: &Path) -> std::collections::HashSet<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.
/// Reports the shared 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
/// 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) {
///
/// The session's own record is what is polled, and the router's is what says
/// whether its going was asked for. Both, because they answer different halves:
/// this session is only watching while it has a record of its own, and "somebody
/// stopped the machine's llama-server" is not news of a crash.
fn watch(shared: Arc<Shared>, router: Arc<Router>) {
if shared.watching.swap(true, Ordering::SeqCst) {
return;
}
std::thread::spawn(move || {
let session_dir = shared.session_dir.clone();
let sink = shared.sink.clone();
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,
None => break,
Some((_, process::Liveness::Dead)) => {
if !process::stopping(&session_dir) {
if !process::stopping(router.dir()) {
let _ = sink.send(Event::Error {
message: "llama-server exited".to_string(),
});
@@ -880,7 +818,12 @@ fn watch(session_dir: PathBuf, sink: EventSink) {
state: SessionStatus::Exited,
});
process::clear(&session_dir);
return;
// Whatever was waiting on this model is waiting on a
// server that has gone, and nothing else will wake it.
shared.settle(Serving::Failed(
"the machine's llama-server is no longer running.".to_string(),
));
break;
}
Some((_, process::Liveness::Unknown)) => {
let _ = sink.send(Event::Status {
@@ -889,9 +832,10 @@ fn watch(session_dir: PathBuf, sink: EventSink) {
}
}
if sink.is_closed() {
return;
break;
}
}
shared.watching.store(false, Ordering::SeqCst);
});
}
@@ -943,7 +887,7 @@ impl LlamaDriver {
// 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)) => {
Ok((serves, tools)) => {
shared.emit(Event::Status {
state: SessionStatus::Running,
});
@@ -953,7 +897,7 @@ impl LlamaDriver {
// 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) {
if let Err(err) = converse(&shared, &serves, &tools, messages) {
shared.emit(Event::Error {
message: format!("{err:#}"),
});
@@ -987,7 +931,7 @@ impl LlamaDriver {
/// thing this one built.
fn converse(
shared: &Arc<Shared>,
endpoint: &str,
serves: &Serves,
tools: &Tools,
mut messages: Vec<Message>,
) -> Result<()> {
@@ -998,7 +942,7 @@ fn converse(
// Read per call rather than per turn, so a sampling change made while
// a long turn is running reaches the rest of it.
let sampling = shared.sampling.lock().unwrap().clone();
let reply = generate(endpoint, &messages, tools, &sampling, shared)?;
let reply = generate(serves, &messages, tools, &sampling, shared)?;
let calls = reply.calls;
messages.push(Message {
tool_calls: calls.iter().map(Call::wire).collect(),
@@ -1038,7 +982,7 @@ fn run_call(shared: &Arc<Shared>, tools: &Tools, call: &Call) -> String {
tool: call.name.clone(),
input: arguments.clone(),
});
let output = if !tools.knows(&call.name) {
let output = if !tools.knows(&call.name, &shared.tools_wanted.lock().unwrap()) {
// The model invented one. Told plainly, because the alternative is a
// silent empty result it reads as the tool having done nothing.
format!(
@@ -1050,7 +994,12 @@ fn run_call(shared: &Arc<Shared>, tools: &Tools, call: &Call) -> String {
} else if !permitted(shared, call) {
tools::REFUSED.to_string()
} else {
match tools.execute(&call.name, &arguments, shared.cwd.as_deref()) {
match tools.execute(
&call.name,
&arguments,
shared.cwd.as_deref(),
&shared.tools_wanted.lock().unwrap(),
) {
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
@@ -1190,49 +1139,23 @@ impl Driver for LlamaDriver {
// is called, and the rename has already happened where the name lives.
fn set_title(&self, _title: &str) {}
/// Takes new settings: the sampling half now, and says so about the rest.
/// Takes new settings, all of which ride on the next request.
///
/// The split is what [`crate::config::ParamSpec::restart`] describes, and
/// it is said out loud rather than left to the screen, because the screen
/// can only say what a setting *usually* does -- this is the one place
/// that knows whether this session's server was started with the old
/// value. A session already stopped needs no such note: its next start
/// will read all of them.
/// Nothing here waits for a restart, which is a property of what a session
/// now owns rather than a coincidence: everything that decides how a model
/// is *loaded* belongs to the model on its machine, because one loaded
/// copy is what several sessions are talking to. See
/// [`crate::config::LLAMA_MODEL_PARAMS`].
///
/// Which tools it offers is stored rather than asked for: the catalog is
/// the machine's and does not change, and what a session shows its model
/// is decided when the request is built.
fn set_params(&self, params: &std::collections::BTreeMap<String, String>) {
*self.shared.sampling.lock().unwrap() = sampling_from(params);
*self.shared.thinking.lock().unwrap() = chosen_thinking(params);
*self.shared.tools_wanted.lock().unwrap() =
Chosen::from(params.get(TOOLS).map(String::as_str));
self.shared.note_unusable_thinking();
// Only the settings that actually differ from what this session's
// server was started with. Listing every restart-only one on every
// save would be a wall of text about nothing having changed.
let waiting: Vec<&str> = crate::config::DriverKind::LlamaCpp
.params()
.iter()
.filter(|spec| {
spec.restart && params.get(spec.key) != self.respawn.meta.params.get(spec.key)
})
.map(|spec| spec.label)
.collect();
// Nothing to say to a session with no server: its next start reads all
// of them, which is what the note would have been asking for.
let running = matches!(&*self.shared.serving.lock().unwrap(), Serving::Ready { .. });
if !waiting.is_empty() && running {
let one = waiting.len() == 1;
self.shared.emit(Event::Error {
message: format!(
"{} {} saved. {} when this session's server next starts -- stop and start \
the session, or change its model, to load {} now.",
waiting.join(", "),
if one { "is" } else { "are" },
if one {
"It takes effect"
} else {
"They take effect"
},
if one { "it" } else { "them" },
),
});
}
}
fn set_permission_mode(&self, mode: &str) {
@@ -1254,13 +1177,15 @@ impl Driver for LlamaDriver {
});
}
/// Puts this session on a different model, by loading one.
/// Puts this session on a different model, by asking for that 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.
/// Nothing is stopped: the machine's server holds whichever models it has
/// been asked for, and the one this session is leaving may be somebody
/// else's. What it costs is a load where nobody had that model open, and a
/// round trip where somebody did. The conversation survives either because
/// it 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
@@ -1275,10 +1200,6 @@ impl Driver for LlamaDriver {
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
@@ -1323,25 +1244,22 @@ impl Driver for LlamaDriver {
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*.
/// Stops generating and leaves the machine's server alone.
fn detach(&self) {
self.shared.cancel.store(true, Ordering::SeqCst);
self.shared.abandon_questions();
}
/// Ends this session's claim on the model and nothing else.
///
/// The server holding it is the machine's, shared with every other session
/// on it, so stopping one session unloads nothing -- see
/// [`process::Detail::Shared`], which is what makes that structural rather
/// than a rule to remember. A model is taken out of memory from the
/// machine's provider settings, where what it costs everybody is visible.
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);
}
}
@@ -1470,7 +1388,7 @@ fn model_path(models_dir: &Path, key: &str) -> Result<PathBuf> {
/// 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 {
pub struct Model {
/// Absolute, on that machine.
path: String,
/// Whether it carries a multi-token-prediction head, which decides one
@@ -1566,73 +1484,13 @@ fn mtp_in_prefix(head: &str) -> bool {
.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;
/// How many tokens this server can hold, from the server itself.
///
/// `/props` reports the per-slot context, which is the whole of it because
/// this driver always starts one slot -- see the `-np` argument. `None` for a
/// server that would not answer, which draws as no ceiling rather than as a
/// guessed one.
fn context_window(endpoint: &str) -> Option<u64> {
ureq::get(format!("{endpoint}/props"))
/// `/props` reports the per-slot context, which is the whole of it for a model
/// loaded with one slot -- see `router`'s `parallel`. `None` for a server that
/// would not answer, which draws as no ceiling rather than as a guessed one.
fn context_window(serves: &Serves) -> Option<u64> {
ureq::get(serves.query("/props"))
.call()
.ok()?
.body_mut()
@@ -1658,9 +1516,9 @@ fn context_window(endpoint: &str) -> Option<u64> {
/// An empty answer is a model that takes neither, which is a thing to say
/// rather than a failure; a server that will not answer gives the same, since
/// a setting nobody can check is one nobody should be told worked.
fn thinking_options(endpoint: &str) -> Vec<String> {
fn thinking_options(serves: &Serves) -> Vec<String> {
let mut options = Vec::new();
let Some(props) = ureq::get(format!("{endpoint}/props"))
let Some(props) = ureq::get(serves.query("/props"))
.call()
.ok()
.and_then(|mut response| response.body_mut().read_json::<Value>().ok())
@@ -1670,8 +1528,8 @@ fn thinking_options(endpoint: &str) -> Vec<String> {
// Both renders have to have worked: a template that *refuses*
// `enable_thinking` also differs from the plain render, and reading that as
// support would offer an "off" that fails every turn.
let plain = render_template(endpoint, &json!({}));
let off = render_template(endpoint, &json!({"enable_thinking": false}));
let plain = render_template(serves, &json!({}));
let off = render_template(serves, &json!({"enable_thinking": false}));
if plain.is_some() && off.is_some() && off != plain {
options.push(THINKING_OFF.to_string());
}
@@ -1684,7 +1542,7 @@ fn thinking_options(endpoint: &str) -> Vec<String> {
THINKING_LEVELS
.iter()
.filter(|level| {
render_template(endpoint, &json!({"reasoning_effort": level})).is_some()
render_template(serves, &json!({"reasoning_effort": level})).is_some()
})
.map(|level| (*level).to_string()),
);
@@ -1698,12 +1556,13 @@ fn thinking_options(endpoint: &str) -> Vec<String> {
/// `/apply-template` is the cheap half of a request: it renders and returns,
/// with no model involved, so asking it seven questions at load time costs
/// nothing anybody waits for.
fn render_template(endpoint: &str, kwargs: &Value) -> Option<String> {
ureq::post(format!("{endpoint}/apply-template"))
fn render_template(serves: &Serves, kwargs: &Value) -> Option<String> {
ureq::post(serves.url("/apply-template"))
.config()
.http_status_as_error(false)
.build()
.send_json(json!({
"model": serves.model,
"messages": [{"role": "user", "content": "hi"}],
"chat_template_kwargs": kwargs,
}))
@@ -1752,19 +1611,22 @@ fn thinking_kwargs(shared: &Shared) -> Option<Value> {
/// [`Event::ThinkingDone`] the moment the model says something else, which is
/// how long it thought for.
fn generate(
endpoint: &str,
serves: &Serves,
messages: &[Message],
tools: &Tools,
sampling: &serde_json::Map<String, Value>,
shared: &Shared,
) -> Result<Reply> {
let mut body = json!({
// Which model, because one `llama-server` is serving every model this
// machine has loaded and this is how a request says which it means.
"model": serves.model,
"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() {
if let Some(offered) = tools.offered(&shared.tools_wanted.lock().unwrap()) {
map.insert("tools".to_string(), json!(offered));
}
for (key, value) in sampling {
@@ -1780,7 +1642,7 @@ fn generate(
shared.emit(Event::Status {
state: SessionStatus::Reading,
});
let mut response = ureq::post(format!("{endpoint}/v1/chat/completions"))
let mut response = ureq::post(serves.url("/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
server/src/session/llama/router.rs
+857
View File
@@ -0,0 +1,857 @@
//! One `llama-server` per machine, in **router mode**: the front door to every
//! model that machine serves, shared by every session on it.
//!
//! A router holds no weights itself. It reads a preset file naming models and
//! their flags, and starts a child `llama-server` per model that is asked for
//! -- so "one server per model, with that model's own settings" is what a
//! machine ends up running, and one process, one port and one record is what
//! this backend has to keep track of. That is the whole reason it is here:
//! before 2026-09-19 each session started its own `llama-server`, so two
//! sessions on one model held two copies of it in memory and a model change
//! cost a load that only that session benefited from.
//!
//! **A router outlives this backend, and nothing here stops it on its own.**
//! It is recorded the way a session's process is ([`process`]), adopted again
//! on the way back up, and ended only when somebody asks for that in the
//! machine's provider settings. A loaded model is minutes of disk and
//! gigabytes of memory; letting the last session to be closed throw that away
//! would make the shared server pointless.
//!
//! **The preset file is the configuration, and it lives on the serving
//! machine.** Flags that decide how a model is loaded -- context size, layers
//! on the GPU, slots, the draft head -- are per model rather than per session,
//! because one loaded model is what several sessions are now talking to. They
//! are written into a section named by the model's key, which is also the name
//! a request routes by, so nothing has to translate between the two.
//!
//! What is deliberately *not* here: which tools a session offers, how hard it
//! thinks, and the sampling settings. Those ride on each request, so they stay
//! the session's own and need no reload -- see `super`'s module comment.
use std::collections::{BTreeMap, HashMap};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use anyhow::{Context, Result, bail};
use serde::Serialize;
use serde_json::{Value, json};
use super::Model;
use crate::config::{MachineConfig, ProviderConfig};
use crate::session::process;
use crate::session::transport::{Launch, Streams, Transport};
/// Where a router's own output goes, both streams into one file. For a remote
/// machine this is the local end of the ssh connection, so it carries what the
/// far `llama-server` said -- including what a child model instance said,
/// which is the only account of a model that would not load.
const LOG: &str = "llama-router.log";
/// The preset file's name in the router's own directory, for a router on this
/// machine. One on another machine keeps it over there instead, at
/// [`REMOTE_PRESET`], since that is the only side that can read it.
const PRESET: &str = "models.ini";
/// The preset file's first line, which `llama-server` refuses a file without.
const VERSION: &str = "version = 1\n";
/// 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 model that will never
/// answer rather than one that is slow.
const LOAD_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300);
/// How long to wait for the router process itself. Short: it loads nothing,
/// so anything beyond a second or two is a port it cannot bind or a program
/// too old for one of these flags.
const START_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
/// How often either of those is checked.
const POLL: std::time::Duration = std::time::Duration::from_millis(250);
/// How much of the router's log to carry into a message somebody reads on a
/// phone.
const LOG_TAIL_LINES: usize = 6;
/// How many models a router keeps loaded at once before evicting the least
/// recently used. One by default because the machine serving models usually
/// has one GPU: a second model loaded beside the first is the case where
/// neither fits.
const DEFAULT_MAX_LOADED: u32 = 1;
/// Every machine's router, so that two sessions on one machine reach one
/// process rather than starting two.
///
/// A registry rather than a field on each session: the sharing *is* the
/// point, and a router that two drivers could each own is one that both would
/// start.
pub struct Routers {
dir: PathBuf,
/// The runtime a router is started on and reaped into.
///
/// Captured here because everything that starts one runs on a *blocking*
/// thread -- loading a model is minutes of disk, so it cannot be on the
/// runtime -- and tokio's `Command::spawn` registers the child with the
/// reactor, so calling it outside a runtime context panics. That panic is
/// silent: it kills the loading thread and leaves the session saying
/// "loading" for ever, with nothing in the log, which is exactly how it
/// was found.
runtime: Option<tokio::runtime::Handle>,
inner: Mutex<HashMap<String, Arc<Router>>>,
}
impl Routers {
/// `dir` is where each router's record, log and (for this machine) preset
/// file live -- beside the session directories, since a router is shared
/// by sessions and belongs to none of them.
///
/// Made on the runtime that will outlive it; `None` is a test with no
/// runtime at all, where there is nothing to reap into either.
pub fn new(dir: PathBuf) -> Self {
Self {
dir,
runtime: tokio::runtime::Handle::try_current().ok(),
inner: Mutex::new(HashMap::new()),
}
}
/// This machine-and-provider's router, made if this is the first ask.
///
/// The machine and provider are re-read every time rather than captured,
/// because both are editable while sessions are running: a renamed
/// machine, a re-probed program, a changed `maxLoaded`. What a *running*
/// router was started with is whatever it was started with; the new value
/// reaches the next start, which is the same rule every other launch flag
/// follows.
pub fn of(&self, machine: &MachineConfig, provider: &ProviderConfig) -> Arc<Router> {
let key = format!("{}/{}", machine.id, provider.name);
let spec = Spec {
transport: Transport::for_machine(machine),
program: provider.program().to_string(),
max_loaded: provider.max_loaded.unwrap_or(DEFAULT_MAX_LOADED),
};
let mut routers = self.inner.lock().unwrap_or_else(|e| e.into_inner());
let router = routers.entry(key.clone()).or_insert_with(|| {
Arc::new(Router {
dir: self.dir.join(key.replace('/', "-")),
spec: Mutex::new(spec.clone()),
runtime: self.runtime.clone(),
gate: Mutex::new(()),
})
});
*router.spec.lock().unwrap_or_else(|e| e.into_inner()) = spec;
Arc::clone(router)
}
}
/// What it takes to start a router, as its machine currently describes it.
#[derive(Clone)]
struct Spec {
transport: Transport,
program: String,
max_loaded: u32,
}
pub struct Router {
/// Its record and log, on this machine whichever machine it serves from.
dir: PathBuf,
spec: Mutex<Spec>,
/// See [`Routers::runtime`].
runtime: Option<tokio::runtime::Handle>,
/// Held while a router is started and while the preset file is edited --
/// the two things that go wrong when two sessions do them at once. Never
/// held across a model load, which takes minutes.
gate: Mutex<()>,
}
impl Router {
/// Where its record and log are, for a session that wants to watch the
/// process it is talking through or read what it last said.
pub fn dir(&self) -> &Path {
&self.dir
}
/// Its process record, if one is running -- which is also how a session
/// records the process it reaches its model through.
pub fn record(&self) -> Option<process::Record> {
process::live(&self.dir)
}
/// The same process, recorded as one a session reaches but does not own.
///
/// The one place that conversion happens, so that a session directory can
/// never come to hold a record saying it owns the machine's router -- see
/// [`process::Detail::Shared`].
pub fn shared_record(&self) -> Option<process::Record> {
let record = self.record()?;
match record.detail {
process::Detail::Http { port } => Some(process::Record {
detail: process::Detail::Shared { port },
..record
}),
process::Detail::Stdio { .. } | process::Detail::Shared { .. } => None,
}
}
/// Where to reach it, or `None` when nothing is running.
pub fn endpoint(&self) -> Option<String> {
match self.record()?.detail {
process::Detail::Http { port } => Some(format!("http://127.0.0.1:{port}")),
process::Detail::Stdio { .. } | process::Detail::Shared { .. } => None,
}
}
/// Puts `model` in memory and says where to talk to it, starting the
/// router first if it is not up.
///
/// Blocking, and slow on purpose: a model that has to come off disk takes
/// as long as it takes. The caller is the driver's loading thread, which
/// is what [`super::Serving::Loading`] exists to describe.
///
/// Nothing is held across that wait. Two sessions load through one router,
/// and the second one wanting a model already in memory must not queue
/// behind the first one's cold load of a different model -- which is most
/// of what sharing a server was for.
pub fn load(
&self,
key: &str,
found: &Model,
settings: &BTreeMap<String, String>,
) -> Result<String> {
let endpoint = self.start_if_down()?;
self.describe_model(key, found, settings)?;
// Nothing to ask for where it is already in memory, which is the case
// this whole module exists to produce: a second session naming a model
// somebody else loaded is a round trip rather than a load. Asking
// anyway is not harmless -- `POST /models/load` answers **400** for a
// model that is already loaded, which arrived as a session that
// refused to start next to one happily using that same model.
if !self.is_ready(key) {
let asked = self
.post("/models/load", json!({ "model": key }))
.with_context(|| format!("asking llama-server to load {key}"));
match (asked, self.wait_loaded(key)) {
// Loaded, whatever the request said: something else may have
// asked for it in the meantime, and what is in memory is the
// answer rather than what one request made of being told to
// put it there.
(_, Ok(())) => {}
// It did not load, and a refusal of the request itself says
// more about why than "it never appeared" does.
(Err(refused), Err(_)) => return Err(refused),
(Ok(_), Err(never)) => return Err(never),
}
}
Ok(endpoint)
}
/// Writes what this model is and how to load it into the preset file, and
/// has the router re-read it.
///
/// Also the path a settings change takes, which is why it is separate from
/// [`load`](Self::load): a model whose entry has changed is **unloaded**
/// by the re-read, and that is the change taking effect rather than a
/// side effect -- the sessions using it load it again, with the new
/// settings, on their next message. What must not happen is the same
/// thing to an unrelated model, which is why the file is written and the
/// re-read asked for only when the text actually differs.
pub fn describe_model(
&self,
key: &str,
found: &Model,
settings: &BTreeMap<String, String>,
) -> Result<()> {
// Read, edit, write: under the gate because two of those at once lose
// one of the two sections.
let _one_at_a_time = self.gate.lock().unwrap_or_else(|e| e.into_inner());
let existing = self.preset()?;
let updated = upsert(&existing, key, &section(found, settings));
if updated == existing {
return Ok(());
}
self.write_preset(&updated)?;
// Only meaningful against a running router; one that is down reads the
// file when it starts.
if self.endpoint().is_some() {
self.get("/models?reload=1")
.context("asking llama-server to re-read its models")?;
}
Ok(())
}
/// Whether this model is in memory now.
fn is_ready(&self, key: &str) -> bool {
self.loaded()
.into_iter()
.any(|model| model.model == key && model.ready)
}
/// Every model this router knows about and what each is doing, or an empty
/// list when it is not running.
pub fn loaded(&self) -> Vec<RouterModel> {
let Ok(answer) = self.get("/models") else {
return Vec::new();
};
answer
.get("data")
.and_then(Value::as_array)
.map(|models| models.iter().filter_map(RouterModel::read).collect())
.unwrap_or_default()
}
/// Takes one model out of memory, leaving the router and every other model
/// alone.
pub fn unload(&self, key: &str) -> Result<()> {
self.post("/models/unload", json!({ "model": key }))
.with_context(|| format!("asking llama-server to unload {key}"))?;
Ok(())
}
/// Ends the router and every model it is holding.
///
/// The only thing that does: a session closing, being deleted, or this
/// backend shutting down all leave it running. Sessions using it will see
/// their process go and report `exited`, which is true -- the model they
/// were talking to is no longer in memory.
///
/// Neither the record nor the mark is cleared here, and that is what tells
/// those sessions this was asked for rather than a crash. A session's
/// watcher looks a couple of seconds later, so anything removed now is
/// removed before the only reader of it has looked -- which is how the
/// first version of this put "llama-server exited" in three transcripts
/// belonging to somebody who had just pressed Stop. The record describes a
/// dead process, which every reader already handles, and starting a new
/// router is what clears both.
pub fn stop(&self) -> Result<()> {
let Some(record) = self.record() else {
return Ok(());
};
process::mark_stopping(&self.dir)?;
process::stop(&record, process::STOP_GRACE);
Ok(())
}
/// The end of the router's 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 machine it is
/// the only half, since nobody reading the phone can open a file over
/// there. Bounded, because this ends up in an event a phone draws.
pub fn log_tail(&self) -> String {
let Ok(text) = std::fs::read_to_string(self.dir.join(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(" / ")
)
}
/// Adopts the running router or starts one, and waits for it to answer.
///
/// Under the gate, and asked again inside it: two sessions starting at once
/// would otherwise both find no record, both start a router, and bind two
/// ports to the same models.
fn start_if_down(&self) -> Result<String> {
if let Some(endpoint) = self.endpoint() {
return Ok(endpoint);
}
let _one_at_a_time = self.gate.lock().unwrap_or_else(|e| e.into_inner());
if let Some(endpoint) = self.endpoint() {
return Ok(endpoint);
}
let spec = self.spec.lock().unwrap_or_else(|e| e.into_inner()).clone();
wg_app_link::private::create_dir(&self.dir)?;
// Whatever the last router left behind, including the mark saying its
// end was asked for -- see [`stop`](Self::stop). From here on, a
// process that goes away is news.
process::clear(&self.dir);
// Written before the router starts, because it is read at startup and
// a router with no preset file at all lists nothing.
let preset = self.write_preset(&self.preset()?)?;
let forward = spec
.transport
.reserve_port()
.context("finding a port for llama-server")?;
let args = vec![
// Loopback there, whichever machine there is: what reaches it from
// outside that machine is the ssh tunnel and nothing else.
"--host".to_string(),
"127.0.0.1".to_string(),
"--port".to_string(),
forward.there.to_string(),
// No `-m`: a `llama-server` given no model is a router.
"--models-preset".to_string(),
preset,
"--models-max".to_string(),
spec.max_loaded.to_string(),
// The built-in agent tools -- read, search, edit, shell. Hosted by
// the router itself, which is what makes them one set for the
// machine rather than one per model. Which of them a *session*
// offers its model is decided here in the backend, per request, so
// there is nothing per-session to pass through: see `super::tools`.
"--tools".to_string(),
"all".to_string(),
];
let launch = Launch::new(&spec.program, args, None).reaching(forward);
// Starting and reaping both happen inside the runtime, though this is
// a blocking thread: tokio's `Command::spawn` registers the child with
// the reactor, so calling it outside a runtime context panics -- and
// that panic kills only this thread, leaving a session that says
// "loading" for ever with nothing in the log. See [`Routers::runtime`].
let _inside = self.runtime.as_ref().map(tokio::runtime::Handle::enter);
// Its output goes to a file, not a pipe. 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 with no sign of why.
let child = spec.transport.spawn(
&launch,
Streams::Detached {
stdin: std::process::Stdio::null(),
stdout: log_file(&self.dir.join(LOG))?.into(),
stderr: log_file(&self.dir.join(LOG))?.into(),
},
)?;
let pid = child
.id()
.context("llama-server exited before it could be recorded")?;
tracing::info!(
"running {} in router mode {} on 127.0.0.1:{} there, reached at 127.0.0.1:{} here, \
as pid {pid}",
spec.program,
spec.transport.describe(),
forward.there,
forward.here,
);
// Reaped so it does not become a zombie while this server is still
// its parent; the record and the health poll are what say whether it
// is alive, because after a restart there is no `Child` to ask.
if let Some(runtime) = &self.runtime {
runtime.spawn(async move {
let mut child = child;
let _ = child.wait().await;
});
}
// The *near* port, because that is the one anything reaching this
// router has to dial -- including a later run of this backend, which
// adopts the record without knowing which machine it is on. For a
// remote machine 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 router 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(&self.dir, &record);
let endpoint = format!("http://127.0.0.1:{}", forward.here);
self.wait_answering(&endpoint)?;
Ok(endpoint)
}
/// Polls until the router answers, watching the process as well as the
/// port: a port already taken, or a `llama-server` too old for one of
/// these flags, exits within a second and would otherwise be waited out.
fn wait_answering(&self, endpoint: &str) -> Result<()> {
let deadline = std::time::Instant::now() + START_TIMEOUT;
let url = format!("{endpoint}/health");
loop {
if let Ok(response) = ureq::get(&url).call()
&& response.status() == 200
{
return Ok(());
}
match process::recorded(&self.dir) {
Some((_, process::Liveness::Alive | process::Liveness::Unknown)) => {}
Some((_, process::Liveness::Dead)) | None => {
process::clear(&self.dir);
bail!("llama-server exited before it answered.{}", self.log_tail());
}
}
if std::time::Instant::now() > deadline {
process::clear(&self.dir);
bail!(
"llama-server did not answer within {}s.{}",
START_TIMEOUT.as_secs(),
self.log_tail()
);
}
std::thread::sleep(POLL);
}
}
/// Polls until the named model is in memory, or says why it will never be.
///
/// A model that will not load is the common failure and it is fast: the
/// child exits, the router reports it unloaded with an exit code, and this
/// says so rather than waiting out the timeout -- which is what turned "it
/// said the file was corrupt" into "gave up after 300s".
fn wait_loaded(&self, key: &str) -> Result<()> {
let deadline = std::time::Instant::now() + LOAD_TIMEOUT;
loop {
match self.loaded().into_iter().find(|model| model.model == key) {
Some(model) if model.ready => return Ok(()),
Some(model) if model.failed => {
bail!("{key} would not load.{}", self.log_tail())
}
// Still loading, or not listed yet after a reload.
Some(_) | None => {}
}
if process::live(&self.dir).is_none() {
bail!(
"llama-server went away while loading {key}.{}",
self.log_tail()
);
}
if std::time::Instant::now() > deadline {
bail!(
"{key} was still loading after {}s.{}",
LOAD_TIMEOUT.as_secs(),
self.log_tail()
);
}
std::thread::sleep(POLL);
}
}
/// One request to the router's management endpoints. The generation ones
/// are the driver's and stream, so they are not these.
fn get(&self, path: &str) -> Result<Value> {
let url = format!("{}{path}", self.answering()?);
Self::read(
ureq::get(&url)
.call()
.with_context(|| format!("GET {path}"))?,
path,
)
}
fn post(&self, path: &str, body: Value) -> Result<Value> {
let url = format!("{}{path}", self.answering()?);
Self::read(
ureq::post(&url)
.send_json(body)
.with_context(|| format!("POST {path}"))?,
path,
)
}
/// Where to reach a router that is running, or the failure saying it is
/// not -- which is what every one of these requests needs first.
fn answering(&self) -> Result<String> {
self.endpoint().context("no llama-server is running")
}
fn read(mut response: ureq::http::Response<ureq::Body>, path: &str) -> Result<Value> {
response
.body_mut()
.read_json()
.with_context(|| format!("reading what {path} answered"))
}
/// The preset file as it stands on the machine that serves the models, or
/// empty where there is none yet.
///
/// Read back rather than remembered, for one reason that matters after a
/// restart: a router adopted from a previous run is already serving models
/// whose sections this process has never seen, and rewriting the file
/// without them would unload them at the next reload.
fn preset(&self) -> Result<String> {
let spec = self.spec.lock().unwrap_or_else(|e| e.into_inner()).clone();
let text = match &spec.transport {
Transport::Here => {
Ok(std::fs::read_to_string(self.dir.join(PRESET)).unwrap_or_default())
}
Transport::Ssh { name, .. } => {
let script = format!("p={REMOTE_PRESET}; cat \"$p\" 2>/dev/null || true");
let launch = Launch::new("sh", vec!["-c".to_string(), script], None);
spec.transport
.capture_blocking(&launch)
.with_context(|| format!("reading the model settings on {name}"))
}
}?;
// A file that is not there yet reads as a new one rather than as
// nothing: `llama-server` refuses a preset with no version line, so
// "empty" is not a state this can hand back or write.
Ok(if text.trim().is_empty() {
VERSION.to_string()
} else {
text
})
}
/// Writes the preset file where the router will read it, and says where
/// that is -- which is the path the router is given, so the two cannot
/// disagree.
fn write_preset(&self, text: &str) -> Result<String> {
let spec = self.spec.lock().unwrap_or_else(|e| e.into_inner()).clone();
match &spec.transport {
Transport::Here => {
let path = self.dir.join(PRESET);
wg_app_link::private::write_file(&path, text.as_bytes())?;
Ok(path.to_string_lossy().into_owned())
}
Transport::Ssh { name, .. } => {
use base64::Engine as _;
// Base64 rather than a heredoc: the text goes through a shell
// on the far side, and an INI value is not something to trust
// to quoting rules twice over.
let encoded = base64::engine::general_purpose::STANDARD.encode(text);
let script = format!(
"p={REMOTE_PRESET}; mkdir -p \"$(dirname \"$p\")\" && \
printf %s \"$1\" | base64 -d > \"$p\" && printf '%s\\n' \"$p\""
);
let launch = Launch::new(
"sh",
vec!["-c".to_string(), script, "sh".to_string(), encoded],
None,
);
let answer = spec
.transport
.capture_blocking(&launch)
.with_context(|| format!("writing the model settings on {name}"))?;
match answer.trim() {
"" => bail!("{name} did not say where it wrote the model settings"),
path => Ok(path.to_string()),
}
}
}
}
}
/// Where the preset file goes on a machine that is not this one, as a shell
/// word the far side expands: under that machine's state directory, beside
/// whatever else belongs to this app there.
///
/// `$HOME` is resolved over there because only that machine knows what it is.
const REMOTE_PRESET: &str = "\"${XDG_STATE_HOME:-$HOME/.local/state}/ai-app/llama-models.ini\"";
/// One model the router knows about, as the provider view and the load poll
/// both read it.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct RouterModel {
pub model: String,
/// The router's own word: `unloaded`, `loading`, `loaded`, `sleeping`.
/// Carried through rather than reduced to a boolean, because the phone
/// draws it and llama.cpp is the authority on what states there are.
pub status: String,
pub ready: bool,
/// Unloaded *and* something went wrong, which is not the same as unloaded.
#[serde(skip)]
pub failed: bool,
}
impl RouterModel {
fn read(entry: &Value) -> Option<Self> {
let model = entry.get("id")?.as_str()?.to_string();
let status = entry
.pointer("/status/value")
.and_then(Value::as_str)
.unwrap_or("unknown")
.to_string();
let exited = entry
.get("exit_code")
.and_then(Value::as_i64)
.unwrap_or_default();
Some(Self {
ready: status == "loaded" || status == "sleeping",
failed: status == "unloaded" && exited != 0,
model,
status,
})
}
}
/// What a model's section says, from the file it is and the settings it has.
///
/// The keys are `llama-server`'s own argument names without their dashes,
/// which is what a preset section is: `--n-gpu-layers 20` is
/// `n-gpu-layers = 20`. So adding a setting is a row here and a row in
/// [`crate::config::LLAMA_MODEL_PARAMS`], and nothing in between.
fn section(found: &Model, settings: &BTreeMap<String, String>) -> String {
let mut lines = vec![format!("model = {}", found.path)];
for (key, flag) in [
("contextSize", "ctx-size"),
("gpuLayers", "n-gpu-layers"),
("threads", "threads"),
// How far ahead the draft head guesses. Not defaulted: 2 measured 7%
// faster than llama.cpp's 3 on this machine's GPU, once, which is a
// reason to make the knob reachable and not a reason to move it for
// everybody.
("specDraftNMax", "spec-draft-n-max"),
] {
if let Some(value) = settings.get(key).map(|value| value.trim())
&& !value.is_empty()
{
lines.push(format!("{flag} = {value}"));
}
}
// One slot unless this model is told otherwise. A session is one
// conversation making one request at a time, and a second session's turn
// waits rather than splitting the cache: measured 2026-09-19 on the 27B
// here, 41.5 tok/s plain at any slot count, **61.4** with the draft head
// at one slot, and **28** with the head at four. Speculating against a
// split KV cache is slower than not speculating at all.
let slots = settings
.get("slots")
.map(|value| value.trim())
.filter(|value| !value.is_empty())
.unwrap_or("1");
lines.push(format!("parallel = {slots}"));
// A model carrying 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. 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"). See `Model::mtp`.
if found.mtp && settings.get("speculative").map(String::as_str) != Some("off") {
lines.push("spec-type = draft-mtp".to_string());
}
lines.join("\n")
}
/// The preset file with `name`'s section replaced by `body`, added at the end
/// if it was not there.
///
/// Text in and text out, rather than a parsed model, because the file belongs
/// to `llama-server` rather than to this: anything in it that this does not
/// understand -- a `[*]` section, a key added by a later version, a comment
/// somebody wrote -- has to survive being edited.
fn upsert(existing: &str, name: &str, body: &str) -> String {
let header = format!("[{name}]");
let mut out = String::new();
let mut skipping = false;
let mut replaced = false;
for line in existing.lines() {
let trimmed = line.trim();
if trimmed.starts_with('[') && trimmed.ends_with(']') {
skipping = trimmed == header;
if skipping {
replaced = true;
push_section(&mut out, &header, body);
continue;
}
}
if !skipping {
out.push_str(line);
out.push('\n');
}
}
if out.is_empty() {
out.push_str(VERSION);
}
if !replaced {
push_section(&mut out, &header, body);
}
out
}
fn push_section(out: &mut String, header: &str, body: &str) {
if !out.ends_with("\n\n") && !out.is_empty() {
out.push('\n');
}
out.push_str(header);
out.push('\n');
out.push_str(body.trim_end());
out.push_str("\n\n");
}
/// An owner-only log opened for appending, so the two streams pointed at it do
/// not overwrite each other and an adopted router 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()))
}
#[cfg(test)]
mod tests {
use super::*;
fn settings(pairs: &[(&str, &str)]) -> BTreeMap<String, String> {
pairs
.iter()
.map(|(key, value)| ((*key).to_string(), (*value).to_string()))
.collect()
}
#[test]
fn a_section_names_the_file_and_the_flags_that_were_set() {
let found = Model {
path: "/models/a.gguf".to_string(),
mtp: true,
};
let text = section(
&found,
&settings(&[("contextSize", "8192"), ("threads", " 6 ")]),
);
assert_eq!(
text,
"model = /models/a.gguf\nctx-size = 8192\nthreads = 6\nparallel = 1\n\
spec-type = draft-mtp"
);
// A blank is not a value: it is the setting being unset, and passing
// it on is a child that exits on an empty argument.
let text = section(&found, &settings(&[("contextSize", " ")]));
assert!(!text.contains("ctx-size"), "{text}");
// The draft head is asked for only where the file has one, and can be
// turned off for a machine where it does not pay.
let plain = Model {
mtp: false,
..found.clone()
};
assert!(!section(&plain, &settings(&[])).contains("spec-type"));
assert!(!section(&found, &settings(&[("speculative", "off")])).contains("spec-type"));
}
#[test]
fn a_section_replaces_its_own_and_leaves_every_other_line_alone() {
let first = upsert("", "repo/a.gguf", "model = /models/a.gguf");
assert!(first.starts_with("version = 1\n"), "{first}");
assert!(
first.contains("[repo/a.gguf]\nmodel = /models/a.gguf\n"),
"{first}"
);
// A second model is added rather than replacing the first, because a
// reload of a file that lost a section unloads that model -- which
// would be one session taking another's model out of memory.
let both = upsert(&first, "repo/b.gguf", "model = /models/b.gguf");
assert!(both.contains("[repo/a.gguf]"), "{both}");
assert!(both.contains("[repo/b.gguf]"), "{both}");
// Editing one rewrites only its own keys, and keeps what llama.cpp's
// own file has that this does not know about.
let with_global = format!(
"version = 1\n\n[*]\njinja = true\n\n{}",
both.trim_start_matches("version = 1\n")
);
let edited = upsert(
&with_global,
"repo/a.gguf",
"model = /models/a.gguf\nctx-size = 4096",
);
assert!(edited.contains("[*]\njinja = true"), "{edited}");
assert!(edited.contains("ctx-size = 4096"), "{edited}");
assert_eq!(edited.matches("[repo/a.gguf]").count(), 1, "{edited}");
assert!(
edited.contains("[repo/b.gguf]\nmodel = /models/b.gguf"),
"{edited}"
);
}
#[test]
fn writing_the_same_settings_twice_changes_nothing() {
// What keeps an unrelated session's model in memory: the file is only
// written, and the router only told to re-read it, when the text
// actually differs.
let once = upsert("", "repo/a.gguf", "model = /models/a.gguf");
assert_eq!(upsert(&once, "repo/a.gguf", "model = /models/a.gguf"), once);
}
}
server/src/session/llama/tools.rs
+142 -14
View File
@@ -1,8 +1,13 @@
//! What a llama session can do besides talk, and who runs it.
//!
//! Two sources, one list. `llama-server` started with `--tools` runs a set of
//! its own -- reading, searching, editing, a shell -- and publishes them at
//! Two sources, one list. The machine's `llama-server` runs a set of its own
//! -- reading, searching, editing, a shell -- and publishes them at
//! `GET /tools` in the shape a model is given, with `POST /tools` to run one.
//! That set belongs to the server rather than to any one session, so **which
//! of them a session offers is a [`Chosen`] applied at each request**: a
//! filter rather than a flag, which is what keeps one session's choice off
//! every other session sharing that server -- and what makes changing it take
//! effect on the next message rather than on a reload.
//! Anything else comes from an MCP server this backend is connected to (see
//! [`super::mcp`]). Both arrive here as a definition to offer and a way to
//! call, and nothing downstream of [`Tools::execute`] knows which a tool was.
@@ -42,8 +47,9 @@ pub struct Tools {
/// changes, because that is a different server on a different port.
endpoint: String,
/// What the model is given, in the order it is offered: the server's own
/// tools first, then each MCP server's.
definitions: Vec<Value>,
/// tools first, then each MCP server's. Named as well, because which of
/// them a session offers is decided per request -- see [`offered`](Self::offered).
definitions: Vec<(String, Value)>,
/// The server's tools, and whether each is run relative to a working
/// directory. Only the ones that say so are sent one -- a tool that
/// ignores it would still have its cache keyed on it.
@@ -55,7 +61,8 @@ pub struct Tools {
impl Tools {
/// Asks a ready `llama-server` what it offers and adds what the MCP
/// servers offered.
/// servers offered. Everything it has, rather than what one session wants:
/// the choosing is [`Chosen`]'s, per request.
///
/// A session with no built-in tools -- or none at all -- is a perfectly
/// good session, so nothing here treats "no tools" as a failure. What *is*
@@ -99,11 +106,11 @@ impl Tools {
.and_then(Value::as_bool)
.unwrap_or(false),
);
definitions.push(definition.clone());
definitions.push((name.to_string(), definition.clone()));
}
for connected in &mcp {
for tool in connected.lock().unwrap().tools() {
definitions.push(tool.definition.clone());
definitions.push((tool.qualified.clone(), tool.definition.clone()));
}
}
Ok(Self {
@@ -121,14 +128,30 @@ impl Tools {
/// template branches on whether tools were given, and an empty list
/// renders the whole "you may call one or more functions" preamble with no
/// functions under it.
pub fn offered(&self) -> Option<&[Value]> {
(!self.definitions.is_empty()).then_some(&self.definitions)
///
/// `chosen` names which of the *server's* tools this session offers; an
/// MCP server's are the session's own to begin with, since they were
/// configured against this provider rather than found on the machine.
pub fn offered(&self, chosen: &Chosen) -> Option<Vec<&Value>> {
let offered: Vec<&Value> = self
.definitions
.iter()
.filter(|(name, _)| !self.server.contains_key(name) || chosen.takes(name))
.map(|(_, definition)| definition)
.collect();
(!offered.is_empty()).then_some(offered)
}
/// Whether this is a tool at all, which decides what to do about a call
/// naming something else.
pub fn knows(&self, name: &str) -> bool {
self.server.contains_key(name) || self.mcp_for(name).is_some()
/// Whether this is a tool this session has, which decides what to do
/// about a call naming something else.
///
/// `chosen` for the same reason [`offered`](Self::offered) takes it: a
/// tool the session did not offer is one the model invented, and being
/// told that is better than being asked for permission to run something
/// that would then be refused. [`execute`](Self::execute) checks it as
/// well, because that is where running it is actually prevented.
pub fn knows(&self, name: &str, chosen: &Chosen) -> bool {
(self.server.contains_key(name) && chosen.takes(name)) || self.mcp_for(name).is_some()
}
/// The MCP server that offered `name`, if one did.
@@ -149,13 +172,24 @@ impl Tools {
/// say they use one. A session with no working directory sends none, and
/// `llama-server` falls back to its own -- which is the honest outcome:
/// this server has no better answer for where "here" is.
pub fn execute(&self, name: &str, arguments: &Value, cwd: Option<&str>) -> Result<String> {
pub fn execute(
&self,
name: &str,
arguments: &Value,
cwd: Option<&str>,
chosen: &Chosen,
) -> Result<String> {
if let Some(server) = self.mcp_for(name) {
return server.lock().unwrap().call(name, arguments);
}
// A tool this session did not offer is not one it may run, even where
// the server has it. A model that names one anyway is guessing, and a
// session whose whole setting was to have no tools must not get a
// shell out of a guess.
let uses_cwd = *self
.server
.get(name)
.filter(|_| chosen.takes(name))
.with_context(|| format!("no tool called {name}"))?;
let mut request = ureq::post(format!("{}/tools", self.endpoint))
.config()
@@ -185,6 +219,46 @@ impl Tools {
}
}
/// Which of the server's own tools a session offers its model.
///
/// Held by the session rather than settled at discovery, so that changing it
/// takes effect on the next message: one shared server has one set of tools,
/// and which of them go into a request is this.
///
/// Three cases rather than a list of names, because two of them are what
/// people actually write: everything, nothing, or these. "Nothing" is the one
/// that has to be sayable at all -- an empty list would be indistinguishable
/// from the setting being unset, which is what `all` means.
pub enum Chosen {
All,
None,
Named(std::collections::HashSet<String>),
}
impl Chosen {
pub fn from(wanted: Option<&str>) -> Self {
match wanted.map(str::trim) {
None | Some("") | Some("all") => Self::All,
Some("none") => Self::None,
Some(list) => Self::Named(
list.split(',')
.map(str::trim)
.filter(|name| !name.is_empty())
.map(str::to_string)
.collect(),
),
}
}
pub fn takes(&self, name: &str) -> bool {
match self {
Self::All => true,
Self::None => false,
Self::Named(names) => names.contains(name),
}
}
}
/// `POST /tools`'s answer as the text a model is given.
///
/// The server answers `plain_text_response` for a tool that ran and `error`
@@ -271,4 +345,58 @@ mod tests {
fn anything_else_is_handed_over_as_itself() {
assert_eq!(result_text(&json!({"rows": 2})), "{\"rows\":2}");
}
/// The vocabulary is `llama-server`'s own, and "none" has to mean it:
/// before the router, that word was a flag the server refused and so a
/// session that never started.
#[test]
fn which_tools_a_session_offers_reads_the_servers_own_words() {
for unset in [None, Some(""), Some(" all ")] {
assert!(Chosen::from(unset).takes("read_file"), "{unset:?}");
}
assert!(!Chosen::from(Some("none")).takes("read_file"));
let two = Chosen::from(Some("read_file, grep_search"));
assert!(two.takes("read_file"));
assert!(two.takes("grep_search"));
assert!(!two.takes("exec_shell_command"));
}
/// The filter is over the *machine's* tools. An MCP server's belong to the
/// session already -- they were configured against this provider rather
/// than found on the machine -- so "none" is a session that still searches
/// the web, and a request with nothing left in it offers no `tools` key at
/// all rather than an empty one.
#[test]
fn a_filtered_catalog_keeps_the_mcp_tools_and_vanishes_when_empty() {
let tools = Tools {
endpoint: "http://127.0.0.1:1".to_string(),
definitions: vec![
("read_file".to_string(), json!({"name": "read_file"})),
("exec_shell_command".to_string(), json!({"name": "shell"})),
("exa_web_search_exa".to_string(), json!({"name": "search"})),
],
server: [
("read_file".to_string(), true),
("exec_shell_command".to_string(), false),
]
.into_iter()
.collect(),
mcp: Vec::new(),
};
let offered = |wanted| {
tools
.offered(&Chosen::from(wanted))
.map(|offered| offered.len())
};
assert_eq!(offered(None), Some(3));
assert_eq!(offered(Some("read_file")), Some(2));
assert_eq!(offered(Some("none")), Some(1));
let nothing = Tools {
definitions: Vec::new(),
..tools
};
assert_eq!(nothing.offered(&Chosen::All), None);
}
}
server/src/session/mod.rs
+99 -2
View File
@@ -20,7 +20,7 @@ pub mod subagent;
pub mod transcript;
pub mod transport;
use std::collections::{HashMap, VecDeque};
use std::collections::{BTreeMap, HashMap, VecDeque};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, RwLock};
use std::time::{SystemTime, UNIX_EPOCH};
@@ -45,6 +45,11 @@ use subagent::Subagents;
use transcript::{SeqEvent, Transcript};
use transport::Transport;
/// Where each machine's shared `llama-server` keeps its record, log and model
/// settings: one directory beside the session directories, since a session id
/// is what everything in there is named by and a router is not a session.
const ROUTERS: &str = "routers";
/// Fan-out buffer per session. A subscriber further behind than this is
/// caught up from the transcript file instead, so the size only bounds
/// memory, not correctness.
@@ -644,6 +649,10 @@ pub struct SessionManager {
/// every echo driver this manager builds is handed a clone -- see
/// [`SessionManager::reporting_usage_fixture`].
usage_fixture: crate::usage::Fixture,
/// Every machine's shared `llama-server` -- see [`llama::router`]. On the
/// manager because the sharing is the point: a registry each session
/// carried its own copy of would be one router per session again.
routers: Arc<llama::router::Routers>,
inner: RwLock<Inner>,
}
@@ -701,6 +710,9 @@ impl SessionManager {
// this manager builds gets a clone, including the ones built
// below, so it has to exist before the first session does.
let usage_fixture = crate::usage::Fixture::new();
// Beside the session directories rather than inside one, because a
// router belongs to a machine and is shared by every session on it.
let routers = Arc::new(llama::router::Routers::new(data_dir.join(ROUTERS)));
let mut live = HashMap::new();
for meta in &config.sessions {
// One unlaunchable session -- a corrupt transcript, an
@@ -715,6 +727,7 @@ impl SessionManager {
data_dir: &data_dir,
models_dir: &models_dir,
usage: &usage_fixture,
routers: &routers,
},
announce.clone(),
// Nothing is started here; see `Launching`.
@@ -737,6 +750,7 @@ impl SessionManager {
pending: Arc::new(pending::Registry::default()),
spawn_throwaway: false,
usage_fixture,
routers,
inner: RwLock::new(Inner { config, live }),
};
Ok(manager)
@@ -756,6 +770,7 @@ impl SessionManager {
data_dir: &self.data_dir,
models_dir: &self.models_dir,
usage: &self.usage_fixture,
routers: &self.routers,
}
}
@@ -898,7 +913,17 @@ impl SessionManager {
if let Some(name) = name {
machine.name = name.trim().to_string();
}
if let Some(providers) = providers {
if let Some(mut providers) = providers {
// A re-probe answers "what is installed here", which is not an
// answer about how this machine's models are loaded: settings
// kept against a provider survive it, by name. Without this,
// pressing Rediscover silently emptied every model's settings.
for provider in &mut providers {
if let Some(old) = machine.provider(&provider.name) {
provider.model_settings = old.model_settings.clone();
provider.max_loaded = old.max_loaded;
}
}
machine.providers = providers;
}
Ok(machine.clone())
@@ -1182,6 +1207,72 @@ impl SessionManager {
self.inner.read().unwrap().config.machines.clone()
}
/// The shared `llama-server` behind a machine's provider, for a provider
/// that has one.
///
/// `None` for every other kind rather than an empty router, because "this
/// provider has no server of its own to look at" is what the machine's
/// provider view has to say, and a router that exists but is never running
/// says something else.
pub fn router_for(
&self,
machine: &MachineConfig,
provider: &ProviderConfig,
) -> Option<Arc<llama::router::Router>> {
match provider.kind {
DriverKind::LlamaCpp => Some(self.routers.of(machine, provider)),
DriverKind::Echo | DriverKind::ClaudeCli | DriverKind::CodexCli => None,
}
}
/// Records how a machine loads one of its models, or how many it keeps
/// loaded at once.
///
/// One function for two routes because it is one edit to one provider
/// entry, and because the alternative is two `update` closures that have
/// to find the same provider the same way. Storage only: telling the
/// running server is `llama::apply_model_settings`, deliberately after
/// this, so a save that reached the config is one the phone can rely on
/// having made whether or not the machine could be reached.
pub fn set_provider_settings(
&self,
machine_id: &str,
provider_name: &str,
max_loaded: Option<Option<u32>>,
model: Option<(String, BTreeMap<String, String>)>,
) -> Result<MachineConfig> {
self.update(|config| {
let machine = config
.machines
.iter_mut()
.find(|machine| machine.id == machine_id)
.with_context(|| format!("no machine with id \"{machine_id}\""))?;
let provider = machine
.providers
.iter_mut()
.find(|provider| provider.name == provider_name)
.with_context(|| format!("no provider \"{provider_name}\" on that machine"))?;
if let Some(max_loaded) = max_loaded {
provider.max_loaded = max_loaded;
}
if let Some((model, params)) = model {
// Removed rather than stored empty: an entry of nothing and no
// entry mean the same thing, and only one of them leaves the
// config file describing models nobody has settings for.
let params: BTreeMap<String, String> = params
.into_iter()
.filter(|(_, value)| !value.trim().is_empty())
.collect();
if params.is_empty() {
provider.model_settings.remove(&model);
} else {
provider.model_settings.insert(model, params);
}
}
Ok(machine.clone())
})
}
pub fn spawn_session(&self, spec: SpawnSpec) -> Result<SessionInfo> {
self.spawn_seeded(spec, None)
}
@@ -2362,6 +2453,9 @@ struct Env<'a> {
data_dir: &'a Path,
models_dir: &'a Path,
usage: &'a crate::usage::Fixture,
/// Every machine's shared `llama-server`, so two sessions on one machine
/// reach one process -- see [`llama::router`].
routers: &'a Arc<llama::router::Routers>,
}
/// Creates the session directory, opens its transcript (continuing the
@@ -2590,6 +2684,7 @@ fn make_driver(
provider,
&Transport::for_machine(machine),
env.models_dir,
env.routers.of(machine, provider),
transcript_path,
dir,
sink.clone(),
@@ -4043,6 +4138,8 @@ mod tests {
command: Some(command.to_string_lossy().into_owned()),
models: Vec::new(),
mcp_servers: Vec::new(),
model_settings: BTreeMap::new(),
max_loaded: None,
},
])],
..Config::default()
server/src/session/process.rs
+65 -9
View File
@@ -56,6 +56,16 @@ pub enum Detail {
/// Spoken to over HTTP on a loopback port, which is all it takes to find
/// it again -- there is no stream to be partway through.
Http { port: u16 },
/// The same, for a process this session reaches but does not own: the
/// llama.cpp router serving every session on its machine.
///
/// A variant rather than a flag because of what it forbids. Liveness is
/// the identical question -- a session whose router has gone has no model
/// -- but ending it is not this session's to ask, and [`signal`] is where
/// that is enforced: stopping, deleting or cleaning up after a session
/// must not take a model out of memory for every other session on that
/// machine.
Shared { port: u16 },
}
/// Whether a recorded process is still there.
@@ -82,6 +92,15 @@ impl Record {
})
}
/// Whether this server may end that process.
///
/// False for the one it shares -- see [`Detail::Shared`]. Liveness is the
/// identical question for both, which is why this is separate from it:
/// "is it there?" and "is it mine to end?" are asked in different places.
pub fn ours(&self) -> bool {
!matches!(self.detail, Detail::Shared { .. })
}
pub fn liveness(&self) -> Liveness {
match stat_of(self.pid) {
// A different start time is a reused pid, so definitely not ours.
@@ -256,10 +275,10 @@ pub const STOP_GRACE: std::time::Duration = std::time::Duration::from_secs(5);
/// a SIGKILL would cost whatever it had not flushed; SIGKILL after the grace
/// period because a session the phone has deleted must not still be running.
pub fn stop(record: &Record, grace: std::time::Duration) {
if record.liveness() != Liveness::Alive {
if !record.ours() || record.liveness() != Liveness::Alive {
return;
}
signal(record.pid, libc::SIGTERM);
signal(record, libc::SIGTERM);
let record = record.clone();
tokio::spawn(async move {
tokio::time::sleep(grace).await;
@@ -285,6 +304,10 @@ pub fn wait_gone(records: &[Record], grace: std::time::Duration) {
let deadline = std::time::Instant::now() + grace;
for record in records {
// Never asked to stop, so there is nothing to wait out.
if !record.ours() {
continue;
}
while record.liveness() == Liveness::Alive && std::time::Instant::now() < deadline {
std::thread::sleep(LOOK);
}
@@ -303,17 +326,25 @@ fn kill_if_still_there(record: &Record, grace: std::time::Duration) {
record.pid,
grace
);
signal(record.pid, libc::SIGKILL);
signal(record, libc::SIGKILL);
}
}
fn signal(pid: u32, signal: libc::c_int) {
/// The one place a session's process is signalled, which is why the refusal to
/// signal a shared one lives here rather than at each caller: every path out of
/// a session -- stopped, deleted, cleaned up on the way down -- ends in this
/// function, and the one that forgot would be a model unloaded under somebody
/// else's turn.
fn signal(record: &Record, signal: libc::c_int) {
if !record.ours() {
return;
}
// SAFETY: `kill` with a positive pid touches only that process, and the pid
// came from a record whose start time was just confirmed to match -- so it
// is still the process this server started, not a reused number. A failure
// (already gone) is nothing to act on.
unsafe {
libc::kill(pid as libc::pid_t, signal);
libc::kill(record.pid as libc::pid_t, signal);
}
}
@@ -424,10 +455,12 @@ mod tests {
write(dir.path(), &record);
assert_eq!(live(dir.path()), Some(record.clone()));
// And the other shape round trips through the same file.
record.detail = Detail::Http { port: 8080 };
write(dir.path(), &record);
assert_eq!(live(dir.path()), Some(record));
// And the other shapes round trip through the same file.
for detail in [Detail::Http { port: 8080 }, Detail::Shared { port: 8080 }] {
record.detail = detail;
write(dir.path(), &record);
assert_eq!(live(dir.path()), Some(record.clone()));
}
mark_stopping(dir.path()).expect("mark stopping");
assert!(stopping(dir.path()));
@@ -463,6 +496,29 @@ mod tests {
assert!(stray.is_empty(), "left behind {stray:?}");
}
/// The whole of what [`Detail::Shared`] is for: a session ending must not
/// take the machine's llama.cpp router with it.
#[tokio::test]
async fn a_shared_process_is_not_stopped_with_the_session_that_reached_it() {
let mut child = std::process::Command::new("sleep")
.arg("30")
.spawn()
.expect("spawn sleep");
let shared = Record::of(child.id(), Detail::Shared { port: 1 }).expect("start time");
stop(&shared, std::time::Duration::from_millis(50));
std::thread::sleep(std::time::Duration::from_millis(200));
assert_eq!(shared.liveness(), Liveness::Alive, "the router was killed");
// The same process, recorded as one this session owns, does stop.
let owned = Record {
detail: Detail::Http { port: 1 },
..shared
};
stop(&owned, std::time::Duration::from_millis(50));
let _ = child.wait();
assert_eq!(owned.liveness(), Liveness::Dead);
}
#[test]
fn a_dead_or_unreadable_record_is_not_live() {
let dir = tempfile::tempdir().expect("tempdir");
server/src/usage.rs
+4
View File
@@ -1109,6 +1109,8 @@ mod tests {
command: None,
models: vec![],
mcp_servers: Vec::new(),
model_settings: Default::default(),
max_loaded: None,
}],
}
}
@@ -1173,6 +1175,8 @@ mod tests {
command: None,
models: vec![],
mcp_servers: Vec::new(),
model_settings: Default::default(),
max_loaded: None,
}];
// A machine with no Claude on it has no Claude limits, and a row
// reporting on it would be a fact about nothing. Echo included: