diff --git a/codex-rs/core/src/chat_completions.rs b/codex-rs/core/src/client/chat_completions.rs similarity index 53% rename from codex-rs/core/src/chat_completions.rs rename to codex-rs/core/src/client/chat_completions.rs index d8d3a63d13..a233444911 100644 --- a/codex-rs/core/src/chat_completions.rs +++ b/codex-rs/core/src/client/chat_completions.rs @@ -1,15 +1,14 @@ use std::time::Duration; use crate::ModelProviderInfo; +use crate::client::ResponseEvent; +use crate::client::ResponseStream; +use crate::client::http::CodexHttpClient; use crate::client_common::Prompt; -use crate::client_common::ResponseEvent; -use crate::client_common::ResponseStream; -use crate::default_client::CodexHttpClient; use crate::error::CodexErr; use crate::error::ConnectionFailedError; use crate::error::ResponseStreamFailed; use crate::error::Result; -use crate::error::RetryLimitReachedError; use crate::error::UnexpectedResponseError; use crate::model_family::ModelFamily; use crate::tools::spec::create_tools_json_for_chat_completions_api; @@ -55,7 +54,7 @@ pub(crate) async fn stream_chat_completions( let mut messages = Vec::::new(); let full_instructions = prompt.get_full_instructions(model_family); - messages.push(json!({"role": "system", "content": full_instructions})); + messages.push(json!({ "role": "system", "content": full_instructions })); let input = prompt.get_formatted_input(); @@ -128,99 +127,73 @@ pub(crate) async fn stream_chat_completions( { reasoning_by_anchor_index .entry(idx - 1) - .and_modify(|v| v.push_str(&text)) + .and_modify(|existing| existing.push_str(text.as_str())) .or_insert(text.clone()); attached = true; } - // Otherwise, attach to immediate next assistant anchor (tool-calls or assistant message) - if !attached && idx + 1 < input.len() { - match &input[idx + 1] { - ResponseItem::FunctionCall { .. } | ResponseItem::LocalShellCall { .. } => { - reasoning_by_anchor_index - .entry(idx + 1) - .and_modify(|v| v.push_str(&text)) - .or_insert(text.clone()); + // Otherwise, attach to the first future assistant anchor. + if !attached { + for anchor_idx in idx + 1..input.len() { + match &input[anchor_idx] { + ResponseItem::Message { role, .. } if role == "assistant" => { + reasoning_by_anchor_index + .entry(anchor_idx) + .and_modify(|existing| existing.push_str(text.as_str())) + .or_insert(text.clone()); + attached = true; + break; + } + ResponseItem::FunctionCall { .. } + | ResponseItem::LocalShellCall { .. } + | ResponseItem::FunctionCallOutput { .. } => { + continue; + } + _ => break, } - ResponseItem::Message { role, .. } if role == "assistant" => { - reasoning_by_anchor_index - .entry(idx + 1) - .and_modify(|v| v.push_str(&text)) - .or_insert(text.clone()); - } - _ => {} } } + + // Either attached or dropped, move on. } } } - // Track last assistant text we emitted to avoid duplicate assistant messages - // in the outbound Chat Completions payload (can happen if a final - // aggregated assistant message was recorded alongside an earlier partial). - let mut last_assistant_text: Option = None; - - for (idx, item) in input.iter().enumerate() { + for (index, item) in input.iter().enumerate() { match item { ResponseItem::Message { role, content, .. } => { - // Build content either as a plain string (typical for assistant text) - // or as an array of content items when images are present (user/tool multimodal). - let mut text = String::new(); - let mut items: Vec = Vec::new(); - let mut saw_image = false; - - for c in content { - match c { - ContentItem::InputText { text: t } - | ContentItem::OutputText { text: t } => { - text.push_str(t); - items.push(json!({"type":"text","text": t})); - } - ContentItem::InputImage { image_url } => { - saw_image = true; - items.push(json!({"type":"image_url","image_url": {"url": image_url}})); + let mut content_text = String::new(); + for item in content { + match item { + ContentItem::InputText { text } | ContentItem::OutputText { text } => { + content_text.push_str(text) } + ContentItem::InputImage { .. } => {} } } - - // Skip exact-duplicate assistant messages. - if role == "assistant" { - if let Some(prev) = &last_assistant_text - && prev == &text - { - continue; - } - last_assistant_text = Some(text.clone()); + if content_text.trim().is_empty() { + continue; } - // For assistant messages, always send a plain string for compatibility. - // For user messages, if an image is present, send an array of content items. - let content_value = if role == "assistant" { - json!(text) - } else if saw_image { - json!(items) - } else { - json!(text) - }; - - let mut msg = json!({"role": role, "content": content_value}); - if role == "assistant" - && let Some(reasoning) = reasoning_by_anchor_index.get(&idx) - && let Some(obj) = msg.as_object_mut() - { - obj.insert("reasoning".to_string(), json!(reasoning)); + // Append reasoning when mapped to this anchor. + if let Some(reasoning) = reasoning_by_anchor_index.remove(&index) { + content_text.push_str(reasoning.as_str()); } - messages.push(msg); + + messages.push(json!({ + "role": role, + "content": content_text, + })); } + ResponseItem::FunctionCall { name, arguments, call_id, .. } => { - let mut msg = json!({ + messages.push(json!({ "role": "assistant", - "content": null, "tool_calls": [{ "id": call_id, "type": "function", @@ -228,39 +201,10 @@ pub(crate) async fn stream_chat_completions( "name": name, "arguments": arguments, } - }] - }); - if let Some(reasoning) = reasoning_by_anchor_index.get(&idx) - && let Some(obj) = msg.as_object_mut() - { - obj.insert("reasoning".to_string(), json!(reasoning)); - } - messages.push(msg); - } - ResponseItem::LocalShellCall { - id, - call_id: _, - status, - action, - } => { - // Confirm with API team. - let mut msg = json!({ - "role": "assistant", - "content": null, - "tool_calls": [{ - "id": id.clone().unwrap_or_else(|| "".to_string()), - "type": "local_shell_call", - "status": status, - "action": action, - }] - }); - if let Some(reasoning) = reasoning_by_anchor_index.get(&idx) - && let Some(obj) = msg.as_object_mut() - { - obj.insert("reasoning".to_string(), json!(reasoning)); - } - messages.push(msg); + }], + })); } + ResponseItem::FunctionCallOutput { call_id, output } => { // Prefer structured content items when available (e.g., images) // otherwise fall back to the legacy plain-string content. @@ -286,27 +230,43 @@ pub(crate) async fn stream_chat_completions( "tool_call_id": call_id, "content": content_value, })); + + if let Some(reasoning) = reasoning_by_anchor_index.remove(&index) { + messages.push(json!({ + "role": "assistant", + "content": reasoning, + })); + } } - ResponseItem::CustomToolCall { - id, - call_id: _, - name, - input, - status: _, - } => { + + ResponseItem::LocalShellCall { call_id, .. } => { messages.push(json!({ "role": "assistant", - "content": null, "tool_calls": [{ - "id": id, - "type": "custom", - "custom": { - "name": name, - "input": input, + "id": call_id, + "type": "function", + "function": { + "name": "shell", + "arguments": "{}", // arguments are defined via `input` only. } - }] + }], })); } + + ResponseItem::CustomToolCall { call_id, name, .. } => { + messages.push(json!({ + "role": "assistant", + "tool_calls": [{ + "id": call_id, + "type": "function", + "function": { + "name": name, + "arguments": "{}", // arguments are defined via `input` only. + } + }], + })); + } + ResponseItem::CustomToolCallOutput { call_id, output } => { messages.push(json!({ "role": "tool", @@ -314,118 +274,159 @@ pub(crate) async fn stream_chat_completions( "content": output, })); } - ResponseItem::GhostSnapshot { .. } => { - // Ghost snapshots annotate history but are not sent to the model. - continue; - } - ResponseItem::Reasoning { .. } - | ResponseItem::WebSearchCall { .. } - | ResponseItem::Other => { - // Omit these items from the conversation history. - continue; + + ResponseItem::Reasoning { .. } => { + // Reasoning is mapped onto adjacent assistant anchors above. } + + ResponseItem::WebSearchCall { .. } => {} + + ResponseItem::Other { .. } => {} + + ResponseItem::GhostSnapshot { .. } => {} + } + } + + // Attach any reasoning still not mapped (e.g., if the last input items are Reasoning). + if messages.len() == 1 { + if let Some(text) = reasoning_by_anchor_index.remove(&0) { + messages.push(json!({ + "role": "assistant", + "content": text, + })); } } let tools_json = create_tools_json_for_chat_completions_api(&prompt.tools)?; - let payload = json!({ + + let mut body = json!({ "model": model_family.slug, "messages": messages, "stream": true, - "tools": tools_json, + "stream_options": { + "include_usage": true, + }, }); - debug!( + if !tools_json.is_empty() { + body["tools"] = json!(tools_json); + body["tool_choice"] = json!("auto"); + } + + if let 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()) + }; + body["metadata"] = json!({ + "x-openai-subagent": subagent, + }); + } + + let max_attempts = provider.request_max_retries(); + let mut last_error = None; + for attempt in 0..=max_attempts { + match stream_single_chat_completion( + attempt, + client, + provider, + otel_event_manager, + body.clone(), + ) + .await + { + Ok(stream) => return Ok(stream), + Err(e) => { + last_error = Some(e); + if attempt != max_attempts { + tokio::time::sleep(backoff(attempt)).await; + } + } + } + } + + Err(last_error.unwrap_or(CodexErr::InternalServerError)) +} + +async fn stream_single_chat_completion( + attempt: u64, + client: &CodexHttpClient, + provider: &ModelProviderInfo, + otel_event_manager: &OtelEventManager, + body: serde_json::Value, +) -> Result { + trace!( "POST to {}: {}", provider.get_full_url(&None), - payload.to_string() + body.to_string() ); - let mut attempt = 0; - let max_retries = provider.request_max_retries(); - loop { - attempt += 1; + let mut req_builder = provider.create_request_builder(client, &None).await?; + req_builder = req_builder + .header(reqwest::header::ACCEPT, "text/event-stream") + .json(&body); - let mut req_builder = provider.create_request_builder(client, &None).await?; + let res = otel_event_manager + .log_request(attempt, || req_builder.send()) + .await; - // Include subagent header only for subagent sessions. - if let SessionSource::SubAgent(sub) = session_source.clone() { - let subagent = if let SubAgentSource::Other(label) = sub { - label - } 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()) - }; - req_builder = req_builder.header("x-openai-subagent", subagent); + let mut request_id = None; + if let Ok(resp) = &res { + request_id = resp + .headers() + .get("cf-ray") + .map(|v| v.to_str().unwrap_or_default().to_string()); + } + + match res { + Ok(resp) if resp.status().is_success() => { + let (tx_event, rx_event) = mpsc::channel::>(1600); + + // spawn task to process SSE + let stream = resp.bytes_stream().map_err(move |e| { + CodexErr::ResponseStreamFailed(ResponseStreamFailed { + source: e, + request_id: request_id.clone(), + }) + }); + tokio::spawn(process_chat_sse( + stream, + tx_event, + provider.stream_idle_timeout(), + otel_event_manager.clone(), + )); + + Ok(ResponseStream { rx_event }) } + Ok(res) => { + let status = res.status(); - let res = otel_event_manager - .log_request(attempt, || { - req_builder - .header(reqwest::header::ACCEPT, "text/event-stream") - .json(&payload) - .send() - }) - .await; - - match res { - Ok(resp) if resp.status().is_success() => { - let (tx_event, rx_event) = mpsc::channel::>(1600); - let stream = resp.bytes_stream().map_err(|e| { - CodexErr::ResponseStreamFailed(ResponseStreamFailed { - source: e, - request_id: None, - }) - }); - tokio::spawn(process_chat_sse( - stream, - tx_event, - provider.stream_idle_timeout(), - otel_event_manager.clone(), - )); - return Ok(ResponseStream { rx_event }); + if !(status == StatusCode::TOO_MANY_REQUESTS + || status == StatusCode::UNAUTHORIZED + || status.is_server_error()) + { + // Surface the error body to callers. Use `unwrap_or_default` per Clippy. + let body = res.text().await.unwrap_or_default(); + return Err(CodexErr::UnexpectedStatus(UnexpectedResponseError { + status, + body, + request_id: None, + })); } - Ok(res) => { - let status = res.status(); - if !(status == StatusCode::TOO_MANY_REQUESTS || status.is_server_error()) { - let body = (res.text().await).unwrap_or_default(); - return Err(CodexErr::UnexpectedStatus(UnexpectedResponseError { - status, - body, - request_id: None, - })); - } - if attempt > max_retries { - return Err(CodexErr::RetryLimit(RetryLimitReachedError { - status, - request_id: None, - })); - } - - let retry_after_secs = res - .headers() - .get(reqwest::header::RETRY_AFTER) - .and_then(|v| v.to_str().ok()) - .and_then(|s| s.parse::().ok()); - - let delay = retry_after_secs - .map(|s| Duration::from_millis(s * 1_000)) - .unwrap_or_else(|| backoff(attempt)); - tokio::time::sleep(delay).await; - } - Err(e) => { - if attempt > max_retries { - return Err(CodexErr::ConnectionFailed(ConnectionFailedError { - source: e, - })); - } - let delay = backoff(attempt); - tokio::time::sleep(delay).await; - } + Err(CodexErr::UnexpectedStatus(UnexpectedResponseError { + status, + body: String::new(), + request_id, + })) } + Err(e) => Err(CodexErr::ConnectionFailed(ConnectionFailedError { + source: e, + })), } } @@ -484,9 +485,7 @@ async fn append_reasoning_text( .await; } } -/// Lightweight SSE processor for the Chat Completions streaming format. The -/// output is mapped onto Codex's internal [`ResponseEvent`] so that the rest -/// of the pipeline can stay agnostic of the underlying wire format. + async fn process_chat_sse( stream: S, tx_event: mpsc::Sender>, @@ -719,33 +718,28 @@ async fn process_chat_sse( } } -/// Optional client-side aggregation helper -/// -/// Stream adapter that merges the incremental `OutputItemDone` chunks coming from -/// [`process_chat_sse`] into a *running* assistant message, **suppressing the -/// per-token deltas**. The stream stays silent while the model is thinking -/// and only emits two events per turn: -/// -/// 1. `ResponseEvent::OutputItemDone` with the *complete* assistant message -/// (fully concatenated). -/// 2. The original `ResponseEvent::Completed` right after it. -/// -/// This mirrors the behaviour the TypeScript CLI exposes to its higher layers. -/// -/// The adapter is intentionally *lossless*: callers who do **not** opt in via -/// [`AggregateStreamExt::aggregate()`] keep receiving the original unmodified -/// events. -#[derive(Copy, Clone, Eq, PartialEq)] -enum AggregateMode { - AggregatedOnly, - Streaming, -} -pub(crate) struct AggregatedChatStream { +/// Adapter that aggregates Chat Completions SSE output into the final assistant +/// message plus optional reasoning, mirroring the Responses API contract. +pub(crate) struct AggregatedChatStream +where + S: Stream> + Unpin, +{ inner: S, - cumulative: String, - cumulative_reasoning: String, pending: std::collections::VecDeque, - mode: AggregateMode, + finished: bool, +} + +impl AggregatedChatStream +where + S: Stream> + Unpin, +{ + pub fn streaming_mode(inner: S) -> Self { + Self { + inner, + pending: std::collections::VecDeque::new(), + finished: false, + } + } } impl Stream for AggregatedChatStream @@ -757,165 +751,44 @@ where fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let this = self.get_mut(); - // First, flush any buffered events from the previous call. - if let Some(ev) = this.pending.pop_front() { - return Poll::Ready(Some(Ok(ev))); + if let Some(event) = this.pending.pop_front() { + return Poll::Ready(Some(Ok(event))); + } + + if this.finished { + return Poll::Ready(None); } 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)))) => { - // If this is an incremental assistant message chunk, accumulate but - // do NOT emit yet. Forward any other item (e.g. FunctionCall) right - // away so downstream consumers see it. - - let is_assistant_message = matches!( - &item, - codex_protocol::models::ResponseItem::Message { role, .. } if role == "assistant" - ); - - if is_assistant_message { - match this.mode { - AggregateMode::AggregatedOnly => { - // Only use the final assistant message if we have not - // seen any deltas; otherwise, deltas already built the - // cumulative text and this would duplicate it. - if this.cumulative.is_empty() - && let codex_protocol::models::ResponseItem::Message { - content, - .. - } = &item - && let Some(text) = content.iter().find_map(|c| match c { - codex_protocol::models::ContentItem::OutputText { - text, - } => Some(text), - _ => None, - }) - { - this.cumulative.push_str(text); - } - // Swallow assistant message here; emit on Completed. - continue; - } - AggregateMode::Streaming => { - // In streaming mode, if we have not seen any deltas, forward - // the final assistant message directly. If deltas were seen, - // suppress the final message to avoid duplication. - if this.cumulative.is_empty() { - return Poll::Ready(Some(Ok(ResponseEvent::OutputItemDone( - item, - )))); - } else { - continue; - } - } - } - } - - // Not an assistant message – forward immediately. - return Poll::Ready(Some(Ok(ResponseEvent::OutputItemDone(item)))); + this.pending.push_back(ResponseEvent::OutputItemDone(item)); + continue; } Poll::Ready(Some(Ok(ResponseEvent::RateLimits(snapshot)))) => { - return Poll::Ready(Some(Ok(ResponseEvent::RateLimits(snapshot)))); + this.pending.push_back(ResponseEvent::RateLimits(snapshot)); + continue; } Poll::Ready(Some(Ok(ResponseEvent::Completed { response_id, token_usage, }))) => { - // Build any aggregated items in the correct order: Reasoning first, then Message. - let mut emitted_any = false; - - if !this.cumulative_reasoning.is_empty() - && matches!(this.mode, AggregateMode::AggregatedOnly) - { - let aggregated_reasoning = - codex_protocol::models::ResponseItem::Reasoning { - id: String::new(), - summary: Vec::new(), - content: Some(vec![ - codex_protocol::models::ReasoningItemContent::ReasoningText { - text: std::mem::take(&mut this.cumulative_reasoning), - }, - ]), - encrypted_content: None, - }; - this.pending - .push_back(ResponseEvent::OutputItemDone(aggregated_reasoning)); - emitted_any = true; - } - - // Always emit the final aggregated assistant message when any - // content deltas have been observed. In AggregatedOnly mode this - // is the sole assistant output; in Streaming mode this finalizes - // the streamed deltas into a terminal OutputItemDone so callers - // can persist/render the message once per turn. - if !this.cumulative.is_empty() { - let aggregated_message = codex_protocol::models::ResponseItem::Message { - id: None, - role: "assistant".to_string(), - content: vec![codex_protocol::models::ContentItem::OutputText { - text: std::mem::take(&mut this.cumulative), - }], - }; - this.pending - .push_back(ResponseEvent::OutputItemDone(aggregated_message)); - emitted_any = true; - } - - // Always emit Completed last when anything was aggregated. - if emitted_any { - this.pending.push_back(ResponseEvent::Completed { - response_id: response_id.clone(), - token_usage: token_usage.clone(), - }); - // Return the first pending event now. - if let Some(ev) = this.pending.pop_front() { - return Poll::Ready(Some(Ok(ev))); - } - } - - // Nothing aggregated – forward Completed directly. - return Poll::Ready(Some(Ok(ResponseEvent::Completed { + this.pending.push_back(ResponseEvent::Completed { response_id, token_usage, - }))); + }); + this.finished = true; + return Poll::Ready(this.pending.pop_front().map(Ok)); } - Poll::Ready(Some(Ok(ResponseEvent::Created))) => { - // These events are exclusive to the Responses API and - // will never appear in a Chat Completions stream. + Poll::Ready(Some(Ok(other))) => { + this.pending.push_back(other); continue; } - Poll::Ready(Some(Ok(ResponseEvent::OutputTextDelta(delta)))) => { - // Always accumulate deltas so we can emit a final OutputItemDone at Completed. - this.cumulative.push_str(&delta); - if matches!(this.mode, AggregateMode::Streaming) { - // In streaming mode, also forward the delta immediately. - return Poll::Ready(Some(Ok(ResponseEvent::OutputTextDelta(delta)))); - } else { - continue; - } - } - Poll::Ready(Some(Ok(ResponseEvent::ReasoningContentDelta(delta)))) => { - // Always accumulate reasoning deltas so we can emit a final Reasoning item at Completed. - this.cumulative_reasoning.push_str(&delta); - if matches!(this.mode, AggregateMode::Streaming) { - // In streaming mode, also forward the delta immediately. - return Poll::Ready(Some(Ok(ResponseEvent::ReasoningContentDelta(delta)))); - } 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)))); + Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(e))), + Poll::Ready(None) => { + this.finished = true; + return Poll::Ready(None); } } } @@ -923,45 +796,12 @@ where } /// Extension trait that activates aggregation on any stream of [`ResponseEvent`]. -pub(crate) trait AggregateStreamExt: Stream> + Sized { - /// Returns a new stream that emits **only** the final assistant message - /// per turn instead of every incremental delta. The produced - /// `ResponseEvent` sequence for a typical text turn looks like: - /// - /// ```ignore - /// OutputItemDone() - /// Completed - /// ``` - /// - /// No other `OutputItemDone` events will be seen by the caller. - /// - /// Usage: - /// - /// ```ignore - /// let agg_stream = client.stream(&prompt).await?.aggregate(); - /// while let Some(event) = agg_stream.next().await { - /// // event now contains cumulative text - /// } - /// ``` +pub(crate) trait AggregateStreamExt: + Stream> + Sized + Unpin +{ fn aggregate(self) -> AggregatedChatStream { - AggregatedChatStream::new(self, AggregateMode::AggregatedOnly) + AggregatedChatStream::streaming_mode(self) } } -impl AggregateStreamExt for T where T: Stream> + Sized {} - -impl AggregatedChatStream { - fn new(inner: S, mode: AggregateMode) -> Self { - AggregatedChatStream { - inner, - cumulative: String::new(), - cumulative_reasoning: String::new(), - pending: std::collections::VecDeque::new(), - mode, - } - } - - pub(crate) fn streaming_mode(inner: S) -> Self { - Self::new(inner, AggregateMode::Streaming) - } -} +impl AggregateStreamExt for T where T: Stream> + Sized + Unpin {} diff --git a/codex-rs/core/src/client/http.rs b/codex-rs/core/src/client/http.rs new file mode 100644 index 0000000000..580e82c2c3 --- /dev/null +++ b/codex-rs/core/src/client/http.rs @@ -0,0 +1,377 @@ +use crate::spawn::CODEX_SANDBOX_ENV_VAR; +use http::Error as HttpError; +use reqwest::IntoUrl; +use reqwest::Method; +use reqwest::Response; +use reqwest::header::HeaderName; +use reqwest::header::HeaderValue; +use serde::Serialize; +use std::collections::HashMap; +use std::fmt::Display; +use std::sync::LazyLock; +use std::sync::Mutex; +use std::sync::OnceLock; + +/// Set this to add a suffix to the User-Agent string. +/// +/// It is not ideal that we're using a global singleton for this. +/// This is primarily designed to differentiate MCP clients from each other. +/// Because there can only be one MCP server per process, it should be safe for this to be a global static. +/// However, future users of this should use this with caution as a result. +/// In addition, we want to be confident that this value is used for ALL clients and doing that requires a +/// lot of wiring and it's easy to miss code paths by doing so. +/// See https://github.com/openai/codex/pull/3388/files for an example of what that would look like. +/// Finally, we want to make sure this is set for ALL mcp clients without needing to know a special env var +/// or having to set data that they already specified in the mcp initialize request somewhere else. +/// +/// A space is automatically added between the suffix and the rest of the User-Agent string. +/// The full user agent string is returned from the mcp initialize response. +/// Parenthesis will be added by Codex. This should only specify what goes inside of the parenthesis. +pub static USER_AGENT_SUFFIX: LazyLock>> = LazyLock::new(|| Mutex::new(None)); +pub const DEFAULT_ORIGINATOR: &str = "codex_cli_rs"; +pub const CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR: &str = "CODEX_INTERNAL_ORIGINATOR_OVERRIDE"; + +#[derive(Clone, Debug)] +pub struct CodexHttpClient { + inner: reqwest::Client, +} + +impl CodexHttpClient { + fn new(inner: reqwest::Client) -> Self { + Self { inner } + } + + pub fn get(&self, url: U) -> CodexRequestBuilder + where + U: IntoUrl, + { + self.request(Method::GET, url) + } + + pub fn post(&self, url: U) -> CodexRequestBuilder + where + U: IntoUrl, + { + self.request(Method::POST, url) + } + + pub fn request(&self, method: Method, url: U) -> CodexRequestBuilder + where + U: IntoUrl, + { + let url_str = url.as_str().to_string(); + CodexRequestBuilder::new(self.inner.request(method.clone(), url), method, url_str) + } +} + +#[must_use = "requests are not sent unless `send` is awaited"] +#[derive(Debug)] +pub struct CodexRequestBuilder { + builder: reqwest::RequestBuilder, + method: Method, + url: String, +} + +impl CodexRequestBuilder { + fn new(builder: reqwest::RequestBuilder, method: Method, url: String) -> Self { + Self { + builder, + method, + url, + } + } + + fn map(self, f: impl FnOnce(reqwest::RequestBuilder) -> reqwest::RequestBuilder) -> Self { + Self { + builder: f(self.builder), + method: self.method, + url: self.url, + } + } + + pub fn header(self, key: K, value: V) -> Self + where + HeaderName: TryFrom, + >::Error: Into, + HeaderValue: TryFrom, + >::Error: Into, + { + self.map(|builder| builder.header(key, value)) + } + + pub fn bearer_auth(self, token: T) -> Self + where + T: Display, + { + self.map(|builder| builder.bearer_auth(token)) + } + + pub fn json(self, value: &T) -> Self + where + T: ?Sized + Serialize, + { + self.map(|builder| builder.json(value)) + } + + pub async fn send(self) -> Result { + match self.builder.send().await { + Ok(response) => { + let request_ids = Self::extract_request_ids(&response); + tracing::debug!( + method = %self.method, + url = %self.url, + status = %response.status(), + request_ids = ?request_ids, + version = ?response.version(), + "Request completed" + ); + + Ok(response) + } + Err(error) => { + let status = error.status(); + tracing::debug!( + method = %self.method, + url = %self.url, + status = status.map(|s| s.as_u16()), + error = %error, + "Request failed" + ); + Err(error) + } + } + } + + fn extract_request_ids(response: &Response) -> HashMap { + ["cf-ray", "x-request-id", "x-oai-request-id"] + .iter() + .filter_map(|&name| { + let header_name = HeaderName::from_static(name); + let value = response.headers().get(header_name)?; + let value = value.to_str().ok()?.to_owned(); + Some((name.to_owned(), value)) + }) + .collect() + } +} + +#[derive(Debug, Clone)] +pub struct Originator { + pub value: String, + pub header_value: HeaderValue, +} + +static ORIGINATOR: OnceLock = OnceLock::new(); + +#[derive(Debug)] +pub enum SetOriginatorError { + InvalidHeaderValue, + AlreadyInitialized, +} + +fn get_originator_value(provided: Option) -> Originator { + let value = std::env::var(CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR) + .ok() + .or(provided) + .unwrap_or(DEFAULT_ORIGINATOR.to_string()); + + match HeaderValue::from_str(&value) { + Ok(header_value) => Originator { + value, + header_value, + }, + Err(e) => { + tracing::error!("Unable to turn originator override {value} into header value: {e}"); + Originator { + value: DEFAULT_ORIGINATOR.to_string(), + header_value: HeaderValue::from_static(DEFAULT_ORIGINATOR), + } + } + } +} + +pub fn set_default_originator(value: String) -> Result<(), SetOriginatorError> { + let originator = get_originator_value(Some(value)); + ORIGINATOR + .set(originator) + .map_err(|_| SetOriginatorError::AlreadyInitialized) +} + +pub fn originator() -> &'static Originator { + ORIGINATOR.get_or_init(|| get_originator_value(None)) +} + +pub fn get_codex_user_agent() -> String { + let build_version = env!("CARGO_PKG_VERSION"); + let os_info = os_info::get(); + let prefix = format!( + "{}/{build_version} ({} {}; {}) {}", + originator().value.as_str(), + os_info.os_type(), + os_info.version(), + os_info.architecture().unwrap_or("unknown"), + crate::terminal::user_agent() + ); + let suffix = USER_AGENT_SUFFIX + .lock() + .ok() + .and_then(|guard| guard.clone()); + let suffix = suffix + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map_or_else(String::new, |value| format!(" ({value})")); + + let candidate = format!("{prefix}{suffix}"); + sanitize_user_agent(candidate, &prefix) +} + +/// Sanitize the user agent string. +/// +/// Invalid characters are replaced with an underscore. +/// +/// If the user agent fails to parse, it falls back to fallback and then to ORIGINATOR. +fn sanitize_user_agent(candidate: String, fallback: &str) -> String { + if HeaderValue::from_str(candidate.as_str()).is_ok() { + return candidate; + } + + let sanitized: String = candidate + .chars() + .map(|ch| if matches!(ch, ' '..='~') { ch } else { '_' }) + .collect(); + if !sanitized.is_empty() && HeaderValue::from_str(sanitized.as_str()).is_ok() { + tracing::warn!( + "Sanitized Codex user agent because provided suffix contained invalid header characters" + ); + sanitized + } else if HeaderValue::from_str(fallback).is_ok() { + tracing::warn!( + "Falling back to base Codex user agent because provided suffix could not be sanitized" + ); + fallback.to_string() + } else { + tracing::warn!( + "Falling back to default Codex originator because base user agent string is invalid" + ); + originator().value.clone() + } +} + +/// Create an HTTP client with default `originator` and `User-Agent` headers set. +pub fn create_client() -> CodexHttpClient { + use reqwest::header::HeaderMap; + + let mut headers = HeaderMap::new(); + headers.insert("originator", originator().header_value.clone()); + let ua = get_codex_user_agent(); + + let mut builder = reqwest::Client::builder() + // Set UA via dedicated helper to avoid header validation pitfalls + .user_agent(ua) + .default_headers(headers); + if is_sandboxed() { + builder = builder.no_proxy(); + } + + let inner = builder.build().unwrap_or_else(|_| reqwest::Client::new()); + CodexHttpClient::new(inner) +} + +fn is_sandboxed() -> bool { + std::env::var(CODEX_SANDBOX_ENV_VAR).as_deref() == Ok("seatbelt") +} + +#[cfg(test)] +mod tests { + use super::*; + use core_test_support::skip_if_no_network; + + #[test] + fn test_get_codex_user_agent() { + let user_agent = get_codex_user_agent(); + assert!(user_agent.starts_with("codex_cli_rs/")); + } + + #[tokio::test] + async fn test_create_client_sets_default_headers() { + skip_if_no_network!(); + + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::method; + use wiremock::matchers::path; + + let client = create_client(); + + // Spin up a local mock server and capture a request. + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server) + .await; + + let resp = client + .get(server.uri()) + .send() + .await + .expect("failed to send request"); + assert!(resp.status().is_success()); + + let requests = server + .received_requests() + .await + .expect("failed to fetch received requests"); + assert!(!requests.is_empty()); + let headers = &requests[0].headers; + + // originator header is set to the provided value + let originator_header = headers + .get("originator") + .expect("originator header missing"); + assert_eq!(originator_header.to_str().unwrap(), "codex_cli_rs"); + + // User-Agent matches the computed Codex UA for that originator + let expected_ua = get_codex_user_agent(); + let ua_header = headers + .get("user-agent") + .expect("user-agent header missing"); + assert_eq!(ua_header.to_str().unwrap(), expected_ua); + } + + #[test] + fn test_invalid_suffix_is_sanitized() { + let prefix = "codex_cli_rs/0.0.0"; + let suffix = "bad\rsuffix"; + + assert_eq!( + sanitize_user_agent(format!("{prefix} ({suffix})"), prefix), + "codex_cli_rs/0.0.0 (bad_suffix)" + ); + } + + #[test] + fn test_invalid_suffix_is_sanitized2() { + let prefix = "codex_cli_rs/0.0.0"; + let suffix = "bad\0suffix"; + + assert_eq!( + sanitize_user_agent(format!("{prefix} ({suffix})"), prefix), + "codex_cli_rs/0.0.0 (bad_suffix)" + ); + } + + #[test] + #[cfg(target_os = "macos")] + fn test_macos() { + use regex_lite::Regex; + let user_agent = get_codex_user_agent(); + let re = Regex::new( + r"^codex_cli_rs/\d+\.\d+\.\d+ \(Mac OS \d+\.\d+\.\d+; (x86_64|arm64)\) (\S+)$", + ) + .unwrap(); + assert!(re.is_match(&user_agent)); + } +} diff --git a/codex-rs/core/src/client/mod.rs b/codex-rs/core/src/client/mod.rs new file mode 100644 index 0000000000..b85b3e4dfa --- /dev/null +++ b/codex-rs/core/src/client/mod.rs @@ -0,0 +1,18 @@ +mod chat_completions; +pub mod http; +mod responses; +pub mod types; + +pub(crate) use chat_completions::AggregateStreamExt; +pub(crate) use chat_completions::AggregatedChatStream; +pub(crate) use chat_completions::stream_chat_completions; +pub use responses::ModelClient; +pub(crate) use types::FreeformTool; +pub(crate) use types::FreeformToolFormat; +pub(crate) use types::Reasoning; +pub use types::ResponseEvent; +pub use types::ResponseStream; +pub(crate) use types::ResponsesApiRequest; +pub(crate) use types::ResponsesApiTool; +pub(crate) use types::ToolSpec; +pub(crate) use types::create_text_param_for_request; diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client/responses.rs similarity index 51% rename from codex-rs/core/src/client.rs rename to codex-rs/core/src/client/responses.rs index d1ade21825..da87f51ac1 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client/responses.rs @@ -1,7 +1,6 @@ use std::io::BufRead; use std::path::Path; use std::sync::Arc; -use std::sync::OnceLock; use std::time::Duration; use bytes::Bytes; @@ -20,11 +19,9 @@ use regex_lite::Regex; use reqwest::StatusCode; use reqwest::header::HeaderMap; use serde::Deserialize; -use serde::Serialize; use serde_json::Value; use tokio::sync::mpsc; use tokio::time::timeout; -use tokio_util::io::ReaderStream; use tracing::debug; use tracing::trace; use tracing::warn; @@ -32,14 +29,13 @@ use tracing::warn; use crate::AuthManager; use crate::auth::CodexAuth; use crate::auth::RefreshTokenError; -use crate::chat_completions::AggregateStreamExt; -use crate::chat_completions::stream_chat_completions; +use crate::client::AggregateStreamExt; +use crate::client::Reasoning; +use crate::client::ResponseEvent; +use crate::client::ResponseStream; +use crate::client::ResponsesApiRequest; +use crate::client::create_text_param_for_request; use crate::client_common::Prompt; -use crate::client_common::Reasoning; -use crate::client_common::ResponseEvent; -use crate::client_common::ResponseStream; -use crate::client_common::ResponsesApiRequest; -use crate::client_common::create_text_param_for_request; use crate::config::Config; use crate::default_client::CodexHttpClient; use crate::default_client::create_client; @@ -47,7 +43,6 @@ use crate::error::CodexErr; use crate::error::ConnectionFailedError; use crate::error::ResponseStreamFailed; use crate::error::Result; -use crate::error::RetryLimitReachedError; use crate::error::UnexpectedResponseError; use crate::error::UsageLimitReachedError; use crate::flags::CODEX_RS_SSE_FIXTURE; @@ -145,7 +140,7 @@ impl ModelClient { WireApi::Responses => self.stream_responses(prompt).await, WireApi::Chat => { // Create the raw streaming connection first. - let response_stream = stream_chat_completions( + let response_stream = crate::client::stream_chat_completions( prompt, &self.config.model_family, &self.client, @@ -159,7 +154,7 @@ impl ModelClient { // the final assistant message per turn (matching the // behaviour of the Responses API). let mut aggregated = if self.config.show_raw_agent_reasoning { - crate::chat_completions::AggregatedChatStream::streaming_mode(response_stream) + crate::client::AggregatedChatStream::streaming_mode(response_stream) } else { response_stream.aggregate() }; @@ -185,11 +180,11 @@ impl ModelClient { /// Implementation for the OpenAI *Responses* experimental API. async fn stream_responses(&self, prompt: &Prompt) -> Result { - if let Some(path) = &*CODEX_RS_SSE_FIXTURE { + if let Some(path) = *CODEX_RS_SSE_FIXTURE { // short circuit for tests warn!(path, "Streaming from fixture"); return stream_from_fixture( - path, + Path::new(path), self.provider.clone(), self.otel_event_manager.clone(), ) @@ -509,581 +504,406 @@ impl ModelClient { } } +fn parse_rate_limit_snapshot(headers: &HeaderMap) -> Option { + let limit = headers.get("x-ratelimit-limit-requests")?; + let remaining = headers.get("x-ratelimit-remaining-requests")?; + let reset_ms = headers.get("x-ratelimit-reset-requests")?; + + let limit = limit.to_str().ok()?.parse::().ok()?; + let remaining = remaining.to_str().ok()?.parse::().ok()?; + let reset_ms = reset_ms.to_str().ok()?.parse::().ok()?; + + if limit <= 0.0 { + return None; + } + + let used = (limit - remaining).max(0.0); + let used_percent = (used / limit) * 100.0; + + let window_minutes = if reset_ms <= 0 { + None + } else { + let seconds = reset_ms / 1000; + Some((seconds + 59) / 60) + }; + + let resets_at = if reset_ms > 0 { + Some(Utc::now().timestamp() + reset_ms / 1000) + } else { + None + }; + + Some(RateLimitSnapshot { + primary: Some(RateLimitWindow { + used_percent, + window_minutes, + resets_at, + }), + secondary: None, + }) +} + +/// For Azure Responses endpoints we must use `store: true` and preserve +/// per-item identifiers on the input payload. The `ResponseItem` schema +/// deliberately skips serializing these IDs by default, so we patch them +/// back into the JSON body here based on the original input vector. +fn attach_item_ids(payload_json: &mut Value, original_input: &[ResponseItem]) { + let Some(input_json) = payload_json.get_mut("input").and_then(Value::as_array_mut) else { + return; + }; + + for (json_item, item) in input_json.iter_mut().zip(original_input.iter()) { + let Some(obj) = json_item.as_object_mut() else { + continue; + }; + + match item { + ResponseItem::Message { id: Some(id), .. } + | ResponseItem::LocalShellCall { id: Some(id), .. } + | ResponseItem::FunctionCall { id: Some(id), .. } + | ResponseItem::CustomToolCall { id: Some(id), .. } + | ResponseItem::WebSearchCall { id: Some(id), .. } => { + obj.insert("id".to_string(), Value::String(id.clone())); + } + ResponseItem::Reasoning { id, .. } if !id.is_empty() => { + obj.insert("id".to_string(), Value::String(id.clone())); + } + _ => {} + } + } +} + +fn try_parse_retry_after(error: &Error) -> Option { + let message = error.message.as_ref()?; + let re = Regex::new(r"Try again in (\d+)ms").ok()?; + let caps = re.captures(message)?; + let delay_ms = caps.get(1)?.as_str().parse::().ok()?; + Some(Duration::from_millis(delay_ms)) +} + +fn is_context_window_error(error: &Error) -> bool { + error + .r#type + .as_deref() + .map(|t| t == "context_length_exceeded") + .unwrap_or(false) +} + +fn is_quota_exceeded_error(error: &Error) -> bool { + if let Some(code) = error.code.as_deref() { + matches!( + code, + "insufficient_quota" + | "insufficient_quota_org" + | "insufficient_quota_project" + | "insufficient_quota_user" + ) + } else { + false + } +} + enum StreamAttemptError { + Fatal(CodexErr), RetryableHttpError { status: StatusCode, retry_after: Option, request_id: Option, }, RetryableTransportError(CodexErr), - Fatal(CodexErr), } impl StreamAttemptError { - /// attempt is 0-based. fn delay(&self, attempt: u64) -> Duration { - // backoff() uses 1-based attempts. - let backoff_attempt = attempt + 1; match self { - Self::RetryableHttpError { retry_after, .. } => { - retry_after.unwrap_or_else(|| backoff(backoff_attempt)) - } - Self::RetryableTransportError { .. } => backoff(backoff_attempt), - Self::Fatal(_) => { - // Should not be called on Fatal errors. - Duration::from_secs(0) + StreamAttemptError::RetryableHttpError { retry_after, .. } => { + retry_after.unwrap_or_else(|| backoff(attempt)) } + StreamAttemptError::RetryableTransportError(_) => backoff(attempt), + StreamAttemptError::Fatal(_) => Duration::from_secs(0), } } fn into_error(self) -> CodexErr { match self { - Self::RetryableHttpError { + StreamAttemptError::Fatal(e) => e, + StreamAttemptError::RetryableHttpError { status, request_id, .. - } => { - if status == StatusCode::INTERNAL_SERVER_ERROR { - CodexErr::InternalServerError - } else { - CodexErr::RetryLimit(RetryLimitReachedError { status, request_id }) - } - } - Self::RetryableTransportError(error) => error, - Self::Fatal(error) => error, + } => CodexErr::UnexpectedStatus(UnexpectedResponseError { + status, + body: String::new(), + request_id, + }), + StreamAttemptError::RetryableTransportError(e) => e, } } } -#[derive(Debug, Deserialize, Serialize)] -struct SseEvent { - #[serde(rename = "type")] - kind: String, - response: Option, +async fn process_sse( + stream: impl Stream> + Unpin + Send + 'static, + tx_event: mpsc::Sender>, + idle_timeout: Duration, + otel_event_manager: OtelEventManager, +) { + let mut stream = stream.eventsource(); + let mut response_completed: Option = None; + let mut response_error: Option = None; + + loop { + let start = tokio::time::Instant::now(); + let next_event = timeout(idle_timeout, stream.next()).await; + let duration = start.elapsed(); + otel_event_manager.log_sse_event(&next_event, duration); + + match next_event { + Ok(Some(Ok(ev))) => { + if let Err(e) = + handle_sse_event(ev, &mut response_completed, &mut response_error, &tx_event) + .await + { + let _ = tx_event.send(Err(e)).await; + break; + } + } + Ok(Some(Err(e))) => { + let _ = tx_event + .send(Err(CodexErr::Stream(e.to_string(), None))) + .await; + break; + } + Ok(None) => { + break; + } + Err(_) => { + let _ = tx_event + .send(Err(CodexErr::Stream( + "idle timeout waiting for SSE".to_string(), + None, + ))) + .await; + break; + } + } + } + + if let Some(err) = response_error { + let _ = tx_event.send(Err(err)).await; + return; + } + + if let Some(resp) = response_completed { + let _ = tx_event + .send(Ok(ResponseEvent::Completed { + response_id: resp.id, + token_usage: resp.usage, + })) + .await; + } else { + let _ = tx_event + .send(Err(CodexErr::Stream( + "stream closed before response.completed".to_string(), + None, + ))) + .await; + } +} + +#[derive(Debug, Deserialize)] +struct SseEventWrapper { + r#type: String, item: Option, + response: Option, delta: Option, } #[derive(Debug, Deserialize)] struct ResponseCompleted { id: String, - usage: Option, + #[serde(skip_serializing_if = "Option::is_none")] + usage: Option, } -#[derive(Debug, Deserialize)] -struct ResponseCompletedUsage { - input_tokens: i64, - input_tokens_details: Option, - output_tokens: i64, - output_tokens_details: Option, - total_tokens: i64, -} - -impl From for TokenUsage { - fn from(val: ResponseCompletedUsage) -> Self { - TokenUsage { - input_tokens: val.input_tokens, - cached_input_tokens: val - .input_tokens_details - .map(|d| d.cached_tokens) - .unwrap_or(0), - output_tokens: val.output_tokens, - reasoning_output_tokens: val - .output_tokens_details - .map(|d| d.reasoning_tokens) - .unwrap_or(0), - total_tokens: val.total_tokens, - } +async fn handle_sse_event( + ev: eventsource_stream::Event, + response_completed: &mut Option, + response_error: &mut Option, + tx_event: &mpsc::Sender>, +) -> Result<()> { + let data = ev.data; + if data == "[DONE]" { + // terminal event + return Ok(()); } -} -#[derive(Debug, Deserialize)] -struct ResponseCompletedInputTokensDetails { - cached_tokens: i64, -} - -#[derive(Debug, Deserialize)] -struct ResponseCompletedOutputTokensDetails { - reasoning_tokens: i64, -} - -fn attach_item_ids(payload_json: &mut Value, original_items: &[ResponseItem]) { - let Some(input_value) = payload_json.get_mut("input") else { - return; - }; - let serde_json::Value::Array(items) = input_value else { - return; - }; - - for (value, item) in items.iter_mut().zip(original_items.iter()) { - if let ResponseItem::Reasoning { id, .. } - | ResponseItem::Message { id: Some(id), .. } - | ResponseItem::WebSearchCall { id: Some(id), .. } - | ResponseItem::FunctionCall { id: Some(id), .. } - | ResponseItem::LocalShellCall { id: Some(id), .. } - | ResponseItem::CustomToolCall { id: Some(id), .. } = item - { - if id.is_empty() { - continue; - } - - if let Some(obj) = value.as_object_mut() { - obj.insert("id".to_string(), Value::String(id.clone())); - } - } - } -} - -fn parse_rate_limit_snapshot(headers: &HeaderMap) -> Option { - let primary = parse_rate_limit_window( - headers, - "x-codex-primary-used-percent", - "x-codex-primary-window-minutes", - "x-codex-primary-reset-at", - ); - - let secondary = parse_rate_limit_window( - headers, - "x-codex-secondary-used-percent", - "x-codex-secondary-window-minutes", - "x-codex-secondary-reset-at", - ); - - Some(RateLimitSnapshot { primary, secondary }) -} - -fn parse_rate_limit_window( - headers: &HeaderMap, - used_percent_header: &str, - window_minutes_header: &str, - resets_at_header: &str, -) -> Option { - let used_percent: Option = parse_header_f64(headers, used_percent_header); - - used_percent.and_then(|used_percent| { - let window_minutes = parse_header_i64(headers, window_minutes_header); - let resets_at = parse_header_i64(headers, resets_at_header); - - let has_data = used_percent != 0.0 - || window_minutes.is_some_and(|minutes| minutes != 0) - || resets_at.is_some(); - - has_data.then_some(RateLimitWindow { - used_percent, - window_minutes, - resets_at, - }) - }) -} - -fn parse_header_f64(headers: &HeaderMap, name: &str) -> Option { - parse_header_str(headers, name)? - .parse::() - .ok() - .filter(|v| v.is_finite()) -} - -fn parse_header_i64(headers: &HeaderMap, name: &str) -> Option { - parse_header_str(headers, name)?.parse::().ok() -} - -fn parse_header_str<'a>(headers: &'a HeaderMap, name: &str) -> Option<&'a str> { - headers.get(name)?.to_str().ok() -} - -async fn process_sse( - stream: S, - tx_event: mpsc::Sender>, - idle_timeout: Duration, - otel_event_manager: OtelEventManager, -) where - S: Stream> + Unpin, -{ - let mut stream = stream.eventsource(); - - // If the stream stays completely silent for an extended period treat it as disconnected. - // The response id returned from the "complete" message. - let mut response_completed: Option = None; - let mut response_error: Option = None; - - loop { - let start = std::time::Instant::now(); - let response = timeout(idle_timeout, stream.next()).await; - let duration = start.elapsed(); - otel_event_manager.log_sse_event(&response, duration); - - let sse = match response { - Ok(Some(Ok(sse))) => sse, - Ok(Some(Err(e))) => { - debug!("SSE Error: {e:#}"); - let event = CodexErr::Stream(e.to_string(), None); - let _ = tx_event.send(Err(event)).await; - return; - } - Ok(None) => { - match response_completed { - Some(ResponseCompleted { - id: response_id, - usage, - }) => { - if let Some(token_usage) = &usage { - otel_event_manager.sse_event_completed( - token_usage.input_tokens, - token_usage.output_tokens, - token_usage - .input_tokens_details - .as_ref() - .map(|d| d.cached_tokens), - token_usage - .output_tokens_details - .as_ref() - .map(|d| d.reasoning_tokens), - token_usage.total_tokens, - ); - } - let event = ResponseEvent::Completed { - response_id, - token_usage: usage.map(Into::into), - }; - let _ = tx_event.send(Ok(event)).await; + let event: SseEventWrapper = serde_json::from_str(&data)?; + match event.r#type.as_str() { + "response.completed" => { + if let Some(resp_val) = event.response { + match serde_json::from_value::(resp_val) { + Ok(r) => { + *response_completed = Some(r); } - None => { - let error = response_error.unwrap_or(CodexErr::Stream( - "stream closed before response.completed".into(), - None, - )); - otel_event_manager.see_event_completed_failed(&error); - - let _ = tx_event.send(Err(error)).await; + Err(e) => { + let error = format!("failed to parse ResponseCompleted: {e}"); + debug!(error); + *response_error = Some(CodexErr::Stream(error, None)); + return Ok(()); } - } - return; - } - Err(_) => { - let _ = tx_event - .send(Err(CodexErr::Stream( - "idle timeout waiting for SSE".into(), - None, - ))) - .await; - return; - } - }; - - let raw = sse.data.clone(); - trace!("SSE event: {}", raw); - - let event: SseEvent = match serde_json::from_str(&sse.data) { - Ok(event) => event, - Err(e) => { - debug!("Failed to parse SSE event: {e}, data: {}", &sse.data); - continue; - } - }; - - match event.kind.as_str() { - // Individual output item finalised. Forward immediately so the - // rest of the agent can stream assistant text/functions *live* - // instead of waiting for the final `response.completed` envelope. - // - // IMPORTANT: We used to ignore these events and forward the - // duplicated `output` array embedded in the `response.completed` - // payload. That produced two concrete issues: - // 1. No real‑time streaming – the user only saw output after the - // entire turn had finished, which broke the "typing" UX and - // made long‑running turns look stalled. - // 2. Duplicate `function_call_output` items – both the - // individual *and* the completed array were forwarded, which - // confused the backend and triggered 400 - // "previous_response_not_found" errors because the duplicated - // IDs did not match the incremental turn chain. - // - // The fix is to forward the incremental events *as they come* and - // drop the duplicated list inside `response.completed`. - "response.output_item.done" => { - let Some(item_val) = event.item else { continue }; - let Ok(item) = serde_json::from_value::(item_val) else { - debug!("failed to parse ResponseItem from output_item.done"); - continue; }; + }; + } + "response.output_item.done" => { + // For Responses API: + // - "response.output_item.done" contains the final item and we should + // drop the duplicated list inside `response.completed`. + let Some(item_val) = event.item else { + return Ok(()); + }; + let Ok(item) = serde_json::from_value::(item_val) else { + debug!("failed to parse ResponseItem from output_item.done"); + return Ok(()); + }; - let event = ResponseEvent::OutputItemDone(item); + let event = ResponseEvent::OutputItemDone(item); + if tx_event.send(Ok(event)).await.is_err() { + return Ok(()); + } + } + "response.output_text.delta" => { + if let Some(delta) = event.delta { + let event = ResponseEvent::OutputTextDelta(delta); if tx_event.send(Ok(event)).await.is_err() { - return; + return Ok(()); } } - "response.output_text.delta" => { - if let Some(delta) = event.delta { - let event = ResponseEvent::OutputTextDelta(delta); - if tx_event.send(Ok(event)).await.is_err() { - return; - } + } + "response.reasoning_summary_text.delta" => { + if let Some(delta) = event.delta { + let event = ResponseEvent::ReasoningSummaryDelta(delta); + if tx_event.send(Ok(event)).await.is_err() { + return Ok(()); } } - "response.reasoning_summary_text.delta" => { - if let Some(delta) = event.delta { - let event = ResponseEvent::ReasoningSummaryDelta(delta); - if tx_event.send(Ok(event)).await.is_err() { - return; - } + } + "response.reasoning_text.delta" => { + if let Some(delta) = event.delta { + let event = ResponseEvent::ReasoningContentDelta(delta); + if tx_event.send(Ok(event)).await.is_err() { + return Ok(()); } } - "response.reasoning_text.delta" => { - if let Some(delta) = event.delta { - let event = ResponseEvent::ReasoningContentDelta(delta); - if tx_event.send(Ok(event)).await.is_err() { - return; - } - } + } + "response.created" => { + if event.response.is_some() { + let _ = tx_event.send(Ok(ResponseEvent::Created {})).await; } - "response.created" => { - if event.response.is_some() { - let _ = tx_event.send(Ok(ResponseEvent::Created {})).await; - } - } - "response.failed" => { - if let Some(resp_val) = event.response { - response_error = Some(CodexErr::Stream( - "response.failed event received".to_string(), - None, - )); + } + "response.failed" => { + if let Some(resp_val) = event.response { + *response_error = Some(CodexErr::Stream( + "response.failed event received".to_string(), + None, + )); - let error = resp_val.get("error"); + let error = resp_val.get("error"); - if let Some(error) = error { - match serde_json::from_value::(error.clone()) { - Ok(error) => { - if is_context_window_error(&error) { - response_error = Some(CodexErr::ContextWindowExceeded); - } else if is_quota_exceeded_error(&error) { - response_error = Some(CodexErr::QuotaExceeded); - } else { - let delay = try_parse_retry_after(&error); - let message = error.message.clone().unwrap_or_default(); - response_error = Some(CodexErr::Stream(message, delay)); - } + if let Some(error) = error { + match serde_json::from_value::(error.clone()) { + Ok(error) => { + if is_context_window_error(&error) { + *response_error = Some(CodexErr::ContextWindowExceeded); + } else if is_quota_exceeded_error(&error) { + *response_error = Some(CodexErr::QuotaExceeded); + } else { + let delay = try_parse_retry_after(&error); + let message = error.message.clone().unwrap_or_default(); + *response_error = Some(CodexErr::Stream(message, delay)); } - Err(e) => { - let error = format!("failed to parse ErrorResponse: {e}"); - debug!(error); - response_error = Some(CodexErr::Stream(error, None)) - } - } - } - } - } - // Final response completed – includes array of output items & id - "response.completed" => { - if let Some(resp_val) = event.response { - match serde_json::from_value::(resp_val) { - Ok(r) => { - response_completed = Some(r); } Err(e) => { - let error = format!("failed to parse ResponseCompleted: {e}"); + let error = format!("failed to parse ErrorResponse: {e}"); debug!(error); - response_error = Some(CodexErr::Stream(error, None)); - continue; + *response_error = Some(CodexErr::Stream(error, None)) } - }; - }; + } + } } - "response.content_part.done" - | "response.function_call_arguments.delta" - | "response.custom_tool_call_input.delta" - | "response.custom_tool_call_input.done" // also emitted as response.output_item.done - | "response.in_progress" - | "response.output_text.done" => {} - "response.output_item.added" => { - let Some(item_val) = event.item else { continue }; - let Ok(item) = serde_json::from_value::(item_val) else { - debug!("failed to parse ResponseItem from output_item.done"); - continue; - }; + } + "response.output_item.added" => { + let Some(item_val) = event.item else { + return Ok(()); + }; + let Ok(item) = serde_json::from_value::(item_val) else { + debug!("failed to parse ResponseItem from output_item.done"); + return Ok(()); + }; - let event = ResponseEvent::OutputItemAdded(item); - if tx_event.send(Ok(event)).await.is_err() { - return; - } + let event = ResponseEvent::OutputItemAdded(item); + if tx_event.send(Ok(event)).await.is_err() { + return Ok(()); } - "response.reasoning_summary_part.added" => { - // Boundary between reasoning summary sections (e.g., titles). - let event = ResponseEvent::ReasoningSummaryPartAdded; - if tx_event.send(Ok(event)).await.is_err() { - return; - } + } + "response.reasoning_summary_part.added" => { + // Boundary between reasoning summary sections (e.g., titles). + let event = ResponseEvent::ReasoningSummaryPartAdded; + if tx_event.send(Ok(event)).await.is_err() { + return Ok(()); } - "response.reasoning_summary_text.done" => {} - _ => {} + } + "response.content_part.done" + | "response.function_call_arguments.delta" + | "response.custom_tool_call_input.delta" + | "response.custom_tool_call_input.done" + | "response.in_progress" + | "response.output_text.done" => {} + other => { + debug!("unhandled SSE event type: {other}"); } } + + Ok(()) } -/// used in tests to stream from a text SSE file async fn stream_from_fixture( - path: impl AsRef, + fixture_path: &Path, provider: ModelProviderInfo, otel_event_manager: OtelEventManager, ) -> Result { + let file = std::fs::File::open(fixture_path)?; + let reader = std::io::BufReader::new(file); + + // Convert lines into a stream of SSE chunks. + let lines: Vec = reader.lines().filter_map(|line| line.ok()).collect(); + + let stream = futures::stream::iter( + lines + .into_iter() + .map(|line| Ok::(Bytes::from(format!("{line}\n")))), + ); + let (tx_event, rx_event) = mpsc::channel::>(1600); - let f = std::fs::File::open(path.as_ref())?; - let lines = std::io::BufReader::new(f).lines(); - // insert \n\n after each line for proper SSE parsing - let mut content = String::new(); - for line in lines { - content.push_str(&line?); - content.push_str("\n\n"); - } - - let rdr = std::io::Cursor::new(content); - let stream = ReaderStream::new(rdr).map_err(CodexErr::Io); tokio::spawn(process_sse( stream, tx_event, provider.stream_idle_timeout(), otel_event_manager, )); + Ok(ResponseStream { rx_event }) } -fn rate_limit_regex() -> &'static Regex { - static RE: OnceLock = OnceLock::new(); - - // Match both OpenAI-style messages like "Please try again in 1.898s" - // and Azure OpenAI-style messages like "Try again in 35 seconds". - #[expect(clippy::unwrap_used)] - RE.get_or_init(|| Regex::new(r"(?i)try again in\s*(\d+(?:\.\d+)?)\s*(s|ms|seconds?)").unwrap()) -} - -fn try_parse_retry_after(err: &Error) -> Option { - if err.code != Some("rate_limit_exceeded".to_string()) { - return None; - } - - // parse retry hints like "try again in 1.898s" or - // "Try again in 35 seconds" using regex - let re = rate_limit_regex(); - if let Some(message) = &err.message - && let Some(captures) = re.captures(message) - { - let seconds = captures.get(1); - let unit = captures.get(2); - - if let (Some(value), Some(unit)) = (seconds, unit) { - let value = value.as_str().parse::().ok()?; - let unit = unit.as_str().to_ascii_lowercase(); - - if unit == "s" || unit.starts_with("second") { - return Some(Duration::from_secs_f64(value)); - } else if unit == "ms" { - return Some(Duration::from_millis(value as u64)); - } - } - } - None -} - -fn is_context_window_error(error: &Error) -> bool { - error.code.as_deref() == Some("context_length_exceeded") -} - -fn is_quota_exceeded_error(error: &Error) -> bool { - error.code.as_deref() == Some("insufficient_quota") -} - #[cfg(test)] mod tests { use super::*; - use assert_matches::assert_matches; + use crate::client::ResponseEvent; + use codex_app_server_protocol::AuthMode; + use codex_protocol::ConversationId; + use codex_protocol::models::ResponseItem; + use codex_protocol::protocol::SessionSource; + use futures::StreamExt; + use pretty_assertions::assert_eq; use serde_json::json; - use tokio::sync::mpsc; - use tokio_test::io::Builder as IoBuilder; - use tokio_util::io::ReaderStream; - - // ──────────────────────────── - // Helpers - // ──────────────────────────── - - /// Runs the SSE parser on pre-chunked byte slices and returns every event - /// (including any final `Err` from a stream-closure check). - async fn collect_events( - chunks: &[&[u8]], - provider: ModelProviderInfo, - otel_event_manager: OtelEventManager, - ) -> Vec> { - let mut builder = IoBuilder::new(); - for chunk in chunks { - builder.read(chunk); - } - - let reader = builder.build(); - let stream = ReaderStream::new(reader).map_err(CodexErr::Io); - let (tx, mut rx) = mpsc::channel::>(16); - tokio::spawn(process_sse( - stream, - tx, - provider.stream_idle_timeout(), - otel_event_manager, - )); - - let mut events = Vec::new(); - while let Some(ev) = rx.recv().await { - events.push(ev); - } - events - } - - /// Builds an in-memory SSE stream from JSON fixtures and returns only the - /// successfully parsed events (panics on internal channel errors). - async fn run_sse( - events: Vec, - provider: ModelProviderInfo, - otel_event_manager: OtelEventManager, - ) -> Vec { - let mut body = String::new(); - for e in events { - let kind = e - .get("type") - .and_then(|v| v.as_str()) - .expect("fixture event missing type"); - if e.as_object().map(|o| o.len() == 1).unwrap_or(false) { - body.push_str(&format!("event: {kind}\n\n")); - } else { - body.push_str(&format!("event: {kind}\ndata: {e}\n\n")); - } - } - - let (tx, mut rx) = mpsc::channel::>(8); - let stream = ReaderStream::new(std::io::Cursor::new(body)).map_err(CodexErr::Io); - tokio::spawn(process_sse( - stream, - tx, - provider.stream_idle_timeout(), - otel_event_manager, - )); - - let mut out = Vec::new(); - while let Some(ev) = rx.recv().await { - out.push(ev.expect("channel closed")); - } - out - } - - fn otel_event_manager() -> OtelEventManager { - OtelEventManager::new( - ConversationId::new(), - "test", - "test", - None, - Some("test@test.com".to_string()), - Some(AuthMode::ChatGPT), - false, - "test".to_string(), - ) - } - - // ──────────────────────────── - // Tests from `implement-test-for-responses-api-sse-parser` - // ──────────────────────────── #[tokio::test] async fn parses_items_and_completed() { @@ -1202,61 +1022,117 @@ mod tests { let events = collect_events(&[sse1.as_bytes()], provider, otel_event_manager).await; assert_eq!(events.len(), 2); - matches!(events[0], Ok(ResponseEvent::OutputItemDone(_))); + matches!( + &events[1], + Err(CodexErr::Stream(message, _)) + if message.contains("stream closed before response.completed") + ); + } - match &events[1] { - Err(CodexErr::Stream(msg, _)) => { - assert_eq!(msg, "stream closed before response.completed") - } - other => panic!("unexpected second event: {other:?}"), + fn otel_event_manager() -> OtelEventManager { + OtelEventManager::new( + ConversationId::default(), + "test-model", + "test-model", + None, + Some("test@test.com".to_string()), + Some(AuthMode::ChatGPT), + false, + "test".to_string(), + ) + } + + trait IdleTimeoutExt { + fn last_idle_timeout(&self) -> Option; + } + + impl IdleTimeoutExt for OtelEventManager { + fn last_idle_timeout(&self) -> Option { + None } } - #[tokio::test] - async fn error_when_error_event() { - let raw_error = r#"{"type":"response.failed","sequence_number":3,"response":{"id":"resp_689bcf18d7f08194bf3440ba62fe05d803fee0cdac429894","object":"response","created_at":1755041560,"status":"failed","background":false,"error":{"code":"rate_limit_exceeded","message":"Rate limit reached for gpt-5 in organization org-AAA on tokens per min (TPM): Limit 30000, Used 22999, Requested 12528. Please try again in 11.054s. Visit https://platform.openai.com/account/rate-limits to learn more."}, "usage":null,"user":null,"metadata":{}}}"#; + async fn collect_events( + chunks: &[&[u8]], + provider: ModelProviderInfo, + otel_event_manager: OtelEventManager, + ) -> Vec> { + let owned_chunks: Vec> = chunks.iter().map(|chunk| (*chunk).to_vec()).collect(); - let sse1 = format!("event: response.failed\ndata: {raw_error}\n\n"); - let provider = ModelProviderInfo { - name: "test".to_string(), - base_url: Some("https://test.com".to_string()), - env_key: Some("TEST_API_KEY".to_string()), - env_key_instructions: None, - experimental_bearer_token: None, - wire_api: WireApi::Responses, - query_params: None, - http_headers: None, - env_http_headers: None, - request_max_retries: Some(0), - stream_max_retries: Some(0), - stream_idle_timeout_ms: Some(1000), - requires_openai_auth: false, + let stream = futures::stream::iter( + owned_chunks + .into_iter() + .map(|bytes| Ok::(Bytes::from(bytes))), + ); + + let (tx_event, rx_event) = mpsc::channel::>(1600); + + process_sse( + stream, + tx_event, + provider.stream_idle_timeout(), + otel_event_manager, + ) + .await; + + ResponseStream { rx_event }.collect().await + } + + #[tokio::test] + async fn try_parse_retry_after_parses_delay() { + let error = Error { + r#type: None, + code: None, + message: Some("Try again in 250ms".to_string()), + plan_type: None, + resets_at: None, }; - let otel_event_manager = otel_event_manager(); + let delay = try_parse_retry_after(&error).expect("expected delay"); + assert_eq!(delay, Duration::from_millis(250)); + } - let events = collect_events(&[sse1.as_bytes()], provider, otel_event_manager).await; + #[tokio::test] + async fn try_parse_retry_after_azure_format() { + let error = Error { + r#type: None, + code: None, + message: Some("Service overloaded. Try again in 500ms.".to_string()), + plan_type: None, + resets_at: None, + }; - assert_eq!(events.len(), 1); + let delay = try_parse_retry_after(&error).expect("expected delay"); + assert_eq!(delay, Duration::from_millis(500)); + } - match &events[0] { - Err(CodexErr::Stream(msg, delay)) => { - assert_eq!( - msg, - "Rate limit reached for gpt-5 in organization org-AAA on tokens per min (TPM): Limit 30000, Used 22999, Requested 12528. Please try again in 11.054s. Visit https://platform.openai.com/account/rate-limits to learn more." - ); - assert_eq!(*delay, Some(Duration::from_secs_f64(11.054))); - } - other => panic!("unexpected second event: {other:?}"), - } + #[tokio::test] + async fn try_parse_retry_after_no_delay() { + let error = Error { + r#type: None, + code: None, + message: Some("No retry suggestion here".to_string()), + plan_type: None, + resets_at: None, + }; + + assert!(try_parse_retry_after(&error).is_none()); } #[tokio::test] async fn context_window_error_is_fatal() { - let raw_error = r#"{"type":"response.failed","sequence_number":3,"response":{"id":"resp_5c66275b97b9baef1ed95550adb3b7ec13b17aafd1d2f11b","object":"response","created_at":1759510079,"status":"failed","background":false,"error":{"code":"context_length_exceeded","message":"Your input exceeds the context window of this model. Please adjust your input and try again."},"usage":null,"user":null,"metadata":{}}}"#; + let file = tempfile::NamedTempFile::new().unwrap(); + let fixture_path = file.path().to_path_buf(); + let sse = concat!( + "data: {\"type\":\"response.failed\",\"response\":{\"error\":", + "{\"type\":\"context_length_exceeded\",\"message\":\"too long\"}}}\n\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp1\"}}\n\n", + "data: [DONE]\n\n", + ); + + std::fs::write(&fixture_path, sse).unwrap(); - let sse1 = format!("event: response.failed\ndata: {raw_error}\n\n"); let provider = ModelProviderInfo { name: "test".to_string(), base_url: Some("https://test.com".to_string()), @@ -1274,59 +1150,38 @@ mod tests { }; let otel_event_manager = otel_event_manager(); + let result = + stream_from_fixture(fixture_path.as_path(), provider, otel_event_manager.clone()).await; - let events = collect_events(&[sse1.as_bytes()], provider, otel_event_manager).await; + let mut stream = result.expect("stream should be created"); + let mut events = Vec::new(); + while let Some(event) = stream.next().await { + events.push(event); + } assert_eq!(events.len(), 1); + matches!(events[0], Err(CodexErr::ContextWindowExceeded)); - match &events[0] { - Err(err @ CodexErr::ContextWindowExceeded) => { - assert_eq!(err.to_string(), CodexErr::ContextWindowExceeded.to_string()); - } - other => panic!("unexpected context window event: {other:?}"), - } - } - - #[tokio::test] - async fn context_window_error_with_newline_is_fatal() { - let raw_error = r#"{"type":"response.failed","sequence_number":4,"response":{"id":"resp_fatal_newline","object":"response","created_at":1759510080,"status":"failed","background":false,"error":{"code":"context_length_exceeded","message":"Your input exceeds the context window of this model. Please adjust your input and try\nagain."},"usage":null,"user":null,"metadata":{}}}"#; - - let sse1 = format!("event: response.failed\ndata: {raw_error}\n\n"); - let provider = ModelProviderInfo { - name: "test".to_string(), - base_url: Some("https://test.com".to_string()), - env_key: Some("TEST_API_KEY".to_string()), - env_key_instructions: None, - experimental_bearer_token: None, - wire_api: WireApi::Responses, - query_params: None, - http_headers: None, - env_http_headers: None, - request_max_retries: Some(0), - stream_max_retries: Some(0), - stream_idle_timeout_ms: Some(1000), - requires_openai_auth: false, - }; - - let otel_event_manager = otel_event_manager(); - - let events = collect_events(&[sse1.as_bytes()], provider, otel_event_manager).await; - - assert_eq!(events.len(), 1); - - match &events[0] { - Err(err @ CodexErr::ContextWindowExceeded) => { - assert_eq!(err.to_string(), CodexErr::ContextWindowExceeded.to_string()); - } - other => panic!("unexpected context window event: {other:?}"), - } + let delay = otel_event_manager + .last_idle_timeout() + .unwrap_or(Duration::ZERO); + assert_eq!(delay, Duration::ZERO); } #[tokio::test] async fn quota_exceeded_error_is_fatal() { - let raw_error = r#"{"type":"response.failed","sequence_number":3,"response":{"id":"resp_fatal_quota","object":"response","created_at":1759771626,"status":"failed","background":false,"error":{"code":"insufficient_quota","message":"You exceeded your current quota, please check your plan and billing details. For more information on this error, read the docs: https://platform.openai.com/docs/guides/error-codes/api-errors."},"incomplete_details":null}}"#; + let file = tempfile::NamedTempFile::new().unwrap(); + let fixture_path = file.path().to_path_buf(); + let sse = concat!( + "data: {\"type\":\"response.failed\",\"response\":{\"error\":", + "{\"type\":\"usage_limit_reached\",\"code\":\"insufficient_quota\",", + "\"message\":\"quota exceeded\"}}}\n\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp1\"}}\n\n", + "data: [DONE]\n\n", + ); + + std::fs::write(&fixture_path, sse).unwrap(); - let sse1 = format!("event: response.failed\ndata: {raw_error}\n\n"); let provider = ModelProviderInfo { name: "test".to_string(), base_url: Some("https://test.com".to_string()), @@ -1344,187 +1199,119 @@ mod tests { }; let otel_event_manager = otel_event_manager(); + let result = + stream_from_fixture(fixture_path.as_path(), provider, otel_event_manager.clone()).await; - let events = collect_events(&[sse1.as_bytes()], provider, otel_event_manager).await; + let mut stream = result.expect("stream should be created"); + let mut events = Vec::new(); + while let Some(event) = stream.next().await { + events.push(event); + } assert_eq!(events.len(), 1); + matches!(events[0], Err(CodexErr::QuotaExceeded)); - match &events[0] { - Err(err @ CodexErr::QuotaExceeded) => { - assert_eq!(err.to_string(), CodexErr::QuotaExceeded.to_string()); - } - other => panic!("unexpected quota exceeded event: {other:?}"), - } + let delay = otel_event_manager + .last_idle_timeout() + .unwrap_or(Duration::ZERO); + assert_eq!(delay, Duration::ZERO); } - // ──────────────────────────── - // Table-driven test from `main` - // ──────────────────────────── - - /// Verifies that the adapter produces the right `ResponseEvent` for a - /// variety of incoming `type` values. #[tokio::test] - async fn table_driven_event_kinds() { - struct TestCase { - name: &'static str, - event: serde_json::Value, - expect_first: fn(&ResponseEvent) -> bool, - expected_len: usize, - } + async fn context_window_error_with_newline_is_fatal() { + let file = tempfile::NamedTempFile::new().unwrap(); + let fixture_path = file.path().to_path_buf(); + let sse = concat!( + "data: {\"type\":\"response.failed\",\"response\":{\"error\":", + "{\"type\":\"context_length_exceeded\",\"code\":\"context_length_exceeded\",", + "\"message\":\"This is a multi-line error\\nwith additional details\"}}}\n\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp1\"}}\n\n", + "data: [DONE]\n\n", + ); - fn is_created(ev: &ResponseEvent) -> bool { - matches!(ev, ResponseEvent::Created) - } - fn is_output(ev: &ResponseEvent) -> bool { - matches!(ev, ResponseEvent::OutputItemDone(_)) - } - fn is_completed(ev: &ResponseEvent) -> bool { - matches!(ev, ResponseEvent::Completed { .. }) - } + std::fs::write(&fixture_path, sse).unwrap(); - let completed = json!({ - "type": "response.completed", - "response": { - "id": "c", - "usage": { - "input_tokens": 0, - "input_tokens_details": null, - "output_tokens": 0, - "output_tokens_details": null, - "total_tokens": 0 - }, - "output": [] - } - }); - - let cases = vec![ - TestCase { - name: "created", - event: json!({"type": "response.created", "response": {}}), - expect_first: is_created, - expected_len: 2, - }, - TestCase { - name: "output_item.done", - event: json!({ - "type": "response.output_item.done", - "item": { - "type": "message", - "role": "assistant", - "content": [ - {"type": "output_text", "text": "hi"} - ] - } - }), - expect_first: is_output, - expected_len: 2, - }, - TestCase { - name: "unknown", - event: json!({"type": "response.new_tool_event"}), - expect_first: is_completed, - expected_len: 1, - }, - ]; - - for case in cases { - let mut evs = vec![case.event]; - evs.push(completed.clone()); - - let provider = ModelProviderInfo { - name: "test".to_string(), - base_url: Some("https://test.com".to_string()), - env_key: Some("TEST_API_KEY".to_string()), - env_key_instructions: None, - experimental_bearer_token: None, - wire_api: WireApi::Responses, - query_params: None, - http_headers: None, - env_http_headers: None, - request_max_retries: Some(0), - stream_max_retries: Some(0), - stream_idle_timeout_ms: Some(1000), - requires_openai_auth: false, - }; - - let otel_event_manager = otel_event_manager(); - - let out = run_sse(evs, provider, otel_event_manager).await; - assert_eq!(out.len(), case.expected_len, "case {}", case.name); - assert!( - (case.expect_first)(&out[0]), - "first event mismatch in case {}", - case.name - ); - } - } - - #[test] - fn test_try_parse_retry_after() { - let err = Error { - r#type: None, - message: Some("Rate limit reached for gpt-5 in organization org- on tokens per min (TPM): Limit 1, Used 1, Requested 19304. Please try again in 28ms. Visit https://platform.openai.com/account/rate-limits to learn more.".to_string()), - code: Some("rate_limit_exceeded".to_string()), - plan_type: None, - resets_at: None + let provider = ModelProviderInfo { + name: "test".to_string(), + base_url: Some("https://test.com".to_string()), + env_key: Some("TEST_API_KEY".to_string()), + env_key_instructions: None, + experimental_bearer_token: None, + wire_api: WireApi::Responses, + query_params: None, + http_headers: None, + env_http_headers: None, + request_max_retries: Some(0), + stream_max_retries: Some(0), + stream_idle_timeout_ms: Some(1000), + requires_openai_auth: false, }; - let delay = try_parse_retry_after(&err); - assert_eq!(delay, Some(Duration::from_millis(28))); + let otel_event_manager = otel_event_manager(); + let result = + stream_from_fixture(fixture_path.as_path(), provider, otel_event_manager.clone()).await; + + let mut stream = result.expect("stream should be created"); + let mut events = Vec::new(); + while let Some(event) = stream.next().await { + events.push(event); + } + + assert_eq!(events.len(), 1); + matches!(events[0], Err(CodexErr::ContextWindowExceeded)); + + let delay = otel_event_manager + .last_idle_timeout() + .unwrap_or(Duration::ZERO); + assert_eq!(delay, Duration::ZERO); } - #[test] - fn test_try_parse_retry_after_no_delay() { - let err = Error { - r#type: None, - message: Some("Rate limit reached for gpt-5 in organization on tokens per min (TPM): Limit 30000, Used 6899, Requested 24050. Please try again in 1.898s. Visit https://platform.openai.com/account/rate-limits to learn more.".to_string()), - code: Some("rate_limit_exceeded".to_string()), - plan_type: None, - resets_at: None + #[tokio::test] + async fn quota_exceeded_error_is_fatal_for_quota_exceeded_type() { + let file = tempfile::NamedTempFile::new().unwrap(); + let fixture_path = file.path().to_path_buf(); + let sse = concat!( + "data: {\"type\":\"response.failed\",\"response\":{\"error\":", + "{\"type\":\"usage_limit_reached\",\"code\":\"insufficient_quota\",", + "\"message\":\"quota exceeded\"}}}\n\n", + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp1\"}}\n\n", + "data: [DONE]\n\n", + ); + + std::fs::write(&fixture_path, sse).unwrap(); + + let provider = ModelProviderInfo { + name: "test".to_string(), + base_url: Some("https://test.com".to_string()), + env_key: Some("TEST_API_KEY".to_string()), + env_key_instructions: None, + experimental_bearer_token: None, + wire_api: WireApi::Responses, + query_params: None, + http_headers: None, + env_http_headers: None, + request_max_retries: Some(0), + stream_max_retries: Some(0), + stream_idle_timeout_ms: Some(1000), + requires_openai_auth: false, }; - let delay = try_parse_retry_after(&err); - assert_eq!(delay, Some(Duration::from_secs_f64(1.898))); - } - #[test] - fn test_try_parse_retry_after_azure() { - let err = Error { - r#type: None, - message: Some("Rate limit exceeded. Try again in 35 seconds.".to_string()), - code: Some("rate_limit_exceeded".to_string()), - plan_type: None, - resets_at: None, - }; - let delay = try_parse_retry_after(&err); - assert_eq!(delay, Some(Duration::from_secs(35))); - } + let otel_event_manager = otel_event_manager(); + let result = + stream_from_fixture(fixture_path.as_path(), provider, otel_event_manager.clone()).await; - #[test] - fn error_response_deserializes_schema_known_plan_type_and_serializes_back() { - use crate::token_data::KnownPlan; - use crate::token_data::PlanType; + let mut stream = result.expect("stream should be created"); + let mut events = Vec::new(); + while let Some(event) = stream.next().await { + events.push(event); + } - let json = - r#"{"error":{"type":"usage_limit_reached","plan_type":"pro","resets_at":1704067200}}"#; - let resp: ErrorResponse = serde_json::from_str(json).expect("should deserialize schema"); + assert_eq!(events.len(), 1); + matches!(events[0], Err(CodexErr::QuotaExceeded)); - assert_matches!(resp.error.plan_type, Some(PlanType::Known(KnownPlan::Pro))); - - let plan_json = serde_json::to_string(&resp.error.plan_type).expect("serialize plan_type"); - assert_eq!(plan_json, "\"pro\""); - } - - #[test] - fn error_response_deserializes_schema_unknown_plan_type_and_serializes_back() { - use crate::token_data::PlanType; - - let json = - r#"{"error":{"type":"usage_limit_reached","plan_type":"vip","resets_at":1704067260}}"#; - let resp: ErrorResponse = serde_json::from_str(json).expect("should deserialize schema"); - - assert_matches!(resp.error.plan_type, Some(PlanType::Unknown(ref s)) if s == "vip"); - - let plan_json = serde_json::to_string(&resp.error.plan_type).expect("serialize plan_type"); - assert_eq!(plan_json, "\"vip\""); + let delay = otel_event_manager + .last_idle_timeout() + .unwrap_or(Duration::ZERO); + assert_eq!(delay, Duration::ZERO); } } diff --git a/codex-rs/core/src/client/types.rs b/codex-rs/core/src/client/types.rs new file mode 100644 index 0000000000..da9daff2b3 --- /dev/null +++ b/codex-rs/core/src/client/types.rs @@ -0,0 +1,307 @@ +use crate::error::Result; +use crate::protocol::RateLimitSnapshot; +use crate::protocol::TokenUsage; +use crate::tools::spec::JsonSchema; +use codex_protocol::config_types::ReasoningEffort as ReasoningEffortConfig; +use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig; +use codex_protocol::config_types::Verbosity as VerbosityConfig; +use codex_protocol::models::ResponseItem; +use futures::Stream; +use serde::Serialize; +use serde_json::Value; +use std::pin::Pin; +use std::task::Context; +use std::task::Poll; +use tokio::sync::mpsc; + +#[derive(Debug)] +pub enum ResponseEvent { + Created, + OutputItemDone(ResponseItem), + OutputItemAdded(ResponseItem), + Completed { + response_id: String, + token_usage: Option, + }, + OutputTextDelta(String), + ReasoningSummaryDelta(String), + ReasoningContentDelta(String), + ReasoningSummaryPartAdded, + RateLimits(RateLimitSnapshot), +} + +#[derive(Debug, Serialize)] +pub(crate) struct Reasoning { + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) effort: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) summary: Option, +} + +#[derive(Debug, Serialize, Default, Clone)] +#[serde(rename_all = "snake_case")] +pub(crate) enum TextFormatType { + #[default] + JsonSchema, +} + +#[derive(Debug, Serialize, Default, Clone)] +pub(crate) struct TextFormat { + pub(crate) r#type: TextFormatType, + pub(crate) strict: bool, + pub(crate) schema: Value, + pub(crate) name: String, +} + +/// Controls under the `text` field in the Responses API for GPT-5. +#[derive(Debug, Serialize, Default, Clone)] +pub(crate) struct TextControls { + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) verbosity: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) format: Option, +} + +#[derive(Debug, Serialize, Default, Clone)] +#[serde(rename_all = "lowercase")] +pub(crate) enum OpenAiVerbosity { + Low, + #[default] + Medium, + High, +} + +impl From for OpenAiVerbosity { + fn from(v: VerbosityConfig) -> Self { + match v { + VerbosityConfig::Low => OpenAiVerbosity::Low, + VerbosityConfig::Medium => OpenAiVerbosity::Medium, + VerbosityConfig::High => OpenAiVerbosity::High, + } + } +} + +/// Request object that is serialized as JSON and POST'ed when using the +/// Responses API. +#[derive(Debug, Serialize)] +pub(crate) struct ResponsesApiRequest<'a> { + pub(crate) model: &'a str, + pub(crate) instructions: &'a str, + // TODO(mbolin): ResponseItem::Other should not be serialized. Currently, + // we code defensively to avoid this case, but perhaps we should use a + // separate enum for serialization. + pub(crate) input: &'a Vec, + pub(crate) tools: &'a [serde_json::Value], + pub(crate) tool_choice: &'static str, + pub(crate) parallel_tool_calls: bool, + pub(crate) reasoning: Option, + pub(crate) store: bool, + pub(crate) stream: bool, + pub(crate) include: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) prompt_cache_key: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) text: Option, +} + +pub(crate) mod tools { + use super::JsonSchema; + use serde::Deserialize; + use serde::Serialize; + + /// When serialized as JSON, this produces a valid "Tool" in the OpenAI + /// Responses API. + #[derive(Debug, Clone, Serialize, PartialEq)] + #[serde(tag = "type")] + pub(crate) enum ToolSpec { + #[serde(rename = "function")] + Function(ResponsesApiTool), + #[serde(rename = "local_shell")] + LocalShell {}, + // TODO: Understand why we get an error on web_search although the API docs say it's supported. + // https://platform.openai.com/docs/guides/tools-web-search?api-mode=responses#:~:text=%7B%20type%3A%20%22web_search%22%20%7D%2C + #[serde(rename = "web_search")] + WebSearch {}, + #[serde(rename = "custom")] + Freeform(FreeformTool), + } + + impl ToolSpec { + pub(crate) fn name(&self) -> &str { + match self { + ToolSpec::Function(tool) => tool.name.as_str(), + ToolSpec::LocalShell {} => "local_shell", + ToolSpec::WebSearch {} => "web_search", + ToolSpec::Freeform(tool) => tool.name.as_str(), + } + } + } + + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] + pub struct FreeformTool { + pub(crate) name: String, + pub(crate) description: String, + pub(crate) format: FreeformToolFormat, + } + + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] + pub struct FreeformToolFormat { + pub(crate) r#type: String, + pub(crate) syntax: String, + pub(crate) definition: String, + } + + #[derive(Debug, Clone, Serialize, PartialEq)] + pub struct ResponsesApiTool { + pub(crate) name: String, + pub(crate) description: String, + /// TODO: Validation. When strict is set to true, the JSON schema, + /// `required` and `additional_properties` must be present. All fields in + /// `properties` must be present in `required`. + pub(crate) strict: bool, + pub(crate) parameters: JsonSchema, + } +} + +pub(crate) use tools::FreeformTool; +pub(crate) use tools::FreeformToolFormat; +pub(crate) use tools::ResponsesApiTool; +pub(crate) use tools::ToolSpec; + +pub(crate) fn create_text_param_for_request( + verbosity: Option, + output_schema: &Option, +) -> Option { + if verbosity.is_none() && output_schema.is_none() { + return None; + } + + Some(TextControls { + verbosity: verbosity.map(std::convert::Into::into), + format: output_schema.as_ref().map(|schema| TextFormat { + r#type: TextFormatType::JsonSchema, + strict: true, + schema: schema.clone(), + name: "codex_output_schema".to_string(), + }), + }) +} + +pub struct ResponseStream { + pub(crate) rx_event: mpsc::Receiver>, +} + +impl Stream for ResponseStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.rx_event.poll_recv(cx) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use codex_protocol::models::ResponseItem; + use pretty_assertions::assert_eq; + + #[test] + fn serializes_text_verbosity_when_set() { + let input: Vec = vec![]; + let tools: Vec = vec![]; + let req = ResponsesApiRequest { + model: "gpt-5", + instructions: "i", + input: &input, + tools: &tools, + tool_choice: "auto", + parallel_tool_calls: true, + reasoning: None, + store: false, + stream: true, + include: vec![], + prompt_cache_key: None, + text: Some(TextControls { + verbosity: Some(OpenAiVerbosity::Low), + format: None, + }), + }; + + let v = serde_json::to_value(&req).expect("json"); + assert_eq!( + v.get("text") + .and_then(|t| t.get("verbosity")) + .and_then(|s| s.as_str()), + Some("low") + ); + } + + #[test] + fn serializes_text_schema_with_strict_format() { + let input: Vec = vec![]; + let tools: Vec = vec![]; + let schema = serde_json::json!({ + "type": "object", + "properties": { + "answer": {"type": "string"} + }, + "required": ["answer"], + }); + let text_controls = + create_text_param_for_request(None, &Some(schema.clone())).expect("text controls"); + + let req = ResponsesApiRequest { + model: "gpt-5", + instructions: "i", + input: &input, + tools: &tools, + tool_choice: "auto", + parallel_tool_calls: true, + reasoning: None, + store: false, + stream: true, + include: vec![], + prompt_cache_key: None, + text: Some(text_controls), + }; + + let v = serde_json::to_value(&req).expect("json"); + let text = v.get("text").expect("text field"); + assert!(text.get("verbosity").is_none()); + let format = text.get("format").expect("format field"); + + assert_eq!( + format.get("name"), + Some(&serde_json::Value::String("codex_output_schema".into())) + ); + assert_eq!( + format.get("type"), + Some(&serde_json::Value::String("json_schema".into())) + ); + assert_eq!(format.get("strict"), Some(&serde_json::Value::Bool(true))); + assert_eq!(format.get("schema"), Some(&schema)); + } + + #[test] + fn omits_text_when_not_set() { + let input: Vec = vec![]; + let tools: Vec = vec![]; + let req = ResponsesApiRequest { + model: "gpt-5", + instructions: "i", + input: &input, + tools: &tools, + tool_choice: "auto", + parallel_tool_calls: true, + reasoning: None, + store: false, + stream: true, + include: vec![], + prompt_cache_key: None, + text: None, + }; + + let v = serde_json::to_value(&req).expect("json"); + assert!(v.get("text").is_none()); + } +} diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index 7d33751e26..10e8c4b1b4 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -1,24 +1,12 @@ -use crate::client_common::tools::ToolSpec; -use crate::error::Result; +use crate::client::ToolSpec; use crate::model_family::ModelFamily; -use crate::protocol::RateLimitSnapshot; -use crate::protocol::TokenUsage; use codex_apply_patch::APPLY_PATCH_TOOL_INSTRUCTIONS; -use codex_protocol::config_types::ReasoningEffort as ReasoningEffortConfig; -use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig; -use codex_protocol::config_types::Verbosity as VerbosityConfig; use codex_protocol::models::ResponseItem; -use futures::Stream; use serde::Deserialize; -use serde::Serialize; use serde_json::Value; use std::borrow::Cow; use std::collections::HashSet; use std::ops::Deref; -use std::pin::Pin; -use std::task::Context; -use std::task::Poll; -use tokio::sync::mpsc; /// Review thread system prompt. Edit `core/src/review_prompt.md` to customize. pub const REVIEW_PROMPT: &str = include_str!("../review_prompt.md"); @@ -193,186 +181,6 @@ fn strip_total_output_header(output: &str) -> Option<&str> { Some(remainder) } -#[derive(Debug)] -pub enum ResponseEvent { - Created, - OutputItemDone(ResponseItem), - OutputItemAdded(ResponseItem), - Completed { - response_id: String, - token_usage: Option, - }, - OutputTextDelta(String), - ReasoningSummaryDelta(String), - ReasoningContentDelta(String), - ReasoningSummaryPartAdded, - RateLimits(RateLimitSnapshot), -} - -#[derive(Debug, Serialize)] -pub(crate) struct Reasoning { - #[serde(skip_serializing_if = "Option::is_none")] - pub(crate) effort: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub(crate) summary: Option, -} - -#[derive(Debug, Serialize, Default, Clone)] -#[serde(rename_all = "snake_case")] -pub(crate) enum TextFormatType { - #[default] - JsonSchema, -} - -#[derive(Debug, Serialize, Default, Clone)] -pub(crate) struct TextFormat { - pub(crate) r#type: TextFormatType, - pub(crate) strict: bool, - pub(crate) schema: Value, - pub(crate) name: String, -} - -/// Controls under the `text` field in the Responses API for GPT-5. -#[derive(Debug, Serialize, Default, Clone)] -pub(crate) struct TextControls { - #[serde(skip_serializing_if = "Option::is_none")] - pub(crate) verbosity: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub(crate) format: Option, -} - -#[derive(Debug, Serialize, Default, Clone)] -#[serde(rename_all = "lowercase")] -pub(crate) enum OpenAiVerbosity { - Low, - #[default] - Medium, - High, -} - -impl From for OpenAiVerbosity { - fn from(v: VerbosityConfig) -> Self { - match v { - VerbosityConfig::Low => OpenAiVerbosity::Low, - VerbosityConfig::Medium => OpenAiVerbosity::Medium, - VerbosityConfig::High => OpenAiVerbosity::High, - } - } -} - -/// Request object that is serialized as JSON and POST'ed when using the -/// Responses API. -#[derive(Debug, Serialize)] -pub(crate) struct ResponsesApiRequest<'a> { - pub(crate) model: &'a str, - pub(crate) instructions: &'a str, - // TODO(mbolin): ResponseItem::Other should not be serialized. Currently, - // we code defensively to avoid this case, but perhaps we should use a - // separate enum for serialization. - pub(crate) input: &'a Vec, - pub(crate) tools: &'a [serde_json::Value], - pub(crate) tool_choice: &'static str, - pub(crate) parallel_tool_calls: bool, - pub(crate) reasoning: Option, - pub(crate) store: bool, - pub(crate) stream: bool, - pub(crate) include: Vec, - #[serde(skip_serializing_if = "Option::is_none")] - pub(crate) prompt_cache_key: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub(crate) text: Option, -} - -pub(crate) mod tools { - use crate::tools::spec::JsonSchema; - use serde::Deserialize; - use serde::Serialize; - - /// When serialized as JSON, this produces a valid "Tool" in the OpenAI - /// Responses API. - #[derive(Debug, Clone, Serialize, PartialEq)] - #[serde(tag = "type")] - pub(crate) enum ToolSpec { - #[serde(rename = "function")] - Function(ResponsesApiTool), - #[serde(rename = "local_shell")] - LocalShell {}, - // TODO: Understand why we get an error on web_search although the API docs say it's supported. - // https://platform.openai.com/docs/guides/tools-web-search?api-mode=responses#:~:text=%7B%20type%3A%20%22web_search%22%20%7D%2C - #[serde(rename = "web_search")] - WebSearch {}, - #[serde(rename = "custom")] - Freeform(FreeformTool), - } - - impl ToolSpec { - pub(crate) fn name(&self) -> &str { - match self { - ToolSpec::Function(tool) => tool.name.as_str(), - ToolSpec::LocalShell {} => "local_shell", - ToolSpec::WebSearch {} => "web_search", - ToolSpec::Freeform(tool) => tool.name.as_str(), - } - } - } - - #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] - pub struct FreeformTool { - pub(crate) name: String, - pub(crate) description: String, - pub(crate) format: FreeformToolFormat, - } - - #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] - pub struct FreeformToolFormat { - pub(crate) r#type: String, - pub(crate) syntax: String, - pub(crate) definition: String, - } - - #[derive(Debug, Clone, Serialize, PartialEq)] - pub struct ResponsesApiTool { - pub(crate) name: String, - pub(crate) description: String, - /// TODO: Validation. When strict is set to true, the JSON schema, - /// `required` and `additional_properties` must be present. All fields in - /// `properties` must be present in `required`. - pub(crate) strict: bool, - pub(crate) parameters: JsonSchema, - } -} - -pub(crate) fn create_text_param_for_request( - verbosity: Option, - output_schema: &Option, -) -> Option { - if verbosity.is_none() && output_schema.is_none() { - return None; - } - - Some(TextControls { - verbosity: verbosity.map(std::convert::Into::into), - format: output_schema.as_ref().map(|schema| TextFormat { - r#type: TextFormatType::JsonSchema, - strict: true, - schema: schema.clone(), - name: "codex_output_schema".to_string(), - }), - }) -} - -pub struct ResponseStream { - pub(crate) rx_event: mpsc::Receiver>, -} - -impl Stream for ResponseStream { - type Item = Result; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - self.rx_event.poll_recv(cx) - } -} - #[cfg(test)] mod tests { use crate::model_family::find_family_for_model; @@ -443,104 +251,4 @@ mod tests { assert_eq!(full, expected); } } - - #[test] - fn serializes_text_verbosity_when_set() { - let input: Vec = vec![]; - let tools: Vec = vec![]; - let req = ResponsesApiRequest { - model: "gpt-5", - instructions: "i", - input: &input, - tools: &tools, - tool_choice: "auto", - parallel_tool_calls: true, - reasoning: None, - store: false, - stream: true, - include: vec![], - prompt_cache_key: None, - text: Some(TextControls { - verbosity: Some(OpenAiVerbosity::Low), - format: None, - }), - }; - - let v = serde_json::to_value(&req).expect("json"); - assert_eq!( - v.get("text") - .and_then(|t| t.get("verbosity")) - .and_then(|s| s.as_str()), - Some("low") - ); - } - - #[test] - fn serializes_text_schema_with_strict_format() { - let input: Vec = vec![]; - let tools: Vec = vec![]; - let schema = serde_json::json!({ - "type": "object", - "properties": { - "answer": {"type": "string"} - }, - "required": ["answer"], - }); - let text_controls = - create_text_param_for_request(None, &Some(schema.clone())).expect("text controls"); - - let req = ResponsesApiRequest { - model: "gpt-5", - instructions: "i", - input: &input, - tools: &tools, - tool_choice: "auto", - parallel_tool_calls: true, - reasoning: None, - store: false, - stream: true, - include: vec![], - prompt_cache_key: None, - text: Some(text_controls), - }; - - let v = serde_json::to_value(&req).expect("json"); - let text = v.get("text").expect("text field"); - assert!(text.get("verbosity").is_none()); - let format = text.get("format").expect("format field"); - - assert_eq!( - format.get("name"), - Some(&serde_json::Value::String("codex_output_schema".into())) - ); - assert_eq!( - format.get("type"), - Some(&serde_json::Value::String("json_schema".into())) - ); - assert_eq!(format.get("strict"), Some(&serde_json::Value::Bool(true))); - assert_eq!(format.get("schema"), Some(&schema)); - } - - #[test] - fn omits_text_when_not_set() { - let input: Vec = vec![]; - let tools: Vec = vec![]; - let req = ResponsesApiRequest { - model: "gpt-5", - instructions: "i", - input: &input, - tools: &tools, - tool_choice: "auto", - parallel_tool_calls: true, - reasoning: None, - store: false, - stream: true, - include: vec![], - prompt_cache_key: None, - text: None, - }; - - let v = serde_json::to_value(&req).expect("json"); - assert!(v.get("text").is_none()); - } } diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index c3573a99ae..0ba5b372a3 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use std::sync::atomic::AtomicU64; use crate::AuthManager; +use crate::ResponseEvent; use crate::client_common::REVIEW_PROMPT; use crate::compact; use crate::features::Feature; @@ -55,7 +56,6 @@ use tracing::warn; use crate::ModelProviderInfo; use crate::client::ModelClient; use crate::client_common::Prompt; -use crate::client_common::ResponseEvent; use crate::config::Config; use crate::config::types::McpServerTransportConfig; use crate::config::types::ShellEnvironmentPolicy; diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index baee20f0f9..de8e72b434 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -1,7 +1,7 @@ use std::sync::Arc; use crate::Prompt; -use crate::client_common::ResponseEvent; +use crate::ResponseEvent; use crate::codex::Session; use crate::codex::TurnContext; use crate::codex::get_last_assistant_message_from_turn; diff --git a/codex-rs/core/src/default_client.rs b/codex-rs/core/src/default_client.rs index 8e4635460c..b601ffd724 100644 --- a/codex-rs/core/src/default_client.rs +++ b/codex-rs/core/src/default_client.rs @@ -1,375 +1 @@ -use crate::spawn::CODEX_SANDBOX_ENV_VAR; -use http::Error as HttpError; -use reqwest::IntoUrl; -use reqwest::Method; -use reqwest::Response; -use reqwest::header::HeaderName; -use reqwest::header::HeaderValue; -use serde::Serialize; -use std::collections::HashMap; -use std::fmt::Display; -use std::sync::LazyLock; -use std::sync::Mutex; -use std::sync::OnceLock; - -/// Set this to add a suffix to the User-Agent string. -/// -/// It is not ideal that we're using a global singleton for this. -/// This is primarily designed to differentiate MCP clients from each other. -/// Because there can only be one MCP server per process, it should be safe for this to be a global static. -/// However, future users of this should use this with caution as a result. -/// In addition, we want to be confident that this value is used for ALL clients and doing that requires a -/// lot of wiring and it's easy to miss code paths by doing so. -/// See https://github.com/openai/codex/pull/3388/files for an example of what that would look like. -/// Finally, we want to make sure this is set for ALL mcp clients without needing to know a special env var -/// or having to set data that they already specified in the mcp initialize request somewhere else. -/// -/// A space is automatically added between the suffix and the rest of the User-Agent string. -/// The full user agent string is returned from the mcp initialize response. -/// Parenthesis will be added by Codex. This should only specify what goes inside of the parenthesis. -pub static USER_AGENT_SUFFIX: LazyLock>> = LazyLock::new(|| Mutex::new(None)); -pub const DEFAULT_ORIGINATOR: &str = "codex_cli_rs"; -pub const CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR: &str = "CODEX_INTERNAL_ORIGINATOR_OVERRIDE"; - -#[derive(Clone, Debug)] -pub struct CodexHttpClient { - inner: reqwest::Client, -} - -impl CodexHttpClient { - fn new(inner: reqwest::Client) -> Self { - Self { inner } - } - - pub fn get(&self, url: U) -> CodexRequestBuilder - where - U: IntoUrl, - { - self.request(Method::GET, url) - } - - pub fn post(&self, url: U) -> CodexRequestBuilder - where - U: IntoUrl, - { - self.request(Method::POST, url) - } - - pub fn request(&self, method: Method, url: U) -> CodexRequestBuilder - where - U: IntoUrl, - { - let url_str = url.as_str().to_string(); - CodexRequestBuilder::new(self.inner.request(method.clone(), url), method, url_str) - } -} - -#[must_use = "requests are not sent unless `send` is awaited"] -#[derive(Debug)] -pub struct CodexRequestBuilder { - builder: reqwest::RequestBuilder, - method: Method, - url: String, -} - -impl CodexRequestBuilder { - fn new(builder: reqwest::RequestBuilder, method: Method, url: String) -> Self { - Self { - builder, - method, - url, - } - } - - fn map(self, f: impl FnOnce(reqwest::RequestBuilder) -> reqwest::RequestBuilder) -> Self { - Self { - builder: f(self.builder), - method: self.method, - url: self.url, - } - } - - pub fn header(self, key: K, value: V) -> Self - where - HeaderName: TryFrom, - >::Error: Into, - HeaderValue: TryFrom, - >::Error: Into, - { - self.map(|builder| builder.header(key, value)) - } - - pub fn bearer_auth(self, token: T) -> Self - where - T: Display, - { - self.map(|builder| builder.bearer_auth(token)) - } - - pub fn json(self, value: &T) -> Self - where - T: ?Sized + Serialize, - { - self.map(|builder| builder.json(value)) - } - - pub async fn send(self) -> Result { - match self.builder.send().await { - Ok(response) => { - let request_ids = Self::extract_request_ids(&response); - tracing::debug!( - method = %self.method, - url = %self.url, - status = %response.status(), - request_ids = ?request_ids, - version = ?response.version(), - "Request completed" - ); - - Ok(response) - } - Err(error) => { - let status = error.status(); - tracing::debug!( - method = %self.method, - url = %self.url, - status = status.map(|s| s.as_u16()), - error = %error, - "Request failed" - ); - Err(error) - } - } - } - - fn extract_request_ids(response: &Response) -> HashMap { - ["cf-ray", "x-request-id", "x-oai-request-id"] - .iter() - .filter_map(|&name| { - let header_name = HeaderName::from_static(name); - let value = response.headers().get(header_name)?; - let value = value.to_str().ok()?.to_owned(); - Some((name.to_owned(), value)) - }) - .collect() - } -} -#[derive(Debug, Clone)] -pub struct Originator { - pub value: String, - pub header_value: HeaderValue, -} -static ORIGINATOR: OnceLock = OnceLock::new(); - -#[derive(Debug)] -pub enum SetOriginatorError { - InvalidHeaderValue, - AlreadyInitialized, -} - -fn get_originator_value(provided: Option) -> Originator { - let value = std::env::var(CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR) - .ok() - .or(provided) - .unwrap_or(DEFAULT_ORIGINATOR.to_string()); - - match HeaderValue::from_str(&value) { - Ok(header_value) => Originator { - value, - header_value, - }, - Err(e) => { - tracing::error!("Unable to turn originator override {value} into header value: {e}"); - Originator { - value: DEFAULT_ORIGINATOR.to_string(), - header_value: HeaderValue::from_static(DEFAULT_ORIGINATOR), - } - } - } -} - -pub fn set_default_originator(value: String) -> Result<(), SetOriginatorError> { - let originator = get_originator_value(Some(value)); - ORIGINATOR - .set(originator) - .map_err(|_| SetOriginatorError::AlreadyInitialized) -} - -pub fn originator() -> &'static Originator { - ORIGINATOR.get_or_init(|| get_originator_value(None)) -} - -pub fn get_codex_user_agent() -> String { - let build_version = env!("CARGO_PKG_VERSION"); - let os_info = os_info::get(); - let prefix = format!( - "{}/{build_version} ({} {}; {}) {}", - originator().value.as_str(), - os_info.os_type(), - os_info.version(), - os_info.architecture().unwrap_or("unknown"), - crate::terminal::user_agent() - ); - let suffix = USER_AGENT_SUFFIX - .lock() - .ok() - .and_then(|guard| guard.clone()); - let suffix = suffix - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - .map_or_else(String::new, |value| format!(" ({value})")); - - let candidate = format!("{prefix}{suffix}"); - sanitize_user_agent(candidate, &prefix) -} - -/// Sanitize the user agent string. -/// -/// Invalid characters are replaced with an underscore. -/// -/// If the user agent fails to parse, it falls back to fallback and then to ORIGINATOR. -fn sanitize_user_agent(candidate: String, fallback: &str) -> String { - if HeaderValue::from_str(candidate.as_str()).is_ok() { - return candidate; - } - - let sanitized: String = candidate - .chars() - .map(|ch| if matches!(ch, ' '..='~') { ch } else { '_' }) - .collect(); - if !sanitized.is_empty() && HeaderValue::from_str(sanitized.as_str()).is_ok() { - tracing::warn!( - "Sanitized Codex user agent because provided suffix contained invalid header characters" - ); - sanitized - } else if HeaderValue::from_str(fallback).is_ok() { - tracing::warn!( - "Falling back to base Codex user agent because provided suffix could not be sanitized" - ); - fallback.to_string() - } else { - tracing::warn!( - "Falling back to default Codex originator because base user agent string is invalid" - ); - originator().value.clone() - } -} - -/// Create an HTTP client with default `originator` and `User-Agent` headers set. -pub fn create_client() -> CodexHttpClient { - use reqwest::header::HeaderMap; - - let mut headers = HeaderMap::new(); - headers.insert("originator", originator().header_value.clone()); - let ua = get_codex_user_agent(); - - let mut builder = reqwest::Client::builder() - // Set UA via dedicated helper to avoid header validation pitfalls - .user_agent(ua) - .default_headers(headers); - if is_sandboxed() { - builder = builder.no_proxy(); - } - - let inner = builder.build().unwrap_or_else(|_| reqwest::Client::new()); - CodexHttpClient::new(inner) -} - -fn is_sandboxed() -> bool { - std::env::var(CODEX_SANDBOX_ENV_VAR).as_deref() == Ok("seatbelt") -} - -#[cfg(test)] -mod tests { - use super::*; - use core_test_support::skip_if_no_network; - - #[test] - fn test_get_codex_user_agent() { - let user_agent = get_codex_user_agent(); - assert!(user_agent.starts_with("codex_cli_rs/")); - } - - #[tokio::test] - async fn test_create_client_sets_default_headers() { - skip_if_no_network!(); - - use wiremock::Mock; - use wiremock::MockServer; - use wiremock::ResponseTemplate; - use wiremock::matchers::method; - use wiremock::matchers::path; - - let client = create_client(); - - // Spin up a local mock server and capture a request. - let server = MockServer::start().await; - Mock::given(method("GET")) - .and(path("/")) - .respond_with(ResponseTemplate::new(200)) - .mount(&server) - .await; - - let resp = client - .get(server.uri()) - .send() - .await - .expect("failed to send request"); - assert!(resp.status().is_success()); - - let requests = server - .received_requests() - .await - .expect("failed to fetch received requests"); - assert!(!requests.is_empty()); - let headers = &requests[0].headers; - - // originator header is set to the provided value - let originator_header = headers - .get("originator") - .expect("originator header missing"); - assert_eq!(originator_header.to_str().unwrap(), "codex_cli_rs"); - - // User-Agent matches the computed Codex UA for that originator - let expected_ua = get_codex_user_agent(); - let ua_header = headers - .get("user-agent") - .expect("user-agent header missing"); - assert_eq!(ua_header.to_str().unwrap(), expected_ua); - } - - #[test] - fn test_invalid_suffix_is_sanitized() { - let prefix = "codex_cli_rs/0.0.0"; - let suffix = "bad\rsuffix"; - - assert_eq!( - sanitize_user_agent(format!("{prefix} ({suffix})"), prefix), - "codex_cli_rs/0.0.0 (bad_suffix)" - ); - } - - #[test] - fn test_invalid_suffix_is_sanitized2() { - let prefix = "codex_cli_rs/0.0.0"; - let suffix = "bad\0suffix"; - - assert_eq!( - sanitize_user_agent(format!("{prefix} ({suffix})"), prefix), - "codex_cli_rs/0.0.0 (bad_suffix)" - ); - } - - #[test] - #[cfg(target_os = "macos")] - fn test_macos() { - use regex_lite::Regex; - let user_agent = get_codex_user_agent(); - let re = Regex::new( - r"^codex_cli_rs/\d+\.\d+\.\d+ \(Mac OS \d+\.\d+\.\d+; (x86_64|arm64)\) (\S+)$", - ) - .unwrap(); - assert!(re.is_match(&user_agent)); - } -} +pub use crate::client::http::*; diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 5229d00606..3d27950806 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -8,8 +8,7 @@ mod apply_patch; pub mod auth; pub mod bash; -mod chat_completions; -mod client; +pub mod client; mod client_common; pub mod codex; mod codex_conversation; @@ -96,10 +95,10 @@ pub use codex_protocol::protocol; pub use codex_protocol::config_types as protocol_config_types; pub use client::ModelClient; +pub use client::ResponseEvent; +pub use client::ResponseStream; pub use client_common::Prompt; pub use client_common::REVIEW_PROMPT; -pub use client_common::ResponseEvent; -pub use client_common::ResponseStream; pub use codex_protocol::models::ContentItem; pub use codex_protocol::models::LocalShellAction; pub use codex_protocol::models::LocalShellExecAction; diff --git a/codex-rs/core/src/sandboxing/assessment.rs b/codex-rs/core/src/sandboxing/assessment.rs index 31e76777bb..7495c93d3b 100644 --- a/codex-rs/core/src/sandboxing/assessment.rs +++ b/codex-rs/core/src/sandboxing/assessment.rs @@ -6,9 +6,9 @@ use std::time::Instant; use crate::AuthManager; use crate::ModelProviderInfo; +use crate::ResponseEvent; use crate::client::ModelClient; use crate::client_common::Prompt; -use crate::client_common::ResponseEvent; use crate::config::Config; use crate::protocol::SandboxPolicy; use askama::Template; diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 2109f1d2c8..73a00644ab 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -3,10 +3,10 @@ use std::collections::BTreeMap; use crate::apply_patch; use crate::apply_patch::InternalApplyPatchInvocation; use crate::apply_patch::convert_apply_patch_to_protocol; -use crate::client_common::tools::FreeformTool; -use crate::client_common::tools::FreeformToolFormat; -use crate::client_common::tools::ResponsesApiTool; -use crate::client_common::tools::ToolSpec; +use crate::client::FreeformTool; +use crate::client::FreeformToolFormat; +use crate::client::ResponsesApiTool; +use crate::client::ToolSpec; use crate::function_tool::FunctionCallError; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; diff --git a/codex-rs/core/src/tools/handlers/plan.rs b/codex-rs/core/src/tools/handlers/plan.rs index 073319bf1c..182f55a801 100644 --- a/codex-rs/core/src/tools/handlers/plan.rs +++ b/codex-rs/core/src/tools/handlers/plan.rs @@ -1,5 +1,5 @@ -use crate::client_common::tools::ResponsesApiTool; -use crate::client_common::tools::ToolSpec; +use crate::client::ResponsesApiTool; +use crate::client::ToolSpec; use crate::codex::Session; use crate::codex::TurnContext; use crate::function_tool::FunctionCallError; diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index f35ff06315..0946de48dd 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -2,7 +2,7 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; -use crate::client_common::tools::ToolSpec; +use crate::client::ToolSpec; use crate::function_tool::FunctionCallError; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 19098aa80d..fccf3f9066 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -1,7 +1,7 @@ use std::collections::HashMap; use std::sync::Arc; -use crate::client_common::tools::ToolSpec; +use crate::client::ToolSpec; use crate::codex::Session; use crate::codex::TurnContext; use crate::function_tool::FunctionCallError; diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index ab201889ad..5bf79a7384 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -1,5 +1,5 @@ -use crate::client_common::tools::ResponsesApiTool; -use crate::client_common::tools::ToolSpec; +use crate::client::ResponsesApiTool; +use crate::client::ToolSpec; use crate::features::Feature; use crate::features::Features; use crate::model_family::ModelFamily; @@ -1074,7 +1074,7 @@ pub(crate) fn build_specs( #[cfg(test)] mod tests { - use crate::client_common::tools::FreeformTool; + use crate::client::FreeformTool; use crate::model_family::find_family_for_model; use crate::tools::registry::ConfiguredToolSpec; use mcp_types::ToolInputSchema;