boxmaker/crates/loopd/tests/channel.rs

//! Tests for the channel protocol on `loop.sock`. Do not edit.
//!
//! A real `channel::serve` runs on a socket in a temporary home, with the fake inference server
//! behind it. The tests speak the frame protocol to it directly.

mod support;

use loopd::channel::{Context, serve};
use loopd::llama::Client;
use loopd::tools::Registry;
use proto::{
    Envelope, ErrorCode, Message, PROTOCOL_VERSION, SessionId, Turn, TurnEvent, read_frame,
    write_frame,
};
use std::os::unix::net::{UnixListener, UnixStream};
use std::path::PathBuf;
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use support::{FakeServer, Home, Reply, ScriptedPort};

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

struct Loopd {
    home: Home,
    server: FakeServer,
    socket: PathBuf,
}

fn start(replies: Vec<proto::ToolResponse>) -> Loopd {
    let home = Home::new();
    let server = FakeServer::start();
    let cfg = home.config(&server.socket);
    let socket = cfg.channel_socket();
    let listener = UnixListener::bind(&socket).unwrap();
    let ctx = Arc::new(Context::new(
        cfg.clone(),
        Client::new(cfg),
        Box::new(ScriptedPort::new(replies)),
        Registry::m2b(),
    ));
    thread::spawn(move || serve(listener, ctx));
    Loopd {
        home,
        server,
        socket,
    }
}

fn id(s: &str) -> SessionId {
    SessionId::new(s).unwrap()
}

/// Sends one turn and collects every frame that comes back.
fn turn(socket: &PathBuf, session: &str, content: &str, resume: bool) -> Vec<Envelope> {
    let mut stream = UnixStream::connect(socket).unwrap();
    stream
        .set_read_timeout(Some(Duration::from_secs(10)))
        .unwrap();
    let msg = Message::Turn(Turn {
        session: id(session),
        content: content.to_string(),
        resume,
    });
    write_frame(
        &mut stream,
        &Envelope {
            v: PROTOCOL_VERSION,
            id: 42,
            r#final: true,
            msg,
        },
    )
    .unwrap();
    let mut frames = Vec::new();
    loop {
        let frame = read_frame(&mut stream).unwrap();
        let last = frame.r#final;
        frames.push(frame);
        if last {
            break;
        }
    }
    frames
}

fn error_code(frames: &[Envelope]) -> Option<ErrorCode> {
    match &frames.last()?.msg {
        Message::Error(e) => Some(e.code),
        _ => None,
    }
}

#[test]
fn a_turn_streams_events_and_ends_with_turn_done() {
    let l = start(vec![]);
    l.server.route(CHAT, vec![Reply::fixture("thinking")]);
    let frames = turn(&l.socket, "a", "what is 17 * 23?", false);
    assert!(frames.len() > 3, "{frames:?}");
    assert!(
        frames.iter().all(|f| f.id == 42),
        "every frame carries the request id"
    );
    assert!(frames.iter().all(|f| f.v == PROTOCOL_VERSION));
    let (last, events) = frames.split_last().unwrap();
    assert!(
        events
            .iter()
            .all(|f| !f.r#final && matches!(f.msg, Message::TurnEvent(_)))
    );
    assert!(
        events
            .iter()
            .any(|f| matches!(f.msg, Message::TurnEvent(TurnEvent::Reasoning { .. })))
    );
    assert!(
        events
            .iter()
            .any(|f| matches!(f.msg, Message::TurnEvent(TurnEvent::Content { .. })))
    );
    match &last.msg {
        Message::TurnDone(done) => {
            assert_eq!(
                done.content,
                support::expected("thinking")["content"].as_str().unwrap()
            );
            assert_eq!(done.usage.reasoning_tokens, 49);
            assert_eq!(done.usage.prompt_n, 40);
        }
        other => panic!("{other:?}"),
    }
    // The content events, concatenated, are the answer.
    let streamed: String = events
        .iter()
        .filter_map(|f| match &f.msg {
            Message::TurnEvent(TurnEvent::Content { text }) => Some(text.as_str()),
            _ => None,
        })
        .collect();
    let Message::TurnDone(done) = &last.msg else {
        unreachable!()
    };
    assert_eq!(streamed, done.content);
    // And the session is on disk.
    assert_eq!(l.home.records("a").len(), 4);
}

#[test]
fn tool_calls_are_reported_by_name_only() {
    let l = start(vec![support::ok_result("straylight\n")]);
    l.server.route(
        CHAT,
        vec![Reply::fixture("tool_call"), Reply::fixture("plain")],
    );
    let frames = turn(&l.socket, "a", "hostname?", false);
    let tool_events: Vec<&TurnEvent> = frames
        .iter()
        .filter_map(|f| match &f.msg {
            Message::TurnEvent(
                e @ (TurnEvent::ToolCallStarted { .. } | TurnEvent::ToolResult { .. }),
            ) => Some(e),
            _ => None,
        })
        .collect();
    assert_eq!(tool_events.len(), 2);
    assert_eq!(
        tool_events[0],
        &TurnEvent::ToolCallStarted {
            name: "read_file".to_string()
        }
    );
    assert_eq!(
        tool_events[1],
        &TurnEvent::ToolResult {
            name: "read_file".to_string(),
            class: proto::DataClass::Private,
            truncated: false
        }
    );
    let text = serde_json::to_string(&frames).unwrap();
    assert!(
        !text.contains("straylight"),
        "the result body is not on the channel: {text}"
    );
}

#[test]
fn resume_and_create_are_checked() {
    let l = start(vec![]);
    l.server
        .route(CHAT, vec![Reply::fixture("turn1"), Reply::fixture("turn2")]);
    assert_eq!(
        error_code(&turn(&l.socket, "a", "x", true)),
        Some(ErrorCode::NoSuchSession),
        "resume needs an existing session"
    );
    assert!(matches!(
        turn(&l.socket, "a", "one", false).last().unwrap().msg,
        Message::TurnDone(_)
    ));
    assert_eq!(
        error_code(&turn(&l.socket, "a", "x", false)),
        Some(ErrorCode::SessionExists),
        "create needs a new one"
    );
    let frames = turn(&l.socket, "a", "two", true);
    assert!(
        matches!(frames.last().unwrap().msg, Message::TurnDone(_)),
        "{frames:?}"
    );
    let sent = l.server.requests_to(CHAT);
    assert_eq!(sent.len(), 2);
    assert_eq!(
        sent[1].json()["messages"].as_array().unwrap().len(),
        4,
        "the second turn carried the first"
    );
}

#[test]
fn a_busy_session_is_refused_at_once_and_another_session_is_not() {
    let l = start(vec![]);
    let size = support::fixture_bytes("http", "plain.http").len();
    l.server
        .route(CHAT, vec![Reply::fixture("plain").trickle(size / 4, 120)]);
    let socket = l.socket.clone();
    let first = thread::spawn(move || turn(&socket, "a", "slow", false));
    thread::sleep(Duration::from_millis(100));
    let started = std::time::Instant::now();
    assert_eq!(
        error_code(&turn(&l.socket, "a", "again", true)),
        Some(ErrorCode::SessionBusy)
    );
    assert!(
        started.elapsed() < Duration::from_millis(100),
        "refused at once"
    );
    assert!(
        matches!(
            turn(&l.socket, "b", "other", false).last().unwrap().msg,
            Message::TurnDone(_)
        ),
        "another session runs"
    );
    assert!(matches!(
        first.join().unwrap().last().unwrap().msg,
        Message::TurnDone(_)
    ));
}

#[test]
fn limits_and_server_errors_come_back_as_error_codes() {
    let l = start(vec![]);
    l.server.route(CHAT, vec![Reply::fixture("context_full")]);
    assert_eq!(
        error_code(&turn(&l.socket, "a", "x", false)),
        Some(ErrorCode::SessionFull)
    );

    l.server.route(CHAT, vec![Reply::fixture("tool_call")]);
    let frames = turn(&l.socket, "b", "x", false);
    assert_eq!(
        error_code(&frames),
        Some(ErrorCode::TurnLimit),
        "a repeated call twice"
    );
    assert!(
        frames
            .iter()
            .any(|f| matches!(f.msg, Message::TurnEvent(TurnEvent::ToolCallStarted { .. }))),
        "events before the error were delivered"
    );

    l.server.route(CHAT, vec![Reply::fixture("bad_request")]);
    let frames = turn(&l.socket, "c", "x", false);
    assert_eq!(error_code(&frames), Some(ErrorCode::Inference));
    let Message::Error(e) = &frames.last().unwrap().msg else {
        unreachable!()
    };
    assert!(e.detail.contains("400"), "{}", e.detail);
}

#[test]
fn bad_frames_get_an_error_and_a_close() {
    let l = start(vec![]);
    // A frame that is not a turn.
    let mut stream = UnixStream::connect(&l.socket).unwrap();
    stream
        .set_read_timeout(Some(Duration::from_secs(5)))
        .unwrap();
    let msg = Message::Error(proto::WireError {
        code: ErrorCode::Internal,
        detail: String::new(),
    });
    write_frame(
        &mut stream,
        &Envelope {
            v: PROTOCOL_VERSION,
            id: 1,
            r#final: true,
            msg,
        },
    )
    .unwrap();
    let reply = read_frame(&mut stream).unwrap();
    assert!(reply.r#final);
    assert!(
        matches!(
            reply.msg,
            Message::Error(proto::WireError {
                code: ErrorCode::BadMessage,
                ..
            })
        ),
        "{reply:?}"
    );
    assert!(
        matches!(read_frame(&mut stream), Err(proto::FrameError::Closed)),
        "then the connection is closed"
    );

    // Bytes that are not a frame at all.
    let mut stream = UnixStream::connect(&l.socket).unwrap();
    stream
        .set_read_timeout(Some(Duration::from_secs(5)))
        .unwrap();
    std::io::Write::write_all(&mut stream, &[0, 0, 0, 3, b'{', b'{', b'{']).unwrap();
    let reply = read_frame(&mut stream).unwrap();
    assert!(
        matches!(
            reply.msg,
            Message::Error(proto::WireError {
                code: ErrorCode::BadMessage,
                ..
            })
        ),
        "{reply:?}"
    );

    // A wrong protocol version.
    let mut stream = UnixStream::connect(&l.socket).unwrap();
    stream
        .set_read_timeout(Some(Duration::from_secs(5)))
        .unwrap();
    let msg = Message::Turn(Turn {
        session: id("a"),
        content: "x".to_string(),
        resume: false,
    });
    write_frame(
        &mut stream,
        &Envelope {
            v: 2,
            id: 1,
            r#final: true,
            msg,
        },
    )
    .unwrap();
    let reply = read_frame(&mut stream).unwrap();
    assert!(
        matches!(
            reply.msg,
            Message::Error(proto::WireError {
                code: ErrorCode::BadVersion,
                ..
            })
        ),
        "{reply:?}"
    );

    // A client that connects and leaves.
    drop(UnixStream::connect(&l.socket).unwrap());
    thread::sleep(Duration::from_millis(50));
    l.server.route(CHAT, vec![Reply::fixture("plain")]);
    assert!(matches!(
        turn(&l.socket, "z", "still up", false).last().unwrap().msg,
        Message::TurnDone(_)
    ));
}