chore: proper client extraction (#6996)

codex-rs/codex-api/src/endpoint/chat.rs

@@ -0,0 +1,266 @@
+use crate::ChatRequest;
+use crate::auth::AuthProvider;
+use crate::common::Prompt as ApiPrompt;
+use crate::common::ResponseEvent;
+use crate::common::ResponseStream;
+use crate::endpoint::streaming::StreamingClient;
+use crate::error::ApiError;
+use crate::provider::Provider;
+use crate::provider::WireApi;
+use crate::sse::chat::spawn_chat_stream;
+use crate::telemetry::SseTelemetry;
+use codex_client::HttpTransport;
+use codex_client::RequestTelemetry;
+use codex_protocol::models::ContentItem;
+use codex_protocol::models::ReasoningItemContent;
+use codex_protocol::models::ResponseItem;
+use codex_protocol::protocol::SessionSource;
+use futures::Stream;
+use http::HeaderMap;
+use serde_json::Value;
+use std::collections::VecDeque;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::Context;
+use std::task::Poll;
+
+pub struct ChatClient<T: HttpTransport, A: AuthProvider> {
+    streaming: StreamingClient<T, A>,
+}
+
+impl<T: HttpTransport, A: AuthProvider> ChatClient<T, A> {
+    pub fn new(transport: T, provider: Provider, auth: A) -> Self {
+        Self {
+            streaming: StreamingClient::new(transport, provider, auth),
+        }
+    }
+
+    pub fn with_telemetry(
+        self,
+        request: Option<Arc<dyn RequestTelemetry>>,
+        sse: Option<Arc<dyn SseTelemetry>>,
+    ) -> Self {
+        Self {
+            streaming: self.streaming.with_telemetry(request, sse),
+        }
+    }
+
+    pub async fn stream_request(&self, request: ChatRequest) -> Result<ResponseStream, ApiError> {
+        self.stream(request.body, request.headers).await
+    }
+
+    pub async fn stream_prompt(
+        &self,
+        model: &str,
+        prompt: &ApiPrompt,
+        conversation_id: Option<String>,
+        session_source: Option<SessionSource>,
+    ) -> Result<ResponseStream, ApiError> {
+        use crate::requests::ChatRequestBuilder;
+
+        let request =
+            ChatRequestBuilder::new(model, &prompt.instructions, &prompt.input, &prompt.tools)
+                .conversation_id(conversation_id)
+                .session_source(session_source)
+                .build(self.streaming.provider())?;
+
+        self.stream_request(request).await
+    }
+
+    fn path(&self) -> &'static str {
+        match self.streaming.provider().wire {
+            WireApi::Chat => "chat/completions",
+            _ => "responses",
+        }
+    }
+
+    pub async fn stream(
+        &self,
+        body: Value,
+        extra_headers: HeaderMap,
+    ) -> Result<ResponseStream, ApiError> {
+        self.streaming
+            .stream(self.path(), body, extra_headers, spawn_chat_stream)
+            .await
+    }
+}
+
+#[derive(Copy, Clone, Eq, PartialEq)]
+pub enum AggregateMode {
+    AggregatedOnly,
+    Streaming,
+}
+
+/// Stream adapter that merges token deltas into a single assistant message per turn.
+pub struct AggregatedStream {
+    inner: ResponseStream,
+    cumulative: String,
+    cumulative_reasoning: String,
+    pending: VecDeque<ResponseEvent>,
+    mode: AggregateMode,
+}
+
+impl Stream for AggregatedStream {
+    type Item = Result<ResponseEvent, ApiError>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.get_mut();
+
+        if let Some(ev) = this.pending.pop_front() {
+            return Poll::Ready(Some(Ok(ev)));
+        }
+
+        loop {
+            match Pin::new(&mut this.inner).poll_next(cx) {
+                Poll::Pending => return Poll::Pending,
+                Poll::Ready(None) => return Poll::Ready(None),
+                Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(e))),
+                Poll::Ready(Some(Ok(ResponseEvent::OutputItemDone(item)))) => {
+                    let is_assistant_message = matches!(
+                        &item,
+                        ResponseItem::Message { role, .. } if role == "assistant"
+                    );
+
+                    if is_assistant_message {
+                        match this.mode {
+                            AggregateMode::AggregatedOnly => {
+                                if this.cumulative.is_empty()
+                                    && let ResponseItem::Message { content, .. } = &item
+                                    && let Some(text) = content.iter().find_map(|c| match c {
+                                        ContentItem::OutputText { text } => Some(text),
+                                        _ => None,
+                                    })
+                                {
+                                    this.cumulative.push_str(text);
+                                }
+                                continue;
+                            }
+                            AggregateMode::Streaming => {
+                                if this.cumulative.is_empty() {
+                                    return Poll::Ready(Some(Ok(ResponseEvent::OutputItemDone(
+                                        item,
+                                    ))));
+                                } else {
+                                    continue;
+                                }
+                            }
+                        }
+                    }
+
+                    return Poll::Ready(Some(Ok(ResponseEvent::OutputItemDone(item))));
+                }
+                Poll::Ready(Some(Ok(ResponseEvent::RateLimits(snapshot)))) => {
+                    return Poll::Ready(Some(Ok(ResponseEvent::RateLimits(snapshot))));
+                }
+                Poll::Ready(Some(Ok(ResponseEvent::Completed {
+                    response_id,
+                    token_usage,
+                }))) => {
+                    let mut emitted_any = false;
+
+                    if !this.cumulative_reasoning.is_empty() {
+                        let aggregated_reasoning = ResponseItem::Reasoning {
+                            id: String::new(),
+                            summary: Vec::new(),
+                            content: Some(vec![ReasoningItemContent::ReasoningText {
+                                text: std::mem::take(&mut this.cumulative_reasoning),
+                            }]),
+                            encrypted_content: None,
+                        };
+                        this.pending
+                            .push_back(ResponseEvent::OutputItemDone(aggregated_reasoning));
+                        emitted_any = true;
+                    }
+
+                    if !this.cumulative.is_empty() {
+                        let aggregated_message = ResponseItem::Message {
+                            id: None,
+                            role: "assistant".to_string(),
+                            content: vec![ContentItem::OutputText {
+                                text: std::mem::take(&mut this.cumulative),
+                            }],
+                        };
+                        this.pending
+                            .push_back(ResponseEvent::OutputItemDone(aggregated_message));
+                        emitted_any = true;
+                    }
+
+                    if emitted_any {
+                        this.pending.push_back(ResponseEvent::Completed {
+                            response_id: response_id.clone(),
+                            token_usage: token_usage.clone(),
+                        });
+                        if let Some(ev) = this.pending.pop_front() {
+                            return Poll::Ready(Some(Ok(ev)));
+                        }
+                    }
+
+                    return Poll::Ready(Some(Ok(ResponseEvent::Completed {
+                        response_id,
+                        token_usage,
+                    })));
+                }
+                Poll::Ready(Some(Ok(ResponseEvent::Created))) => {
+                    continue;
+                }
+                Poll::Ready(Some(Ok(ResponseEvent::OutputTextDelta(delta)))) => {
+                    this.cumulative.push_str(&delta);
+                    if matches!(this.mode, AggregateMode::Streaming) {
+                        return Poll::Ready(Some(Ok(ResponseEvent::OutputTextDelta(delta))));
+                    } else {
+                        continue;
+                    }
+                }
+                Poll::Ready(Some(Ok(ResponseEvent::ReasoningContentDelta {
+                    delta,
+                    content_index,
+                }))) => {
+                    this.cumulative_reasoning.push_str(&delta);
+                    if matches!(this.mode, AggregateMode::Streaming) {
+                        return Poll::Ready(Some(Ok(ResponseEvent::ReasoningContentDelta {
+                            delta,
+                            content_index,
+                        })));
+                    } else {
+                        continue;
+                    }
+                }
+                Poll::Ready(Some(Ok(ResponseEvent::ReasoningSummaryDelta { .. }))) => continue,
+                Poll::Ready(Some(Ok(ResponseEvent::ReasoningSummaryPartAdded { .. }))) => {
+                    continue;
+                }
+                Poll::Ready(Some(Ok(ResponseEvent::OutputItemAdded(item)))) => {
+                    return Poll::Ready(Some(Ok(ResponseEvent::OutputItemAdded(item))));
+                }
+            }
+        }
+    }
+}
+
+pub trait AggregateStreamExt {
+    fn aggregate(self) -> AggregatedStream;
+
+    fn streaming_mode(self) -> ResponseStream;
+}
+
+impl AggregateStreamExt for ResponseStream {
+    fn aggregate(self) -> AggregatedStream {
+        AggregatedStream::new(self, AggregateMode::AggregatedOnly)
+    }
+
+    fn streaming_mode(self) -> ResponseStream {
+        self
+    }
+}
+
+impl AggregatedStream {
+    fn new(inner: ResponseStream, mode: AggregateMode) -> Self {
+        AggregatedStream {
+            inner,
+            cumulative: String::new(),
+            cumulative_reasoning: String::new(),
+            pending: VecDeque::new(),
+            mode,
+        }
+    }
+}

codex-rs/codex-api/src/endpoint/compact.rs

@@ -0,0 +1,162 @@
+use crate::auth::AuthProvider;
+use crate::auth::add_auth_headers;
+use crate::common::CompactionInput;
+use crate::error::ApiError;
+use crate::provider::Provider;
+use crate::provider::WireApi;
+use crate::telemetry::run_with_request_telemetry;
+use codex_client::HttpTransport;
+use codex_client::RequestTelemetry;
+use codex_protocol::models::ResponseItem;
+use http::HeaderMap;
+use http::Method;
+use serde::Deserialize;
+use serde_json::to_value;
+use std::sync::Arc;
+
+pub struct CompactClient<T: HttpTransport, A: AuthProvider> {
+    transport: T,
+    provider: Provider,
+    auth: A,
+    request_telemetry: Option<Arc<dyn RequestTelemetry>>,
+}
+
+impl<T: HttpTransport, A: AuthProvider> CompactClient<T, A> {
+    pub fn new(transport: T, provider: Provider, auth: A) -> Self {
+        Self {
+            transport,
+            provider,
+            auth,
+            request_telemetry: None,
+        }
+    }
+
+    pub fn with_telemetry(mut self, request: Option<Arc<dyn RequestTelemetry>>) -> Self {
+        self.request_telemetry = request;
+        self
+    }
+
+    fn path(&self) -> Result<&'static str, ApiError> {
+        match self.provider.wire {
+            WireApi::Compact | WireApi::Responses => Ok("responses/compact"),
+            WireApi::Chat => Err(ApiError::Stream(
+                "compact endpoint requires responses wire api".to_string(),
+            )),
+        }
+    }
+
+    pub async fn compact(
+        &self,
+        body: serde_json::Value,
+        extra_headers: HeaderMap,
+    ) -> Result<Vec<ResponseItem>, ApiError> {
+        let path = self.path()?;
+        let builder = || {
+            let mut req = self.provider.build_request(Method::POST, path);
+            req.headers.extend(extra_headers.clone());
+            req.body = Some(body.clone());
+            add_auth_headers(&self.auth, req)
+        };
+
+        let resp = run_with_request_telemetry(
+            self.provider.retry.to_policy(),
+            self.request_telemetry.clone(),
+            builder,
+            |req| self.transport.execute(req),
+        )
+        .await?;
+        let parsed: CompactHistoryResponse =
+            serde_json::from_slice(&resp.body).map_err(|e| ApiError::Stream(e.to_string()))?;
+        Ok(parsed.output)
+    }
+
+    pub async fn compact_input(
+        &self,
+        input: &CompactionInput<'_>,
+        extra_headers: HeaderMap,
+    ) -> Result<Vec<ResponseItem>, ApiError> {
+        let body = to_value(input)
+            .map_err(|e| ApiError::Stream(format!("failed to encode compaction input: {e}")))?;
+        self.compact(body, extra_headers).await
+    }
+}
+
+#[derive(Debug, Deserialize)]
+struct CompactHistoryResponse {
+    output: Vec<ResponseItem>,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::provider::RetryConfig;
+    use async_trait::async_trait;
+    use codex_client::Request;
+    use codex_client::Response;
+    use codex_client::StreamResponse;
+    use codex_client::TransportError;
+    use http::HeaderMap;
+    use std::time::Duration;
+
+    #[derive(Clone, Default)]
+    struct DummyTransport;
+
+    #[async_trait]
+    impl HttpTransport for DummyTransport {
+        async fn execute(&self, _req: Request) -> Result<Response, TransportError> {
+            Err(TransportError::Build("execute should not run".to_string()))
+        }
+
+        async fn stream(&self, _req: Request) -> Result<StreamResponse, TransportError> {
+            Err(TransportError::Build("stream should not run".to_string()))
+        }
+    }
+
+    #[derive(Clone, Default)]
+    struct DummyAuth;
+
+    impl AuthProvider for DummyAuth {
+        fn bearer_token(&self) -> Option<String> {
+            None
+        }
+    }
+
+    fn provider(wire: WireApi) -> Provider {
+        Provider {
+            name: "test".to_string(),
+            base_url: "https://example.com/v1".to_string(),
+            query_params: None,
+            wire,
+            headers: HeaderMap::new(),
+            retry: RetryConfig {
+                max_attempts: 1,
+                base_delay: Duration::from_millis(1),
+                retry_429: false,
+                retry_5xx: true,
+                retry_transport: true,
+            },
+            stream_idle_timeout: Duration::from_secs(1),
+        }
+    }
+
+    #[tokio::test]
+    async fn errors_when_wire_is_chat() {
+        let client = CompactClient::new(DummyTransport, provider(WireApi::Chat), DummyAuth);
+        let input = CompactionInput {
+            model: "gpt-test",
+            input: &[],
+            instructions: "inst",
+        };
+        let err = client
+            .compact_input(&input, HeaderMap::new())
+            .await
+            .expect_err("expected wire mismatch to fail");
+
+        match err {
+            ApiError::Stream(msg) => {
+                assert_eq!(msg, "compact endpoint requires responses wire api");
+            }
+            other => panic!("unexpected error: {other:?}"),
+        }
+    }
+}

codex-rs/codex-api/src/endpoint/mod.rs

@@ -0,0 +1,4 @@
+pub mod chat;
+pub mod compact;
+pub mod responses;
+mod streaming;

codex-rs/codex-api/src/endpoint/responses.rs

@@ -0,0 +1,107 @@
+use crate::auth::AuthProvider;
+use crate::common::Prompt as ApiPrompt;
+use crate::common::Reasoning;
+use crate::common::ResponseStream;
+use crate::common::TextControls;
+use crate::endpoint::streaming::StreamingClient;
+use crate::error::ApiError;
+use crate::provider::Provider;
+use crate::provider::WireApi;
+use crate::requests::ResponsesRequest;
+use crate::requests::ResponsesRequestBuilder;
+use crate::sse::spawn_response_stream;
+use crate::telemetry::SseTelemetry;
+use codex_client::HttpTransport;
+use codex_client::RequestTelemetry;
+use codex_protocol::protocol::SessionSource;
+use http::HeaderMap;
+use serde_json::Value;
+use std::sync::Arc;
+
+pub struct ResponsesClient<T: HttpTransport, A: AuthProvider> {
+    streaming: StreamingClient<T, A>,
+}
+
+#[derive(Default)]
+pub struct ResponsesOptions {
+    pub reasoning: Option<Reasoning>,
+    pub include: Vec<String>,
+    pub prompt_cache_key: Option<String>,
+    pub text: Option<TextControls>,
+    pub store_override: Option<bool>,
+    pub conversation_id: Option<String>,
+    pub session_source: Option<SessionSource>,
+}
+
+impl<T: HttpTransport, A: AuthProvider> ResponsesClient<T, A> {
+    pub fn new(transport: T, provider: Provider, auth: A) -> Self {
+        Self {
+            streaming: StreamingClient::new(transport, provider, auth),
+        }
+    }
+
+    pub fn with_telemetry(
+        self,
+        request: Option<Arc<dyn RequestTelemetry>>,
+        sse: Option<Arc<dyn SseTelemetry>>,
+    ) -> Self {
+        Self {
+            streaming: self.streaming.with_telemetry(request, sse),
+        }
+    }
+
+    pub async fn stream_request(
+        &self,
+        request: ResponsesRequest,
+    ) -> Result<ResponseStream, ApiError> {
+        self.stream(request.body, request.headers).await
+    }
+
+    pub async fn stream_prompt(
+        &self,
+        model: &str,
+        prompt: &ApiPrompt,
+        options: ResponsesOptions,
+    ) -> Result<ResponseStream, ApiError> {
+        let ResponsesOptions {
+            reasoning,
+            include,
+            prompt_cache_key,
+            text,
+            store_override,
+            conversation_id,
+            session_source,
+        } = options;
+
+        let request = ResponsesRequestBuilder::new(model, &prompt.instructions, &prompt.input)
+            .tools(&prompt.tools)
+            .parallel_tool_calls(prompt.parallel_tool_calls)
+            .reasoning(reasoning)
+            .include(include)
+            .prompt_cache_key(prompt_cache_key)
+            .text(text)
+            .conversation(conversation_id)
+            .session_source(session_source)
+            .store_override(store_override)
+            .build(self.streaming.provider())?;
+
+        self.stream_request(request).await
+    }
+
+    fn path(&self) -> &'static str {
+        match self.streaming.provider().wire {
+            WireApi::Responses | WireApi::Compact => "responses",
+            WireApi::Chat => "chat/completions",
+        }
+    }
+
+    pub async fn stream(
+        &self,
+        body: Value,
+        extra_headers: HeaderMap,
+    ) -> Result<ResponseStream, ApiError> {
+        self.streaming
+            .stream(self.path(), body, extra_headers, spawn_response_stream)
+            .await
+    }
+}

codex-rs/codex-api/src/endpoint/streaming.rs

@@ -0,0 +1,82 @@
+use crate::auth::AuthProvider;
+use crate::auth::add_auth_headers;
+use crate::common::ResponseStream;
+use crate::error::ApiError;
+use crate::provider::Provider;
+use crate::telemetry::SseTelemetry;
+use crate::telemetry::run_with_request_telemetry;
+use codex_client::HttpTransport;
+use codex_client::RequestTelemetry;
+use codex_client::StreamResponse;
+use http::HeaderMap;
+use http::Method;
+use serde_json::Value;
+use std::sync::Arc;
+use std::time::Duration;
+
+pub(crate) struct StreamingClient<T: HttpTransport, A: AuthProvider> {
+    transport: T,
+    provider: Provider,
+    auth: A,
+    request_telemetry: Option<Arc<dyn RequestTelemetry>>,
+    sse_telemetry: Option<Arc<dyn SseTelemetry>>,
+}
+
+impl<T: HttpTransport, A: AuthProvider> StreamingClient<T, A> {
+    pub(crate) fn new(transport: T, provider: Provider, auth: A) -> Self {
+        Self {
+            transport,
+            provider,
+            auth,
+            request_telemetry: None,
+            sse_telemetry: None,
+        }
+    }
+
+    pub(crate) fn with_telemetry(
+        mut self,
+        request: Option<Arc<dyn RequestTelemetry>>,
+        sse: Option<Arc<dyn SseTelemetry>>,
+    ) -> Self {
+        self.request_telemetry = request;
+        self.sse_telemetry = sse;
+        self
+    }
+
+    pub(crate) fn provider(&self) -> &Provider {
+        &self.provider
+    }
+
+    pub(crate) async fn stream(
+        &self,
+        path: &str,
+        body: Value,
+        extra_headers: HeaderMap,
+        spawner: fn(StreamResponse, Duration, Option<Arc<dyn SseTelemetry>>) -> ResponseStream,
+    ) -> Result<ResponseStream, ApiError> {
+        let builder = || {
+            let mut req = self.provider.build_request(Method::POST, path);
+            req.headers.extend(extra_headers.clone());
+            req.headers.insert(
+                http::header::ACCEPT,
+                http::HeaderValue::from_static("text/event-stream"),
+            );
+            req.body = Some(body.clone());
+            add_auth_headers(&self.auth, req)
+        };
+
+        let stream_response = run_with_request_telemetry(
+            self.provider.retry.to_policy(),
+            self.request_telemetry.clone(),
+            builder,
+            |req| self.transport.stream(req),
+        )
+        .await?;
+
+        Ok(spawner(
+            stream_response,
+            self.provider.stream_idle_timeout,
+            self.sse_telemetry.clone(),
+        ))
+    }
+}