use async_trait::async_trait;
use futures::{Stream, StreamExt, stream};
use crate::{
Chunk, CompletionRequest, LlmProvider, StopReason, Usage, error::DummyError, request::ToolCall,
};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TurnStreamEvent {
TextDelta(String),
ReasoningDelta(String),
ToolStarted {
id: String,
name: String,
},
}
#[derive(Debug, Default, Clone)]
pub struct TurnOutput {
pub text: String,
pub reasoning: String,
pub tool_calls: Vec<ToolCall>,
pub usage: Usage,
pub grounded: bool,
pub stop: Option<StopReason>,
}
pub async fn collect_turn<S, E>(stream: S) -> Result<TurnOutput, E>
where
S: Stream<Item = Result<Chunk, E>> + Unpin,
{
collect_turn_observed(stream, async |_| {}).await
}
pub async fn collect_turn_observed<S, E, F>(mut stream: S, mut on_event: F) -> Result<TurnOutput, E>
where
S: Stream<Item = Result<Chunk, E>> + Unpin,
F: AsyncFnMut(TurnStreamEvent),
{
let mut out = TurnOutput::default();
let mut pending: Vec<ToolCall> = Vec::new();
while let Some(item) = stream.next().await {
match item? {
Chunk::TextDelta(s) => {
on_event(TurnStreamEvent::TextDelta(s.clone())).await;
out.text.push_str(&s);
}
Chunk::ReasoningDelta(s) => {
on_event(TurnStreamEvent::ReasoningDelta(s.clone())).await;
out.reasoning.push_str(&s);
}
Chunk::ToolCallStart {
id,
name,
signature,
} => {
on_event(TurnStreamEvent::ToolStarted {
id: id.clone(),
name: name.clone(),
})
.await;
pending.push(ToolCall {
id,
name,
args_json: String::new(),
signature,
});
}
Chunk::ToolCallArgsDelta {
id,
args_json_delta,
} => {
if let Some(tc) = pending.iter_mut().find(|tc| tc.id == id) {
tc.args_json.push_str(&args_json_delta);
}
}
Chunk::ToolCallEnd { id } => {
if let Some(pos) = pending.iter().position(|tc| tc.id == id) {
out.tool_calls.push(pending.remove(pos));
}
}
Chunk::Usage(u) => out.usage = u,
Chunk::Grounded => out.grounded = true,
Chunk::Stop(r) => {
let keep_tool_use =
out.stop == Some(StopReason::ToolUse) && matches!(r, StopReason::EndTurn);
if !keep_tool_use {
out.stop = Some(r);
}
}
}
}
out.tool_calls.append(&mut pending);
Ok(out)
}
pub const STUB_TOOL_CALL_ENV: &str = "POLYCHROME_STUB_TOOL_CALL";
fn stub_tool_name() -> Option<String> {
std::env::var(STUB_TOOL_CALL_ENV)
.ok()
.filter(|s| !s.is_empty())
}
fn stub_tool_sequence() -> Vec<String> {
stub_tool_name()
.into_iter()
.flat_map(|s| {
s.split(',')
.map(str::trim)
.filter(|s| !s.is_empty())
.map(str::to_owned)
.collect::<Vec<_>>()
})
.collect()
}
fn stub_tool_args_for(name: &str, seq_len: usize) -> String {
let Some(raw) = std::env::var(STUB_TOOL_ARGS_ENV)
.ok()
.filter(|s| !s.is_empty())
else {
return "{}".to_owned();
};
if seq_len <= 1 {
return raw;
}
match serde_json::from_str::<serde_json::Value>(&raw) {
Ok(serde_json::Value::Object(map)) => map
.get(name)
.map_or_else(|| "{}".to_owned(), ToString::to_string),
_ => "{}".to_owned(),
}
}
pub const STUB_TOOL_ARGS_ENV: &str = "POLYCHROME_STUB_TOOL_ARGS";
#[derive(Clone, Copy, Default)]
pub struct StubProvider;
#[async_trait]
impl LlmProvider for StubProvider {
type Error = DummyError;
async fn complete(
&self,
req: CompletionRequest,
) -> Result<futures::stream::BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error> {
let sequence = stub_tool_sequence();
if !sequence.is_empty() {
let results_seen = req
.messages
.iter()
.flat_map(|m| m.content.iter())
.filter(|c| matches!(c, crate::Content::ToolResult(_)))
.count();
if let Some(tool_name) = sequence.get(results_seen) {
let id = format!("stub-call-{}", results_seen + 1);
let chunks = vec![
Ok(Chunk::tool_call_start(&id, tool_name)),
Ok(Chunk::tool_call_args_delta(
&id,
stub_tool_args_for(tool_name, sequence.len()),
)),
Ok(Chunk::tool_call_end(&id)),
Ok(Chunk::Stop(StopReason::ToolUse)),
];
return Ok(stream::iter(chunks).boxed());
}
}
let chunks = vec![
Ok(Chunk::text_delta("Hello from the ")),
Ok(Chunk::text_delta("stub provider.")),
Ok(Chunk::Usage(Usage {
input_tokens: 5,
output_tokens: 4,
..Default::default()
})),
Ok(Chunk::Stop(StopReason::EndTurn)),
];
Ok(stream::iter(chunks).boxed())
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::pedantic, clippy::nursery, missing_docs)]
use super::*;
#[tokio::test]
async fn stub_provider_collects_into_text() {
let stream = StubProvider
.complete(CompletionRequest::new("stub"))
.await
.expect("stream opens");
let out = collect_turn(stream).await.expect("collect");
assert_eq!(out.text, "Hello from the stub provider.");
assert!(out.tool_calls.is_empty());
assert_eq!(out.usage.output_tokens, 4);
assert_eq!(out.stop, Some(StopReason::EndTurn));
}
#[tokio::test]
async fn stub_provider_emits_the_configured_tool_args() {
temp_env::async_with_vars(
[
(STUB_TOOL_CALL_ENV, Some("search")),
(STUB_TOOL_ARGS_ENV, Some(r#"{"q":"rust"}"#)),
],
async {
let stream = StubProvider
.complete(CompletionRequest::new("stub"))
.await
.expect("stream opens");
let out = collect_turn(stream).await.expect("collect");
assert_eq!(out.tool_calls.len(), 1);
assert_eq!(out.tool_calls[0].args_json, r#"{"q":"rust"}"#);
},
)
.await;
}
#[tokio::test]
async fn stub_provider_defaults_tool_args_to_empty_object() {
temp_env::async_with_vars(
[
(STUB_TOOL_CALL_ENV, Some("search")),
(STUB_TOOL_ARGS_ENV, None),
],
async {
let stream = StubProvider
.complete(CompletionRequest::new("stub"))
.await
.expect("stream opens");
let out = collect_turn(stream).await.expect("collect");
assert_eq!(out.tool_calls.len(), 1);
assert_eq!(out.tool_calls[0].args_json, "{}");
},
)
.await;
}
#[tokio::test]
async fn stub_provider_canned_text_path_is_unchanged_by_the_args_knob() {
temp_env::async_with_vars(
[
(STUB_TOOL_CALL_ENV, None),
(STUB_TOOL_ARGS_ENV, Some(r#"{"q":"rust"}"#)),
],
async {
let stream = StubProvider
.complete(CompletionRequest::new("stub"))
.await
.expect("stream opens");
let out = collect_turn(stream).await.expect("collect");
assert_eq!(out.text, "Hello from the stub provider.");
assert!(out.tool_calls.is_empty());
},
)
.await;
}
#[tokio::test]
async fn collect_assembles_tool_call_from_deltas() {
let chunks: Vec<Result<Chunk, DummyError>> = vec![
Ok(Chunk::text_delta("calling ")),
Ok(Chunk::tool_call_start("c1", "search")),
Ok(Chunk::tool_call_args_delta("c1", r#"{"q":"#)),
Ok(Chunk::tool_call_args_delta("c1", r#""rust"}"#)),
Ok(Chunk::tool_call_end("c1")),
Ok(Chunk::Stop(StopReason::ToolUse)),
];
let out = collect_turn(stream::iter(chunks)).await.expect("collect");
assert_eq!(out.text, "calling ");
assert_eq!(out.tool_calls.len(), 1);
assert_eq!(out.tool_calls[0].name, "search");
assert_eq!(out.tool_calls[0].args_json, r#"{"q":"rust"}"#);
assert_eq!(out.stop, Some(StopReason::ToolUse));
}
#[tokio::test]
async fn collect_keeps_parallel_tool_calls_with_deferred_ends() {
let chunks: Vec<Result<Chunk, DummyError>> = vec![
Ok(Chunk::tool_call_start("c0", "search")),
Ok(Chunk::tool_call_args_delta("c0", r#"{"q":"a"}"#)),
Ok(Chunk::tool_call_start("c1", "fetch")),
Ok(Chunk::tool_call_args_delta("c1", r#"{"u":"b"}"#)),
Ok(Chunk::tool_call_end("c0")),
Ok(Chunk::tool_call_end("c1")),
Ok(Chunk::Stop(StopReason::ToolUse)),
];
let out = collect_turn(stream::iter(chunks)).await.expect("collect");
assert_eq!(out.tool_calls.len(), 2, "both parallel calls preserved");
assert_eq!(out.tool_calls[0].id, "c0");
assert_eq!(out.tool_calls[0].name, "search");
assert_eq!(out.tool_calls[0].args_json, r#"{"q":"a"}"#);
assert_eq!(out.tool_calls[1].id, "c1");
assert_eq!(out.tool_calls[1].name, "fetch");
assert_eq!(out.tool_calls[1].args_json, r#"{"u":"b"}"#);
assert_eq!(out.stop, Some(StopReason::ToolUse));
}
#[tokio::test]
async fn collect_flushes_a_call_left_open_at_eof() {
let chunks: Vec<Result<Chunk, DummyError>> = vec![
Ok(Chunk::tool_call_start("c0", "search")),
Ok(Chunk::tool_call_args_delta("c0", r#"{"q":"a"}"#)),
Ok(Chunk::Stop(StopReason::ToolUse)),
];
let out = collect_turn(stream::iter(chunks)).await.expect("collect");
assert_eq!(out.tool_calls.len(), 1);
assert_eq!(out.tool_calls[0].args_json, r#"{"q":"a"}"#);
}
#[tokio::test]
async fn tool_use_stop_is_sticky_against_later_end_turn() {
let chunks: Vec<Result<Chunk, DummyError>> = vec![
Ok(Chunk::tool_call_start("c1", "search")),
Ok(Chunk::tool_call_end("c1")),
Ok(Chunk::Stop(StopReason::ToolUse)),
Ok(Chunk::Stop(StopReason::EndTurn)),
];
let out = collect_turn(stream::iter(chunks)).await.expect("collect");
assert_eq!(out.stop, Some(StopReason::ToolUse));
}
#[tokio::test]
async fn hard_stop_wins_over_earlier_tool_use() {
let chunks: Vec<Result<Chunk, DummyError>> = vec![
Ok(Chunk::tool_call_start("c1", "search")),
Ok(Chunk::tool_call_end("c1")),
Ok(Chunk::Stop(StopReason::ToolUse)),
Ok(Chunk::Stop(StopReason::MaxTokens)),
];
let out = collect_turn(stream::iter(chunks)).await.expect("collect");
assert_eq!(out.stop, Some(StopReason::MaxTokens));
}
#[tokio::test]
async fn collect_folds_reasoning_separately_from_text() {
let chunks: Vec<Result<Chunk, DummyError>> = vec![
Ok(Chunk::reasoning_delta("first ")),
Ok(Chunk::reasoning_delta("thought")),
Ok(Chunk::text_delta("the answer")),
Ok(Chunk::Stop(StopReason::EndTurn)),
];
let out = collect_turn(stream::iter(chunks)).await.expect("collect");
assert_eq!(out.reasoning, "first thought");
assert_eq!(out.text, "the answer");
}
#[tokio::test]
async fn observed_reasoning_deltas_are_emitted() {
let chunks: Vec<Result<Chunk, DummyError>> = vec![
Ok(Chunk::reasoning_delta("hmm")),
Ok(Chunk::text_delta("ok")),
Ok(Chunk::Stop(StopReason::EndTurn)),
];
let mut events = Vec::new();
let out = collect_turn_observed(stream::iter(chunks), async |e| events.push(e))
.await
.expect("collect");
assert_eq!(out.reasoning, "hmm");
assert!(events.contains(&TurnStreamEvent::ReasoningDelta("hmm".to_owned())));
assert!(events.contains(&TurnStreamEvent::TextDelta("ok".to_owned())));
}
#[tokio::test]
async fn collect_folds_grounded_evidence() {
let chunks: Vec<Result<Chunk, DummyError>> = vec![
Ok(Chunk::text_delta("per recent sources, ")),
Ok(Chunk::grounded()),
Ok(Chunk::text_delta("it's sunny.")),
Ok(Chunk::Stop(StopReason::EndTurn)),
];
let out = collect_turn(stream::iter(chunks)).await.expect("collect");
assert!(
out.grounded,
"a Grounded chunk anywhere in the stream must fold to true"
);
}
#[tokio::test]
async fn collect_defaults_grounded_to_false() {
let chunks: Vec<Result<Chunk, DummyError>> = vec![
Ok(Chunk::text_delta("the answer")),
Ok(Chunk::Stop(StopReason::EndTurn)),
];
let out = collect_turn(stream::iter(chunks)).await.expect("collect");
assert!(!out.grounded);
}
#[tokio::test]
async fn collect_propagates_error() {
let chunks: Vec<Result<Chunk, DummyError>> = vec![
Ok(Chunk::text_delta("partial")),
Err(DummyError::Other("mid-stream fault".to_owned())),
];
let res = collect_turn(stream::iter(chunks)).await;
assert!(res.is_err());
}
#[tokio::test]
async fn collect_turn_observed_backpressures_on_a_full_bounded_channel() {
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::task::{Context, Poll};
use futures::SinkExt;
struct CountingStream<S> {
inner: S,
polls: Arc<AtomicUsize>,
}
impl<S: Stream + Unpin> Stream for CountingStream<S> {
type Item = S::Item;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.polls.fetch_add(1, Ordering::SeqCst);
let this = self.get_mut();
Pin::new(&mut this.inner).poll_next(cx)
}
}
let chunks: Vec<Result<Chunk, DummyError>> = vec![
Ok(Chunk::text_delta("a")),
Ok(Chunk::text_delta("b")),
Ok(Chunk::text_delta("c")),
Ok(Chunk::Stop(StopReason::EndTurn)),
];
let polls = Arc::new(AtomicUsize::new(0));
let stream = CountingStream {
inner: stream::iter(chunks),
polls: polls.clone(),
};
let (tx, mut rx) = futures::channel::mpsc::channel::<TurnStreamEvent>(1);
let handle = tokio::spawn(async move {
let mut tx = tx;
collect_turn_observed(stream, async move |ev| {
let _ = tx.send(ev).await;
})
.await
});
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
assert!(
!handle.is_finished(),
"fold must not complete while the channel is full and undrained"
);
let stalled_at = polls.load(Ordering::SeqCst);
assert!(
stalled_at < 4,
"stream must not have been fully drained while the channel is full"
);
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
assert_eq!(
polls.load(Ordering::SeqCst),
stalled_at,
"poll count must plateau while the channel is full — proof the stall is real"
);
let mut texts = Vec::new();
while texts.len() < 3 {
match rx.next().await {
Some(TurnStreamEvent::TextDelta(s)) => texts.push(s),
Some(_) => {}
None => break,
}
}
let out = tokio::time::timeout(std::time::Duration::from_secs(5), handle)
.await
.expect("fold must complete once the channel drains")
.expect("task join")
.expect("collect");
assert_eq!(out.text, "abc");
assert_eq!(texts, vec!["a".to_owned(), "b".to_owned(), "c".to_owned()]);
}
}