From df7b070e715e42c14498237b24656e8beec84552 Mon Sep 17 00:00:00 2001 From: Chris Guidry Date: Fri, 31 Jul 2026 10:33:01 -0400 Subject: [PATCH] Dispatch tool calls as their fragments complete, not after the round Dm::turn drained the whole response, then dispatched every tool call at once: the player watched static narration while the model generated JSON for several calls, then all the results appeared together. Tool-call fragments arrive in non-decreasing index order over the wire, so a call is complete the moment a fragment for a higher index arrives, and can roll and reach the screen while the model still generates the next one. ChatStream's iterator now yields StreamItem::Text and StreamItem::Call as each completes, replacing the old text-only iterator plus into_tool_calls collected after the stream ended. Dm::turn dispatches a Call item immediately on a normal round, and still assembles the same assistant-then-tool-results messages and the same request wire format once the round finishes; a withheld round drops any Call it sees, unchanged from before. Request/response shape, history assembly, and cancel semantics are unchanged; only the timing of dispatch and display moves earlier. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01CUXjWo1zGFhdJUig1hcQGf --- src/chat/client.rs | 254 ++++++++++++++++++++++++---------- src/chat/mod.rs | 2 +- src/dm/dm_tool_round_tests.rs | 50 +++++++ src/dm/mod.rs | 90 ++++++------ 4 files changed, 286 insertions(+), 110 deletions(-) diff --git a/src/chat/client.rs b/src/chat/client.rs index 0835cc3..f818804 100644 --- a/src/chat/client.rs +++ b/src/chat/client.rs @@ -1,4 +1,4 @@ -use std::collections::BTreeMap; +use std::collections::VecDeque; use std::fmt; use std::io; @@ -64,19 +64,33 @@ impl Client { } } -/// A streamed chat completion response, yielding each text delta as it -/// arrives. +/// One piece of a streamed chat completion response. +#[derive(Debug, Clone, PartialEq)] +pub enum StreamItem { + /// A piece of the reply's narration, as it arrives. + Text(String), + /// A tool call, once every fragment for it has arrived. + Call(ToolCall), +} + +/// A streamed chat completion response, yielding text as it arrives and +/// each tool call as soon as its fragments are complete. /// /// The stream ends when the server sends `[DONE]` or closes the connection. -/// Alongside the text, it assembles any tool calls the response made; -/// collect them with `into_tool_calls` once the stream ends. +/// The OpenAI streaming contract sends a tool call's fragments in +/// non-decreasing index order, so a call is complete, and ready to yield, +/// the moment a fragment for a higher index arrives; the call at the +/// highest index seen completes when the stream ends. pub struct ChatStream { events: SseReader>, - tool_calls: BTreeMap, + /// The call at the highest index seen so far, still being assembled. + current_call: Option<(usize, PartialToolCall)>, + /// Items ready to yield before the next chunk is read: one chunk can + /// both complete a call and carry text. + queue: VecDeque, } -/// A tool call being assembled from fragments, keyed by its index in the -/// response. +/// A tool call being assembled from fragments. #[derive(Default)] struct PartialToolCall { id: String, @@ -84,38 +98,49 @@ struct PartialToolCall { arguments: String, } +impl PartialToolCall { + fn into_tool_call(self) -> ToolCall { + ToolCall { + id: self.id, + kind: "function".to_string(), + function: ToolCallFunction { + name: self.name, + arguments: self.arguments, + }, + } + } +} + impl ChatStream { fn new(reader: ureq::BodyReader<'static>) -> Self { Self { events: SseReader::new(reader), - tool_calls: BTreeMap::new(), + current_call: None, + queue: VecDeque::new(), } } - /// Consumes the stream and returns its assembled tool calls in index - /// order. - /// - /// Empty when the response carried no tool calls, such as a plain text - /// reply. Call this once the stream has run to completion. - pub fn into_tool_calls(self) -> Vec { - self.tool_calls - .into_values() - .map(|partial| ToolCall { - id: partial.id, - kind: "function".to_string(), - function: ToolCallFunction { - name: partial.name, - arguments: partial.arguments, - }, - }) - .collect() - } - - /// Merges one chunk's tool-call fragments into the calls assembled so - /// far. + /// Merges one chunk's tool-call fragments into the call being + /// assembled, queuing the call at the previous index as soon as a + /// fragment names a higher one. fn absorb_tool_call_fragments(&mut self, fragments: Vec) { for fragment in fragments { - let partial = self.tool_calls.entry(fragment.index).or_default(); + if self + .current_call + .as_ref() + .is_some_and(|(index, _)| *index != fragment.index) + { + let (_, partial) = self + .current_call + .take() + .expect("current_call is Some, just checked above"); + self.queue + .push_back(StreamItem::Call(partial.into_tool_call())); + } + + let (_, partial) = self + .current_call + .get_or_insert_with(|| (fragment.index, PartialToolCall::default())); if let Some(id) = fragment.id { partial.id = id; } @@ -133,13 +158,23 @@ impl ChatStream { } impl Iterator for ChatStream { - type Item = Result; + type Item = Result; fn next(&mut self) -> Option { loop { - let payload = match self.events.next()? { - Ok(payload) => payload, - Err(error) => return Some(Err(ChatError::Io(error))), + if let Some(item) = self.queue.pop_front() { + return Some(Ok(item)); + } + + let payload = match self.events.next() { + None => { + return self + .current_call + .take() + .map(|(_, partial)| Ok(StreamItem::Call(partial.into_tool_call()))); + } + Some(Ok(payload)) => payload, + Some(Err(error)) => return Some(Err(ChatError::Io(error))), }; let chunk: ChatChunk = match serde_json::from_str(&payload) { Ok(chunk) => chunk, @@ -151,13 +186,11 @@ impl Iterator for ChatStream { if let Some(fragments) = choice.delta.tool_calls { self.absorb_tool_call_fragments(fragments); } - let Some(content) = choice.delta.content else { - continue; - }; - if content.is_empty() { - continue; + if let Some(content) = choice.delta.content + && !content.is_empty() + { + self.queue.push_back(StreamItem::Text(content)); } - return Some(Ok(content)); } } } @@ -315,13 +348,19 @@ mod tests { let (url, requests, server) = fake_server(sse_response(body)); let client = client_for(url); - let deltas: Vec = client + let items: Vec = client .stream(&a_message(), &[]) .unwrap() - .map(|delta| delta.unwrap()) + .map(|item| item.unwrap()) .collect(); - assert_eq!(deltas, vec!["Hello".to_string(), " world".to_string()]); + assert_eq!( + items, + vec![ + StreamItem::Text("Hello".to_string()), + StreamItem::Text(" world".to_string()), + ] + ); server.join().unwrap(); let request = requests.recv().unwrap(); assert_eq!(request.method, "POST"); @@ -403,27 +442,30 @@ mod tests { let (url, _requests, server) = fake_server(sse_response(body)); let client = client_for(url); - let deltas: Vec = client + let items: Vec = client .stream(&a_message(), &[]) .unwrap() - .map(|delta| delta.unwrap()) + .map(|item| item.unwrap()) .collect(); - assert_eq!(deltas, vec!["Hi".to_string()]); + assert_eq!(items, vec![StreamItem::Text("Hi".to_string())]); server.join().unwrap(); } #[test] - fn a_plain_text_response_assembles_to_no_tool_calls() { + fn a_plain_text_response_yields_only_text_items() { let body = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n\n\ data: [DONE]\n\n"; let (url, _requests, server) = fake_server(sse_response(body)); let client = client_for(url); - let mut stream = client.stream(&a_message(), &[]).unwrap(); - stream.by_ref().for_each(drop); + let items: Vec = client + .stream(&a_message(), &[]) + .unwrap() + .map(|item| item.unwrap()) + .collect(); - assert_eq!(stream.into_tool_calls(), vec![]); + assert_eq!(items, vec![StreamItem::Text("Hello".to_string())]); server.join().unwrap(); } @@ -436,19 +478,22 @@ mod tests { let (url, _requests, server) = fake_server(sse_response(body)); let client = client_for(url); - let mut stream = client.stream(&a_message(), &[]).unwrap(); - stream.by_ref().for_each(drop); + let items: Vec = client + .stream(&a_message(), &[]) + .unwrap() + .map(|item| item.unwrap()) + .collect(); assert_eq!( - stream.into_tool_calls(), - vec![ToolCall { + items, + vec![StreamItem::Call(ToolCall { id: "call_1".to_string(), kind: "function".to_string(), function: ToolCallFunction { name: "roll".to_string(), arguments: "{\"notation\":\"1d20\"}".to_string(), }, - }] + })] ); server.join().unwrap(); } @@ -461,19 +506,22 @@ mod tests { let (url, _requests, server) = fake_server(sse_response(body)); let client = client_for(url); - let mut stream = client.stream(&a_message(), &[]).unwrap(); - stream.by_ref().for_each(drop); + let items: Vec = client + .stream(&a_message(), &[]) + .unwrap() + .map(|item| item.unwrap()) + .collect(); assert_eq!( - stream.into_tool_calls(), - vec![ToolCall { + items, + vec![StreamItem::Call(ToolCall { id: "call_1".to_string(), kind: "function".to_string(), function: ToolCallFunction { name: "roll".to_string(), arguments: "{}".to_string(), }, - }] + })] ); server.join().unwrap(); } @@ -488,15 +536,75 @@ mod tests { let (url, _requests, server) = fake_server(sse_response(body)); let client = client_for(url); - let mut stream = client.stream(&a_message(), &[]).unwrap(); - stream.by_ref().for_each(drop); + let items: Vec = client + .stream(&a_message(), &[]) + .unwrap() + .map(|item| item.unwrap()) + .collect(); + + assert_eq!( + items, + vec![ + StreamItem::Call(ToolCall { + id: "call_1".to_string(), + kind: "function".to_string(), + function: ToolCallFunction { + name: "roll".to_string(), + arguments: "{}".to_string(), + }, + }), + StreamItem::Call(ToolCall { + id: "call_2".to_string(), + kind: "function".to_string(), + function: ToolCallFunction { + name: "roll".to_string(), + arguments: "{}".to_string(), + }, + }), + ] + ); + server.join().unwrap(); + } + + #[test] + fn text_and_completed_calls_interleave_in_stream_order() { + let body = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n\n\ + data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"function\":{\"name\":\"roll\",\"arguments\":\"{}\"}}]},\"finish_reason\":null}]}\n\n\ + data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":1,\"id\":\"call_2\",\"function\":{\"name\":\"roll\",\"arguments\":\"{}\"}}]},\"finish_reason\":null}]}\n\n\ + data: {\"choices\":[{\"delta\":{\"content\":\"world\"},\"finish_reason\":null}]}\n\n\ + data: [DONE]\n\n"; + let (url, _requests, server) = fake_server(sse_response(body)); + let client = client_for(url); - let ids: Vec = stream - .into_tool_calls() - .into_iter() - .map(|call| call.id) + let items: Vec = client + .stream(&a_message(), &[]) + .unwrap() + .map(|item| item.unwrap()) .collect(); - assert_eq!(ids, vec!["call_1".to_string(), "call_2".to_string()]); + + assert_eq!( + items, + vec![ + StreamItem::Text("Hello".to_string()), + StreamItem::Call(ToolCall { + id: "call_1".to_string(), + kind: "function".to_string(), + function: ToolCallFunction { + name: "roll".to_string(), + arguments: "{}".to_string(), + }, + }), + StreamItem::Text("world".to_string()), + StreamItem::Call(ToolCall { + id: "call_2".to_string(), + kind: "function".to_string(), + function: ToolCallFunction { + name: "roll".to_string(), + arguments: "{}".to_string(), + }, + }), + ] + ); server.join().unwrap(); } @@ -509,7 +617,10 @@ mod tests { let mut stream = client.stream(&a_message(), &[]).unwrap(); - assert_eq!(stream.next().unwrap().unwrap(), "Hi"); + assert_eq!( + stream.next().unwrap().unwrap(), + StreamItem::Text("Hi".to_string()) + ); assert!(stream.next().is_none()); server.join().unwrap(); } @@ -533,7 +644,10 @@ mod tests { let mut stream = client.stream(&a_message(), &[]).unwrap(); - assert_eq!(stream.next().unwrap().unwrap(), "Hi"); + assert_eq!( + stream.next().unwrap().unwrap(), + StreamItem::Text("Hi".to_string()) + ); assert!(matches!(stream.next(), Some(Err(ChatError::Io(_))))); server.join().unwrap(); } diff --git a/src/chat/mod.rs b/src/chat/mod.rs index 4044530..f193b7d 100644 --- a/src/chat/mod.rs +++ b/src/chat/mod.rs @@ -2,7 +2,7 @@ pub mod client; pub mod sse; pub mod wire; -pub use client::{ChatError, ChatStream, Client}; +pub use client::{ChatError, ChatStream, Client, StreamItem}; pub use sse::SseReader; pub use wire::{ ChatChunk, ChatRequest, ChunkChoice, Delta, FunctionDelta, Message, Role, ToolCall, diff --git a/src/dm/dm_tool_round_tests.rs b/src/dm/dm_tool_round_tests.rs index 28669f1..a827e58 100644 --- a/src/dm/dm_tool_round_tests.rs +++ b/src/dm/dm_tool_round_tests.rs @@ -97,6 +97,23 @@ fn reply_response(text: &str) -> Vec { round_response(text, &[]) } +/// One chunk carrying a single content delta. +fn content_chunk(text: &str) -> String { + format!( + "data: {{\"choices\":[{{\"delta\":{{\"content\":{}}},\"finish_reason\":null}}]}}\n\n", + json!(text) + ) +} + +/// One chunk carrying a single tool-call fragment at `index`, with `id`, +/// `name`, and `arguments` all present in this one fragment. +fn tool_call_chunk(index: usize, id: &str, name: &str, arguments: &str) -> String { + format!( + "data: {{\"choices\":[{{\"delta\":{{\"tool_calls\":[{{\"index\":{index},\"id\":\"{id}\",\"function\":{{\"name\":\"{name}\",\"arguments\":{}}}}}]}},\"finish_reason\":null}}]}}\n\n", + json!(arguments) + ) +} + #[test] fn a_tool_round_then_a_reply_the_second_request_carries_the_tool_result() { let (url, requests, server) = fake_server(vec![ @@ -260,6 +277,39 @@ fn a_multi_call_round_dispatches_in_order() { assert_eq!(messages[4]["tool_call_id"], "call_2"); } +#[test] +fn a_completed_calls_tool_event_arrives_before_the_stream_finishes_generating_the_next_one() { + let body = format!( + "{}{}{}{}data: [DONE]\n\n", + content_chunk("You search the room. "), + tool_call_chunk(0, "call_1", "roll", &roll_args("1d20", "public")), + tool_call_chunk(1, "call_2", "roll", &roll_args("1d6", "public")), + content_chunk(" Nothing happens."), + ); + let (url, _requests, server) = + fake_server(vec![sse_response(&body), reply_response("You move on.")]); + let mut dm = dm_with_seeded_dice(url, 0); + let mut event_kinds: Vec<&'static str> = Vec::new(); + + dm.turn("I search and act.", &mut |event| { + event_kinds.push(match event { + TurnDelta::Text(_) => "text", + TurnDelta::Tool(_) => "tool", + }); + ControlFlow::Continue(()) + }) + .unwrap(); + + // Call 1 completes, and dispatches, the moment call 2's fragment + // arrives; call 2 itself completes only once the round's stream + // ends. So the tool event for call 1 lands between the round's two + // narration deltas, rather than after both calls have finished + // streaming. The final "text" is the next round's reply, which ends + // the turn. + assert_eq!(event_kinds, vec!["text", "tool", "text", "tool", "text"]); + server.join().unwrap(); +} + #[test] fn a_narrated_round_resets_the_tool_only_counter() { let before = MAX_TOOL_ONLY_ROUNDS - 1; diff --git a/src/dm/mod.rs b/src/dm/mod.rs index 0fefb52..07a4dd3 100644 --- a/src/dm/mod.rs +++ b/src/dm/mod.rs @@ -7,7 +7,7 @@ use rand::rngs::StdRng; use ratatui::text::Text; use serde_json::Value; -use crate::chat::{ChatError, Client, Message, Role}; +use crate::chat::{ChatError, Client, Message, Role, StreamItem}; use crate::config::Config; use tools::Toolbox; use tools::dice::DiceTool; @@ -106,13 +106,16 @@ impl Dm { /// and tool events through `on_delta` as they arrive. /// /// Each round sends the history, `input`, and the rounds so far in - /// this turn, streams a response, and checks whether it made any tool - /// calls. A round with no calls ends the turn: `input`, every - /// message from an earlier round, and the final reply join the - /// history in order, and this returns `Turn::Reply` carrying the - /// final round's narration. A round with calls runs each one in - /// order through the toolbox, appends the assistant's call and each - /// tool's result to the turn, and goes again. + /// this turn, then streams the response. A tool call dispatches + /// through the toolbox the moment its fragments finish streaming, + /// rather than waiting for the round to end, so the player sees a + /// call's result while the model is still generating the next one. + /// A round with no calls ends the turn: `input`, every message from + /// an earlier round, and the final reply join the history in order, + /// and this returns `Turn::Reply` carrying the final round's + /// narration. A round with calls appends the assistant's message, + /// carrying the narration and every call, and each tool's result, in + /// call order, to the turn, and goes again. /// /// After [`MAX_TOOL_ONLY_ROUNDS`] rounds in a row that called tools /// and narrated nothing, or after [`MAX_ROUNDS`] rounds regardless of @@ -163,14 +166,46 @@ impl Dm { let mut stream = self.client.stream(&messages, &definitions)?; let mut narration = String::new(); - for delta in &mut stream { - let delta = delta?; - narration.push_str(&delta); - if on_delta(TurnDelta::Text(delta)).is_break() { - return Ok(Turn::Cancelled(narration)); + let mut calls = Vec::new(); + let mut tool_results = Vec::new(); + + for item in &mut stream { + match item? { + StreamItem::Text(text) => { + narration.push_str(&text); + if on_delta(TurnDelta::Text(text)).is_break() { + return Ok(Turn::Cancelled(narration)); + } + } + // A withheld round asked the model not to call tools, so any + // call it sends anyway is not dispatched and not recorded. + StreamItem::Call(_) if withheld => {} + StreamItem::Call(call) => { + calls.push(call.clone()); + let args = match serde_json::from_str::(&call.function.arguments) { + Ok(args) => args, + Err(_) => { + tool_results.push(tool_message( + &call.id, + format!( + "arguments for `{}` were not valid JSON: `{}`", + call.function.name, call.function.arguments + ), + )); + continue; + } + }; + + let outcome = self.toolbox.call(&call.function.name, &args); + if let Some(display) = outcome.display + && on_delta(TurnDelta::Tool(display)).is_break() + { + return Ok(Turn::Cancelled(narration)); + } + tool_results.push(tool_message(&call.id, outcome.for_model)); + } } } - let calls = stream.into_tool_calls(); if withheld || calls.is_empty() { self.history.push(user_message); @@ -182,33 +217,10 @@ impl Dm { round_messages.push(Message { role: Role::Assistant, content: narration.clone(), - tool_calls: Some(calls.clone()), + tool_calls: Some(calls), tool_call_id: None, }); - - for call in &calls { - let args = match serde_json::from_str::(&call.function.arguments) { - Ok(args) => args, - Err(_) => { - round_messages.push(tool_message( - &call.id, - format!( - "arguments for `{}` were not valid JSON: `{}`", - call.function.name, call.function.arguments - ), - )); - continue; - } - }; - - let outcome = self.toolbox.call(&call.function.name, &args); - if let Some(display) = outcome.display - && on_delta(TurnDelta::Tool(display)).is_break() - { - return Ok(Turn::Cancelled(narration)); - } - round_messages.push(tool_message(&call.id, outcome.for_model)); - } + round_messages.extend(tool_results); tool_only_rounds = if narration.is_empty() { tool_only_rounds + 1 -- 2.51.2