| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| 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 crate::CodexResponsesHeaders; |
| use async_channel::Sender; |
| use codex_api::AgentIdentityTelemetry; |
| use codex_api::ApiError; |
| use codex_api::AuthProvider; |
| 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::ClientRedirectPolicy; |
| 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::ResponseItemId; |
| 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::ResponseItem; |
| use codex_protocol::openai_models::ModelInfo; |
| use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; |
| use codex_protocol::protocol::AuthRecoveryEvent; |
| use codex_protocol::protocol::Event as ProtocolEvent; |
| use codex_protocol::protocol::EventMsg; |
| use codex_protocol::protocol::InternalSessionSource; |
| use codex_protocol::protocol::SessionSource; |
| use codex_protocol::protocol::W3cTraceContext; |
| use codex_rollout_trace::InferenceTraceAttempt; |
| use codex_rollout_trace::InferenceTraceContext; |
| use codex_tools::create_tools_json_for_responses_api; |
| use codex_tools::create_tools_json_for_responses_lite; |
| use codex_tools::create_tools_raw_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; |
| 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 uuid::Uuid; |
|
|
| 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::context::BaseInstructionsFragment; |
| use crate::context::ContextualUserFragment; |
| use crate::cyber_access_program; |
| 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::ProviderUnauthorizedRecovery; |
| use codex_model_provider::ResponsesConnectionKey; |
| use codex_model_provider::SharedModelProvider; |
| use codex_model_provider::WorkspaceRoutingContext; |
| 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_ROUTING_HINT_HEADER: &str = "x-codex-routing-hint"; |
| 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 MEMORIES_SUMMARIZE_ENDPOINT: &str = "/memories/trace_summarize"; |
| #[cfg(test)] |
| pub(crate) const WEBSOCKET_CONNECT_TIMEOUT: Duration = |
| Duration::from_millis(DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS); |
|
|
| 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, |
| workspace_routing: WorkspaceRoutingContext, |
| auth_env_telemetry: AuthEnvTelemetry, |
| session_source: SessionSource, |
| originator: String, |
| model_verbosity: Option<VerbosityConfig>, |
| content_item_kinds_enabled: bool, |
| 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>, |
| } |
|
|
| enum ClientRouting { |
| Workspace, |
| ConfiguredProvider, |
| } |
|
|
| |
| |
| |
| |
| struct CurrentClientSetup { |
| auth: Option<CodexAuth>, |
| auth_owner_generation: Option<u64>, |
| auth_revision: Option<u64>, |
| api_provider: ApiProvider, |
| redirect_policy: ClientRedirectPolicy, |
| 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>, |
| codex_responses_headers: Option<Arc<CodexResponsesHeaders>>, |
| event_sender: Option<Sender<ProtocolEvent>>, |
| 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>, |
| responses_headers: ApiHeaderMap, |
| |
| auth_owner_generation: Option<u64>, |
| connection_key: Option<ResponsesConnectionKey>, |
| 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: _, |
| access_programs: _, |
| } = 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: _, |
| access_programs: _, |
| } = 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 |
| } |
|
|
| fn response_items_equal_ignoring_internal_metadata( |
| previous: &ResponseItem, |
| current: &ResponseItem, |
| ) -> bool { |
| if previous == current { |
| return true; |
| } |
|
|
| let mut previous = previous.clone(); |
| previous.clear_internal_chat_message_metadata_passthrough(); |
| let mut current = current.clone(); |
| current.clear_internal_chat_message_metadata_passthrough(); |
| previous == current |
| } |
|
|
| 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>, |
| content_item_kinds_enabled: bool, |
| 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, |
| workspace_routing: WorkspaceRoutingContext, |
| ) -> 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, |
| workspace_routing, |
| auth_env_telemetry, |
| session_source, |
| originator, |
| model_verbosity, |
| content_item_kinds_enabled, |
| 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, |
| codex_responses_headers: None, |
| event_sender: None, |
| http_client_factory, |
| } |
| } |
|
|
| pub(crate) fn with_session_context( |
| mut self, |
| prompt_cache_key_override: Option<String>, |
| event_sender: Sender<ProtocolEvent>, |
| codex_responses_headers: Option<Arc<CodexResponsesHeaders>>, |
| ) -> Self { |
| self.prompt_cache_key_override = prompt_cache_key_override; |
| self.event_sender = Some(event_sender); |
| self.codex_responses_headers = codex_responses_headers; |
| self |
| } |
|
|
| fn prompt_cache_key(&self, responses_metadata: &CodexResponsesMetadata) -> String { |
| if let Some(prompt_cache_key) = &self.prompt_cache_key_override { |
| return prompt_cache_key.clone(); |
| } |
|
|
| if let SessionSource::Internal(source) = &self.state.session_source |
| && let Some(parent_thread_id) = responses_metadata.parent_thread_id |
| { |
| return format!("{source}:{parent_thread_id}"); |
| } |
|
|
| responses_metadata.session_id.clone() |
| } |
|
|
| |
| |
| fn responses_session_id(&self, metadata: &CodexResponsesMetadata) -> String { |
| if self.state.session_source.is_non_root_agent() { |
| metadata.session_id.clone() |
| } else { |
| self.prompt_cache_key(metadata) |
| } |
| } |
|
|
| |
| |
| |
| |
| pub fn new_session(&self) -> ModelClientSession { |
| let auth_owner_generation = self.auth_owner_generation(); |
| let mut websocket_session = self.take_cached_websocket_session(); |
| if websocket_session.auth_owner_generation != auth_owner_generation { |
| |
| websocket_session = WebsocketSession { |
| auth_owner_generation, |
| ..Default::default() |
| }; |
| } |
| ModelClientSession { |
| client: self.clone(), |
| websocket_session, |
| turn_state: Arc::new(OnceLock::new()), |
| } |
| } |
|
|
| pub(crate) fn auth_manager(&self) -> Option<Arc<AuthManager>> { |
| self.state.provider.auth_manager() |
| } |
|
|
| fn auth_owner_generation(&self) -> Option<u64> { |
| self.auth_manager().map(|manager| { |
| manager |
| .auth_change_state_receiver() |
| .borrow() |
| .owner_generation |
| }) |
| } |
|
|
| 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 |
| } |
|
|
| 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(ClientRouting::ConfiguredProvider) |
| .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, |
| client_setup.redirect_policy, |
| )?; |
| 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(crate) async fn realtime_sideband_headers( |
| &self, |
| mut extra_headers: ApiHeaderMap, |
| ) -> Result<ApiHeaderMap> { |
| let client_setup = self |
| .current_client_setup(ClientRouting::ConfiguredProvider) |
| .await?; |
| if let Some(header_value) = self.generate_attestation_header_for().await { |
| extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value); |
| } |
| extra_headers.extend(sideband_websocket_auth_headers( |
| client_setup.api_auth.as_ref(), |
| )); |
| Ok(extra_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(ClientRouting::ConfiguredProvider) |
| .await?; |
| let transport = self.build_api_transport( |
| &client_setup.api_provider, |
| MEMORIES_SUMMARIZE_ENDPOINT, |
| client_setup.redirect_policy, |
| )?; |
| 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(|effort| model_info.resolve_reasoning_effort(effort)) |
| .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( |
| &self, |
| model_info: &ModelInfo, |
| effort: Option<ReasoningEffortConfig>, |
| summary: ReasoningSummaryConfig, |
| ) -> Reasoning { |
| Reasoning { |
| effort: effort |
| .or_else(|| model_info.default_reasoning_level.clone()) |
| .map(|effort| model_info.resolve_reasoning_effort(effort)), |
| summary: (model_info.supports_reasoning_summary_parameter |
| && summary != ReasoningSummaryConfig::None) |
| .then_some(summary), |
| |
| |
| context: model_info |
| .use_responses_lite |
| .then_some(ReasoningContext::AllTurns), |
| } |
| } |
|
|
| pub(crate) fn build_responses_request( |
| &self, |
| 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); |
| let is_openai = self.state.provider.info().is_openai(); |
| let (instructions, tools) = if model_info.use_responses_lite { |
| |
| |
| let prefix_namespace = Uuid::new_v5( |
| &Uuid::NAMESPACE_OID, |
| self.state.thread_id.to_string().as_bytes(), |
| ); |
| let tools = if self.state.provider.capabilities().namespace_tools { |
| create_tools_json_for_responses_lite(&prompt.tools)? |
| } else { |
| create_tools_json_for_responses_api(&prompt.tools)? |
| }; |
| let mut prefix = vec![ResponseItem::AdditionalTools { |
| id: Some(ResponseItemId::with_suffix( |
| "at", |
| Uuid::new_v5(&prefix_namespace, &serde_json::to_vec(&tools)?), |
| )), |
| role: "developer".to_string(), |
| tools, |
| }]; |
| if !prompt.base_instructions.text.is_empty() { |
| let mut instructions = ContextualUserFragment::into(BaseInstructionsFragment( |
| prompt.base_instructions.text.clone(), |
| )); |
| instructions.set_id(Some(ResponseItemId::with_suffix( |
| "msg", |
| Uuid::new_v5(&prefix_namespace, prompt.base_instructions.text.as_bytes()), |
| ))); |
| prefix.push(instructions); |
| } |
| input.splice(0..0, prefix); |
| (String::new(), None) |
| } else { |
| ( |
| prompt.base_instructions.text.clone(), |
| Some(create_tools_raw_json_for_responses_api(&prompt.tools)?.into()), |
| ) |
| }; |
| if !is_openai { |
| for item in &mut input { |
| item.clear_internal_chat_message_metadata_passthrough(); |
| if let ResponseItem::FunctionCall { |
| encrypted_function_args, |
| .. |
| } = item |
| { |
| *encrypted_function_args = None; |
| } |
| } |
| } |
| 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: false, |
| stream: true, |
| stream_options, |
| include, |
| service_tier, |
| prompt_cache_key, |
| text, |
| client_metadata: Some(responses_metadata.client_metadata()), |
| access_programs: None, |
| }; |
| Ok(request) |
| } |
|
|
| fn filter_tool_result_metadata(input: &mut [ResponseItem], api_provider: &ApiProvider) { |
| |
| |
| let result_metadata_allowed = |
| url::Url::parse(&api_provider.base_url) |
| .ok() |
| .is_some_and(|url| { |
| url.scheme() == "https" |
| && url.host_str().is_some_and(|host| { |
| host == "api.openai.com" |
| || codex_http_client::is_allowed_chatgpt_host(host) |
| }) |
| }); |
| if !result_metadata_allowed { |
| for item in input { |
| item.clear_tool_result_metadata(); |
| } |
| } |
| } |
|
|
| 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); |
| } |
| if !self.state.content_item_kinds_enabled { |
| item.clear_content_item_kinds(); |
| } |
| } |
| } |
|
|
| |
| |
| |
| 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, routing: ClientRouting) -> Result<CurrentClientSetup> { |
| |
| |
| let auth_owner_generation = self.auth_owner_generation(); |
| let auth_changes = self |
| .state |
| .provider |
| .auth_manager() |
| .map(|manager| manager.auth_change_receiver()); |
| loop { |
| let revision = auth_changes.as_ref().map(|changes| *changes.borrow()); |
| let auth = self.state.provider.auth().await; |
| let (api_provider, redirect_policy) = match routing { |
| ClientRouting::Workspace => { |
| let resolved = self |
| .state |
| .provider |
| .responses_api_provider(&self.state.workspace_routing) |
| .await?; |
| (resolved.provider, resolved.redirect_policy) |
| } |
| ClientRouting::ConfiguredProvider => ( |
| self.state.provider.api_provider().await?, |
| ClientRedirectPolicy::Default, |
| ), |
| }; |
| 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?; |
| |
| |
| if self.state.provider.info().auth.is_none() { |
| if self.auth_owner_generation() != auth_owner_generation { |
| return Err(std::io::Error::other( |
| "account changed while preparing model request", |
| ) |
| .into()); |
| } |
| |
| if auth_changes.as_ref().map(|changes| *changes.borrow()) != revision { |
| continue; |
| } |
| } |
| return Ok(CurrentClientSetup { |
| auth, |
| auth_owner_generation, |
| auth_revision: revision, |
| api_provider, |
| redirect_policy, |
| api_auth: resolved_auth.auth, |
| agent_identity_telemetry: resolved_auth.agent_identity_telemetry, |
| }); |
| } |
| } |
|
|
| fn responses_headers(&self, auth: Option<&CodexAuth>, model: &str) -> ApiHeaderMap { |
| self.codex_responses_headers |
| .as_ref() |
| .filter(|config| { |
| config.model == model |
| && self.uses_codex_backend(auth) |
| && self.state.provider.info().supports_codex_backend_routes() |
| }) |
| .map(|config| config.headers.clone()) |
| .unwrap_or_default() |
| } |
|
|
| fn uses_codex_backend(&self, auth: Option<&CodexAuth>) -> bool { |
| let provider = self.state.provider.info(); |
| auth.is_some_and(CodexAuth::uses_codex_backend) |
| && provider.is_openai() |
| && provider.requires_openai_auth |
| && provider.env_key.is_none() |
| && provider.experimental_bearer_token.is_none() |
| && provider.auth.is_none() |
| && provider.aws.is_none() |
| } |
|
|
| fn set_guardian_metadata( |
| &self, |
| metadata: &mut Option<HashMap<String, String>>, |
| parent_response_id: Option<&str>, |
| auth: Option<&CodexAuth>, |
| responses_headers: &ApiHeaderMap, |
| ) { |
| if let Some(metadata) = metadata.as_mut() { |
| metadata.remove("guardian_credits_requested"); |
| metadata.remove("parent_response_id"); |
| } |
| let guardian_reviewer = responses_headers |
| .get("x-codex-guardian") |
| .is_some_and(|value| value == "reviewer"); |
| if guardian_reviewer && let Some(parent_response_id) = parent_response_id { |
| metadata.get_or_insert_with(HashMap::new).insert( |
| "parent_response_id".to_owned(), |
| parent_response_id.to_owned(), |
| ); |
| } |
| if !guardian_reviewer |
| && !crate::guardian::is_basic_session_source(&self.state.session_source) |
| && matches!( |
| auth, |
| Some(CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_)) |
| ) |
| && self.uses_codex_backend(auth) |
| && self.state.provider.info().supports_codex_backend_routes() |
| { |
| metadata |
| .get_or_insert_with(HashMap::new) |
| .insert("guardian_credits_requested".to_owned(), "true".to_owned()); |
| } |
| } |
|
|
| fn build_routing_hint_header( |
| &self, |
| auth: Option<&CodexAuth>, |
| model: &str, |
| service_tier: Option<&str>, |
| ) -> Option<HeaderValue> { |
| if !self.uses_codex_backend(auth) { |
| return None; |
| } |
|
|
| let routing_hint = match service_tier { |
| Some(tier) => format!("model={model};tier={tier}"), |
| None => format!("model={model}"), |
| }; |
| HeaderValue::from_str(&routing_hint).ok() |
| } |
|
|
| fn build_api_transport( |
| &self, |
| api_provider: &ApiProvider, |
| endpoint: &str, |
| redirect_policy: ClientRedirectPolicy, |
| ) -> Result<ReqwestTransport> { |
| let redirect_policy = if api_provider |
| .headers |
| .contains_key(codex_model_provider::ACCOUNT_ROUTING_HEADER) |
| { |
| ClientRedirectPolicy::Reject |
| } else { |
| redirect_policy |
| }; |
| let client = create_client_for_route( |
| &self.http_client_factory, |
| &api_provider.url_for_path(endpoint), |
| ClientRouteClass::Api, |
| redirect_policy, |
| ) |
| .map_err(std::io::Error::from)?; |
| Ok(ReqwestTransport::from_http_client(client)) |
| } |
|
|
| pub(crate) async fn prewarm_auth(&self) -> Result<()> { |
| self.current_client_setup(ClientRouting::Workspace) |
| .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, |
| responses_headers: &ApiHeaderMap, |
| ) -> std::result::Result<ApiWebSocketConnection, ApiError> { |
| let mut headers = self.build_websocket_headers(responses_metadata).await; |
| headers.extend(responses_headers.clone()); |
| 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(self.responses_session_id(responses_metadata)), |
| Some(responses_metadata.thread_id.to_string()), |
| )); |
| headers.extend(self.build_responses_compatibility_headers(responses_metadata)); |
| if let Some(routing_hint) = &responses_metadata.routing_hint { |
| headers.insert(X_CODEX_ROUTING_HINT_HEADER, routing_hint.clone()); |
| } |
| 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 { |
| fn reset_websocket_session(&mut self) { |
| self.websocket_session.connection = None; |
| self.websocket_session.responses_headers.clear(); |
| self.websocket_session.connection_key = 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(self.client.responses_session_id(responses_metadata)), |
| 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 response_items = |
| last_response.map_or(&[][..], |response| response.items_added.as_slice()); |
| let previous_items_len = previous_request |
| .input |
| .len() |
| .checked_add(response_items.len())?; |
| 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 previous_items = previous_request.input.iter().chain(response_items); |
| if !previous_items |
| .zip(request_items_to_compare) |
| .all(|(previous, current)| { |
| response_items_equal_ignoring_internal_metadata(previous, current) |
| }) |
| { |
| 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, |
| model_info: &ModelInfo, |
| session_telemetry: &SessionTelemetry, |
| responses_metadata: &CodexResponsesMetadata, |
| ) -> std::result::Result<(), ApiError> { |
| if !self.client.responses_websocket_enabled() { |
| return Ok(()); |
| } |
| let client_setup = self |
| .client |
| .current_client_setup(ClientRouting::Workspace) |
| .await |
| .map_err(|err| { |
| ApiError::Stream(format!( |
| "failed to build websocket prewarm client setup: {err}" |
| )) |
| })?; |
| let connection_key = |
| ResponsesConnectionKey::new(&client_setup.api_provider, client_setup.auth_revision); |
| if self.websocket_session.connection.is_some() |
| && self.websocket_session.connection_key.as_ref() == Some(&connection_key) |
| { |
| return Ok(()); |
| } |
| self.reset_websocket_session(); |
| 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 responses_headers = self |
| .client |
| .responses_headers(client_setup.auth.as_ref(), &model_info.slug); |
| self.websocket_connection(WebsocketConnectParams { |
| session_telemetry, |
| api_provider: client_setup.api_provider, |
| auth_revision: client_setup.auth_revision, |
| api_auth: client_setup.api_auth, |
| auth_owner_generation: client_setup.auth_owner_generation, |
| responses_metadata, |
| auth_context, |
| request_route_telemetry: RequestRouteTelemetry::for_endpoint("/responses"), |
| responses_headers: &responses_headers, |
| }) |
| .await?; |
| 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, |
| auth_revision, |
| api_auth, |
| auth_owner_generation, |
| responses_metadata, |
| auth_context, |
| request_route_telemetry, |
| responses_headers, |
| } = params; |
| let connection_key = ResponsesConnectionKey::new(&api_provider, auth_revision); |
| let needs_new = match self.websocket_session.connection.as_ref() { |
| Some(conn) => { |
| self.websocket_session.responses_headers != *responses_headers |
| || self.websocket_session.connection_key.as_ref() != Some(&connection_key) |
| || conn.is_closed().await |
| } |
| None => true, |
| }; |
| |
| let owner_changed = self.websocket_session.auth_owner_generation != auth_owner_generation |
| || self.client.auth_owner_generation() != auth_owner_generation; |
| if owner_changed { |
| self.turn_state = Arc::new(OnceLock::new()); |
| } |
|
|
| if needs_new || owner_changed { |
| self.reset_websocket_session(); |
| let new_conn = match self |
| .client |
| .connect_websocket( |
| session_telemetry, |
| api_provider, |
| api_auth, |
| responses_metadata, |
| auth_context, |
| request_route_telemetry, |
| responses_headers, |
| ) |
| .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.responses_headers = responses_headers.clone(); |
| self.websocket_session.auth_owner_generation = auth_owner_generation; |
| self.websocket_session.connection_key = Some(connection_key); |
| 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 = tracing::field::Empty, |
| 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 provider_auth_recovery_attempted = false; |
| let mut pending_retry = PendingUnauthorizedRetry::default(); |
| loop { |
| let client_setup = self |
| .client |
| .current_client_setup(ClientRouting::Workspace) |
| .await?; |
| let responses_headers = self |
| .client |
| .responses_headers(client_setup.auth.as_ref(), &model_info.slug); |
| tracing::Span::current().record("api.path", "/responses"); |
| let transport = self.client.build_api_transport( |
| &client_setup.api_provider, |
| "/responses", |
| client_setup.redirect_policy, |
| )?; |
| 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"), |
| 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( |
| prompt, |
| model_info, |
| effort.clone(), |
| summary, |
| service_tier.clone(), |
| responses_metadata, |
| )?; |
| ModelClient::filter_tool_result_metadata( |
| &mut request.input, |
| &client_setup.api_provider, |
| ); |
| self.client.set_guardian_metadata( |
| &mut request.client_metadata, |
| responses_metadata.parent_response_id.as_deref(), |
| client_setup.auth.as_ref(), |
| &responses_headers, |
| ); |
| let guardian_reviewer = responses_headers |
| .get("x-codex-guardian") |
| .is_some_and(|value| value == "reviewer"); |
| if guardian_reviewer { |
| request.service_tier = None; |
| } |
| if !guardian_reviewer |
| && let Some(header_value) = self.client.build_routing_hint_header( |
| client_setup.auth.as_ref(), |
| &request.model, |
| request.service_tier.as_deref(), |
| ) |
| { |
| options |
| .extra_headers |
| .insert(X_CODEX_ROUTING_HINT_HEADER, header_value); |
| } |
| request.access_programs = cyber_access_program::for_auth( |
| client_setup.auth.as_ref(), |
| prompt.cyber_access_program, |
| ); |
| self.client |
| .prepare_response_items_for_request(&mut request.input); |
| if crate::guardian::is_basic_session_source(&self.client.state.session_source) { |
| crate::guardian::observe_guardian_request(session_telemetry, &request); |
| } |
| let request_session_telemetry = |
| session_telemetry_for_request(session_telemetry, &request); |
| options.extra_headers.extend(responses_headers); |
| 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)) |
| if self |
| .client |
| .state |
| .provider |
| .is_recoverable_auth_error(&unauthorized_transport) => |
| { |
| 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, |
| &mut provider_auth_recovery_attempted, |
| session_telemetry, |
| &self.client.state.provider, |
| self.client.event_sender.as_ref(), |
| responses_metadata.turn_id.as_deref(), |
| ) |
| .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 = tracing::field::Empty, |
| 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 provider = Arc::clone(&self.client.state.provider); |
| let auth_manager = provider.auth_manager(); |
|
|
| let mut auth_recovery = auth_manager |
| .as_ref() |
| .map(AuthManager::unauthorized_recovery); |
| let mut provider_auth_recovery_attempted = false; |
| let mut pending_retry = PendingUnauthorizedRetry::default(); |
| loop { |
| let client_setup = self |
| .client |
| .current_client_setup(ClientRouting::Workspace) |
| .await?; |
| let responses_headers = self |
| .client |
| .responses_headers(client_setup.auth.as_ref(), &model_info.slug); |
| tracing::Span::current().record("api.path", "/responses"); |
| 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( |
| prompt, |
| model_info, |
| effort.clone(), |
| summary, |
| service_tier.clone(), |
| responses_metadata, |
| )?; |
| ModelClient::filter_tool_result_metadata( |
| &mut request.input, |
| &client_setup.api_provider, |
| ); |
| let guardian_reviewer = responses_headers |
| .get("x-codex-guardian") |
| .is_some_and(|value| value == "reviewer"); |
| if guardian_reviewer { |
| request.service_tier = None; |
| } |
| request.access_programs = cyber_access_program::for_auth( |
| client_setup.auth.as_ref(), |
| prompt.cyber_access_program, |
| ); |
| let mut websocket_metadata = responses_metadata.clone(); |
| websocket_metadata.routing_hint = if !guardian_reviewer { |
| self.client.build_routing_hint_header( |
| client_setup.auth.as_ref(), |
| &request.model, |
| request.service_tier.as_deref(), |
| ) |
| } else { |
| None |
| }; |
| let request_session_telemetry = if warmup { |
| |
| session_telemetry.clone() |
| } else { |
| session_telemetry_for_request(session_telemetry, &request) |
| }; |
| match self |
| .websocket_connection(WebsocketConnectParams { |
| session_telemetry, |
| api_provider: client_setup.api_provider, |
| auth_revision: client_setup.auth_revision, |
| api_auth: client_setup.api_auth, |
| auth_owner_generation: client_setup.auth_owner_generation, |
| responses_metadata: &websocket_metadata, |
| auth_context: request_auth_context, |
| request_route_telemetry: RequestRouteTelemetry::for_endpoint("/responses"), |
| responses_headers: &responses_headers, |
| }) |
| .await |
| { |
| Ok(_) => {} |
| Err(ApiError::Transport(TransportError::Http { status, .. })) |
| if status == StatusCode::UPGRADE_REQUIRED => |
| { |
| return Ok(WebsocketStreamOutcome::FallbackToHttp); |
| } |
| Err(ApiError::Transport(unauthorized_transport)) |
| if provider.is_recoverable_auth_error(&unauthorized_transport) => |
| { |
| pending_retry = PendingUnauthorizedRetry::from_recovery( |
| handle_unauthorized( |
| unauthorized_transport, |
| &mut auth_recovery, |
| &mut provider_auth_recovery_attempted, |
| session_telemetry, |
| &provider, |
| self.client.event_sender.as_ref(), |
| responses_metadata.turn_id.as_deref(), |
| ) |
| .await?, |
| ); |
| continue; |
| } |
| Err(err) => return Err(provider.map_api_error(err)), |
| } |
|
|
| |
| if !warmup |
| && crate::guardian::is_basic_session_source(&self.client.state.session_source) |
| { |
| crate::guardian::observe_guardian_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()); |
| } |
| 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 mut 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) |
| }; |
| self.client.set_guardian_metadata( |
| &mut ws_payload.client_metadata, |
| responses_metadata.parent_response_id.as_deref(), |
| client_setup.auth.as_ref(), |
| &responses_headers, |
| ); |
| 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, |
| usage_metadata, |
| 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, |
| usage_metadata, |
| 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 | AuthMode::BedrockAccessKeys => { |
| "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, |
| auth_revision: Option<u64>, |
| api_auth: SharedAuthProvider, |
| auth_owner_generation: Option<u64>, |
| responses_metadata: &'a CodexResponsesMetadata, |
| auth_context: AuthRequestTelemetryContext, |
| request_route_telemetry: RequestRouteTelemetry, |
| responses_headers: &'a ApiHeaderMap, |
| } |
|
|
| fn emit_auth_recovery_event( |
| event_sender: Option<&Sender<ProtocolEvent>>, |
| turn_id: Option<&str>, |
| provider: &SharedModelProvider, |
| message: &str, |
| event: fn(AuthRecoveryEvent) -> EventMsg, |
| ) { |
| if let (Some(sender), Some(turn_id)) = (event_sender, turn_id) { |
| let _ = sender.try_send(ProtocolEvent { |
| id: turn_id.to_string(), |
| msg: event(AuthRecoveryEvent { |
| provider: provider.info().name.clone(), |
| message: message.to_string(), |
| }), |
| }); |
| } |
| } |
|
|
| async fn handle_unauthorized( |
| transport: TransportError, |
| auth_recovery: &mut Option<UnauthorizedRecovery>, |
| provider_auth_recovery_attempted: &mut bool, |
| session_telemetry: &SessionTelemetry, |
| provider: &SharedModelProvider, |
| event_sender: Option<&Sender<ProtocolEvent>>, |
| turn_id: Option<&str>, |
| ) -> Result<UnauthorizedRecoveryExecution> { |
| let debug = extract_response_debug_context(&transport); |
| if !*provider_auth_recovery_attempted { |
| *provider_auth_recovery_attempted = true; |
| let messages = provider.auth_recovery_messages(); |
| if let Some(messages) = messages { |
| emit_auth_recovery_event( |
| event_sender, |
| turn_id, |
| provider, |
| messages.started, |
| EventMsg::AuthRecoveryStarted, |
| ); |
| } |
| match provider.recover_from_unauthorized().await { |
| Ok(ProviderUnauthorizedRecovery::Recovered) => { |
| if let Some(messages) = messages { |
| emit_auth_recovery_event( |
| event_sender, |
| turn_id, |
| provider, |
| messages.succeeded, |
| EventMsg::AuthRecoveryCompleted, |
| ); |
| } |
| return Ok(UnauthorizedRecoveryExecution { |
| mode: "provider", |
| phase: "provider_refresh", |
| }); |
| } |
| Ok(ProviderUnauthorizedRecovery::NotConfigured) => {} |
| Err(error) => { |
| let original = provider.map_api_error(ApiError::Transport(transport)); |
| warn!( |
| error = %error, |
| original_error = %original, |
| "provider authentication recovery failed" |
| ); |
| return Err(if error.is_retryable() { |
| original |
| } else { |
| error |
| }); |
| } |
| } |
| } |
|
|
| 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<StatusCode>, |
| 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; |
|
|