mirror of
https://github.com/openai/codex.git
synced 2026-11-07 20:30:33 +00:00
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:
+296
-179
@@ -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(¶ms.api_provider, params.auth_revision);
|
||||
let matches_setup = self.matches_setup(¶ms, &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, ¶ms)
|
||||
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>,
|
||||
|
||||
@@ -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(¶ms.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));
|
||||
}
|
||||
@@ -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();
|
||||
|
||||
@@ -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
|
||||
{
|
||||
|
||||
@@ -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(());
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user