diff --git a/codex-rs/api-client/src/chat.rs b/codex-rs/api-client/src/chat.rs index 4edf267eb6..5f727baf1e 100644 --- a/codex-rs/api-client/src/chat.rs +++ b/codex-rs/api-client/src/chat.rs @@ -9,6 +9,7 @@ use crate::error::Error; use crate::error::Result; use crate::stream::ResponseEvent; use crate::stream::ResponseStream; +use crate::stream::WireResponseStream; use codex_provider_config::ModelProviderInfo; #[derive(Clone)] @@ -25,6 +26,7 @@ pub struct ChatCompletionsApiClientConfig { pub model: String, pub otel_event_manager: OtelEventManager, pub session_source: SessionSource, + pub extra_headers: Vec<(String, String)>, } #[derive(Clone)] @@ -56,12 +58,18 @@ impl PayloadClient for ChatCompletionsApiClient { } let auth = crate::client::http::resolve_auth(&None).await; + let extra_headers: Vec<(&str, String)> = self + .config + .extra_headers + .iter() + .map(|(k, v)| (k.as_str(), v.clone())) + .collect(); let mut req_builder = crate::client::http::build_request( &self.config.http_client, &self.config.provider, &auth, session_source, - &[], + &extra_headers, ) .await?; @@ -95,3 +103,25 @@ impl PayloadClient for ChatCompletionsApiClient { Ok(crate::stream::EventStream::from_receiver(rx_event)) } } + +impl ChatCompletionsApiClient { + pub async fn stream_payload_wire( + &self, + payload_json: &serde_json::Value, + session_source: Option<&codex_protocol::protocol::SessionSource>, + ) -> Result { + use futures::StreamExt; + let legacy = self.stream_payload(payload_json, session_source).await?; + let (tx, rx) = tokio::sync::mpsc::channel(1600); + tokio::spawn(async move { + futures::pin_mut!(legacy); + while let Some(item) = legacy.next().await { + let converted = item.and_then(|ev| crate::wire::map_response_event_to_wire(ev)); + if tx.send(converted).await.is_err() { + break; + } + } + }); + Ok(crate::stream::EventStream::from_receiver(rx)) + } +} diff --git a/codex-rs/api-client/src/lib.rs b/codex-rs/api-client/src/lib.rs index ab3240d624..63550d6b7a 100644 --- a/codex-rs/api-client/src/lib.rs +++ b/codex-rs/api-client/src/lib.rs @@ -1,4 +1,3 @@ -pub mod aggregate; pub mod api; pub mod auth; pub mod chat; @@ -10,8 +9,8 @@ pub mod error; pub mod responses; pub mod routed_client; pub mod stream; +mod wire; -pub use crate::aggregate::AggregateStreamExt; pub use crate::auth::AuthContext; pub use crate::auth::AuthProvider; pub use crate::chat::ChatCompletionsApiClient; @@ -32,9 +31,5 @@ pub use crate::stream::TextFormat; pub use crate::stream::TextFormatType; pub use crate::stream::WireEvent; pub use crate::stream::WireResponseStream; -pub use codex_provider_config::BUILT_IN_OSS_MODEL_PROVIDER_ID; pub use codex_provider_config::ModelProviderInfo; pub use codex_provider_config::WireApi; -pub use codex_provider_config::built_in_model_providers; -pub use codex_provider_config::create_oss_provider; -pub use codex_provider_config::create_oss_provider_with_base_url; diff --git a/codex-rs/api-client/src/responses.rs b/codex-rs/api-client/src/responses.rs index c26f9f9959..4e07a2d99d 100644 --- a/codex-rs/api-client/src/responses.rs +++ b/codex-rs/api-client/src/responses.rs @@ -16,6 +16,7 @@ use crate::error::Error; use crate::error::Result; use crate::stream::ResponseEvent; use crate::stream::ResponseStream; +use crate::stream::WireResponseStream; use codex_provider_config::ModelProviderInfo; #[derive(Clone)] @@ -34,6 +35,7 @@ pub struct ResponsesApiClientConfig { pub conversation_id: ConversationId, pub auth_provider: Option>, pub otel_event_manager: OtelEventManager, + pub extra_headers: Vec<(String, String)>, } #[derive(Clone)] @@ -77,10 +79,21 @@ impl PayloadClient for ResponsesApiClient { .unwrap_or_else(|_| "".to_string()) ); - let extra_headers = vec![ - ("conversation_id", self.config.conversation_id.to_string()), - ("session_id", self.config.conversation_id.to_string()), + let mut owned_headers: Vec<(String, String)> = vec![ + ( + "conversation_id".to_string(), + self.config.conversation_id.to_string(), + ), + ( + "session_id".to_string(), + self.config.conversation_id.to_string(), + ), ]; + owned_headers.extend(self.config.extra_headers.iter().cloned()); + let extra_headers: Vec<(&str, String)> = owned_headers + .iter() + .map(|(k, v)| (k.as_str(), v.clone())) + .collect(); let mut req_builder = crate::client::http::build_request( &self.config.http_client, &self.config.provider, @@ -140,3 +153,25 @@ impl PayloadClient for ResponsesApiClient { Ok(crate::stream::EventStream::from_receiver(rx_event)) } } + +impl ResponsesApiClient { + pub async fn stream_payload_wire( + &self, + payload_json: &Value, + session_source: Option<&codex_protocol::protocol::SessionSource>, + ) -> Result { + use futures::StreamExt; + let legacy = self.stream_payload(payload_json, session_source).await?; + let (tx, rx) = tokio::sync::mpsc::channel(1600); + tokio::spawn(async move { + futures::pin_mut!(legacy); + while let Some(item) = legacy.next().await { + let converted = item.and_then(|ev| crate::wire::map_response_event_to_wire(ev)); + if tx.send(converted).await.is_err() { + break; + } + } + }); + Ok(crate::stream::EventStream::from_receiver(rx)) + } +} diff --git a/codex-rs/api-client/src/routed_client.rs b/codex-rs/api-client/src/routed_client.rs index 0ca4def2b6..e772961d35 100644 --- a/codex-rs/api-client/src/routed_client.rs +++ b/codex-rs/api-client/src/routed_client.rs @@ -17,7 +17,6 @@ use crate::WireResponseStream; use crate::api::PayloadClient; use crate::auth::AuthProvider; use crate::client::fixtures::stream_from_fixture; -use crate::stream::WireTokenUsage; use codex_provider_config::ModelProviderInfo; /// Dispatches to the appropriate API client implementation based on the provider wire API. @@ -62,6 +61,7 @@ impl RoutedApiClient { conversation_id: self.config.conversation_id, auth_provider: self.config.auth_provider.clone(), otel_event_manager: self.config.otel_event_manager.clone(), + extra_headers: vec![], }; let client = ::new(cfg)?; client @@ -75,6 +75,7 @@ impl RoutedApiClient { model: self.config.model.clone(), otel_event_manager: self.config.otel_event_manager.clone(), session_source: self.config.session_source.clone(), + extra_headers: vec![], }; let client = ::new(cfg)?; client @@ -122,52 +123,5 @@ impl PayloadClient for RoutedApiClient { } fn map_response_event_to_wire(ev: crate::stream::ResponseEvent) -> Result { - Ok(match ev { - crate::stream::ResponseEvent::Created => WireEvent::Created, - crate::stream::ResponseEvent::OutputItemDone(item) => { - WireEvent::OutputItemDone(serde_json::to_value(item).unwrap_or(serde_json::Value::Null)) - } - crate::stream::ResponseEvent::OutputItemAdded(item) => WireEvent::OutputItemAdded( - serde_json::to_value(item).unwrap_or(serde_json::Value::Null), - ), - crate::stream::ResponseEvent::Completed { - response_id, - token_usage, - } => { - let mapped = token_usage.map(|u| WireTokenUsage { - input_tokens: u.input_tokens, - cached_input_tokens: u.cached_input_tokens, - output_tokens: u.output_tokens, - reasoning_output_tokens: u.reasoning_output_tokens, - total_tokens: u.total_tokens, - }); - WireEvent::Completed { - response_id, - token_usage: mapped, - } - } - crate::stream::ResponseEvent::OutputTextDelta(s) => WireEvent::OutputTextDelta(s), - crate::stream::ResponseEvent::ReasoningSummaryDelta(s) => { - WireEvent::ReasoningSummaryDelta(s) - } - crate::stream::ResponseEvent::ReasoningContentDelta(s) => { - WireEvent::ReasoningContentDelta(s) - } - crate::stream::ResponseEvent::ReasoningSummaryPartAdded => { - WireEvent::ReasoningSummaryPartAdded - } - crate::stream::ResponseEvent::RateLimits(s) => { - let to_win = |w: Option| -> Option { - w.map(|w| crate::stream::WireRateLimitWindow { - used_percent: Some(w.used_percent), - window_minutes: w.window_minutes, - resets_at: w.resets_at, - }) - }; - WireEvent::RateLimits(crate::stream::WireRateLimitSnapshot { - primary: to_win(s.primary), - secondary: to_win(s.secondary), - }) - } - }) + crate::wire::map_response_event_to_wire(ev) } diff --git a/codex-rs/api-client/src/wire.rs b/codex-rs/api-client/src/wire.rs new file mode 100644 index 0000000000..c526043af6 --- /dev/null +++ b/codex-rs/api-client/src/wire.rs @@ -0,0 +1,55 @@ +use crate::error::Result; +use crate::stream::WireEvent; +use crate::stream::WireRateLimitSnapshot; +use crate::stream::WireRateLimitWindow; + +pub fn map_response_event_to_wire(ev: crate::stream::ResponseEvent) -> Result { + Ok(match ev { + crate::stream::ResponseEvent::Created => WireEvent::Created, + crate::stream::ResponseEvent::OutputItemDone(item) => { + WireEvent::OutputItemDone(serde_json::to_value(item).unwrap_or(serde_json::Value::Null)) + } + crate::stream::ResponseEvent::OutputItemAdded(item) => WireEvent::OutputItemAdded( + serde_json::to_value(item).unwrap_or(serde_json::Value::Null), + ), + crate::stream::ResponseEvent::Completed { + response_id, + token_usage, + } => { + let mapped = token_usage.map(|u| crate::stream::WireTokenUsage { + input_tokens: u.input_tokens, + cached_input_tokens: u.cached_input_tokens, + output_tokens: u.output_tokens, + reasoning_output_tokens: u.reasoning_output_tokens, + total_tokens: u.total_tokens, + }); + WireEvent::Completed { + response_id, + token_usage: mapped, + } + } + crate::stream::ResponseEvent::OutputTextDelta(s) => WireEvent::OutputTextDelta(s), + crate::stream::ResponseEvent::ReasoningSummaryDelta(s) => { + WireEvent::ReasoningSummaryDelta(s) + } + crate::stream::ResponseEvent::ReasoningContentDelta(s) => { + WireEvent::ReasoningContentDelta(s) + } + crate::stream::ResponseEvent::ReasoningSummaryPartAdded => { + WireEvent::ReasoningSummaryPartAdded + } + crate::stream::ResponseEvent::RateLimits(s) => { + let to_win = |w: Option| -> Option { + w.map(|w| WireRateLimitWindow { + used_percent: Some(w.used_percent), + window_minutes: w.window_minutes, + resets_at: w.resets_at, + }) + }; + WireEvent::RateLimits(WireRateLimitSnapshot { + primary: to_win(s.primary), + secondary: to_win(s.secondary), + }) + } + }) +} diff --git a/codex-rs/api-client/src/aggregate.rs b/codex-rs/core/src/aggregate.rs similarity index 98% rename from codex-rs/api-client/src/aggregate.rs rename to codex-rs/core/src/aggregate.rs index 44c02e1f88..09ce8fefdc 100644 --- a/codex-rs/api-client/src/aggregate.rs +++ b/codex-rs/core/src/aggregate.rs @@ -3,12 +3,12 @@ use std::pin::Pin; use std::task::Context; use std::task::Poll; -use codex_protocol::models::ContentItem; -use codex_protocol::models::ResponseItem; +use crate::ContentItem; +use crate::ResponseItem; use futures::Stream; +use crate::ResponseEvent; use crate::error::Result; -use crate::stream::ResponseEvent; pub trait AggregateStreamExt: Stream> + Sized { fn aggregate(self) -> AggregatedChatStream diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 88b340d462..07899c8be9 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -44,6 +44,7 @@ pub use codex_provider_config::ModelProviderInfo; pub use codex_provider_config::WireApi; pub use codex_provider_config::built_in_model_providers; pub use codex_provider_config::create_oss_provider_with_base_url; +mod aggregate; mod conversation_manager; mod event_mapping; pub mod review_format; @@ -94,6 +95,7 @@ pub use codex_protocol::protocol; // as those in the protocol crate when constructing protocol messages. pub use codex_protocol::config_types as protocol_config_types; +pub use aggregate::AggregateStreamExt; pub use client::ModelClient; pub use client_common::Prompt; pub use client_common::REVIEW_PROMPT;