Recover WebSocket transport after transient HTTP fallback (#58950)

## Why

A transient WebSocket failure previously kept the session on HTTP for its remaining lifetime, even if WebSocket connectivity recovered.

## What changed

- Probe WebSocket connectivity in the background during HTTP fallback, allowing a subsequent request to use a ready connection without delaying HTTP requests.
- Verify recovery with a real request and bound its first-response wait to five seconds. Keep HTTP 426 responses and recovery timeouts on HTTP for the rest of the session.
- Keep retries of an HTTP request on HTTP. If a recovered connection fails before output, replay the request over HTTP without consuming its ordinary retry budget; failures after output use normal retry accounting.
- Rate-limit recovery probes and honor `Retry-After` deadlines.

## Testing

Extend integration coverage for successful recovery within a turn, disconnects before and after output, full-context replay followed by incremental WebSocket requests, persistent HTTP 426 fallback, and retry deadlines.

GitOrigin-RevId: 1c139a4bbce1a67442db5622fd70ff0cc8056c9c
This commit is contained in:
Charlie Marsh
2026-11-02 19:02:50 +00:00
committed by copyberry
parent 00df7f65ab
commit 0a2dda4b5e
10 changed files with 959 additions and 260 deletions
+296 -179
View File
@@ -29,11 +29,15 @@ 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;
#[path = "client_tool_metadata.rs"]
mod tool_metadata;
#[path = "client_websocket_recovery.rs"]
mod websocket_recovery;
use websocket_recovery::Connection;
use websocket_recovery::RecoveryAttempt;
use websocket_recovery::WebsocketRoutes;
use crate::CodexResponsesHeaders;
use crate::PromptCacheKey;
@@ -138,6 +142,7 @@ use crate::cyber_access_program;
use crate::feedback_tags;
use crate::responses_metadata::CodexResponsesMetadata;
use crate::responses_metadata::subagent_header_value;
use crate::responses_retry::ResponsesRequestTransport;
use crate::util::emit_feedback_auth_recovery_tags;
use codex_feedback::FeedbackRequestTags;
use codex_feedback::emit_feedback_request_tags_with_auth_env;
@@ -223,7 +228,6 @@ struct ModelClientState {
output_token_replay_enabled: bool,
include_attestation: bool,
attestation_provider: Option<Arc<dyn AttestationProvider>>,
disable_websockets: AtomicBool,
agent_identity_session_fallback: AgentIdentitySessionFallback,
cached_websocket_session: StdMutex<WebsocketSession>,
/// Last full tool list used for sampling, retained across turns and connection resets.
@@ -289,8 +293,8 @@ impl RequestRouteTelemetry {
/// 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.
/// HTTP fallback is shared across turns. After a transient failure, a subsequent request can
/// take over a freshly connected socket. An HTTP 426 keeps the session on HTTP.
///
/// 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
@@ -298,6 +302,7 @@ impl RequestRouteTelemetry {
#[derive(Debug, Clone)]
pub struct ModelClient {
state: Arc<ModelClientState>,
websocket_routes: WebsocketRoutes,
agent_identity_policy: AgentIdentityAuthPolicy,
prompt_cache_key_policy: PromptCacheKey,
session_id_override: Option<String>,
@@ -539,9 +544,19 @@ impl WebsocketSession {
}
}
enum WebsocketStreamOutcome {
Stream(ResponseStream),
FallbackToHttp,
pub(crate) enum HttpFallbackPolicy {
StayOnHttp,
AllowProbes,
}
enum WebsocketHandshake {
Connected,
UpgradeRequired,
}
enum WebsocketPreconnectMode<'a> {
ForegroundWithAuthRecovery(&'a mut Arc<OnceLock<String>>),
ProbeWithoutAuthRecovery,
}
/// Result of opening a WebRTC Realtime call.
@@ -632,11 +647,11 @@ impl ModelClient {
output_token_replay_enabled,
include_attestation,
attestation_provider,
disable_websockets: AtomicBool::new(false),
agent_identity_session_fallback: AgentIdentitySessionFallback::default(),
cached_websocket_session: StdMutex::new(WebsocketSession::default()),
last_inference_tools: StdMutex::new(None),
}),
websocket_routes: WebsocketRoutes::default(),
agent_identity_policy,
prompt_cache_key_policy: PromptCacheKey::default(),
session_id_override: None,
@@ -788,27 +803,6 @@ impl ModelClient {
.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,
@@ -1188,13 +1182,7 @@ impl ModelClient {
///
/// 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
self.state.provider.info().supports_websockets && self.websocket_routes.enabled()
}
/// Returns auth + provider configuration resolved from the current session auth state.
@@ -1699,7 +1687,43 @@ impl ModelClientSession {
if !self.client.responses_websocket_enabled() {
return Ok(());
}
let provider = Arc::clone(&self.client.state.provider);
if let WebsocketHandshake::UpgradeRequired = self
.websocket_session
.preconnect(
&self.client,
WebsocketPreconnectMode::ForegroundWithAuthRecovery(&mut self.turn_state),
model_info,
service_tier,
session_telemetry,
responses_metadata,
)
.await?
{
self.try_switch_fallback_transport(session_telemetry, HttpFallbackPolicy::StayOnHttp);
}
Ok(())
}
}
impl WebsocketSession {
/// Connect without sending a model request. A speculative connection leaves auth recovery
/// to the foreground request and has no turn routing state to update.
async fn preconnect(
&mut self,
client: &ModelClient,
mode: WebsocketPreconnectMode<'_>,
model_info: &ModelInfo,
service_tier: Option<String>,
session_telemetry: &SessionTelemetry,
responses_metadata: &CodexResponsesMetadata,
) -> Result<WebsocketHandshake> {
let (mut turn_state, recover_auth) = match mode {
WebsocketPreconnectMode::ForegroundWithAuthRecovery(turn_state) => {
(Some(turn_state), true)
}
WebsocketPreconnectMode::ProbeWithoutAuthRecovery => (None, false),
};
let provider = Arc::clone(&client.state.provider);
let auth_manager = provider.auth_manager();
let mut auth_recovery = auth_manager
.as_ref()
@@ -1707,8 +1731,7 @@ impl ModelClientSession {
let mut provider_auth_recovery_attempted = false;
let mut pending_retry = PendingUnauthorizedRetry::default();
loop {
let client_setup = self
.client
let client_setup = client
.current_client_setup(ClientRouting::Workspace)
.await?;
let auth_context = AuthRequestTelemetryContext::new(
@@ -1718,11 +1741,10 @@ impl ModelClientSession {
pending_retry,
)
.map_err(|error| provider.map_api_error(error))?;
let responses_headers = self
.client
.responses_headers(client_setup.auth.as_ref(), &model_info.slug);
let responses_headers =
client.responses_headers(client_setup.auth.as_ref(), &model_info.slug);
let mut websocket_metadata = responses_metadata.clone();
websocket_metadata.routing_hint = self.client.build_routing_hint_header(
websocket_metadata.routing_hint = client.build_routing_hint_header(
client_setup.auth.as_ref(),
&responses_headers,
&model_info.slug,
@@ -1731,30 +1753,34 @@ impl ModelClientSession {
.as_deref(),
);
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,
http_client_factory: client_setup.http_client_factory,
auth_owner_generation: client_setup.auth_owner_generation,
responses_metadata: &websocket_metadata,
auth_context,
request_route_telemetry: RequestRouteTelemetry::for_endpoint("/responses"),
responses_headers: &responses_headers,
runtime_id: client_setup.runtime_id,
})
.connection(
client,
turn_state.as_deref_mut(),
WebsocketConnectParams {
session_telemetry,
api_provider: client_setup.api_provider,
auth_revision: client_setup.auth_revision,
api_auth: client_setup.api_auth,
http_client_factory: client_setup.http_client_factory,
auth_owner_generation: client_setup.auth_owner_generation,
responses_metadata: &websocket_metadata,
auth_context,
request_route_telemetry: RequestRouteTelemetry::for_endpoint("/responses"),
responses_headers: &responses_headers,
runtime_id: client_setup.runtime_id,
},
)
.await
{
Ok(_) => return Ok(()),
Ok(_) => return Ok(WebsocketHandshake::Connected),
Err(ApiError::Transport(TransportError::Http { status, .. }))
if status == StatusCode::UPGRADE_REQUIRED =>
{
self.try_switch_fallback_transport(session_telemetry, model_info);
return Ok(());
return Ok(WebsocketHandshake::UpgradeRequired);
}
Err(ApiError::Transport(unauthorized_transport))
if provider.is_recoverable_auth_error(&unauthorized_transport) =>
if recover_auth
&& provider.is_recoverable_auth_error(&unauthorized_transport) =>
{
pending_retry = PendingUnauthorizedRetry::from_recovery(
handle_unauthorized(
@@ -1763,7 +1789,7 @@ impl ModelClientSession {
&mut provider_auth_recovery_attempted,
session_telemetry,
&provider,
self.client.event_sender.as_ref(),
client.event_sender.as_ref(),
responses_metadata.turn_id.as_deref(),
)
.await?,
@@ -1774,27 +1800,44 @@ impl ModelClientSession {
}
}
fn matches_setup(
&self,
params: &WebsocketConnectParams<'_>,
key: &ResponsesConnectionKey,
) -> bool {
self.connection_key.as_ref() == Some(key)
&& self.responses_headers == *params.responses_headers
&& self.runtime_id == params.runtime_id
&& self.parent_routing_affinity.as_ref() == params.parent_routing_affinity()
}
/// Returns a websocket connection for this turn.
#[instrument(
name = "model_client.websocket_connection",
level = "info",
skip_all,
fields(
provider = %self.client.state.provider.info().name,
wire_api = %self.client.state.provider.info().wire_api,
provider = %client.state.provider.info().name,
wire_api = %client.state.provider.info().wire_api,
transport = "responses_websocket",
api.path = "/responses",
turn.has_metadata_header = params.responses_metadata.has_turn_metadata()
)
)]
async fn websocket_connection(
async fn connection(
&mut self,
client: &ModelClient,
turn_state: Option<&mut Arc<OnceLock<String>>>,
params: WebsocketConnectParams<'_>,
) -> std::result::Result<&ApiWebSocketConnection, ApiError> {
) -> std::result::Result<(), ApiError> {
let connection_key =
ResponsesConnectionKey::new(&params.api_provider, params.auth_revision);
let matches_setup = self.matches_setup(&params, &connection_key);
let parent_routing_affinity = params.parent_routing_affinity().cloned();
let WebsocketConnectParams {
session_telemetry,
api_provider,
auth_revision,
auth_revision: _,
api_auth,
http_client_factory,
auth_owner_generation,
@@ -1804,41 +1847,26 @@ impl ModelClientSession {
responses_headers,
runtime_id,
} = params;
let parent_routing_affinity = if is_guardian_reviewer(responses_headers) {
responses_metadata.parent_routing_affinity.clone()
} else {
None
};
let connection_key = ResponsesConnectionKey::new(&api_provider, auth_revision);
let reset_reason = match self.websocket_session.connection.as_ref() {
Some(_)
if self.websocket_session.responses_headers != *responses_headers
|| self.websocket_session.connection_key.as_ref() != Some(&connection_key)
|| self.websocket_session.runtime_id != runtime_id
|| self.websocket_session.parent_routing_affinity
!= parent_routing_affinity =>
{
Some("other")
}
let reset_reason = match self.connection.as_ref() {
Some(_) if !matches_setup => Some("other"),
Some(conn) if conn.is_closed().await => Some("connection_closed"),
Some(_) | None => None,
};
let needs_new = self.websocket_session.connection.is_none() || reset_reason.is_some();
let needs_new = self.connection.is_none() || reset_reason.is_some();
// 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());
let owner_changed = self.auth_owner_generation != auth_owner_generation
|| client.auth_owner_generation() != auth_owner_generation;
if owner_changed && let Some(turn_state) = turn_state {
*turn_state = Arc::new(OnceLock::new());
}
if needs_new || owner_changed {
self.websocket_session.reset(if owner_changed {
self.reset(if owner_changed {
Some("other")
} else {
reset_reason
});
let new_conn = self
.client
let new_conn = client
.connect_websocket(
session_telemetry,
api_provider,
@@ -1851,25 +1879,18 @@ impl ModelClientSession {
parent_routing_affinity.as_ref(),
)
.await?;
self.websocket_session.connection = Some(new_conn);
self.websocket_session.parent_routing_affinity = parent_routing_affinity;
self.websocket_session.responses_headers = responses_headers.clone();
self.websocket_session.runtime_id = runtime_id;
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);
self.connection = Some(new_conn);
self.parent_routing_affinity = parent_routing_affinity;
self.responses_headers = responses_headers.clone();
self.runtime_id = runtime_id;
self.auth_owner_generation = auth_owner_generation;
self.connection_key = Some(connection_key);
self.set_connection_reused(/*connection_reused*/ false);
} else {
self.websocket_session
.set_connection_reused(/*connection_reused*/ true);
self.set_connection_reused(/*connection_reused*/ true);
}
self.websocket_session
.connection
.as_ref()
.ok_or(ApiError::Stream(
"websocket connection is unavailable".to_string(),
))
Ok(())
}
}
@@ -2036,7 +2057,7 @@ impl ResponsesHttpExecutor<'_> {
inference_trace_attempt,
Arc::clone(&self.client.state.provider),
interceptors,
/*last_response_tx*/ None,
/*websocket*/ None,
));
}
Err(ApiError::Transport(unauthorized_transport))
@@ -2083,6 +2104,11 @@ impl ResponsesHttpExecutor<'_> {
}
}
enum WebsocketRequestMode<'a> {
Warmup,
Generate(&'a mut ResponsesRequestTransport),
}
impl ModelClientSession {
/// Streams a turn via the Responses API over WebSocket transport.
#[allow(clippy::too_many_arguments)]
@@ -2096,7 +2122,7 @@ impl ModelClientSession {
transport = "responses_websocket",
api.path = tracing::field::Empty,
turn.has_metadata_header = responses_metadata.has_turn_metadata(),
websocket.warmup = warmup
websocket.warmup = matches!(mode, WebsocketRequestMode::Warmup)
)
)]
async fn stream_responses_websocket(
@@ -2108,10 +2134,11 @@ impl ModelClientSession {
summary: ReasoningSummaryConfig,
service_tier: Option<String>,
responses_metadata: &CodexResponsesMetadata,
warmup: bool,
mut mode: WebsocketRequestMode<'_>,
request_trace: Option<W3cTraceContext>,
inference_trace: &InferenceTraceContext,
) -> Result<WebsocketStreamOutcome> {
) -> Result<Option<ResponseStream>> {
let warmup = matches!(mode, WebsocketRequestMode::Warmup);
let after_prewarm = !warmup && self.websocket_session.last_response_from_untraced_warmup;
let provider = Arc::clone(&self.client.state.provider);
let auth_manager = provider.auth_manager();
@@ -2142,6 +2169,23 @@ impl ModelClientSession {
pending_retry,
)
.map_err(|error| self.client.state.provider.map_api_error(error))?;
let params = WebsocketConnectParams {
session_telemetry,
api_provider: client_setup.api_provider,
auth_revision: client_setup.auth_revision,
api_auth: client_setup.api_auth,
http_client_factory: client_setup.http_client_factory,
auth_owner_generation: client_setup.auth_owner_generation,
runtime_id: client_setup.runtime_id,
responses_metadata,
auth_context: request_auth_context,
request_route_telemetry: RequestRouteTelemetry::for_endpoint("/responses"),
responses_headers: &responses_headers,
};
let Some(connection) = websocket_recovery::select_for_request(self, &mut mode, &params)
else {
return Ok(None);
};
let mut request = self.client.build_responses_request(
prompt,
model_info,
@@ -2159,40 +2203,45 @@ impl ModelClientSession {
prompt.cyber_access_program,
self.client.api_key_cyber_access_programs,
)?;
let mut websocket_metadata = responses_metadata.clone();
websocket_metadata.routing_hint = self.client.build_routing_hint_header(
client_setup.auth.as_ref(),
&responses_headers,
&request.model,
request.service_tier.as_deref(),
);
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,
http_client_factory: client_setup.http_client_factory,
auth_owner_generation: client_setup.auth_owner_generation,
runtime_id: client_setup.runtime_id,
responses_metadata: &websocket_metadata,
auth_context: request_auth_context,
request_route_telemetry: RequestRouteTelemetry::for_endpoint("/responses"),
responses_headers: &responses_headers,
})
.await
{
Ok(_) => {}
let connected = match connection {
Connection::Recovered(attempt) => Ok(Some(attempt)),
Connection::Websocket => {
let mut websocket_metadata = responses_metadata.clone();
websocket_metadata.routing_hint = self.client.build_routing_hint_header(
client_setup.auth.as_ref(),
&responses_headers,
&request.model,
request.service_tier.as_deref(),
);
self.websocket_session
.connection(
&self.client,
Some(&mut self.turn_state),
WebsocketConnectParams {
responses_metadata: &websocket_metadata,
..params
},
)
.await
.map(|()| None)
}
};
let mut recovery = match connected {
Ok(recovery) => recovery,
Err(ApiError::Transport(TransportError::Http { status, .. }))
if status == StatusCode::UPGRADE_REQUIRED =>
{
return Ok(WebsocketStreamOutcome::FallbackToHttp);
self.try_switch_fallback_transport(
session_telemetry,
HttpFallbackPolicy::StayOnHttp,
);
return Ok(None);
}
Err(ApiError::Transport(unauthorized_transport))
if provider.is_recoverable_auth_error(&unauthorized_transport) =>
@@ -2212,7 +2261,7 @@ impl ModelClientSession {
continue;
}
Err(err) => return Err(provider.map_api_error(err)),
}
};
let mut client_metadata = self.client.build_ws_client_metadata(
responses_metadata,
@@ -2350,10 +2399,13 @@ impl ModelClientSession {
],
);
let stream_result = websocket_connection
.stream_request(
.stream_request_with_ack_timeout(
ws_request,
self.websocket_session.connection_reused(),
Some(Arc::clone(&self.turn_state)),
recovery
.as_ref()
.map(|_| websocket_recovery::RECOVERY_ACK_TIMEOUT),
)
.await;
if let Some(original_item_ids) = original_item_ids {
@@ -2366,6 +2418,9 @@ impl ModelClientSession {
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);
if let Some(attempt) = recovery.take() {
attempt.failed(&err);
}
inference_trace_attempt.record_failed(
&err,
response_debug_context.request_id.as_deref(),
@@ -2380,10 +2435,13 @@ impl ModelClientSession {
inference_trace_attempt,
Arc::clone(&self.client.state.provider),
interceptors,
Some(last_response_tx),
Some(WebsocketResponse {
last_response_tx,
recovery,
}),
);
self.websocket_session.last_response_rx = Some(last_response_rx);
return Ok(WebsocketStreamOutcome::Stream(stream));
return Ok(Some(stream));
}
}
@@ -2462,13 +2520,13 @@ impl ModelClientSession {
summary,
service_tier,
responses_metadata,
/*warmup*/ true,
WebsocketRequestMode::Warmup,
current_span_w3c_trace_context(),
&disabled_trace,
)
.await
{
Ok(WebsocketStreamOutcome::Stream(mut stream)) => {
Ok(Some(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 {
@@ -2479,10 +2537,7 @@ impl ModelClientSession {
}
Ok(())
}
Ok(WebsocketStreamOutcome::FallbackToHttp) => {
self.try_switch_fallback_transport(session_telemetry, model_info);
Ok(())
}
Ok(None) => Ok(()),
Err(err) => Err(err),
}
}
@@ -2506,13 +2561,42 @@ impl ModelClientSession {
service_tier: Option<String>,
responses_metadata: &CodexResponsesMetadata,
inference_trace: &InferenceTraceContext,
) -> Result<ResponseStream> {
self.stream_attempt(
prompt,
model_info,
session_telemetry,
effort,
summary,
service_tier,
responses_metadata,
inference_trace,
&mut ResponsesRequestTransport::Auto,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn stream_attempt(
&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,
transport: &mut ResponsesRequestTransport,
) -> 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
let websocket_available = matches!(transport, ResponsesRequestTransport::Auto)
&& self.client.state.provider.info().supports_websockets
&& self.client.websocket_routes.can_stream();
if websocket_available
&& let Some(stream) = self
.stream_responses_websocket(
prompt,
model_info,
@@ -2521,19 +2605,23 @@ impl ModelClientSession {
summary,
service_tier.clone(),
responses_metadata,
/*warmup*/ false,
request_trace,
WebsocketRequestMode::Generate(transport),
current_span_w3c_trace_context(),
inference_trace,
)
.await?
{
WebsocketStreamOutcome::Stream(stream) => return Ok(stream),
WebsocketStreamOutcome::FallbackToHttp => {
self.try_switch_fallback_transport(session_telemetry, model_info);
}
}
{
return Ok(stream);
}
*transport = ResponsesRequestTransport::HttpOnly;
websocket_recovery::reconnect_if_due(
&self.client,
model_info,
service_tier.as_deref(),
session_telemetry,
responses_metadata,
);
ResponsesHttpExecutor::new(&self.client)
.stream_responses_http(
&self.turn_state,
@@ -2553,21 +2641,25 @@ impl ModelClientSession {
}
}
/// 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.
/// Returns whether this call moved the session from WebSocket to HTTP.
pub(crate) fn try_switch_fallback_transport(
&mut self,
session_telemetry: &SessionTelemetry,
model_info: &ModelInfo,
policy: HttpFallbackPolicy,
) -> bool {
let activated = self
.client
.force_http_fallback(session_telemetry, model_info);
self.websocket_session = WebsocketSession::default();
let activated = self.client.state.provider.info().supports_websockets
&& self.client.websocket_routes.fall_back(policy);
if activated {
warn!("falling back to HTTP");
session_telemetry.counter(
"codex.transport.fallback_to_http",
/*inc*/ 1,
&[("from_wire_api", "responses_websocket")],
);
}
self.client
.store_cached_websocket_session(WebsocketSession::default());
self.websocket_session.reset(/*reason*/ None);
activated
}
}
@@ -2631,13 +2723,18 @@ fn add_responses_lite_header(headers: &mut ApiHeaderMap, use_responses_lite: boo
const RESPONSE_STREAM_CHANNEL_CAPACITY: usize = 1600;
const STREAM_DROPPED_REASON: &str = "response stream dropped before provider terminal event";
struct WebsocketResponse {
last_response_tx: oneshot::Sender<LastResponse>,
recovery: Option<RecoveryAttempt>,
}
fn map_response_stream(
mut api_stream: codex_api::ResponseStream,
session_telemetry: SessionTelemetry,
inference_trace_attempt: InferenceTraceAttempt,
provider: SharedModelProvider,
interceptors: Vec<Box<dyn codex_extension_api::ModelResponseInterceptor>>,
last_response_tx: Option<oneshot::Sender<LastResponse>>,
websocket: Option<WebsocketResponse>,
) -> ResponseStream {
let routing_affinity = api_stream.routing_affinity.take();
let upstream_request_id = api_stream.upstream_request_id.take();
@@ -2648,7 +2745,7 @@ fn map_response_stream(
session_telemetry,
inference_trace_attempt,
provider,
last_response_tx,
websocket,
);
stream.interrupt = interrupt;
stream.routing_affinity = routing_affinity;
@@ -2661,7 +2758,7 @@ fn map_response_events<S>(
session_telemetry: SessionTelemetry,
inference_trace_attempt: InferenceTraceAttempt,
provider: SharedModelProvider,
mut last_response_tx: Option<oneshot::Sender<LastResponse>>,
mut websocket: Option<WebsocketResponse>,
) -> ResponseStream
where
S: futures::Stream<Item = std::result::Result<ResponseEvent, ApiError>>
@@ -2675,9 +2772,10 @@ where
let consumer_dropped_for_stream = consumer_dropped.clone();
tokio::spawn(async move {
let mut recovery = websocket.as_mut().and_then(|ws| ws.recovery.take());
let mut logged_error = false;
let collect_output_items =
last_response_tx.is_some() || inference_trace_attempt.needs_output_items();
websocket.is_some() || inference_trace_attempt.needs_output_items();
let mut items_added: Vec<ResponseItem> = Vec::new();
let mut active_item_id: Option<ResponseItemId> = None;
let (request_start, mut ttft_ms) = (Instant::now(), None);
@@ -2701,6 +2799,9 @@ where
let Some(event) = event else {
break;
};
if let (Some(attempt), Ok(event)) = (&mut recovery, &event) {
attempt.observe(event);
}
match event {
Ok(ResponseEvent::OutputItemDone(mut item)) => {
// Assign IDs after interceptors, which may replace provider items.
@@ -2745,8 +2846,11 @@ where
&token_usage,
&items_added,
);
if let Some(sender) = last_response_tx.take() {
let _ = sender.send(LastResponse {
if let Some(attempt) = recovery.take() {
attempt.completed(&session_telemetry);
}
if let Some(ws) = websocket.take() {
let _ = ws.last_response_tx.send(LastResponse {
response_id: response_id.clone(),
items_added,
});
@@ -2795,6 +2899,9 @@ where
feedback_tags!(last_model_request_id = upstream_request_id);
}
let mapped = provider.map_api_error(err);
if let Some(attempt) = recovery.take() {
attempt.failed(&mapped);
}
// Terminal response output is trace evidence. Only streamed done items
// belong to the continuation cache or reach tool dispatch.
let completed_item_count = items_added.len();
@@ -2829,11 +2936,11 @@ where
}
}
}
inference_trace_attempt.record_failed(
"stream closed before response.completed",
upstream_request_id,
&items_added,
);
let reason = "stream closed before response.completed";
if let Some(attempt) = recovery {
attempt.failed(&CodexErr::Stream(reason.to_owned()));
}
inference_trace_attempt.record_failed(reason, upstream_request_id, &items_added);
});
ResponseStream {
@@ -2929,6 +3036,16 @@ struct WebsocketConnectParams<'a> {
responses_headers: &'a ApiHeaderMap,
}
impl WebsocketConnectParams<'_> {
fn parent_routing_affinity(&self) -> Option<&codex_http_client::ChatGptRoutingAffinity> {
if is_guardian_reviewer(self.responses_headers) {
self.responses_metadata.parent_routing_affinity.as_ref()
} else {
None
}
}
}
fn emit_auth_recovery_event(
event_sender: Option<&Sender<ProtocolEvent>>,
turn_id: Option<&str>,
+11 -8
View File
@@ -1676,7 +1676,7 @@ async fn untraced_http_stream_releases_completed_output_before_response_ends() {
test_session_telemetry(),
InferenceTraceAttempt::disabled(),
test_model_provider(),
/*last_response_tx*/ None,
/*websocket*/ None,
);
assert!(matches!(
@@ -1718,7 +1718,7 @@ async fn incomplete_response_traces_attempted_call_without_streaming_it() -> any
test_session_telemetry(),
attempt,
test_model_provider(),
/*last_response_tx*/ None,
/*websocket*/ None,
);
assert!(matches!(stream.next().await, Some(Err(_))));
assert!(stream.next().await.is_none());
@@ -1763,7 +1763,7 @@ async fn dropped_response_stream_traces_cancelled_partial_output() -> anyhow::Re
test_session_telemetry(),
attempt,
test_model_provider(),
/*last_response_tx*/ None,
/*websocket*/ None,
);
let observed = stream
@@ -1820,7 +1820,7 @@ async fn response_stream_records_last_model_feedback_ids() {
test_session_telemetry(),
InferenceTraceAttempt::disabled(),
test_model_provider(),
/*last_response_tx*/ None,
/*websocket*/ None,
);
while stream.next().await.is_some() {}
@@ -2044,7 +2044,7 @@ async fn dropped_backpressured_response_stream_traces_cancelled_partial_output()
test_session_telemetry(),
attempt,
test_model_provider(),
/*last_response_tx*/ None,
/*websocket*/ None,
);
// Fill the mapper channel with non-terminal events, then yield one output
@@ -2347,7 +2347,10 @@ async fn intercepted_output_reaches_trace_and_websocket_bookkeeping() -> anyhow:
attempt,
test_model_provider(),
vec![Box::new(ReplaceOutput)],
Some(last_response_tx),
Some(super::WebsocketResponse {
last_response_tx,
recovery: None,
}),
);
let mut delivered = Vec::new();
while let Some(event) = stream.next().await {
@@ -2373,7 +2376,6 @@ async fn intercepted_output_reaches_trace_and_websocket_bookkeeping() -> anyhow:
async fn inference_tools_changes_follow_full_specs_across_turns() {
let client = test_model_client(SessionSource::Cli);
let telemetry = test_session_telemetry();
let model_info = test_model_info();
let alpha = ResponsesApiTool {
name: "alpha".into(),
description: "Original".into(),
@@ -2405,7 +2407,8 @@ async fn inference_tools_changes_follow_full_specs_across_turns() {
.enumerate()
{
if index == 4 {
session.try_switch_fallback_transport(&telemetry, &model_info);
session
.try_switch_fallback_transport(&telemetry, super::HttpFallbackPolicy::StayOnHttp);
drop(session);
session = client.new_session();
}
@@ -0,0 +1,294 @@
//! Routes requests during HTTP fallback and publishes at most one background handshake.
use super::HttpFallbackPolicy;
use super::ModelClient;
use super::ModelClientSession;
use super::ResponseEvent;
use super::ResponsesConnectionKey;
use super::WebsocketConnectParams;
use super::WebsocketHandshake;
use super::WebsocketPreconnectMode;
use super::WebsocketRequestMode;
use super::WebsocketSession;
use crate::responses_metadata::CodexResponsesMetadata;
use crate::responses_retry::ResponsesRequestTransport;
use codex_http_client::RetryAfter;
use codex_otel::SessionTelemetry;
use codex_otel::auth_storage::AuthStorageOriginator;
use codex_protocol::error::CodexErr;
use codex_protocol::error::CodexErrorDetails;
use codex_protocol::openai_models::ModelInfo;
use futures::FutureExt;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::MutexGuard;
use std::sync::OnceLock;
use std::sync::Weak;
use std::time::Duration;
use tokio::sync::oneshot;
use tokio::time::Instant;
use tokio_util::task::AbortOnDropHandle;
const RETRY_INTERVAL: Duration = Duration::from_secs(30);
pub(super) const RECOVERY_ACK_TIMEOUT: Duration = Duration::from_secs(5);
/// Foreground clients own probe cancellation. A handshake can only publish to a live owner.
#[derive(Debug, Clone)]
pub(super) enum WebsocketRoutes {
Session(Arc<RouteOwner>),
Probe,
}
#[derive(Debug, Default)]
pub(super) struct RouteOwner(Mutex<WebsocketTransport>);
impl RouteOwner {
fn lock(&self) -> MutexGuard<'_, WebsocketTransport> {
self.0
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
}
impl Default for WebsocketRoutes {
fn default() -> Self {
Self::Session(Arc::default())
}
}
impl WebsocketRoutes {
fn owner(&self) -> Option<&Arc<RouteOwner>> {
match self {
Self::Session(owner) => Some(owner),
Self::Probe => None,
}
}
pub(super) fn enabled(&self) -> bool {
self.owner()
.is_some_and(|owner| matches!(owner.lock().route, Route::Websocket))
}
pub(super) fn can_stream(&self) -> bool {
self.owner()
.is_some_and(|owner| matches!(owner.lock().route, Route::Websocket | Route::Ready(_)))
}
pub(super) fn fall_back(&self, policy: HttpFallbackPolicy) -> bool {
let Some(owner) = self.owner() else {
return false;
};
let mut transport = owner.lock();
let activated = matches!(transport.route, Route::Websocket);
match policy {
HttpFallbackPolicy::StayOnHttp => transport.route = Route::HttpOnly,
HttpFallbackPolicy::AllowProbes if activated => transport.route = Route::HttpFallback,
HttpFallbackPolicy::AllowProbes => {}
}
activated
}
}
#[derive(Debug, Default)]
struct WebsocketTransport {
route: Route,
// Even a proxy that accepts upgrades can consistently reject request frames.
next_probe: Option<Instant>,
}
#[derive(Debug, Default)]
enum Route {
#[default]
Websocket,
HttpOnly,
HttpFallback,
Connecting(AbortOnDropHandle<()>),
Ready(Box<WebsocketSession>),
Verifying,
}
pub(super) enum Connection {
Websocket,
Recovered(RecoveryAttempt),
}
/// Owns the provisional route until its first real request completes or is dropped.
pub(super) struct RecoveryAttempt {
owner: Weak<RouteOwner>,
retry_over_http: Option<oneshot::Sender<()>>,
}
impl RecoveryAttempt {
pub(super) fn observe(&mut self, event: &ResponseEvent) {
if matches!(
event,
ResponseEvent::GuardianAssessment(_)
| ResponseEvent::OutputItemAdded(_)
| ResponseEvent::OutputItemDone(_)
| ResponseEvent::OutputTextDelta(_)
| ResponseEvent::ToolCallInputDelta { .. }
| ResponseEvent::ReasoningSummaryDelta { .. }
| ResponseEvent::ReasoningSummaryDone { .. }
| ResponseEvent::ReasoningContentDelta { .. }
| ResponseEvent::ReasoningSummaryPartAdded { .. }
) {
self.retry_over_http = None;
}
}
pub(super) fn failed(mut self, error: &CodexErr) {
let route = match error.details() {
CodexErrorDetails::RequestTimeout => Route::HttpOnly,
_ => Route::HttpFallback,
};
self.finish(route, error.retry_after());
if let Some(retry) = self.retry_over_http.take() {
let _ = retry.send(());
}
}
pub(super) fn completed(mut self, telemetry: &SessionTelemetry) {
if self.finish(Route::Websocket, /*retry_after*/ None) {
telemetry.counter("codex.transport.recovered_websocket", /*inc*/ 1, &[]);
tracing::info!("resumed WebSocket after HTTP fallback");
}
}
fn finish(&mut self, route: Route, retry_after: Option<RetryAfter>) -> bool {
// Consume the owner reference so Drop only acts on an abandoned attempt.
let Some(owner) = std::mem::take(&mut self.owner).upgrade() else {
return false;
};
let mut transport = owner.lock();
if !matches!(transport.route, Route::Verifying) {
return false;
}
if matches!(route, Route::HttpFallback) {
transport.next_probe =
Some(Instant::now() + RETRY_INTERVAL).max(retry_after.map(RetryAfter::deadline));
}
transport.route = route;
true
}
}
impl Drop for RecoveryAttempt {
fn drop(&mut self) {
self.finish(Route::HttpFallback, /*retry_after*/ None);
}
}
/// Select the transport before the caller builds a request for it.
pub(super) fn select_for_request(
session: &mut ModelClientSession,
mode: &mut WebsocketRequestMode<'_>,
params: &WebsocketConnectParams<'_>,
) -> Option<Connection> {
let owner = session.client.websocket_routes.owner()?;
let mut transport = owner.lock();
match (&mut transport.route, mode) {
(Route::Websocket, _) => Some(Connection::Websocket),
(Route::Ready(socket), WebsocketRequestMode::Generate(request_transport)) => {
let key = ResponsesConnectionKey::new(&params.api_provider, params.auth_revision);
// A background socket has no outstanding requests. If it is closed or busy,
// keep using HTTP rather than opening or waiting for another socket here.
if socket.auth_owner_generation != params.auth_owner_generation
|| socket.auth_owner_generation != session.client.auth_owner_generation()
|| !socket.matches_setup(params, &key)
|| socket
.connection
.as_ref()
.is_none_or(|connection| connection.is_closed().now_or_never() != Some(false))
{
transport.route = Route::HttpFallback;
return None;
}
if session.websocket_session.auth_owner_generation != socket.auth_owner_generation {
session.turn_state = Arc::new(OnceLock::new());
}
session.websocket_session = std::mem::take(socket.as_mut());
session
.websocket_session
.set_connection_reused(/*connection_reused*/ true);
let (retry, failed_before_output) = oneshot::channel();
**request_transport = ResponsesRequestTransport::Verifying(failed_before_output);
transport.route = Route::Verifying;
Some(Connection::Recovered(RecoveryAttempt {
owner: Arc::downgrade(owner),
retry_over_http: Some(retry),
}))
}
_ => None,
}
}
/// Start one background handshake alongside HTTP, without sending a prompt or a turn-state token.
pub(super) fn reconnect_if_due(
client: &ModelClient,
model_info: &ModelInfo,
service_tier: Option<&str>,
telemetry: &SessionTelemetry,
metadata: &CodexResponsesMetadata,
) {
let Some(owner) = client.websocket_routes.owner() else {
return;
};
let mut transport = owner.lock();
let now = Instant::now();
match &transport.route {
Route::HttpFallback => {}
Route::Connecting(task) if task.is_finished() => {}
_ => return,
}
if transport.next_probe.is_some_and(|at| now < at) {
return;
}
transport.next_probe = Some(now + RETRY_INTERVAL);
let route_owner = Arc::downgrade(owner);
let client = ModelClient {
websocket_routes: WebsocketRoutes::Probe,
..client.clone()
};
let model_info = model_info.clone();
let service_tier = service_tier.map(str::to_owned);
let telemetry = telemetry.clone();
let metadata = metadata.clone();
let timeout = client.state.provider.info().websocket_connect_timeout();
let task = tokio::spawn(AuthStorageOriginator::current().scope(async move {
let mut socket = WebsocketSession::default();
let handshake = socket.preconnect(
&client,
WebsocketPreconnectMode::ProbeWithoutAuthRecovery,
&model_info,
service_tier,
&telemetry,
&metadata,
);
let result = tokio::time::timeout(timeout, handshake).await;
let Some(owner) = route_owner.upgrade() else {
return;
};
let mut transport = owner.lock();
// A foreground 426 takes precedence over any in-flight probe.
if !matches!(transport.route, Route::Connecting(_)) {
return;
}
transport.route = match result {
Ok(Ok(WebsocketHandshake::Connected)) => Route::Ready(Box::new(socket)),
Ok(Ok(WebsocketHandshake::UpgradeRequired)) => Route::HttpOnly,
Ok(Err(error)) => {
transport.next_probe = transport
.next_probe
.max(error.retry_after().map(|advice| advice.deadline()));
Route::HttpFallback
}
Err(_) => Route::HttpFallback,
};
tracing::debug!(
name: "codex.websocket_probe_finished",
ready = matches!(transport.route, Route::Ready(_)),
"background WebSocket handshake finished"
);
}));
transport.route = Route::Connecting(AbortOnDropHandle::new(task));
}
+9 -1
View File
@@ -19,6 +19,7 @@ use crate::hook_runtime::run_post_compact_hooks;
use crate::hook_runtime::run_pre_compact_hooks;
use crate::responses_metadata::CodexResponsesMetadata;
use crate::responses_metadata::CompactionTurnMetadata;
use crate::responses_retry::ResponsesRequestTransport;
use crate::session::RequestEffortUsage;
use crate::session::session::Session;
use crate::session::step_context::StepContext;
@@ -292,6 +293,7 @@ async fn run_compact_task_inner_impl(
let max_retries = turn_context.provider.info().stream_max_retries();
let mut retries = 0;
let mut transport = ResponsesRequestTransport::Auto;
// Reuse one client session so turn-scoped state (sticky routing and websocket incremental
// request tracking) survives retries within this compact turn.
let mut client_session = sess.services.model_client.new_session();
@@ -332,6 +334,7 @@ async fn run_compact_task_inner_impl(
&sess,
turn_context.as_ref(),
&mut client_session,
&mut transport,
&responses_metadata,
&prompt,
compaction_metadata.phase(),
@@ -373,6 +376,9 @@ async fn run_compact_task_inner_impl(
return Err(e);
}
Err(e) => {
if transport.retry_failed_recovery(&e).await {
continue;
}
if retries < max_retries {
retries += 1;
let delay = backoff(retries);
@@ -826,13 +832,14 @@ async fn drain_to_completed(
sess: &Session,
turn_context: &TurnContext,
client_session: &mut ModelClientSession,
transport: &mut ResponsesRequestTransport,
responses_metadata: &CodexResponsesMetadata,
prompt: &Prompt,
phase: CompactionPhase,
) -> CodexResult<CompactionResponse> {
let mut includes_reasoning = turn_context.model_info().use_responses_lite;
let mut stream = client_session
.stream(
.stream_attempt(
prompt,
turn_context.model_info(),
&turn_context.session_telemetry,
@@ -847,6 +854,7 @@ async fn drain_to_completed(
// Rollout tracing currently models remote compaction only; local compaction streams
// are left untraced until the reducer has a first-class local compaction lifecycle.
&InferenceTraceContext::disabled(),
transport,
)
.await?;
let mut output = Vec::new();
+2 -1
View File
@@ -330,7 +330,7 @@ async fn run_remote_compaction_request_v2(
let mut retry_state = ResponsesStreamRetryState::default();
loop {
let result = match client_session
.stream(
.stream_attempt(
prompt,
&settings.model_info,
&turn_context.session_telemetry,
@@ -340,6 +340,7 @@ async fn run_remote_compaction_request_v2(
step_context.settings.service_tier.clone(),
responses_metadata,
&InferenceTraceContext::disabled(),
&mut retry_state.transport,
)
.await
{
+33 -1
View File
@@ -4,6 +4,7 @@
use std::time::Duration;
use crate::client::HttpFallbackPolicy;
use crate::client::ModelClientSession;
use crate::context::ContentFilterGuidance;
use crate::context::ContextualUserFragment;
@@ -29,8 +30,33 @@ pub(crate) enum ResponsesStreamRequest {
RemoteCompactionV2,
}
/// Once a logical request uses HTTP, its retries must stay on HTTP too.
pub(crate) enum ResponsesRequestTransport {
Auto,
Verifying(tokio::sync::oneshot::Receiver<()>),
HttpOnly,
}
impl ResponsesRequestTransport {
/// A recovered socket gets one request to prove that it can carry frames.
pub(crate) async fn retry_failed_recovery(&mut self, err: &CodexErr) -> bool {
if let Self::Verifying(failed_before_output) = self
&& failed_before_output.try_recv().is_ok()
&& err.retry_delay(1).is_some()
{
*self = Self::HttpOnly;
if let Some(advice) = err.retry_after() {
tokio::time::sleep_until(advice.deadline()).await;
}
return true;
}
false
}
}
pub(crate) struct ResponsesStreamRetryState {
retries: u64,
pub(crate) transport: ResponsesRequestTransport,
connection_retries: u64,
connection_retry_delay: Duration,
}
@@ -39,6 +65,7 @@ impl Default for ResponsesStreamRetryState {
fn default() -> Self {
Self {
retries: 0,
transport: ResponsesRequestTransport::Auto,
connection_retries: 0,
connection_retry_delay: INITIAL_CONNECTION_RETRY_DELAY,
}
@@ -89,6 +116,9 @@ pub(crate) async fn handle_response_stream_error(
let Some(delay) = err.retry_delay(retry_count) else {
return Err(err);
};
if retry_state.transport.retry_failed_recovery(&err).await {
return Ok(());
}
let retry_after = err.retry_after();
if turn_context
@@ -119,9 +149,10 @@ pub(crate) async fn handle_response_stream_error(
}
if retry_state.retries >= max_retries
&& matches!(retry_state.transport, ResponsesRequestTransport::Auto)
&& client_session.try_switch_fallback_transport(
&turn_context.session_telemetry,
turn_context.model_info(),
HttpFallbackPolicy::AllowProbes,
)
{
// Changing transport must not bypass the server's retry deadline.
@@ -135,6 +166,7 @@ pub(crate) async fn handle_response_stream_error(
}),
)
.await;
retry_state.transport = ResponsesRequestTransport::HttpOnly;
retry_state.retries = 0;
return Ok(());
}
+5 -1
View File
@@ -36,6 +36,7 @@ use crate::mentions::collect_tool_mentions_from_messages;
use crate::plugins::build_plugin_injections;
use crate::responses_metadata::CodexResponsesMetadata;
use crate::responses_metadata::CodexResponsesRequestKind;
use crate::responses_retry::ResponsesRequestTransport;
use crate::responses_retry::ResponsesStreamRequest;
use crate::responses_retry::ResponsesStreamRetryState;
use crate::responses_retry::handle_response_stream_error;
@@ -1880,6 +1881,7 @@ async fn run_sampling_request(
Arc::clone(&step_context),
Arc::clone(&turn_store),
client_session,
&mut retry_state.transport,
&responses_metadata,
Arc::clone(&turn_diff_tracker),
prompt,
@@ -2733,6 +2735,7 @@ async fn try_run_sampling_request(
step_context: Arc<StepContext>,
turn_store: Arc<codex_extension_api::ExtensionData>,
client_session: &mut ModelClientSession,
transport: &mut ResponsesRequestTransport,
responses_metadata: &CodexResponsesMetadata,
turn_diff_tracker: SharedTurnDiffTracker,
prompt: Prompt,
@@ -2831,7 +2834,7 @@ async fn try_run_sampling_request(
None
};
let stream = client_session
.stream(
.stream_attempt(
&prompt,
&step_context.settings.model_info,
&step_context.session_telemetry,
@@ -2840,6 +2843,7 @@ async fn try_run_sampling_request(
step_context.settings.service_tier.clone(),
responses_metadata,
&inference_trace,
transport,
)
.instrument(trace_span!("stream_request"))
.or_cancel(&cancellation_token)