boxmaker/crates/loopd/tests/session_slot.rs

//! Tests that a live turn uses the session's slot. Do not edit.

mod support;

use bxcore::SessionId;
use loopd::baseline::Baseline;
use loopd::llama::Client;
use loopd::session::Session;
use loopd::tools::Registry;
use loopd::turn::{Runtime, run_turn};
use support::{FakeServer, Home, Reply};

const CHAT: &str = "/v1/chat/completions";

fn sse() -> Reply {
    let data1 = r#"data: {"choices":[{"finish_reason":null,"index":0,"delta":{"content":"ok"}}],"id":"t","timings":{"cache_n":0,"prompt_n":10,"predicted_n":1}}"#;
    let data2 = r#"data: {"choices":[{"finish_reason":"stop","index":0,"delta":{}}],"id":"t","timings":{"cache_n":0,"prompt_n":10,"predicted_n":1}}"#;
    let body = format!("{data1}\n\n{data2}\n\ndata: [DONE]\n\n");
    Reply::raw(format!(
        "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
        body.len()
    ))
}

fn runtime<'a>(
    cfg: &'a loopd::config::Config,
    client: &'a Client,
    port: &'a support::ScriptedPort,
    registry: &'a Registry,
) -> Runtime<'a> {
    Runtime {
        cfg,
        client,
        port,
        registry,
    }
}

#[test]
fn a_session_created_on_the_background_slot_sends_that_slot() {
    let home = Home::new();
    let server = FakeServer::start();
    let cfg = home.config(&server.socket);
    server.route(CHAT, vec![sse()]);
    let client = Client::new(cfg.clone());
    let port = support::ScriptedPort::new(vec![]);
    let registry = Registry::m2b();
    let baseline = Baseline::assemble(&cfg, &registry).unwrap();
    let mut s = Session::create(
        &home.dir,
        SessionId::new("heartbeat").unwrap(),
        baseline,
        cfg.slots.background,
    )
    .unwrap();
    let mut events = Vec::new();
    run_turn(
        &mut s,
        &runtime(&cfg, &client, &port, &registry),
        "pulse",
        None,
        &mut |e| events.push(e.clone()),
    )
    .unwrap();
    let sent = server.requests_to(CHAT);
    assert_eq!(sent.len(), 1);
    assert_eq!(sent[0].json()["id_slot"], cfg.slots.background);
}

#[test]
fn a_main_slot_session_still_sends_main() {
    let home = Home::new();
    let server = FakeServer::start();
    let cfg = home.config(&server.socket);
    server.route(CHAT, vec![sse()]);
    let client = Client::new(cfg.clone());
    let port = support::ScriptedPort::new(vec![]);
    let registry = Registry::m2b();
    let baseline = Baseline::assemble(&cfg, &registry).unwrap();
    let mut s = Session::create(
        &home.dir,
        SessionId::new("mm-thread").unwrap(),
        baseline,
        cfg.slots.main,
    )
    .unwrap();
    let mut events = Vec::new();
    run_turn(
        &mut s,
        &runtime(&cfg, &client, &port, &registry),
        "hi",
        None,
        &mut |e| events.push(e.clone()),
    )
    .unwrap();
    assert_eq!(
        server.requests_to(CHAT)[0].json()["id_slot"],
        cfg.slots.main
    );
}

#[test]
fn the_request_slot_follows_session_start_not_the_config_main() {
    let home = Home::new();
    let server = FakeServer::start();
    let mut cfg = home.config(&server.socket);
    cfg.slots.main = 0;
    cfg.slots.background = 1;
    server.route(CHAT, vec![sse()]);
    let client = Client::new(cfg.clone());
    let port = support::ScriptedPort::new(vec![]);
    let registry = Registry::m2b();
    let baseline = Baseline::assemble(&cfg, &registry).unwrap();
    let mut s = Session::create(&home.dir, SessionId::new("a").unwrap(), baseline, 7).unwrap();
    let mut events = Vec::new();
    run_turn(
        &mut s,
        &runtime(&cfg, &client, &port, &registry),
        "x",
        None,
        &mut |e| events.push(e.clone()),
    )
    .unwrap();
    assert_eq!(server.requests_to(CHAT)[0].json()["id_slot"], 7);
}

#[test]
fn a_heartbeat_turn_between_two_main_turns_does_not_change_the_main_prefix() {
    let home = Home::new();
    let server = FakeServer::start();
    let cfg = home.config(&server.socket);
    server.route(CHAT, vec![sse(), sse(), sse()]);
    let client = Client::new(cfg.clone());
    let port = support::ScriptedPort::new(vec![]);
    let registry = Registry::m2b();
    let mut main = Session::create(
        &home.dir,
        SessionId::new("mm-a").unwrap(),
        Baseline::assemble(&cfg, &registry).unwrap(),
        cfg.slots.main,
    )
    .unwrap();
    let mut beat = Session::create(
        &home.dir,
        SessionId::new("heartbeat").unwrap(),
        Baseline::assemble(&cfg, &registry).unwrap(),
        cfg.slots.background,
    )
    .unwrap();
    let rt = runtime(&cfg, &client, &port, &registry);
    let mut events = Vec::new();
    run_turn(&mut main, &rt, "one", None, &mut |e| events.push(e.clone())).unwrap();
    run_turn(&mut beat, &rt, "pulse", None, &mut |e| {
        events.push(e.clone())
    })
    .unwrap();
    run_turn(&mut main, &rt, "two", None, &mut |e| events.push(e.clone())).unwrap();
    let sent = server.requests_to(CHAT);
    assert_eq!(sent.len(), 3);
    assert_eq!(sent[0].json()["id_slot"], 0);
    assert_eq!(sent[1].json()["id_slot"], 1);
    assert_eq!(sent[2].json()["id_slot"], 0);
    let first = sent[0].json()["messages"].as_array().unwrap().clone();
    let third = sent[2].json()["messages"].as_array().unwrap().clone();
    assert!(
        third.len() > first.len(),
        "the main session's second turn is a longer prefix"
    );
    assert_eq!(&third[..first.len()], &first[..]);
}