use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::OnceLock;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use codex_api::AgentIdentityTelemetry;
use codex_api::ApiError;
use codex_api::AuthProvider;
use codex_api::CompactClient as ApiCompactClient;
use codex_api::CompactionInput as ApiCompactionInput;
use codex_api::Compression;
use codex_api::MemoriesClient as ApiMemoriesClient;
use codex_api::MemorySummarizeInput as ApiMemorySummarizeInput;
use codex_api::MemorySummarizeOutput as ApiMemorySummarizeOutput;
use codex_api::Provider as ApiProvider;
use codex_api::RawMemory as ApiRawMemory;
use codex_api::RealtimeCallClient as ApiRealtimeCallClient;
use codex_api::RealtimeSessionConfig as ApiRealtimeSessionConfig;
use codex_api::Reasoning;
use codex_api::ReasoningContext;
use codex_api::RequestTelemetry;
use codex_api::ReqwestTransport;
use codex_api::ResponseCreateWsRequest;
use codex_api::ResponsesApiRequest;
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::ResponsesWsRequest;
use codex_api::SharedAuthProvider;
use codex_api::SseTelemetry;
use codex_api::StreamOptions;
use codex_api::TransportError;
use codex_api::WebsocketTelemetry;
use codex_api::auth_header_telemetry;
use codex_api::build_session_headers;
use codex_api::create_text_param_for_request;
use codex_api::response_create_client_metadata;
use codex_http_client::ClientRouteClass;
use codex_http_client::HttpClientFactory;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_login::RefreshTokenError;
use codex_login::UnauthorizedRecovery;
use codex_login::default_client::add_originator_header;
use codex_login::default_client::create_client_for_route;
use codex_otel::SessionTelemetry;
use codex_otel::current_span_w3c_trace_context;
use codex_protocol::auth::AuthMode;
use codex_protocol::ThreadId;
use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig;
use codex_protocol::config_types::Verbosity as VerbosityConfig;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig;
use codex_protocol::protocol::InternalSessionSource;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::W3cTraceContext;
use codex_rollout_trace::CompactionTraceContext;
use codex_rollout_trace::InferenceTraceAttempt;
use codex_rollout_trace::InferenceTraceContext;
use codex_tools::create_tools_json_for_responses_api;
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 std::time::Duration;
use std::time::Instant;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::sync::oneshot::error::TryRecvError;
use tokio_tungstenite::tungstenite::Error;
use tokio_tungstenite::tungstenite::Message;
use tokio_util::sync::CancellationToken;
use tracing::instrument;
use tracing::trace;
use tracing::warn;
use crate::attestation::AttestationContext;
use crate::attestation::AttestationProvider;
use crate::attestation::X_OAI_ATTESTATION_HEADER;
use crate::client_common::Prompt;
use crate::client_common::ResponseEvent;
use crate::client_common::ResponseStream;
use crate::feedback_tags;
use crate::responses_metadata::CodexResponsesMetadata;
use crate::responses_metadata::subagent_header_value;
use crate::util::emit_feedback_auth_recovery_tags;
use codex_feedback::FeedbackRequestTags;
use codex_feedback::emit_feedback_request_tags_with_auth_env;
use codex_login::auth::AgentIdentityAuthPolicy;
use codex_login::auth_env_telemetry::AuthEnvTelemetry;
use codex_login::auth_env_telemetry::collect_auth_env_telemetry;
use codex_model_provider::AgentIdentitySessionFallback;
use codex_model_provider::ProviderAuthScope;
use codex_model_provider::SharedModelProvider;
use codex_model_provider::create_model_provider;
#[cfg(test)]
use codex_model_provider_info::DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS;
use codex_model_provider_info::ModelProviderInfo;
use codex_model_provider_info::WireApi;
use codex_protocol::error::CodexErr;
use codex_protocol::error::Result;
use codex_response_debug_context::extract_response_debug_context;
use codex_response_debug_context::extract_response_debug_context_from_api_error;
use codex_response_debug_context::telemetry_api_error_message;
use codex_response_debug_context::telemetry_transport_error_message;
pub const OPENAI_BETA_HEADER: &str = "OpenAI-Beta";
pub const X_CODEX_INSTALLATION_ID_HEADER: &str = "x-codex-installation-id";
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_CODEX_PARENT_THREAD_ID_HEADER: &str = "x-codex-parent-thread-id";
pub const X_CODEX_WINDOW_ID_HEADER: &str = "x-codex-window-id";
pub const X_OPENAI_MEMGEN_REQUEST_HEADER: &str = "x-openai-memgen-request";
pub const X_OPENAI_SUBAGENT_HEADER: &str = "x-openai-subagent";
pub const X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER: &str =
"x-responsesapi-include-timing-metrics";
const X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY: &str =
"x-codex-ws-stream-request-start-ms";
const WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY: &str =
"ws_request_header_x_openai_internal_codex_responses_lite";
const RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE: &str = "responses_websockets=2026-02-06";
const X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER: &str =
"x-openai-internal-codex-responses-lite";
const REALTIME_CALLS_ENDPOINT: &str = "/realtime/calls";
const RESPONSES_ENDPOINT: &str = "/responses";
const RESPONSES_COMPACT_ENDPOINT: &str = "/responses/compact";
const COMPACT_REQUEST_TIMEOUT_IDLE_MULTIPLIER: u32 = 4;
const MEMORIES_SUMMARIZE_ENDPOINT: &str = "/memories/trace_summarize";
#[cfg(test)]
pub(crate) const WEBSOCKET_CONNECT_TIMEOUT: Duration =
Duration::from_millis(DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS);
pub(crate) struct CompactConversationRequestSettings {
pub(crate) effort: Option<ReasoningEffortConfig>,
pub(crate) summary: ReasoningSummaryConfig,
pub(crate) service_tier: Option<String>,
}
fn reasoning_effort_for_request(effort: ReasoningEffortConfig) -> ReasoningEffortConfig {
match effort {
ReasoningEffortConfig::Ultra => ReasoningEffortConfig::Max,
effort => effort,
}
}
fn session_telemetry_for_request(
session_telemetry: &SessionTelemetry,
request: &ResponsesApiRequest,
) -> SessionTelemetry {
session_telemetry.clone().with_inference_request(
request.service_tier.as_deref(),
request
.reasoning
.as_ref()
.and_then(|reasoning| reasoning.effort.as_ref()),
)
}
#[derive(Debug)]
struct ModelClientState {
thread_id: ThreadId,
provider: SharedModelProvider,
auth_env_telemetry: AuthEnvTelemetry,
session_source: SessionSource,
originator: String,
model_verbosity: Option<VerbosityConfig>,
enable_request_compression: bool,
include_timing_metrics: bool,
beta_features_header: Option<String>,
concurrent_reasoning_summaries_enabled: bool,
include_attestation: bool,
attestation_provider: Option<Arc<dyn AttestationProvider>>,
disable_websockets: AtomicBool,
agent_identity_session_fallback: AgentIdentitySessionFallback,
cached_websocket_session: StdMutex<WebsocketSession>,
}
struct CurrentClientSetup {
auth: Option<CodexAuth>,
api_provider: ApiProvider,
api_auth: SharedAuthProvider,
agent_identity_telemetry: Option<AgentIdentityTelemetry>,
}
#[derive(Clone, Copy)]
struct RequestRouteTelemetry {
endpoint: &'static str,
}
impl RequestRouteTelemetry {
fn for_endpoint(endpoint: &'static str) -> Self {
Self { endpoint }
}
}
#[derive(Debug, Clone)]
pub struct ModelClient {
state: Arc<ModelClientState>,
agent_identity_policy: AgentIdentityAuthPolicy,
prompt_cache_key_override: Option<String>,
http_client_factory: HttpClientFactory,
}
pub struct ModelClientSession {
client: ModelClient,
websocket_session: WebsocketSession,
turn_state: Arc<OnceLock<String>>,
}
#[derive(Debug, Clone)]
struct LastResponse {
response_id: String,
items_added: Vec<ResponseItem>,
}
#[derive(Debug, Default)]
struct WebsocketSession {
connection: Option<ApiWebSocketConnection>,
last_request: Option<ResponsesApiRequest>,
last_response_rx: Option<oneshot::Receiver<LastResponse>>,
last_response_from_untraced_warmup: bool,
connection_reused: StdMutex<bool>,
}
fn responses_request_properties_match(
previous: &ResponsesApiRequest,
current: &ResponsesApiRequest,
) -> bool {
let ResponsesApiRequest {
model: previous_model,
instructions: previous_instructions,
input: _,
tools: previous_tools,
tool_choice: previous_tool_choice,
parallel_tool_calls: previous_parallel_tool_calls,
reasoning: previous_reasoning,
store: previous_store,
stream: previous_stream,
stream_options: _,
include: previous_include,
service_tier: previous_service_tier,
prompt_cache_key: previous_prompt_cache_key,
text: previous_text,
client_metadata: _,
} = previous;
let ResponsesApiRequest {
model: current_model,
instructions: current_instructions,
input: _,
tools: current_tools,
tool_choice: current_tool_choice,
parallel_tool_calls: current_parallel_tool_calls,
reasoning: current_reasoning,
store: current_store,
stream: current_stream,
stream_options: _,
include: current_include,
service_tier: current_service_tier,
prompt_cache_key: current_prompt_cache_key,
text: current_text,
client_metadata: _,
} = current;
previous_model == current_model
&& previous_instructions == current_instructions
&& previous_tools == current_tools
&& previous_tool_choice == current_tool_choice
&& previous_parallel_tool_calls == current_parallel_tool_calls
&& previous_reasoning == current_reasoning
&& previous_store == current_store
&& previous_stream == current_stream
&& previous_include == current_include
&& previous_service_tier == current_service_tier
&& previous_prompt_cache_key == current_prompt_cache_key
&& previous_text == current_text
}
impl WebsocketSession {
fn set_connection_reused(&self, connection_reused: bool) {
*self
.connection_reused
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = connection_reused;
}
fn connection_reused(&self) -> bool {
*self
.connection_reused
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
}
enum WebsocketStreamOutcome {
Stream(ResponseStream),
FallbackToHttp,
}
pub(crate) struct RealtimeWebrtcCallStart {
pub(crate) sdp: String,
pub(crate) call_id: String,
pub(crate) sideband_headers: ApiHeaderMap,
}
fn sideband_websocket_auth_headers(api_auth: &dyn AuthProvider) -> ApiHeaderMap {
let mut headers = ApiHeaderMap::new();
api_auth.add_auth_headers(&mut headers);
headers
}
impl ModelClient {
#[allow(clippy::too_many_arguments)]
pub fn new(
auth_manager: Option<Arc<AuthManager>>,
agent_identity_policy: AgentIdentityAuthPolicy,
thread_id: ThreadId,
provider_info: ModelProviderInfo,
session_source: SessionSource,
originator: String,
model_verbosity: Option<VerbosityConfig>,
enable_request_compression: bool,
include_timing_metrics: bool,
beta_features_header: Option<String>,
concurrent_reasoning_summaries_enabled: bool,
attestation_provider: Option<Arc<dyn AttestationProvider>>,
http_client_factory: HttpClientFactory,
) -> Self {
let model_provider = create_model_provider(provider_info, auth_manager);
let codex_api_key_env_enabled = model_provider
.auth_manager()
.as_ref()
.is_some_and(|manager| manager.codex_api_key_env_enabled());
let auth_env_telemetry =
collect_auth_env_telemetry(model_provider.info(), codex_api_key_env_enabled);
let include_attestation = model_provider.supports_attestation();
Self {
state: Arc::new(ModelClientState {
thread_id,
provider: model_provider,
auth_env_telemetry,
session_source,
originator,
model_verbosity,
enable_request_compression,
include_timing_metrics,
beta_features_header,
concurrent_reasoning_summaries_enabled,
include_attestation,
attestation_provider,
disable_websockets: AtomicBool::new(false),
agent_identity_session_fallback: AgentIdentitySessionFallback::default(),
cached_websocket_session: StdMutex::new(WebsocketSession::default()),
}),
agent_identity_policy,
prompt_cache_key_override: None,
http_client_factory,
}
}
pub(crate) fn with_prompt_cache_key_override(
mut self,
prompt_cache_key_override: Option<String>,
) -> Self {
self.prompt_cache_key_override = prompt_cache_key_override;
self
}
fn prompt_cache_key(&self, responses_metadata: &CodexResponsesMetadata) -> String {
self.prompt_cache_key_override
.clone()
.unwrap_or_else(|| responses_metadata.session_id.clone())
}
pub fn new_session(&self) -> ModelClientSession {
ModelClientSession {
client: self.clone(),
websocket_session: self.take_cached_websocket_session(),
turn_state: Arc::new(OnceLock::new()),
}
}
pub(crate) fn auth_manager(&self) -> Option<Arc<AuthManager>> {
self.state.provider.auth_manager()
}
fn take_cached_websocket_session(&self) -> WebsocketSession {
let mut cached_websocket_session = self
.state
.cached_websocket_session
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
std::mem::take(&mut *cached_websocket_session)
}
fn store_cached_websocket_session(&self, websocket_session: WebsocketSession) {
*self
.state
.cached_websocket_session
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = websocket_session;
}
pub(crate) fn force_http_fallback(
&self,
session_telemetry: &SessionTelemetry,
_model_info: &ModelInfo,
) -> bool {
let websocket_enabled = self.responses_websocket_enabled();
let activated =
websocket_enabled && !self.state.disable_websockets.swap(true, Ordering::Relaxed);
if activated {
warn!("falling back to HTTP");
session_telemetry.counter(
"codex.transport.fallback_to_http",
1,
&[("from_wire_api", "responses_websocket")],
);
}
self.store_cached_websocket_session(WebsocketSession::default());
activated
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn compact_conversation_history(
&self,
prompt: &Prompt,
model_info: &ModelInfo,
turn_state: Option<Arc<OnceLock<String>>>,
settings: CompactConversationRequestSettings,
session_telemetry: &SessionTelemetry,
compaction_trace: &CompactionTraceContext,
responses_metadata: &CodexResponsesMetadata,
) -> Result<Vec<ResponseItem>> {
if prompt.input.is_empty() {
return Ok(Vec::new());
}
let client_setup = self.current_client_setup().await?;
let transport =
self.build_api_transport(&client_setup.api_provider, RESPONSES_COMPACT_ENDPOINT)?;
let request_telemetry = Self::build_request_telemetry(
session_telemetry,
AuthRequestTelemetryContext::new(
client_setup.auth.as_ref().map(CodexAuth::auth_mode),
client_setup.api_auth.as_ref(),
client_setup.agent_identity_telemetry.clone(),
PendingUnauthorizedRetry::default(),
),
RequestRouteTelemetry::for_endpoint(RESPONSES_COMPACT_ENDPOINT),
self.state.auth_env_telemetry.clone(),
);
let request = self.build_responses_request(
&client_setup.api_provider,
prompt,
model_info,
settings.effort,
settings.summary,
settings.service_tier,
responses_metadata,
)?;
let ResponsesApiRequest {
model,
instructions,
mut input,
tools,
parallel_tool_calls,
reasoning,
service_tier,
prompt_cache_key,
text,
..
} = request;
self.prepare_response_items_for_request(&mut input);
let payload = ApiCompactionInput {
model: &model,
input: &input,
instructions: &instructions,
tools,
parallel_tool_calls,
reasoning,
service_tier: service_tier.as_deref(),
prompt_cache_key: prompt_cache_key.as_deref(),
text,
};
let mut extra_headers = ApiHeaderMap::new();
if let Ok(header_value) = HeaderValue::from_str(&responses_metadata.installation_id) {
extra_headers.insert(X_CODEX_INSTALLATION_ID_HEADER, header_value);
}
extra_headers.extend(build_responses_headers(
self.state.beta_features_header.as_deref(),
turn_state.as_ref(),
));
add_originator_header(&mut extra_headers, self.state.originator.as_str());
extra_headers.extend(self.build_responses_compatibility_headers(responses_metadata));
extra_headers.extend(build_session_headers(
Some(responses_metadata.session_id.to_string()),
Some(responses_metadata.thread_id.to_string()),
));
if let Some(header_value) = self.generate_attestation_header_for().await {
extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value);
}
add_responses_lite_header(&mut extra_headers, model_info.use_responses_lite);
let compact_request_timeout = client_setup
.api_provider
.stream_idle_timeout
.saturating_mul(COMPACT_REQUEST_TIMEOUT_IDLE_MULTIPLIER);
let client =
ApiCompactClient::new(transport, client_setup.api_provider, client_setup.api_auth)
.with_telemetry(Some(request_telemetry));
let trace_attempt = compaction_trace.start_attempt(&payload);
let result = client
.compact_input(
&payload,
extra_headers,
compact_request_timeout,
turn_state.as_deref(),
)
.await
.map_err(|error| self.state.provider.map_api_error(error));
trace_attempt.record_result(result.as_deref());
result
}
pub(crate) async fn create_realtime_call_with_headers(
&self,
sdp: String,
session_config: ApiRealtimeSessionConfig,
mut extra_headers: ApiHeaderMap,
api_provider_override: Option<ApiProvider>,
) -> Result<RealtimeWebrtcCallStart> {
let client_setup = self.current_client_setup().await?;
if let Some(header_value) = self.generate_attestation_header_for().await {
extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value);
}
let mut sideband_headers = extra_headers.clone();
sideband_headers.extend(sideband_websocket_auth_headers(
client_setup.api_auth.as_ref(),
));
let api_provider = api_provider_override.unwrap_or(client_setup.api_provider);
let transport = self.build_api_transport(&api_provider, REALTIME_CALLS_ENDPOINT)?;
let response = ApiRealtimeCallClient::new(transport, api_provider, client_setup.api_auth)
.create_with_session_and_headers(sdp, session_config, extra_headers)
.await
.map_err(|error| self.state.provider.map_api_error(error))?;
Ok(RealtimeWebrtcCallStart {
sdp: response.sdp,
call_id: response.call_id,
sideband_headers,
})
}
pub async fn summarize_memories(
&self,
raw_memories: Vec<ApiRawMemory>,
model_info: &ModelInfo,
effort: Option<ReasoningEffortConfig>,
session_telemetry: &SessionTelemetry,
) -> Result<Vec<ApiMemorySummarizeOutput>> {
if raw_memories.is_empty() {
return Ok(Vec::new());
}
let client_setup = self.current_client_setup().await?;
let transport =
self.build_api_transport(&client_setup.api_provider, MEMORIES_SUMMARIZE_ENDPOINT)?;
let request_telemetry = Self::build_request_telemetry(
session_telemetry,
AuthRequestTelemetryContext::new(
client_setup.auth.as_ref().map(CodexAuth::auth_mode),
client_setup.api_auth.as_ref(),
client_setup.agent_identity_telemetry.clone(),
PendingUnauthorizedRetry::default(),
),
RequestRouteTelemetry::for_endpoint(MEMORIES_SUMMARIZE_ENDPOINT),
self.state.auth_env_telemetry.clone(),
);
let client =
ApiMemoriesClient::new(transport, client_setup.api_provider, client_setup.api_auth)
.with_telemetry(Some(request_telemetry));
let payload = ApiMemorySummarizeInput {
model: model_info.slug.clone(),
raw_memories,
reasoning: effort
.map(reasoning_effort_for_request)
.map(|effort| Reasoning {
effort: Some(effort),
summary: None,
context: None,
}),
};
client
.summarize_input(&payload, self.build_subagent_headers())
.await
.map_err(|error| self.state.provider.map_api_error(error))
}
fn build_subagent_headers(&self) -> ApiHeaderMap {
let mut extra_headers = ApiHeaderMap::new();
add_originator_header(&mut extra_headers, self.state.originator.as_str());
if let Some(subagent) = subagent_header_value(&self.state.session_source)
&& let Ok(val) = HeaderValue::from_str(&subagent)
{
extra_headers.insert(X_OPENAI_SUBAGENT_HEADER, val);
}
if matches!(
self.state.session_source,
SessionSource::Internal(InternalSessionSource::MemoryConsolidation)
) {
extra_headers.insert(
X_OPENAI_MEMGEN_REQUEST_HEADER,
HeaderValue::from_static("true"),
);
}
extra_headers
}
fn build_responses_compatibility_headers(
&self,
responses_metadata: &CodexResponsesMetadata,
) -> ApiHeaderMap {
let mut extra_headers = responses_metadata.compatibility_headers();
if matches!(
self.state.session_source,
SessionSource::Internal(InternalSessionSource::MemoryConsolidation)
) {
extra_headers.insert(
X_OPENAI_MEMGEN_REQUEST_HEADER,
HeaderValue::from_static("true"),
);
}
extra_headers
}
fn build_ws_client_metadata(
&self,
responses_metadata: &CodexResponsesMetadata,
use_responses_lite: bool,
) -> HashMap<String, String> {
let mut client_metadata = responses_metadata.client_metadata();
if use_responses_lite {
client_metadata.insert(
WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY.to_string(),
"true".to_string(),
);
}
client_metadata
}
async fn generate_attestation_header_for(&self) -> Option<HeaderValue> {
if !self.state.include_attestation {
return None;
}
self.state
.attestation_provider
.as_ref()?
.header_for_request(AttestationContext {
thread_id: self.state.thread_id,
})
.await
}
fn build_request_telemetry(
session_telemetry: &SessionTelemetry,
auth_context: AuthRequestTelemetryContext,
request_route_telemetry: RequestRouteTelemetry,
auth_env_telemetry: AuthEnvTelemetry,
) -> Arc<dyn RequestTelemetry> {
let telemetry = Arc::new(ApiTelemetry::new(
session_telemetry.clone(),
auth_context,
request_route_telemetry,
auth_env_telemetry,
));
let request_telemetry: Arc<dyn RequestTelemetry> = telemetry;
request_telemetry
}
fn build_reasoning(
model_info: &ModelInfo,
effort: Option<ReasoningEffortConfig>,
summary: ReasoningSummaryConfig,
) -> Reasoning {
Reasoning {
effort: effort
.or_else(|| model_info.default_reasoning_level.clone())
.map(reasoning_effort_for_request),
summary: (model_info.supports_reasoning_summary_parameter
&& summary != ReasoningSummaryConfig::None)
.then_some(summary),
context: model_info
.use_responses_lite
.then_some(ReasoningContext::AllTurns),
}
}
#[allow(clippy::too_many_arguments)]
fn build_responses_request(
&self,
provider: &codex_api::Provider,
prompt: &Prompt,
model_info: &ModelInfo,
effort: Option<ReasoningEffortConfig>,
summary: ReasoningSummaryConfig,
service_tier: Option<String>,
responses_metadata: &CodexResponsesMetadata,
) -> Result<ResponsesApiRequest> {
let mut input = prompt.get_formatted_input_for_request(model_info.use_responses_lite);
let is_openai = self.state.provider.info().is_openai();
if !is_openai {
input
.iter_mut()
.for_each(ResponseItem::clear_internal_chat_message_metadata_passthrough);
}
let tools = create_tools_json_for_responses_api(&prompt.tools)?;
let (instructions, tools) = if model_info.use_responses_lite {
let mut prefix = vec![ResponseItem::AdditionalTools {
id: None,
role: "developer".to_string(),
tools,
}];
if !prompt.base_instructions.text.is_empty() {
prefix.push(ResponseItem::Message {
id: None,
role: "developer".to_string(),
content: vec![ContentItem::InputText {
text: prompt.base_instructions.text.clone(),
}],
phase: None,
internal_chat_message_metadata_passthrough: None,
});
}
input.splice(0..0, prefix);
(String::new(), None)
} else {
(prompt.base_instructions.text.clone(), Some(tools))
};
let reasoning = Self::build_reasoning(model_info, effort, summary);
let stream_options = (self.state.concurrent_reasoning_summaries_enabled
&& is_openai
&& reasoning.summary.is_some())
.then_some(StreamOptions {
reasoning_summary_delivery: codex_api::ReasoningSummaryDelivery::SequentialCutoff,
});
let include = vec!["reasoning.encrypted_content".to_string()];
let verbosity = if model_info.support_verbosity {
self.state.model_verbosity.or(model_info.default_verbosity)
} else {
if self.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,
prompt.output_schema_strict,
);
let prompt_cache_key = Some(self.prompt_cache_key(responses_metadata));
let service_tier = model_info.service_tier_for_request(service_tier);
let request = ResponsesApiRequest {
model: model_info.slug.clone(),
instructions,
input,
tools,
tool_choice: "auto".to_string(),
parallel_tool_calls: prompt.parallel_tool_calls && !model_info.use_responses_lite,
reasoning: Some(reasoning),
store: provider.is_azure_responses_endpoint(),
stream: true,
stream_options,
include,
service_tier,
prompt_cache_key,
text,
client_metadata: Some(responses_metadata.client_metadata()),
};
Ok(request)
}
fn prepare_response_items_for_request(&self, input: &mut [ResponseItem]) {
for item in input {
if item.id().is_some_and(|id| !id.is_prefixed()) {
item.set_id( None);
}
}
}
pub fn responses_websocket_enabled(&self) -> bool {
if !self.state.provider.info().supports_websockets
|| self.state.disable_websockets.load(Ordering::Relaxed)
{
return false;
}
true
}
async fn current_client_setup(&self) -> Result<CurrentClientSetup> {
let auth = self.state.provider.auth().await;
let api_provider = self.state.provider.api_provider().await?;
let resolved_auth = self
.state
.provider
.api_auth_for_scope(ProviderAuthScope {
agent_identity_policy: self.agent_identity_policy,
session_source: self.state.session_source.clone(),
agent_identity_session_fallback: self.state.agent_identity_session_fallback.clone(),
})
.await?;
Ok(CurrentClientSetup {
auth,
api_provider,
api_auth: resolved_auth.auth,
agent_identity_telemetry: resolved_auth.agent_identity_telemetry,
})
}
fn build_api_transport(
&self,
api_provider: &ApiProvider,
endpoint: &str,
) -> Result<ReqwestTransport> {
let request_url = api_provider.url_for_path(endpoint);
let client = create_client_for_route(
&self.http_client_factory,
&request_url,
ClientRouteClass::Api,
)
.map_err(std::io::Error::from)?;
Ok(ReqwestTransport::from_http_client(client))
}
pub(crate) async fn prewarm_auth(&self) -> Result<()> {
self.current_client_setup().await.map(|_| ())
}
#[allow(clippy::too_many_arguments)]
async fn connect_websocket(
&self,
session_telemetry: &SessionTelemetry,
api_provider: codex_api::Provider,
api_auth: SharedAuthProvider,
responses_metadata: &CodexResponsesMetadata,
auth_context: AuthRequestTelemetryContext,
request_route_telemetry: RequestRouteTelemetry,
) -> std::result::Result<ApiWebSocketConnection, ApiError> {
let headers = self.build_websocket_headers(responses_metadata).await;
let websocket_telemetry = ModelClientSession::build_websocket_telemetry(
session_telemetry,
auth_context.clone(),
request_route_telemetry,
self.state.auth_env_telemetry.clone(),
);
let websocket_connect_timeout = self.state.provider.info().websocket_connect_timeout();
let start = Instant::now();
let result = match tokio::time::timeout(
websocket_connect_timeout,
ApiWebSocketResponsesClient::new(api_provider, api_auth).connect(
&self.http_client_factory,
headers,
codex_login::default_client::default_headers(),
None,
Some(websocket_telemetry),
),
)
.await
{
Ok(result) => result,
Err(_) => Err(ApiError::Transport(TransportError::Timeout)),
};
let error_message = result.as_ref().err().map(telemetry_api_error_message);
let response_debug = result
.as_ref()
.err()
.map(extract_response_debug_context_from_api_error)
.unwrap_or_default();
let status = result.as_ref().err().and_then(api_error_http_status);
session_telemetry.record_websocket_connect(
start.elapsed(),
status,
error_message.as_deref(),
auth_context.auth_header_attached,
auth_context.auth_header_name,
auth_context.retry_after_unauthorized,
auth_context.recovery_mode,
auth_context.recovery_phase,
request_route_telemetry.endpoint,
false,
response_debug.request_id.as_deref(),
response_debug.cf_ray.as_deref(),
response_debug.auth_error.as_deref(),
response_debug.auth_error_code.as_deref(),
auth_context.agent_identity_telemetry(),
);
emit_feedback_request_tags_with_auth_env(
&FeedbackRequestTags {
endpoint: request_route_telemetry.endpoint,
auth_header_attached: auth_context.auth_header_attached,
auth_header_name: auth_context.auth_header_name,
auth_mode: auth_context.auth_mode,
auth_retry_after_unauthorized: Some(auth_context.retry_after_unauthorized),
auth_recovery_mode: auth_context.recovery_mode,
auth_recovery_phase: auth_context.recovery_phase,
auth_connection_reused: Some(false),
auth_request_id: response_debug.request_id.as_deref(),
auth_cf_ray: response_debug.cf_ray.as_deref(),
auth_error: response_debug.auth_error.as_deref(),
auth_error_code: response_debug.auth_error_code.as_deref(),
auth_recovery_followup_success: auth_context
.retry_after_unauthorized
.then_some(result.is_ok()),
auth_recovery_followup_status: auth_context
.retry_after_unauthorized
.then_some(status)
.flatten(),
},
&self.state.auth_env_telemetry,
);
result
}
async fn build_websocket_headers(
&self,
responses_metadata: &CodexResponsesMetadata,
) -> ApiHeaderMap {
let mut headers = build_responses_headers(
self.state.beta_features_header.as_deref(),
None,
);
add_originator_header(&mut headers, self.state.originator.as_str());
if let Ok(header_value) = HeaderValue::from_str(&responses_metadata.thread_id) {
headers.insert("x-client-request-id", header_value);
}
headers.extend(build_session_headers(
Some(responses_metadata.session_id.to_string()),
Some(responses_metadata.thread_id.to_string()),
));
headers.extend(self.build_responses_compatibility_headers(responses_metadata));
if let Some(header_value) = self.generate_attestation_header_for().await {
headers.insert(X_OAI_ATTESTATION_HEADER, header_value);
}
headers.insert(
OPENAI_BETA_HEADER,
HeaderValue::from_static(RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE),
);
if self.state.include_timing_metrics {
headers.insert(
X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER,
HeaderValue::from_static("true"),
);
}
headers
}
}
impl Drop for ModelClientSession {
fn drop(&mut self) {
let websocket_session = std::mem::take(&mut self.websocket_session);
self.client
.store_cached_websocket_session(websocket_session);
}
}
impl ModelClientSession {
pub(crate) fn turn_state(&self) -> Arc<OnceLock<String>> {
Arc::clone(&self.turn_state)
}
fn reset_websocket_session(&mut self) {
self.websocket_session.connection = None;
self.websocket_session.last_request = None;
self.websocket_session.last_response_rx = None;
self.websocket_session.last_response_from_untraced_warmup = false;
self.websocket_session
.set_connection_reused( false);
}
#[allow(clippy::too_many_arguments)]
async fn build_responses_options(
&self,
responses_metadata: &CodexResponsesMetadata,
compression: Compression,
use_responses_lite: bool,
) -> ApiResponsesOptions {
ApiResponsesOptions {
session_id: Some(responses_metadata.session_id.to_string()),
thread_id: Some(responses_metadata.thread_id.to_string()),
session_source: Some(self.client.state.session_source.clone()),
extra_headers: {
let mut headers = build_responses_headers(
self.client.state.beta_features_header.as_deref(),
Some(&self.turn_state),
);
add_originator_header(&mut headers, self.client.state.originator.as_str());
headers.extend(
self.client
.build_responses_compatibility_headers(responses_metadata),
);
if let Some(header_value) = self.client.generate_attestation_header_for().await {
headers.insert(X_OAI_ATTESTATION_HEADER, header_value);
}
add_responses_lite_header(&mut headers, use_responses_lite);
headers
},
compression,
turn_state: Some(Arc::clone(&self.turn_state)),
}
}
fn get_incremental_items(
&self,
request: &ResponsesApiRequest,
last_response: Option<&LastResponse>,
allow_empty_delta: bool,
) -> Option<Vec<ResponseItem>> {
let previous_request = self.websocket_session.last_request.as_ref()?;
if !responses_request_properties_match(previous_request, request) {
trace!("incremental request failed, websocket reuse properties didn't match");
return None;
}
let mut previous_items = previous_request.input.clone();
if let Some(response) = last_response {
previous_items.extend_from_slice(&response.items_added);
}
previous_items
.iter_mut()
.for_each(ResponseItem::clear_internal_chat_message_metadata_passthrough);
let Some((request_items_to_compare, incremental_items)) =
request.input.split_at_checked(previous_items.len())
else {
trace!("incremental request failed, incompatible request length");
return None;
};
let mut request_prefix = request_items_to_compare.to_vec();
request_prefix
.iter_mut()
.for_each(ResponseItem::clear_internal_chat_message_metadata_passthrough);
if previous_items != request_prefix {
trace!("incremental request failed, items didn't match");
return None;
}
if !allow_empty_delta && incremental_items.is_empty() {
return None;
}
Some(incremental_items.to_vec())
}
fn get_last_response(&mut self) -> Option<LastResponse> {
self.websocket_session
.last_response_rx
.take()
.and_then(|mut receiver| match receiver.try_recv() {
Ok(last_response) => Some(last_response),
Err(TryRecvError::Closed) | Err(TryRecvError::Empty) => None,
})
}
fn prepare_websocket_request(
&mut self,
request: &ResponsesApiRequest,
) -> (Option<(String, Vec<ResponseItem>)>, bool) {
let Some(last_response) = self.get_last_response() else {
return (None, false);
};
let previous_response_id_from_untraced_warmup =
self.websocket_session.last_response_from_untraced_warmup;
let Some(incremental_items) = self.get_incremental_items(
request,
Some(&last_response),
true,
) else {
return (None, false);
};
if last_response.response_id.is_empty() {
trace!("incremental request failed, no previous response id");
return (None, false);
}
(
Some((last_response.response_id, incremental_items)),
previous_response_id_from_untraced_warmup,
)
}
pub async fn preconnect_websocket(
&mut self,
session_telemetry: &SessionTelemetry,
responses_metadata: &CodexResponsesMetadata,
) -> std::result::Result<(), ApiError> {
if !self.client.responses_websocket_enabled() {
return Ok(());
}
if self.websocket_session.connection.is_some() {
return Ok(());
}
let client_setup = self.client.current_client_setup().await.map_err(|err| {
ApiError::Stream(format!(
"failed to build websocket prewarm client setup: {err}"
))
})?;
let auth_context = AuthRequestTelemetryContext::new(
client_setup.auth.as_ref().map(CodexAuth::auth_mode),
client_setup.api_auth.as_ref(),
client_setup.agent_identity_telemetry.clone(),
PendingUnauthorizedRetry::default(),
);
let connection = self
.client
.connect_websocket(
session_telemetry,
client_setup.api_provider,
client_setup.api_auth,
responses_metadata,
auth_context,
RequestRouteTelemetry::for_endpoint(RESPONSES_ENDPOINT),
)
.await?;
self.websocket_session.connection = Some(connection);
self.websocket_session
.set_connection_reused( false);
Ok(())
}
#[instrument(
name = "model_client.websocket_connection",
level = "info",
skip_all,
fields(
provider = %self.client.state.provider.info().name,
wire_api = %self.client.state.provider.info().wire_api,
transport = "responses_websocket",
api.path = "responses",
turn.has_metadata_header = params.responses_metadata.has_turn_metadata()
)
)]
async fn websocket_connection(
&mut self,
params: WebsocketConnectParams<'_>,
) -> std::result::Result<&ApiWebSocketConnection, ApiError> {
let WebsocketConnectParams {
session_telemetry,
api_provider,
api_auth,
responses_metadata,
auth_context,
request_route_telemetry,
} = params;
let needs_new = match self.websocket_session.connection.as_ref() {
Some(conn) => conn.is_closed().await,
None => true,
};
if needs_new {
self.websocket_session.last_request = None;
self.websocket_session.last_response_rx = None;
self.websocket_session.last_response_from_untraced_warmup = false;
let new_conn = match self
.client
.connect_websocket(
session_telemetry,
api_provider,
api_auth,
responses_metadata,
auth_context,
request_route_telemetry,
)
.await
{
Ok(new_conn) => new_conn,
Err(err) => {
if matches!(err, ApiError::Transport(TransportError::Timeout)) {
self.reset_websocket_session();
}
return Err(err);
}
};
self.websocket_session.connection = Some(new_conn);
self.websocket_session
.set_connection_reused( false);
} else {
self.websocket_session
.set_connection_reused( true);
}
self.websocket_session
.connection
.as_ref()
.ok_or(ApiError::Stream(
"websocket connection is unavailable".to_string(),
))
}
fn responses_request_compression(&self, auth: Option<&CodexAuth>) -> Compression {
if self.client.state.enable_request_compression
&& auth.is_some_and(CodexAuth::uses_codex_backend)
&& self.client.state.provider.info().is_openai()
{
Compression::Zstd
} else {
Compression::None
}
}
#[allow(clippy::too_many_arguments)]
#[instrument(
name = "model_client.stream_responses_api",
level = "info",
skip_all,
fields(
model = %model_info.slug,
wire_api = %self.client.state.provider.info().wire_api,
transport = "responses_http",
http.method = "POST",
api.path = "responses",
turn.has_metadata_header = responses_metadata.has_turn_metadata()
)
)]
async fn stream_responses_api(
&self,
prompt: &Prompt,
model_info: &ModelInfo,
session_telemetry: &SessionTelemetry,
effort: Option<ReasoningEffortConfig>,
summary: ReasoningSummaryConfig,
service_tier: Option<String>,
responses_metadata: &CodexResponsesMetadata,
inference_trace: &InferenceTraceContext,
) -> Result<ResponseStream> {
let auth_manager = self.client.state.provider.auth_manager();
let mut auth_recovery = auth_manager
.as_ref()
.map(AuthManager::unauthorized_recovery);
let mut pending_retry = PendingUnauthorizedRetry::default();
loop {
let client_setup = self.client.current_client_setup().await?;
let transport = self
.client
.build_api_transport(&client_setup.api_provider, RESPONSES_ENDPOINT)?;
let request_auth_context = AuthRequestTelemetryContext::new(
client_setup.auth.as_ref().map(CodexAuth::auth_mode),
client_setup.api_auth.as_ref(),
client_setup.agent_identity_telemetry.clone(),
pending_retry,
);
let (request_telemetry, sse_telemetry) = Self::build_streaming_telemetry(
session_telemetry,
request_auth_context,
RequestRouteTelemetry::for_endpoint(RESPONSES_ENDPOINT),
self.client.state.auth_env_telemetry.clone(),
);
let compression = self.responses_request_compression(client_setup.auth.as_ref());
let mut options = self
.build_responses_options(
responses_metadata,
compression,
model_info.use_responses_lite,
)
.await;
let mut request = self.client.build_responses_request(
&client_setup.api_provider,
prompt,
model_info,
effort.clone(),
summary,
service_tier.clone(),
responses_metadata,
)?;
self.client
.prepare_response_items_for_request(&mut request.input);
let request_session_telemetry =
session_telemetry_for_request(session_telemetry, &request);
let inference_trace_attempt = inference_trace.start_attempt();
inference_trace_attempt.add_request_headers(&mut options.extra_headers);
inference_trace_attempt.record_started(&request);
let client = ApiResponsesClient::new(
transport,
client_setup.api_provider,
client_setup.api_auth,
)
.with_telemetry(Some(request_telemetry), Some(sse_telemetry));
let stream_result = client.stream_request(request, options).await;
match stream_result {
Ok(stream) => {
let (stream, _) = map_response_stream(
stream,
request_session_telemetry,
inference_trace_attempt,
Arc::clone(&self.client.state.provider),
);
return Ok(stream);
}
Err(ApiError::Transport(
unauthorized_transport @ TransportError::Http { status, .. },
)) if status == StatusCode::UNAUTHORIZED => {
let response_debug_context =
extract_response_debug_context(&unauthorized_transport);
inference_trace_attempt.record_failed(
&unauthorized_transport,
response_debug_context.request_id.as_deref(),
&[],
);
pending_retry = PendingUnauthorizedRetry::from_recovery(
handle_unauthorized(
unauthorized_transport,
&mut auth_recovery,
session_telemetry,
&self.client.state.provider,
)
.await?,
);
continue;
}
Err(err) => {
let response_debug_context =
extract_response_debug_context_from_api_error(&err);
let err = self.client.state.provider.map_api_error(err);
inference_trace_attempt.record_failed(
&err,
response_debug_context.request_id.as_deref(),
&[],
);
return Err(err);
}
}
}
}
#[allow(clippy::too_many_arguments)]
#[instrument(
name = "model_client.stream_responses_websocket",
level = "info",
skip_all,
fields(
model = %model_info.slug,
wire_api = %self.client.state.provider.info().wire_api,
transport = "responses_websocket",
api.path = "responses",
turn.has_metadata_header = responses_metadata.has_turn_metadata(),
websocket.warmup = warmup
)
)]
async fn stream_responses_websocket(
&mut self,
prompt: &Prompt,
model_info: &ModelInfo,
session_telemetry: &SessionTelemetry,
effort: Option<ReasoningEffortConfig>,
summary: ReasoningSummaryConfig,
service_tier: Option<String>,
responses_metadata: &CodexResponsesMetadata,
warmup: bool,
request_trace: Option<W3cTraceContext>,
inference_trace: &InferenceTraceContext,
) -> Result<WebsocketStreamOutcome> {
let auth_manager = self.client.state.provider.auth_manager();
let mut auth_recovery = auth_manager
.as_ref()
.map(AuthManager::unauthorized_recovery);
let mut pending_retry = PendingUnauthorizedRetry::default();
loop {
let client_setup = self.client.current_client_setup().await?;
let request_auth_context = AuthRequestTelemetryContext::new(
client_setup.auth.as_ref().map(CodexAuth::auth_mode),
client_setup.api_auth.as_ref(),
client_setup.agent_identity_telemetry.clone(),
pending_retry,
);
let mut request = self.client.build_responses_request(
&client_setup.api_provider,
prompt,
model_info,
effort.clone(),
summary,
service_tier.clone(),
responses_metadata,
)?;
let request_session_telemetry = if warmup {
session_telemetry.clone()
} else {
session_telemetry_for_request(session_telemetry, &request)
};
let mut client_metadata = self
.client
.build_ws_client_metadata(responses_metadata, model_info.use_responses_lite);
if let Some(turn_state) = self.turn_state.get() {
client_metadata.insert(X_CODEX_TURN_STATE_HEADER.to_string(), turn_state.clone());
}
match self
.websocket_connection(WebsocketConnectParams {
session_telemetry,
api_provider: client_setup.api_provider,
api_auth: client_setup.api_auth,
responses_metadata,
auth_context: request_auth_context,
request_route_telemetry: RequestRouteTelemetry::for_endpoint(
RESPONSES_ENDPOINT,
),
})
.await
{
Ok(_) => {}
Err(ApiError::Transport(TransportError::Http { status, .. }))
if status == StatusCode::UPGRADE_REQUIRED =>
{
return Ok(WebsocketStreamOutcome::FallbackToHttp);
}
Err(ApiError::Transport(
unauthorized_transport @ TransportError::Http { status, .. },
)) if status == StatusCode::UNAUTHORIZED => {
pending_retry = PendingUnauthorizedRetry::from_recovery(
handle_unauthorized(
unauthorized_transport,
&mut auth_recovery,
session_telemetry,
&self.client.state.provider,
)
.await?,
);
continue;
}
Err(err) => return Err(self.client.state.provider.map_api_error(err)),
}
let (incremental_request, previous_response_id_from_untraced_warmup) =
self.prepare_websocket_request(&request);
let inference_trace_attempt = if warmup {
InferenceTraceAttempt::disabled()
} else {
inference_trace.start_attempt()
};
if previous_response_id_from_untraced_warmup {
inference_trace_attempt.record_started(&request);
}
let (previous_response_id, mut incremental_items) = match incremental_request {
Some((response_id, items)) => (Some(response_id), Some(items)),
None => (None, None),
};
let original_item_ids = if let Some(incremental_items) = &mut incremental_items {
self.client
.prepare_response_items_for_request(incremental_items);
None
} else {
let original_item_ids = request
.input
.iter()
.map(|item| item.id().cloned())
.collect::<Vec<_>>();
self.client
.prepare_response_items_for_request(&mut request.input);
Some(original_item_ids)
};
let ws_payload = ResponseCreateWsRequest {
previous_response_id,
input: incremental_items.as_deref().unwrap_or(&request.input),
generate: if warmup { Some(false) } else { None },
client_metadata: response_create_client_metadata(
Some(client_metadata),
request_trace.as_ref(),
),
..ResponseCreateWsRequest::from(&request)
};
let mut ws_request = ResponsesWsRequest::ResponseCreate(ws_payload);
stamp_ws_stream_request_start_ms(&mut ws_request);
if !previous_response_id_from_untraced_warmup {
inference_trace_attempt.record_started(&ws_request);
}
let websocket_connection =
self.websocket_session.connection.as_ref().ok_or_else(|| {
self.client.state.provider.map_api_error(ApiError::Stream(
"websocket connection is unavailable".to_string(),
))
})?;
let stream_result = websocket_connection
.stream_request(
ws_request,
self.websocket_session.connection_reused(),
Some(Arc::clone(&self.turn_state)),
)
.await;
if let Some(original_item_ids) = original_item_ids {
for (item, original_item_id) in request.input.iter_mut().zip(original_item_ids) {
item.set_id(original_item_id);
}
}
self.websocket_session.last_request = Some(request);
self.websocket_session.last_response_from_untraced_warmup = warmup;
let stream_result = stream_result.map_err(|err| {
let response_debug_context = extract_response_debug_context_from_api_error(&err);
let err = self.client.state.provider.map_api_error(err);
inference_trace_attempt.record_failed(
&err,
response_debug_context.request_id.as_deref(),
&[],
);
err
})?;
let (stream, last_request_rx) = map_response_stream(
stream_result,
request_session_telemetry,
inference_trace_attempt,
Arc::clone(&self.client.state.provider),
);
self.websocket_session.last_response_rx = Some(last_request_rx);
return Ok(WebsocketStreamOutcome::Stream(stream));
}
}
fn build_streaming_telemetry(
session_telemetry: &SessionTelemetry,
auth_context: AuthRequestTelemetryContext,
request_route_telemetry: RequestRouteTelemetry,
auth_env_telemetry: AuthEnvTelemetry,
) -> (Arc<dyn RequestTelemetry>, Arc<dyn SseTelemetry>) {
let telemetry = Arc::new(ApiTelemetry::new(
session_telemetry.clone(),
auth_context,
request_route_telemetry,
auth_env_telemetry,
));
let request_telemetry: Arc<dyn RequestTelemetry> = telemetry.clone();
let sse_telemetry: Arc<dyn SseTelemetry> = telemetry;
(request_telemetry, sse_telemetry)
}
fn build_websocket_telemetry(
session_telemetry: &SessionTelemetry,
auth_context: AuthRequestTelemetryContext,
request_route_telemetry: RequestRouteTelemetry,
auth_env_telemetry: AuthEnvTelemetry,
) -> Arc<dyn WebsocketTelemetry> {
let telemetry = Arc::new(ApiTelemetry::new(
session_telemetry.clone(),
auth_context,
request_route_telemetry,
auth_env_telemetry,
));
let websocket_telemetry: Arc<dyn WebsocketTelemetry> = telemetry;
websocket_telemetry
}
#[allow(clippy::too_many_arguments)]
pub async fn prewarm_websocket(
&mut self,
prompt: &Prompt,
model_info: &ModelInfo,
session_telemetry: &SessionTelemetry,
effort: Option<ReasoningEffortConfig>,
summary: ReasoningSummaryConfig,
service_tier: Option<String>,
responses_metadata: &CodexResponsesMetadata,
) -> Result<()> {
if !self.client.responses_websocket_enabled() {
return Ok(());
}
if self.websocket_session.last_request.is_some() {
return Ok(());
}
let disabled_trace = InferenceTraceContext::disabled();
match self
.stream_responses_websocket(
prompt,
model_info,
session_telemetry,
effort,
summary,
service_tier,
responses_metadata,
true,
current_span_w3c_trace_context(),
&disabled_trace,
)
.await
{
Ok(WebsocketStreamOutcome::Stream(mut stream)) => {
while let Some(event) = stream.next().await {
match event {
Ok(ResponseEvent::Completed { .. }) => break,
Err(err) => return Err(err),
_ => {}
}
}
Ok(())
}
Ok(WebsocketStreamOutcome::FallbackToHttp) => {
self.try_switch_fallback_transport(session_telemetry, model_info);
Ok(())
}
Err(err) => Err(err),
}
}
#[allow(clippy::too_many_arguments)]
pub async fn stream(
&mut self,
prompt: &Prompt,
model_info: &ModelInfo,
session_telemetry: &SessionTelemetry,
effort: Option<ReasoningEffortConfig>,
summary: ReasoningSummaryConfig,
service_tier: Option<String>,
responses_metadata: &CodexResponsesMetadata,
inference_trace: &InferenceTraceContext,
) -> Result<ResponseStream> {
let wire_api = self.client.state.provider.info().wire_api;
match wire_api {
WireApi::Responses => {
if self.client.responses_websocket_enabled() {
let request_trace = current_span_w3c_trace_context();
match self
.stream_responses_websocket(
prompt,
model_info,
session_telemetry,
effort.clone(),
summary,
service_tier.clone(),
responses_metadata,
false,
request_trace,
inference_trace,
)
.await?
{
WebsocketStreamOutcome::Stream(stream) => return Ok(stream),
WebsocketStreamOutcome::FallbackToHttp => {
self.try_switch_fallback_transport(session_telemetry, model_info);
}
}
}
self.stream_responses_api(
prompt,
model_info,
session_telemetry,
effort,
summary,
service_tier,
responses_metadata,
inference_trace,
)
.await
}
}
}
pub(crate) fn try_switch_fallback_transport(
&mut self,
session_telemetry: &SessionTelemetry,
model_info: &ModelInfo,
) -> bool {
let activated = self
.client
.force_http_fallback(session_telemetry, model_info);
self.websocket_session = WebsocketSession::default();
activated
}
}
fn stamp_ws_stream_request_start_ms(request: &mut ResponsesWsRequest<'_>) {
let ResponsesWsRequest::ResponseCreate(payload) = request;
payload
.client_metadata
.get_or_insert_with(HashMap::new)
.insert(
X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY.to_string(),
crate::turn_timing::now_unix_timestamp_ms().to_string(),
);
}
fn build_responses_headers(
beta_features_header: Option<&str>,
turn_state: Option<&Arc<OnceLock<String>>>,
) -> 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);
}
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);
}
headers
}
fn add_responses_lite_header(headers: &mut ApiHeaderMap, use_responses_lite: bool) {
if use_responses_lite {
headers.insert(
X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER,
HeaderValue::from_static("true"),
);
}
}
const RESPONSE_STREAM_CHANNEL_CAPACITY: usize = 1600;
const STREAM_DROPPED_REASON: &str = "response stream dropped before provider terminal event";
fn map_response_stream(
api_stream: codex_api::ResponseStream,
session_telemetry: SessionTelemetry,
inference_trace_attempt: InferenceTraceAttempt,
provider: SharedModelProvider,
) -> (ResponseStream, oneshot::Receiver<LastResponse>) {
let codex_api::ResponseStream {
rx_event,
upstream_request_id,
} = api_stream;
let api_stream = codex_api::ResponseStream {
rx_event,
upstream_request_id: None,
};
map_response_events(
upstream_request_id,
api_stream,
session_telemetry,
inference_trace_attempt,
provider,
)
}
fn map_response_events<S>(
upstream_request_id: Option<String>,
api_stream: S,
session_telemetry: SessionTelemetry,
inference_trace_attempt: InferenceTraceAttempt,
provider: SharedModelProvider,
) -> (ResponseStream, oneshot::Receiver<LastResponse>)
where
S: futures::Stream<Item = std::result::Result<ResponseEvent, ApiError>>
+ Unpin
+ Send
+ 'static,
{
let (tx_event, rx_event) =
mpsc::channel::<Result<ResponseEvent>>(RESPONSE_STREAM_CHANNEL_CAPACITY);
let (tx_last_response, rx_last_response) = oneshot::channel::<LastResponse>();
let consumer_dropped = CancellationToken::new();
let consumer_dropped_for_stream = consumer_dropped.clone();
tokio::spawn(async move {
let mut logged_error = false;
let mut tx_last_response = Some(tx_last_response);
let mut items_added: Vec<ResponseItem> = Vec::new();
let (request_start, mut ttft_ms) = (Instant::now(), None);
let mut api_stream = api_stream;
let upstream_request_id = upstream_request_id.as_deref();
if let Some(upstream_request_id) = upstream_request_id {
feedback_tags!(last_model_request_id = upstream_request_id);
}
loop {
let event = tokio::select! {
_ = consumer_dropped.cancelled() => {
inference_trace_attempt.record_cancelled(
STREAM_DROPPED_REASON,
upstream_request_id,
&items_added,
);
return;
}
event = api_stream.next() => event,
};
let Some(event) = event else {
break;
};
match event {
Ok(ResponseEvent::OutputItemDone(item)) => {
items_added.push(item.clone());
if tx_event
.send(Ok(ResponseEvent::OutputItemDone(item)))
.await
.is_err()
{
inference_trace_attempt.record_cancelled(
STREAM_DROPPED_REASON,
upstream_request_id,
&items_added,
);
return;
}
}
Ok(ResponseEvent::Completed {
response_id,
token_usage,
end_turn,
}) => {
feedback_tags!(last_model_response_id = &response_id);
if let Some(usage) = &token_usage {
session_telemetry.sse_event_completed(usage, ttft_ms);
}
inference_trace_attempt.record_completed(
&response_id,
upstream_request_id,
&token_usage,
&items_added,
);
if let Some(sender) = tx_last_response.take() {
let _ = sender.send(LastResponse {
response_id: response_id.clone(),
items_added: std::mem::take(&mut items_added),
});
}
if tx_event
.send(Ok(ResponseEvent::Completed {
response_id,
token_usage,
end_turn,
}))
.await
.is_err()
{
return;
}
}
Ok(event) => {
if matches!(&event, ResponseEvent::OutputItemAdded(_)) && ttft_ms.is_none() {
ttft_ms = Some(
i64::try_from(request_start.elapsed().as_millis()).unwrap_or(i64::MAX),
);
}
if tx_event.send(Ok(event)).await.is_err() {
inference_trace_attempt.record_cancelled(
STREAM_DROPPED_REASON,
upstream_request_id,
&items_added,
);
return;
}
}
Err(err) => {
let response_debug_context =
extract_response_debug_context_from_api_error(&err);
let upstream_request_id =
upstream_request_id.or(response_debug_context.request_id.as_deref());
if let Some(upstream_request_id) = upstream_request_id {
feedback_tags!(last_model_request_id = upstream_request_id);
}
let mapped = provider.map_api_error(err);
inference_trace_attempt.record_failed(
&mapped,
upstream_request_id,
&items_added,
);
if !logged_error {
session_telemetry.see_event_completed_failed(&mapped);
logged_error = true;
}
if tx_event.send(Err(mapped)).await.is_err() {
return;
}
}
}
}
inference_trace_attempt.record_failed(
"stream closed before response.completed",
upstream_request_id,
&items_added,
);
});
(
ResponseStream {
rx_event,
consumer_dropped: consumer_dropped_for_stream,
},
rx_last_response,
)
}
#[derive(Clone, Copy, Debug)]
struct UnauthorizedRecoveryExecution {
mode: &'static str,
phase: &'static str,
}
#[derive(Clone, Copy, Debug, Default)]
struct PendingUnauthorizedRetry {
retry_after_unauthorized: bool,
recovery_mode: Option<&'static str>,
recovery_phase: Option<&'static str>,
}
impl PendingUnauthorizedRetry {
fn from_recovery(recovery: UnauthorizedRecoveryExecution) -> Self {
Self {
retry_after_unauthorized: true,
recovery_mode: Some(recovery.mode),
recovery_phase: Some(recovery.phase),
}
}
}
#[derive(Clone, Debug, Default)]
struct AuthRequestTelemetryContext {
auth_mode: Option<&'static str>,
auth_header_attached: bool,
auth_header_name: Option<&'static str>,
agent_identity_telemetry: Option<AgentIdentityTelemetry>,
retry_after_unauthorized: bool,
recovery_mode: Option<&'static str>,
recovery_phase: Option<&'static str>,
}
impl AuthRequestTelemetryContext {
fn new(
auth_mode: Option<AuthMode>,
api_auth: &dyn AuthProvider,
agent_identity_telemetry: Option<AgentIdentityTelemetry>,
retry: PendingUnauthorizedRetry,
) -> Self {
let auth_telemetry = auth_header_telemetry(api_auth);
Self {
auth_mode: auth_mode.map(|mode| match mode {
AuthMode::ApiKey | AuthMode::BedrockApiKey => "ApiKey",
AuthMode::Chatgpt
| AuthMode::ChatgptAuthTokens
| AuthMode::Headers
| AuthMode::AgentIdentity
| AuthMode::PersonalAccessToken => "Chatgpt",
}),
auth_header_attached: auth_telemetry.attached,
auth_header_name: auth_telemetry.name,
agent_identity_telemetry,
retry_after_unauthorized: retry.retry_after_unauthorized,
recovery_mode: retry.recovery_mode,
recovery_phase: retry.recovery_phase,
}
}
fn agent_identity_telemetry(&self) -> Option<&AgentIdentityTelemetry> {
self.agent_identity_telemetry.as_ref()
}
}
struct WebsocketConnectParams<'a> {
session_telemetry: &'a SessionTelemetry,
api_provider: codex_api::Provider,
api_auth: SharedAuthProvider,
responses_metadata: &'a CodexResponsesMetadata,
auth_context: AuthRequestTelemetryContext,
request_route_telemetry: RequestRouteTelemetry,
}
async fn handle_unauthorized(
transport: TransportError,
auth_recovery: &mut Option<UnauthorizedRecovery>,
session_telemetry: &SessionTelemetry,
provider: &SharedModelProvider,
) -> Result<UnauthorizedRecoveryExecution> {
let debug = extract_response_debug_context(&transport);
if let Some(recovery) = auth_recovery
&& recovery.has_next()
{
let mode = recovery.mode_name();
let phase = recovery.step_name();
return match recovery.next().await {
Ok(step_result) => {
session_telemetry.record_auth_recovery(
mode,
phase,
"recovery_succeeded",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
None,
step_result.auth_state_changed(),
);
emit_feedback_auth_recovery_tags(
mode,
phase,
"recovery_succeeded",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
);
Ok(UnauthorizedRecoveryExecution { mode, phase })
}
Err(RefreshTokenError::Permanent(failed)) => {
session_telemetry.record_auth_recovery(
mode,
phase,
"recovery_failed_permanent",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
None,
None,
);
emit_feedback_auth_recovery_tags(
mode,
phase,
"recovery_failed_permanent",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
);
Err(CodexErr::RefreshTokenFailed(failed))
}
Err(RefreshTokenError::Transient(other)) => {
session_telemetry.record_auth_recovery(
mode,
phase,
"recovery_failed_transient",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
None,
None,
);
emit_feedback_auth_recovery_tags(
mode,
phase,
"recovery_failed_transient",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
);
Err(CodexErr::Io(other))
}
};
}
let (mode, phase, recovery_reason) = match auth_recovery.as_ref() {
Some(recovery) => (
recovery.mode_name(),
recovery.step_name(),
Some(recovery.unavailable_reason()),
),
None => ("none", "none", Some("auth_manager_missing")),
};
session_telemetry.record_auth_recovery(
mode,
phase,
"recovery_not_run",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
recovery_reason,
None,
);
emit_feedback_auth_recovery_tags(
mode,
phase,
"recovery_not_run",
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
);
Err(provider.map_api_error(ApiError::Transport(transport)))
}
fn api_error_http_status(error: &ApiError) -> Option<u16> {
match error {
ApiError::Transport(TransportError::Http { status, .. }) => Some(status.as_u16()),
_ => None,
}
}
struct ApiTelemetry {
session_telemetry: SessionTelemetry,
auth_context: AuthRequestTelemetryContext,
request_route_telemetry: RequestRouteTelemetry,
auth_env_telemetry: AuthEnvTelemetry,
}
impl ApiTelemetry {
fn new(
session_telemetry: SessionTelemetry,
auth_context: AuthRequestTelemetryContext,
request_route_telemetry: RequestRouteTelemetry,
auth_env_telemetry: AuthEnvTelemetry,
) -> Self {
Self {
session_telemetry,
auth_context,
request_route_telemetry,
auth_env_telemetry,
}
}
}
impl RequestTelemetry for ApiTelemetry {
fn on_request(
&self,
attempt: u64,
status: Option<HttpStatusCode>,
error: Option<&TransportError>,
duration: Duration,
) {
let error_message = error.map(telemetry_transport_error_message);
let status = status.map(|s| s.as_u16());
let debug = error
.map(extract_response_debug_context)
.unwrap_or_default();
self.session_telemetry.record_api_request(
attempt,
status,
error_message.as_deref(),
duration,
self.auth_context.auth_header_attached,
self.auth_context.auth_header_name,
self.auth_context.retry_after_unauthorized,
self.auth_context.recovery_mode,
self.auth_context.recovery_phase,
self.request_route_telemetry.endpoint,
debug.request_id.as_deref(),
debug.cf_ray.as_deref(),
debug.auth_error.as_deref(),
debug.auth_error_code.as_deref(),
self.auth_context.agent_identity_telemetry(),
);
emit_feedback_request_tags_with_auth_env(
&FeedbackRequestTags {
endpoint: self.request_route_telemetry.endpoint,
auth_header_attached: self.auth_context.auth_header_attached,
auth_header_name: self.auth_context.auth_header_name,
auth_mode: self.auth_context.auth_mode,
auth_retry_after_unauthorized: Some(self.auth_context.retry_after_unauthorized),
auth_recovery_mode: self.auth_context.recovery_mode,
auth_recovery_phase: self.auth_context.recovery_phase,
auth_connection_reused: None,
auth_request_id: debug.request_id.as_deref(),
auth_cf_ray: debug.cf_ray.as_deref(),
auth_error: debug.auth_error.as_deref(),
auth_error_code: debug.auth_error_code.as_deref(),
auth_recovery_followup_success: self
.auth_context
.retry_after_unauthorized
.then_some(error.is_none()),
auth_recovery_followup_status: self
.auth_context
.retry_after_unauthorized
.then_some(status)
.flatten(),
},
&self.auth_env_telemetry,
);
}
}
impl SseTelemetry for ApiTelemetry {
fn on_sse_poll(
&self,
result: &std::result::Result<
Option<std::result::Result<Event, EventStreamError<TransportError>>>,
tokio::time::error::Elapsed,
>,
duration: Duration,
) {
self.session_telemetry.log_sse_event(result, duration);
}
}
impl WebsocketTelemetry for ApiTelemetry {
fn on_ws_request(&self, duration: Duration, error: Option<&ApiError>, connection_reused: bool) {
let error_message = error.map(telemetry_api_error_message);
let status = error.and_then(api_error_http_status);
let debug = error
.map(extract_response_debug_context_from_api_error)
.unwrap_or_default();
self.session_telemetry.record_websocket_request(
duration,
error_message.as_deref(),
connection_reused,
self.auth_context.agent_identity_telemetry(),
);
emit_feedback_request_tags_with_auth_env(
&FeedbackRequestTags {
endpoint: self.request_route_telemetry.endpoint,
auth_header_attached: self.auth_context.auth_header_attached,
auth_header_name: self.auth_context.auth_header_name,
auth_mode: self.auth_context.auth_mode,
auth_retry_after_unauthorized: Some(self.auth_context.retry_after_unauthorized),
auth_recovery_mode: self.auth_context.recovery_mode,
auth_recovery_phase: self.auth_context.recovery_phase,
auth_connection_reused: Some(connection_reused),
auth_request_id: debug.request_id.as_deref(),
auth_cf_ray: debug.cf_ray.as_deref(),
auth_error: debug.auth_error.as_deref(),
auth_error_code: debug.auth_error_code.as_deref(),
auth_recovery_followup_success: self
.auth_context
.retry_after_unauthorized
.then_some(error.is_none()),
auth_recovery_followup_status: self
.auth_context
.retry_after_unauthorized
.then_some(status)
.flatten(),
},
&self.auth_env_telemetry,
);
}
fn on_ws_event(
&self,
result: &std::result::Result<Option<std::result::Result<Message, Error>>, ApiError>,
duration: Duration,
) {
self.session_telemetry
.record_websocket_event(result, duration);
}
}
#[cfg(test)]
#[path = "client_tests.rs"]
mod tests;