boxmaker/crates/loopd/tests/assemble.rs

//! Tests for folding streamed chunks into a completion. Do not edit.
//!
//! The streams are recorded from the real server. The expected results in
//! `tests/fixtures/expected/` were derived from the same recordings by a separate script, so
//! this checks the Rust code against an independent reading of the same bytes.

mod support;

use loopd::llama::assemble::Assembler;
use loopd::llama::{ChatEvent, Completion, FinishReason, InferError};

const ALL: [&str; 7] = [
    "plain",
    "thinking",
    "tool_call",
    "progress",
    "capped",
    "turn1",
    "turn2",
];

/// The `data:` payloads of a recorded stream, without `[DONE]`.
fn payloads(name: &str) -> Vec<String> {
    let raw = String::from_utf8(support::fixture_bytes("http", &format!("{name}.http"))).unwrap();
    raw.lines()
        .filter_map(|l| l.strip_prefix("data: "))
        .filter(|p| p.starts_with('{'))
        .map(str::to_string)
        .collect()
}

fn assemble(name: &str) -> (Completion, Vec<ChatEvent>) {
    let mut a = Assembler::new();
    let mut events = Vec::new();
    for p in payloads(name) {
        events.extend(a.push(&p).unwrap());
    }
    assert!(a.is_finished(), "{name}");
    (a.finish(false).unwrap(), events)
}

fn opt(v: &serde_json::Value) -> Option<String> {
    v.as_str().map(str::to_string)
}

#[test]
fn every_recording_assembles_to_what_the_independent_script_found() {
    for name in ALL {
        let want = support::expected(name);
        let (got, events) = assemble(name);
        assert_eq!(Some(got.id.clone()), opt(&want["id"]), "{name}: id");
        assert_eq!(got.content, opt(&want["content"]), "{name}: content");
        assert_eq!(
            got.reasoning_content,
            opt(&want["reasoning_content"]),
            "{name}: reasoning"
        );
        let finish = match got.finish_reason {
            FinishReason::Stop => "stop",
            FinishReason::ToolCalls => "tool_calls",
            FinishReason::Length => "length",
        };
        assert_eq!(finish, want["finish_reason"], "{name}: finish reason");
        assert_eq!(got.timings.cache_n, want["cache_n"], "{name}");
        assert_eq!(got.timings.prompt_n, want["prompt_n"], "{name}");
        assert_eq!(got.timings.predicted_n, want["predicted_n"], "{name}");
        assert_eq!(
            got.reasoning_tokens, want["reasoning_tokens"],
            "{name}: reasoning tokens"
        );
        assert!(!got.thinking_capped, "{name}");

        let calls: Vec<serde_json::Value> = got
            .tool_calls
            .iter()
            .map(|c| serde_json::json!({"id": c.id, "name": c.name, "arguments": c.arguments}))
            .collect();
        assert_eq!(
            serde_json::Value::Array(calls),
            want["tool_calls"],
            "{name}: tool calls"
        );

        let count = |f: fn(&ChatEvent) -> bool| events.iter().filter(|e| f(e)).count();
        assert_eq!(
            count(|e| matches!(e, ChatEvent::Reasoning(_))),
            want["reasoning_events"],
            "{name}"
        );
        assert_eq!(
            count(|e| matches!(e, ChatEvent::Content(_))),
            want["content_events"],
            "{name}"
        );
        assert_eq!(
            count(|e| matches!(e, ChatEvent::ToolCallDelta { .. })),
            want["tool_events"],
            "{name}"
        );
        let progress: Vec<serde_json::Value> = events
            .iter()
            .filter_map(|e| match e {
                ChatEvent::Progress {
                    total,
                    cache,
                    processed,
                } => Some(serde_json::json!([total, cache, processed])),
                _ => None,
            })
            .collect();
        assert_eq!(
            serde_json::Value::Array(progress),
            want["progress"],
            "{name}: progress"
        );
    }
}

#[test]
fn the_events_carry_the_same_text_as_the_completion() {
    for name in ALL {
        let (got, events) = assemble(name);
        let mut content = String::new();
        let mut reasoning = String::new();
        let mut arguments = String::new();
        for e in &events {
            match e {
                ChatEvent::Content(t) => content.push_str(t),
                ChatEvent::Reasoning(t) => reasoning.push_str(t),
                ChatEvent::ToolCallDelta { arguments: a, .. } => arguments.push_str(a),
                _ => {}
            }
        }
        assert_eq!(got.content.unwrap_or_default(), content, "{name}");
        assert_eq!(
            got.reasoning_content.unwrap_or_default(),
            reasoning,
            "{name}"
        );
        let all_arguments: String = got
            .tool_calls
            .iter()
            .map(|c| c.arguments.as_str())
            .collect();
        assert_eq!(all_arguments, arguments, "{name}");
    }
}

#[test]
fn the_tool_call_arguments_are_valid_json() {
    let (got, events) = assemble("tool_call");
    let args: serde_json::Value = serde_json::from_str(&got.tool_calls[0].arguments).unwrap();
    assert_eq!(args, serde_json::json!({"path": "/etc/hostname"}));
    // Only the first delta of a call carries its id and name.
    let named = events
        .iter()
        .filter(|e| matches!(e, ChatEvent::ToolCallDelta { name: Some(_), .. }))
        .count();
    assert_eq!(named, 1);
}

#[test]
fn reasoning_state_follows_the_stream() {
    let mut a = Assembler::new();
    assert!(!a.in_reasoning());
    assert_eq!(a.reasoning_tokens(), 0);
    let mut seen_reasoning = false;
    let mut last = 0;
    for p in payloads("thinking") {
        let events = a.push(&p).unwrap();
        if events.iter().any(|e| matches!(e, ChatEvent::Reasoning(_))) {
            seen_reasoning = true;
            assert!(a.in_reasoning());
            assert!(
                a.reasoning_tokens() > last,
                "the count must grow with each reasoning chunk"
            );
            last = a.reasoning_tokens();
        }
        if events.iter().any(|e| matches!(e, ChatEvent::Content(_))) {
            assert!(!a.in_reasoning(), "content ends the reasoning block");
            assert_eq!(a.reasoning_tokens(), last, "and the count stops");
        }
    }
    assert!(seen_reasoning);
    assert_eq!(last, 49);
    assert!(a.finish(true).unwrap().thinking_capped);
}

#[test]
fn two_tool_calls_in_one_completion() {
    let chunk = |delta: &str, finish: &str| {
        format!(
            r#"{{"id":"c1","choices":[{{"index":0,"finish_reason":{finish},"delta":{delta}}}]}}"#
        )
    };
    let mut a = Assembler::new();
    for delta in [
        r#"{"tool_calls":[{"index":0,"id":"a","type":"function","function":{"name":"one","arguments":"{\"x\""}}]}"#,
        r#"{"tool_calls":[{"index":0,"function":{"arguments":":1}"}}]}"#,
        r#"{"tool_calls":[{"index":1,"id":"b","type":"function","function":{"name":"two","arguments":"{}"}}]}"#,
    ] {
        a.push(&chunk(delta, "null")).unwrap();
    }
    a.push(&chunk("{}", "\"tool_calls\"")).unwrap();
    let got = a.finish(false).unwrap();
    let calls: Vec<(&str, &str, &str)> = got
        .tool_calls
        .iter()
        .map(|c| (c.id.as_str(), c.name.as_str(), c.arguments.as_str()))
        .collect();
    assert_eq!(calls, vec![("a", "one", "{\"x\":1}"), ("b", "two", "{}")]);
}

/// The server sends fields this code does not use, and newer builds may send more. Unknown
/// fields in the server's chunks must be ignored. (Our own formats reject unknown fields; the
/// server's format is not ours.)
#[test]
fn unknown_fields_from_the_server_are_ignored() {
    let mut a = Assembler::new();
    let chunk = r#"{"id":"c1","zz_new":1,"choices":[{"index":0,"zz":2,"finish_reason":"stop","delta":{"content":"hi","zz":3}}],"timings":{"cache_n":1,"prompt_n":2,"predicted_n":3,"zz_ms":0.5}}"#;
    a.push(chunk).unwrap();
    let got = a.finish(false).unwrap();
    assert_eq!(got.content.as_deref(), Some("hi"));
    assert_eq!(
        (
            got.timings.cache_n,
            got.timings.prompt_n,
            got.timings.predicted_n
        ),
        (1, 2, 3)
    );
}

#[test]
fn bad_streams_are_errors_not_panics() {
    let mut a = Assembler::new();
    assert!(matches!(a.push("not json"), Err(InferError::Protocol(_))));
    assert!(matches!(a.push("[1,2]"), Err(InferError::Protocol(_))));

    let mut a = Assembler::new();
    let odd = r#"{"id":"c1","choices":[{"index":0,"finish_reason":"content_filter","delta":{}}]}"#;
    assert!(
        matches!(a.push(odd), Err(InferError::Protocol(_))),
        "an unknown finish reason"
    );

    let mut a = Assembler::new();
    let skip = r#"{"id":"c1","choices":[{"index":0,"finish_reason":null,"delta":{"tool_calls":[{"index":5,"id":"a","function":{"name":"n","arguments":""}}]}}]}"#;
    assert!(
        matches!(a.push(skip), Err(InferError::Protocol(_))),
        "a tool call index that skips ahead"
    );

    let mut a = Assembler::new();
    let nameless = r#"{"id":"c1","choices":[{"index":0,"finish_reason":"tool_calls","delta":{"tool_calls":[{"index":0,"function":{"arguments":"{}"}}]}}]}"#;
    a.push(nameless).unwrap();
    assert!(
        matches!(a.finish(false), Err(InferError::Protocol(_))),
        "a tool call without id or name"
    );

    // A stream that stops before any finish_reason.
    let mut a = Assembler::new();
    for p in payloads("plain").iter().take(5) {
        a.push(p).unwrap();
    }
    assert!(!a.is_finished());
    assert!(matches!(
        a.finish(false),
        Err(InferError::StreamClosedEarly)
    ));

    let mut a = Assembler::new();
    a.push(r#"{"choices":[{"index":0,"finish_reason":"stop","delta":{}}]}"#)
        .unwrap();
    assert!(
        matches!(a.finish(false), Err(InferError::Protocol(_))),
        "no completion id"
    );
}