boxmaker/crates/loopd/tests/chat.rs

//! Tests for one request: events, waits, liveness and errors. Do not edit.
//!
//! Limits come from `support::test_config`: poll 40 ms, busy wait 400 ms, load wait 300 ms,
//! idle grace 200 ms, liveness 150 ms.

mod support;

use loopd::llama::{ChatEvent, ChatMessage, ChatRequest, Client, FinishReason, InferError};
use std::time::{Duration, Instant};
use support::{FakeServer, Reply};

const CHAT: &str = "/v1/chat/completions";
const IDLE: &str = r#"[{"id":0,"is_processing":false},{"id":1,"is_processing":false}]"#;
const SLOT0_BUSY: &str = r#"[{"id":0,"is_processing":true},{"id":1,"is_processing":false}]"#;
const SLOT1_BUSY: &str = r#"[{"id":0,"is_processing":false},{"id":1,"is_processing":true}]"#;

fn request() -> ChatRequest {
    ChatRequest {
        slot: 0,
        messages: vec![ChatMessage::User {
            content: "hi".to_string(),
        }],
        tools: vec![],
        thinking: false,
    }
}

fn run(
    server: &FakeServer,
) -> (
    Result<loopd::llama::Completion, InferError>,
    Vec<ChatEvent>,
    Duration,
) {
    let client = Client::new(support::test_config(&server.socket));
    let mut events = Vec::new();
    let started = Instant::now();
    let result = client.chat(&request(), &mut |e| events.push(e.clone()));
    (result, events, started.elapsed())
}

fn waiting(events: &[ChatEvent]) -> Vec<bool> {
    events
        .iter()
        .filter_map(|e| match e {
            ChatEvent::Waiting { slot_busy } => Some(*slot_busy),
            _ => None,
        })
        .collect()
}

#[test]
fn a_recorded_completion_comes_back_whole() {
    let server = FakeServer::start();
    server.route(CHAT, vec![Reply::fixture("tool_call")]);
    let (result, events, _) = run(&server);
    let done = result.unwrap();
    let want = support::expected("tool_call");
    assert_eq!(done.finish_reason, FinishReason::ToolCalls);
    assert_eq!(done.content.as_deref(), want["content"].as_str());
    assert_eq!(done.tool_calls[0].arguments, r#"{"path":"/etc/hostname"}"#);
    assert_eq!(done.timings.prompt_n, 312);
    assert!(!done.thinking_capped);
    assert!(
        waiting(&events).is_empty(),
        "the answer came at once, so nothing was polled"
    );
    assert!(matches!(
        events.first(),
        Some(ChatEvent::Progress { total: 312, .. })
    ));

    let sent = &server.requests_to(CHAT)[0];
    assert_eq!(sent.method, "POST");
    assert_eq!(sent.json()["messages"][0]["content"], "hi");
    assert_eq!(sent.json()["id_slot"], 0);
}

#[test]
fn the_result_does_not_depend_on_how_the_bytes_arrive() {
    for piece in [1, 17, 4096] {
        let server = FakeServer::start();
        server.route(CHAT, vec![Reply::fixture("thinking").trickle(piece, 0)]);
        let done = run(&server).0.unwrap();
        let want = support::expected("thinking");
        assert_eq!(
            done.content.as_deref(),
            want["content"].as_str(),
            "pieces of {piece}"
        );
        assert_eq!(
            done.reasoning_content.as_deref(),
            want["reasoning_content"].as_str()
        );
        assert_eq!(done.reasoning_tokens, 49);
    }
}

#[test]
fn a_busy_slot_is_waited_out() {
    let server = FakeServer::start();
    server.route(CHAT, vec![Reply::fixture("plain").head_delay(250)]);
    server.route("/slots", vec![Reply::json(200, SLOT0_BUSY)]);
    let (result, events, took) = run(&server);
    assert!(result.is_ok(), "{result:?}");
    let polls = waiting(&events);
    assert!(polls.len() >= 4, "250 ms at 40 ms per poll: {polls:?}");
    assert!(polls.iter().all(|busy| *busy));
    assert!(took >= Duration::from_millis(250));
    assert!(server.requests_to("/slots").len() >= 4);
    assert_eq!(
        server.requests_to("/slots")[0].target,
        "/slots?model=test-model"
    );
}

#[test]
fn a_slot_that_stays_busy_is_a_wait_timeout() {
    let server = FakeServer::start();
    server.route(CHAT, vec![Reply::fixture("plain").head_delay(5_000)]);
    server.route("/slots", vec![Reply::json(200, SLOT0_BUSY)]);
    let (result, events, took) = run(&server);
    assert!(matches!(result, Err(InferError::WaitTimeout)), "{result:?}");
    assert!(
        took >= Duration::from_millis(400) && took < Duration::from_millis(1500),
        "{took:?}"
    );
    assert!(waiting(&events).iter().all(|busy| *busy));
}

#[test]
fn only_the_requests_own_slot_counts_as_busy() {
    // Slot 1 is busy, slot 0 (ours) is idle and silent: that is a stall, not a queue.
    let server = FakeServer::start();
    server.route(CHAT, vec![Reply::fixture("plain").head_delay(5_000)]);
    server.route("/slots", vec![Reply::json(200, SLOT1_BUSY)]);
    let (result, events, took) = run(&server);
    assert!(matches!(result, Err(InferError::Stalled)), "{result:?}");
    assert!(
        took >= Duration::from_millis(200) && took < Duration::from_millis(1000),
        "{took:?}"
    );
    assert!(waiting(&events).iter().all(|busy| !*busy));
}

#[test]
fn an_unavailable_server_is_a_load_timeout() {
    for slots_reply in [
        Reply::json(503, r#"{"error":"loading"}"#),
        Reply::raw("").cut_after(0),
    ] {
        let server = FakeServer::start();
        server.route(CHAT, vec![Reply::fixture("plain").head_delay(5_000)]);
        server.route("/slots", vec![slots_reply]);
        let (result, _, took) = run(&server);
        assert!(matches!(result, Err(InferError::LoadTimeout)), "{result:?}");
        assert!(
            took >= Duration::from_millis(300) && took < Duration::from_millis(1500),
            "{took:?}"
        );
    }
}

#[test]
fn the_wait_clocks_restart_when_the_state_changes() {
    // Busy for about 5 polls, then idle: the idle grace starts counting from there, so the
    // request outlives 200 ms of busy plus most of the 200 ms grace and then completes.
    let server = FakeServer::start();
    server.route(CHAT, vec![Reply::fixture("plain").head_delay(330)]);
    let busy = Reply::json(200, SLOT0_BUSY);
    server.route(
        "/slots",
        vec![
            busy.clone(),
            busy.clone(),
            busy.clone(),
            busy.clone(),
            busy,
            Reply::json(200, IDLE),
        ],
    );
    let (result, events, _) = run(&server);
    assert!(result.is_ok(), "{result:?}");
    let polls = waiting(&events);
    assert!(
        polls.starts_with(&[true, true, true, true, true]),
        "{polls:?}"
    );
    assert_eq!(polls.last(), Some(&false), "{polls:?}");
}

#[test]
fn silence_in_mid_stream_is_a_stall() {
    let server = FakeServer::start();
    let reply = Reply::fixture("thinking");
    let cut = reply.offset_after_events(10);
    server.route(CHAT, vec![reply.hang_after(cut, 5_000)]);
    let (result, events, took) = run(&server);
    assert!(matches!(result, Err(InferError::Stalled)), "{result:?}");
    assert!(
        took >= Duration::from_millis(150) && took < Duration::from_millis(1000),
        "{took:?}"
    );
    assert!(
        events.iter().any(|e| matches!(e, ChatEvent::Reasoning(_))),
        "events before the stall were delivered"
    );
}

#[test]
fn a_slow_but_steady_stream_is_not_a_stall() {
    // 60 ms between pieces is well inside the 150 ms liveness limit, however long the whole takes.
    let server = FakeServer::start();
    let size = support::fixture_bytes("http", "turn1.http").len();
    server.route(CHAT, vec![Reply::fixture("turn1").trickle(size / 8, 60)]);
    let (result, _, took) = run(&server);
    assert_eq!(result.unwrap().content.as_deref(), Some("Blue"));
    assert!(
        took >= Duration::from_millis(400),
        "the stream took longer than the liveness limit: {took:?}"
    );
}

#[test]
fn a_stream_that_is_cut_is_closed_early_wherever_the_cut_falls() {
    let reply = Reply::fixture("plain");
    let whole = support::fixture_bytes("http", "plain.http").len();
    let after_five = reply.offset_after_events(5);
    // After the head only; in the middle of an event; between events; just before [DONE].
    for cut in [300, after_five + 20, after_five, whole - 30] {
        let server = FakeServer::start();
        server.route(CHAT, vec![reply.clone().cut_after(cut)]);
        let (result, _, _) = run(&server);
        assert!(
            matches!(result, Err(InferError::StreamClosedEarly)),
            "cut at {cut}: {result:?}"
        );
    }
    // Cut before the head is complete.
    let server = FakeServer::start();
    server.route(CHAT, vec![reply.cut_after(40)]);
    assert!(matches!(run(&server).0, Err(InferError::StreamClosedEarly)));
}

#[test]
fn error_statuses_keep_their_body() {
    let server = FakeServer::start();
    server.route(CHAT, vec![Reply::fixture("bad_request")]);
    match run(&server).0 {
        Err(InferError::Http { status: 400, body }) => {
            assert!(body.contains("Expected 'messages'"), "{body}")
        }
        other => panic!("{other:?}"),
    }
    let server = FakeServer::start();
    server.route(
        CHAT,
        vec![Reply::json(503, r#"{"error":{"message":"Loading model"}}"#)],
    );
    assert!(matches!(
        run(&server).0,
        Err(InferError::Http { status: 503, .. })
    ));
}

#[test]
fn garbage_in_the_stream_is_a_protocol_error() {
    let server = FakeServer::start();
    let raw = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\r\ndata: {not json}\n\n";
    server.route(CHAT, vec![Reply::raw(raw)]);
    assert!(matches!(run(&server).0, Err(InferError::Protocol(_))));

    let server = FakeServer::start();
    let huge = format!(
        "HTTP/1.1 200 OK\r\n\r\ndata: {}\n\n",
        "x".repeat(2 * 1024 * 1024)
    );
    server.route(CHAT, vec![Reply::raw(huge)]);
    assert!(
        matches!(run(&server).0, Err(InferError::Protocol(_))),
        "a line over 1 MiB"
    );
}

#[test]
fn no_server_is_a_connect_error() {
    let server = FakeServer::start();
    let mut cfg = support::test_config(&server.socket);
    cfg.infer.socket = std::env::temp_dir().join("loopd-no-such-socket.sock");
    let result = Client::new(cfg).chat(&request(), &mut |_| {});
    assert!(matches!(result, Err(InferError::Connect(_))), "{result:?}");
    let e: Box<dyn std::error::Error> = Box::new(InferError::Stalled);
    assert!(!e.to_string().is_empty());
}