mod identity;
mod parts;
use crate::completion::{CompletionError, CompletionResponse, Usage};
use crate::message::{
AssistantContent, Reasoning, ReasoningContent, Text, ToolCall, ToolFunction, ToolResult,
};
use crate::wasm_compat::WasmCompatSend;
use futures::stream::{AbortHandle, Abortable};
use futures::{Stream, StreamExt};
pub use identity::{MintKind, StreamPartId, SyntheticIds, WireId};
use parts::PartsAccumulator;
use serde::{Deserialize, Serialize};
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::AtomicBool;
use std::task::{Context, Poll};
use tokio::sync::watch;
pub struct PauseControl {
pub(crate) paused_tx: watch::Sender<bool>,
pub(crate) paused_rx: watch::Receiver<bool>,
}
impl PauseControl {
pub fn new() -> Self {
let (paused_tx, paused_rx) = watch::channel(false);
Self {
paused_tx,
paused_rx,
}
}
pub fn pause(&self) {
let _ = self.paused_tx.send(true);
}
pub fn resume(&self) {
let _ = self.paused_tx.send(false);
}
pub fn is_paused(&self) -> bool {
*self.paused_rx.borrow()
}
}
impl Default for PauseControl {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
pub enum ToolCallDeltaContent {
Name(String),
Delta(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum UnparseableToolInput {
Drop,
EmptyObject,
Error,
Keep,
}
#[derive(Debug, Clone)]
pub struct ToolInputEnd {
pub id: StreamPartId,
pub tool_id: Option<WireId>,
pub name: Option<String>,
pub arguments: Option<serde_json::Value>,
pub call_id: Option<String>,
pub signature: Option<String>,
pub additional_params: Option<serde_json::Value>,
pub on_unparseable: UnparseableToolInput,
}
#[derive(Debug, Clone)]
pub struct ToolCallDecoration {
pub tool_id: String,
pub signature: Option<String>,
pub additional_params: Option<serde_json::Value>,
}
impl ToolInputEnd {
pub fn new(id: impl Into<StreamPartId>, on_unparseable: UnparseableToolInput) -> Self {
Self {
id: id.into(),
tool_id: None,
name: None,
arguments: None,
call_id: None,
signature: None,
additional_params: None,
on_unparseable,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StreamFinalKind {
Final,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(from = "StreamFinalRepr")]
pub struct StreamFinal {
pub kind: StreamFinalKind,
pub usage: Usage,
#[serde(default)]
pub finish_reason: Option<crate::completion::FinishReason>,
#[serde(default)]
pub message_id: Option<String>,
#[serde(default)]
pub response_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_request_id: Option<String>,
pub provider: String,
#[serde(default)]
pub model: Option<String>,
#[serde(default, skip_serializing_if = "serde_json::Value::is_null")]
pub raw: serde_json::Value,
}
impl StreamFinal {
pub fn new(provider: impl Into<String>, usage: Usage) -> Self {
Self {
kind: StreamFinalKind::Final,
usage,
finish_reason: None,
message_id: None,
response_id: None,
provider_request_id: None,
provider: provider.into(),
model: None,
raw: serde_json::Value::Null,
}
}
pub fn with_finish_reason(self, finish_reason: crate::completion::FinishReason) -> Self {
self.with_optional_finish_reason(Some(finish_reason))
}
pub fn with_optional_finish_reason(
mut self,
finish_reason: Option<crate::completion::FinishReason>,
) -> Self {
self.finish_reason = finish_reason;
self
}
pub fn identity(&self) -> crate::completion::ResponseIdentity {
crate::completion::ResponseIdentity {
message_id: self.message_id.clone(),
response_id: self.response_id.clone(),
provider_request_id: self.provider_request_id.clone(),
}
}
}
crate::provider_response::response_metadata_setters!(StreamFinal);
#[derive(Deserialize)]
struct StreamFinalRepr {
kind: StreamFinalKind,
usage: Usage,
#[serde(default)]
finish_reason: Option<crate::completion::FinishReason>,
#[serde(default)]
message_id: Option<String>,
#[serde(default)]
response_id: Option<String>,
#[serde(default)]
provider_request_id: Option<String>,
provider: String,
#[serde(default)]
model: Option<String>,
#[serde(default)]
raw: serde_json::Value,
}
impl From<StreamFinalRepr> for StreamFinal {
fn from(repr: StreamFinalRepr) -> Self {
let StreamFinalRepr {
kind,
usage,
finish_reason,
message_id,
response_id,
provider_request_id,
provider,
model,
raw,
} = repr;
let StreamFinalKind::Final = kind;
Self::new(provider, usage)
.with_optional_finish_reason(finish_reason)
.with_optional_message_id(message_id)
.with_optional_response_id(response_id)
.with_optional_provider_request_id(provider_request_id)
.with_optional_model(model)
.with_raw(raw)
}
}
#[derive(Clone, PartialEq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct UnknownPayload(serde_json::Value);
impl UnknownPayload {
pub fn new(value: serde_json::Value) -> Self {
Self(value)
}
pub fn value(&self) -> &serde_json::Value {
&self.0
}
}
impl std::fmt::Debug for UnknownPayload {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let bytes = serde_json::to_vec(&self.0)
.map(|json| json.len())
.unwrap_or(0);
write!(f, "UnknownPayload({bytes} bytes redacted)")
}
}
impl From<serde_json::Value> for UnknownPayload {
fn from(value: serde_json::Value) -> Self {
Self(value)
}
}
#[cfg(test)]
mod unknown_payload_tests {
use super::UnknownPayload;
#[test]
fn debug_output_never_contains_payload_content() {
let payload = UnknownPayload::new(serde_json::json!({
"secret_field": "SENSITIVE-CONTENT",
}));
let rendered = format!("{payload:?}");
assert!(!rendered.contains("SENSITIVE-CONTENT"));
assert!(!rendered.contains("secret_field"));
assert!(rendered.contains("redacted"));
}
#[test]
fn serde_round_trip_is_transparent() {
let value = serde_json::json!({"type": "future_event", "n": 1});
let payload = UnknownPayload::new(value.clone());
let encoded = serde_json::to_string(&payload).expect("serializes");
assert_eq!(encoded, serde_json::to_string(&value).expect("serializes"));
let decoded: UnknownPayload = serde_json::from_str(&encoded).expect("deserializes");
assert_eq!(decoded, payload);
}
}
#[derive(Debug, Clone)]
pub enum RawStreamingChoice<R = StreamFinal> {
Message(String),
TextStart {
id: StreamPartId,
additional_params: Option<crate::message::AdditionalParams>,
},
TextAdditionalParams(crate::message::AdditionalParams),
ToolCall(RawStreamingToolCall),
ToolCallDelta {
id: StreamPartId,
content: ToolCallDeltaContent,
},
ToolInputEnd(ToolInputEnd),
Reasoning {
id: StreamPartId,
provider_id: Option<WireId>,
content: ReasoningContent,
},
ReasoningStart {
id: StreamPartId,
provider_id: Option<WireId>,
},
ReasoningEnd {
id: StreamPartId,
reasoning: Option<Reasoning>,
signature: Option<String>,
wire_sent: bool,
},
TextEnd {
id: StreamPartId,
},
ReasoningDelta {
id: StreamPartId,
provider_id: Option<WireId>,
reasoning: String,
},
FinalResponse(R),
MessageId(String),
Unknown(UnknownPayload),
}
impl<R> RawStreamingChoice<R> {
pub fn try_map_final<S>(
self,
map: impl FnOnce(R) -> Result<S, CompletionError>,
) -> Result<RawStreamingChoice<S>, CompletionError> {
Ok(match self {
Self::Message(text) => RawStreamingChoice::Message(text),
Self::TextStart {
id,
additional_params,
} => RawStreamingChoice::TextStart {
id,
additional_params,
},
Self::TextAdditionalParams(params) => RawStreamingChoice::TextAdditionalParams(params),
Self::ToolCall(call) => RawStreamingChoice::ToolCall(call),
Self::ToolCallDelta { id, content } => {
RawStreamingChoice::ToolCallDelta { id, content }
}
Self::ToolInputEnd(end) => RawStreamingChoice::ToolInputEnd(end),
Self::Reasoning {
id,
provider_id,
content,
} => RawStreamingChoice::Reasoning {
id,
provider_id,
content,
},
Self::ReasoningDelta {
id,
provider_id,
reasoning,
} => RawStreamingChoice::ReasoningDelta {
id,
provider_id,
reasoning,
},
Self::ReasoningStart { id, provider_id } => {
RawStreamingChoice::ReasoningStart { id, provider_id }
}
Self::ReasoningEnd {
id,
reasoning,
signature,
wire_sent,
} => RawStreamingChoice::ReasoningEnd {
id,
reasoning,
signature,
wire_sent,
},
Self::TextEnd { id } => RawStreamingChoice::TextEnd { id },
Self::FinalResponse(response) => RawStreamingChoice::FinalResponse(map(response)?),
Self::MessageId(id) => RawStreamingChoice::MessageId(id),
Self::Unknown(value) => RawStreamingChoice::Unknown(value),
})
}
}
#[derive(Debug, Clone)]
pub struct RawStreamingToolCall {
pub id: StreamPartId,
pub tool_id: Option<WireId>,
pub internal_call_id: String,
pub call_id: Option<String>,
pub name: String,
pub arguments: serde_json::Value,
pub signature: Option<String>,
pub additional_params: Option<serde_json::Value>,
}
impl RawStreamingToolCall {
pub fn empty() -> Self {
Self {
id: StreamPartId::minted(MintKind::Tool, u64::MAX),
tool_id: None,
internal_call_id: crate::id::generate(),
call_id: None,
name: String::new(),
arguments: serde_json::Value::Null,
signature: None,
additional_params: None,
}
}
pub fn new(id: impl Into<StreamPartId>, name: String, arguments: serde_json::Value) -> Self {
let id = id.into();
let tool_id = id.wire_str().and_then(WireId::new);
Self {
id,
tool_id,
internal_call_id: crate::id::generate(),
call_id: None,
name,
arguments,
signature: None,
additional_params: None,
}
}
pub fn with_call_id(mut self, call_id: String) -> Self {
self.call_id = Some(call_id);
self
}
pub fn with_signature(mut self, signature: Option<String>) -> Self {
self.signature = signature;
self
}
pub fn with_additional_params(mut self, additional_params: Option<serde_json::Value>) -> Self {
self.additional_params = additional_params;
self
}
}
impl From<RawStreamingToolCall> for ToolCall {
fn from(tool_call: RawStreamingToolCall) -> Self {
let provider = crate::message::ProviderCallId::from_optional_wire(
tool_call.call_id,
tool_call.tool_id.map(WireId::into_string),
);
let id = crate::message::ToolCallId::for_provider(provider.as_ref());
ToolCall {
id,
provider,
function: ToolFunction {
name: tool_call.name,
arguments: tool_call.arguments,
},
signature: tool_call.signature,
additional_params: tool_call.additional_params,
}
}
}
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
pub type RawStreamingResult<R> =
Pin<Box<dyn Stream<Item = Result<RawStreamingChoice<R>, CompletionError>> + Send>>;
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
pub type RawStreamingResult<R> =
Pin<Box<dyn Stream<Item = Result<RawStreamingChoice<R>, CompletionError>>>>;
pub type StreamingResult = RawStreamingResult<StreamFinal>;
pub fn normalize_stream<R, F>(stream: RawStreamingResult<R>, mut map: F) -> StreamingResult
where
R: Serialize + 'static,
F: FnMut(R) -> Result<StreamFinal, CompletionError> + WasmCompatSend + 'static,
{
let mut emitted_tool_call = false;
Box::pin(stream.map(move |item| {
item.and_then(|choice| {
if matches!(&choice, RawStreamingChoice::ToolCall(_)) {
emitted_tool_call = true;
}
choice.try_map_final(|response| {
let raw = serde_json::to_value(&response)?;
let mut response = map(response)?.with_raw(raw);
response.finish_reason = response
.finish_reason
.map(|reason| reason.reconcile_with_output(emitted_tool_call));
Ok(response)
})
})
}))
}
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
type ResumeWait = Pin<Box<dyn Future<Output = ()> + Send>>;
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
type ResumeWait = Pin<Box<dyn Future<Output = ()>>>;
pub struct StreamingCompletionResponse {
pub(crate) inner: Abortable<StreamingResult>,
pub(crate) abort_handle: AbortHandle,
pub(crate) pause_control: PauseControl,
parts: PartsAccumulator,
provider: String,
pub choice: Vec<AssistantContent>,
finished: bool,
resume_wait: Option<ResumeWait>,
reasoning_correlators: std::collections::HashMap<StreamPartId, String>,
finished_reasoning_correlators: std::collections::HashMap<StreamPartId, String>,
pub response: Option<StreamFinal>,
pub final_response_yielded: AtomicBool,
pub message_id: Option<String>,
}
impl StreamingCompletionResponse {
pub fn stream(provider: impl Into<String>, inner: StreamingResult) -> Self {
let (abort_handle, abort_registration) = AbortHandle::new_pair();
let abortable_stream = Abortable::new(inner, abort_registration);
let pause_control = PauseControl::new();
Self {
inner: abortable_stream,
abort_handle,
pause_control,
parts: PartsAccumulator::new(),
provider: provider.into(),
choice: Vec::new(),
finished: false,
resume_wait: None,
reasoning_correlators: std::collections::HashMap::new(),
finished_reasoning_correlators: std::collections::HashMap::new(),
response: None,
final_response_yielded: AtomicBool::new(false),
message_id: None,
}
}
pub fn provider(&self) -> &str {
&self.provider
}
fn reasoning_end_correlator(&mut self, id: StreamPartId, restated: bool) -> String {
match self.reasoning_correlators.remove(&id) {
Some(taken) => {
self.finished_reasoning_correlators
.insert(id, taken.clone());
taken
}
None if restated => {
let minted = crate::id::generate();
self.finished_reasoning_correlators
.insert(id, minted.clone());
minted
}
None => self
.finished_reasoning_correlators
.entry(id)
.or_insert_with(crate::id::generate)
.clone(),
}
}
pub fn cancel(&mut self) {
self.abort_handle.abort();
let (abort_handle, abort_registration) = AbortHandle::new_pair();
let empty: StreamingResult = Box::pin(futures::stream::poll_fn(|_| Poll::Ready(None)));
self.inner = Abortable::new(empty, abort_registration);
self.abort_handle = abort_handle;
self.pause_control.resume();
}
pub fn pause(&self) {
self.pause_control.pause();
}
pub fn resume(&self) {
self.pause_control.resume();
}
pub fn is_paused(&self) -> bool {
self.pause_control.is_paused()
}
pub fn usage(&self) -> Usage {
self.response
.as_ref()
.map(|response| response.usage)
.unwrap_or_default()
}
pub fn identity(&self) -> crate::completion::ResponseIdentity {
crate::completion::ResponseIdentity {
message_id: self.message_id.clone(),
..self
.response
.as_ref()
.map(StreamFinal::identity)
.unwrap_or_default()
}
}
}
impl From<StreamingCompletionResponse> for CompletionResponse {
fn from(value: StreamingCompletionResponse) -> CompletionResponse {
let terminal = value.response.as_ref();
CompletionResponse::new(
value.choice,
terminal.map(|response| response.usage).unwrap_or_default(),
value.provider,
)
.with_optional_message_id(
value
.message_id
.or_else(|| terminal.and_then(|response| response.message_id.clone())),
)
.with_optional_response_id(terminal.and_then(|response| response.response_id.clone()))
.with_optional_provider_request_id(
terminal.and_then(|response| response.provider_request_id.clone()),
)
.with_optional_finish_reason(terminal.and_then(|response| response.finish_reason.clone()))
.with_optional_model(terminal.and_then(|response| response.model.clone()))
}
}
impl Stream for StreamingCompletionResponse {
type Item = Result<StreamedAssistantContent, CompletionError>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let stream = self.get_mut();
if stream.finished {
return Poll::Ready(None);
}
if stream.is_paused() {
let wait = match stream.resume_wait.as_mut() {
Some(wait) => wait,
None => {
let mut paused_rx = stream.pause_control.paused_rx.clone();
stream.resume_wait.insert(Box::pin(async move {
let _ = paused_rx.wait_for(|paused| !*paused).await;
}))
}
};
if wait.as_mut().poll(cx).is_pending() {
return Poll::Pending;
}
stream.resume_wait = None;
}
loop {
return match Pin::new(&mut stream.inner).poll_next(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(None) => {
let finished = stream.parts.finish();
if !finished.is_empty() {
stream.choice = finished;
}
stream.finished = true;
Poll::Ready(None)
}
Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(err))),
Poll::Ready(Some(Ok(choice))) => match choice {
RawStreamingChoice::Message(text) => {
stream.parts.text_delta(&text);
Poll::Ready(Some(Ok(StreamedAssistantContent::text(&text))))
}
RawStreamingChoice::TextStart {
id,
additional_params,
} => {
stream.parts.text_start(&id, additional_params);
continue;
}
RawStreamingChoice::TextAdditionalParams(additional_params) => {
stream.parts.text_additional_params(additional_params);
continue;
}
RawStreamingChoice::ToolCallDelta { id, content } => {
let internal_call_id = match &content {
ToolCallDeltaContent::Name(name) => {
stream.parts.tool_name_delta(&id, name)
}
ToolCallDeltaContent::Delta(fragment) => {
stream.parts.tool_args_delta(&id, fragment)
}
};
Poll::Ready(Some(Ok(StreamedAssistantContent::ToolCallDelta {
internal_call_id,
content,
})))
}
RawStreamingChoice::ToolInputEnd(end) => match stream.parts.tool_input_end(end)
{
Ok(Some((tool_call, internal_call_id))) => {
Poll::Ready(Some(Ok(StreamedAssistantContent::ToolCall {
tool_call,
internal_call_id,
})))
}
Ok(None) => continue,
Err(err) => Poll::Ready(Some(Err(err))),
},
RawStreamingChoice::Reasoning {
id,
provider_id,
content,
} => {
let restatement = Reasoning {
id: provider_id.map(WireId::into_string),
content: vec![content],
};
let completed = stream.parts.reasoning_end(&id, Some(restatement), None);
let correlator = stream.reasoning_end_correlator(id, true);
match completed {
Some(completed) => {
Poll::Ready(Some(Ok(StreamedAssistantContent::Reasoning {
reasoning: completed,
id: correlator,
})))
}
None => continue,
}
}
RawStreamingChoice::ReasoningStart { id, provider_id } => {
if stream.parts.reasoning_start(&id, provider_id.as_ref()) {
stream
.reasoning_correlators
.insert(id, crate::id::generate());
}
continue;
}
RawStreamingChoice::ReasoningEnd {
id,
reasoning,
signature,
wire_sent,
} => {
let authoritative = reasoning.is_some() || signature.is_some() || wire_sent;
let restated = reasoning.is_some();
let completed = stream.parts.reasoning_end(&id, reasoning, signature);
let correlator = stream.reasoning_end_correlator(id, restated);
match completed {
Some(completed) if authoritative => {
Poll::Ready(Some(Ok(StreamedAssistantContent::Reasoning {
reasoning: completed,
id: correlator,
})))
}
_ => continue,
}
}
RawStreamingChoice::TextEnd { id } => {
stream.parts.text_end(&id);
continue;
}
RawStreamingChoice::ReasoningDelta {
id,
provider_id,
reasoning,
} => {
stream
.parts
.reasoning_delta(&id, provider_id.as_ref(), &reasoning);
let correlator = stream
.reasoning_correlators
.entry(id)
.or_insert_with(crate::id::generate)
.clone();
Poll::Ready(Some(Ok(StreamedAssistantContent::ReasoningDelta {
id: correlator,
provider_id: provider_id.map(WireId::into_string),
reasoning,
})))
}
RawStreamingChoice::ToolCall(raw_tool_call) => {
let minted_internal_call_id = raw_tool_call.internal_call_id.clone();
let part_id = raw_tool_call.id.clone();
let tool_call: ToolCall = raw_tool_call.into();
let internal_call_id = stream.parts.tool_call(
&part_id,
tool_call.clone(),
minted_internal_call_id,
);
Poll::Ready(Some(Ok(StreamedAssistantContent::ToolCall {
tool_call,
internal_call_id,
})))
}
RawStreamingChoice::FinalResponse(mut response) => {
response.finish_reason = response.finish_reason.map(|reason| {
reason.reconcile_with_output(stream.parts.saw_tool_call())
});
if stream
.final_response_yielded
.load(std::sync::atomic::Ordering::SeqCst)
{
continue;
} else {
if stream.message_id.is_none() {
stream.message_id = response.message_id.clone();
}
stream.response = Some(response.clone());
stream
.final_response_yielded
.store(true, std::sync::atomic::Ordering::SeqCst);
let final_response = StreamedAssistantContent::final_response(response);
Poll::Ready(Some(Ok(final_response)))
}
}
RawStreamingChoice::MessageId(id) => {
stream.message_id = Some(id);
continue;
}
RawStreamingChoice::Unknown(value) => {
Poll::Ready(Some(Ok(StreamedAssistantContent::Unknown(value))))
}
},
};
}
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::*;
use crate::completion::FinishReason;
use async_stream::stream;
use tokio::time::sleep;
const TEST_PROVIDER: &str = "test-provider";
fn fixture_params(value: serde_json::Value) -> crate::message::AdditionalParams {
crate::message::AdditionalParams::try_from_value(value)
.expect("fixture params must be a JSON object")
.expect("fixture params must carry data")
}
fn mock_final_with_total_tokens(total_tokens: u64) -> StreamFinal {
let mut usage = Usage::new();
usage.total_tokens = total_tokens;
StreamFinal::new(TEST_PROVIDER, usage)
}
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
fn to_stream_result(
stream: impl futures::Stream<Item = Result<RawStreamingChoice, CompletionError>>
+ Send
+ 'static,
) -> StreamingResult {
Box::pin(stream)
}
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
fn to_stream_result(
stream: impl futures::Stream<Item = Result<RawStreamingChoice, CompletionError>> + 'static,
) -> StreamingResult {
Box::pin(stream)
}
fn create_mock_stream() -> StreamingCompletionResponse {
let stream = stream! {
yield Ok(RawStreamingChoice::Message("hello 1".to_string()));
sleep(Duration::from_millis(100)).await;
yield Ok(RawStreamingChoice::Message("hello 2".to_string()));
sleep(Duration::from_millis(100)).await;
yield Ok(RawStreamingChoice::Message("hello 3".to_string()));
sleep(Duration::from_millis(100)).await;
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(15)));
};
StreamingCompletionResponse::stream(TEST_PROVIDER, to_stream_result(stream))
}
#[tokio::test]
async fn a_long_run_of_non_yielding_events_does_not_grow_the_stack() {
let raw = stream! {
for n in 0..50_000u32 {
yield Ok(RawStreamingChoice::MessageId(format!("msg_{n}")));
}
yield Ok(RawStreamingChoice::Message("done".to_string()));
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(1)));
};
let mut stream = StreamingCompletionResponse::stream(TEST_PROVIDER, to_stream_result(raw));
let mut texts = Vec::new();
while let Some(item) = stream.next().await {
if let Ok(StreamedAssistantContent::Text(text)) = item {
texts.push(text.text);
}
}
assert_eq!(texts, vec!["done".to_string()]);
assert_eq!(stream.message_id.as_deref(), Some("msg_49999"));
}
#[tokio::test]
async fn stream_identity_falls_back_to_the_terminal_records_ids() {
let raw = stream! {
yield Ok(RawStreamingChoice::Message("done".to_string()));
yield Ok(RawStreamingChoice::FinalResponse(
mock_final_with_total_tokens(1)
.with_message_id("msg_terminal")
.with_response_id("resp_1")
.with_provider_request_id("req_1"),
));
};
let mut stream = StreamingCompletionResponse::stream(TEST_PROVIDER, to_stream_result(raw));
while stream.next().await.is_some() {}
assert_eq!(
stream.identity(),
crate::completion::ResponseIdentity {
message_id: Some("msg_terminal".to_string()),
response_id: Some("resp_1".to_string()),
provider_request_id: Some("req_1".to_string()),
}
);
}
#[tokio::test]
async fn stream_identity_prefers_an_explicit_message_id_event() {
let raw = stream! {
yield Ok(RawStreamingChoice::MessageId("msg_event".to_string()));
yield Ok(RawStreamingChoice::Message("done".to_string()));
yield Ok(RawStreamingChoice::FinalResponse(
mock_final_with_total_tokens(1)
.with_message_id("msg_terminal")
.with_response_id("resp_1"),
));
};
let mut stream = StreamingCompletionResponse::stream(TEST_PROVIDER, to_stream_result(raw));
while stream.next().await.is_some() {}
assert_eq!(
stream.identity(),
crate::completion::ResponseIdentity {
message_id: Some("msg_event".to_string()),
response_id: Some("resp_1".to_string()),
provider_request_id: None,
}
);
}
fn create_reasoning_stream() -> StreamingCompletionResponse {
let stream = stream! {
yield Ok(RawStreamingChoice::Reasoning { id: StreamPartId::wire("rs_1"),
provider_id: WireId::new("rs_1"),
content: ReasoningContent::Text {
text: "step one".to_string(),
signature: Some("sig_1".to_string()),
},
});
yield Ok(RawStreamingChoice::Message("final answer".to_string()));
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(5)));
};
StreamingCompletionResponse::stream(TEST_PROVIDER, to_stream_result(stream))
}
fn create_reasoning_only_stream() -> StreamingCompletionResponse {
let stream = stream! {
yield Ok(RawStreamingChoice::Reasoning { id: StreamPartId::wire("rs_only"),
provider_id: WireId::new("rs_only"),
content: ReasoningContent::Summary("hidden summary".to_string()),
});
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(2)));
};
StreamingCompletionResponse::stream(TEST_PROVIDER, to_stream_result(stream))
}
fn create_interleaved_stream() -> StreamingCompletionResponse {
let stream = stream! {
yield Ok(RawStreamingChoice::Reasoning { id: StreamPartId::wire("rs_interleaved"),
provider_id: WireId::new("rs_interleaved"),
content: ReasoningContent::Text {
text: "chain-of-thought".to_string(),
signature: None,
},
});
yield Ok(RawStreamingChoice::Message("final-text".to_string()));
yield Ok(RawStreamingChoice::ToolCall(
RawStreamingToolCall::new(
"tool_1".to_string(),
"mock_tool".to_string(),
serde_json::json!({"arg": 1}),
),
));
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(3)));
};
StreamingCompletionResponse::stream(TEST_PROVIDER, to_stream_result(stream))
}
fn create_text_tool_text_stream() -> StreamingCompletionResponse {
let stream = stream! {
yield Ok(RawStreamingChoice::Message("first".to_string()));
yield Ok(RawStreamingChoice::ToolCall(
RawStreamingToolCall::new(
"tool_split".to_string(),
"mock_tool".to_string(),
serde_json::json!({"arg": "x"}),
),
));
yield Ok(RawStreamingChoice::Message("second".to_string()));
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(3)));
};
StreamingCompletionResponse::stream(TEST_PROVIDER, to_stream_result(stream))
}
fn create_text_metadata_stream() -> StreamingCompletionResponse {
let stream = stream! {
yield Ok(RawStreamingChoice::TextStart {
id: StreamPartId::wire("block-0"),
additional_params: None,
});
yield Ok(RawStreamingChoice::Message("first".to_string()));
yield Ok(RawStreamingChoice::TextAdditionalParams(fixture_params(serde_json::json!({
"citations": [{
"type": "char_location",
"cited_text": "First citation.",
"document_index": 0,
"start_char_index": 0,
"end_char_index": 15
}]
}))));
yield Ok(RawStreamingChoice::TextAdditionalParams(fixture_params(serde_json::json!({
"citations": [{
"type": "char_location",
"cited_text": "Second citation.",
"document_index": 0,
"start_char_index": 16,
"end_char_index": 32
}]
}))));
yield Ok(RawStreamingChoice::TextStart {
id: StreamPartId::wire("block-1"),
additional_params: crate::message::AdditionalParams::try_from_value(serde_json::json!({
"block": 2
})).expect("object params"),
});
yield Ok(RawStreamingChoice::Message("second".to_string()));
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(3)));
};
StreamingCompletionResponse::stream(TEST_PROVIDER, to_stream_result(stream))
}
#[tokio::test]
async fn into_completion_response_derives_usage_from_final_response() {
let mut stream = create_mock_stream();
while stream.next().await.is_some() {}
assert_eq!(stream.usage().total_tokens, 15);
let response: CompletionResponse = stream.into();
assert_eq!(response.usage.total_tokens, 15);
assert_eq!(response.provider, TEST_PROVIDER);
}
#[tokio::test]
async fn into_completion_response_carries_the_terminal_request_id() {
let mut stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::Message("hi".to_string()));
yield Ok(RawStreamingChoice::FinalResponse(
StreamFinal::new(TEST_PROVIDER, Usage::new())
.with_response_id("resp_1")
.with_provider_request_id("req_transport_1"),
));
}),
);
while stream.next().await.is_some() {}
let response: CompletionResponse = stream.into();
assert_eq!(response.response_id.as_deref(), Some("resp_1"));
assert_eq!(
response.provider_request_id.as_deref(),
Some("req_transport_1")
);
}
#[tokio::test]
async fn a_stream_without_a_terminal_record_still_names_its_provider() {
let mut stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::Message("truncated".to_string()));
}),
);
while stream.next().await.is_some() {}
assert!(stream.response.is_none());
let response: CompletionResponse = stream.into();
assert_eq!(response.provider, TEST_PROVIDER);
assert_eq!(response.usage, Usage::new());
assert_eq!(response.finish_reason(), None);
assert_eq!(response.model, None);
}
#[tokio::test]
async fn a_stream_that_errors_mid_stream_keeps_content_and_omits_the_terminal() {
let mut stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::Message("partial".to_string()));
yield Err(CompletionError::ProviderError(
"connection reset".to_string(),
));
}),
);
let mut saw_error = false;
while let Some(item) = stream.next().await {
if item.is_err() {
saw_error = true;
}
}
assert!(saw_error, "the mid-stream error must be forwarded");
assert!(stream.response.is_none());
assert_eq!(
stream.choice.first(),
Some(&AssistantContent::text("partial".to_string())),
);
}
#[tokio::test]
async fn normalize_stream_upgrades_a_stop_that_carried_a_tool_call() {
let raw: RawStreamingResult<Usage> = Box::pin(stream! {
yield Ok(RawStreamingChoice::ToolCall(RawStreamingToolCall {
tool_id: WireId::new("call_1"),
id: StreamPartId::wire("call_1"),
call_id: None,
internal_call_id: "internal_1".to_string(),
name: "lookup".to_string(),
arguments: serde_json::json!({}),
signature: None,
additional_params: None,
}));
yield Ok(RawStreamingChoice::FinalResponse(Usage::new()));
});
let normalized = normalize_stream(raw, |usage| {
Ok(StreamFinal::new(TEST_PROVIDER, usage).with_finish_reason(FinishReason::Stop))
});
let mut stream = StreamingCompletionResponse::stream(TEST_PROVIDER, normalized);
while stream.next().await.is_some() {}
assert_eq!(
stream
.response
.as_ref()
.and_then(|final_record| final_record.finish_reason.clone()),
Some(FinishReason::ToolCalls),
);
}
#[tokio::test]
async fn normalize_stream_leaves_a_stop_without_tool_calls_alone() {
let raw: RawStreamingResult<Usage> = Box::pin(stream! {
yield Ok(RawStreamingChoice::Message("done".to_string()));
yield Ok(RawStreamingChoice::FinalResponse(Usage::new()));
});
let normalized = normalize_stream(raw, |usage| {
Ok(StreamFinal::new(TEST_PROVIDER, usage).with_finish_reason(FinishReason::Stop))
});
let mut stream = StreamingCompletionResponse::stream(TEST_PROVIDER, normalized);
while stream.next().await.is_some() {}
assert_eq!(
stream
.response
.as_ref()
.and_then(|final_record| final_record.finish_reason.clone()),
Some(FinishReason::Stop),
);
}
#[test]
fn stream_final_round_trips_and_is_distinguishable_from_unknown_content() {
let final_record = StreamFinal::new(
"example",
Usage {
input_tokens: 4,
output_tokens: 6,
total_tokens: 10,
cached_input_tokens: 1,
cache_creation_input_tokens: 2,
tool_use_prompt_tokens: 3,
reasoning_tokens: 4,
},
)
.with_finish_reason(FinishReason::Other("future_reason".to_owned()))
.with_message_id("msg_123")
.with_model("provider-model-v2");
let encoded = serde_json::to_value(StreamedAssistantContent::Final(final_record.clone()))
.expect("serialize final item");
assert_eq!(encoded["kind"], serde_json::json!("final"));
let decoded = serde_json::from_value::<StreamedAssistantContent>(encoded)
.expect("deserialize final item");
assert_eq!(decoded, StreamedAssistantContent::Final(final_record));
let provider_item = serde_json::json!({
"provider_native_event": "future_terminal",
"usage": {"total_tokens": 10}
});
let decoded = serde_json::from_value::<StreamedAssistantContent>(provider_item.clone())
.expect("deserialize unknown item");
assert_eq!(
decoded,
StreamedAssistantContent::Unknown(provider_item.into())
);
}
#[test]
fn deserializing_stream_final_filters_empty_identifiers() {
let decoded = serde_json::from_value::<StreamFinal>(serde_json::json!({
"kind": "final",
"usage": Usage::new(),
"message_id": "",
"response_id": "",
"model": "",
"provider": "example",
}))
.expect("deserialize terminal record");
assert_eq!(decoded.message_id, None);
assert_eq!(decoded.response_id, None);
assert_eq!(decoded.model, None);
}
#[derive(Debug, Serialize, Deserialize, PartialEq)]
struct ProviderTerminal {
usage: Usage,
provider_only: String,
}
fn provider_terminal_stream() -> RawStreamingResult<ProviderTerminal> {
Box::pin(stream! {
yield Ok(RawStreamingChoice::Message("done".to_string()));
yield Ok(RawStreamingChoice::FinalResponse(ProviderTerminal {
usage: Usage {
input_tokens: 3,
output_tokens: 5,
total_tokens: 8,
..Usage::new()
},
provider_only: "kept".to_string(),
}));
})
}
async fn drain(normalized: StreamingResult) -> StreamFinal {
let mut stream = StreamingCompletionResponse::stream(TEST_PROVIDER, normalized);
while stream.next().await.is_some() {}
stream
.response
.expect("stream should end with a terminal record")
}
#[tokio::test]
async fn normalize_stream_captures_the_terminal_record() {
let normalized = normalize_stream(provider_terminal_stream(), |terminal| {
Ok(StreamFinal::new(TEST_PROVIDER, terminal.usage))
});
let final_record = drain(normalized).await;
let raw = &final_record.raw;
let typed = ProviderTerminal::deserialize(raw).expect("raw is the provider's terminal");
assert_eq!(typed.provider_only, "kept");
assert_eq!(&serde_json::to_value(&typed).expect("re-serialize"), raw);
assert_eq!(final_record.usage.total_tokens, 8);
assert_eq!(final_record.provider, TEST_PROVIDER);
assert_eq!(final_record.finish_reason, None);
}
#[tokio::test]
async fn normalize_stream_reconciles_finish_reason_with_raw_attached() {
let raw: RawStreamingResult<Usage> = Box::pin(stream! {
yield Ok(RawStreamingChoice::ToolCall(RawStreamingToolCall {
tool_id: WireId::new("call_1"),
id: StreamPartId::wire("call_1"),
call_id: None,
internal_call_id: "internal_1".to_string(),
name: "lookup".to_string(),
arguments: serde_json::json!({}),
signature: None,
additional_params: None,
}));
yield Ok(RawStreamingChoice::FinalResponse(Usage::new()));
});
let normalized = normalize_stream(raw, |usage| {
Ok(StreamFinal::new(TEST_PROVIDER, usage).with_finish_reason(FinishReason::Stop))
});
let final_record = drain(normalized).await;
assert_eq!(final_record.finish_reason, Some(FinishReason::ToolCalls));
assert!(!final_record.raw.is_null());
}
#[test]
fn stream_final_raw_round_trips_through_serde_mirror() {
let payload = serde_json::json!({
"usage": {"total_tokens": 8},
"provider_only": "kept"
});
let final_record = StreamFinal::new("example", Usage::new())
.with_message_id("msg_123")
.with_raw(payload.clone());
let encoded = serde_json::to_value(&final_record).expect("serialize");
assert_eq!(encoded["raw"], payload);
let decoded = serde_json::from_value::<StreamFinal>(encoded.clone()).expect("deserialize");
assert_eq!(decoded.raw, payload);
assert_eq!(decoded, final_record);
assert_eq!(
serde_json::to_value(&decoded).expect("re-serialize"),
encoded
);
let wrapped = StreamedAssistantContent::Final(final_record.clone());
let encoded = serde_json::to_value(&wrapped).expect("serialize wrapped");
let decoded = serde_json::from_value::<StreamedAssistantContent>(encoded)
.expect("deserialize wrapped");
assert_eq!(decoded, wrapped);
let legacy = serde_json::json!({
"kind": "final",
"usage": serde_json::to_value(Usage::new()).unwrap(),
"provider": "example"
});
let decoded = serde_json::from_value::<StreamFinal>(legacy).expect("legacy loads");
assert!(decoded.raw.is_null());
let bare = serde_json::to_value(StreamFinal::new("example", Usage::new())).unwrap();
assert!(bare.get("raw").is_none());
}
#[test]
fn stream_final_serde_round_trip_is_identity() {
let final_record = StreamFinal::new(
"example",
Usage {
input_tokens: 4,
output_tokens: 6,
total_tokens: 10,
cached_input_tokens: 1,
cache_creation_input_tokens: 2,
tool_use_prompt_tokens: 3,
reasoning_tokens: 4,
},
)
.with_finish_reason(FinishReason::Stop)
.with_message_id("msg_123")
.with_response_id("resp_456")
.with_model("provider-model-v2");
let encoded = serde_json::to_value(&final_record).expect("serialize terminal record");
assert_eq!(encoded["kind"], serde_json::json!("final"));
let decoded = serde_json::from_value::<StreamFinal>(encoded.clone()).expect("deserialize");
assert_eq!(decoded, final_record);
assert_eq!(
serde_json::to_value(&decoded).expect("re-serialize"),
encoded
);
}
#[tokio::test]
async fn usage_is_zero_sentinel_before_final_response() {
let stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::Message("no final response".to_string()));
}),
);
assert_eq!(stream.usage().total_tokens, 0);
}
#[tokio::test]
async fn test_stream_cancellation() {
let mut stream = create_mock_stream();
println!("Response: ");
let mut chunk_count = 0;
while let Some(chunk) = stream.next().await {
match chunk {
Ok(StreamedAssistantContent::Text(text)) => {
print!("{}", text.text);
std::io::Write::flush(&mut std::io::stdout()).unwrap();
chunk_count += 1;
}
Ok(StreamedAssistantContent::ToolCall {
tool_call,
internal_call_id,
}) => {
println!("\nTool Call: {tool_call:?}, internal_call_id={internal_call_id:?}");
chunk_count += 1;
}
Ok(StreamedAssistantContent::ToolCallDelta {
internal_call_id,
content,
}) => {
println!(
"\nTool Call delta: internal_call_id={internal_call_id:?}, content={content:?}"
);
chunk_count += 1;
}
Ok(StreamedAssistantContent::Final(res)) => {
println!("\nFinal response: {res:?}");
}
Ok(StreamedAssistantContent::Reasoning { reasoning, .. }) => {
let reasoning = reasoning.display_text();
print!("{reasoning}");
std::io::Write::flush(&mut std::io::stdout()).unwrap();
}
Ok(StreamedAssistantContent::ReasoningDelta { reasoning, .. }) => {
println!("Reasoning delta: {reasoning}");
chunk_count += 1;
}
Ok(StreamedAssistantContent::Unknown(value)) => {
println!("\nUnknown item: {value:?}");
chunk_count += 1;
}
Err(e) => {
eprintln!("Error: {e:?}");
break;
}
}
if chunk_count >= 2 {
println!("\nCancelling stream...");
stream.cancel();
println!("Stream cancelled.");
break;
}
}
let next_chunk = stream.next().await;
assert!(
next_chunk.is_none(),
"Expected no further chunks after cancellation, got {next_chunk:?}"
);
}
#[tokio::test]
async fn test_stream_pause_resume() {
let stream = create_mock_stream();
stream.pause();
assert!(stream.is_paused());
stream.resume();
assert!(!stream.is_paused());
}
#[tokio::test]
async fn a_paused_stream_parks_until_resume_instead_of_busy_waking() {
let stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::Message("hello".to_string()));
}),
);
let resume = stream.pause_control.paused_tx.clone();
stream.pause();
let mut task = tokio_test::task::spawn(stream);
assert!(
task.poll_next().is_pending(),
"a paused stream yields nothing"
);
assert!(
!task.is_woken(),
"a paused stream must idle, not re-wake itself"
);
resume.send(false).expect("resume");
assert!(task.is_woken(), "resuming must wake the parked stream");
assert!(matches!(
task.poll_next(),
Poll::Ready(Some(Ok(StreamedAssistantContent::Text(text)))) if text.text == "hello"
));
}
#[tokio::test]
async fn cancelling_a_paused_stream_terminates_instead_of_deadlocking() {
let mut stream = create_mock_stream();
stream.pause();
stream.cancel();
assert!(
!stream.is_paused(),
"cancel must lift the pause so the termination is observable"
);
assert!(
stream.next().await.is_none(),
"a cancelled stream terminates"
);
}
#[tokio::test]
async fn re_polling_a_drained_stream_preserves_the_aggregated_choice() {
let mut stream = create_mock_stream();
while stream.next().await.is_some() {}
let drained: Vec<AssistantContent> = stream.choice.clone().into_iter().collect();
assert_eq!(
drained,
vec![AssistantContent::text("hello 1hello 2hello 3")]
);
for _ in 0..3 {
assert!(
stream.next().await.is_none(),
"a drained stream stays drained"
);
}
assert_eq!(
stream.choice.clone().into_iter().collect::<Vec<_>>(),
drained,
"re-polling must not re-run the destructive finish()"
);
let response: CompletionResponse = stream.into();
assert_eq!(response.choice.into_iter().collect::<Vec<_>>(), drained);
}
#[tokio::test]
async fn a_provider_error_mentioning_aborted_reaches_the_consumer() {
let mut stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::Message("partial".to_string()));
yield Err(CompletionError::ProviderError(
"upstream aborted the request".to_string(),
));
}),
);
let mut errors = Vec::new();
while let Some(item) = stream.next().await {
if let Err(err) = item {
errors.push(err.to_string());
}
}
assert_eq!(errors.len(), 1, "the error must not be swallowed");
assert!(errors[0].contains("upstream aborted the request"));
assert_eq!(
stream.choice.first(),
Some(&AssistantContent::text("partial".to_string()))
);
assert!(stream.response.is_none());
}
#[tokio::test]
async fn a_full_tool_call_correlates_with_the_deltas_of_the_same_id() {
let mut stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::ToolCallDelta {
id: StreamPartId::wire("tc1"),
content: ToolCallDeltaContent::Name("add".to_string()),
});
yield Ok(RawStreamingChoice::ToolCallDelta {
id: StreamPartId::wire("tc1"),
content: ToolCallDeltaContent::Delta("{\"x\":1}".to_string()),
});
yield Ok(RawStreamingChoice::ToolCall(RawStreamingToolCall::new(
"tc1".to_string(),
"add".to_string(),
serde_json::json!({"x": 1}),
)));
yield Ok(RawStreamingChoice::ToolInputEnd(ToolInputEnd::new(
"tc1",
UnparseableToolInput::Drop,
)));
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(1)));
}),
);
let mut delta_ids = Vec::new();
let mut completed_ids = Vec::new();
while let Some(item) = stream.next().await {
match item.expect("stream item should be Ok") {
StreamedAssistantContent::ToolCallDelta {
internal_call_id, ..
} => delta_ids.push(internal_call_id),
StreamedAssistantContent::ToolCall {
internal_call_id, ..
} => completed_ids.push(internal_call_id),
_ => {}
}
}
assert_eq!(delta_ids.len(), 2);
assert_eq!(delta_ids[0], delta_ids[1], "one call, one internal id");
assert_eq!(
completed_ids,
vec![delta_ids[0].clone()],
"the completed call must carry the id its deltas published"
);
let tool_calls: Vec<&ToolCall> = stream
.choice
.iter()
.filter_map(|item| match item {
AssistantContent::ToolCall(tool_call) => Some(tool_call),
_ => None,
})
.collect();
assert_eq!(tool_calls.len(), 1, "got {:?}", stream.choice);
}
#[tokio::test]
async fn test_stream_aggregates_reasoning_content() {
let mut stream = create_reasoning_stream();
while stream.next().await.is_some() {}
let choice_items: Vec<AssistantContent> = stream.choice.clone().into_iter().collect();
assert!(choice_items.iter().any(|item| matches!(
item,
AssistantContent::Reasoning(Reasoning {
id: Some(id),
content
}) if id == "rs_1"
&& matches!(
content.first(),
Some(ReasoningContent::Text {
text,
signature: Some(signature)
}) if text == "step one" && signature == "sig_1"
)
)));
}
#[tokio::test]
async fn full_reasoning_block_supersedes_its_accumulated_deltas() {
let mut stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::ReasoningDelta {
id: StreamPartId::wire("rs_1"),
provider_id: WireId::new("rs_1"),
reasoning: "partial ".to_string(),
});
yield Ok(RawStreamingChoice::Reasoning { id: StreamPartId::wire("rs_1"),
provider_id: WireId::new("rs_1"),
content: ReasoningContent::Text {
text: "the complete chain".to_string(),
signature: Some("sig_1".to_string()),
},
});
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(2)));
}),
);
while stream.next().await.is_some() {}
let choice_items: Vec<AssistantContent> = stream.choice.clone().into_iter().collect();
let reasoning_items: Vec<&Reasoning> = choice_items
.iter()
.filter_map(|item| match item {
AssistantContent::Reasoning(reasoning) => Some(reasoning),
_ => None,
})
.collect();
assert_eq!(reasoning_items.len(), 1, "got {choice_items:?}");
let reasoning = reasoning_items.first().expect("one reasoning item");
assert_eq!(reasoning.id.as_deref(), Some("rs_1"));
assert!(matches!(
reasoning.content.first(),
Some(ReasoningContent::Text { text, signature: Some(signature) })
if text == "the complete chain" && signature == "sig_1"
));
}
#[tokio::test]
async fn full_reasoning_block_with_a_different_id_appends() {
let mut stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::ReasoningDelta {
id: StreamPartId::wire("rs_1"),
provider_id: WireId::new("rs_1"),
reasoning: "first item deltas".to_string(),
});
yield Ok(RawStreamingChoice::Reasoning { id: StreamPartId::wire("rs_2"),
provider_id: WireId::new("rs_2"),
content: ReasoningContent::Text {
text: "a different item".to_string(),
signature: None,
},
});
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(2)));
}),
);
while stream.next().await.is_some() {}
let choice_items: Vec<AssistantContent> = stream.choice.clone().into_iter().collect();
let reasoning_ids: Vec<Option<&str>> = choice_items
.iter()
.filter_map(|item| match item {
AssistantContent::Reasoning(reasoning) => Some(reasoning.id.as_deref()),
_ => None,
})
.collect();
assert_eq!(reasoning_ids, vec![Some("rs_1"), Some("rs_2")]);
}
#[tokio::test]
async fn wire_sent_bare_end_yields_the_completed_block_synthesized_stays_silent() {
let run = |wire_sent: bool| async move {
let mut stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::ReasoningDelta {
id: StreamPartId::minted(MintKind::Block, 0),
provider_id: None,
reasoning: "unsigned thoughts".to_string(),
});
yield Ok(RawStreamingChoice::ReasoningEnd {
id: StreamPartId::minted(MintKind::Block, 0),
reasoning: None,
signature: None,
wire_sent,
});
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(2)));
}),
);
let mut completed = Vec::new();
while let Some(item) = stream.next().await {
if let Ok(StreamedAssistantContent::Reasoning { reasoning, .. }) = item {
completed.push(reasoning);
}
}
completed
};
let wire = run(true).await;
assert_eq!(wire.len(), 1, "a wire-sent end announces the boundary");
assert!(matches!(
wire[0].content.first(),
Some(ReasoningContent::Text { text, signature: None }) if text == "unsigned thoughts"
));
let synthesized = run(false).await;
assert!(
synthesized.is_empty(),
"a synthesized bare end fabricates nothing: {synthesized:?}"
);
}
#[tokio::test]
async fn reused_key_after_end_mints_a_fresh_delta_correlator() {
let key = || StreamPartId::minted(MintKind::Reasoning, 0);
let mut stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::ReasoningDelta {
id: key(),
provider_id: None,
reasoning: "block A".to_string(),
});
yield Ok(RawStreamingChoice::ReasoningEnd {
id: key(),
reasoning: None,
signature: None,
wire_sent: false,
});
yield Ok(RawStreamingChoice::Message("interleaved".to_string()));
yield Ok(RawStreamingChoice::ReasoningDelta {
id: key(),
provider_id: None,
reasoning: "block B".to_string(),
});
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(2)));
}),
);
let mut delta_ids = Vec::new();
while let Some(item) = stream.next().await {
if let Ok(StreamedAssistantContent::ReasoningDelta { id, .. }) = item {
delta_ids.push(id);
}
}
assert_eq!(delta_ids.len(), 2, "one delta per block");
assert_ne!(
delta_ids[0], delta_ids[1],
"distinct parts must not share a correlator"
);
}
#[tokio::test]
async fn completed_reasoning_restates_the_delta_correlator() {
let key = || StreamPartId::minted(MintKind::Block, 0);
let mut stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::ReasoningDelta {
id: key(),
provider_id: None,
reasoning: "unsigned thoughts".to_string(),
});
yield Ok(RawStreamingChoice::ReasoningEnd {
id: key(),
reasoning: None,
signature: None,
wire_sent: true,
});
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(2)));
}),
);
let mut delta_ids = Vec::new();
let mut completed = Vec::new();
while let Some(item) = stream.next().await {
match item {
Ok(StreamedAssistantContent::ReasoningDelta { id, .. }) => delta_ids.push(id),
Ok(StreamedAssistantContent::Reasoning { reasoning, id }) => {
completed.push((reasoning, id));
}
_ => {}
}
}
let (reasoning, correlator) = completed.first().expect("one completed block");
assert_eq!(
Some(correlator),
delta_ids.first(),
"the completed block restates its deltas' correlator"
);
assert_eq!(
reasoning.id, None,
"no provider handle exists on this wire; the correlator must not leak into it"
);
}
#[tokio::test]
async fn completed_reasoning_keeps_correlator_and_provider_handle_distinct() {
let mut stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::ReasoningDelta {
id: StreamPartId::wire("rs_1"),
provider_id: WireId::new("rs_1"),
reasoning: "signed thoughts".to_string(),
});
yield Ok(RawStreamingChoice::ReasoningEnd {
id: StreamPartId::wire("rs_1"),
reasoning: None,
signature: Some("sig_1".to_string()),
wire_sent: true,
});
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(2)));
}),
);
let mut delta_ids = Vec::new();
let mut completed = Vec::new();
while let Some(item) = stream.next().await {
match item {
Ok(StreamedAssistantContent::ReasoningDelta { id, .. }) => delta_ids.push(id),
Ok(StreamedAssistantContent::Reasoning { reasoning, id }) => {
completed.push((reasoning, id));
}
_ => {}
}
}
let (reasoning, correlator) = completed.first().expect("one completed block");
assert_eq!(Some(correlator), delta_ids.first());
assert_eq!(reasoning.id.as_deref(), Some("rs_1"));
assert_ne!(
correlator.as_str(),
"rs_1",
"the rig correlator and the provider handle are separate values"
);
}
#[tokio::test]
async fn late_signature_after_synthesized_end_restates_the_delta_correlator() {
let key = || StreamPartId::minted(MintKind::Reasoning, 0);
let mut stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::ReasoningDelta {
id: key(),
provider_id: None,
reasoning: "hidden thoughts".to_string(),
});
yield Ok(RawStreamingChoice::ReasoningEnd {
id: key(),
reasoning: None,
signature: None,
wire_sent: false,
});
yield Ok(RawStreamingChoice::Message("visible".to_string()));
yield Ok(RawStreamingChoice::ReasoningEnd {
id: key(),
reasoning: None,
signature: Some("sig_late".to_string()),
wire_sent: true,
});
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(2)));
}),
);
let mut delta_ids = Vec::new();
let mut completed = Vec::new();
while let Some(item) = stream.next().await {
match item {
Ok(StreamedAssistantContent::ReasoningDelta { id, .. }) => delta_ids.push(id),
Ok(StreamedAssistantContent::Reasoning { reasoning, id }) => {
completed.push((reasoning, id));
}
_ => {}
}
}
assert_eq!(completed.len(), 1, "one signed completion, no duplicate");
let (reasoning, correlator) = completed.first().expect("one completed block");
assert_eq!(
Some(correlator),
delta_ids.first(),
"the signed completion restates the correlator its deltas carried"
);
assert!(
reasoning.content.iter().any(|content| matches!(
content,
ReasoningContent::Text { signature: Some(sig), .. } if sig == "sig_late"
)),
"the trailing signature landed on the completed part"
);
}
#[tokio::test]
async fn a_delta_less_start_under_a_reused_key_mints_a_fresh_correlator() {
let key = || StreamPartId::minted(MintKind::Reasoning, 0);
let mut stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::ReasoningDelta {
id: key(),
provider_id: None,
reasoning: "part one".to_string(),
});
yield Ok(RawStreamingChoice::ReasoningEnd {
id: key(),
reasoning: None,
signature: None,
wire_sent: true,
});
yield Ok(RawStreamingChoice::ReasoningStart {
id: key(),
provider_id: None,
});
yield Ok(RawStreamingChoice::ReasoningEnd {
id: key(),
reasoning: None,
signature: Some("sig2".to_string()),
wire_sent: true,
});
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(2)));
}),
);
let mut completed_ids = Vec::new();
while let Some(item) = stream.next().await {
if let Ok(StreamedAssistantContent::Reasoning { id, .. }) = item {
completed_ids.push(id);
}
}
assert_eq!(completed_ids.len(), 2, "two distinct parts complete");
assert_ne!(
completed_ids.first(),
completed_ids.get(1),
"distinct parts must not share a public correlator"
);
}
#[tokio::test]
async fn reused_accumulation_key_mints_a_fresh_correlator_after_an_end() {
let key = || StreamPartId::minted(MintKind::Reasoning, 0);
let mut stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::ReasoningDelta {
id: key(),
provider_id: None,
reasoning: "first part".to_string(),
});
yield Ok(RawStreamingChoice::ReasoningEnd {
id: key(),
reasoning: None,
signature: None,
wire_sent: true,
});
yield Ok(RawStreamingChoice::ReasoningDelta {
id: key(),
provider_id: None,
reasoning: "second part".to_string(),
});
yield Ok(RawStreamingChoice::ReasoningEnd {
id: key(),
reasoning: None,
signature: None,
wire_sent: true,
});
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(2)));
}),
);
let mut completed_ids = Vec::new();
while let Some(item) = stream.next().await {
if let Ok(StreamedAssistantContent::Reasoning { id, .. }) = item {
completed_ids.push(id);
}
}
assert_eq!(completed_ids.len(), 2, "two parts under the reused key");
assert_ne!(
completed_ids.first(),
completed_ids.get(1),
"a reused key opens a new part with a fresh correlator"
);
}
#[tokio::test]
async fn whole_block_reasoning_mints_a_unique_correlator() {
let mut stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::Reasoning {
id: StreamPartId::wire("rs_1"),
provider_id: WireId::new("rs_1"),
content: ReasoningContent::Text {
text: "first".to_string(),
signature: None,
},
});
yield Ok(RawStreamingChoice::Reasoning {
id: StreamPartId::wire("rs_2"),
provider_id: WireId::new("rs_2"),
content: ReasoningContent::Text {
text: "second".to_string(),
signature: None,
},
});
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(2)));
}),
);
let mut correlators = Vec::new();
while let Some(item) = stream.next().await {
if let Ok(StreamedAssistantContent::Reasoning { id, .. }) = item {
correlators.push(id);
}
}
assert_eq!(correlators.len(), 2);
assert!(correlators.iter().all(|id| !id.is_empty()));
assert_ne!(
correlators[0], correlators[1],
"distinct parts must not share a correlator"
);
}
#[tokio::test]
async fn full_reasoning_block_supersedes_deltas_across_interleaved_output() {
let mut stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::ReasoningDelta {
id: StreamPartId::wire("rs_1"),
provider_id: WireId::new("rs_1"),
reasoning: "partial ".to_string(),
});
yield Ok(RawStreamingChoice::ToolCall(RawStreamingToolCall::new(
"call_1".to_string(),
"probe".to_string(),
serde_json::json!({}),
)));
yield Ok(RawStreamingChoice::Reasoning { id: StreamPartId::wire("rs_1"),
provider_id: WireId::new("rs_1"),
content: ReasoningContent::Text {
text: "the full block".to_string(),
signature: None,
},
});
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(2)));
}),
);
while stream.next().await.is_some() {}
let choice_items: Vec<AssistantContent> = stream.choice.clone().into_iter().collect();
let reasoning_items: Vec<&Reasoning> = choice_items
.iter()
.filter_map(|item| match item {
AssistantContent::Reasoning(reasoning) => Some(reasoning),
_ => None,
})
.collect();
assert_eq!(
reasoning_items.len(),
1,
"the full block must replace the delta-built item, not join it"
);
let only = reasoning_items.first().expect("one reasoning item");
assert_eq!(only.id.as_deref(), Some("rs_1"));
assert!(
only.content.iter().any(|content| matches!(
content,
ReasoningContent::Text { text, .. } if text == "the full block"
)),
"the surviving item must carry the full block's content"
);
}
#[tokio::test]
async fn minted_id_full_reasoning_block_does_not_clobber_a_wire_id_item() {
let mut stream = StreamingCompletionResponse::stream(
TEST_PROVIDER,
to_stream_result(stream! {
yield Ok(RawStreamingChoice::ReasoningDelta {
id: StreamPartId::wire("rs_1"),
provider_id: WireId::new("rs_1"),
reasoning: "identified deltas".to_string(),
});
yield Ok(RawStreamingChoice::Reasoning {
id: StreamPartId::wire("reasoning-0"),
provider_id: WireId::new("reasoning-0"),
content: ReasoningContent::Text {
text: "anonymous block".to_string(),
signature: None,
},
});
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(2)));
}),
);
while stream.next().await.is_some() {}
let choice_items: Vec<AssistantContent> = stream.choice.clone().into_iter().collect();
let reasoning_ids: Vec<Option<&str>> = choice_items
.iter()
.filter_map(|item| match item {
AssistantContent::Reasoning(reasoning) => Some(reasoning.id.as_deref()),
_ => None,
})
.collect();
assert_eq!(reasoning_ids, vec![Some("rs_1"), Some("reasoning-0")]);
}
#[tokio::test]
async fn test_stream_reasoning_only_does_not_inject_empty_text() {
let mut stream = create_reasoning_only_stream();
while stream.next().await.is_some() {}
let choice_items: Vec<AssistantContent> = stream.choice.clone().into_iter().collect();
assert_eq!(choice_items.len(), 1);
assert!(matches!(
choice_items.first(),
Some(AssistantContent::Reasoning(Reasoning { id: Some(id), .. })) if id == "rs_only"
));
}
#[tokio::test]
async fn test_stream_aggregates_assistant_items_in_arrival_order() {
let mut stream = create_interleaved_stream();
while stream.next().await.is_some() {}
let choice_items: Vec<AssistantContent> = stream.choice.clone().into_iter().collect();
assert_eq!(choice_items.len(), 3);
assert!(matches!(
choice_items.first(),
Some(AssistantContent::Reasoning(Reasoning { id: Some(id), .. })) if id == "rs_interleaved"
));
assert!(matches!(
choice_items.get(1),
Some(AssistantContent::Text(Text { text, .. })) if text == "final-text"
));
assert!(matches!(
choice_items.get(2),
Some(AssistantContent::ToolCall(ToolCall { id, .. })) if id == "tool_1"
));
}
#[tokio::test]
async fn unknown_choice_reaches_consumer_but_not_aggregated_choice() {
let unknown = serde_json::json!({
"type": "web_search_call",
"id": "ws_1",
"status": "completed",
});
let yielded = unknown.clone();
let stream = stream! {
yield Ok(RawStreamingChoice::Unknown(yielded.into()));
yield Ok(RawStreamingChoice::Message("done".to_string()));
yield Ok(RawStreamingChoice::FinalResponse(mock_final_with_total_tokens(1)));
};
let mut stream =
StreamingCompletionResponse::stream(TEST_PROVIDER, to_stream_result(stream));
let mut consumer_unknown = None;
let mut consumer_text = String::new();
while let Some(item) = stream.next().await {
match item.expect("stream item should be Ok") {
StreamedAssistantContent::Unknown(value) => consumer_unknown = Some(value),
StreamedAssistantContent::Text(text) => consumer_text.push_str(&text.text),
_ => {}
}
}
assert_eq!(consumer_unknown.as_ref(), Some(&unknown.into()));
assert_eq!(consumer_text, "done");
let choice_items: Vec<AssistantContent> = stream.choice.clone().into_iter().collect();
assert_eq!(choice_items.len(), 1);
assert!(matches!(
choice_items.first(),
Some(AssistantContent::Text(Text { text, .. })) if text == "done"
));
}
#[tokio::test]
async fn test_stream_keeps_non_contiguous_text_chunks_split_by_tool_call() {
let mut stream = create_text_tool_text_stream();
while stream.next().await.is_some() {}
let choice_items: Vec<AssistantContent> = stream.choice.clone().into_iter().collect();
assert_eq!(choice_items.len(), 3);
assert!(matches!(
choice_items.first(),
Some(AssistantContent::Text(Text { text, .. })) if text == "first"
));
assert!(matches!(
choice_items.get(1),
Some(AssistantContent::ToolCall(ToolCall { id, .. })) if id == "tool_split"
));
assert!(matches!(
choice_items.get(2),
Some(AssistantContent::Text(Text { text, .. })) if text == "second"
));
}
#[tokio::test]
async fn test_stream_preserves_text_additional_params() {
let mut stream = create_text_metadata_stream();
while stream.next().await.is_some() {}
let choice_items: Vec<AssistantContent> = stream.choice.clone().into_iter().collect();
assert_eq!(choice_items.len(), 2);
let Some(AssistantContent::Text(Text {
text,
additional_params: Some(additional_params),
})) = choice_items.first()
else {
panic!("expected first text item with metadata");
};
assert_eq!(text, "first");
assert_eq!(
additional_params["citations"]
.as_array()
.expect("citations should be an array")
.len(),
2
);
let Some(AssistantContent::Text(Text {
text,
additional_params: Some(additional_params),
})) = choice_items.get(1)
else {
panic!("expected second text item with metadata");
};
assert_eq!(text, "second");
assert_eq!(additional_params["block"], 2);
}
}
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(untagged)]
pub enum StreamedAssistantContent {
Text(Text),
ToolCall {
tool_call: ToolCall,
internal_call_id: String,
},
ToolCallDelta {
internal_call_id: String,
content: ToolCallDeltaContent,
},
Reasoning {
reasoning: Reasoning,
id: String,
},
ReasoningDelta {
id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
provider_id: Option<String>,
reasoning: String,
},
Final(StreamFinal),
Unknown(UnknownPayload),
}
impl StreamedAssistantContent {
pub fn text(text: &str) -> Self {
Self::Text(Text::new(text.to_string()))
}
pub fn final_response(res: StreamFinal) -> Self {
Self::Final(res)
}
}
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(untagged)]
pub enum StreamedUserContent {
ToolResult {
tool_result: ToolResult,
internal_call_id: String,
},
}
impl StreamedUserContent {
pub fn tool_result(tool_result: ToolResult, internal_call_id: String) -> Self {
Self::ToolResult {
tool_result,
internal_call_id,
}
}
}