Download codex-rs/core/src/client.rs from SaylorTwift/codex: direct link, hf CLI and curl.
- Browser
- Download file 111 kB
-
https://huggingface.co/SaylorTwift/codex/resolve/main/codex-rs/core/src/client.rs
- Command line
-
hf download hf://SaylorTwift/codex/codex-rs/core/src/client.rs
-
curl -L -o client.rs https://huggingface.co/SaylorTwift/codex/resolve/main/codex-rs/core/src/client.rs
111 kB
| //! Session- and turn-scoped helpers for talking to model provider APIs. | |
| //! | |
| //! `ModelClient` is intended to live for the lifetime of a Codex session and holds the stable | |
| //! configuration and state needed to talk to a provider (auth, provider selection, conversation id, | |
| //! and transport fallback state). | |
| //! | |
| //! Per-turn settings (model selection, reasoning controls, telemetry context, and turn metadata) | |
| //! are passed explicitly to streaming and unary methods so that the turn lifetime is visible at the | |
| //! call site. | |
| //! | |
| //! A [`ModelClientSession`] is created per turn and is used to stream one or more Responses API | |
| //! requests during that turn. It caches a Responses WebSocket connection (opened lazily) and stores | |
| //! per-turn state such as the `x-codex-turn-state` token used for sticky routing. | |
| //! Cached connections, incremental response state, and turn routing are discarded when auth | |
| //! ownership changes. | |
| //! | |
| //! WebSocket prewarm is a v2-only `response.create` with `generate=false`; it waits for completion | |
| //! so the next request can reuse the same connection and `previous_response_id`. | |
| //! | |
| //! Turn execution performs prewarm as a best-effort step before the first stream request so the | |
| //! subsequent request can reuse the same connection. | |
| //! | |
| //! ## Retry-Budget Tradeoff | |
| //! | |
| //! WebSocket prewarm is treated as the first websocket connection attempt for a turn. If it | |
| //! fails, normal stream retry/fallback logic handles recovery on the same turn. | |
| 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; | |
| 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"; | |
| 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()), | |
| ) | |
| } | |
| /// Session-scoped state shared by all [`ModelClient`] clones. | |
| /// | |
| /// This is intentionally kept minimal so `ModelClient` does not need to hold a full `Config`. Most | |
| /// configuration is per turn and is passed explicitly to streaming/unary methods. | |
| 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, | |
| } | |
| /// Resolved API client setup for a single request attempt. | |
| /// | |
| /// Keeping this as a single bundle ensures prewarm and normal request paths | |
| /// share the same auth/provider setup flow. | |
| 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>, | |
| } | |
| struct RequestRouteTelemetry { | |
| endpoint: &'static str, | |
| } | |
| impl RequestRouteTelemetry { | |
| fn for_endpoint(endpoint: &'static str) -> Self { | |
| Self { endpoint } | |
| } | |
| } | |
| /// A session-scoped client for model-provider API calls. | |
| /// | |
| /// This holds configuration and state that should be shared across turns within a Codex session | |
| /// (auth, provider selection, thread id, and transport fallback state). | |
| /// | |
| /// WebSocket fallback is session-scoped: once a turn activates the HTTP fallback, subsequent turns | |
| /// will also use HTTP for the remainder of the session. | |
| /// | |
| /// Turn-scoped settings (model selection, reasoning controls, telemetry context, and turn | |
| /// metadata) are passed explicitly to the relevant methods to keep turn lifetime visible at the | |
| /// call site. | |
| 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, | |
| } | |
| /// A turn-scoped streaming session created from a [`ModelClient`]. | |
| /// | |
| /// The session establishes a Responses WebSocket connection lazily and reuses it across multiple | |
| /// requests within the turn. It also caches per-turn state: | |
| /// | |
| /// - The last full request, so subsequent calls can reuse incremental websocket request payloads | |
| /// only when the current request is an incremental extension of the previous one. | |
| /// - The `x-codex-turn-state` sticky-routing token, which must be replayed for all requests within | |
| /// the same turn. | |
| /// | |
| /// Create a fresh `ModelClientSession` for each Codex turn. Reusing it across turns would replay | |
| /// the previous turn's sticky-routing token into the next turn, which violates the client/server | |
| /// contract and can cause routing bugs. | |
| pub struct ModelClientSession { | |
| client: ModelClient, | |
| websocket_session: WebsocketSession, | |
| /// Turn state for sticky routing. | |
| /// | |
| /// This is an `OnceLock` that stores the turn state value received from the server | |
| /// on turn start via the `x-codex-turn-state` response header. Once set, this value | |
| /// should be sent back to the server in the `x-codex-turn-state` request header for | |
| /// all subsequent requests within the same turn to maintain sticky routing. | |
| /// | |
| /// This is a contract between the client and server: we receive it at turn start, | |
| /// keep sending it unchanged between turn requests (e.g., for retries, incremental | |
| /// appends, or continuation requests), and must not send it between different turns. | |
| /// An auth ownership change clears it so the new owner gets fresh routing state. | |
| turn_state: Arc<OnceLock<String>>, | |
| } | |
| struct LastResponse { | |
| response_id: String, | |
| items_added: Vec<ResponseItem>, | |
| } | |
| struct WebsocketSession { | |
| connection: Option<ApiWebSocketConnection>, | |
| responses_headers: ApiHeaderMap, | |
| /// Owner of the cached state, including before a connection is opened. | |
| 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>, | |
| } | |
| // This is intentionally not a `PartialEq` implementation: request equality includes `input` and | |
| // `client_metadata`, while websocket reuse compares the input separately and ignores metadata. | |
| // Access programs are authorized per response, including continuations, without replaying input. | |
| // Keep the destructuring exhaustive so new request fields require an explicit reuse decision. | |
| 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 | |
| // Stream options control delivery for this response, not the context | |
| // referenced by `previous_response_id`. | |
| && 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, | |
| } | |
| /// Result of opening a WebRTC Realtime call. | |
| /// | |
| /// The SDP answer goes back to the client. The call id and auth headers stay on the server so the | |
| /// ordinary Realtime WebSocket machinery can join the same in-progress call as a sideband | |
| /// controller. | |
| pub(crate) struct RealtimeWebrtcCallStart { | |
| pub(crate) sdp: String, | |
| pub(crate) call_id: String, | |
| pub(crate) sideband_headers: ApiHeaderMap, | |
| } | |
| /// Reuses the API-auth material that created the WebRTC call for the sideband WebSocket join. | |
| /// | |
| /// API-key sessions send that API bearer. ChatGPT-auth sessions send their bearer plus account id; | |
| /// transceiver is responsible for accepting that same call-create identity on the direct | |
| /// `api.openai.com` sideband path. | |
| 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 { | |
| /// Creates a new session-scoped `ModelClient`. | |
| /// | |
| /// All arguments are expected to be stable for the lifetime of a Codex session. Per-turn values | |
| /// are passed to [`ModelClientSession::stream`] (and other turn-scoped methods) explicitly. The | |
| /// HTTP client factory must come from the effective session configuration so every transport | |
| /// observes the resolved outbound proxy policy. | |
| 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() | |
| } | |
| // ChatGPT derives cache affinity from the Responses session-id header. Keep the | |
| // actual session identity in turn metadata, hooks, and history/notes requests. | |
| 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) | |
| } | |
| } | |
| /// Creates a fresh turn-scoped streaming session. | |
| /// | |
| /// This constructor does not perform network I/O itself; the session opens a websocket lazily | |
| /// when the first stream request is issued. | |
| 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 { | |
| // Drop the old owner's cache before this turn can establish fresh routing state. | |
| 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", | |
| /*inc*/ 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> { | |
| // Create the media call over HTTP first, then retain matching auth so realtime can attach | |
| // the server-side control WebSocket to the call id from that HTTP response. | |
| 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) | |
| } | |
| /// Builds memory summaries for each provided normalized raw memory. | |
| /// | |
| /// This is a unary call (no streaming) to `/v1/memories/trace_summarize`. | |
| /// | |
| /// The model selection, reasoning effort, and telemetry context are passed explicitly to keep | |
| /// `ModelClient` session-scoped. | |
| pub async fn summarize_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 | |
| } | |
| /// Builds request telemetry for unary API calls (e.g., Compact endpoint). | |
| 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), | |
| // When Responses Lite is disabled, omit context so Responses uses the default, | |
| // which is currently `current_turn`. | |
| 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 { | |
| // These prompt-only items are rebuilt on every request. Hash their visible payloads | |
| // within the thread so retries and resumed sessions preserve their identity. | |
| 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) { | |
| // Check the resolved destination only when sending, not for local budget estimates. | |
| // HTTP and WS (including v2 compaction) share this raw-metadata-only filter. | |
| 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(/*new_id*/ None); | |
| } | |
| if !self.state.content_item_kinds_enabled { | |
| item.clear_content_item_kinds(); | |
| } | |
| } | |
| } | |
| /// Returns whether the Responses-over-WebSocket transport is active for this session. | |
| /// | |
| /// WebSocket use is controlled by provider capability and session-scoped fallback state. | |
| pub fn responses_websocket_enabled(&self) -> bool { | |
| if !self.state.provider.info().supports_websockets | |
| || self.state.disable_websockets.load(Ordering::Relaxed) | |
| { | |
| return false; | |
| } | |
| true | |
| } | |
| /// Returns auth + provider configuration resolved from the current session auth state. | |
| /// | |
| /// This centralizes setup used by both prewarm and normal request paths so they stay in | |
| /// lockstep when auth/provider resolution changes. | |
| async fn current_client_setup(&self, routing: ClientRouting) -> Result<CurrentClientSetup> { | |
| // Capture before resolving credentials so an account switch during setup cannot label | |
| // an old connection with the new owner's revision. | |
| 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?; | |
| // Command-backed bearer refreshes cannot change workspace routing. Keep the captured | |
| // revisions so a refresh during setup still invalidates cached WebSocket state. | |
| 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()); | |
| } | |
| // Rebuild the entire bundle after a credential-only refresh, including routing. | |
| 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(|_| ()) | |
| } | |
| /// Opens a websocket connection using the same header and telemetry wiring as normal turns. | |
| /// | |
| /// Both startup prewarm and in-turn `needs_new` reconnects call this path so handshake | |
| /// behavior remains consistent across both flows. | |
| 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(), | |
| /*turn_state*/ 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, | |
| /*connection_reused*/ 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 | |
| } | |
| /// Builds websocket handshake headers for both prewarm and turn-time reconnect. | |
| async fn build_websocket_headers( | |
| &self, | |
| responses_metadata: &CodexResponsesMetadata, | |
| ) -> ApiHeaderMap { | |
| let mut headers = build_responses_headers( | |
| self.state.beta_features_header.as_deref(), | |
| /*turn_state*/ 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(/*connection_reused*/ false); | |
| } | |
| /// Builds shared Responses API transport options and request-body options. | |
| /// | |
| /// Keeping option construction in one place ensures request-scoped headers are consistent | |
| /// regardless of transport choice. | |
| 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)), | |
| } | |
| } | |
| /// Checks whether the current request is an incremental extension of the previous request. | |
| /// We only reuse an incremental input delta when non-input request fields are unchanged and | |
| /// `input` is a strict extension of the previous known input. Server-returned output items | |
| /// are treated as part of the baseline so we do not resend them. | |
| 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), | |
| /*allow_empty_delta*/ 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, | |
| ) | |
| } | |
| /// Opportunistically preconnects a websocket for this turn-scoped client session. | |
| /// | |
| /// This performs only connection setup; it never sends prompt payloads. | |
| 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(()) | |
| } | |
| /// Returns a websocket connection for this turn. | |
| 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, | |
| }; | |
| // Resolving an external auth provider can change ownership during client setup. | |
| 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(/*connection_reused*/ false); | |
| } else { | |
| self.websocket_session | |
| .set_connection_reused(/*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 | |
| } | |
| } | |
| /// Streams a turn via the OpenAI Responses API. | |
| /// | |
| /// Handles reasoning summaries, verbosity, and the `text` controls used for output schemas. | |
| 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(), | |
| /*output_items*/ &[], | |
| ); | |
| 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(), | |
| /*output_items*/ &[], | |
| ); | |
| return Err(err); | |
| } | |
| } | |
| } | |
| } | |
| /// Streams a turn via the Responses API over WebSocket transport. | |
| 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 { | |
| // `generate=false` prewarm is connection setup, not an inference request. | |
| 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)), | |
| } | |
| // Measure the complete logical request, not only the websocket delta. | |
| 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 { | |
| // Prewarm sends `generate=false`; it is connection setup, not a | |
| // model inference attempt that should appear in rollout traces. | |
| InferenceTraceAttempt::disabled() | |
| } else { | |
| inference_trace.start_attempt() | |
| }; | |
| if previous_response_id_from_untraced_warmup { | |
| // The transport can reuse an untraced warmup response id and omit the | |
| // already-sent input, but rollout replay needs the logical model-visible | |
| // request rather than the compressed websocket delta. | |
| 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(), | |
| /*output_items*/ &[], | |
| ); | |
| 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)); | |
| } | |
| } | |
| /// Builds request and SSE telemetry for streaming API calls. | |
| 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) | |
| } | |
| /// Builds telemetry for the Responses API WebSocket transport. | |
| 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 | |
| } | |
| 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, | |
| /*warmup*/ true, | |
| current_span_w3c_trace_context(), | |
| &disabled_trace, | |
| ) | |
| .await | |
| { | |
| Ok(WebsocketStreamOutcome::Stream(mut stream)) => { | |
| // Wait for the v2 warmup request to complete before sending the first turn request. | |
| 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), | |
| } | |
| } | |
| /// Streams a single model request within the current turn. | |
| /// | |
| /// The caller is responsible for passing per-turn settings explicitly (model selection, | |
| /// reasoning settings, telemetry context, and turn metadata). This method will prefer the | |
| /// Responses WebSocket transport when the provider supports it and it remains healthy, and will | |
| /// fall back to the HTTP Responses API transport otherwise. The trace context may be enabled or | |
| /// disabled, but is always explicit so transport paths do not need separate trace/no-trace | |
| /// branches. | |
| 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, | |
| /*warmup*/ 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 | |
| } | |
| } | |
| } | |
| /// Permanently disables WebSockets for this Codex session and resets WebSocket state. | |
| /// | |
| /// This is used after exhausting the provider retry budget, to force subsequent requests onto | |
| /// the HTTP transport. | |
| /// | |
| /// Returns `true` if this call activated fallback, or `false` if fallback was already active. | |
| pub(crate) fn try_switch_fallback_transport( | |
| &mut self, | |
| session_telemetry: &SessionTelemetry, | |
| model_info: &ModelInfo, | |
| ) -> bool { | |
| let activated = self | |
| .client | |
| .force_http_fallback(session_telemetry, model_info); | |
| self.websocket_session = WebsocketSession::default(); | |
| activated | |
| } | |
| } | |
| /// Stamp a ResponsesWsRequest with the current time. | |
| /// | |
| /// Meant to be called just before sending the request over the socket, to capture realistic | |
| /// transport timing. | |
| 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(), | |
| ); | |
| } | |
| /// Builds the extra headers attached to Responses API requests. | |
| /// | |
| /// These headers implement Codex-specific conventions: | |
| /// | |
| /// - `x-codex-beta-features`: comma-separated beta feature keys enabled for the session. | |
| /// - `x-codex-turn-state`: sticky routing token captured earlier in the turn. | |
| 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, | |
| ) | |
| } | |
| /// Handles a 401 response by optionally refreshing ChatGPT tokens once. | |
| /// | |
| /// When refresh succeeds, the caller should retry the API call; otherwise | |
| /// the mapped `CodexErr` is returned to the caller. | |
| struct UnauthorizedRecoveryExecution { | |
| mode: &'static str, | |
| phase: &'static str, | |
| } | |
| 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), | |
| } | |
| } | |
| } | |
| 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(), | |
| /*recovery_reason*/ 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(), | |
| /*recovery_reason*/ None, | |
| /*auth_state_changed*/ 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(), | |
| /*recovery_reason*/ None, | |
| /*auth_state_changed*/ 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, | |
| /*auth_state_changed*/ 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); | |
| } | |
| } | |
| mod tests; | |