Add the assembler that folds streamed chunks into a completion

Implemented-By: GLM-5.3 (z.ai, default settings) via OpenCode
This commit is contained in:
2026-09-18 02:56:31 -07:00
parent f48999306e
commit 3d876d4fa9
11 changed files with 762 additions and 0 deletions
+221
View File
@@ -0,0 +1,221 @@
//! Folding the server's streamed chunks into one completion and the `ChatEvent`s they stand for.
use serde::Deserialize;
use super::{ChatEvent, Completion, FinishReason, InferError, Timings};
// These structs parse the server's format, not ours, so unknown fields are ignored on purpose:
// newer servers send fields this code does not use.
#[derive(Deserialize)]
struct Chunk {
id: Option<String>,
choices: Vec<Choice>,
timings: Option<Timings>,
prompt_progress: Option<PromptProgress>,
}
#[derive(Deserialize)]
struct Choice {
finish_reason: Option<String>,
delta: Delta,
}
#[derive(Deserialize)]
struct Delta {
reasoning_content: Option<String>,
content: Option<String>,
tool_calls: Option<Vec<ToolCallPiece>>,
}
#[derive(Deserialize)]
struct ToolCallPiece {
index: u64,
id: Option<String>,
function: Option<FunctionPiece>,
}
#[derive(Deserialize)]
struct FunctionPiece {
name: Option<String>,
arguments: Option<String>,
}
#[derive(Deserialize)]
struct PromptProgress {
total: u64,
cache: u64,
processed: u64,
}
/// A tool call being assembled from its pieces; `id` and `name` wait for the first piece
/// that carries them.
#[derive(Default)]
struct PendingCall {
id: Option<String>,
name: Option<String>,
arguments: String,
}
#[derive(Default)]
pub struct Assembler {
id: Option<String>,
content: Option<String>,
reasoning_content: Option<String>,
reasoning_tokens: u64,
in_reasoning: bool,
finish_reason: Option<FinishReason>,
timings: Timings,
tool_calls: Vec<PendingCall>,
}
impl Assembler {
pub fn new() -> Self {
Self::default()
}
/// Takes one `data:` payload (JSON text) and returns the events it stands for, in order.
pub fn push(&mut self, data: &str) -> Result<Vec<ChatEvent>, InferError> {
let chunk: Chunk = serde_json::from_str(data)
.map_err(|e| InferError::Protocol(format!("unreadable chunk: {e}")))?;
// Timings first: the reasoning-token count below reads the latest one.
if let Some(t) = chunk.timings {
self.timings = t;
}
if let Some(id) = chunk.id {
self.id = Some(id);
}
let mut events = Vec::new();
if let Some(p) = chunk.prompt_progress {
events.push(ChatEvent::Progress {
total: p.total,
cache: p.cache,
processed: p.processed,
});
}
let mut carries_reasoning = false;
for choice in chunk.choices {
if let Some(reason) = choice.finish_reason {
self.finish_reason = Some(match reason.as_str() {
"stop" => FinishReason::Stop,
"tool_calls" => FinishReason::ToolCalls,
"length" => FinishReason::Length,
other => {
return Err(InferError::Protocol(format!(
"unknown finish_reason {other:?}"
)));
}
});
}
if let Some(text) = choice.delta.reasoning_content
&& !text.is_empty()
{
carries_reasoning = true;
self.in_reasoning = true;
self.reasoning_content
.get_or_insert_with(String::new)
.push_str(&text);
events.push(ChatEvent::Reasoning(text));
}
if let Some(text) = choice.delta.content
&& !text.is_empty()
{
self.in_reasoning = false;
self.content.get_or_insert_with(String::new).push_str(&text);
events.push(ChatEvent::Content(text));
}
for piece in choice.delta.tool_calls.unwrap_or_default() {
self.in_reasoning = false;
let index = u32::try_from(piece.index).map_err(|_| {
InferError::Protocol("tool call index out of range".to_string())
})?;
let (name, arguments) = match piece.function {
Some(f) => (f.name, f.arguments.unwrap_or_default()),
None => (None, String::new()),
};
events.push(ChatEvent::ToolCallDelta {
index,
id: piece.id.clone(),
name: name.clone(),
arguments: arguments.clone(),
});
let slot = usize::try_from(index).map_err(|_| {
InferError::Protocol("tool call index out of range".to_string())
})?;
if slot == self.tool_calls.len() {
self.tool_calls.push(PendingCall {
id: piece.id,
name,
arguments,
});
} else {
let Some(call) = self.tool_calls.get_mut(slot) else {
return Err(InferError::Protocol(format!(
"tool call index {index} skips ahead"
)));
};
if call.id.is_none() {
call.id = piece.id;
}
if call.name.is_none() {
call.name = name;
}
call.arguments.push_str(&arguments);
}
}
}
if carries_reasoning {
self.reasoning_tokens = self.timings.predicted_n;
}
Ok(events)
}
pub fn id(&self) -> Option<&str> {
self.id.as_deref()
}
pub fn reasoning_tokens(&self) -> u64 {
self.reasoning_tokens
}
pub fn in_reasoning(&self) -> bool {
self.in_reasoning
}
pub fn is_finished(&self) -> bool {
self.finish_reason.is_some()
}
pub fn finish(self, thinking_capped: bool) -> Result<Completion, InferError> {
let finish_reason = self.finish_reason.ok_or(InferError::StreamClosedEarly)?;
let id = self
.id
.ok_or_else(|| InferError::Protocol("no chunk carried a completion id".to_string()))?;
let mut tool_calls = Vec::new();
for call in self.tool_calls {
let (id, name) = match (call.id, call.name) {
(Some(id), Some(name)) => (id, name),
_ => {
return Err(InferError::Protocol(
"a tool call never got an id or a name".to_string(),
));
}
};
tool_calls.push(proto::ToolCall {
id,
name,
arguments: call.arguments,
});
}
Ok(Completion {
id,
content: self.content,
reasoning_content: self.reasoning_content,
tool_calls,
finish_reason,
timings: self.timings,
reasoning_tokens: self.reasoning_tokens,
thinking_capped,
})
}
}
+1
View File
@@ -1,5 +1,6 @@
//! The llama inference server's chat-completions types and the request builder.
pub mod assemble;
pub mod request;
#[derive(Debug, Clone, PartialEq, Eq)]
+281
View File
@@ -0,0 +1,281 @@
//! 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"
);
}
+31
View File
@@ -0,0 +1,31 @@
{
"id": "chatcmpl-1x092OMeEwmeoSzn7cE1GLkJlskrslR0",
"content": "The number of ways is the number of derangements of 8 elements, which equals **40320**. This counts permutations of 8 rooks (one per row and column) where none lands on the main diagonal (a fixed point would place a rook on the diagonal).",
"reasoning_content": "The question asks about 8 rooks on a chessboard that don't attack each other, with none on the main diagonal.\n\nRooks that don't attack each other means exactly one rook per row and one rook per column. This is equivalent to a permutation matrix. So the number of ways to",
"tool_calls": [],
"finish_reason": "stop",
"cache_n": 15,
"prompt_n": 55,
"predicted_n": 126,
"reasoning_tokens": 63,
"progress": [
[
70,
15,
15
],
[
70,
15,
66
],
[
70,
15,
70
]
],
"reasoning_events": 63,
"content_events": 60,
"tool_events": 0
}
+36
View File
@@ -0,0 +1,36 @@
{
"id": "chatcmpl-LMYge0IefEaKdgrXAjafeR0nrNUicPF2",
"content": "The box is made. (run 6aac4728)",
"reasoning_content": null,
"tool_calls": [],
"finish_reason": "stop",
"cache_n": 0,
"prompt_n": 46,
"predicted_n": 16,
"reasoning_tokens": 0,
"progress": [
[
46,
0,
0
],
[
46,
0,
15
],
[
46,
0,
42
],
[
46,
0,
46
]
],
"reasoning_events": 0,
"content_events": 15,
"tool_events": 0
}
+51
View File
@@ -0,0 +1,51 @@
{
"id": "chatcmpl-l28a5WTzw2ZnKeMg92QlUvh4jZ1Tys3u",
"content": "ok (run 6aac47",
"reasoning_content": null,
"tool_calls": [],
"finish_reason": "length",
"cache_n": 15,
"prompt_n": 7029,
"predicted_n": 8,
"reasoning_tokens": 0,
"progress": [
[
7044,
15,
15
],
[
7044,
15,
2063
],
[
7044,
15,
4111
],
[
7044,
15,
6159
],
[
7044,
15,
6528
],
[
7044,
15,
7040
],
[
7044,
15,
7044
]
],
"reasoning_events": 0,
"content_events": 8,
"tool_events": 0
}
+31
View File
@@ -0,0 +1,31 @@
{
"id": "chatcmpl-oSb4bEgcGEBGNn7Bn6LXqGrh20qMbo4f",
"content": "17 * 23 equals 391.",
"reasoning_content": "17 * 23. Let me compute: 17 * 23 = 17 * 20 + 17 * 3 = 340 + 51 = 391.\n",
"tool_calls": [],
"finish_reason": "stop",
"cache_n": 15,
"prompt_n": 40,
"predicted_n": 64,
"reasoning_tokens": 49,
"progress": [
[
55,
15,
15
],
[
55,
15,
51
],
[
55,
15,
55
]
],
"reasoning_events": 49,
"content_events": 12,
"tool_events": 0
}
+42
View File
@@ -0,0 +1,42 @@
{
"id": "chatcmpl-e9aFhGwvvU0YGs2GTOy83IQWDpVtNvcP",
"content": "I'll read the `/etc/hostname` file for you.\n\n",
"reasoning_content": null,
"tool_calls": [
{
"id": "wgE8iFI58Zni4WCTiCMNp4TzCcM8ou7F",
"name": "read_file",
"arguments": "{\"path\":\"/etc/hostname\"}"
}
],
"finish_reason": "tool_calls",
"cache_n": 0,
"prompt_n": 312,
"predicted_n": 41,
"reasoning_tokens": 0,
"progress": [
[
312,
0,
0
],
[
312,
0,
278
],
[
312,
0,
308
],
[
312,
0,
312
]
],
"reasoning_events": 0,
"content_events": 14,
"tool_events": 7
}
+31
View File
@@ -0,0 +1,31 @@
{
"id": "chatcmpl-1OEWBOVSZ1NthqfpaLOboUWbY9OJixgN",
"content": "Blue",
"reasoning_content": null,
"tool_calls": [],
"finish_reason": "stop",
"cache_n": 15,
"prompt_n": 29,
"predicted_n": 2,
"reasoning_tokens": 0,
"progress": [
[
44,
15,
15
],
[
44,
15,
40
],
[
44,
15,
44
]
],
"reasoning_events": 0,
"content_events": 1,
"tool_events": 0
}
+36
View File
@@ -0,0 +1,36 @@
{
"id": "chatcmpl-Po3MDTUPhb5xQOmzad8o6J8hupMG8HcE",
"content": "Green",
"reasoning_content": null,
"tool_calls": [],
"finish_reason": "stop",
"cache_n": 45,
"prompt_n": 30,
"predicted_n": 2,
"reasoning_tokens": 0,
"progress": [
[
75,
45,
45
],
[
75,
45,
47
],
[
75,
45,
71
],
[
75,
45,
75
]
],
"reasoning_events": 0,
"content_events": 1,
"tool_events": 0
}
+1
View File
@@ -20,6 +20,7 @@ reviewer adds findings under "Reviews" once per milestone.
| M2a/04-loopd-http | 2026-09-18 | done | 2 | fail | none | Added `pub mod http;` to crates/loopd/src/lib.rs and wrote crates/loopd/src/http.rs (423 lines): Request/Head/HttpError, send (exact header order, no `Content-Length` for GET), read_head (buffers across timeouts; `Timeout`/`Closed`/`TooLarge`/`Malformed`), parse_status+parse_head (`HTTP/1.1`/`HTTP/1.0`, status 100-599, lowercased names, duplicates kept, trimmed values), body (chunked/length/close; chunk extensions + trailers skipped), and read_capped. Two fixes: the chunk `Data` phase jumped to `Crlf` on `take==want`, but `want` was capped by the caller's buffer so it switched mid-chunk and returned `malformed chunk` on the recorded fixture — changed to switch on `chunk_remaining==0`; `read_length` reads straight into the caller buffer instead of an 8192 temp slice that would panic for readers larger than 8192. First gate failed on three clippy lints (needless borrows of `format!` results and `map_or` -> `is_some_and`), fixed on the second run. All 15 tests pass and `make gate` prints `gate: ok`. |
| M2a/05-loopd-sse | 2026-09-18 | done | 2 | fail | none | Added `pub mod sse;` to crates/loopd/src/lib.rs and wrote crates/loopd/src/sse.rs: SseItem (Data, Done), SseError (Io, Timeout, Truncated, LineTooLong, NotUtf8) with Display/std::error::Error, and Events<R> which reads a line in bounded 4096-byte chunks, skipping blank/comment/other-field lines and stripping `data:` plus one leading space, decoding UTF-8 only once a whole line has arrived. Two fixes: `drain(..pos)` left the newline in the buffer so blank lines never advanced — changed to `drain(..=pos)` and pop the endings; `process_line` returns `Ok(None)` for a skipped line, which collided with `next_item`'s "stream ended" `Ok(None)` — restructured so a skip continues the loop and only a clean EOF sets `ended`. Both LineTooLong checks (mid-arrival and whole-read) verified by checking the accumulated length before reading and the finished line length. All 7 tests pass; first gate run failed on an unused import, fixed. |
| M2a/06-llama-request | 2026-09-18 | done | 1 | pass | none | Added `pub mod llama;` to crates/loopd/src/lib.rs and wrote crates/loopd/src/llama/mod.rs (ChatMessage, ToolSchema, ChatRequest, ChatEvent, FinishReason, Timings with the server's deserialize shape, Completion, InferError with Display/std::error::Error, and a Client holding the config) and crates/loopd/src/llama/request.rs. build_body serializes the body from `#[derive(Serialize)]` structs so key order is fixed; each message kind is serialized with its own struct (the assistant renders `content` as `""` when `None`, and leaves `reasoning_content`/`tool_calls` out when absent; the top-level `tools` array is omitted when empty; `type` comes from `#[serde(rename = "type")]`); the first `cargo build` after writing the structs missed the `role` field on every message struct, caught by the failing test compile, added. 5 request tests pass; `make gate` prints `gate: ok`. |
| M2a/07-llama-assemble | 2026-09-18 | done | 3 | fail | none | Wrote crates/loopd/src/llama/assemble.rs (Assembler plus server-format Chunk/Choice/Delta/ToolCallPiece/FunctionPiece/PromptProgress structs with no deny_unknown_fields) and registered `pub mod assemble;`. Accumulation: text goes through get_or_insert_with so content/reasoning stay None until a non-empty piece arrives; tool-call pieces land by index via u32::try_from then usize::try_from and get_mut, a new call only at exactly the end, a skip-ahead or an out-of-range index is Protocol; timings update before the reasoning-token count reads predicted_n; finish checks finish_reason (StreamClosedEarly), then id, then every tool call has id and name. First gate failed on two clippy collapsible_if lints; rewrote the two nested ifs as edition-2024 let-chains and re-ran. All 7 assemble tests pass; `make gate` prints `gate: ok`. |
## Reviews