starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-06 08:06:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(openai): recover websocket sessions after ambiguous failures

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-09-03 23:33:26 +08:00
co-authored by Codex
parent a709acb149
commit 4220f4e095
10 changed files with 1086 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/providers/openai/src/credential/mod.rs
+2 -2
View File
@@ -16,8 +16,8 @@ pub mod token_client;
mod types;
pub(crate) use affinity::{
CODEX_ROOT_SESSION_TTL, CodexSessionAffinity, derive_codex_cyber_policy_session_key,
derive_codex_session_affinity, derive_previous_response_id_hash,
CodexSessionAffinity, derive_codex_cyber_policy_session_key, derive_codex_session_affinity,
derive_previous_response_id_hash,
};
pub(crate) use oauth::oauth_owner_ref;
pub(crate) use types::parse_access_token_expiration;
backend/crates/providers/openai/src/provider/execution.rs
+47 -5
View File
@@ -147,6 +147,7 @@ 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>,
@@ -155,7 +156,8 @@ pub(super) struct ColdResponse {
pub(super) output_started_at: Instant,
pub(super) session_affinity_key: Option<ProviderSessionAffinityKey>,
pub(super) session_affinity_key_hash: Option<String>,
pub(super) session_transport_fallbacks: CodexSessionTransportFallbacks,
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>,
@@ -292,6 +294,7 @@ 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);
@@ -304,6 +307,7 @@ pub(super) async fn create_response_attempt(
request,
request_context,
Some(account_id),
connection_preference,
) => response.map_err(CodexHandshakeAttemptError::Client),
}
}
@@ -472,6 +476,7 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
request,
upstream_model,
transport_policy,
websocket_connection_preference,
context,
selector,
quota,
@@ -480,7 +485,8 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
output_started_at,
session_affinity_key,
session_affinity_key_hash,
session_transport_fallbacks,
session_transport_recovery,
mut session_recovery_probe,
websocket_retry_count,
stream_max_retries,
mut session_capture,
@@ -530,6 +536,7 @@ 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 {
@@ -563,7 +570,8 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
account_id: active_account.id().as_str(),
session_affinity_key: session_affinity_key.as_ref(),
session_affinity_key_hash: session_affinity_key_hash.as_deref(),
session_transport_fallbacks: &session_transport_fallbacks,
session_transport_recovery: &session_transport_recovery,
session_recovery_probe: &mut session_recovery_probe,
},
);
}
@@ -701,7 +709,10 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
Err(mut failure) => {
if failure.websocket_transport_retryable
&& response_transport == CodexBackendTransport::WebSocket
&& !pre_commit_events.is_committed()
&& (matches!(
failure.error.send_state(),
UpstreamSendState::Sent | UpstreamSendState::Ambiguous
) || !pre_commit_events.is_committed())
{
apply_websocket_recovery_policy(
&mut failure,
@@ -715,7 +726,8 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
account_id: active_account.id().as_str(),
session_affinity_key: session_affinity_key.as_ref(),
session_affinity_key_hash: session_affinity_key_hash.as_deref(),
session_transport_fallbacks: &session_transport_fallbacks,
session_transport_recovery: &session_transport_recovery,
session_recovery_probe: &mut session_recovery_probe,
},
);
}
@@ -808,6 +820,21 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
.any(|event| matches!(event, GatewayEvent::Completed(_)));
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()
{
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"
);
}
if allows_account_state_mutation && (completed || terminal_failure.is_some()) {
synchronize_passive_quota(
&quota,
@@ -939,6 +966,21 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
.any(|event| matches!(event, GatewayEvent::Completed(_)));
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()
{
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"
);
}
if allows_account_state_mutation && completed && terminal_failure.is_none() {
// 同上：尾部 finish() 也可能产出 completed，亲和记录必须先于任何下游 yield。
selector
backend/crates/providers/openai/src/provider/failure.rs
+42 -13
View File
@@ -458,7 +458,8 @@ pub(super) struct WebSocketRecoveryContext<'a> {
pub(super) account_id: &'a str,
pub(super) session_affinity_key: Option<&'a ProviderSessionAffinityKey>,
pub(super) session_affinity_key_hash: Option<&'a str>,
pub(super) session_transport_fallbacks: &'a CodexSessionTransportFallbacks,
pub(super) session_transport_recovery: &'a CodexSessionTransportRecovery,
pub(super) session_recovery_probe: &'a mut Option<CodexSessionWebSocketProbe>,
}
pub(super) fn websocket_client_failure_policy(
@@ -509,10 +510,18 @@ pub(super) fn apply_websocket_recovery_policy(
UpstreamSendState::Sent | UpstreamSendState::Ambiguous
) {
// payload 已经发送或发送结果不确定时，代理不能用同一业务输入自动重放。
// 只为客户端下一次完整请求记住 HTTP 偏好，当前请求保持终态失败。
let sticky_state_transition = context
.session_affinity_key
.map(|key| context.session_transport_fallbacks.disable_websocket(key));
// 当前请求保持终态失败；恢复状态只裁决客户端主动发起的下一次请求。
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));
tracing::warn!(
request_id = context.request_id,
attempt_index = context.attempt_index,
@@ -527,8 +536,12 @@ pub(super) fn apply_websocket_recovery_policy(
continuation_recovery_action = "client_replay_required",
session_affinity_present = context.session_affinity_key.is_some(),
session_affinity_key_hash = context.session_affinity_key_hash.unwrap_or(""),
sticky_http_enabled = sticky_state_transition.is_some(),
sticky_state_transition = sticky_state_transition.unwrap_or(false),
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;
@@ -576,9 +589,21 @@ pub(super) fn apply_websocket_recovery_policy(
}
failure.error.set_pre_delivery_transport_fallback();
let sticky_state_transition = context
.session_affinity_key
.map(|key| context.session_transport_fallbacks.disable_websocket(key));
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));
tracing::warn!(
request_id = context.request_id,
attempt_index = context.attempt_index,
@@ -598,9 +623,13 @@ pub(super) fn apply_websocket_recovery_policy(
},
session_affinity_present = context.session_affinity_key.is_some(),
session_affinity_key_hash = context.session_affinity_key_hash.unwrap_or(""),
sticky_http_enabled = sticky_state_transition.is_some(),
sticky_state_transition = sticky_state_transition.unwrap_or(false),
"OpenAI upstream WebSocket disabled for the current session"
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"
);
}
backend/crates/providers/openai/src/provider/mod.rs
+48 -18
View File
@@ -61,7 +61,10 @@ use crate::credential::{
derive_codex_cyber_policy_session_key, derive_codex_session_affinity,
derive_previous_response_id_hash,
};
use crate::session_transport::CodexSessionTransportFallbacks;
use crate::session_transport::{
CodexSessionRecoveryTransition, CodexSessionTransportDecision, CodexSessionTransportRecovery,
CodexSessionWebSocketProbe, SessionWebSocketFallback,
};
use crate::transport::canonical::{
CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalOutcome,
};
@@ -90,7 +93,8 @@ use crate::transport::{
CODEX_RESPONSES_PATH, CodexAccountSelectionTelemetry, CodexBackendClient,
CodexBackendJsonResponse, CodexBackendStreamingResponse, CodexBackendTransport,
CodexClientError, CodexRateLimitUpdates, CodexRequestContext, CodexResponseMetadata,
CodexTransportMetrics, CodexUpstreamDiagnostics, CodexWebSocketPool, endpoint_url,
CodexTransportMetrics, CodexUpstreamDiagnostics, CodexWebSocketPool,
WebSocketConnectionPreference, endpoint_url,
};
mod execution;
@@ -145,7 +149,7 @@ pub struct CodexProvider {
image_edits_url: Url,
search_url: Url,
session_identity: Option<CodexSessionIdentity>,
session_transport_fallbacks: CodexSessionTransportFallbacks,
session_transport_recovery: CodexSessionTransportRecovery,
stream_max_retries: u32,
}
@@ -185,7 +189,7 @@ impl CodexProvider {
image_edits_url,
search_url,
session_identity: None,
session_transport_fallbacks: CodexSessionTransportFallbacks::default(),
session_transport_recovery: CodexSessionTransportRecovery::default(),
stream_max_retries,
})
}
@@ -442,26 +446,46 @@ impl Provider for CodexProvider {
);
let requirement = transport_requirement(&upstream_request);
let requested_transport = selected_transport(&upstream_request);
let sticky_http_fallback = context.transport() == AttemptTransport::Default
&& requested_transport == CodexProviderTransport::PreferWebSocket
&& requirement.allows_sticky_http_fallback()
&& session_affinity.as_ref().is_some_and(|affinity| {
self.session_transport_fallbacks
.is_http_only(affinity.key())
});
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 transport = if requirement.requires_websocket() {
CodexProviderTransport::PreferWebSocket
} else {
match context.transport() {
AttemptTransport::Default if sticky_http_fallback => {
CodexProviderTransport::HttpOnly
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::Default => requested_transport,
AttemptTransport::Retry(_) => requested_transport,
AttemptTransport::Fallback => CodexProviderTransport::HttpOnly,
}
};
if sticky_http_fallback {
if let Some(session_transport_action) = session_transport_action {
tracing::info!(
request_id = %context.request_id(),
attempt_index = context.attempt_index().get(),
@@ -469,7 +493,11 @@ impl Provider for CodexProvider {
.as_ref()
.map_or("", CodexSessionAffinity::key_hash),
transport_requirement = requirement.as_str(),
"Reusing session-sticky OpenAI upstream HTTP fallback"
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);
@@ -518,7 +546,9 @@ impl Provider for CodexProvider {
output_started_at,
session_affinity_key,
session_affinity_key_hash,
session_transport_fallbacks: self.session_transport_fallbacks.clone(),
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,
backend/crates/providers/openai/src/session_transport.rs
+406 -53
View File
@@ -1,104 +1,457 @@
//! OpenAI 根会话级上游传输降级状态。
//! OpenAI 根会话级上游传输恢复状态。
use std::collections::HashMap;
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::{Duration, Instant};
use std::time::Duration;
use gateway_core::provider_ports::ProviderSessionAffinityKey;
use tokio::time::Instant;
use uuid::Uuid;
use crate::credential::CODEX_ROOT_SESSION_TTL;
/// 未成功恢复的失败链需要跨过最长冷却，避免 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_HTTP_FALLBACKS: usize = 16_384;
const MAX_SESSION_RECOVERIES: usize = 16_384;
#[derive(Clone)]
pub(crate) struct CodexSessionTransportFallbacks {
state: Arc<Mutex<SessionTransportFallbackState>>,
ttl: Duration,
pub(crate) struct CodexSessionTransportRecovery {
state: Arc<Mutex<SessionTransportRecoveryState>>,
http_cooldowns: [Duration; 4],
retention: Duration,
capacity: usize,
}
#[derive(Default)]
struct SessionTransportFallbackState {
http_only: HashMap<ProviderSessionAffinityKey, Instant>,
struct SessionTransportRecoveryState {
sessions: HashMap<ProviderSessionAffinityKey, SessionCircuit>,
}
impl Default for CodexSessionTransportFallbacks {
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::with_limits(CODEX_ROOT_SESSION_TTL, MAX_SESSION_HTTP_FALLBACKS)
Self {
state: Arc::new(Mutex::new(SessionTransportRecoveryState::default())),
http_cooldowns: HTTP_COOLDOWNS,
retention: SESSION_RECOVERY_RETENTION,
capacity: MAX_SESSION_RECOVERIES,
}
}
}
impl CodexSessionTransportFallbacks {
fn with_limits(ttl: Duration, capacity: usize) -> Self {
debug_assert!(!ttl.is_zero());
debug_assert!(capacity > 0);
#[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 {
state: Arc::new(Mutex::new(SessionTransportFallbackState::default())),
ttl,
capacity,
recovery,
key,
probe_id,
armed: true,
}
}
/// 返回该根会话是否已粘性切到 HTTP，并为活跃会话续期。
pub(crate) fn is_http_only(&self, key: &ProviderSessionAffinityKey) -> bool {
self.is_http_only_at(key, Instant::now())
pub(crate) fn succeed(mut self) -> bool {
let recovered = self.recovery.record_probe_success(&self.key, self.probe_id);
self.armed = false;
recovered
}
/// 在当前 WebSocket 重试预算耗尽后，为根会话启用粘性 HTTP。
///
/// 返回 `true` 表示本次发生了禁用状态转换；已禁用会话只续期。
pub(crate) fn disable_websocket(&self, key: &ProviderSessionAffinityKey) -> bool {
self.disable_websocket_at(key, Instant::now())
pub(crate) fn post_send_failed(mut self) -> CodexSessionRecoveryTransition {
let transition = self.recovery.record_post_send_failure(&self.key);
self.armed = false;
transition
}
fn is_http_only_at(&self, key: &ProviderSessionAffinityKey, now: Instant) -> bool {
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);
}
}
}
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 Some(last_seen) = state.http_only.get_mut(key) else {
return false;
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;
};
if now.saturating_duration_since(*last_seen) >= self.ttl {
state.http_only.remove(key);
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),
}
}
}
*last_seen = now;
true
}
fn disable_websocket_at(&self, key: &ProviderSessionAffinityKey, now: Instant) -> bool {
fn record_post_send_failure_at(
&self,
key: &ProviderSessionAffinityKey,
now: Instant,
) -> CodexSessionRecoveryTransition {
let mut state = self.lock();
let already_disabled = state
.http_only
.get(key)
.is_some_and(|last_seen| now.saturating_duration_since(*last_seen) < self.ttl);
if !already_disabled && !state.http_only.contains_key(key) {
let failure_count = self.next_failure_count(state.sessions.get(key), now);
if !state.sessions.contains_key(key) {
self.make_room(&mut state, now);
}
state.http_only.insert(key.clone(), now);
!already_disabled
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 make_room(&self, state: &mut SessionTransportFallbackState, now: Instant) {
if state.http_only.len() < self.capacity {
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
}
})
}
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]
}
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
}
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;
}
state
.http_only
.retain(|_, last_seen| now.saturating_duration_since(*last_seen) < self.ttl);
if state.http_only.len() < self.capacity {
let failure_count = *failure_count;
let last_failure_at = *last_failure_at;
state.sessions.insert(
key.clone(),
SessionCircuit::FreshPending {
failure_count,
last_failure_at,
},
);
}
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
.http_only
.sessions
.iter()
.min_by_key(|(_, last_seen)| **last_seen)
.min_by_key(|(_, circuit)| circuit.last_failure_at())
.map(|(key, _)| key.clone());
if let Some(oldest) = oldest {
state.http_only.remove(&oldest);
state.sessions.remove(&oldest);
}
}
fn lock(&self) -> MutexGuard<'_, SessionTransportFallbackState> {
fn lock(&self) -> MutexGuard<'_, SessionTransportRecoveryState> {
self.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
backend/crates/providers/openai/src/transport/client_sse.rs
+40 -10
View File
@@ -47,6 +47,16 @@ 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(
@@ -188,9 +198,16 @@ impl CodexBackendClient {
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, false, true)
.prepare_response_transport(
request,
context,
pool_account_id,
connection_preference,
true,
)
.await?;
self.create_response_stream_with_prepared(request, context, prepared)
.await
@@ -204,18 +221,24 @@ impl CodexBackendClient {
context: CodexRequestContext<'_>,
pool_account_id: Option<&str>,
) -> CodexClientResult<PreparedResponseTransport> {
self.prepare_response_transport(request, context, pool_account_id, false, false)
.await
self.prepare_response_transport(
request,
context,
pool_account_id,
WebSocketConnectionPreference::ReuseOrConnect,
false,
)
.await
}
/// 与 [`Self::prepare_response_transport_with_pool_account`] 相同；
/// `force_fresh` 时绕过连接池，为连接寿命限制重试强制新建连接。
/// fresh 偏好会绕过连接池，为受控恢复探针强制新建连接。
async fn prepare_response_transport(
&self,
request: &CodexResponsesRequest,
context: CodexRequestContext<'_>,
pool_account_id: Option<&str>,
force_fresh: bool,
connection_preference: WebSocketConnectionPreference,
defer_websocket_recovery: bool,
) -> CodexClientResult<PreparedResponseTransport> {
let requirement = transport_requirement(request);
@@ -256,10 +279,11 @@ 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 = if force_fresh {
None
} else {
self.websocket_pool.as_deref().zip(pool_key)
let pool = match connection_preference {
WebSocketConnectionPreference::ReuseOrConnect => {
self.websocket_pool.as_deref().zip(pool_key)
}
WebSocketConnectionPreference::Fresh => None,
};
let fast_path_budget = match requirement {
TransportRequirement::PersistedContinuation | TransportRequirement::NewChain => {
@@ -434,7 +458,13 @@ impl CodexBackendClient {
);
drop(exchange);
match self
.prepare_response_transport(request, context, None, true, false)
.prepare_response_transport(
request,
context,
None,
WebSocketConnectionPreference::Fresh,
false,
)
.await
{
Ok(PreparedResponseTransport {
backend/crates/providers/openai/src/transport/mod.rs
+1
View File
@@ -21,6 +21,7 @@ pub mod tls;
pub mod usage;
pub mod websocket;
pub(crate) use self::client_sse::WebSocketConnectionPreference;
pub use self::{
canonical::{CodexCanonicalDecoder, CodexCanonicalError},
catalog::{
backend/crates/providers/openai/src/transport/protocol/responses.rs
+2 -2
View File
@@ -162,8 +162,8 @@ impl TransportRequirement {
matches!(self, Self::NewChain)
}
/// 会话级 sticky 状态是否可以把首选 WebSocket 改为 HTTP。
pub fn allows_sticky_http_fallback(self) -> bool {
/// 会话级恢复状态是否可以为请求选择 fresh WebSocket 或临时 HTTP。
pub fn allows_session_transport_recovery(self) -> bool {
matches!(self, Self::NewChain)
}
backend/crates/providers/openai/tests/provider/contract.rs
+494 -45
View File
@@ -24,7 +24,7 @@ use gateway_core::engine::{
use gateway_core::error::{
ContinuationFailure, ContinuationRecoveryDisposition, PreDeliveryRetry, ProviderErrorKind,
};
use gateway_core::event::GatewayEvent;
use gateway_core::event::{GatewayEvent, WebSocketPoolKind};
use gateway_core::lifecycle::CancellationToken;
use gateway_core::operation::{
CapabilityRequirements, GenerateRequest, ImageRequest, ImageRequestKind, Operation,
@@ -321,6 +321,25 @@ fn generate_with_session_context(
)
}
fn generate_with_persisted_session_context(
account_id: &str,
conversation_id: &str,
session_id: &str,
thread_id: &str,
) -> GenerateRequest {
generate_with_session_context(session_id, Some(thread_id), None).with_provider_session_state(
ProviderSessionState::new(
"openai",
Map::from_iter([
("account_id".to_owned(), json!(account_id)),
("conversation_id".to_owned(), json!(conversation_id)),
("continuation_scope".to_owned(), json!("persisted")),
]),
)
.expect("provider session state"),
)
}
fn http_generate_operation() -> Operation {
let payload = ProtocolPayload::json_object(
"openai",
@@ -439,6 +458,7 @@ fn contract_account_scope() -> Arc<FrozenAccountScope> {
"acct_truncated_stream",
"acct_usage_limit_request_path",
"acct_websocket_close",
"acct_websocket_cooldown",
"acct_websocket_busy_replay",
"acct_websocket_turn_state",
]
@@ -1295,9 +1315,13 @@ async fn selection_infrastructure_errors_have_a_distinct_classification() {
}
#[tokio::test]
async fn websocket_close_after_delivery_preserves_details_without_enabling_sticky_http() {
async fn websocket_close_after_delivery_preserves_details_and_forces_a_fresh_client_retry() {
const ACCOUNT_ID: &str = "acct_websocket_close";
const CONVERSATION_ID: &str = "conversation-websocket-close-after-delivery";
const SESSION_ID: &str = "committed-websocket-session";
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_websocket_close").await;
create_account(&store, ACCOUNT_ID).await;
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind WebSocket listener");
@@ -1316,6 +1340,20 @@ async fn websocket_close_after_delivery_preserves_details_without_enabling_stick
.await
.expect("WebSocket request")
.expect("valid WebSocket request");
websocket
.send(Message::Text(
json!({
"type": "response.created",
"response": {
"id": "resp_before_committed_close",
"model": "gpt-5.4"
}
})
.to_string()
.into(),
))
.await
.expect("start WebSocket response");
websocket
.send(Message::Text(
json!({
@@ -1345,6 +1383,20 @@ async fn websocket_close_after_delivery_preserves_details_without_enabling_stick
.await
.expect("next-turn WebSocket request")
.expect("valid next-turn WebSocket request");
websocket
.send(Message::Text(
json!({
"type": "response.created",
"response": {
"id": "resp_after_committed_close",
"model": "gpt-5.4"
}
})
.to_string()
.into(),
))
.await
.expect("start next-turn WebSocket response");
websocket
.send(Message::Text(
json!({
@@ -1369,10 +1421,11 @@ async fn websocket_close_after_delivery_preserves_details_without_enabling_stick
});
let provider = provider_with_base_url(&store, base_url);
let first_operation = Operation::Generate(generate_with_session_context(
"committed-websocket-session",
Some("thread-first"),
None,
let first_operation = Operation::Generate(generate_with_persisted_session_context(
ACCOUNT_ID,
CONVERSATION_ID,
SESSION_ID,
"thread-first",
));
let mut stream = provider
.execute(
@@ -1409,10 +1462,11 @@ async fn websocket_close_after_delivery_preserves_details_without_enabling_stick
Some("websocket_close_1009")
);
let second_operation = Operation::Generate(generate_with_session_context(
"committed-websocket-session",
Some("thread-second"),
None,
let second_operation = Operation::Generate(generate_with_persisted_session_context(
ACCOUNT_ID,
CONVERSATION_ID,
SESSION_ID,
"thread-second",
));
let mut next_stream = provider
.execute(
@@ -1420,18 +1474,29 @@ async fn websocket_close_after_delivery_preserves_details_without_enabling_stick
context("req_after_committed_close", CancellationToken::new()),
)
.await
.expect("post-delivery close must not enable sticky HTTP fallback");
.expect("post-delivery close must allow a fresh WebSocket client retry");
assert_eq!(next_stream.metadata().transport().as_str(), "websocket");
let mut retry_pool_observation = None;
while let Some(event) = next_stream.next().await {
event.expect("next turn WebSocket response");
let event = event.expect("next turn WebSocket response");
if let Some(observation) = event.response_observation()
&& observation.transport().as_str() == "websocket"
{
retry_pool_observation = Some(observation.websocket_pool());
}
}
assert_eq!(retry_pool_observation, Some(None));
server.await.expect("WebSocket server");
}
#[tokio::test]
async fn ambiguous_websocket_close_only_sticks_new_chains_in_the_current_session() {
async fn ambiguous_websocket_close_probes_a_fresh_websocket_for_the_next_new_chain() {
const ACCOUNT_ID: &str = "acct_websocket_close";
const CONVERSATION_ID: &str = "conversation-websocket-ambiguous-close";
const SESSION_ID: &str = "sticky-websocket-session";
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_websocket_close").await;
create_account(&store, ACCOUNT_ID).await;
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind WebSocket listener");
@@ -1511,19 +1576,31 @@ async fn ambiguous_websocket_close_only_sticks_new_chains_in_the_current_session
.await
.expect("close WebSocket before terminal event");
let (mut http, _) = listener.accept().await.expect("accept sticky 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 sticky HTTP response");
let (stream, _) = listener
.accept()
.await
.expect("accept fresh WebSocket probe");
let mut websocket = accept_codex_test_websocket(stream).await;
let _request = websocket
.next()
.await
.expect("fresh WebSocket request")
.expect("valid fresh WebSocket request");
websocket
.send(Message::Text(
json!({
"type": "response.created",
"response": {"id": "resp_fresh_retry", "model": "gpt-5.4"}
})
.to_string()
.into(),
))
.await
.expect("start fresh WebSocket response");
websocket
.send(completed_response("resp_fresh_retry"))
.await
.expect("complete fresh WebSocket request");
let (stream, _) = listener
.accept()
@@ -1568,10 +1645,11 @@ async fn ambiguous_websocket_close_only_sticks_new_chains_in_the_current_session
let provider = provider_with_base_url_and_retry_budget(&store, base_url, 2);
let mut saw_websocket_observation = false;
let operation = Operation::Generate(generate_with_session_context(
"sticky-websocket-session",
Some("thread-first"),
None,
let operation = Operation::Generate(generate_with_persisted_session_context(
ACCOUNT_ID,
CONVERSATION_ID,
SESSION_ID,
"thread-first",
));
let mut stream = provider
.execute(
@@ -1640,23 +1718,31 @@ async fn ambiguous_websocket_close_only_sticks_new_chains_in_the_current_session
assert_eq!(connection.exit_reason(), "normal_close");
assert!(connection.age_ms() >= connection.idle_ms());
let second_operation = Operation::Generate(generate_with_session_context(
"sticky-websocket-session",
Some("thread-second"),
None,
let second_operation = Operation::Generate(generate_with_persisted_session_context(
ACCOUNT_ID,
CONVERSATION_ID,
SESSION_ID,
"thread-second",
));
let mut sticky_stream = provider
let mut fresh_stream = provider
.execute(
planned_request("openai", second_operation),
context("req_websocket_sticky_http", CancellationToken::new()),
context("req_websocket_fresh_retry", CancellationToken::new()),
)
.await
.expect("same session should prepare an HTTP stream");
assert_eq!(sticky_stream.metadata().transport().as_str(), "http_sse");
while let Some(event) = sticky_stream.next().await {
event.expect("sticky HTTP response");
.expect("same session should prepare a fresh WebSocket stream");
assert_eq!(fresh_stream.metadata().transport().as_str(), "websocket");
let mut fresh_pool_observation = None;
while let Some(event) = fresh_stream.next().await {
let event = event.expect("fresh WebSocket response");
if let Some(observation) = event.response_observation()
&& observation.transport().as_str() == "websocket"
{
fresh_pool_observation = Some(observation.websocket_pool());
}
}
drop(sticky_stream);
assert_eq!(fresh_pool_observation, Some(None));
drop(fresh_stream);
let continuation_operation = Operation::Generate(generate_with_session_context(
"sticky-websocket-session",
@@ -1669,7 +1755,7 @@ async fn ambiguous_websocket_close_only_sticks_new_chains_in_the_current_session
external_continuation_context("req_sticky_external_continuation"),
)
.await
.expect("sticky state must not force a continuation onto HTTP");
.expect("successful fresh probe must restore normal continuation routing");
assert_eq!(
continuation_stream.metadata().transport().as_str(),
"websocket"
@@ -1719,12 +1805,375 @@ async fn ambiguous_websocket_close_only_sticks_new_chains_in_the_current_session
),
)
.await
.expect("required warmup must bypass sticky HTTP fallback");
.expect("required warmup must remain on WebSocket");
assert_eq!(warmup_stream.metadata().transport().as_str(), "websocket");
drop(warmup_stream);
server.await.expect("WebSocket and HTTP server");
}
#[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";
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, ACCOUNT_ID).await;
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind WebSocket 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");
let mut websocket = accept_codex_test_websocket(stream).await;
let _request = 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");
websocket
.close(Some(CloseFrame {
code: CloseCode::Normal,
reason: "".into(),
}))
.await
.expect("close WebSocket before terminal event");
}
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");
}
});
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
.execute(
planned_request(
"openai",
operation(&format!("thread-active-cooldown-{index}")),
),
context(
&format!("req_active_cooldown_{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());
}
}
Some(Err(error)) => break error,
None => panic!("half-open WebSocket reset must surface a provider error"),
}
};
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();
}
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());
}
}
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");
}
#[tokio::test]
async fn websocket_upgrade_required_immediately_enables_session_http_fallback() {
let store = Arc::new(MemoryAccountStore::default());
backend/crates/providers/openai/tests/transport/protocol/responses.rs
+4 -4
View File
@@ -295,8 +295,8 @@ fn new_chain_should_allow_pre_delivery_http_fallback() {
}
#[test]
fn sticky_http_fallback_should_only_apply_to_new_chains() {
assert!(TransportRequirement::NewChain.allows_sticky_http_fallback());
fn session_transport_recovery_should_only_apply_to_new_chains() {
assert!(TransportRequirement::NewChain.allows_session_transport_recovery());
for requirement in [
TransportRequirement::HttpRequired,
@@ -306,8 +306,8 @@ fn sticky_http_fallback_should_only_apply_to_new_chains() {
TransportRequirement::ExternalUnknown,
] {
assert!(
!requirement.allows_sticky_http_fallback(),
"sticky HTTP must not override {}",
!requirement.allows_session_transport_recovery(),
"session recovery must not override {}",
requirement.as_str()
);
}