boxmaker/docs/plans/M2a/files/crates/loopd/tests/device.rs

//! Checks against the real server. Ignored unless run by `make verify-device`. Do not edit.
//!
//! They need two environment variables:
//!   BOXMAKER_INFERPROXY  path to the built `inferproxy` binary
//!   BOXMAKER_UPSTREAM    host:port of the llama-server router, for example straylight:11434
//! and, optionally, BOXMAKER_MODEL (default ornith-1.5-35b-a3b). They use slot 0.
//!
//! They never touch the real `llama-server` process. The "server dies" test kills and restarts
//! its own `inferproxy`.

use loopd::config::Config;
use loopd::llama::info::{CacheOutcome, cache_outcome};
use loopd::llama::{ChatEvent, ChatMessage, ChatRequest, Client};
use std::path::{Path, PathBuf};
use std::process::{Child, Command};
use std::sync::mpsc;
use std::thread;
use std::time::Duration;

struct Proxy {
    child: Child,
    socket: PathBuf,
}

impl Proxy {
    fn start(socket: &Path) -> Proxy {
        let binary = std::env::var("BOXMAKER_INFERPROXY").expect("BOXMAKER_INFERPROXY is not set");
        let upstream = std::env::var("BOXMAKER_UPSTREAM").expect("BOXMAKER_UPSTREAM is not set");
        let child = Command::new(binary)
            .arg("--listen")
            .arg(socket)
            .arg("--upstream")
            .arg(upstream)
            .spawn()
            .expect("cannot start inferproxy");
        let mut proxy = Proxy {
            child,
            socket: socket.to_path_buf(),
        };
        for _ in 0..100 {
            if socket.exists() {
                return proxy;
            }
            thread::sleep(Duration::from_millis(20));
        }
        proxy.kill();
        panic!("inferproxy did not create {}", socket.display());
    }

    fn kill(&mut self) {
        let _ = self.child.kill();
        let _ = self.child.wait();
    }
}

impl Drop for Proxy {
    fn drop(&mut self) {
        self.kill();
    }
}

fn socket_path(name: &str) -> PathBuf {
    let dir = std::env::temp_dir().join(format!("loopd-device-{}-{name}", std::process::id()));
    std::fs::create_dir_all(&dir).unwrap();
    dir.join("infer.sock")
}

fn config(socket: &Path) -> Config {
    let model =
        std::env::var("BOXMAKER_MODEL").unwrap_or_else(|_| "ornith-1.5-35b-a3b".to_string());
    let text = format!(
        r#"
[infer]
socket = "{}"
model = "{model}"
[slots]
main = 0
background = 1
[expect]
template_sha256 = "f55f52930aa8bf44ab5cb85f99370fcc3c56e9a85640b812086d5330bce5d86b"
n_ctx = 131072
slots = 2
"#,
        socket.display()
    );
    Config::parse(&text).unwrap()
}

fn user(text: &str) -> ChatMessage {
    // A different prompt each run, so that an earlier run's cache cannot make a check pass.
    let nonce = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_nanos();
    ChatMessage::User {
        content: format!("{text} (run {nonce})"),
    }
}

#[test]
#[ignore = "needs the real server; run with make verify-device"]
fn the_startup_self_test_passes() {
    let socket = socket_path("selftest");
    let _proxy = Proxy::start(&socket);
    let client = Client::new(config(&socket));
    loopd::selftest::run(&client, &mut |step| eprintln!("selftest: {step}")).unwrap();
}

#[test]
#[ignore = "needs the real server; run with make verify-device"]
fn a_capped_thinking_block_ends_and_the_answer_arrives() {
    let socket = socket_path("cap");
    let _proxy = Proxy::start(&socket);
    let mut cfg = config(&socket);
    cfg.limits.thinking_cap = 60;
    let client = Client::new(cfg);
    let ask = "Think at great length about how many ways 8 rooks fit on a chessboard without attacking \
               each other and with none on the main diagonal. Then answer in one sentence.";
    let req = ChatRequest {
        slot: 0,
        messages: vec![user(ask)],
        tools: vec![],
        thinking: true,
    };
    let mut capped = Vec::new();
    let done = client
        .chat_with_retry(&req, &mut |e| {
            if let ChatEvent::ThinkingCapped { tokens } = e {
                capped.push(*tokens);
            }
        })
        .unwrap();
    assert!(done.thinking_capped);
    assert_eq!(capped.len(), 1);
    assert!(capped[0] >= 60 && capped[0] < 70, "capped at {capped:?}");
    assert!(
        done.reasoning_tokens < 60 + 256,
        "thinking went on to {}",
        done.reasoning_tokens
    );
    assert!(
        done.content.is_some_and(|c| !c.trim().is_empty()),
        "an answer followed the forced end"
    );
}

#[test]
#[ignore = "needs the real server; run with make verify-device"]
fn a_request_survives_its_proxy_being_killed_and_restarted() {
    let socket = socket_path("restart");
    let mut proxy = Proxy::start(&socket);
    let mut cfg = config(&socket);
    cfg.limits.retry_backoff_ms = vec![1_500];
    let client = Client::new(cfg);
    let ask = "Write about 300 words on the history of cork.";
    let req = ChatRequest {
        slot: 0,
        messages: vec![user(ask)],
        tools: vec![],
        thinking: false,
    };

    let (tx, rx) = mpsc::channel();
    let worker = thread::spawn(move || {
        let mut retries = 0;
        let mut told = false;
        let result = client.chat_with_retry(&req, &mut |e| match e {
            ChatEvent::Content(_) if !told => {
                told = true;
                let _ = tx.send(());
            }
            ChatEvent::Retrying { .. } => retries += 1,
            _ => {}
        });
        (result, retries)
    });

    rx.recv_timeout(Duration::from_secs(120))
        .expect("no content arrived");
    proxy.kill(); // the stream dies in the middle of the answer
    thread::sleep(Duration::from_millis(300));
    let _proxy = Proxy::start(&proxy.socket);

    let (result, retries) = worker.join().unwrap();
    let done = result.expect("the retry should have succeeded");
    assert!(retries >= 1, "the request was retried");
    assert!(
        done.content
            .is_some_and(|c| c.split_whitespace().count() > 100)
    );
}

#[test]
#[ignore = "needs the real server; run with make verify-device"]
fn a_second_turn_reuses_the_first_turns_cache() {
    let socket = socket_path("cache");
    let _proxy = Proxy::start(&socket);
    let client = Client::new(config(&socket));
    let mut messages = vec![user(
        "What is 17 * 23? Think briefly, then answer in one short sentence.",
    )];
    let req = ChatRequest {
        slot: 0,
        messages: messages.clone(),
        tools: vec![],
        thinking: true,
    };
    let turn1 = client.chat_with_retry(&req, &mut |_| {}).unwrap();
    assert!(
        turn1.reasoning_content.is_some(),
        "this check is about replaying a thinking block"
    );
    messages.push(ChatMessage::Assistant {
        content: turn1.content.clone(),
        reasoning_content: turn1.reasoning_content.clone(),
        tool_calls: turn1.tool_calls.clone(),
    });
    messages.push(user("And 17 * 24?"));
    let req = ChatRequest {
        slot: 0,
        messages,
        tools: vec![],
        thinking: true,
    };
    let turn2 = client.chat_with_retry(&req, &mut |_| {}).unwrap();
    assert_eq!(
        cache_outcome(&turn1.timings, &turn2.timings),
        CacheOutcome::Hit,
        "{:?} then {:?}",
        turn1.timings,
        turn2.timings
    );
}