use futures::Stream;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::borrow::Cow;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use tokio_util::sync::CancellationToken;
pub use crate::stream_error_kind::StreamErrorKind;
use crate::types::{
AgentContext, AssistantMessage, ContentBlock, Cost, ModelSpec, StopReason, Usage,
};
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StreamTransport {
#[default]
Sse,
}
#[non_exhaustive]
#[derive(Debug, Clone, Default)]
pub enum CacheStrategy {
#[default]
None,
Auto,
Anthropic,
Google {
ttl: Duration,
},
}
pub type OnRawPayload = Arc<dyn Fn(&str) + Send + Sync>;
#[non_exhaustive]
#[derive(Clone, Debug, Default, PartialEq)]
pub struct RateLimitSnapshot {
pub used_percent: Option<f32>,
pub remaining_requests: Option<u64>,
pub remaining_tokens: Option<u64>,
pub resets_in: Option<Duration>,
pub window: Option<Duration>,
pub plan: Option<String>,
pub raw: std::collections::BTreeMap<String, String>,
}
impl RateLimitSnapshot {
pub fn from_headers<'a>(headers: impl IntoIterator<Item = (&'a str, &'a str)>) -> Self {
let mut snapshot = Self::default();
for (name, value) in headers {
let name = name.to_ascii_lowercase();
let value = value.trim();
if !is_rate_limit_header(&name) {
continue;
}
match name.as_str() {
"x-codex-primary-used-percent" => snapshot.used_percent = value.parse().ok(),
"x-ratelimit-remaining-requests" | "anthropic-ratelimit-requests-remaining" => {
snapshot.remaining_requests = value.parse().ok();
}
"x-ratelimit-remaining-tokens" | "anthropic-ratelimit-tokens-remaining" => {
snapshot.remaining_tokens = value.parse().ok();
}
"x-codex-primary-reset-after-seconds" | "x-ratelimit-reset-requests" => {
snapshot.resets_in = parse_reset_duration(value);
}
"retry-after" if snapshot.resets_in.is_none() => {
snapshot.resets_in = parse_reset_duration(value);
}
"x-codex-primary-window-minutes" => {
snapshot.window = value
.parse::<u64>()
.ok()
.map(|m| Duration::from_secs(m * 60));
}
"x-codex-plan-type" => snapshot.plan = Some(value.to_owned()),
_ => {}
}
snapshot.raw.insert(name, value.to_owned());
}
snapshot
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.raw.is_empty()
}
}
fn is_rate_limit_header(name: &str) -> bool {
name.contains("ratelimit")
|| name.contains("rate-limit")
|| name.starts_with("x-codex-")
|| name == "retry-after"
}
fn parse_reset_duration(value: &str) -> Option<Duration> {
if let Ok(secs) = value.parse::<f64>() {
return (secs.is_finite() && secs >= 0.0).then(|| Duration::from_secs_f64(secs));
}
let mut total = Duration::ZERO;
let mut number = String::new();
let mut saw_unit = false;
let mut chars = value.chars().peekable();
while let Some(c) = chars.next() {
if c.is_ascii_digit() || c == '.' {
number.push(c);
continue;
}
let amount: f64 = number.parse().ok()?;
number.clear();
let unit = match c {
'm' if chars.peek() == Some(&'s') => {
chars.next();
0.001
}
'h' => 3600.0,
'm' => 60.0,
's' => 1.0,
_ => return None,
};
total += Duration::from_secs_f64(amount * unit);
saw_unit = true;
}
(saw_unit && number.is_empty()).then_some(total)
}
pub type OnRateLimit = Arc<dyn Fn(&RateLimitSnapshot) + Send + Sync>;
#[non_exhaustive]
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct ServingOptions {
pub context_length: Option<u64>,
pub top_p: Option<f64>,
pub keep_alive: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub format: Option<ResponseFormat>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reasoning_effort: Option<ReasoningEffort>,
#[serde(default, skip_serializing_if = "std::collections::BTreeMap::is_empty")]
pub extra: std::collections::BTreeMap<String, Value>,
}
#[non_exhaustive]
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ResponseFormat {
Json,
Schema(Value),
}
#[non_exhaustive]
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ReasoningEffort {
Off,
Minimal,
Low,
Medium,
High,
#[serde(rename = "xhigh", alias = "x_high", alias = "extra_high")]
XHigh,
Max,
}
impl ServingOptions {
#[must_use]
pub const fn with_context_length(mut self, context_length: u64) -> Self {
self.context_length = Some(context_length);
self
}
#[must_use]
pub const fn with_top_p(mut self, top_p: f64) -> Self {
self.top_p = Some(top_p);
self
}
#[must_use]
pub fn with_keep_alive(mut self, keep_alive: impl Into<String>) -> Self {
self.keep_alive = Some(keep_alive.into());
self
}
#[must_use]
pub fn with_format(mut self, format: ResponseFormat) -> Self {
self.format = Some(format);
self
}
#[must_use]
pub const fn with_reasoning_effort(mut self, reasoning_effort: ReasoningEffort) -> Self {
self.reasoning_effort = Some(reasoning_effort);
self
}
#[must_use]
pub fn with_extra(mut self, extra: std::collections::BTreeMap<String, Value>) -> Self {
self.extra = extra;
self
}
pub fn is_default(&self) -> bool {
*self == Self::default()
}
#[must_use]
pub fn unsupported_fields(&self, support: ServingOptionSupport) -> Vec<&'static str> {
let mut dropped = Vec::new();
if self.context_length.is_some() && !support.context_length {
dropped.push("context_length");
}
if self.top_p.is_some() && !support.top_p {
dropped.push("top_p");
}
if self.keep_alive.is_some() && !support.keep_alive {
dropped.push("keep_alive");
}
if self.format.is_some() && !support.format {
dropped.push("format");
}
if self.reasoning_effort.is_some() && !support.reasoning_effort {
dropped.push("reasoning_effort");
}
if !self.extra.is_empty() && !support.extra {
dropped.push("extra");
}
dropped
}
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[allow(clippy::struct_excessive_bools)]
pub struct ServingOptionSupport {
pub context_length: bool,
pub top_p: bool,
pub keep_alive: bool,
pub format: bool,
pub reasoning_effort: bool,
pub extra: bool,
}
impl ServingOptionSupport {
#[must_use]
pub const fn all() -> Self {
Self {
context_length: true,
top_p: true,
keep_alive: true,
format: true,
reasoning_effort: true,
extra: true,
}
}
#[must_use]
pub const fn none() -> Self {
Self {
context_length: false,
top_p: false,
keep_alive: false,
format: false,
reasoning_effort: false,
extra: false,
}
}
#[must_use]
pub const fn with_context_length(mut self, supported: bool) -> Self {
self.context_length = supported;
self
}
#[must_use]
pub const fn with_top_p(mut self, supported: bool) -> Self {
self.top_p = supported;
self
}
#[must_use]
pub const fn with_keep_alive(mut self, supported: bool) -> Self {
self.keep_alive = supported;
self
}
#[must_use]
pub const fn with_format(mut self, supported: bool) -> Self {
self.format = supported;
self
}
#[must_use]
pub const fn with_reasoning_effort(mut self, supported: bool) -> Self {
self.reasoning_effort = supported;
self
}
#[must_use]
pub const fn with_extra(mut self, supported: bool) -> Self {
self.extra = supported;
self
}
}
#[non_exhaustive]
#[derive(Clone, Default)]
pub struct StreamOptions {
pub temperature: Option<f64>,
pub max_tokens: Option<u64>,
pub session_id: Option<String>,
pub api_key: Option<String>,
pub transport: StreamTransport,
pub cache_strategy: CacheStrategy,
pub on_raw_payload: Option<OnRawPayload>,
pub on_rate_limit: Option<OnRateLimit>,
pub serving: ServingOptions,
}
impl StreamOptions {
#[must_use]
pub const fn with_temperature(mut self, temperature: f64) -> Self {
self.temperature = Some(temperature);
self
}
#[must_use]
pub const fn with_max_tokens(mut self, max_tokens: u64) -> Self {
self.max_tokens = Some(max_tokens);
self
}
#[must_use]
pub fn with_session_id(mut self, session_id: impl Into<String>) -> Self {
self.session_id = Some(session_id.into());
self
}
#[must_use]
pub fn with_api_key(mut self, api_key: impl Into<String>) -> Self {
self.api_key = Some(api_key.into());
self
}
#[must_use]
pub const fn with_transport(mut self, transport: StreamTransport) -> Self {
self.transport = transport;
self
}
#[must_use]
pub fn with_cache_strategy(mut self, cache_strategy: CacheStrategy) -> Self {
self.cache_strategy = cache_strategy;
self
}
#[must_use]
pub fn with_on_raw_payload(mut self, on_raw_payload: OnRawPayload) -> Self {
self.on_raw_payload = Some(on_raw_payload);
self
}
#[must_use]
pub fn with_on_rate_limit(mut self, on_rate_limit: OnRateLimit) -> Self {
self.on_rate_limit = Some(on_rate_limit);
self
}
#[must_use]
pub fn with_serving(mut self, serving: ServingOptions) -> Self {
self.serving = serving;
self
}
}
impl std::fmt::Debug for StreamOptions {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("StreamOptions")
.field("temperature", &self.temperature)
.field("max_tokens", &self.max_tokens)
.field("session_id", &self.session_id)
.field("api_key", &self.api_key.as_ref().map(|_| "[REDACTED]"))
.field("transport", &self.transport)
.field("cache_strategy", &self.cache_strategy)
.field(
"on_raw_payload",
&self.on_raw_payload.as_ref().map(|_| "<callback>"),
)
.field(
"on_rate_limit",
&self.on_rate_limit.as_ref().map(|_| "<callback>"),
)
.field("serving", &self.serving)
.finish()
}
}
#[non_exhaustive]
#[derive(Debug, Clone)]
pub enum AssistantMessageEvent {
Start,
TextStart { content_index: usize },
TextDelta { content_index: usize, delta: String },
TextEnd { content_index: usize },
ThinkingStart { content_index: usize },
ThinkingDelta { content_index: usize, delta: String },
ThinkingEnd {
content_index: usize,
signature: Option<String>,
},
ToolCallStart {
content_index: usize,
id: String,
name: String,
},
ToolCallDelta { content_index: usize, delta: String },
ToolCallEnd { content_index: usize },
Done {
stop_reason: StopReason,
usage: Usage,
cost: Cost,
},
Error {
stop_reason: StopReason,
error_message: String,
usage: Option<Usage>,
error_kind: Option<StreamErrorKind>,
retry_after: Option<std::time::Duration>,
},
}
impl AssistantMessageEvent {
pub fn error(message: impl Into<String>) -> Self {
Self::Error {
stop_reason: StopReason::Error,
error_message: message.into(),
usage: None,
error_kind: None,
retry_after: None,
}
}
pub fn error_throttled(message: impl Into<String>) -> Self {
Self::Error {
stop_reason: StopReason::Error,
error_message: message.into(),
usage: None,
error_kind: Some(StreamErrorKind::Throttled),
retry_after: None,
}
}
pub fn error_context_overflow(message: impl Into<String>) -> Self {
Self::Error {
stop_reason: StopReason::Error,
error_message: message.into(),
usage: None,
error_kind: Some(StreamErrorKind::ContextWindowExceeded),
retry_after: None,
}
}
pub fn error_auth(message: impl Into<String>) -> Self {
Self::Error {
stop_reason: StopReason::Error,
error_message: message.into(),
usage: None,
error_kind: Some(StreamErrorKind::Auth),
retry_after: None,
}
}
pub fn error_network(message: impl Into<String>) -> Self {
Self::Error {
stop_reason: StopReason::Error,
error_message: message.into(),
usage: None,
error_kind: Some(StreamErrorKind::Network),
retry_after: None,
}
}
pub fn error_content_filtered(message: impl Into<String>) -> Self {
Self::Error {
stop_reason: StopReason::Error,
error_message: message.into(),
usage: None,
error_kind: Some(StreamErrorKind::ContentFiltered),
retry_after: None,
}
}
pub fn error_model_retired(message: impl Into<String>) -> Self {
Self::Error {
stop_reason: StopReason::Error,
error_message: message.into(),
usage: None,
error_kind: Some(StreamErrorKind::ModelRetired),
retry_after: None,
}
}
pub fn text_response(text: &str) -> Vec<Self> {
vec![
Self::Start,
Self::TextStart { content_index: 0 },
Self::TextDelta {
content_index: 0,
delta: text.to_string(),
},
Self::TextEnd { content_index: 0 },
Self::Done {
stop_reason: StopReason::Stop,
usage: Usage::default(),
cost: Cost::default(),
},
]
}
}
#[non_exhaustive]
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AssistantMessageDelta {
Text {
content_index: usize,
delta: Cow<'static, str>,
},
Thinking {
content_index: usize,
delta: Cow<'static, str>,
},
ToolCall {
content_index: usize,
delta: Cow<'static, str>,
},
}
pub trait StreamFn: Send + Sync {
fn stream<'a>(
&'a self,
model: &'a ModelSpec,
context: &'a AgentContext,
options: &'a StreamOptions,
cancellation_token: CancellationToken,
) -> Pin<Box<dyn Stream<Item = AssistantMessageEvent> + Send + 'a>>;
fn supported_serving_options(&self) -> ServingOptionSupport {
ServingOptionSupport::all()
}
}
#[must_use]
pub fn stream_owned(
stream_fn: Arc<dyn StreamFn>,
model: ModelSpec,
context: AgentContext,
options: StreamOptions,
cancellation_token: CancellationToken,
) -> Pin<Box<dyn Stream<Item = AssistantMessageEvent> + Send + 'static>> {
let (tx, rx) = tokio::sync::mpsc::channel(16);
tokio::spawn(async move {
use futures::StreamExt as _;
let mut events = stream_fn.stream(&model, &context, &options, cancellation_token);
while let Some(event) = events.next().await {
if tx.send(event).await.is_err() {
break;
}
}
});
Box::pin(tokio_stream::wrappers::ReceiverStream::new(rx))
}
pub type MappedOptions = Result<StreamOptions, Vec<AssistantMessageEvent>>;
pub struct MapOptionsStreamFn<F> {
inner: Arc<dyn StreamFn>,
map: F,
}
impl<F> MapOptionsStreamFn<F>
where
F: Fn(&ModelSpec, &AgentContext, StreamOptions) -> MappedOptions + Send + Sync,
{
pub fn new(inner: Arc<dyn StreamFn>, map: F) -> Self {
Self { inner, map }
}
}
impl<F> StreamFn for MapOptionsStreamFn<F>
where
F: Fn(&ModelSpec, &AgentContext, StreamOptions) -> MappedOptions + Send + Sync,
{
fn stream<'a>(
&'a self,
model: &'a ModelSpec,
context: &'a AgentContext,
options: &'a StreamOptions,
cancellation_token: CancellationToken,
) -> Pin<Box<dyn Stream<Item = AssistantMessageEvent> + Send + 'a>> {
match (self.map)(model, context, options.clone()) {
Ok(mapped) => stream_owned(
Arc::clone(&self.inner),
model.clone(),
context.clone_for_send(),
mapped,
cancellation_token,
),
Err(events) => Box::pin(futures::stream::iter(events)),
}
}
fn supported_serving_options(&self) -> ServingOptionSupport {
self.inner.supported_serving_options()
}
}
pub fn sanitize_incomplete_tool_calls(message: &mut AssistantMessage) -> usize {
let mut fixed = 0;
for block in &mut message.content {
if let ContentBlock::ToolCall {
arguments,
partial_json,
..
} = block
{
let needs_fix = partial_json.is_some() || !arguments.is_object();
if needs_fix {
*arguments = Value::Object(serde_json::Map::new());
*partial_json = None;
fixed += 1;
}
}
}
fixed
}
#[allow(clippy::too_many_lines)]
pub fn accumulate_message(
events: Vec<AssistantMessageEvent>,
provider: &str,
model_id: &str,
) -> Result<AssistantMessage, String> {
fn ensure_block_open(
open_blocks: &[bool],
content_index: usize,
event_name: &str,
) -> Result<(), String> {
match open_blocks.get(content_index) {
Some(false) => Err(format!(
"{event_name}: block at index {content_index} is already closed"
)),
Some(true) | None => Ok(()),
}
}
fn all_open_blocks_are_tool_calls(content: &[ContentBlock], open_blocks: &[bool]) -> bool {
open_blocks
.iter()
.enumerate()
.filter(|(_, open)| **open)
.all(|(content_index, _)| {
matches!(
content.get(content_index),
Some(ContentBlock::ToolCall { .. })
)
})
}
fn validate_terminal_open_blocks(
event_name: &str,
content: Option<&[ContentBlock]>,
open_blocks: &[bool],
tolerate_truncated_tool_args: bool,
) -> Result<(), String> {
if let Some(idx) = open_blocks.iter().position(|open| *open) {
let content = content.ok_or_else(|| format!("{event_name} before Start"))?;
if tolerate_truncated_tool_args && all_open_blocks_are_tool_calls(content, open_blocks)
{
tracing::debug!(
"{event_name}(Length) with unterminated content block at index {idx} - tolerating for max-tokens recovery"
);
} else {
return Err(format!(
"{event_name} received with unterminated content block at index {idx}"
));
}
}
Ok(())
}
let mut content: Option<Vec<ContentBlock>> = None;
let mut open_blocks: Vec<bool> = Vec::new();
let mut stop_reason: Option<StopReason> = None;
let mut usage: Option<Usage> = None;
let mut cost: Option<Cost> = None;
let mut error_message: Option<String> = None;
let mut error_kind: Option<StreamErrorKind> = None;
let mut saw_start = false;
let mut saw_terminal = false;
let tolerate_truncated_tool_args = events.iter().any(|e| {
matches!(
e,
AssistantMessageEvent::Done {
stop_reason: StopReason::Length,
..
}
)
});
for event in events {
match &event {
AssistantMessageEvent::TextStart { .. }
| AssistantMessageEvent::TextDelta { .. }
| AssistantMessageEvent::TextEnd { .. }
| AssistantMessageEvent::ThinkingStart { .. }
| AssistantMessageEvent::ThinkingDelta { .. }
| AssistantMessageEvent::ThinkingEnd { .. }
| AssistantMessageEvent::ToolCallStart { .. }
| AssistantMessageEvent::ToolCallDelta { .. }
| AssistantMessageEvent::ToolCallEnd { .. } => {
if saw_terminal {
return Err("content event after terminal event".into());
}
}
AssistantMessageEvent::Done { .. } | AssistantMessageEvent::Error { .. } => {
if saw_terminal {
return Err("duplicate terminal event".into());
}
}
AssistantMessageEvent::Start => {
if saw_terminal {
return Err("Start event after terminal event".into());
}
}
}
match event {
AssistantMessageEvent::Start => {
if saw_start {
return Err("duplicate Start event".into());
}
saw_start = true;
content = Some(Vec::new());
}
AssistantMessageEvent::TextStart { content_index } => {
let blocks = content.as_mut().ok_or("TextStart before Start")?;
if content_index != blocks.len() {
return Err(format!(
"TextStart content_index {content_index} != content length {}",
blocks.len()
));
}
blocks.push(ContentBlock::Text {
text: String::new(),
});
open_blocks.push(true);
}
AssistantMessageEvent::TextDelta {
content_index,
delta,
} => {
let blocks = content.as_mut().ok_or("TextDelta before Start")?;
ensure_block_open(&open_blocks, content_index, "TextDelta")?;
let block = blocks
.get_mut(content_index)
.ok_or_else(|| format!("TextDelta: invalid content_index {content_index}"))?;
match block {
ContentBlock::Text { text } => text.push_str(&delta),
_ => {
return Err(format!(
"TextDelta: block at index {content_index} is not Text"
));
}
}
}
AssistantMessageEvent::TextEnd { content_index } => {
let blocks = content.as_ref().ok_or("TextEnd before Start")?;
let block = blocks
.get(content_index)
.ok_or_else(|| format!("TextEnd: invalid content_index {content_index}"))?;
if !matches!(block, ContentBlock::Text { .. }) {
return Err(format!(
"TextEnd: block at index {content_index} is not Text"
));
}
ensure_block_open(&open_blocks, content_index, "TextEnd")?;
if let Some(open) = open_blocks.get_mut(content_index) {
*open = false;
}
}
AssistantMessageEvent::ThinkingStart { content_index } => {
let blocks = content.as_mut().ok_or("ThinkingStart before Start")?;
if content_index != blocks.len() {
return Err(format!(
"ThinkingStart content_index {content_index} != content length {}",
blocks.len()
));
}
blocks.push(ContentBlock::Thinking {
thinking: String::new(),
signature: None,
});
open_blocks.push(true);
}
AssistantMessageEvent::ThinkingDelta {
content_index,
delta,
} => {
let blocks = content.as_mut().ok_or("ThinkingDelta before Start")?;
ensure_block_open(&open_blocks, content_index, "ThinkingDelta")?;
let block = blocks.get_mut(content_index).ok_or_else(|| {
format!("ThinkingDelta: invalid content_index {content_index}")
})?;
match block {
ContentBlock::Thinking { thinking, .. } => thinking.push_str(&delta),
_ => {
return Err(format!(
"ThinkingDelta: block at index {content_index} is not Thinking"
));
}
}
}
AssistantMessageEvent::ThinkingEnd {
content_index,
signature,
} => {
let blocks = content.as_mut().ok_or("ThinkingEnd before Start")?;
ensure_block_open(&open_blocks, content_index, "ThinkingEnd")?;
let block = blocks
.get_mut(content_index)
.ok_or_else(|| format!("ThinkingEnd: invalid content_index {content_index}"))?;
match block {
ContentBlock::Thinking { signature: sig, .. } => *sig = signature,
_ => {
return Err(format!(
"ThinkingEnd: block at index {content_index} is not Thinking"
));
}
}
if let Some(open) = open_blocks.get_mut(content_index) {
*open = false;
}
}
AssistantMessageEvent::ToolCallStart {
content_index,
id,
name,
} => {
let blocks = content.as_mut().ok_or("ToolCallStart before Start")?;
if content_index != blocks.len() {
return Err(format!(
"ToolCallStart content_index {content_index} != content length {}",
blocks.len()
));
}
blocks.push(ContentBlock::ToolCall {
id,
name,
arguments: Value::Null,
partial_json: Some(String::new()),
});
open_blocks.push(true);
}
AssistantMessageEvent::ToolCallDelta {
content_index,
delta,
} => {
let blocks = content.as_mut().ok_or("ToolCallDelta before Start")?;
ensure_block_open(&open_blocks, content_index, "ToolCallDelta")?;
let block = blocks.get_mut(content_index).ok_or_else(|| {
format!("ToolCallDelta: invalid content_index {content_index}")
})?;
match block {
ContentBlock::ToolCall { partial_json, .. } => {
let pj = partial_json
.as_mut()
.ok_or("ToolCallDelta: partial_json already consumed")?;
pj.push_str(&delta);
}
_ => {
return Err(format!(
"ToolCallDelta: block at index {content_index} is not ToolCall"
));
}
}
}
AssistantMessageEvent::ToolCallEnd { content_index } => {
let blocks = content.as_mut().ok_or("ToolCallEnd before Start")?;
let block = blocks
.get_mut(content_index)
.ok_or_else(|| format!("ToolCallEnd: invalid content_index {content_index}"))?;
ensure_block_open(&open_blocks, content_index, "ToolCallEnd")?;
match block {
ContentBlock::ToolCall {
arguments,
partial_json,
..
} => {
let json_str = partial_json
.as_ref()
.ok_or("ToolCallEnd: partial_json already consumed")?
.clone();
if json_str.is_empty() {
*arguments = Value::Object(serde_json::Map::new());
*partial_json = None;
} else {
match serde_json::from_str::<Value>(&json_str) {
Ok(v) => {
*arguments = v;
*partial_json = None;
}
Err(e) => {
if tolerate_truncated_tool_args {
} else {
return Err(format!(
"ToolCallEnd: failed to parse arguments JSON: {e}"
));
}
}
}
}
}
_ => {
return Err(format!(
"ToolCallEnd: block at index {content_index} is not ToolCall"
));
}
}
if let Some(open) = open_blocks.get_mut(content_index) {
*open = false;
}
}
AssistantMessageEvent::Done {
stop_reason: sr,
usage: u,
cost: c,
} => {
validate_terminal_open_blocks(
"Done",
content.as_deref(),
&open_blocks,
tolerate_truncated_tool_args,
)?;
stop_reason = Some(sr);
usage = Some(u);
cost = Some(c);
saw_terminal = true;
}
AssistantMessageEvent::Error {
stop_reason: sr,
error_message: em,
usage: u,
error_kind: ek,
retry_after: _,
} => {
validate_terminal_open_blocks("Error", content.as_deref(), &open_blocks, false)?;
stop_reason = Some(sr);
error_message = Some(em);
error_kind = ek;
if let Some(u) = u {
usage = Some(u);
}
saw_terminal = true;
}
}
}
let content = content.ok_or("no Start event found")?;
let stop_reason = stop_reason.ok_or("no terminal event (Done or Error) found")?;
let timestamp = crate::util::now_timestamp();
Ok(AssistantMessage {
content,
provider: provider.to_owned(),
model_id: model_id.to_owned(),
usage: usage.unwrap_or_default(),
cost: cost.unwrap_or_default(),
stop_reason,
error_message,
error_kind,
timestamp,
cache_hint: None,
})
}
const _: () = {
const fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<StreamErrorKind>();
assert_send_sync::<StreamTransport>();
assert_send_sync::<StreamOptions>();
assert_send_sync::<AssistantMessageEvent>();
assert_send_sync::<AssistantMessageDelta>();
};
#[cfg(test)]
#[path = "stream_tests.rs"]
mod tests;
#[cfg(test)]
#[path = "stream_reasoning_effort_tests.rs"]
mod reasoning_effort_tests;
#[cfg(test)]
#[path = "stream_reasoning_effort_alias_tests.rs"]
mod reasoning_effort_alias_tests;