//! Session- and turn-scoped helpers for talking to model provider APIs. //! //! `ModelClient` is intended to live for the lifetime of a Codex session and holds the stable //! configuration and state needed to talk to a provider (auth, provider selection, conversation id, //! and feature-gated request behavior). //! //! Per-turn settings (model selection, reasoning controls, telemetry context, web search //! eligibility, and turn metadata) are passed explicitly to streaming and unary methods so that the //! turn lifetime is visible at the call site. //! //! A [`ModelClientSession`] is created per turn and is used to stream one or more Responses API //! requests during that turn. It caches a Responses WebSocket connection (opened lazily) and //! stores per-turn state such as the `x-codex-turn-state` token used for sticky routing. use std::sync::Arc; use std::sync::OnceLock; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use crate::api_bridge::CoreAuthProvider; use crate::api_bridge::auth_provider_from_auth; use crate::api_bridge::map_api_error; use crate::auth::UnauthorizedRecovery; use codex_api::CompactClient as ApiCompactClient; use codex_api::CompactionInput as ApiCompactionInput; use codex_api::MemoriesClient as ApiMemoriesClient; use codex_api::MemoryTrace as ApiMemoryTrace; use codex_api::MemoryTraceSummarizeInput as ApiMemoryTraceSummarizeInput; use codex_api::MemoryTraceSummaryOutput as ApiMemoryTraceSummaryOutput; use codex_api::Prompt as ApiPrompt; use codex_api::RequestTelemetry; use codex_api::ReqwestTransport; use codex_api::ResponseAppendWsRequest; use codex_api::ResponseCreateWsRequest; use codex_api::ResponsesClient as ApiResponsesClient; use codex_api::ResponsesOptions as ApiResponsesOptions; use codex_api::ResponsesWebsocketClient as ApiWebSocketResponsesClient; use codex_api::ResponsesWebsocketConnection as ApiWebSocketConnection; use codex_api::SseTelemetry; use codex_api::TransportError; use codex_api::WebsocketTelemetry; use codex_api::build_conversation_headers; use codex_api::common::Reasoning; use codex_api::common::ResponsesWsRequest; use codex_api::create_text_param_for_request; use codex_api::error::ApiError; use codex_api::requests::responses::Compression; use codex_otel::OtelManager; use codex_protocol::ThreadId; use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig; use codex_protocol::config_types::Verbosity as VerbosityConfig; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelInfo; use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; use codex_protocol::protocol::SessionSource; use eventsource_stream::Event; use eventsource_stream::EventStreamError; use futures::StreamExt; use http::HeaderMap as ApiHeaderMap; use http::HeaderValue; use http::StatusCode as HttpStatusCode; use reqwest::StatusCode; use serde_json::Value; use std::time::Duration; use tokio::sync::mpsc; use tokio_tungstenite::tungstenite::Error; use tokio_tungstenite::tungstenite::Message; use tracing::warn; use crate::AuthManager; use crate::auth::CodexAuth; use crate::auth::RefreshTokenError; use crate::client_common::Prompt; use crate::client_common::ResponseEvent; use crate::client_common::ResponseStream; use crate::default_client::build_reqwest_client; use crate::error::CodexErr; use crate::error::Result; use crate::flags::CODEX_RS_SSE_FIXTURE; use crate::model_provider_info::ModelProviderInfo; use crate::model_provider_info::WireApi; use crate::tools::spec::create_tools_json_for_responses_api; pub const WEB_SEARCH_ELIGIBLE_HEADER: &str = "x-oai-web-search-eligible"; pub const X_CODEX_TURN_STATE_HEADER: &str = "x-codex-turn-state"; pub const X_CODEX_TURN_METADATA_HEADER: &str = "x-codex-turn-metadata"; pub const X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER: &str = "x-responsesapi-include-timing-metrics"; /// Session-scoped state shared by all [`ModelClient`] clones. /// /// This is intentionally kept minimal so `ModelClient` does not need to hold a full `Config`. Most /// configuration is per turn and is passed explicitly to streaming/unary methods. #[derive(Debug)] struct ModelClientState { auth_manager: Option>, conversation_id: ThreadId, provider: ModelProviderInfo, session_source: SessionSource, model_verbosity: Option, enable_responses_websockets: bool, enable_request_compression: bool, include_timing_metrics: bool, beta_features_header: Option, disable_websockets: AtomicBool, } /// A session-scoped client for model-provider API calls. /// /// This holds configuration and state that should be shared across turns within a Codex session /// (auth, provider selection, conversation id, feature-gated request behavior, and transport /// fallback state). /// /// WebSocket fallback is session-scoped: once a turn activates the HTTP fallback, subsequent turns /// will also use HTTP for the remainder of the session. /// /// Turn-scoped settings (model selection, reasoning controls, telemetry context, web search /// eligibility, and turn metadata) are passed explicitly to the relevant methods to keep turn /// lifetime visible at the call site. /// /// This type is cheap to clone. #[derive(Debug, Clone)] pub struct ModelClient { state: Arc, } /// A turn-scoped streaming session created from a [`ModelClient`]. /// /// The session lazily establishes a Responses WebSocket connection (and reuses it across multiple /// requests) and caches per-turn state: /// /// - The last request's input items, so subsequent calls can use `response.append` when the input /// is an incremental extension of the previous request. /// - The `x-codex-turn-state` sticky-routing token, which must be replayed for all requests within /// the same turn. /// /// Create a fresh `ModelClientSession` for each Codex turn. Reusing it across turns would replay /// the previous turn's sticky-routing token into the next turn, which violates the client/server /// contract and can cause routing bugs. pub struct ModelClientSession { client: ModelClient, connection: Option, websocket_last_items: Vec, /// Turn state for sticky routing. /// /// This is an `OnceLock` that stores the turn state value received from the server /// on turn start via the `x-codex-turn-state` response header. Once set, this value /// should be sent back to the server in the `x-codex-turn-state` request header for /// all subsequent requests within the same turn to maintain sticky routing. /// /// This is a contract between the client and server: we receive it at turn start, /// keep sending it unchanged between turn requests (e.g., for retries, incremental /// appends, or continuation requests), and must not send it between different turns. turn_state: Arc>, } impl ModelClient { #[allow(clippy::too_many_arguments)] /// Creates a new session-scoped `ModelClient`. /// /// All arguments are expected to be stable for the lifetime of a Codex session. Per-turn values /// are passed to [`ModelClientSession::stream`] (and other turn-scoped methods) explicitly. pub fn new( auth_manager: Option>, conversation_id: ThreadId, provider: ModelProviderInfo, session_source: SessionSource, model_verbosity: Option, enable_responses_websockets: bool, enable_request_compression: bool, include_timing_metrics: bool, beta_features_header: Option, ) -> Self { Self { state: Arc::new(ModelClientState { auth_manager, conversation_id, provider, session_source, model_verbosity, enable_responses_websockets, enable_request_compression, include_timing_metrics, beta_features_header, disable_websockets: AtomicBool::new(false), }), } } /// Creates a fresh turn-scoped streaming session. /// /// This does not open any network connections; the WebSocket connection is established lazily /// when the first WebSocket stream request is issued. pub fn new_session(&self) -> ModelClientSession { ModelClientSession { client: self.clone(), connection: None, websocket_last_items: Vec::new(), turn_state: Arc::new(OnceLock::new()), } } /// Compacts the current conversation history using the Compact endpoint. /// /// This is a unary call (no streaming) that returns a new list of /// `ResponseItem`s representing the compacted transcript. /// /// The model selection and telemetry context are passed explicitly to keep `ModelClient` /// session-scoped. pub async fn compact_conversation_history( &self, prompt: &Prompt, model_info: &ModelInfo, otel_manager: &OtelManager, ) -> Result> { if prompt.input.is_empty() { return Ok(Vec::new()); } let auth_manager = self.state.auth_manager.clone(); let auth = match auth_manager.as_ref() { Some(manager) => manager.auth().await, None => None, }; let api_provider = self .state .provider .to_api_provider(auth.as_ref().map(CodexAuth::internal_auth_mode))?; let api_auth = auth_provider_from_auth(auth.clone(), &self.state.provider)?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_telemetry = Self::build_request_telemetry(otel_manager); let client = ApiCompactClient::new(transport, api_provider, api_auth) .with_telemetry(Some(request_telemetry)); let instructions = prompt.base_instructions.text.clone(); let payload = ApiCompactionInput { model: &model_info.slug, input: &prompt.input, instructions: &instructions, }; let extra_headers = self.build_subagent_headers(); client .compact_input(&payload, extra_headers) .await .map_err(map_api_error) } /// Builds memory summaries for each provided normalized trace. /// /// This is a unary call (no streaming) to `/v1/memories/trace_summarize`. /// /// The model selection, reasoning effort, and telemetry context are passed explicitly to keep /// `ModelClient` session-scoped. pub async fn summarize_memory_traces( &self, traces: Vec, model_info: &ModelInfo, effort: Option, otel_manager: &OtelManager, ) -> Result> { if traces.is_empty() { return Ok(Vec::new()); } let auth_manager = self.state.auth_manager.clone(); let auth = match auth_manager.as_ref() { Some(manager) => manager.auth().await, None => None, }; let api_provider = self .state .provider .to_api_provider(auth.as_ref().map(CodexAuth::internal_auth_mode))?; let api_auth = auth_provider_from_auth(auth, &self.state.provider)?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_telemetry = Self::build_request_telemetry(otel_manager); let client = ApiMemoriesClient::new(transport, api_provider, api_auth) .with_telemetry(Some(request_telemetry)); let payload = ApiMemoryTraceSummarizeInput { model: model_info.slug.clone(), traces, reasoning: effort.map(|effort| Reasoning { effort: Some(effort), summary: None, }), }; client .trace_summarize_input(&payload, self.build_subagent_headers()) .await .map_err(map_api_error) } fn build_subagent_headers(&self) -> ApiHeaderMap { let mut extra_headers = ApiHeaderMap::new(); if let SessionSource::SubAgent(sub) = &self.state.session_source { let subagent = match sub { crate::protocol::SubAgentSource::Review => "review".to_string(), crate::protocol::SubAgentSource::Compact => "compact".to_string(), crate::protocol::SubAgentSource::ThreadSpawn { .. } => "collab_spawn".to_string(), crate::protocol::SubAgentSource::Other(label) => label.clone(), }; if let Ok(val) = HeaderValue::from_str(&subagent) { extra_headers.insert("x-openai-subagent", val); } } extra_headers } /// Builds request telemetry for unary API calls (e.g., Compact endpoint). fn build_request_telemetry(otel_manager: &OtelManager) -> Arc { let telemetry = Arc::new(ApiTelemetry::new(otel_manager.clone())); let request_telemetry: Arc = telemetry; request_telemetry } } impl ModelClientSession { fn disable_websockets(&self) -> bool { self.client.state.disable_websockets.load(Ordering::Relaxed) } fn activate_http_fallback(&self, websocket_enabled: bool) -> bool { websocket_enabled && !self .client .state .disable_websockets .swap(true, Ordering::Relaxed) } fn responses_websocket_enabled(&self) -> bool { self.client.state.provider.supports_websockets && self.client.state.enable_responses_websockets } fn build_responses_request(prompt: &Prompt) -> Result { let instructions = prompt.base_instructions.text.clone(); let tools_json: Vec = create_tools_json_for_responses_api(&prompt.tools)?; Ok(build_api_prompt(prompt, instructions, tools_json)) } #[allow(clippy::too_many_arguments)] fn build_responses_options( &self, prompt: &Prompt, model_info: &ModelInfo, effort: Option, summary: ReasoningSummaryConfig, web_search_eligible: bool, turn_metadata_header: Option<&str>, compression: Compression, ) -> ApiResponsesOptions { let turn_metadata_header = turn_metadata_header.and_then(|value| HeaderValue::from_str(value).ok()); let default_reasoning_effort = model_info.default_reasoning_level; let reasoning = if model_info.supports_reasoning_summaries { Some(Reasoning { effort: effort.or(default_reasoning_effort), summary: if summary == ReasoningSummaryConfig::None { None } else { Some(summary) }, }) } else { None }; let include = if reasoning.is_some() { vec!["reasoning.encrypted_content".to_string()] } else { Vec::new() }; let verbosity = if model_info.support_verbosity { self.client .state .model_verbosity .or(model_info.default_verbosity) } else { if self.client.state.model_verbosity.is_some() { warn!( "model_verbosity is set but ignored as the model does not support verbosity: {}", model_info.slug ); } None }; let text = create_text_param_for_request(verbosity, &prompt.output_schema); let conversation_id = self.client.state.conversation_id.to_string(); ApiResponsesOptions { reasoning, include, prompt_cache_key: Some(conversation_id.clone()), text, store_override: None, conversation_id: Some(conversation_id), session_source: Some(self.client.state.session_source.clone()), extra_headers: build_responses_headers( self.client.state.beta_features_header.as_deref(), web_search_eligible, Some(&self.turn_state), turn_metadata_header.as_ref(), ), compression, turn_state: Some(Arc::clone(&self.turn_state)), } } fn get_incremental_items(&self, input_items: &[ResponseItem]) -> Option> { // Checks whether the current request input is an incremental append to the previous request. // If items in the new request contain all the items from the previous request we build // a response.append request otherwise we start with a fresh response.create request. let previous_len = self.websocket_last_items.len(); let can_append = previous_len > 0 && input_items.starts_with(&self.websocket_last_items) && previous_len < input_items.len(); if can_append { Some(input_items[previous_len..].to_vec()) } else { None } } fn prepare_websocket_request( &self, model_slug: &str, api_prompt: &ApiPrompt, options: &ApiResponsesOptions, ) -> ResponsesWsRequest { if let Some(append_items) = self.get_incremental_items(&api_prompt.input) { return ResponsesWsRequest::ResponseAppend(ResponseAppendWsRequest { input: append_items, }); } let ApiResponsesOptions { reasoning, include, prompt_cache_key, text, store_override, .. } = options; let store = store_override.unwrap_or(false); let payload = ResponseCreateWsRequest { model: model_slug.to_string(), instructions: api_prompt.instructions.clone(), input: api_prompt.input.clone(), tools: api_prompt.tools.clone(), tool_choice: "auto".to_string(), parallel_tool_calls: api_prompt.parallel_tool_calls, reasoning: reasoning.clone(), store, stream: true, include: include.clone(), prompt_cache_key: prompt_cache_key.clone(), text: text.clone(), }; ResponsesWsRequest::ResponseCreate(payload) } async fn websocket_connection( &mut self, otel_manager: &OtelManager, api_provider: codex_api::Provider, api_auth: CoreAuthProvider, options: &ApiResponsesOptions, ) -> std::result::Result<&ApiWebSocketConnection, ApiError> { let needs_new = match self.connection.as_ref() { Some(conn) => conn.is_closed().await, None => true, }; if needs_new { let mut headers = options.extra_headers.clone(); headers.extend(build_conversation_headers(options.conversation_id.clone())); if self.client.state.include_timing_metrics { headers.insert( X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER, HeaderValue::from_static("true"), ); } let websocket_telemetry = Self::build_websocket_telemetry(otel_manager); let new_conn: ApiWebSocketConnection = ApiWebSocketResponsesClient::new(api_provider, api_auth) .connect( headers, options.turn_state.clone(), Some(websocket_telemetry), ) .await?; self.connection = Some(new_conn); } self.connection.as_ref().ok_or(ApiError::Stream( "websocket connection is unavailable".to_string(), )) } fn responses_request_compression(&self, auth: Option<&crate::auth::CodexAuth>) -> Compression { if self.client.state.enable_request_compression && auth.is_some_and(CodexAuth::is_chatgpt_auth) && self.client.state.provider.is_openai() { Compression::Zstd } else { Compression::None } } /// Streams a turn via the OpenAI Responses API. /// /// Handles SSE fixtures, reasoning summaries, verbosity, and the /// `text` controls used for output schemas. #[allow(clippy::too_many_arguments)] async fn stream_responses_api( &self, prompt: &Prompt, model_info: &ModelInfo, otel_manager: &OtelManager, effort: Option, summary: ReasoningSummaryConfig, web_search_eligible: bool, turn_metadata_header: Option<&str>, ) -> Result { if let Some(path) = &*CODEX_RS_SSE_FIXTURE { warn!(path, "Streaming from fixture"); let stream = codex_api::stream_from_fixture( path, self.client.state.provider.stream_idle_timeout(), ) .map_err(map_api_error)?; return Ok(map_response_stream(stream, otel_manager.clone())); } let auth_manager = self.client.state.auth_manager.clone(); let api_prompt = Self::build_responses_request(prompt)?; let mut auth_recovery = auth_manager .as_ref() .map(super::auth::AuthManager::unauthorized_recovery); loop { let auth = match auth_manager.as_ref() { Some(manager) => manager.auth().await, None => None, }; let api_provider = self .client .state .provider .to_api_provider(auth.as_ref().map(CodexAuth::internal_auth_mode))?; let api_auth = auth_provider_from_auth(auth.clone(), &self.client.state.provider)?; let transport = ReqwestTransport::new(build_reqwest_client()); let (request_telemetry, sse_telemetry) = Self::build_streaming_telemetry(otel_manager); let compression = self.responses_request_compression(auth.as_ref()); let client = ApiResponsesClient::new(transport, api_provider, api_auth) .with_telemetry(Some(request_telemetry), Some(sse_telemetry)); let options = self.build_responses_options( prompt, model_info, effort, summary, web_search_eligible, turn_metadata_header, compression, ); let stream_result = client .stream_prompt(&model_info.slug, &api_prompt, options) .await; match stream_result { Ok(stream) => { return Ok(map_response_stream(stream, otel_manager.clone())); } Err(ApiError::Transport( unauthorized_transport @ TransportError::Http { status, .. }, )) if status == StatusCode::UNAUTHORIZED => { handle_unauthorized(unauthorized_transport, &mut auth_recovery).await?; continue; } Err(err) => return Err(map_api_error(err)), } } } /// Streams a turn via the Responses API over WebSocket transport. #[allow(clippy::too_many_arguments)] async fn stream_responses_websocket( &mut self, prompt: &Prompt, model_info: &ModelInfo, otel_manager: &OtelManager, effort: Option, summary: ReasoningSummaryConfig, web_search_eligible: bool, turn_metadata_header: Option<&str>, ) -> Result { let auth_manager = self.client.state.auth_manager.clone(); let api_prompt = Self::build_responses_request(prompt)?; let mut auth_recovery = auth_manager .as_ref() .map(super::auth::AuthManager::unauthorized_recovery); loop { let auth = match auth_manager.as_ref() { Some(manager) => manager.auth().await, None => None, }; let api_provider = self .client .state .provider .to_api_provider(auth.as_ref().map(CodexAuth::internal_auth_mode))?; let api_auth = auth_provider_from_auth(auth.clone(), &self.client.state.provider)?; let compression = self.responses_request_compression(auth.as_ref()); let options = self.build_responses_options( prompt, model_info, effort, summary, web_search_eligible, turn_metadata_header, compression, ); let request = self.prepare_websocket_request(&model_info.slug, &api_prompt, &options); let connection = match self .websocket_connection( otel_manager, api_provider.clone(), api_auth.clone(), &options, ) .await { Ok(connection) => connection, Err(ApiError::Transport( unauthorized_transport @ TransportError::Http { status, .. }, )) if status == StatusCode::UNAUTHORIZED => { handle_unauthorized(unauthorized_transport, &mut auth_recovery).await?; continue; } Err(err) => return Err(map_api_error(err)), }; let stream_result = connection .stream_request(request) .await .map_err(map_api_error)?; self.websocket_last_items = api_prompt.input.clone(); return Ok(map_response_stream(stream_result, otel_manager.clone())); } } /// Builds request and SSE telemetry for streaming API calls. fn build_streaming_telemetry( otel_manager: &OtelManager, ) -> (Arc, Arc) { let telemetry = Arc::new(ApiTelemetry::new(otel_manager.clone())); let request_telemetry: Arc = telemetry.clone(); let sse_telemetry: Arc = telemetry; (request_telemetry, sse_telemetry) } /// Builds telemetry for the Responses API WebSocket transport. fn build_websocket_telemetry(otel_manager: &OtelManager) -> Arc { let telemetry = Arc::new(ApiTelemetry::new(otel_manager.clone())); let websocket_telemetry: Arc = telemetry; websocket_telemetry } #[allow(clippy::too_many_arguments)] /// Streams a single model request within the current turn. /// /// The caller is responsible for passing per-turn settings explicitly (model selection, /// reasoning settings, telemetry context, web search eligibility, and turn metadata). This /// method will prefer the Responses WebSocket transport when enabled and healthy, and will /// fall back to the HTTP Responses API transport otherwise. pub async fn stream( &mut self, prompt: &Prompt, model_info: &ModelInfo, otel_manager: &OtelManager, effort: Option, summary: ReasoningSummaryConfig, web_search_eligible: bool, turn_metadata_header: Option<&str>, ) -> Result { let wire_api = self.client.state.provider.wire_api; match wire_api { WireApi::Responses => { let websocket_enabled = self.responses_websocket_enabled() && !self.disable_websockets(); if websocket_enabled { self.stream_responses_websocket( prompt, model_info, otel_manager, effort, summary, web_search_eligible, turn_metadata_header, ) .await } else { self.stream_responses_api( prompt, model_info, otel_manager, effort, summary, web_search_eligible, turn_metadata_header, ) .await } } } } /// Permanently disables WebSockets for this Codex session and resets WebSocket state. /// /// This is used after exhausting the provider retry budget, to force subsequent requests onto /// the HTTP transport. Returns `true` if this call activated fallback, or `false` if fallback /// was already active. pub(crate) fn try_switch_fallback_transport(&mut self, otel_manager: &OtelManager) -> bool { let websocket_enabled = self.responses_websocket_enabled(); let activated = self.activate_http_fallback(websocket_enabled); if activated { warn!("falling back to HTTP"); otel_manager.counter( "codex.transport.fallback_to_http", 1, &[("from_wire_api", "responses_websocket")], ); self.connection = None; self.websocket_last_items.clear(); } activated } } /// Adapts the core `Prompt` type into the `codex-api` payload shape. fn build_api_prompt(prompt: &Prompt, instructions: String, tools_json: Vec) -> ApiPrompt { ApiPrompt { instructions, input: prompt.get_formatted_input(), tools: tools_json, parallel_tool_calls: prompt.parallel_tool_calls, output_schema: prompt.output_schema.clone(), } } /// Builds the extra headers attached to Responses API requests. /// /// These headers implement Codex-specific conventions: /// /// - `x-codex-beta-features`: comma-separated beta feature keys enabled for the session. /// - `x-oai-web-search-eligible`: whether this turn is allowed to use web search. /// - `x-codex-turn-state`: sticky routing token captured earlier in the turn. /// - `x-codex-turn-metadata`: optional per-turn metadata for observability. fn build_responses_headers( beta_features_header: Option<&str>, web_search_eligible: bool, turn_state: Option<&Arc>>, turn_metadata_header: Option<&HeaderValue>, ) -> ApiHeaderMap { let mut headers = ApiHeaderMap::new(); if let Some(value) = beta_features_header && !value.is_empty() && let Ok(header_value) = HeaderValue::from_str(value) { headers.insert("x-codex-beta-features", header_value); } headers.insert( WEB_SEARCH_ELIGIBLE_HEADER, HeaderValue::from_static(if web_search_eligible { "true" } else { "false" }), ); if let Some(turn_state) = turn_state && let Some(state) = turn_state.get() && let Ok(header_value) = HeaderValue::from_str(state) { headers.insert(X_CODEX_TURN_STATE_HEADER, header_value); } if let Some(header_value) = turn_metadata_header { headers.insert(X_CODEX_TURN_METADATA_HEADER, header_value.clone()); } headers } fn map_response_stream(api_stream: S, otel_manager: OtelManager) -> ResponseStream where S: futures::Stream> + Unpin + Send + 'static, { let (tx_event, rx_event) = mpsc::channel::>(1600); tokio::spawn(async move { let mut logged_error = false; let mut api_stream = api_stream; while let Some(event) = api_stream.next().await { match event { Ok(ResponseEvent::Completed { response_id, token_usage, }) => { if let Some(usage) = &token_usage { otel_manager.sse_event_completed( usage.input_tokens, usage.output_tokens, Some(usage.cached_input_tokens), Some(usage.reasoning_output_tokens), usage.total_tokens, ); } if tx_event .send(Ok(ResponseEvent::Completed { response_id, token_usage, })) .await .is_err() { return; } } Ok(event) => { if tx_event.send(Ok(event)).await.is_err() { return; } } Err(err) => { let mapped = map_api_error(err); if !logged_error { otel_manager.see_event_completed_failed(&mapped); logged_error = true; } if tx_event.send(Err(mapped)).await.is_err() { return; } } } } }); ResponseStream { rx_event } } /// Handles a 401 response by optionally refreshing ChatGPT tokens once. /// /// When refresh succeeds, the caller should retry the API call; otherwise /// the mapped `CodexErr` is returned to the caller. async fn handle_unauthorized( transport: TransportError, auth_recovery: &mut Option, ) -> Result<()> { if let Some(recovery) = auth_recovery && recovery.has_next() { return match recovery.next().await { Ok(_) => Ok(()), Err(RefreshTokenError::Permanent(failed)) => Err(CodexErr::RefreshTokenFailed(failed)), Err(RefreshTokenError::Transient(other)) => Err(CodexErr::Io(other)), }; } Err(map_api_error(ApiError::Transport(transport))) } struct ApiTelemetry { otel_manager: OtelManager, } impl ApiTelemetry { fn new(otel_manager: OtelManager) -> Self { Self { otel_manager } } } impl RequestTelemetry for ApiTelemetry { fn on_request( &self, attempt: u64, status: Option, error: Option<&TransportError>, duration: Duration, ) { let error_message = error.map(std::string::ToString::to_string); self.otel_manager.record_api_request( attempt, status.map(|s| s.as_u16()), error_message.as_deref(), duration, ); } } impl SseTelemetry for ApiTelemetry { fn on_sse_poll( &self, result: &std::result::Result< Option>>, tokio::time::error::Elapsed, >, duration: Duration, ) { self.otel_manager.log_sse_event(result, duration); } } impl WebsocketTelemetry for ApiTelemetry { fn on_ws_request(&self, duration: Duration, error: Option<&ApiError>) { let error_message = error.map(std::string::ToString::to_string); self.otel_manager .record_websocket_request(duration, error_message.as_deref()); } fn on_ws_event( &self, result: &std::result::Result>, ApiError>, duration: Duration, ) { self.otel_manager.record_websocket_event(result, duration); } }