ai-app/server/src/session/llama.rs

//! The llama.cpp driver: a `llama-server` process per session, spoken to over
//! its OpenAI-compatible HTTP API and translated into the common event model.
//!
//! Two things make this shaped differently from the Claude driver.
//!
//! **It is spawned but not spoken to over stdio.** The process is started
//! through the same [`Transport`] as any other and then reached over HTTP on a
//! loopback port. A remote llama-server would need its port forwarded as well
//! as its command wrapped, which is not built, so a session on an ssh host is
//! refused rather than silently talking to the wrong machine.
//!
//! **The server is stateless between requests**, so the whole conversation goes
//! with every one. It is rebuilt from the session's transcript rather than kept
//! in this struct, which is not tidiness: a copy in driver memory is invisible
//! to a second device and gone when this process restarts.
//!
//! That leaves the Claude driver as the odd one out rather than this one -- the
//! CLI's own memory of a conversation is a cache in front of the same
//! transcript. Resolve any inconsistency in this direction.

use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};

use anyhow::{Context, Result, bail};
use serde::{Deserialize, Serialize};
use serde_json::json;

use super::driver::{AttachmentRef, Driver, Event, EventSink, SessionStatus};
use super::process;
use super::transport::{Launch, Streams, Transport};
use crate::config::{ProviderConfig, SessionConfig};

/// How long to wait for a model to load before giving up. Loading is mostly
/// disk, and a large quantised model on a cold cache is genuinely slow, so this
/// is generous -- the failure it exists for is a server that will never answer.
const READY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300);

/// One turn in the conversation this driver keeps on the server's behalf.
#[derive(Debug, Clone, Serialize, Deserialize)]
struct Message {
    role: String,
    content: String,
}

pub struct LlamaDriver {
    sink: EventSink,
    /// Where this session's own llama-server answers.
    endpoint: String,
    /// Where the conversation is read back from, one line per event.
    transcript: PathBuf,
    /// Sampling settings chosen at spawn, sent with every request.
    sampling: serde_json::Map<String, serde_json::Value>,
    /// Set by [`Driver::interrupt`]; the streaming loop checks it between
    /// chunks and stops, leaving what was generated in the transcript.
    cancel: Arc<AtomicBool>,
    /// Where this session's process record lives, so [`Driver::stop`] can find
    /// the server it has to end.
    session_dir: PathBuf,
}

impl LlamaDriver {
    /// Takes charge of this session's `llama-server`: the one already loaded if
    /// there is one, otherwise a new one.
    ///
    /// One entry point, for the reason `ClaudeDriver::launch` gives, expensive
    /// in a different currency: two servers holding the same model is twice the
    /// memory, and the second would bind a different port while the phone kept
    /// talking to the first.
    pub fn launch(
        meta: &SessionConfig,
        provider: &ProviderConfig,
        transport: &Transport,
        models_dir: &Path,
        transcript: &Path,
        session_dir: &Path,
        sink: EventSink,
    ) -> Result<Self> {
        if !matches!(transport, Transport::Here) {
            bail!(
                "llama.cpp sessions can only run on this machine for now: the model is served \
                 over HTTP, and forwarding that port to another host isn't built yet."
            );
        }
        let model = meta.model.as_deref().context(
            "a llama.cpp session needs a model -- one of the downloaded ones, by its key",
        )?;
        let path = model_path(models_dir, model)?;

        // Already loaded and still running: keep talking to it. The health poll
        // below confirms it is really answering, so adopting a pid whose server
        // has wedged still reports as a failure rather than as a session that
        // silently never replies.
        if let Some(process::Record {
            detail: process::Detail::Http { port },
            pid,
            ..
        }) = process::live(session_dir)
        {
            tracing::info!(
                "session {} reattaching to the llama-server it left loaded (pid {pid}, port {port})",
                meta.id
            );
            return Ok(Self::attached(
                format!("http://127.0.0.1:{port}"),
                meta,
                model,
                transcript,
                session_dir,
                sink,
            ));
        }

        let port = free_port().context("finding a port for llama-server")?;
        let mut args: Vec<String> = vec![
            "-m".into(),
            path.to_string_lossy().into_owned(),
            "--host".into(),
            "127.0.0.1".into(),
            "--port".into(),
            port.to_string(),
        ];
        // Settings that belong to the server because they decide how the model
        // is loaded; the sampling ones ride on each request instead, so changing
        // them later needn't reload anything.
        for (key, flag) in [
            ("contextSize", "-c"),
            ("gpuLayers", "-ngl"),
            ("threads", "-t"),
        ] {
            if let Some(value) = meta.params.get(key) {
                args.push(flag.to_string());
                args.push(value.clone());
            }
        }

        let program = provider.command.as_deref().unwrap_or("llama-server");
        let launch = Launch::new(program, args, meta.cwd.as_deref());
        // 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:{port} as pid {pid}",
            meta.id
        );
        // 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;
        });

        let record = process::Record::of(pid, process::Detail::Http { port })
            .context("llama-server was gone before its start time could be read")?;
        process::write(session_dir, &record);

        Ok(Self::attached(
            format!("http://127.0.0.1:{port}"),
            meta,
            model,
            transcript,
            session_dir,
            sink,
        ))
    }

    /// The driver for a `llama-server` at `endpoint`, however it got there.
    ///
    /// Shared by starting one and adopting one, because everything after "there
    /// is a server at this address" is identical -- including waiting for it to
    /// answer, which an adopted one still owes: a recorded pid says a process
    /// exists, not that its model is loaded.
    fn attached(
        endpoint: String,
        meta: &SessionConfig,
        model: &str,
        transcript: &Path,
        session_dir: &Path,
        sink: EventSink,
    ) -> Self {
        // Loading is slow enough to be worth saying so: the session shows as
        // running until the model is in memory, rather than looking ready and
        // refusing the first message.
        let _ = sink.send(Event::Status {
            state: SessionStatus::Running,
        });
        {
            let sink = sink.clone();
            let endpoint = endpoint.clone();
            let model = model.to_string();
            let session_dir = session_dir.to_path_buf();
            std::thread::spawn(move || match wait_until_ready(&endpoint) {
                Ok(()) => {
                    tracing::info!("{model} loaded and answering at {endpoint}");
                    let _ = sink.send(Event::Status {
                        state: SessionStatus::Idle,
                    });
                    watch(session_dir, sink);
                }
                Err(err) => {
                    let _ = sink.send(Event::Error {
                        message: format!("{model} never became ready: {err:#}"),
                    });
                    let _ = sink.send(Event::Status {
                        state: SessionStatus::Exited,
                    });
                    process::clear(&session_dir);
                }
            });
        }

        let mut sampling = serde_json::Map::new();
        for (key, field) in [
            ("temperature", "temperature"),
            ("topP", "top_p"),
            ("topK", "top_k"),
            ("maxTokens", "max_tokens"),
        ] {
            if let Some(raw) = meta.params.get(key)
                && let Ok(number) = raw.parse::<f64>()
            {
                sampling.insert(field.to_string(), json!(number));
            }
        }

        Self {
            sink,
            endpoint,
            transcript: transcript.to_path_buf(),
            sampling,
            cancel: Arc::new(AtomicBool::new(false)),
            session_dir: session_dir.to_path_buf(),
        }
    }
}

/// Where llama-server's own output goes. One file for both streams: it is
/// diagnostics nobody parses, and interleaving them is how it reads in a
/// terminal anyway.
const SERVER_LOG: &str = "llama-server.log";

/// How often a loaded server is checked for still being there. Slower than the
/// Claude driver's stdout poll because nothing is waiting on it: this only has
/// to notice a server that has gone.
const WATCH_INTERVAL: std::time::Duration = std::time::Duration::from_secs(2);

/// An owner-only log opened for appending, so the two streams pointed at
/// it do not overwrite each other and a reattach keeps what came before.
fn log_file(path: &Path) -> Result<std::fs::File> {
    use std::os::unix::fs::OpenOptionsExt;
    std::fs::OpenOptions::new()
        .create(true)
        .append(true)
        .mode(0o600)
        .open(path)
        .with_context(|| format!("opening {}", path.display()))
}

/// Reports the server going away, for as long as the session is there to report
/// it to.
///
/// Polled rather than waited on, for the reason the Claude driver gives: after a
/// restart this server is not the process's parent, so liveness has to be a
/// question asked of the record -- and asking it two different ways is how the
/// two answers come to disagree.
fn watch(session_dir: PathBuf, sink: EventSink) {
    std::thread::spawn(move || {
        loop {
            std::thread::sleep(WATCH_INTERVAL);
            match process::recorded(&session_dir) {
                Some((_, process::Liveness::Alive)) => {}
                // Nothing recorded means the session was stopped or deleted
                // deliberately, and whoever did that has already said so.
                None => return,
                Some((_, process::Liveness::Dead)) => {
                    let _ = sink.send(Event::Error {
                        message: "llama-server exited".to_string(),
                    });
                    let _ = sink.send(Event::Status {
                        state: SessionStatus::Exited,
                    });
                    process::clear(&session_dir);
                    return;
                }
                Some((_, process::Liveness::Unknown)) => {
                    let _ = sink.send(Event::Status {
                        state: SessionStatus::Unknown,
                    });
                }
            }
            if sink.is_closed() {
                return;
            }
        }
    });
}

impl Driver for LlamaDriver {
    fn send_user_message(&self, text: String, attachments: Vec<AttachmentRef>) {
        if !attachments.is_empty() {
            let _ = self.sink.send(Event::Error {
                message: "this model can't be sent attachments or files".to_string(),
            });
        }
        let sink = self.sink.clone();
        let endpoint = self.endpoint.clone();
        let transcript = self.transcript.clone();
        let sampling = self.sampling.clone();
        let cancel = Arc::clone(&self.cancel);
        cancel.store(false, Ordering::Relaxed);

        // Its own thread: the request blocks for as long as the model takes to
        // generate, which is the whole point of streaming it.
        std::thread::spawn(move || {
            // Nothing is ever held back here -- there is no queue to wait in --
            // so the message is taken the moment it arrives. Said anyway,
            // because this is what records it: see `MessageTaken`.
            let _ = sink.send(Event::MessageTaken {
                id: None,
                text: text.clone(),
                // Never any: this driver refuses attachments above.
                attachments: Vec::new(),
            });
            let _ = sink.send(Event::Status {
                state: SessionStatus::Running,
            });
            // Everything before this message, plus this message. Read rather
            // than remembered, and `text` is appended here rather than waited
            // for, because the message's own transcript entry is still on its
            // way when this runs.
            let mut messages = conversation(&transcript);
            messages.push(Message {
                role: "user".into(),
                content: text,
            });
            // The reply is not stored: the deltas below are the durable record,
            // so the next turn reads back exactly what the phone was shown --
            // including a partial one that was interrupted.
            if let Err(err) = generate(&endpoint, &messages, &sampling, &cancel, &sink) {
                let _ = sink.send(Event::Error {
                    message: format!("{err:#}"),
                });
            }
            let _ = sink.send(Event::Status {
                state: SessionStatus::Idle,
            });
        });
    }

    fn answer_question(&self, _id: &str, _answers: &[String]) {
        // Nothing here asks questions: this driver has no tools.
    }

    fn interrupt(&self) {
        self.cancel.store(true, Ordering::Relaxed);
    }

    // Nothing to forward: this process has no notion of what the conversation
    // is called, and the rename has already happened where the name lives.
    fn set_title(&self, _title: &str) {}

    fn set_permission_mode(&self, _mode: &str) {
        let _ = self.sink.send(Event::Error {
            message: "a llama.cpp session runs no tools, so there is nothing for a permission \
                      mode to govern."
                .to_string(),
        });
    }

    fn set_model(&self, _model: &str) {
        let _ = self.sink.send(Event::Error {
            message: "a llama.cpp session's model is fixed when it starts, because the server \
                      loads one model into memory. Spawn another session to use a different one."
                .to_string(),
        });
    }

    fn run_command(&self, text: &str) {
        let _ = self.sink.send(Event::Error {
            message: format!(
                "a llama.cpp session has no commands of its own, so {text} means nothing to it."
            ),
        });
    }

    fn compact(&self) {
        let _ = self.sink.send(Event::Error {
            message: "llama.cpp has no compaction. Clear the session instead, which costs nothing."
                .to_string(),
        });
    }

    fn clear(&self) {
        // All of it. `conversation` folds from the last of these, so recording
        // the marker *is* the reset -- there is no driver state to keep in step
        // with it, which is the same property that makes a second device see the
        // same conversation this one does.
        let _ = self.sink.send(Event::Cleared);
    }

    /// Stops generating and leaves the server loaded.
    ///
    /// Worth being deliberate about, because the cost points the other way from
    /// the Claude driver's: a `llama-server` holds its whole model in memory, so
    /// a leaked one is gigabytes nobody is using. It is left anyway, because the
    /// alternative is unloading and reloading that model on every backend
    /// restart -- minutes of disk, for a session somebody is in the middle of.
    /// The record is what keeps it from being *nobody's*.
    fn detach(&self) {
        self.cancel.store(true, Ordering::Relaxed);
    }

    fn stop(&self) {
        self.cancel.store(true, Ordering::Relaxed);
        if let Some(record) = process::live(&self.session_dir) {
            process::stop(&record, process::STOP_GRACE);
        }
        process::clear(&self.session_dir);
    }
}

/// The conversation so far, folded out of the transcript.
///
/// Consecutive `AssistantText` deltas are one assistant turn, closed by the next
/// user message -- which is also what makes an interrupted reply come back as
/// the partial text the phone actually saw.
///
/// This must stay a pure function of the transcript and must never re-render
/// earlier turns. llama.cpp caches the prompt prefix, so a growing conversation
/// reprocesses almost nothing -- but only while every turn is byte-identical to
/// last time. Changing how an old turn is rendered silently reprocesses the
/// whole history on every message.
fn conversation(path: &Path) -> Vec<Message> {
    let Ok(events) = crate::session::transcript::read_after(path, 0) else {
        return Vec::new();
    };
    let mut messages: Vec<Message> = Vec::new();
    let mut pending = String::new();
    // Everything before the last clear is still in the transcript and is
    // deliberately not in the conversation. Folding from zero would put it back,
    // which is the whole of what clearing had to undo.
    let events = match events.iter().rposition(|e| e.event == Event::Cleared) {
        Some(at) => &events[at + 1..],
        None => &events[..],
    };
    for event in events.iter().cloned() {
        match event.event {
            Event::UserMessage { text, .. } => {
                if !pending.is_empty() {
                    messages.push(Message {
                        role: "assistant".into(),
                        content: std::mem::take(&mut pending),
                    });
                }
                messages.push(Message {
                    role: "user".into(),
                    content: text,
                });
            }
            Event::AssistantText { delta } => pending.push_str(&delta),
            _ => {}
        }
    }
    if !pending.is_empty() {
        messages.push(Message {
            role: "assistant".into(),
            content: pending,
        });
    }
    messages
}

/// Where a model key resolves to on disk, refusing anything that climbs
/// out of the models directory -- the key arrives from a phone.
fn model_path(models_dir: &Path, key: &str) -> Result<PathBuf> {
    let mut path = models_dir.to_path_buf();
    for part in key.split('/') {
        if part.is_empty() || part == "." || part == ".." {
            bail!("\"{key}\" is not a model key this can resolve");
        }
        path.push(part);
    }
    if !path.is_file() {
        bail!("no downloaded model called \"{key}\" -- download it first");
    }
    Ok(path)
}

/// An unused loopback port, by asking the OS for one and letting it go. Racy in
/// principle, but nothing on this machine is hunting for ports, and the
/// alternative -- parsing the port back out of the server's log -- couples us to
/// its output format for no real gain.
fn free_port() -> Result<u16> {
    let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
    Ok(listener.local_addr()?.port())
}

/// Polls until the server says it is ready, or gives up.
fn wait_until_ready(endpoint: &str) -> Result<()> {
    let deadline = std::time::Instant::now() + READY_TIMEOUT;
    let url = format!("{endpoint}/health");
    loop {
        if let Ok(response) = ureq::get(&url).call()
            && response.status() == 200
        {
            return Ok(());
        }
        if std::time::Instant::now() > deadline {
            bail!("gave up after {}s", READY_TIMEOUT.as_secs());
        }
        std::thread::sleep(std::time::Duration::from_millis(250));
    }
}

/// One streamed completion: posts the conversation, emits each delta as it
/// arrives. Emits rather than returns, because the transcript those events land
/// in is what the next turn reads back.
fn generate(
    endpoint: &str,
    messages: &[Message],
    sampling: &serde_json::Map<String, serde_json::Value>,
    cancel: &AtomicBool,
    sink: &EventSink,
) -> Result<()> {
    let mut body = json!({
        "messages": messages,
        "stream": true,
        "stream_options": {"include_usage": true},
    });
    let map = body.as_object_mut().expect("built as an object");
    for (key, value) in sampling {
        map.insert(key.clone(), value.clone());
    }

    let mut response = ureq::post(format!("{endpoint}/v1/chat/completions"))
        .header("Content-Type", "application/json")
        .send_json(&body)
        .context("asking llama-server to generate")?;

    let reader = std::io::BufReader::new(response.body_mut().as_reader());
    let mut tokens = 0u64;
    // The prompt side only, which is what the model is holding -- the same
    // definition the other dialects report, so one word on the phone means one
    // thing whichever kind of session it is.
    let mut context = None;
    for line in std::io::BufRead::lines(reader) {
        if cancel.load(Ordering::Relaxed) {
            break;
        }
        let line = line.context("reading the generation stream")?;
        // Server-sent events: the payload lines are the ones that matter.
        let Some(payload) = line.strip_prefix("data: ") else {
            continue;
        };
        if payload.trim() == "[DONE]" {
            break;
        }
        let Ok(chunk) = serde_json::from_str::<serde_json::Value>(payload) else {
            continue;
        };
        if let Some(usage) = chunk.get("usage") {
            if let Some(total) = usage
                .get("total_tokens")
                .and_then(serde_json::Value::as_u64)
            {
                tokens = total;
            }
            if let Some(prompt) = usage
                .get("prompt_tokens")
                .and_then(serde_json::Value::as_u64)
            {
                context = Some(prompt);
            }
        }
        let delta = chunk
            .get("choices")
            .and_then(|c| c.get(0))
            .and_then(|c| c.get("delta"))
            .and_then(|d| d.get("content"))
            .and_then(serde_json::Value::as_str)
            .unwrap_or_default();
        if !delta.is_empty() {
            let _ = sink.send(Event::AssistantText {
                delta: delta.to_string(),
            });
        }
    }
    if tokens > 0 {
        let _ = sink.send(Event::UsageDelta { tokens, context });
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::session::transcript::Transcript;

    /// Writes a transcript the way the pump does, so the fold is tested against
    /// the real file format rather than a hand-built vector.
    fn transcript_with(events: &[Event]) -> (tempfile::TempDir, PathBuf) {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("transcript.jsonl");
        let mut transcript = Transcript::open(&path).expect("open");
        for event in events {
            transcript.append(event.clone(), 0.0).expect("append");
        }
        (dir, path)
    }

    #[test]
    fn deltas_between_user_messages_are_one_assistant_turn() {
        let (_dir, path) = transcript_with(&[
            Event::UserMessage {
                id: None,
                text: "hello".into(),
                attachments: Vec::new(),
            },
            Event::AssistantText {
                delta: "hi ".into(),
            },
            Event::AssistantText {
                delta: "there".into(),
            },
            Event::Status {
                state: SessionStatus::Idle,
            },
            Event::UserMessage {
                id: None,
                text: "again".into(),
                attachments: Vec::new(),
            },
            Event::AssistantText {
                delta: "yes".into(),
            },
        ]);
        let messages = conversation(&path);
        assert_eq!(
            messages
                .iter()
                .map(|m| (m.role.as_str(), m.content.as_str()))
                .collect::<Vec<_>>(),
            [
                ("user", "hello"),
                ("assistant", "hi there"),
                ("user", "again"),
                ("assistant", "yes")
            ],
        );
    }

    #[test]
    /// The interrupted case, which decides what a resumed conversation is built
    /// from: whatever the phone was shown. The deltas that arrived before the
    /// stop are in the transcript, so they are in the prompt -- the model is
    /// never told it said something the user did not see.
    fn an_interrupted_reply_stays_in_the_conversation() {
        let (_dir, path) = transcript_with(&[
            Event::UserMessage {
                id: None,
                text: "count".into(),
                attachments: Vec::new(),
            },
            Event::AssistantText {
                delta: "one two".into(),
            },
            Event::Status {
                state: SessionStatus::Idle,
            },
        ]);
        let messages = conversation(&path);
        assert_eq!(messages.len(), 2);
        assert_eq!(messages[1].content, "one two");
    }

    #[test]
    /// Events this driver does not produce must not disturb the fold: a
    /// transcript can carry errors and status changes from a session that
    /// was, say, relaunched.
    fn other_events_are_not_part_of_the_conversation() {
        let (_dir, path) = transcript_with(&[
            Event::Status {
                state: SessionStatus::Running,
            },
            Event::UserMessage {
                id: None,
                text: "hello".into(),
                attachments: Vec::new(),
            },
            Event::Error {
                message: "something went wrong".into(),
            },
            Event::AssistantText {
                delta: "still here".into(),
            },
            Event::UsageDelta {
                tokens: 12,
                context: Some(12),
            },
        ]);
        let messages = conversation(&path);
        assert_eq!(messages.len(), 2);
        assert_eq!(messages[0].content, "hello");
        assert_eq!(messages[1].content, "still here");
    }

    #[test]
    /// Clearing decides what the *model* is given, not just what the phone
    /// draws. Everything above the marker stays in the transcript and none of it
    /// is sent.
    fn the_conversation_starts_after_the_last_clear() {
        let (_dir, path) = transcript_with(&[
            Event::UserMessage {
                id: None,
                text: "the long expensive conversation".into(),
                attachments: Vec::new(),
            },
            Event::AssistantText {
                delta: "at length".into(),
            },
            Event::Cleared,
            Event::UserMessage {
                id: None,
                text: "a fresh start".into(),
                attachments: Vec::new(),
            },
            Event::AssistantText {
                delta: "cheaply".into(),
            },
        ]);
        let messages = conversation(&path);
        assert_eq!(messages.len(), 2);
        assert_eq!(messages[0].content, "a fresh start");
        assert_eq!(messages[1].content, "cheaply");
    }

    #[test]
    /// The *last* one, so clearing twice does not resurrect what the
    /// first clear dropped.
    fn only_the_newest_clear_counts() {
        let (_dir, path) = transcript_with(&[
            Event::UserMessage {
                id: None,
                text: "one".into(),
                attachments: Vec::new(),
            },
            Event::Cleared,
            Event::UserMessage {
                id: None,
                text: "two".into(),
                attachments: Vec::new(),
            },
            Event::Cleared,
            Event::UserMessage {
                id: None,
                text: "three".into(),
                attachments: Vec::new(),
            },
        ]);
        let messages = conversation(&path);
        assert_eq!(messages.len(), 1);
        assert_eq!(messages[0].content, "three");
    }

    #[test]
    fn a_model_key_cannot_climb_out_of_the_models_directory() {
        let dir = tempfile::tempdir().expect("tempdir");
        for attempt in ["../../etc/passwd", "unsloth/../../escape.gguf", ""] {
            assert!(
                model_path(dir.path(), attempt).is_err(),
                "{attempt:?} should have been refused",
            );
        }
    }
}