use tokio::sync::mpsc;
use super::error::{ModelError, Result};
use super::reasoning::ReasoningChunk;
use super::tool_call::ToolCall;
use super::types::{FinishReason, ProviderContinuation, TokenUsage};
#[derive(Debug, Clone)]
pub enum StreamEvent {
Text(String),
Reasoning(ReasoningChunk),
ToolCall(ToolCall),
Status(String),
Done {
usage: Option<TokenUsage>,
provider_continuation: Option<ProviderContinuation>,
stop_reason: Option<FinishReason>,
},
}
pub type StreamSink = mpsc::Sender<StreamEvent>;
pub type StatusNotify = std::sync::Arc<dyn Fn(&str) + Send + Sync>;
pub async fn emit(sink: Option<&StreamSink>, event: StreamEvent) -> Result<()> {
let Some(sink) = sink else {
return Ok(());
};
sink.send(event)
.await
.map_err(|_| ModelError::StreamError("stream receiver closed".to_string()))
}
pub async fn emit_all(
sink: Option<&StreamSink>,
events: impl IntoIterator<Item = StreamEvent>,
) -> Result<()> {
for event in events {
emit(sink, event).await?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn stream_event_clone() {
let ev = StreamEvent::Text("hello".to_string());
let cloned = ev.clone();
match (ev, cloned) {
(StreamEvent::Text(a), StreamEvent::Text(b)) => assert_eq!(a, b),
_ => panic!("clone should produce same variant"),
}
}
#[test]
fn stream_event_done_carries_the_whole_terminal_payload() {
let ev = StreamEvent::Done {
usage: Some(TokenUsage::provider(10, 20)),
provider_continuation: Some(ProviderContinuation::Anthropic {
signature: "sig".to_string(),
}),
stop_reason: Some(FinishReason::ToolUse),
};
match ev {
StreamEvent::Done {
usage,
provider_continuation,
stop_reason,
} => {
assert_eq!(usage.expect("usage").total_tokens(), 30);
assert!(matches!(
provider_continuation,
Some(ProviderContinuation::Anthropic { .. })
));
assert_eq!(stop_reason, Some(FinishReason::ToolUse));
},
_ => panic!("expected Done"),
}
}
#[test]
fn stream_event_reasoning_with_chunk() {
let chunk = ReasoningChunk {
text: "weighing options".to_string(),
signature: None,
};
let ev = StreamEvent::Reasoning(chunk.clone());
match ev {
StreamEvent::Reasoning(c) => {
assert_eq!(c.text, chunk.text);
assert_eq!(c.signature, chunk.signature);
},
_ => panic!("expected Reasoning"),
}
}
#[test]
fn sink_is_send_sync() {
fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<StreamSink>();
assert_send_sync::<StatusNotify>();
}
#[tokio::test]
async fn emit_all_preserves_order_and_applies_backpressure() {
let (tx, mut rx) = mpsc::channel::<StreamEvent>(1);
let batch = vec![
StreamEvent::Text("a".to_string()),
StreamEvent::Reasoning(ReasoningChunk {
text: "r".to_string(),
signature: None,
}),
StreamEvent::Text("b".to_string()),
StreamEvent::Done {
usage: None,
provider_continuation: None,
stop_reason: None,
},
];
let producer = tokio::spawn(async move { emit_all(Some(&tx), batch).await });
let mut seen = Vec::new();
while let Some(event) = rx.recv().await {
seen.push(match event {
StreamEvent::Text(s) => s,
StreamEvent::Reasoning(c) => c.text,
StreamEvent::Done { .. } => "done".to_string(),
StreamEvent::ToolCall(_) | StreamEvent::Status(_) => "other".to_string(),
});
}
producer.await.expect("join").expect("emit_all");
assert_eq!(seen, vec!["a", "r", "b", "done"]);
}
#[tokio::test]
async fn emit_stops_the_read_loop_once_the_receiver_is_gone() {
let (tx, rx) = mpsc::channel::<StreamEvent>(4);
drop(rx);
let err = emit(Some(&tx), StreamEvent::Text("x".to_string()))
.await
.expect_err("closed receiver");
assert!(matches!(err, ModelError::StreamError(_)));
}
#[tokio::test]
async fn emit_without_a_sink_is_a_no_op() {
emit(None, StreamEvent::Text("x".to_string()))
.await
.expect("no sink");
}
}