refactor(openai): unify WebSocket recovery and session fallback

Consolidate retry decisions in the provider and share the failure budget with session HTTP fallback. Preserve WebSocket priority and the 800ms cold connection budget.

Remove transport-local retries, fresh recovery probes, and metadata-based replay exceptions; use a 300s stream idle timeout.

Validation: backend tests (1816 passed, 1 ignored), fmt, clippy, and real Codex CLI 10-turn WebSocket continuation run.
This commit is contained in:
Zyy
2026-09-07 00:45:08 +08:00
parent 074c07b5ca
commit aa6bf6497f
29 changed files with 676 additions and 1695 deletions
@@ -223,8 +223,6 @@ struct CurrentAttempt {
index: NonZeroU32,
started_at: SystemTime,
send_observed: bool,
/// Provider 已为当前 attempt 交付可供同请求恢复的不透明状态检查点。
session_checkpoint_observed: bool,
response_observation: Option<ProviderResponseObservation>,
}
@@ -596,7 +594,7 @@ where
.wire_event(wire.protocol(), wire.event_type());
}
if let Some(state) = event.session_update() {
self.observe_session_checkpoint(state);
self.observe_session_update(state);
}
if let Some(observation) = event.take_observation() {
self.observe_response(observation);
@@ -978,7 +976,6 @@ where
index: next_attempt,
started_at: SystemTime::now(),
send_observed: false,
session_checkpoint_observed: false,
response_observation: None,
});
Ok(None)
@@ -1004,14 +1001,13 @@ where
self.mark_send_observed().await;
}
fn observe_session_checkpoint(&mut self, state: &ProviderSessionState) {
fn observe_session_update(&mut self, state: &ProviderSessionState) {
let Some(current) = self.current.as_mut() else {
return;
};
if state.provider() != current.metadata.provider().as_str() {
return;
}
current.session_checkpoint_observed = true;
self.operation.set_provider_session_state(state.clone());
}
@@ -1081,11 +1077,6 @@ where
.await;
}
let provider_proved_replay_safe = provider_proved_replay_safe(&error);
// Ambiguous 仍是默认 fail-closed。唯一例外是 Provider 在同一 attempt 中先
// 交付了不透明恢复检查点,并明确证明从该检查点重放安全;该例外只用于
// 固定原账号的 transport recovery,不放开普通换号或 continuation 重放。
let checkpointed_transport_replay =
current.session_checkpoint_observed && error.replay_is_safe();
let continuation_retry = self.prepare_continuation_retry(
&current,
&error,
@@ -1108,16 +1099,14 @@ where
delay,
}) if self.downstream_committed_at.is_none()
&& !self.delivery_pending
&& (attempt_send_state != UpstreamSendState::Ambiguous
|| checkpointed_transport_replay) =>
&& attempt_send_state != UpstreamSendState::Ambiguous =>
{
Some((AttemptTransport::Retry(retry_index), delay))
}
Some(crate::error::PreDeliveryRetry::SameAccountTransportFallback)
if self.downstream_committed_at.is_none()
&& !self.delivery_pending
&& (attempt_send_state != UpstreamSendState::Ambiguous
|| checkpointed_transport_replay) =>
&& attempt_send_state != UpstreamSendState::Ambiguous =>
{
Some((AttemptTransport::Fallback, Duration::ZERO))
}
@@ -2909,7 +2909,7 @@ fn ambiguous_pre_delivery_retry_marker_does_not_rotate_account() {
}
#[test]
fn ambiguous_transport_retry_uses_provider_session_checkpoint_on_same_account() {
fn ambiguous_transport_retry_is_rejected_even_with_provider_session_state() {
let operation = generate_operation();
let route_plan = plan(&operation);
let checkpoint = ProviderSessionState::new(
@@ -2920,7 +2920,7 @@ fn ambiguous_transport_retry_uses_provider_session_checkpoint_on_same_account()
let mut checkpoint_event = ProviderEvent::observation(ProviderResponseObservation::new(
UpstreamTransport::new("websocket").expect("transport"),
));
checkpoint_event.attach_session_update(checkpoint.clone());
checkpoint_event.attach_session_update(checkpoint);
let retry_index = NonZeroU32::new(1).expect("retry index");
let (coordinator, store, provider) = coordinator(vec![
Script::ObservedStream {
@@ -2949,27 +2949,14 @@ fn ambiguous_transport_retry_uses_provider_session_checkpoint_on_same_account()
CancellationToken::new(),
))
.expect("start execution");
block_on(session.collect_uncommitted()).expect("checkpointed retry succeeds");
block_on(session.commit_downstream(Some(200))).expect("commit winning response");
let original = ProviderAccountId::new("acct_first").expect("account id");
let contexts = provider.contexts.lock().expect("contexts lock");
assert_eq!(contexts.len(), 2);
assert_eq!(
contexts[1].transport(),
AttemptTransport::Retry(retry_index)
);
assert_eq!(contexts[1].required_account(), Some(&original));
drop(contexts);
let operations = provider.operations.lock().expect("operations lock");
assert_eq!(
operations[1].provider_session_state("openai"),
Some(&checkpoint)
);
let error =
block_on(session.collect_uncommitted()).expect_err("ambiguous send cannot be replayed");
assert!(matches!(error, EngineError::Provider(_)));
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 1);
let state = store.state.lock().expect("store lock");
assert_eq!(state.intermediate_failures, 1);
assert_eq!(state.finalizations[0].attempt_count, 2);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
assert_eq!(state.intermediate_failures, 0);
assert_eq!(state.finalizations[0].attempt_count, 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Failed);
}
#[test]
@@ -19,7 +19,7 @@ use crate::{
};
/// 本服务的流式请求默认重试次数。
pub const DEFAULT_STREAM_MAX_RETRIES: u64 = 3;
pub const DEFAULT_STREAM_MAX_RETRIES: u64 = 5;
/// 防止错误配置产生无界隐藏重放;与官方 Codex 的硬上限一致。
pub const MAX_STREAM_MAX_RETRIES: u64 = 100;
@@ -175,7 +175,7 @@ pub struct CodexWebSocketPoolSettings {
pub enabled: bool,
pub max_age_ms: u64,
pub max_connecting: usize,
pub initial_event_timeout_ms: u64,
pub stream_idle_timeout_ms: u64,
}
impl Default for CodexWebSocketPoolSettings {
@@ -184,7 +184,7 @@ impl Default for CodexWebSocketPoolSettings {
enabled: true,
max_age_ms: 55 * 60 * 1000,
max_connecting: 8,
initial_event_timeout_ms: 20_000,
stream_idle_timeout_ms: 300_000,
}
}
}
@@ -202,8 +202,8 @@ impl CodexWebSocketPoolSettings {
enabled: self.enabled,
max_age: Duration::from_millis(self.max_age_ms),
max_connecting: self.max_connecting,
initial_event_timeout: (self.initial_event_timeout_ms != 0)
.then(|| Duration::from_millis(self.initial_event_timeout_ms)),
stream_idle_timeout: (self.stream_idle_timeout_ms != 0)
.then(|| Duration::from_millis(self.stream_idle_timeout_ms)),
..CodexWebSocketPoolConfig::default()
}
}
@@ -147,7 +147,6 @@ pub(super) struct ColdResponse {
pub(super) request: CodexResponsesRequest,
pub(super) upstream_model: UpstreamModelId,
pub(super) transport_policy: CodexProviderTransport,
pub(super) websocket_connection_preference: WebSocketConnectionPreference,
pub(super) context: AttemptContext,
pub(super) selector: Arc<CodexCredentialSelector>,
pub(super) quota: Arc<CodexCredentialQuotaService>,
@@ -157,7 +156,6 @@ pub(super) struct ColdResponse {
pub(super) session_affinity_key: Option<ProviderSessionAffinityKey>,
pub(super) session_affinity_key_hash: Option<String>,
pub(super) session_transport_recovery: CodexSessionTransportRecovery,
pub(super) session_recovery_probe: Option<CodexSessionWebSocketProbe>,
pub(super) websocket_retry_count: u32,
pub(super) stream_max_retries: u32,
pub(super) session_capture: Option<OpenAiSessionCapture>,
@@ -186,13 +184,6 @@ pub(super) struct OpenAiSessionState {
#[serde(default)]
pub(super) client_turn_id: Option<String>,
pub(super) continuation_scope: OpenAiContinuationScope,
/// 该状态来自同一业务请求的 metadata 检查点,只允许被紧随其后的恢复 attempt 消费。
#[serde(default, skip_serializing_if = "is_false")]
pub(super) retry_checkpoint: bool,
}
const fn is_false(value: &bool) -> bool {
!*value
}
#[derive(Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
@@ -241,7 +232,6 @@ pub(super) fn encode_openai_session_state(
fn encode_openai_session_capture(
capture: &OpenAiSessionCapture,
retry_checkpoint: bool,
) -> Result<ProviderSessionState, ProviderError> {
let Some(continuation_scope) = capture.continuation_scope else {
return Err(provider_error(
@@ -255,7 +245,6 @@ fn encode_openai_session_capture(
turn_state: capture.turn_state.clone(),
client_turn_id: capture.client_turn_id.clone(),
continuation_scope,
retry_checkpoint,
})
}
@@ -272,7 +261,7 @@ pub(super) fn attach_openai_session_update(
let Some(capture) = capture.take() else {
return;
};
let Ok(update) = encode_openai_session_capture(&capture, false) else {
let Ok(update) = encode_openai_session_capture(&capture) else {
return;
};
events[terminal_index].attach_session_update(update);
@@ -312,7 +301,6 @@ pub(super) async fn create_response_attempt(
account_id: &str,
deadline: SystemTime,
cancellation: &CancellationToken,
connection_preference: WebSocketConnectionPreference,
) -> Result<CodexBackendStreamingResponse, CodexHandshakeAttemptError> {
let Some(handshake_deadline) = remaining(deadline) else {
return Err(CodexHandshakeAttemptError::Timeout);
@@ -321,11 +309,10 @@ pub(super) async fn create_response_attempt(
biased;
_ = cancellation.cancelled() => Err(CodexHandshakeAttemptError::Cancelled),
_ = tokio::time::sleep(handshake_deadline) => Err(CodexHandshakeAttemptError::Timeout),
response = client.create_response_stream_with_deferred_websocket_recovery(
response = client.create_response_stream_with_pool_account(
request,
request_context,
Some(account_id),
connection_preference,
) => response.map_err(CodexHandshakeAttemptError::Client),
}
}
@@ -496,7 +483,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
request,
upstream_model,
transport_policy,
websocket_connection_preference,
context,
selector,
quota,
@@ -506,7 +492,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
session_affinity_key,
session_affinity_key_hash,
session_transport_recovery,
mut session_recovery_probe,
websocket_retry_count,
stream_max_retries,
mut session_capture,
@@ -557,7 +542,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
active_account.id().as_str(),
context.deadline(),
&cancellation,
websocket_connection_preference,
)
.await;
let websocket_failure_policy = match &response {
@@ -584,7 +568,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
WebSocketRecoveryContext {
policy,
requirement: request_transport_requirement,
replay_checkpoint_available: false,
retry_count: websocket_retry_count,
max_retries: stream_max_retries,
request_id: context.request_id().as_str(),
@@ -593,7 +576,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
session_affinity_key: session_affinity_key.as_ref(),
session_affinity_key_hash: session_affinity_key_hash.as_deref(),
session_transport_recovery: &session_transport_recovery,
session_recovery_probe: &mut session_recovery_probe,
},
);
}
@@ -616,9 +598,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
Err(failure.error)?;
return;
}
let mut turn_state_checkpoint_observed = response.transport
== CodexBackendTransport::WebSocket
&& response.turn_state.is_some();
if let Some(capture) = session_capture.as_mut() {
capture.continuation_scope = Some(if capture.response_store {
OpenAiContinuationScope::Persisted
@@ -744,35 +723,11 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
&mut observation_state,
)
.await;
turn_state_checkpoint_observed |= turn_state_merge.is_some();
let retry_checkpoint_update = if turn_state_checkpoint_observed
&& failure.websocket_checkpoint_retryable
{
session_capture
.as_ref()
.and_then(|capture| encode_openai_session_capture(capture, true).ok())
let observation_event = if rate_limits_changed || turn_state_merge.is_some() {
observation_state.observation().map(ProviderEvent::observation)
} else {
None
};
let mut observation_event = if rate_limits_changed
|| turn_state_merge.is_some()
|| retry_checkpoint_update.is_some()
{
observation_state
.observation()
.map(ProviderEvent::observation)
} else {
None
};
if let (Some(event), Some(update)) =
(observation_event.as_mut(), retry_checkpoint_update)
{
event.attach_session_update(update);
}
let replay_checkpoint_available = observation_event
.as_ref()
.and_then(ProviderEvent::session_update)
.is_some();
if failure.websocket_transport_retryable
&& response_transport == CodexBackendTransport::WebSocket
&& (matches!(
@@ -785,7 +740,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
WebSocketRecoveryContext {
policy: WebSocketFailurePolicy::Budgeted,
requirement: request_transport_requirement,
replay_checkpoint_available,
retry_count: websocket_retry_count,
max_retries: stream_max_retries,
request_id: context.request_id().as_str(),
@@ -794,7 +748,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
session_affinity_key: session_affinity_key.as_ref(),
session_affinity_key_hash: session_affinity_key_hash.as_deref(),
session_transport_recovery: &session_transport_recovery,
session_recovery_probe: &mut session_recovery_probe,
},
);
}
@@ -829,7 +782,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
&mut observation_state,
)
.await;
turn_state_checkpoint_observed |= turn_state_merge.is_some();
let turn_state_changed = turn_state_merge.unwrap_or(false);
let first_event_changed =
observation_state.observe_stream_chunk(&chunk, output_started_at);
@@ -879,19 +831,10 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
let terminal_changed = completed
&& observation_state.mark_completed(terminal_response_is_incomplete(&events));
if response_transport == CodexBackendTransport::WebSocket
&& completed
&& terminal_failure.is_none()
&& let Some(probe) = session_recovery_probe.take()
&& completed && terminal_failure.is_none()
&& let Some(key) = session_affinity_key.as_ref()
{
let session_recovery_cleared = probe.succeed();
tracing::info!(
request_id = %context.request_id(),
attempt_index = context.attempt_index().get(),
session_affinity_key_hash = session_affinity_key_hash.as_deref().unwrap_or(""),
session_transport_action = "websocket_recovered",
session_recovery_cleared,
"OpenAI session fresh WebSocket probe reached a terminal response"
);
session_transport_recovery.websocket_succeeded(key);
}
if allows_account_state_mutation && (completed || terminal_failure.is_some()) {
synchronize_passive_quota(
@@ -1022,19 +965,10 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
let terminal_changed = completed
&& observation_state.mark_completed(terminal_response_is_incomplete(&events));
if response_transport == CodexBackendTransport::WebSocket
&& completed
&& terminal_failure.is_none()
&& let Some(probe) = session_recovery_probe.take()
&& completed && terminal_failure.is_none()
&& let Some(key) = session_affinity_key.as_ref()
{
let session_recovery_cleared = probe.succeed();
tracing::info!(
request_id = %context.request_id(),
attempt_index = context.attempt_index().get(),
session_affinity_key_hash = session_affinity_key_hash.as_deref().unwrap_or(""),
session_transport_action = "websocket_recovered",
session_recovery_cleared,
"OpenAI session fresh WebSocket probe reached a terminal response"
);
session_transport_recovery.websocket_succeeded(key);
}
if allows_account_state_mutation && completed && terminal_failure.is_none() {
// 同上:尾部 finish() 也可能产出 completed,亲和记录必须先于任何下游 yield。
@@ -39,8 +39,6 @@ pub fn openai_failure_affects_account_score(error: &ProviderError) -> bool {
pub(super) struct MappedProviderFailure {
pub(super) error: ProviderError,
pub(super) websocket_transport_retryable: bool,
/// 兼容少数链路在 metadata 后、终态前关闭;仅在捕获 turn-state 时允许安全恢复。
pub(super) websocket_checkpoint_retryable: bool,
pub(super) account_failure: Option<CodexAccountFailure>,
/// 原始上游错误描述,仅在凭据错误状态下持久化。
pub(super) error_message: Option<String>,
@@ -56,7 +54,6 @@ impl MappedProviderFailure {
Self {
error,
websocket_transport_retryable: false,
websocket_checkpoint_retryable: false,
account_failure: None,
error_message: None,
cyber_policy_failure: false,
@@ -454,8 +451,6 @@ pub(super) enum WebSocketFailurePolicy {
pub(super) struct WebSocketRecoveryContext<'a> {
pub(super) policy: WebSocketFailurePolicy,
pub(super) requirement: TransportRequirement,
/// 本 attempt 已从 metadata 捕获可恢复同一请求的 turn-state 检查点。
pub(super) replay_checkpoint_available: bool,
pub(super) retry_count: u32,
pub(super) max_retries: u32,
pub(super) request_id: &'a str,
@@ -464,7 +459,6 @@ pub(super) struct WebSocketRecoveryContext<'a> {
pub(super) session_affinity_key: Option<&'a ProviderSessionAffinityKey>,
pub(super) session_affinity_key_hash: Option<&'a str>,
pub(super) session_transport_recovery: &'a CodexSessionTransportRecovery,
pub(super) session_recovery_probe: &'a mut Option<CodexSessionWebSocketProbe>,
}
pub(super) fn websocket_client_failure_policy(
@@ -497,6 +491,9 @@ pub(super) fn websocket_client_failure_policy(
if !matches!(
error.classified(),
CodexWebSocketExchangeError::InvalidRequest(_)
| CodexWebSocketExchangeError::ContinuationUnavailable { .. }
| CodexWebSocketExchangeError::OriginCircuitOpen
| CodexWebSocketExchangeError::OriginHalfOpenBusy
) =>
{
Some(WebSocketFailurePolicy::Budgeted)
@@ -509,25 +506,18 @@ pub(super) fn apply_websocket_recovery_policy(
failure: &mut MappedProviderFailure,
context: WebSocketRecoveryContext<'_>,
) {
let session_budget_exhausted = context.session_affinity_key.is_some_and(|key| {
context
.session_transport_recovery
.record_websocket_failure(key, context.max_retries)
});
let send_state = failure.error.send_state();
let post_send = matches!(
send_state,
UpstreamSendState::Sent | UpstreamSendState::Ambiguous
);
if post_send && !context.replay_checkpoint_available {
// payload 已经发送或发送结果不确定时,代理不能用同一业务输入自动重放。
// 当前请求保持终态失败;恢复状态只裁决客户端主动发起的下一次请求。
let session_recovery_transition = match context.session_recovery_probe.take() {
Some(probe) => Some(probe.post_send_failed()),
None => context.session_affinity_key.map(|key| {
context
.session_transport_recovery
.record_post_send_failure(key)
}),
};
let session_http_cooldown_ms = session_recovery_transition
.and_then(CodexSessionRecoveryTransition::cooldown)
.map(|duration| u64::try_from(duration.as_millis()).unwrap_or(u64::MAX));
if send_state == UpstreamSendState::Ambiguous || (post_send && !failure.error.replay_is_safe())
{
tracing::warn!(
request_id = context.request_id,
attempt_index = context.attempt_index,
@@ -540,34 +530,27 @@ pub(super) fn apply_websocket_recovery_policy(
upstream_send_state = ?send_state,
transport_requirement = context.requirement.as_str(),
continuation_recovery_action = "client_replay_required",
session_http_fallback = session_budget_exhausted,
session_affinity_present = context.session_affinity_key.is_some(),
session_affinity_key_hash = context.session_affinity_key_hash.unwrap_or(""),
session_transport_action = session_recovery_transition
.map_or("untracked", CodexSessionRecoveryTransition::action),
session_websocket_failure_count = session_recovery_transition
.map_or(0, CodexSessionRecoveryTransition::failure_count),
session_http_cooldown_ms = session_http_cooldown_ms.unwrap_or_default(),
session_http_cooldown_present = session_http_cooldown_ms.is_some(),
"OpenAI upstream WebSocket failed after payload send; proxy replay was suppressed"
);
return;
}
let recovery_is_safe = context.requirement.allows_pre_send_http_fallback();
if !recovery_is_safe {
if matches!(
context.requirement,
TransportRequirement::ExactWebSocketContinuation | TransportRequirement::HttpRequired
) {
return;
}
// 传输恢复必须保持原账号可调度;最终 HTTP attempt 若仍失败,再按真实 HTTP
// 结果更新账号健康度,避免中间 WS 错误把同账号钉选提前冷却掉。
failure.account_failure = None;
if post_send {
// Core 对 Ambiguous 默认仍然 fail-closed;只有同一 attempt 的 session
// checkpoint 与这个 Provider proof 同时存在时,才允许同账号传输恢复。
failure.error.set_replay_safe();
}
let fallback_now = context.policy == WebSocketFailurePolicy::ImmediateFallback
|| context.retry_count >= context.max_retries;
|| context.retry_count >= context.max_retries
|| session_budget_exhausted;
if !fallback_now {
let Some(retry_index) = context.retry_count.checked_add(1).and_then(NonZeroU32::new) else {
return;
@@ -594,28 +577,18 @@ pub(super) fn apply_websocket_recovery_policy(
websocket_max_retries = context.max_retries,
websocket_retry_delay_ms = u64::try_from(delay.as_millis()).unwrap_or(u64::MAX),
transport_requirement = context.requirement.as_str(),
replay_checkpoint_available = context.replay_checkpoint_available,
"OpenAI upstream WebSocket failed; retrying the same account"
);
return;
}
if !context.requirement.allows_pre_send_http_fallback() {
return;
}
failure.error.set_pre_delivery_transport_fallback();
let session_fallback = match context.policy {
WebSocketFailurePolicy::Budgeted => SessionWebSocketFallback::RetryBudgetExhausted,
WebSocketFailurePolicy::ImmediateFallback => SessionWebSocketFallback::UpgradeRequired,
};
let session_recovery_transition = match context.session_recovery_probe.take() {
Some(probe) => Some(probe.fallback(session_fallback)),
None => context.session_affinity_key.map(|key| {
context
.session_transport_recovery
.record_websocket_fallback(key, session_fallback)
}),
};
let session_http_cooldown_ms = session_recovery_transition
.and_then(CodexSessionRecoveryTransition::cooldown)
.map(|duration| u64::try_from(duration.as_millis()).unwrap_or(u64::MAX));
if let Some(key) = context.session_affinity_key {
context.session_transport_recovery.disable_websocket(key);
}
tracing::warn!(
request_id = context.request_id,
attempt_index = context.attempt_index,
@@ -633,16 +606,9 @@ pub(super) fn apply_websocket_recovery_policy(
WebSocketFailurePolicy::Budgeted => "retry_budget_exhausted",
WebSocketFailurePolicy::ImmediateFallback => "upgrade_required",
},
replay_checkpoint_available = context.replay_checkpoint_available,
session_affinity_present = context.session_affinity_key.is_some(),
session_affinity_key_hash = context.session_affinity_key_hash.unwrap_or(""),
session_transport_action =
session_recovery_transition.map_or("untracked", CodexSessionRecoveryTransition::action),
session_websocket_failure_count =
session_recovery_transition.map_or(0, CodexSessionRecoveryTransition::failure_count),
session_http_cooldown_ms = session_http_cooldown_ms.unwrap_or_default(),
session_http_cooldown_present = session_http_cooldown_ms.is_some(),
"OpenAI upstream WebSocket entered session HTTP cooldown"
"OpenAI upstream WebSocket disabled for this session"
);
}
@@ -686,31 +652,14 @@ pub(super) fn continuation_replay_required_error(reason: &'static str) -> Provid
pub(super) fn map_stream_error(error: CodexClientError) -> MappedProviderFailure {
let allows_pre_delivery_retry = stream_transport_allows_pre_delivery_retry(&error);
let websocket_failure = error.transport() == Some(CodexBackendTransport::WebSocket);
let websocket_checkpoint_retryable = websocket_metadata_close_can_resume(&error);
let mut failure = map_client_error(error, UpstreamSendState::Sent, false);
failure.websocket_transport_retryable = allows_pre_delivery_retry && websocket_failure;
failure.websocket_checkpoint_retryable =
allows_pre_delivery_retry && websocket_checkpoint_retryable;
if allows_pre_delivery_retry && !websocket_failure {
failure.error = failure.error.with_pre_delivery_retry();
}
failure
}
fn websocket_metadata_close_can_resume(error: &CodexClientError) -> bool {
let CodexClientError::WebSocket(error) = error else {
return false;
};
let Some(close) = error.close_before_terminal() else {
return false;
};
close.code() == Some(1000)
&& matches!(
close.last_event_type(),
Some("codex.response.metadata" | "response.metadata")
)
}
pub(super) fn stream_transport_allows_pre_delivery_retry(error: &CodexClientError) -> bool {
match error {
CodexClientError::Http(_)
@@ -865,9 +814,13 @@ pub(super) fn map_client_error(
error.classified(),
CodexWebSocketExchangeError::ConnectionLimitReached
) {
failure.error = failure.error.with_upstream_code(OpaqueUpstreamValue::new(
WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE.to_owned(),
));
failure.error =
failure
.error
.with_replay_safe()
.with_upstream_code(OpaqueUpstreamValue::new(
WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE.to_owned(),
));
}
if let Some(client_visible_error) = client_visible_error {
failure.error = failure
@@ -963,9 +916,6 @@ fn websocket_diagnostic(error: &CodexWebSocketExchangeError) -> ProviderDiagnost
CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. } => {
"Reused OpenAI WebSocket died before the first upstream event".to_owned()
}
CodexWebSocketExchangeError::InitialEventTimeout { timeout } => {
format!("OpenAI WebSocket first-event timeout after {timeout:?}")
}
CodexWebSocketExchangeError::ClosedBeforeTerminal(_) => {
unreachable!("close-before-terminal errors are handled before the variant match")
}
@@ -1098,7 +1048,6 @@ pub(super) fn map_upstream_failure(
MappedProviderFailure {
error,
websocket_transport_retryable: false,
websocket_checkpoint_retryable: false,
account_failure: account_failure(
category,
failure.retry_after_seconds,
@@ -1217,8 +1166,9 @@ pub(super) fn websocket_send_state(error: &CodexWebSocketExchangeError) -> Upstr
| CodexWebSocketExchangeError::SendTimeout { .. }
| CodexWebSocketExchangeError::ClosedBeforeTerminal(_)
| CodexWebSocketExchangeError::ReceiveIdleTimeout { .. }
| CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. }
| CodexWebSocketExchangeError::InitialEventTimeout { .. } => UpstreamSendState::Ambiguous,
| CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. } => {
UpstreamSendState::Ambiguous
}
CodexWebSocketExchangeError::ConnectionObserved { .. } => {
unreachable!("classified websocket errors never retain observation wrappers")
}
@@ -1232,8 +1182,7 @@ pub(super) fn websocket_error_kind(error: &CodexWebSocketExchangeError) -> Provi
| CodexWebSocketExchangeError::UnexpectedBinaryEvent => ProviderErrorKind::Protocol,
CodexWebSocketExchangeError::ConnectTimeout { .. }
| CodexWebSocketExchangeError::SendTimeout { .. }
| CodexWebSocketExchangeError::ReceiveIdleTimeout { .. }
| CodexWebSocketExchangeError::InitialEventTimeout { .. } => ProviderErrorKind::Timeout,
| CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } => ProviderErrorKind::Timeout,
CodexWebSocketExchangeError::OriginCircuitOpen
| CodexWebSocketExchangeError::OriginHalfOpenBusy
| CodexWebSocketExchangeError::SharedConnectFailed
@@ -61,10 +61,7 @@ use crate::credential::{
derive_codex_cyber_policy_session_key, derive_codex_session_affinity,
derive_previous_response_id_hash,
};
use crate::session_transport::{
CodexSessionRecoveryTransition, CodexSessionTransportDecision, CodexSessionTransportRecovery,
CodexSessionWebSocketProbe, SessionWebSocketFallback,
};
use crate::session_transport::CodexSessionTransportRecovery;
use crate::transport::canonical::{
CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalOutcome,
};
@@ -94,7 +91,7 @@ use crate::transport::{
CodexBackendJsonResponse, CodexBackendStreamingResponse, CodexBackendTransport,
CodexClientError, CodexRateLimitUpdates, CodexRequestContext, CodexResponseMetadata,
CodexTransportMetrics, CodexTurnStateUpdate, CodexUpstreamDiagnostics, CodexWebSocketPool,
WebSocketConnectionPreference, endpoint_url,
endpoint_url,
};
mod execution;
@@ -321,14 +318,10 @@ impl Provider for CodexProvider {
identity.prepare_local_conversation(&mut upstream_request);
}
if let Some(previous_session) = previous_session.as_ref() {
upstream_request.turn_state = if previous_session.retry_checkpoint
|| same_client_turn(
previous_session.client_turn_id.as_deref(),
upstream_request.client_turn_id.as_deref(),
) {
// 客户端可能携带 metadata-close 失败后才拿到的新状态;同一
// turn 内显式回传值比最后一次成功响应保存的状态更新。同请求
// 内部恢复由一次性 retry checkpoint 明确证明,不依赖客户端 turn ID。
upstream_request.turn_state = if same_client_turn(
previous_session.client_turn_id.as_deref(),
upstream_request.client_turn_id.as_deref(),
) {
upstream_request
.turn_state
.take()
@@ -453,60 +446,17 @@ impl Provider for CodexProvider {
);
let requirement = transport_requirement(&upstream_request);
let requested_transport = selected_transport(&upstream_request);
let session_transport_decision = if matches!(
context.transport(),
AttemptTransport::Default | AttemptTransport::Retry(_)
) && requested_transport
== CodexProviderTransport::PreferWebSocket
&& requirement.allows_session_transport_recovery()
&& let Some(affinity) = session_affinity.as_ref()
{
self.session_transport_recovery.decide(affinity.key())
} else {
CodexSessionTransportDecision::Default
};
let session_transport_action = session_transport_decision.action();
let session_websocket_failure_count = session_transport_decision.failure_count();
let session_http_cooldown_ms = session_transport_decision
.retry_after()
.map(|duration| u64::try_from(duration.as_millis()).unwrap_or(u64::MAX));
let mut websocket_connection_preference = WebSocketConnectionPreference::ReuseOrConnect;
let mut session_recovery_probe = None;
let session_http_fallback = requirement.allows_pre_send_http_fallback()
&& session_affinity
.as_ref()
.is_some_and(|affinity| self.session_transport_recovery.uses_http(affinity.key()));
let transport = if requirement.requires_websocket() {
CodexProviderTransport::PreferWebSocket
} else if context.transport() == AttemptTransport::Fallback || session_http_fallback {
CodexProviderTransport::HttpOnly
} else {
match context.transport() {
AttemptTransport::Default | AttemptTransport::Retry(_) => {
match session_transport_decision {
CodexSessionTransportDecision::Default => requested_transport,
CodexSessionTransportDecision::FreshWebSocket { probe, .. } => {
websocket_connection_preference = WebSocketConnectionPreference::Fresh;
session_recovery_probe = Some(probe);
CodexProviderTransport::PreferWebSocket
}
CodexSessionTransportDecision::HttpSse { .. } => {
CodexProviderTransport::HttpOnly
}
}
}
AttemptTransport::Fallback => CodexProviderTransport::HttpOnly,
}
requested_transport
};
if let Some(session_transport_action) = session_transport_action {
tracing::info!(
request_id = %context.request_id(),
attempt_index = context.attempt_index().get(),
session_affinity_key_hash = session_affinity
.as_ref()
.map_or("", CodexSessionAffinity::key_hash),
transport_requirement = requirement.as_str(),
session_transport_action,
session_websocket_failure_count = session_websocket_failure_count.unwrap_or_default(),
session_http_cooldown_ms = session_http_cooldown_ms.unwrap_or_default(),
session_http_cooldown_present = session_http_cooldown_ms.is_some(),
"Applied OpenAI session transport recovery decision"
);
}
apply_transport(&mut upstream_request, transport);
let metadata = ProviderCallMetadata::new(
provider_kind,
@@ -553,9 +503,7 @@ impl Provider for CodexProvider {
output_started_at,
session_affinity_key,
session_affinity_key_hash,
websocket_connection_preference,
session_transport_recovery: self.session_transport_recovery.clone(),
session_recovery_probe,
websocket_retry_count,
stream_max_retries: self.stream_max_retries,
session_capture,
@@ -1,4 +1,4 @@
//! OpenAI 根会话级上游传输恢复状态。
//! 根会话的连续 WS 失败预算;耗尽后保持 HTTP,空闲过期后释放状态。
use std::collections::HashMap;
use std::sync::{Arc, Mutex, MutexGuard};
@@ -6,453 +6,106 @@ use std::time::Duration;
use gateway_core::provider_ports::ProviderSessionAffinityKey;
use tokio::time::Instant;
use uuid::Uuid;
/// 未成功恢复的失败链需要跨过最长冷却,避免 4 小时档在半开前被清理。
const SESSION_RECOVERY_RETENTION: Duration = Duration::from_secs(8 * 60 * 60);
const HTTP_COOLDOWNS: [Duration; 4] = [
Duration::from_secs(5 * 60),
Duration::from_secs(30 * 60),
Duration::from_secs(60 * 60),
Duration::from_secs(4 * 60 * 60),
];
/// Provider 是进程级单例;限制不可显式观测结束的客户端会话占用。
const MAX_SESSION_RECOVERIES: usize = 16_384;
const SESSION_IDLE_RETENTION: Duration = Duration::from_secs(8 * 60 * 60);
const MAX_SESSION_TRANSPORTS: usize = 16_384;
#[derive(Clone)]
#[derive(Clone, Default)]
pub(crate) struct CodexSessionTransportRecovery {
state: Arc<Mutex<SessionTransportRecoveryState>>,
http_cooldowns: [Duration; 4],
retention: Duration,
capacity: usize,
sessions: Arc<Mutex<HashMap<ProviderSessionAffinityKey, SessionTransport>>>,
}
#[derive(Default)]
struct SessionTransportRecoveryState {
sessions: HashMap<ProviderSessionAffinityKey, SessionCircuit>,
struct SessionTransport {
state: TransportState,
last_used: Instant,
}
enum SessionCircuit {
FreshPending {
failure_count: u8,
last_failure_at: Instant,
},
Open {
failure_count: u8,
last_failure_at: Instant,
until: Instant,
},
HalfOpen {
failure_count: u8,
last_failure_at: Instant,
probe_id: Uuid,
},
}
impl SessionCircuit {
const fn failure_count(&self) -> u8 {
match self {
Self::FreshPending { failure_count, .. }
| Self::Open { failure_count, .. }
| Self::HalfOpen { failure_count, .. } => *failure_count,
}
}
const fn last_failure_at(&self) -> Instant {
match self {
Self::FreshPending {
last_failure_at, ..
}
| Self::Open {
last_failure_at, ..
}
| Self::HalfOpen {
last_failure_at, ..
} => *last_failure_at,
}
}
}
impl Default for CodexSessionTransportRecovery {
fn default() -> Self {
Self {
state: Arc::new(Mutex::new(SessionTransportRecoveryState::default())),
http_cooldowns: HTTP_COOLDOWNS,
retention: SESSION_RECOVERY_RETENTION,
capacity: MAX_SESSION_RECOVERIES,
}
}
}
#[derive(Default)]
pub(crate) enum CodexSessionTransportDecision {
#[default]
Default,
FreshWebSocket {
failure_count: u8,
probe: CodexSessionWebSocketProbe,
},
HttpSse {
reason: SessionHttpFallbackReason,
failure_count: u8,
retry_after: Option<Duration>,
},
}
impl CodexSessionTransportDecision {
pub(crate) const fn action(&self) -> Option<&'static str> {
match self {
Self::Default => None,
Self::FreshWebSocket { .. } => Some("fresh_websocket_probe"),
Self::HttpSse {
reason: SessionHttpFallbackReason::Cooldown,
..
} => Some("http_cooldown"),
Self::HttpSse {
reason: SessionHttpFallbackReason::ProbeInFlight,
..
} => Some("http_probe_in_flight"),
}
}
pub(crate) const fn failure_count(&self) -> Option<u8> {
match self {
Self::Default => None,
Self::FreshWebSocket { failure_count, .. } | Self::HttpSse { failure_count, .. } => {
Some(*failure_count)
}
}
}
pub(crate) const fn retry_after(&self) -> Option<Duration> {
match self {
Self::HttpSse { retry_after, .. } => *retry_after,
Self::Default | Self::FreshWebSocket { .. } => None,
}
}
}
#[derive(Clone, Copy)]
pub(crate) enum SessionWebSocketFallback {
RetryBudgetExhausted,
UpgradeRequired,
}
impl SessionWebSocketFallback {
const fn minimum_failure_count(self) -> u8 {
match self {
Self::RetryBudgetExhausted => 2,
Self::UpgradeRequired => 5,
}
}
}
#[derive(Clone, Copy)]
pub(crate) enum SessionHttpFallbackReason {
Cooldown,
ProbeInFlight,
}
#[derive(Clone, Copy)]
pub(crate) struct CodexSessionRecoveryTransition {
failure_count: u8,
cooldown: Option<Duration>,
}
impl CodexSessionRecoveryTransition {
pub(crate) const fn action(self) -> &'static str {
if self.cooldown.is_some() {
"http_cooldown"
} else {
"fresh_websocket_next"
}
}
pub(crate) const fn failure_count(self) -> u8 {
self.failure_count
}
pub(crate) const fn cooldown(self) -> Option<Duration> {
self.cooldown
}
}
/// 一次根会话级 fresh WebSocket 许可。
///
/// 许可未得到明确结果就被释放时,恢复状态会重新允许下一次半开探测。
pub(crate) struct CodexSessionWebSocketProbe {
recovery: CodexSessionTransportRecovery,
key: ProviderSessionAffinityKey,
probe_id: Uuid,
armed: bool,
}
impl CodexSessionWebSocketProbe {
fn new(
recovery: CodexSessionTransportRecovery,
key: ProviderSessionAffinityKey,
probe_id: Uuid,
) -> Self {
Self {
recovery,
key,
probe_id,
armed: true,
}
}
pub(crate) fn succeed(mut self) -> bool {
let recovered = self.recovery.record_probe_success(&self.key, self.probe_id);
self.armed = false;
recovered
}
pub(crate) fn post_send_failed(mut self) -> CodexSessionRecoveryTransition {
let transition = self.recovery.record_post_send_failure(&self.key);
self.armed = false;
transition
}
pub(crate) fn fallback(
mut self,
fallback: SessionWebSocketFallback,
) -> CodexSessionRecoveryTransition {
let transition = self.recovery.record_websocket_fallback(&self.key, fallback);
self.armed = false;
transition
}
}
impl Drop for CodexSessionWebSocketProbe {
fn drop(&mut self) {
if self.armed {
self.recovery.cancel_probe(&self.key, self.probe_id);
}
}
enum TransportState {
WebSocket { failures: u32 },
Http,
}
impl CodexSessionTransportRecovery {
pub(crate) fn decide(&self, key: &ProviderSessionAffinityKey) -> CodexSessionTransportDecision {
self.decide_at(key, Instant::now())
}
pub(crate) fn record_post_send_failure(
&self,
key: &ProviderSessionAffinityKey,
) -> CodexSessionRecoveryTransition {
self.record_post_send_failure_at(key, Instant::now())
}
pub(crate) fn record_websocket_fallback(
&self,
key: &ProviderSessionAffinityKey,
fallback: SessionWebSocketFallback,
) -> CodexSessionRecoveryTransition {
self.record_websocket_fallback_at(key, fallback, Instant::now())
}
fn decide_at(
&self,
key: &ProviderSessionAffinityKey,
now: Instant,
) -> CodexSessionTransportDecision {
let mut state = self.lock();
let stale = state.sessions.get(key).is_some_and(|circuit| {
now.saturating_duration_since(circuit.last_failure_at()) >= self.retention
});
if stale {
state.sessions.remove(key);
return CodexSessionTransportDecision::Default;
}
let Some(circuit) = state.sessions.get_mut(key) else {
return CodexSessionTransportDecision::Default;
pub(crate) fn uses_http(&self, key: &ProviderSessionAffinityKey) -> bool {
let now = Instant::now();
let mut sessions = self.lock();
let Some(session) = sessions.get_mut(key) else {
return false;
};
match circuit {
SessionCircuit::Open {
failure_count,
until,
..
} if now < *until => CodexSessionTransportDecision::HttpSse {
reason: SessionHttpFallbackReason::Cooldown,
failure_count: *failure_count,
retry_after: Some(until.saturating_duration_since(now)),
},
SessionCircuit::HalfOpen { failure_count, .. } => {
CodexSessionTransportDecision::HttpSse {
reason: SessionHttpFallbackReason::ProbeInFlight,
failure_count: *failure_count,
retry_after: None,
}
}
SessionCircuit::FreshPending {
failure_count,
last_failure_at,
}
| SessionCircuit::Open {
failure_count,
last_failure_at,
..
} => {
let failure_count = *failure_count;
let last_failure_at = *last_failure_at;
let probe_id = Uuid::new_v4();
*circuit = SessionCircuit::HalfOpen {
failure_count,
last_failure_at,
probe_id,
};
CodexSessionTransportDecision::FreshWebSocket {
failure_count,
probe: CodexSessionWebSocketProbe::new(self.clone(), key.clone(), probe_id),
}
}
if now.saturating_duration_since(session.last_used) >= SESSION_IDLE_RETENTION {
sessions.remove(key);
return false;
}
session.last_used = now;
matches!(session.state, TransportState::Http)
}
fn record_post_send_failure_at(
/// 返回连续失败是否已耗尽首发加重试预算。
pub(crate) fn record_websocket_failure(
&self,
key: &ProviderSessionAffinityKey,
now: Instant,
) -> CodexSessionRecoveryTransition {
let mut state = self.lock();
let failure_count = self.next_failure_count(state.sessions.get(key), now);
if !state.sessions.contains_key(key) {
self.make_room(&mut state, now);
}
if failure_count == 1 {
state.sessions.insert(
key.clone(),
SessionCircuit::FreshPending {
failure_count,
last_failure_at: now,
},
);
return CodexSessionRecoveryTransition {
failure_count,
cooldown: None,
};
}
let cooldown = self.cooldown_for(failure_count);
state.sessions.insert(
key.clone(),
SessionCircuit::Open {
failure_count,
last_failure_at: now,
until: now + cooldown,
},
);
CodexSessionRecoveryTransition {
failure_count,
cooldown: Some(cooldown),
}
}
fn record_websocket_fallback_at(
&self,
key: &ProviderSessionAffinityKey,
fallback: SessionWebSocketFallback,
now: Instant,
) -> CodexSessionRecoveryTransition {
let mut state = self.lock();
let failure_count = self
.next_failure_count(state.sessions.get(key), now)
.max(fallback.minimum_failure_count());
if !state.sessions.contains_key(key) {
self.make_room(&mut state, now);
}
let cooldown = self.cooldown_for(failure_count);
state.sessions.insert(
key.clone(),
SessionCircuit::Open {
failure_count,
last_failure_at: now,
until: now + cooldown,
},
);
CodexSessionRecoveryTransition {
failure_count,
cooldown: Some(cooldown),
}
}
fn next_failure_count(&self, circuit: Option<&SessionCircuit>, now: Instant) -> u8 {
circuit.map_or(1, |circuit| {
if now.saturating_duration_since(circuit.last_failure_at()) < self.retention {
circuit.failure_count().saturating_add(1)
} else {
1
max_retries: u32,
) -> bool {
self.update(key, |state| match state {
TransportState::WebSocket { failures } if failures < max_retries => {
TransportState::WebSocket {
failures: failures + 1,
}
}
TransportState::WebSocket { .. } | TransportState::Http => TransportState::Http,
})
}
fn cooldown_for(&self, failure_count: u8) -> Duration {
let index = usize::from(failure_count.saturating_sub(2)).min(self.http_cooldowns.len() - 1);
self.http_cooldowns[index]
pub(crate) fn disable_websocket(&self, key: &ProviderSessionAffinityKey) {
self.update(key, |_| TransportState::Http);
}
fn record_probe_success(&self, key: &ProviderSessionAffinityKey, probe_id: Uuid) -> bool {
let mut state = self.lock();
let owns_probe = matches!(
state.sessions.get(key),
Some(SessionCircuit::HalfOpen {
probe_id: active_probe_id,
..
}) if *active_probe_id == probe_id
);
if owns_probe {
state.sessions.remove(key);
}
owns_probe
pub(crate) fn websocket_succeeded(&self, key: &ProviderSessionAffinityKey) {
self.lock().remove(key);
}
fn cancel_probe(&self, key: &ProviderSessionAffinityKey, probe_id: Uuid) {
let mut state = self.lock();
let Some(SessionCircuit::HalfOpen {
failure_count,
last_failure_at,
probe_id: active_probe_id,
}) = state.sessions.get(key)
else {
return;
};
if *active_probe_id != probe_id {
return;
fn update(
&self,
key: &ProviderSessionAffinityKey,
update: impl FnOnce(TransportState) -> TransportState,
) -> bool {
let now = Instant::now();
let mut sessions = self.lock();
let previous = sessions
.remove(key)
.filter(|session| {
now.saturating_duration_since(session.last_used) < SESSION_IDLE_RETENTION
})
.map_or(TransportState::WebSocket { failures: 0 }, |session| {
session.state
});
if sessions.len() >= MAX_SESSION_TRANSPORTS {
sessions.retain(|_, session| {
now.saturating_duration_since(session.last_used) < SESSION_IDLE_RETENTION
});
if sessions.len() >= MAX_SESSION_TRANSPORTS
&& let Some(oldest) = sessions
.iter()
.min_by_key(|(_, session)| session.last_used)
.map(|(key, _)| key.clone())
{
sessions.remove(&oldest);
}
}
let failure_count = *failure_count;
let last_failure_at = *last_failure_at;
state.sessions.insert(
let state = update(previous);
let uses_http = matches!(state, TransportState::Http);
sessions.insert(
key.clone(),
SessionCircuit::FreshPending {
failure_count,
last_failure_at,
SessionTransport {
state,
last_used: now,
},
);
uses_http
}
fn make_room(&self, state: &mut SessionTransportRecoveryState, now: Instant) {
if state.sessions.len() < self.capacity {
return;
}
state.sessions.retain(|_, circuit| {
now.saturating_duration_since(circuit.last_failure_at()) < self.retention
});
if state.sessions.len() < self.capacity {
return;
}
let oldest = state
.sessions
.iter()
.min_by_key(|(_, circuit)| circuit.last_failure_at())
.map(|(key, _)| key.clone());
if let Some(oldest) = oldest {
state.sessions.remove(&oldest);
}
}
fn lock(&self) -> MutexGuard<'_, SessionTransportRecoveryState> {
self.state
fn lock(&self) -> MutexGuard<'_, HashMap<ProviderSessionAffinityKey, SessionTransport>> {
self.sessions
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
@@ -511,8 +511,6 @@ pub enum CodexTransportDecision {
Http2WebSocketBudgetExhausted,
Http2BreakerOpen,
Http2PoolUnavailable,
Http2PreSendFailure,
Http2PreDeliveryFailure,
}
impl CodexTransportDecision {
@@ -526,8 +524,6 @@ impl CodexTransportDecision {
Self::Http2WebSocketBudgetExhausted => "http2_ws_budget_exhausted",
Self::Http2BreakerOpen => "http2_breaker_open",
Self::Http2PoolUnavailable => "http2_pool_unavailable",
Self::Http2PreSendFailure => "http2_ws_pre_send_failure",
Self::Http2PreDeliveryFailure => "http2_ws_pre_delivery_failure",
}
}
}
@@ -615,7 +611,6 @@ pub(crate) struct PreparedResponseTransport {
pub(super) requirement: TransportRequirement,
pub(super) route: PreparedResponseRoute,
pub(super) metrics: CodexTransportMetrics,
pub(super) defer_websocket_recovery: bool,
}
pub(super) enum PreparedResponseRoute {
@@ -862,19 +857,18 @@ pub(super) fn websocket_success_decision(
}
}
pub(super) fn http_fallback_decision(
pub(super) fn local_http_fallback_decision(
error: &CodexWebSocketExchangeError,
) -> CodexTransportDecision {
) -> Option<CodexTransportDecision> {
match error.classified() {
CodexWebSocketExchangeError::OriginCircuitOpen
| CodexWebSocketExchangeError::OriginHalfOpenBusy => {
CodexTransportDecision::Http2BreakerOpen
Some(CodexTransportDecision::Http2BreakerOpen)
}
CodexWebSocketExchangeError::ContinuationUnavailable { .. }
| CodexWebSocketExchangeError::SharedConnectFailed => {
CodexTransportDecision::Http2PoolUnavailable
CodexWebSocketExchangeError::ContinuationUnavailable { .. } => {
Some(CodexTransportDecision::Http2PoolUnavailable)
}
_ => CodexTransportDecision::Http2PreSendFailure,
_ => None,
}
}
@@ -33,7 +33,7 @@ use crate::transport::{
response_meta,
websocket::{
CodexWebSocketConnection, CodexWebSocketExchangeError, CodexWebSocketPool,
CodexWebSocketPoolKey, CodexWebSocketStreamingExchange, DEFAULT_INITIAL_EVENT_TIMEOUT,
CodexWebSocketPoolKey, CodexWebSocketStreamingExchange, DEFAULT_STREAM_IDLE_TIMEOUT,
WEBSOCKET_FAST_PATH_BUDGET, WebSocketFastPath, WebSocketOriginBreaker,
execute_prepared_response_create_request_stream, post_send_ambiguous,
prepare_response_create_request_with_pool, websocket_audit_dir,
@@ -43,16 +43,6 @@ use crate::transport::{
use super::client::*;
/// 单次 Responses WebSocket opening 是否可以使用会话连接池。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub(crate) enum WebSocketConnectionPreference {
/// 优先复用匹配的空闲连接,否则建立并登记一条新连接。
#[default]
ReuseOrConnect,
/// 绕过连接池建立独立连接,用于根会话恢复探针。
Fresh,
}
impl CodexBackendClient {
/// 构造客户端。
pub fn new(
@@ -216,27 +206,6 @@ impl CodexBackendClient {
.await
}
/// Provider 生产路径:保留所有 WebSocket 失败,交由统一重试预算处理。
pub(crate) async fn create_response_stream_with_deferred_websocket_recovery(
&self,
request: &CodexResponsesRequest,
context: CodexRequestContext<'_>,
pool_account_id: Option<&str>,
connection_preference: WebSocketConnectionPreference,
) -> CodexClientResult<CodexBackendStreamingResponse> {
let prepared = self
.prepare_response_transport(
request,
context,
pool_account_id,
connection_preference,
true,
)
.await?;
self.create_response_stream_with_prepared(request, context, prepared)
.await
}
/// 在发送 payload 前完成 transport 选择和可取消的 WebSocket opening。
#[doc(hidden)]
pub(crate) async fn prepare_response_transport_with_pool_account(
@@ -244,36 +213,18 @@ impl CodexBackendClient {
request: &CodexResponsesRequest,
context: CodexRequestContext<'_>,
pool_account_id: Option<&str>,
) -> CodexClientResult<PreparedResponseTransport> {
self.prepare_response_transport(
request,
context,
pool_account_id,
WebSocketConnectionPreference::ReuseOrConnect,
false,
)
.await
}
/// 与 [`Self::prepare_response_transport_with_pool_account`] 相同;
/// fresh 偏好会绕过连接池,为受控恢复探针强制新建连接。
async fn prepare_response_transport(
&self,
request: &CodexResponsesRequest,
context: CodexRequestContext<'_>,
pool_account_id: Option<&str>,
connection_preference: WebSocketConnectionPreference,
defer_websocket_recovery: bool,
) -> CodexClientResult<PreparedResponseTransport> {
let requirement = transport_requirement(request);
context.trace.cloned().unwrap_or_default().record("transport.preparing", serde_json::json!({
"requirement": requirement.as_str(), "connectionPreference": format!("{connection_preference:?}"),
}));
context.trace.cloned().unwrap_or_default().record(
"transport.preparing",
serde_json::json!({
"requirement": requirement.as_str(),
}),
);
if requirement == TransportRequirement::HttpRequired {
return Ok(PreparedResponseTransport {
requirement,
route: PreparedResponseRoute::Http,
defer_websocket_recovery,
metrics: CodexTransportMetrics {
decision: Some(CodexTransportDecision::HttpRequired),
..CodexTransportMetrics::default()
@@ -315,12 +266,7 @@ impl CodexBackendClient {
let pool_key =
self.websocket_pool_key(request, context, pool_account_id, &connection_profile);
let pool_log_context = pool_key.as_ref().map(WebSocketPoolLogContext::from_key);
let pool = match connection_preference {
WebSocketConnectionPreference::ReuseOrConnect => {
self.websocket_pool.as_deref().zip(pool_key)
}
WebSocketConnectionPreference::Fresh => None,
};
let pool = self.websocket_pool.as_deref().zip(pool_key);
let fast_path_budget = match requirement {
TransportRequirement::PersistedContinuation | TransportRequirement::NewChain => {
Some(WEBSOCKET_FAST_PATH_BUDGET)
@@ -338,7 +284,7 @@ impl CodexBackendClient {
&self.websocket_origin_key,
fast_path_budget,
requirement.requires_websocket(),
Some(DEFAULT_INITIAL_EVENT_TIMEOUT),
Some(DEFAULT_STREAM_IDLE_TIMEOUT),
)
.await;
let prepared = match prepared {
@@ -363,15 +309,12 @@ impl CodexBackendClient {
transport_decision_wait_ms: Some(wait_ms),
..CodexTransportMetrics::default()
},
defer_websocket_recovery,
});
}
Err(error)
if (!defer_websocket_recovery || error.requires_immediate_pool_fallback())
&& requirement.allows_pre_send_http_fallback()
&& error.allows_pre_send_http_fallback() =>
if requirement.allows_pre_send_http_fallback()
&& let Some(decision) = local_http_fallback_decision(&error) =>
{
let decision = http_fallback_decision(&error);
let wait_ms = elapsed_duration_millis(prepare_started_at.elapsed());
context.trace.cloned().unwrap_or_default().record(
"transport.fallback", serde_json::json!({
@@ -390,7 +333,6 @@ impl CodexBackendClient {
transport_decision_wait_ms: Some(wait_ms),
..CodexTransportMetrics::default()
},
defer_websocket_recovery,
});
}
Err(error) => return Err(websocket_exchange_error_to_client_error(error)),
@@ -428,7 +370,6 @@ impl CodexBackendClient {
prepared,
})),
metrics,
defer_websocket_recovery,
})
}
@@ -443,7 +384,6 @@ impl CodexBackendClient {
requirement,
route,
metrics,
defer_websocket_recovery,
} = prepared;
context.trace.cloned().unwrap_or_default().record(
"transport.selected",
@@ -465,8 +405,7 @@ impl CodexBackendClient {
request: websocket_request,
prepared,
} = *route;
let delivery_wait_started_at = Instant::now();
let mut exchange = match execute_prepared_response_create_request_stream(
let mut exchange = execute_prepared_response_create_request_stream(
&websocket_request,
prepared,
context
@@ -476,161 +415,19 @@ impl CodexBackendClient {
.exchange("websocket"),
)
.await
{
Ok(exchange) => exchange,
Err(error)
if !defer_websocket_recovery
&& requirement.allows_pre_delivery_http_fallback()
&& error.allows_pre_delivery_http_fallback() =>
{
return self
.fallback_to_http_before_websocket_delivery(
request,
context,
requirement,
metrics,
delivery_wait_started_at,
error,
)
.await;
}
Err(error) => return Err(websocket_exchange_error_to_client_error(error)),
};
if requirement.allows_pre_delivery_http_fallback() {
let mut fresh_retry_used = false;
loop {
match await_websocket_delivery_boundary(&mut exchange).await {
Ok(DeliveryBoundary::Ready) => break,
Ok(DeliveryBoundary::ConnectionLimitReached)
if defer_websocket_recovery =>
{
return Err(CodexClientError::WebSocket(
CodexWebSocketExchangeError::ConnectionLimitReached,
));
}
Ok(DeliveryBoundary::ConnectionLimitReached) => {
if fresh_retry_used {
return self
.http_fallback_before_delivery(
request,
context,
requirement,
metrics,
delivery_wait_started_at,
"websocket connection limit reached",
)
.await;
}
fresh_retry_used = true;
context.trace.cloned().unwrap_or_default().record(
"transport.retry",
serde_json::json!({
"reason": "websocket_connection_limit_reached",
"connectionPreference": "fresh",
}),
);
drop(exchange);
match self
.prepare_response_transport(
request,
context,
None,
WebSocketConnectionPreference::Fresh,
false,
)
.await
{
Ok(PreparedResponseTransport {
route: PreparedResponseRoute::WebSocket(route),
..
}) => {
match execute_prepared_response_create_request_stream(
&route.request,
route.prepared,
context
.trace
.cloned()
.unwrap_or_default()
.exchange("websocket"),
)
.await
{
Ok(next_exchange) => {
exchange = next_exchange;
continue;
}
Err(error)
if error.allows_pre_delivery_http_fallback() =>
{
return self
.fallback_to_http_before_websocket_delivery(
request,
context,
requirement,
metrics,
delivery_wait_started_at,
error,
)
.await;
}
Err(error) => {
return Err(
websocket_exchange_error_to_client_error(
post_send_ambiguous(error),
),
);
}
}
}
Ok(PreparedResponseTransport {
route: PreparedResponseRoute::Http,
..
}) => {
return self
.create_response_stream_http_sse(request, context)
.await
.map(|mut response| {
merge_preparation_metrics(
&mut response.transport_metrics,
metrics,
);
response
});
}
Err(error) => {
return self
.http_fallback_before_delivery(
request,
context,
requirement,
metrics,
delivery_wait_started_at,
&error.to_string(),
)
.await;
}
}
}
Err(error)
if !defer_websocket_recovery
&& error.allows_pre_delivery_http_fallback() =>
{
return self
.fallback_to_http_before_websocket_delivery(
request,
context,
requirement,
metrics,
delivery_wait_started_at,
error,
)
.await;
}
Err(error) => {
return Err(websocket_exchange_error_to_client_error(
post_send_ambiguous(error),
));
}
.map_err(websocket_exchange_error_to_client_error)?;
if requirement.allows_connection_restart() {
match await_websocket_delivery_boundary(&mut exchange).await {
Ok(DeliveryBoundary::Ready) => {}
Ok(DeliveryBoundary::ConnectionLimitReached) => {
return Err(CodexClientError::WebSocket(
CodexWebSocketExchangeError::ConnectionLimitReached,
));
}
Err(error) => {
return Err(websocket_exchange_error_to_client_error(
post_send_ambiguous(error),
));
}
}
}
@@ -658,59 +455,6 @@ impl CodexBackendClient {
}
}
async fn fallback_to_http_before_websocket_delivery(
&self,
request: &CodexResponsesRequest,
context: CodexRequestContext<'_>,
requirement: TransportRequirement,
metrics: CodexTransportMetrics,
delivery_wait_started_at: Instant,
error: CodexWebSocketExchangeError,
) -> CodexClientResult<CodexBackendStreamingResponse> {
self.http_fallback_before_delivery(
request,
context,
requirement,
metrics,
delivery_wait_started_at,
&error.to_string(),
)
.await
}
async fn http_fallback_before_delivery(
&self,
request: &CodexResponsesRequest,
context: CodexRequestContext<'_>,
requirement: TransportRequirement,
mut metrics: CodexTransportMetrics,
delivery_wait_started_at: Instant,
detail: &str,
) -> CodexClientResult<CodexBackendStreamingResponse> {
let decision_wait_ms = metrics
.transport_decision_wait_ms
.unwrap_or_default()
.saturating_add(elapsed_duration_millis(delivery_wait_started_at.elapsed()));
metrics.decision = Some(CodexTransportDecision::Http2PreDeliveryFailure);
metrics.transport_decision_wait_ms = Some(decision_wait_ms);
context.trace.cloned().unwrap_or_default().record(
"transport.fallback",
serde_json::json!({
"to": "http_sse", "reason": "websocket_pre_delivery_failure",
"requirement": requirement.as_str(),
"decision": CodexTransportDecision::Http2PreDeliveryFailure.as_str(),
"waitMs": decision_wait_ms,
"detail": diagnostic_json(&serde_json::json!({"message": detail})),
}),
);
self.create_response_stream_http_sse(request, context)
.await
.map(|mut response| {
merge_preparation_metrics(&mut response.transport_metrics, metrics);
response
})
}
fn websocket_pool_key(
&self,
request: &CodexResponsesRequest,
@@ -1,10 +1,7 @@
//! 在 transport 边界采集的上游响应诊断信息。
use super::client::CodexClientVisibleUpstreamResponse;
use super::protocol::responses::{
PREVIOUS_RESPONSE_NOT_FOUND_CODE, ResponsesSseFailure,
is_bare_invalid_previous_response_id_error,
};
use super::protocol::responses::ResponsesSseFailure;
use base64::{Engine as _, engine::general_purpose::STANDARD};
use chrono::{DateTime, Utc};
use reqwest::{StatusCode, header::HeaderMap};
@@ -146,7 +143,7 @@ impl CodexUpstreamFailure {
rate_limit_headers: &[(String, String)],
send_phase: CodexUpstreamSendPhase,
) -> Self {
let fields = ParsedUpstreamError::from_http_response(status, body);
let fields = ParsedUpstreamError::from_http_response(body);
let category = classify_upstream_failure(
UpstreamFailureSource::HttpResponse,
Some(status),
@@ -268,7 +265,7 @@ struct ParsedUpstreamError {
}
impl ParsedUpstreamError {
fn from_http_response(status: StatusCode, body: &str) -> Self {
fn from_http_response(body: &str) -> Self {
let Ok(value) = serde_json::from_str::<Value>(body) else {
return Self {
code: None,
@@ -294,16 +291,7 @@ impl ParsedUpstreamError {
.or_else(|| value.get("message"))
.and_then(Value::as_str)
.and_then(non_empty_owned);
let code = if status == StatusCode::BAD_REQUEST
&& is_bare_invalid_previous_response_id_error(
code_value,
error_type.as_deref(),
client_message.as_deref(),
) {
Some(PREVIOUS_RESPONSE_NOT_FOUND_CODE.to_owned())
} else {
code_value.and_then(Value::as_str).and_then(non_empty_owned)
};
let code = code_value.and_then(Value::as_str).and_then(non_empty_owned);
let message = client_message
.clone()
.or_else(|| error.as_str().and_then(non_empty_owned))
@@ -21,7 +21,6 @@ pub mod tls;
pub mod usage;
pub mod websocket;
pub(crate) use self::client_sse::WebSocketConnectionPreference;
pub use self::{
canonical::{CodexCanonicalDecoder, CodexCanonicalError},
catalog::{
@@ -15,19 +15,6 @@ pub(crate) const X_CODEX_TURN_STATE_CLIENT_METADATA_KEY: &str = "x-codex-turn-st
/// 本地生成 history unavailable 错误时使用的官方提示文本。
pub(crate) const PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE: &str =
"Previous response was not found. Retrying the full request.";
const INVALID_PREVIOUS_RESPONSE_ID_MESSAGE: &str = "Invalid `previous_response_id`.";
/// 精确匹配上游省略错误码的旧版 previous response 拒绝形状。
pub(crate) fn is_bare_invalid_previous_response_id_error(
code: Option<&Value>,
error_type: Option<&str>,
message: Option<&str>,
) -> bool {
code.is_none_or(Value::is_null)
&& error_type == Some("invalid_request_error")
&& message == Some(INVALID_PREVIOUS_RESPONSE_ID_MESSAGE)
}
/// Codex Responses 上游请求体。
///
/// 发往上游的 Responses 请求。`body` 持有客户端原始 JSON object,逐字段(含顺序、
@@ -161,13 +148,8 @@ impl TransportRequirement {
)
}
/// WebSocket 已发送 payload、但尚未交付业务事件时,是否允许同账号 HTTP 重放。
pub fn allows_pre_delivery_http_fallback(self) -> bool {
matches!(self, Self::NewChain)
}
/// 会话级恢复状态是否可以为请求选择 fresh WebSocket 或临时 HTTP。
pub fn allows_session_transport_recovery(self) -> bool {
/// 无续接依赖的新请求可以在明确的连接级拒绝后重新建连。
pub fn allows_connection_restart(self) -> bool {
matches!(self, Self::NewChain)
}
@@ -2,10 +2,7 @@ use gateway_protocol::openai::sse::encode_sse_event;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::transport::protocol::responses::{
CodexResponsesRequest, PREVIOUS_RESPONSE_NOT_FOUND_CODE,
is_bare_invalid_previous_response_id_error, transport_requirement,
};
use crate::transport::protocol::responses::{CodexResponsesRequest, transport_requirement};
const REDACTED_PAYLOAD_VALUE: &str = "<redacted>";
@@ -57,49 +54,13 @@ pub fn websocket_event_to_sse_frame(raw: &str) -> Option<String> {
websocket_event_frame(&value, raw)
}
/// 同上,复用已解析的事件 JSON;`data` 段通常逐字节内嵌 `raw` 原文。
///
/// 仅剥离传输层内部帧:`codex.rate_limits` 与承载 turn_state 的
/// `response.metadata` 不下发客户端。唯一的业务帧适配是为上游省略错误码的
/// `Invalid previous_response_id` 补齐官方客户端用于完整历史重放的稳定错误码;
/// 原始 status、type 与 message 均保持不变。
/// 复用已解析的 JSON,剥离内部帧并原样转发公开事件。
pub(crate) fn websocket_event_frame(value: &Value, raw: &str) -> Option<String> {
let event = websocket_event_type(value)?;
if is_internal_websocket_event(event) || event == "response.metadata" {
return None;
}
let normalized = previous_response_not_found_event(value);
Some(encode_sse_event(
event,
normalized.as_deref().unwrap_or(raw),
))
}
fn previous_response_not_found_event(value: &Value) -> Option<String> {
if websocket_event_type(value) != Some("error")
|| value
.get("status")
.or_else(|| value.get("status_code"))
.and_then(Value::as_u64)
!= Some(400)
{
return None;
}
let error = value.get("error")?.as_object()?;
if !is_bare_invalid_previous_response_id_error(
error.get("code"),
error.get("type").and_then(Value::as_str),
error.get("message").and_then(Value::as_str),
) {
return None;
}
let mut normalized = value.clone();
normalized.get_mut("error")?.as_object_mut()?.insert(
"code".to_owned(),
Value::String(PREVIOUS_RESPONSE_NOT_FOUND_CODE.to_owned()),
);
serde_json::to_string(&normalized).ok()
Some(encode_sse_event(event, raw))
}
pub(crate) fn websocket_event_type(value: &Value) -> Option<&str> {
@@ -1,8 +1,14 @@
//! Responses WebSocket pre-send/post-send 与 pool/breaker 编排。
use std::time::{Duration, Instant};
use std::{
future::Future,
time::{Duration, Instant},
};
use tokio::{sync::oneshot, time::timeout};
use tokio::{
sync::oneshot,
time::{Instant as TokioInstant, timeout_at},
};
use super::{
breaker::{
@@ -43,7 +49,7 @@ pub(crate) struct PreparedWebSocket {
binding: PoolBinding,
connect_elapsed: Option<Duration>,
decision_wait_elapsed: Duration,
initial_event_timeout: Option<Duration>,
stream_idle_timeout: Option<Duration>,
}
enum PoolBinding {
@@ -111,7 +117,7 @@ pub(crate) async fn prepare_response_create_request_with_pool(
origin_key: &str,
fast_path_budget: Option<Duration>,
require_pool: bool,
fallback_initial_event_timeout: Option<Duration>,
fallback_stream_idle_timeout: Option<Duration>,
) -> Result<WebSocketFastPath<PreparedWebSocket>, CodexWebSocketExchangeError> {
let decision_started_at = Instant::now();
let Some((pool, key)) = pool else {
@@ -125,7 +131,7 @@ pub(crate) async fn prepare_response_create_request_with_pool(
breaker,
origin_key,
fast_path_budget,
fallback_initial_event_timeout,
fallback_stream_idle_timeout,
decision_started_at,
)
.await;
@@ -159,7 +165,7 @@ pub(crate) async fn prepare_response_create_request_with_pool(
},
connect_elapsed: None,
decision_wait_elapsed: decision_started_at.elapsed(),
initial_event_timeout: pool.initial_event_timeout(),
stream_idle_timeout: pool.stream_idle_timeout(),
}))
}
WebSocketPoolAcquire::Connect(connect_lease) => {
@@ -209,10 +215,10 @@ async fn prepare_unpooled_websocket(
breaker: &WebSocketOriginBreaker,
origin_key: &str,
fast_path_budget: Option<Duration>,
initial_event_timeout: Option<Duration>,
stream_idle_timeout: Option<Duration>,
decision_started_at: Instant,
) -> Result<WebSocketFastPath<PreparedWebSocket>, CodexWebSocketExchangeError> {
let permit = acquire_breaker_permit(breaker, origin_key)?;
let permit = acquire_breaker_permit(breaker, origin_key, fast_path_budget.is_some())?;
let context = pump_log_context_from_connection(request.connection());
let connected = connect_with_budget(
request.connection(),
@@ -224,7 +230,9 @@ async fn prepare_unpooled_websocket(
let connected = match connected {
Ok(WebSocketFastPath::Ready(connection)) => Ok(connection),
Ok(WebSocketFastPath::Missed) => {
permit.fast_timeout();
if let Some(permit) = permit {
permit.fast_timeout();
}
return Ok(WebSocketFastPath::Missed);
}
Err(error) => Err(error),
@@ -235,7 +243,7 @@ async fn prepare_unpooled_websocket(
binding: PoolBinding::Unpooled,
connect_elapsed: Some(connect_elapsed),
decision_wait_elapsed: decision_started_at.elapsed(),
initial_event_timeout,
stream_idle_timeout,
}))
}
@@ -248,7 +256,7 @@ async fn prepare_pooled_websocket(
fast_path_budget: Option<Duration>,
decision_started_at: Instant,
) -> Result<WebSocketFastPath<PreparedWebSocket>, CodexWebSocketExchangeError> {
let permit = match acquire_breaker_permit(breaker, origin_key) {
let permit = match acquire_breaker_permit(breaker, origin_key, fast_path_budget.is_some()) {
Ok(permit) => permit,
Err(error) => {
connect_lease.failed().await;
@@ -271,7 +279,7 @@ async fn prepare_pooled_websocket(
},
connect_elapsed: Some(handoff.connect_elapsed),
decision_wait_elapsed: decision_started_at.elapsed(),
initial_event_timeout: pool.initial_event_timeout(),
stream_idle_timeout: pool.stream_idle_timeout(),
})),
WebSocketFastPath::Missed => Ok(WebSocketFastPath::Missed),
}
@@ -280,7 +288,7 @@ async fn prepare_pooled_websocket(
struct PooledWebSocketConnectWaiter {
started_at: tokio::time::Instant,
receiver: oneshot::Receiver<Result<PooledWebSocketHandoff, CodexWebSocketExchangeError>>,
fast_path_reporter: WebSocketOriginFastPathReporter,
fast_path_reporter: Option<WebSocketOriginFastPathReporter>,
}
struct PooledWebSocketHandoff {
@@ -294,19 +302,16 @@ impl PooledWebSocketConnectWaiter {
&mut self,
fast_path_budget: Option<Duration>,
) -> Result<WebSocketFastPath<PooledWebSocketHandoff>, CodexWebSocketExchangeError> {
let received = match fast_path_budget {
Some(budget) => {
let remaining = budget.saturating_sub(self.started_at.elapsed());
match timeout(remaining, &mut self.receiver).await {
Ok(received) => received,
Err(_) => {
self.fast_path_reporter.missed();
return Ok(WebSocketFastPath::Missed);
let received =
match wait_for_fast_path(self.started_at, fast_path_budget, &mut self.receiver).await {
WebSocketFastPath::Ready(received) => received,
WebSocketFastPath::Missed => {
if let Some(reporter) = &self.fast_path_reporter {
reporter.missed();
}
return Ok(WebSocketFastPath::Missed);
}
}
None => (&mut self.receiver).await,
};
};
received
.map_err(|_| CodexWebSocketExchangeError::SharedConnectFailed)?
.map(WebSocketFastPath::Ready)
@@ -317,12 +322,14 @@ fn start_pooled_websocket_connect(
connection: CodexWebSocketConnection,
pool: CodexWebSocketPool,
connect_lease: WebSocketPoolConnectLease,
permit: WebSocketOriginBreakerPermit,
permit: Option<WebSocketOriginBreakerPermit>,
) -> PooledWebSocketConnectWaiter {
let task_key = connect_lease.key().clone();
let started_at = connect_lease.started_at();
let cancellation = connect_lease.cancellation_token();
let fast_path_reporter = permit.fast_path_reporter();
let fast_path_reporter = permit
.as_ref()
.map(WebSocketOriginBreakerPermit::fast_path_reporter);
let keepalive = pool.keepalive();
let context = PumpLogContext::new(
Some(task_key.account_id().to_owned()),
@@ -333,7 +340,7 @@ fn start_pooled_websocket_connect(
let connected = tokio::select! {
biased;
_ = cancellation.cancelled() => {
permit.cancel();
if let Some(permit) = permit { permit.cancel(); }
let _ = sender.send(Err(CodexWebSocketExchangeError::SharedConnectFailed));
connect_lease.failed().await;
tracing::info!(
@@ -450,21 +457,15 @@ async fn connect_with_budget(
fast_path_budget: Option<Duration>,
context: PumpLogContext,
) -> Result<WebSocketFastPath<(PooledWebSocketConnection, Duration)>, CodexWebSocketExchangeError> {
match fast_path_budget {
Some(budget) => {
let Ok(connected) = timeout(
budget,
connect_websocket_connection(connection, keepalive, context),
)
.await
else {
return Ok(WebSocketFastPath::Missed);
};
connected.map(WebSocketFastPath::Ready)
}
None => connect_websocket_connection(connection, keepalive, context)
.await
.map(WebSocketFastPath::Ready),
match wait_for_fast_path(
TokioInstant::now(),
fast_path_budget,
connect_websocket_connection(connection, keepalive, context),
)
.await
{
WebSocketFastPath::Ready(connected) => connected.map(WebSocketFastPath::Ready),
WebSocketFastPath::Missed => Ok(WebSocketFastPath::Missed),
}
}
@@ -489,9 +490,13 @@ async fn connect_websocket_connection(
fn acquire_breaker_permit(
breaker: &WebSocketOriginBreaker,
origin_key: &str,
) -> Result<WebSocketOriginBreakerPermit, CodexWebSocketExchangeError> {
allow_http_fallback: bool,
) -> Result<Option<WebSocketOriginBreakerPermit>, CodexWebSocketExchangeError> {
if !allow_http_fallback {
return Ok(None);
}
match breaker.try_acquire(origin_key) {
WebSocketOriginBreakerDecision::Allowed(permit) => Ok(permit),
WebSocketOriginBreakerDecision::Allowed(permit) => Ok(Some(permit)),
WebSocketOriginBreakerDecision::Open => Err(CodexWebSocketExchangeError::OriginCircuitOpen),
WebSocketOriginBreakerDecision::HalfOpenBusy => {
Err(CodexWebSocketExchangeError::OriginHalfOpenBusy)
@@ -500,9 +505,12 @@ fn acquire_breaker_permit(
}
fn finish_breaker_attempt(
permit: WebSocketOriginBreakerPermit,
permit: Option<WebSocketOriginBreakerPermit>,
connected: Result<(PooledWebSocketConnection, Duration), CodexWebSocketExchangeError>,
) -> Result<(PooledWebSocketConnection, Duration), CodexWebSocketExchangeError> {
let Some(permit) = permit else {
return connected;
};
match connected {
Ok(connection) => {
permit.succeed();
@@ -524,15 +532,21 @@ async fn wait_for_shared_connect(
waiter: pool::WebSocketPoolConnectWaiter,
fast_path_budget: Option<Duration>,
) -> WebSocketFastPath<WebSocketPoolConnectOutcome> {
match fast_path_budget {
Some(budget) => {
let remaining = waiter.remaining_budget(budget);
match timeout(remaining, waiter.wait()).await {
Ok(outcome) => WebSocketFastPath::Ready(outcome),
Err(_) => WebSocketFastPath::Missed,
}
}
None => WebSocketFastPath::Ready(waiter.wait().await),
wait_for_fast_path(waiter.started_at(), fast_path_budget, waiter.wait()).await
}
// 同一次 opening 的所有等待者共享截止时间;后来的请求不会重置 800ms 预算。
async fn wait_for_fast_path<F: Future>(
started_at: TokioInstant,
budget: Option<Duration>,
future: F,
) -> WebSocketFastPath<F::Output> {
let Some(budget) = budget else {
return WebSocketFastPath::Ready(future.await);
};
match timeout_at(started_at + budget, future).await {
Ok(value) => WebSocketFastPath::Ready(value),
Err(_) => WebSocketFastPath::Missed,
}
}
@@ -564,7 +578,7 @@ pub(crate) async fn execute_prepared_response_create_request_stream(
let PreparedWebSocket {
connection,
binding,
initial_event_timeout,
stream_idle_timeout,
..
} = prepared;
let (lease, reused, pool_decision) = binding.into_parts();
@@ -607,7 +621,7 @@ pub(crate) async fn execute_prepared_response_create_request_stream(
metadata,
pool_return,
reused,
initial_event_timeout,
stream_idle_timeout,
trace,
);
exchange.pool_decision = pool_decision;
@@ -90,12 +90,6 @@ pub enum CodexWebSocketExchangeError {
#[source]
source: Option<Box<CodexWebSocketExchangeError>>,
},
/// 建连并发送后,上游在配置时间内没有产生任何事件。
#[error("websocket first upstream event not received within {timeout:?}")]
InitialEventTimeout {
/// 首个上游事件超时时长。
timeout: Duration,
},
/// 将一个已分类交互错误与物理连接生命周期快照绑定。
#[error("{source}")]
ConnectionObserved {
@@ -325,38 +319,4 @@ impl CodexWebSocketExchangeError {
send_phase,
}))
}
/// opening 阶段只有明确的 transport 可用性失败才能切到同账号 HTTP。
pub(in crate::transport) fn allows_pre_send_http_fallback(&self) -> bool {
if let Self::ConnectionObserved { source, .. } = self {
return source.allows_pre_send_http_fallback();
}
matches!(
self,
Self::Connect(_)
| Self::ConnectTimeout { .. }
| Self::OriginCircuitOpen
| Self::OriginHalfOpenBusy
| Self::SharedConnectFailed
| Self::ContinuationUnavailable { .. }
)
}
/// 精确连接状态不可用属于本地池路由事实,不消耗上游 WS 重试预算。
pub(in crate::transport) fn requires_immediate_pool_fallback(&self) -> bool {
matches!(self.classified(), Self::ContinuationUnavailable { .. })
}
/// 首个业务事件交付前,普通 WebSocket 失败可切到同账号 HTTP。
pub(in crate::transport) fn allows_pre_delivery_http_fallback(&self) -> bool {
match self {
Self::Upstream(_) => false,
Self::PostSendAmbiguous {
source: Some(source),
..
} => source.allows_pre_delivery_http_fallback(),
Self::ConnectionObserved { source, .. } => source.allows_pre_delivery_http_fallback(),
_ => true,
}
}
}
@@ -6,23 +6,7 @@ use tokio::time::timeout;
use tungstenite::Message;
use super::super::pump::PumpedWebSocket;
use super::{
CodexWebSocketExchangeError, WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT,
WEBSOCKET_RECEIVE_IDLE_TIMEOUT,
};
pub(super) fn receive_idle_timeout(
saw_upstream_activity: bool,
initial_event_timeout: Option<Duration>,
) -> Duration {
if saw_upstream_activity {
WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT
} else {
initial_event_timeout
.filter(|timeout| !timeout.is_zero())
.unwrap_or(WEBSOCKET_RECEIVE_IDLE_TIMEOUT)
}
}
use super::CodexWebSocketExchangeError;
pub(super) async fn next_websocket_message(
websocket: &mut PumpedWebSocket,
@@ -4,7 +4,7 @@ mod io;
mod reducer;
mod stream;
use std::{pin::Pin, sync::Arc, time::Duration};
use std::{pin::Pin, sync::Arc};
use bytes::Bytes;
use futures::Stream;
@@ -19,8 +19,6 @@ use crate::transport::{
use super::error::CodexWebSocketExchangeError;
use super::pool::{CodexWebSocketConnectionMetadata, WebSocketPoolDecision};
const WEBSOCKET_RECEIVE_IDLE_TIMEOUT: Duration = Duration::from_secs(20);
const WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
const WEBSOCKET_STREAM_BUFFER: usize = 16;
pub(super) use self::stream::{WebSocketStreamPoolReturn, stream_websocket_response};
@@ -17,7 +17,7 @@ use super::super::{
PumpExitReason, PumpedWebSocket, WebSocketConnectionObservation, transport_metric_reason,
},
};
use super::io::{next_websocket_message, receive_idle_timeout, reused_stream_receive_error};
use super::io::{next_websocket_message, reused_stream_receive_error};
use super::reducer::{ExchangeAction, WebSocketTerminalKind, reduce_websocket_event};
use super::{
CodexWebSocketExchangeError, CodexWebSocketRateLimitUpdates, CodexWebSocketStreamingExchange,
@@ -62,7 +62,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
metadata: CodexWebSocketConnectionMetadata,
pool_return: Option<WebSocketStreamPoolReturn>,
reused_connection: bool,
initial_event_timeout: Option<Duration>,
stream_idle_timeout: Option<Duration>,
trace: TraceContext,
) -> CodexWebSocketStreamingExchange {
let websocket_connection_id = websocket.connection_id();
@@ -85,7 +85,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
metadata,
pool_return,
reused_connection,
initial_event_timeout,
stream_idle_timeout,
trace,
shutdown,
rate_limit_updates: rate_limit_updates_for_task,
@@ -125,7 +125,7 @@ struct WebSocketStreamForwardState {
metadata: CodexWebSocketConnectionMetadata,
pool_return: Option<WebSocketStreamPoolReturn>,
reused_connection: bool,
initial_event_timeout: Option<Duration>,
stream_idle_timeout: Option<Duration>,
shutdown: CancellationToken,
rate_limit_updates: CodexWebSocketRateLimitUpdates,
turn_state_update: CodexWebSocketTurnStateUpdate,
@@ -139,7 +139,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
mut metadata,
pool_return,
reused_connection,
initial_event_timeout,
stream_idle_timeout,
shutdown,
rate_limit_updates,
turn_state_update,
@@ -150,7 +150,6 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
.as_mut()
.map(|pool_return| std::mem::take(&mut pool_return.continuation))
.unwrap_or_default();
let mut saw_upstream_activity = false;
let mut last_event_type = None;
loop {
let message = tokio::select! {
@@ -177,20 +176,12 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
}
message = next_websocket_message(
&mut websocket,
receive_idle_timeout(saw_upstream_activity, initial_event_timeout),
stream_idle_timeout.filter(|timeout| !timeout.is_zero()).unwrap_or(super::super::pool::DEFAULT_STREAM_IDLE_TIMEOUT),
) => message,
};
let message = match message {
Ok(message) => message,
Err(error) => {
let error = match error {
CodexWebSocketExchangeError::ReceiveIdleTimeout { timeout }
if !saw_upstream_activity =>
{
CodexWebSocketExchangeError::InitialEventTimeout { timeout }
}
error => error,
};
trace.record(
"upstream.read.failed",
json!({
@@ -223,10 +214,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
break;
};
let raw = match message {
tungstenite::Message::Text(text) => {
saw_upstream_activity = true;
text.to_string()
}
tungstenite::Message::Text(text) => text.to_string(),
tungstenite::Message::Binary(bytes) => {
trace.capture("upstream.binary", &bytes);
let error = CodexWebSocketExchangeError::UnexpectedBinaryEvent;
@@ -475,7 +463,6 @@ fn exchange_exit_reason(error: &CodexWebSocketExchangeError) -> &'static str {
}
}
CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } => "receive_idle_timeout",
CodexWebSocketExchangeError::InitialEventTimeout { .. } => "initial_event_timeout",
CodexWebSocketExchangeError::UnexpectedBinaryEvent => "unexpected_binary_event",
_ => "exchange_failure",
}
@@ -41,5 +41,5 @@ pub(crate) use self::{
execute_prepared_response_create_request_stream, post_send_ambiguous,
prepare_response_create_request_with_pool,
},
pool::DEFAULT_INITIAL_EVENT_TIMEOUT,
pool::DEFAULT_STREAM_IDLE_TIMEOUT,
};
@@ -1,7 +1,5 @@
//! WebSocket 连接池 reservation、handoff 与 lease 生命周期。
use std::time::Duration;
use tokio::{
sync::{OwnedSemaphorePermit, watch},
time::Instant,
@@ -40,8 +38,8 @@ pub(crate) struct WebSocketPoolConnectWaiter {
}
impl WebSocketPoolConnectWaiter {
pub(crate) fn remaining_budget(&self, budget: Duration) -> Duration {
budget.saturating_sub(self.started_at.elapsed())
pub(crate) fn started_at(&self) -> Instant {
self.started_at
}
pub(crate) async fn wait(mut self) -> WebSocketPoolConnectOutcome {
@@ -37,7 +37,7 @@ const DEFAULT_MAX_AGE: Duration = Duration::from_mins(55);
const DEFAULT_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(25);
const DEFAULT_PING_INTERVAL: Duration = Duration::from_secs(25);
const DEFAULT_PING_TIMEOUT: Duration = Duration::from_secs(5);
pub(crate) const DEFAULT_INITIAL_EVENT_TIMEOUT: Duration = Duration::from_secs(20);
pub(crate) const DEFAULT_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(300);
/// WebSocket 连接池。
#[derive(Clone)]
@@ -73,8 +73,8 @@ pub struct CodexWebSocketPoolConfig {
pub ping_timeout: Duration,
/// idle socket 无活动多久后视为失活。
pub liveness_timeout: Option<Duration>,
/// 建连并发送后首个上游事件到达前的超时;`None` 表示禁用。
pub initial_event_timeout: Option<Duration>,
/// 等待下一条上游消息的空闲超时;`None` 或零值使用默认 300 秒。
pub stream_idle_timeout: Option<Duration>,
}
impl Default for CodexWebSocketPoolConfig {
@@ -87,9 +87,9 @@ impl Default for CodexWebSocketPoolConfig {
ping_interval: Some(DEFAULT_PING_INTERVAL),
ping_timeout: DEFAULT_PING_TIMEOUT,
// idle 连接不设失活截断:靠 ping/pong 保活,只在 max_age(55 分钟)
// 或 ping 失败时关闭,最大化跨轮复用(对齐 Codex CLI 的长连接策略)。
// 或 ping 失败时关闭,维持跨轮可复用连接。
liveness_timeout: None,
initial_event_timeout: Some(DEFAULT_INITIAL_EVENT_TIMEOUT),
stream_idle_timeout: Some(DEFAULT_STREAM_IDLE_TIMEOUT),
}
}
}
@@ -136,9 +136,9 @@ impl CodexWebSocketPool {
self.config.keepalive()
}
/// 建连并发送后首个上游事件到达前的超时;`None` 表示禁用。
pub(crate) fn initial_event_timeout(&self) -> Option<Duration> {
self.config.initial_event_timeout
/// 等待下一条上游消息的空闲超时;`None` 或零值使用默认 300 秒。
pub(crate) fn stream_idle_timeout(&self) -> Option<Duration> {
self.config.stream_idle_timeout
}
/// 注册由连接池生命周期托管的 opening 任务。
@@ -97,7 +97,7 @@ fn openai_config_restricts_upstream_base_url_to_https_or_loopback_http() {
#[test]
fn openai_config_defaults_to_the_provider_owned_operating_values() {
let config = OpenAiConfig::default();
assert_eq!(DEFAULT_STREAM_MAX_RETRIES, 3);
assert_eq!(DEFAULT_STREAM_MAX_RETRIES, 5);
assert_eq!(
(
@@ -105,7 +105,7 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() {
config.ws_pool.enabled,
config.ws_pool.max_age_ms,
config.ws_pool.max_connecting,
config.ws_pool.initial_event_timeout_ms,
config.ws_pool.stream_idle_timeout_ms,
config.quota.refresh_interval_minutes,
config.auth.refresh_enabled,
config.auth.oauth_client_id.as_str(),
@@ -117,7 +117,7 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() {
true,
3_300_000,
8,
20_000,
300_000,
15,
true,
"app_EMoamEEZ73f0CkXaXp7hrann",
@@ -459,7 +459,6 @@ fn contract_account_scope() -> Arc<FrozenAccountScope> {
"acct_truncated_stream",
"acct_usage_limit_request_path",
"acct_websocket_close",
"acct_websocket_cooldown",
"acct_websocket_fast_path",
"acct_websocket_busy_replay",
"acct_websocket_metadata_close",
@@ -666,7 +665,6 @@ async fn capture_turn_state_request(
previous_turn_id: Option<&str>,
current_turn_id: Option<&str>,
client_turn_state: Option<&str>,
retry_checkpoint: bool,
) -> wiremock::Request {
let account_id = "acct_session_affinity";
let store = Arc::new(MemoryAccountStore::default());
@@ -691,9 +689,6 @@ async fn capture_turn_state_request(
if let Some(previous_turn_id) = previous_turn_id {
session_state.insert("client_turn_id".to_owned(), json!(previous_turn_id));
}
if retry_checkpoint {
session_state.insert("retry_checkpoint".to_owned(), json!(true));
}
let mut protocol_context = Map::from_iter([("use_websocket".to_owned(), json!(false))]);
if let Some(current_turn_id) = current_turn_id {
protocol_context.insert("turn_id".to_owned(), json!(current_turn_id));
@@ -1326,7 +1321,7 @@ async fn selection_infrastructure_errors_have_a_distinct_classification() {
}
#[tokio::test]
async fn websocket_close_after_delivery_preserves_details_and_forces_a_fresh_client_retry() {
async fn websocket_close_after_delivery_preserves_details_and_reconnects_through_pool() {
const ACCOUNT_ID: &str = "acct_websocket_close";
const CONVERSATION_ID: &str = "conversation-websocket-close-after-delivery";
const SESSION_ID: &str = "committed-websocket-session";
@@ -1496,7 +1491,7 @@ async fn websocket_close_after_delivery_preserves_details_and_forces_a_fresh_cli
retry_pool_observation = Some(observation.websocket_pool());
}
}
assert_eq!(retry_pool_observation, Some(None));
assert_eq!(retry_pool_observation, Some(Some(WebSocketPoolKind::New)));
server.await.expect("WebSocket server");
}
@@ -1687,7 +1682,7 @@ async fn websocket_fast_path_miss_uses_http_and_keeps_background_preconnect() {
}
#[tokio::test]
async fn ambiguous_websocket_close_probes_a_fresh_websocket_for_the_next_new_chain() {
async fn ambiguous_websocket_close_reconnects_and_retains_native_continuation() {
const ACCOUNT_ID: &str = "acct_websocket_close";
const CONVERSATION_ID: &str = "conversation-websocket-ambiguous-close";
const SESSION_ID: &str = "sticky-websocket-session";
@@ -1776,7 +1771,7 @@ async fn ambiguous_websocket_close_probes_a_fresh_websocket_for_the_next_new_cha
let (stream, _) = listener
.accept()
.await
.expect("accept fresh WebSocket probe");
.expect("accept new pooled WebSocket");
let mut websocket = accept_codex_test_websocket(stream).await;
let _request = websocket
.next()
@@ -1799,11 +1794,6 @@ async fn ambiguous_websocket_close_probes_a_fresh_websocket_for_the_next_new_cha
.await
.expect("complete fresh WebSocket request");
let (stream, _) = listener
.accept()
.await
.expect("accept same-session continuation WebSocket connection");
let mut websocket = accept_codex_test_websocket(stream).await;
let request = websocket
.next()
.await
@@ -1817,7 +1807,7 @@ async fn ambiguous_websocket_close_probes_a_fresh_websocket_for_the_next_new_cha
.expect("same-session continuation request JSON");
assert_eq!(
payload.get("previous_response_id"),
Some(&json!("external-previous-response"))
Some(&json!("resp_fresh_retry"))
);
websocket
.send(completed_response("resp_external_continuation"))
@@ -1930,29 +1920,48 @@ async fn ambiguous_websocket_close_probes_a_fresh_websocket_for_the_next_new_cha
.expect("same session should prepare a fresh WebSocket stream");
assert_eq!(fresh_stream.metadata().transport().as_str(), "websocket");
let mut fresh_pool_observation = None;
let mut recovered_session = None;
while let Some(event) = fresh_stream.next().await {
let event = event.expect("fresh WebSocket response");
if let Some(update) = event.session_update() {
recovered_session = Some(update.clone());
}
if let Some(observation) = event.response_observation()
&& observation.transport().as_str() == "websocket"
{
fresh_pool_observation = Some(observation.websocket_pool());
}
}
assert_eq!(fresh_pool_observation, Some(None));
assert_eq!(fresh_pool_observation, Some(Some(WebSocketPoolKind::New)));
let recovered_session = recovered_session.expect("recovered session");
assert_eq!(
recovered_session.payload().get("continuation_scope"),
Some(&json!("connection_local"))
);
drop(fresh_stream);
let continuation_operation = Operation::Generate(generate_with_session_context(
"sticky-websocket-session",
Some("thread-continuation"),
None,
));
let continuation_operation = Operation::Generate(
generate_with_session_context(
"sticky-websocket-session",
Some("thread-continuation"),
None,
)
.with_provider_session_state(recovered_session),
);
let mut continuation_stream = provider
.execute(
planned_request("openai", continuation_operation),
external_continuation_context("req_sticky_external_continuation"),
pinned_continuation_context(
"req_recovered_continuation",
ACCOUNT_ID,
"resp_fresh_retry",
"resp_fresh_retry",
1,
ContinuationAttempt::Native,
),
)
.await
.expect("successful fresh probe must restore normal continuation routing");
.expect("recovered pooled socket must support native continuation");
assert_eq!(
continuation_stream.metadata().transport().as_str(),
"websocket"
@@ -2009,366 +2018,117 @@ async fn ambiguous_websocket_close_probes_a_fresh_websocket_for_the_next_new_cha
}
#[tokio::test]
async fn repeated_ambiguous_websocket_close_uses_steep_sse_cooldowns_and_one_half_open_probe() {
const ACCOUNT_ID: &str = "acct_websocket_cooldown";
const CONVERSATION_ID: &str = "conversation-websocket-cooldown";
async fn repeated_websocket_failures_exhaust_budget_then_use_http_and_report_http_failure() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, ACCOUNT_ID).await;
create_account(&store, "acct_websocket_close").await;
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind WebSocket listener");
.expect("bind listener");
let base_url = format!(
"http://{}",
listener.local_addr().expect("listener address")
);
let server = tokio::spawn(async move {
let completed_response = |response_id: &str| {
Message::Text(
json!({
"type": "response.completed",
"response": {
"id": response_id,
"model": "gpt-5.4",
"status": "completed",
"output": [],
"usage": {
"input_tokens": 1,
"output_tokens": 1,
"total_tokens": 2
}
}
})
.to_string()
.into(),
)
};
for response_id in ["resp_first_reset", "resp_fresh_probe_reset"] {
let (stream, _) = listener
.accept()
.await
.expect("accept failing WebSocket connection");
for code in [1000, 1001, 1011] {
let (stream, _) = listener.accept().await.expect("accept websocket");
let mut websocket = accept_codex_test_websocket(stream).await;
let _request = websocket
websocket
.next()
.await
.expect("WebSocket request")
.expect("valid WebSocket request");
websocket
.send(Message::Text(
json!({
"type": "response.created",
"response": {"id": response_id, "model": "gpt-5.4"}
})
.to_string()
.into(),
))
.await
.expect("send structural WebSocket event");
.expect("request")
.expect("valid frame");
websocket
.close(Some(CloseFrame {
code: CloseCode::Normal,
code: CloseCode::from(code),
reason: "".into(),
}))
.await
.expect("close WebSocket before terminal event");
.expect("close websocket");
}
for response_id in [
"resp_third_reset",
"resp_fourth_reset",
"resp_fifth_reset",
"resp_sixth_reset",
] {
let (mut http, _) = listener
.accept()
.await
.expect("accept cooldown HTTP request");
let request = capture_http_request(&mut http).await;
let request_head = String::from_utf8_lossy(&request).to_ascii_lowercase();
assert!(!request_head.contains("upgrade: websocket"));
http.write_all(
format!(
"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{CAPTURE_COMPLETED_SSE}",
CAPTURE_COMPLETED_SSE.len()
)
.as_bytes(),
)
.await
.expect("write cooldown HTTP response");
let (stream, _) = listener
.accept()
.await
.expect("accept failing half-open WebSocket connection");
let mut websocket = accept_codex_test_websocket(stream).await;
let _request = websocket
.next()
.await
.expect("half-open WebSocket request")
.expect("valid half-open WebSocket request");
websocket
.send(Message::Text(
json!({
"type": "response.created",
"response": {"id": response_id, "model": "gpt-5.4"}
})
.to_string()
.into(),
))
.await
.expect("send half-open structural event");
websocket
.close(Some(CloseFrame {
code: CloseCode::Normal,
reason: "".into(),
}))
.await
.expect("close half-open WebSocket before terminal event");
}
for _ in 0..2 {
let (mut http, _) = listener
.accept()
.await
.expect("accept final cooldown HTTP request");
let request = capture_http_request(&mut http).await;
let request_head = String::from_utf8_lossy(&request).to_ascii_lowercase();
assert!(!request_head.contains("upgrade: websocket"));
http.write_all(
format!(
"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{CAPTURE_COMPLETED_SSE}",
CAPTURE_COMPLETED_SSE.len()
)
.as_bytes(),
)
.await
.expect("write final cooldown HTTP response");
}
for response_id in ["resp_half_open", "resp_recovered"] {
let (stream, _) = listener
.accept()
.await
.expect("accept recovery WebSocket connection");
let mut websocket = accept_codex_test_websocket(stream).await;
let _request = websocket
.next()
.await
.expect("recovery WebSocket request")
.expect("valid recovery WebSocket request");
websocket
.send(Message::Text(
json!({
"type": "response.created",
"response": {"id": response_id, "model": "gpt-5.4"}
})
.to_string()
.into(),
))
.await
.expect("start recovery WebSocket response");
websocket
.send(completed_response(response_id))
.await
.expect("complete recovery WebSocket request");
for status in [200, 503] {
let (mut http, _) = listener.accept().await.expect("accept HTTP fallback");
assert!(
String::from_utf8_lossy(&capture_http_request(&mut http).await)
.starts_with("POST /codex/responses")
);
let body = if status == 200 {
CAPTURE_COMPLETED_SSE
} else {
r#"{"error":{"code":"server_error","message":"unavailable"}}"#
};
http.write_all(format!("HTTP/1.1 {status} Response\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.expect("HTTP response");
}
});
let provider = provider_with_base_url(&store, base_url);
let operation = |thread_id: &str| {
Operation::Generate(generate_with_persisted_session_context(
ACCOUNT_ID,
CONVERSATION_ID,
"websocket-cooldown-session",
thread_id,
))
};
let mut first = provider
.execute(
planned_request("openai", operation("thread-first-reset")),
context("req_first_reset", CancellationToken::new()),
)
.await
.expect("prepare first WebSocket stream");
let first_error = loop {
match first.next().await {
Some(Ok(_)) => {}
Some(Err(error)) => break error,
None => panic!("first WebSocket reset must surface a provider error"),
}
};
assert_eq!(first_error.send_state(), UpstreamSendState::Ambiguous);
assert_eq!(first_error.pre_delivery_retry(), None);
drop(first);
let mut fresh = provider
.execute(
planned_request("openai", operation("thread-fresh-reset")),
context("req_fresh_reset", CancellationToken::new()),
)
.await
.expect("prepare fresh WebSocket probe");
assert_eq!(fresh.metadata().transport().as_str(), "websocket");
let mut fresh_pool_observation = None;
let second_error = loop {
match fresh.next().await {
Some(Ok(event)) => {
if let Some(observation) = event.response_observation()
&& observation.transport().as_str() == "websocket"
{
fresh_pool_observation = Some(observation.websocket_pool());
}
}
Some(Err(error)) => break error,
None => panic!("fresh WebSocket reset must surface a provider error"),
}
};
assert_eq!(fresh_pool_observation, Some(None));
assert_eq!(second_error.send_state(), UpstreamSendState::Ambiguous);
assert_eq!(second_error.pre_delivery_retry(), None);
drop(fresh);
tokio::time::pause();
for (index, cooldown) in [
Duration::from_secs(5 * 60),
Duration::from_secs(30 * 60),
Duration::from_secs(60 * 60),
Duration::from_secs(4 * 60 * 60),
]
.into_iter()
.enumerate()
{
tokio::time::advance(cooldown - Duration::from_secs(60)).await;
let mut active_cooldown = provider
let provider = provider_with_base_url_and_retry_budget(&store, base_url, 2);
for index in 0..3 {
let operation = Operation::Generate(generate_with_persisted_session_context(
"acct_websocket_close",
"conversation-repeated-close",
"repeated-close",
"turn",
));
let mut stream = provider
.execute(
planned_request(
"openai",
operation(&format!("thread-active-cooldown-{index}")),
),
planned_request("openai", operation),
context(
&format!("req_active_cooldown_{index}"),
&format!("req_repeated_close_{index}"),
CancellationToken::new(),
),
)
.await
.expect("prepare active cooldown HTTP stream");
assert_eq!(active_cooldown.metadata().transport().as_str(), "http_sse");
while let Some(event) = active_cooldown.next().await {
event.expect("active cooldown HTTP response");
}
tokio::time::advance(Duration::from_secs(60)).await;
tokio::time::resume();
let mut failed_probe = provider
.execute(
planned_request("openai", operation(&format!("thread-failed-probe-{index}"))),
context(
&format!("req_failed_probe_{index}"),
CancellationToken::new(),
),
)
.await
.expect("prepare half-open WebSocket probe");
assert_eq!(failed_probe.metadata().transport().as_str(), "websocket");
let mut failed_probe_pool_observation = None;
let probe_error = loop {
match failed_probe.next().await {
Some(Ok(event)) => {
if let Some(observation) = event.response_observation()
&& observation.transport().as_str() == "websocket"
{
failed_probe_pool_observation = Some(observation.websocket_pool());
}
}
.expect("prepare stream");
assert_eq!(stream.metadata().transport().as_str(), "websocket");
let error = loop {
match stream.next().await {
Some(Ok(_)) => {}
Some(Err(error)) => break error,
None => panic!("half-open WebSocket reset must surface a provider error"),
None => panic!("close must fail"),
}
};
assert_eq!(failed_probe_pool_observation, Some(None));
assert_eq!(probe_error.send_state(), UpstreamSendState::Ambiguous);
assert_eq!(probe_error.pre_delivery_retry(), None);
drop(failed_probe);
tokio::time::pause();
assert_eq!(error.send_state(), UpstreamSendState::Ambiguous);
assert_eq!(error.pre_delivery_retry(), None);
assert!(!error.replay_is_safe());
}
tokio::time::advance(Duration::from_secs(4 * 60 * 60 - 60)).await;
let mut capped_cooldown = provider
.execute(
planned_request("openai", operation("thread-capped-cooldown")),
context("req_capped_cooldown", CancellationToken::new()),
)
.await
.expect("prepare capped cooldown HTTP stream");
assert_eq!(capped_cooldown.metadata().transport().as_str(), "http_sse");
while let Some(event) = capped_cooldown.next().await {
event.expect("capped cooldown HTTP response");
}
tokio::time::advance(Duration::from_secs(60)).await;
tokio::time::resume();
let mut half_open = provider
.execute(
planned_request("openai", operation("thread-half-open")),
context("req_half_open", CancellationToken::new()),
)
.await
.expect("prepare half-open WebSocket probe");
assert_eq!(half_open.metadata().transport().as_str(), "websocket");
let mut concurrent = provider
.execute(
planned_request("openai", operation("thread-probe-busy")),
context("req_probe_busy", CancellationToken::new()),
)
.await
.expect("prepare concurrent HTTP stream");
assert_eq!(concurrent.metadata().transport().as_str(), "http_sse");
while let Some(event) = concurrent.next().await {
event.expect("probe-busy HTTP response");
}
let mut half_open_pool_observation = None;
let mut half_open_completed = false;
while let Some(event) = half_open.next().await {
let event = event.expect("half-open WebSocket response");
half_open_completed |= event
.canonical_facts()
.iter()
.any(|event| matches!(event, GatewayEvent::Completed(_)));
if let Some(observation) = event.response_observation()
&& observation.transport().as_str() == "websocket"
{
half_open_pool_observation = Some(observation.websocket_pool());
for index in 0..2 {
let operation = Operation::Generate(generate_with_persisted_session_context(
"acct_websocket_close",
"conversation-repeated-close",
"repeated-close",
"turn",
));
let mut stream = provider
.execute(
planned_request("openai", operation),
context(
&format!("req_http_after_ws_{index}"),
CancellationToken::new(),
),
)
.await
.expect("prepare HTTP");
assert_eq!(stream.metadata().transport().as_str(), "http_sse");
let mut failure = None;
while let Some(event) = stream.next().await {
if let Err(error) = event {
failure = Some(error);
break;
}
}
if index == 0 {
assert!(failure.is_none());
} else {
let failure = failure.expect("HTTP failure is final transport result");
assert_eq!(failure.upstream_status(), Some(503));
assert!(!matches!(
failure.pre_delivery_retry(),
Some(
PreDeliveryRetry::SameAccountTransportRetry { .. }
| PreDeliveryRetry::SameAccountTransportFallback
)
));
}
}
assert!(half_open_completed);
assert_eq!(half_open_pool_observation, Some(None));
let mut recovered = provider
.execute(
planned_request("openai", operation("thread-recovered")),
context("req_recovered", CancellationToken::new()),
)
.await
.expect("prepare recovered WebSocket stream");
assert_eq!(recovered.metadata().transport().as_str(), "websocket");
let mut recovered_pool_observation = None;
while let Some(event) = recovered.next().await {
let event = event.expect("recovered WebSocket response");
if let Some(observation) = event.response_observation()
&& observation.transport().as_str() == "websocket"
{
recovered_pool_observation = Some(observation.websocket_pool());
}
}
assert_eq!(
recovered_pool_observation,
Some(Some(WebSocketPoolKind::New))
);
server.await.expect("WebSocket and HTTP recovery server");
server.await.expect("server");
}
#[tokio::test]
@@ -2452,6 +2212,45 @@ async fn websocket_upgrade_required_immediately_enables_session_http_fallback()
while let Some(event) = second.next().await {
event.expect("HTTP fallback response");
}
tokio::time::pause();
for index in 0..3 {
tokio::time::advance(Duration::from_secs(5 * 60 * 60)).await;
let stream = provider
.execute(
planned_request("openai", operation()),
context(
&format!("req_active_http_session_{index}"),
CancellationToken::new(),
),
)
.await
.expect("prepare active session");
assert_eq!(stream.metadata().transport().as_str(), "http_sse");
}
let unrelated = provider
.execute(
planned_request(
"openai",
Operation::Generate(generate_with_session_context(
"unrelated-session",
None,
None,
)),
),
context("req_unrelated_session", CancellationToken::new()),
)
.await
.expect("prepare unrelated session");
assert_eq!(unrelated.metadata().transport().as_str(), "websocket");
tokio::time::advance(Duration::from_secs(8 * 60 * 60)).await;
let expired = provider
.execute(
planned_request("openai", operation()),
context("req_expired_http_session", CancellationToken::new()),
)
.await
.expect("prepare expired session");
assert_eq!(expired.metadata().transport().as_str(), "websocket");
server.await.expect("upstream server");
}
@@ -2572,7 +2371,7 @@ async fn websocket_turn_state_metadata_is_exposed_through_response_observation()
}
#[tokio::test]
async fn websocket_turn_state_metadata_close_requests_checkpointed_retry() {
async fn websocket_turn_state_metadata_close_does_not_authorize_replay() {
const ACCOUNT_ID: &str = "acct_websocket_metadata_close";
const TURN_STATE: &str = "turn-state-before-close";
@@ -2630,7 +2429,7 @@ async fn websocket_turn_state_metadata_close_requests_checkpointed_retry() {
.await
.expect("prepare WebSocket provider stream");
let mut observed_turn_state = false;
let mut replay_checkpoint = None;
let mut session_update = None;
let error = loop {
match stream.next().await {
Some(Ok(event)) => {
@@ -2641,7 +2440,7 @@ async fn websocket_turn_state_metadata_close_requests_checkpointed_retry() {
});
}
if let Some(update) = event.session_update() {
replay_checkpoint = Some(update.clone());
session_update = Some(update.clone());
}
}
Some(Err(error)) => break error,
@@ -2651,22 +2450,10 @@ async fn websocket_turn_state_metadata_close_requests_checkpointed_retry() {
server.await.expect("WebSocket server");
assert!(observed_turn_state);
let replay_checkpoint = replay_checkpoint.expect("metadata close checkpoint");
assert_eq!(
replay_checkpoint.payload().get("turn_state"),
Some(&json!(TURN_STATE))
);
assert_eq!(
replay_checkpoint.payload().get("retry_checkpoint"),
Some(&json!(true))
);
assert!(session_update.is_none());
assert_eq!(error.send_state(), UpstreamSendState::Ambiguous);
assert!(error.replay_is_safe());
assert!(matches!(
error.pre_delivery_retry(),
Some(PreDeliveryRetry::SameAccountTransportRetry { retry_index, .. })
if retry_index.get() == 1
));
assert!(!error.replay_is_safe());
assert_eq!(error.pre_delivery_retry(), None);
assert_eq!(
error.upstream_code().map(|code| code.as_str()),
Some("websocket_close_1000")
@@ -3114,7 +2901,6 @@ async fn matching_turn_id_should_restore_previous_turn_state() {
Some("turn-same"),
Some("turn-same"),
None,
false,
)
.await;
@@ -3124,17 +2910,6 @@ async fn matching_turn_id_should_restore_previous_turn_state() {
);
}
#[tokio::test]
async fn retry_checkpoint_should_restore_turn_state_without_client_turn_id() {
let request =
capture_turn_state_request("req_restore_retry_checkpoint", None, None, None, true).await;
assert_eq!(
captured_header_values(&request, "x-codex-turn-state"),
vec![b"previous-turn-state".to_vec()]
);
}
#[tokio::test]
async fn matching_turn_id_should_prefer_an_explicit_client_echo_over_saved_provider_state() {
let request = capture_turn_state_request(
@@ -3142,7 +2917,6 @@ async fn matching_turn_id_should_prefer_an_explicit_client_echo_over_saved_provi
Some("turn-same"),
Some("turn-same"),
Some("client-turn-state"),
false,
)
.await;
@@ -3174,7 +2948,6 @@ async fn new_or_unidentified_turn_should_not_restore_previous_turn_state() {
previous_turn_id,
current_turn_id,
client_turn_state,
false,
)
.await;
assert!(captured_header_values(&request, "x-codex-turn-state").is_empty());
@@ -5217,3 +4990,141 @@ async fn provider_routes_catalog_model_when_supported_in_api_is_false() {
.is_some()
);
}
#[tokio::test]
async fn completed_websocket_response_resets_consecutive_failure_budget() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_websocket_close").await;
let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener");
let base_url = format!("http://{}", listener.local_addr().expect("address"));
let server = tokio::spawn(async move {
let mut active = None;
for succeeds in [false, true, false, false, true] {
if active.is_none() {
let (stream, _) = listener.accept().await.expect("accept WS");
active = Some(accept_codex_test_websocket(stream).await);
}
let ws = active.as_mut().expect("active connection");
ws.next().await.expect("request").expect("valid frame");
if succeeds {
ws.send(Message::Text(json!({"type":"response.created","response":{"id":"resp_budget_reset","model":"gpt-5.4"}}).to_string().into())).await.expect("created response");
ws.send(Message::Text(json!({"type":"response.completed","response":{"id":"resp_budget_reset","model":"gpt-5.4","status":"completed","output":[]}}).to_string().into())).await.expect("complete response");
} else {
ws.close(None).await.expect("close connection");
active = None;
}
}
});
let provider = provider_with_base_url_and_retry_budget(&store, base_url, 2);
for (index, succeeds) in [false, true, false, false, true].into_iter().enumerate() {
let operation = Operation::Generate(generate_with_persisted_session_context(
"acct_websocket_close",
"conversation-budget-reset",
"budget-reset",
"turn",
));
let mut stream = provider
.execute(
planned_request("openai", operation),
context(
&format!("req_reset_budget_{index}"),
CancellationToken::new(),
),
)
.await
.expect("prepare request");
assert_eq!(
stream.metadata().transport().as_str(),
"websocket",
"request {index}"
);
let mut failed = false;
let mut completed = false;
while let Some(event) = stream.next().await {
if let Ok(event) = &event {
completed |= event
.canonical_facts()
.iter()
.any(|fact| matches!(fact, GatewayEvent::Completed(_)));
}
if event.is_err() {
failed = true;
break;
}
}
assert_eq!(failed, !succeeds);
assert_eq!(completed, succeeds);
}
server.await.expect("server");
}
#[tokio::test]
async fn connection_limit_rejection_requests_one_provider_retry_and_reconnects_in_pool() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_websocket_close").await;
let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener");
let base_url = format!("http://{}", listener.local_addr().expect("address"));
let server = tokio::spawn(async move {
for first in [true, false] {
let (stream, _) = listener.accept().await.expect("accept WS");
let mut ws = accept_codex_test_websocket(stream).await;
ws.next().await.expect("request").expect("valid frame");
let event = if first {
json!({"type":"error","status":400,"error":{"code":"websocket_connection_limit_reached","type":"invalid_request_error","message":"connection expired"}})
} else {
json!({"type":"response.completed","response":{"id":"resp_reconnected","model":"gpt-5.4","status":"completed","output":[]}})
};
if !first {
ws.send(Message::Text(json!({"type":"response.created","response":{"id":"resp_reconnected","model":"gpt-5.4"}}).to_string().into())).await.expect("created response");
}
ws.send(Message::Text(event.to_string().into()))
.await
.expect("send response");
}
});
let provider = provider_with_base_url_and_retry_budget(&store, base_url, 1);
let operation = || {
Operation::Generate(generate_with_persisted_session_context(
"acct_websocket_close",
"conversation-rejected",
"rejected",
"turn",
))
};
let mut stream = provider
.execute(
planned_request("openai", operation()),
context("req_rejected", CancellationToken::new()),
)
.await
.expect("prepare request");
let error = loop {
match stream.next().await {
Some(Ok(_)) => {}
Some(Err(error)) => break error,
None => panic!("expected rejection"),
}
};
assert_eq!(error.send_state(), UpstreamSendState::Sent);
assert!(error.replay_is_safe());
assert!(
matches!(error.pre_delivery_retry(), Some(PreDeliveryRetry::SameAccountTransportRetry { retry_index, .. }) if retry_index.get() == 1)
);
drop(stream);
let retry_context = context("req_rejected", CancellationToken::new()).with_transport(
AttemptTransport::Retry(NonZeroU32::new(1).expect("retry index")),
);
let mut retry = provider
.execute(planned_request("openai", operation()), retry_context)
.await
.expect("prepare retry");
let mut pooled = false;
while let Some(event) = retry.next().await {
let event = event.expect("retry succeeds");
pooled |= event.response_observation().is_some_and(|observation| {
observation.websocket_pool() == Some(WebSocketPoolKind::New)
});
}
assert!(pooled);
server.await.expect("server");
}
@@ -59,7 +59,7 @@ async fn record_websocket_origin_failures(
let request_id = format!("req_{scenario}_{attempt}");
let error = backend
.create_response(
&explicit_websocket_warmup_request(&conversation_id),
&new_chain_request(&conversation_id),
request_context(&request_id, Some("chatgpt-account")),
)
.await
@@ -998,7 +998,7 @@ async fn concurrent_same_key_should_singleflight_websocket_opening() {
}
#[tokio::test]
async fn websocket_failure_before_first_delivery_should_use_same_account_http() {
async fn websocket_failure_before_first_delivery_should_not_replay_payload() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
@@ -1061,18 +1061,12 @@ async fn websocket_failure_before_first_delivery_should_use_same_account_http()
.await
.unwrap();
let (mut http, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut http).await;
assert!(request.starts_with("POST /codex/responses HTTP/1.1"));
assert_eq!(
read_header_value(&request, "authorization"),
Some("Bearer access-token")
assert!(
timeout(Duration::from_millis(100), listener.accept())
.await
.is_err(),
"transport must not replay over HTTP or WS"
);
assert_eq!(
read_header_value(&request, "chatgpt-account-id"),
Some("chatgpt-account")
);
write_completed_sse_response(&mut http).await;
});
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
@@ -1081,22 +1075,16 @@ async fn websocket_failure_before_first_delivery_should_use_same_account_http()
)
.with_websocket_pool(Arc::new(CodexWebSocketPool::default()));
let response = backend
let error = backend
.create_response(
&new_chain_request("conversation-pre-delivery-fallback"),
request_context("req_pre_delivery_fallback", Some("chatgpt-account")),
)
.await
.expect("pre-delivery WebSocket failure should use same-account HTTP");
.expect_err("sent payload must not be replayed");
server.await.unwrap();
assert_eq!(response.transport, CodexBackendTransport::HttpSse);
assert_eq!(
response.transport_metrics.decision,
Some(CodexTransportDecision::Http2PreDeliveryFailure)
);
assert!(response.body.contains("resp_http_fallback"));
assert!(!response.body.contains("resp_abandoned_websocket"));
assert!(matches!(error, CodexClientError::WebSocket(_)));
}
#[tokio::test]
@@ -1479,3 +1467,63 @@ async fn half_open_upstream_response_should_close_origin_breaker() {
assert!(response.body.contains("resp_after_half_open_success"));
pool.shutdown().await;
}
#[tokio::test]
async fn open_fast_path_breaker_should_allow_required_warmup_and_hot_reuse() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let threshold = WebSocketOriginBreakerConfig::default().failure_threshold;
let server = tokio::spawn(async move {
reject_websocket_openings(&listener, threshold).await;
let (mut http, _) = listener.accept().await.unwrap();
assert!(
read_http_request(&mut http)
.await
.starts_with("POST /codex/responses")
);
write_completed_sse_response(&mut http).await;
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
for id in ["resp_required_warmup", "resp_hot_after_breaker"] {
websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(completed_websocket_response(id, 2, 1).into()))
.await
.unwrap();
}
});
let pool = Arc::new(CodexWebSocketPool::default());
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
test_wire_profile(),
)
.with_websocket_pool(Arc::clone(&pool));
record_websocket_origin_failures(&backend, threshold, "required-warmup").await;
let fallback = backend
.create_response(
&new_chain_request("ordinary-cold"),
request_context("req_cold_suppressed", Some("chatgpt-account")),
)
.await
.expect("open fast path selects HTTP");
assert_eq!(fallback.transport, CodexBackendTransport::HttpSse);
let warmup = backend
.create_response(
&explicit_websocket_warmup_request("required-warmup"),
request_context("req_required_warmup", Some("chatgpt-account")),
)
.await
.expect("required WS bypasses optional fast path breaker");
assert_eq!(warmup.transport, CodexBackendTransport::WebSocket);
let hot = backend
.create_response(
&new_chain_request("required-warmup"),
request_context("req_hot_after_breaker", Some("chatgpt-account")),
)
.await
.expect("hot socket remains reusable");
assert_eq!(hot.websocket_pool_decision.unwrap().kind(), "reuse");
server.await.unwrap();
pool.shutdown().await;
}
@@ -506,7 +506,7 @@ fn websocket_pool_config_for_tests(
ping_interval,
ping_timeout: Duration::from_secs(1),
liveness_timeout,
initial_event_timeout: None,
stream_idle_timeout: None,
}
}
@@ -289,30 +289,11 @@ fn new_chain_should_allow_pre_delivery_http_fallback() {
request.use_websocket = true;
assert!(
transport_requirement(&request).allows_pre_delivery_http_fallback(),
"new chains may follow the TS same-account fallback behavior"
transport_requirement(&request).allows_connection_restart(),
"new chains have no connection-bound state to lose on an explicit rejection"
);
}
#[test]
fn session_transport_recovery_should_only_apply_to_new_chains() {
assert!(TransportRequirement::NewChain.allows_session_transport_recovery());
for requirement in [
TransportRequirement::HttpRequired,
TransportRequirement::ExplicitWebSocketWarmup,
TransportRequirement::ExactWebSocketContinuation,
TransportRequirement::PersistedContinuation,
TransportRequirement::ExternalUnknown,
] {
assert!(
!requirement.allows_session_transport_recovery(),
"session recovery must not override {}",
requirement.as_str()
);
}
}
#[test]
fn transport_should_mark_unknown_previous_response_as_external() {
let mut request = codex_request("gpt-test", "be brief", Vec::new());
@@ -331,7 +312,7 @@ fn previous_response_should_forbid_pre_delivery_http_fallback() {
request.set_previous_response_id(Some("resp_previous".to_owned()));
assert!(
!transport_requirement(&request).allows_pre_delivery_http_fallback(),
!transport_requirement(&request).allows_connection_restart(),
"a sent continuation must not be replayed on another transport"
);
}
@@ -161,7 +161,7 @@ fn websocket_event_to_sse_should_forward_public_events_and_strip_internal_events
}
#[test]
fn websocket_event_to_sse_should_add_missing_previous_response_recovery_code() {
fn websocket_event_to_sse_should_preserve_errors_without_inventing_recovery_codes() {
let event = json!({
"type": "error",
"status": 400,
@@ -174,7 +174,7 @@ fn websocket_event_to_sse_should_add_missing_previous_response_recovery_code() {
let frame = websocket_event_to_sse_frame(&event).expect("public error event");
assert!(frame.contains(r#""code":"previous_response_not_found""#));
assert_eq!(frame, format!("event: error\ndata: {event}\n\n"));
}
#[test]
@@ -1392,12 +1392,12 @@ async fn codex_backend_client_should_timeout_when_upstream_is_silent() {
message,
..
})
if message.contains("20s")
if message.contains("300s")
);
}
#[tokio::test]
async fn websocket_response_created_should_switch_to_active_stream_idle_timeout() {
async fn websocket_stream_should_allow_silence_below_idle_timeout() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let (created_tx, created_rx) = tokio::sync::oneshot::channel();
@@ -1437,7 +1437,7 @@ async fn websocket_response_created_should_switch_to_active_stream_idle_timeout(
let response = response_task
.await
.expect("websocket task should finish")
.expect("structural activity should disable the initial timeout");
.expect("structural events and model output share the same idle timeout");
server.await.unwrap();
assert!(response.body.contains("event: response.created"));
@@ -1445,60 +1445,32 @@ async fn websocket_response_created_should_switch_to_active_stream_idle_timeout(
}
#[tokio::test]
async fn websocket_connection_limit_error_should_retry_on_fresh_connection() {
async fn websocket_connection_limit_should_be_reported_without_transport_retry() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let (attempts_tx, mut attempts_rx) = tokio::sync::mpsc::channel(2);
let server = tokio::spawn(async move {
for attempt in 0..2 {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
attempts_tx.send(attempt).await.unwrap();
if attempt == 0 {
websocket
.send(Message::Text(
json!({
"type": "error",
"status": 400,
"error": {
"type": "invalid_request_error",
"code": "websocket_connection_limit_reached",
"message": "Responses websocket connection limit reached (60 minutes). Create a new websocket connection to continue."
}
})
.to_string()
.into(),
))
.await
.unwrap();
} else {
websocket
.send(Message::Text(
completed_websocket_response("resp_retry_fresh", 3, 1).into(),
))
.await
.unwrap();
}
websocket.close(None).await.unwrap();
}
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
websocket.next().await.unwrap().unwrap();
websocket.send(Message::Text(json!({
"type":"error", "status":400,
"error":{"type":"invalid_request_error", "code":"websocket_connection_limit_reached", "message":"connection expired"}
}).to_string().into())).await.unwrap();
assert!(
timeout(Duration::from_millis(100), listener.accept())
.await
.is_err()
);
});
let prepared = prepared_websocket_request(&format!("http://{addr}"));
let response = execute_response_create_request(&prepared)
let error = execute_response_create_request(&prepared)
.await
.expect("connection limit should retry on a fresh connection");
.expect_err("transport reports rejection");
assert!(matches!(
error,
CodexClientError::WebSocket(CodexWebSocketExchangeError::ConnectionLimitReached)
));
server.await.unwrap();
let mut attempts = Vec::new();
while let Some(attempt) = attempts_rx.recv().await {
attempts.push(attempt);
}
assert_eq!(
attempts,
vec![0, 1],
"client should retry once on a new connection"
);
assert!(response.body.contains("resp_retry_fresh"));
}
#[tokio::test]
@@ -401,7 +401,7 @@ async fn codex_backend_client_stream_should_keep_fresh_socket_after_structural_a
first_websocket.close(None).await.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig {
initial_event_timeout: Some(Duration::from_millis(30)),
stream_idle_timeout: Some(Duration::from_millis(200)),
..websocket_pool_config_for_tests(None, None, None)
}));
let backend = CodexBackendClient::new(
@@ -490,7 +490,7 @@ async fn codex_backend_client_stream_should_keep_reused_socket_after_structural_
first_websocket.close(None).await.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig {
initial_event_timeout: Some(Duration::from_millis(30)),
stream_idle_timeout: Some(Duration::from_millis(200)),
..websocket_pool_config_for_tests(None, None, None)
}));
let backend = CodexBackendClient::new(
@@ -548,7 +548,7 @@ async fn codex_backend_client_stream_should_use_http_when_connecting_limit_is_ex
});
let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig {
max_connecting: 0,
initial_event_timeout: Some(Duration::from_millis(30)),
stream_idle_timeout: Some(Duration::from_millis(200)),
..websocket_pool_config_for_tests(None, None, None)
}));
let backend = CodexBackendClient::new(