diff --git a/codex-rs/api-client/src/aggregate.rs b/codex-rs/api-client/src/aggregate.rs new file mode 100644 index 0000000000..cee283dd2e --- /dev/null +++ b/codex-rs/api-client/src/aggregate.rs @@ -0,0 +1,177 @@ +use std::collections::VecDeque; +use std::pin::Pin; +use std::task::Context; +use std::task::Poll; + +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; +use futures::Stream; + +use crate::error::Result; +use crate::stream::ResponseEvent; + +#[derive(Clone, Copy, Debug)] +pub enum ChatAggregationMode { + AggregatedOnly, + Streaming, +} + +pub trait AggregateStreamExt: Stream> + Sized { + fn aggregate(self) -> AggregatedChatStream + where + Self: Unpin, + { + AggregatedChatStream::new(self, AggregateMode::AggregatedOnly) + } + + fn streaming_mode(self) -> AggregatedChatStream + where + Self: Unpin, + { + AggregatedChatStream::new(self, AggregateMode::Streaming) + } +} + +impl AggregateStreamExt for S where S: Stream> + Sized + Unpin {} + +enum AggregateMode { + AggregatedOnly, + Streaming, +} + +pub struct AggregatedChatStream { + inner: S, + cumulative: String, + cumulative_reasoning: String, + pending: VecDeque, + mode: AggregateMode, +} + +impl AggregatedChatStream +where + S: Stream> + Unpin, +{ + fn new(inner: S, mode: AggregateMode) -> Self { + Self { + inner, + cumulative: String::new(), + cumulative_reasoning: String::new(), + pending: VecDeque::new(), + mode, + } + } +} + +impl Stream for AggregatedChatStream +where + S: Stream> + Unpin, +{ + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if let Some(ev) = self.pending.pop_front() { + return Poll::Ready(Some(Ok(ev))); + } + + loop { + match Pin::new(&mut self.inner).poll_next(cx) { + std::task::Poll::Pending => return Poll::Pending, + std::task::Poll::Ready(None) => return std::task::Poll::Ready(None), + std::task::Poll::Ready(Some(Err(err))) => { + return std::task::Poll::Ready(Some(Err(err))); + } + std::task::Poll::Ready(Some(Ok(ResponseEvent::OutputItemDone(item)))) => { + let is_assistant_message = matches!( + &item, + ResponseItem::Message { role, .. } if role == "assistant" + ); + + if is_assistant_message { + if let ResponseItem::Message { role, content, .. } = item { + let mut text = String::new(); + for c in content { + match c { + ContentItem::InputText { text: t } + | ContentItem::OutputText { text: t } => text.push_str(&t), + ContentItem::InputImage { image_url } => { + text.push_str(&image_url) + } + } + } + self.cumulative.push_str(&text); + if matches!(self.mode, AggregateMode::Streaming) { + let output_item = + ResponseEvent::OutputItemDone(ResponseItem::Message { + id: None, + role, + content: vec![ContentItem::OutputText { + text: self.cumulative.clone(), + }], + }); + self.pending.push_back(output_item); + } + } + } else { + return std::task::Poll::Ready(Some(Ok(ResponseEvent::OutputItemDone( + item, + )))); + } + } + std::task::Poll::Ready(Some(Ok(ResponseEvent::OutputItemAdded(item)))) => { + if !matches!( + &item, + ResponseItem::Message { role, .. } if role == "assistant" + ) { + return std::task::Poll::Ready(Some(Ok(ResponseEvent::OutputItemAdded( + item, + )))); + } + } + std::task::Poll::Ready(Some(Ok(ResponseEvent::ReasoningContentDelta(delta)))) => { + self.cumulative_reasoning.push_str(&delta); + if matches!(self.mode, AggregateMode::Streaming) { + let ev = + ResponseEvent::ReasoningContentDelta(self.cumulative_reasoning.clone()); + self.pending.push_back(ev); + } + } + std::task::Poll::Ready(Some(Ok(ResponseEvent::ReasoningSummaryDelta(delta)))) => { + if matches!(self.mode, AggregateMode::Streaming) { + let ev = ResponseEvent::ReasoningSummaryDelta(delta); + self.pending.push_back(ev); + } + } + std::task::Poll::Ready(Some(Ok(ResponseEvent::Completed { + response_id, + token_usage, + }))) => { + let assistant_event = ResponseEvent::OutputItemDone(ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: self.cumulative.clone(), + }], + }); + let completion_event = ResponseEvent::Completed { + response_id, + token_usage, + }; + + if matches!(self.mode, AggregateMode::Streaming) { + self.pending.push_back(assistant_event); + self.pending.push_back(completion_event); + } else { + return std::task::Poll::Ready(Some(Ok(assistant_event))); + } + } + std::task::Poll::Ready(Some(Ok(ev))) => { + return std::task::Poll::Ready(Some(Ok(ev))); + } + } + + if let Some(ev) = self.pending.pop_front() { + return std::task::Poll::Ready(Some(Ok(ev))); + } + } + } +} diff --git a/codex-rs/api-client/src/common.rs b/codex-rs/api-client/src/common.rs new file mode 100644 index 0000000000..0a08f42326 --- /dev/null +++ b/codex-rs/api-client/src/common.rs @@ -0,0 +1,31 @@ +use std::time::Duration; + +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; + +/// Exponential backoff with a 100ms base and a cap on the exponent to avoid +/// unbounded growth. The attempt number is clamped to [0, 6]. +pub(crate) fn backoff(attempt: i64) -> Duration { + let capped = attempt.clamp(0, 6) as u32; + Duration::from_millis(100 * 2_i64.pow(capped) as u64) +} + +/// Apply the `x-openai-subagent` header when the session source indicates a +/// subagent. Returns the original builder unchanged when not applicable. +pub(crate) fn apply_subagent_header( + mut builder: reqwest::RequestBuilder, + session_source: Option<&SessionSource>, +) -> reqwest::RequestBuilder { + if let Some(SessionSource::SubAgent(sub)) = session_source { + let subagent = if let SubAgentSource::Other(label) = sub { + label.clone() + } else { + serde_json::to_value(sub) + .ok() + .and_then(|v| v.as_str().map(std::string::ToString::to_string)) + .unwrap_or_else(|| "other".to_string()) + }; + builder = builder.header("x-openai-subagent", subagent); + } + builder +}