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): fall back after websocket fast-path misses

Browse Source 
Treat the 800 ms fast-path budget as a routing decision instead of a provider error, while preserving background preconnect for the next request.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-09-04 00:36:50 +08:00
co-authored by Codex
parent 4220f4e095
commit 024fc9fcf0
7 changed files with 256 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/providers/openai/src/provider/failure.rs
-5
View File
@@ -899,9 +899,6 @@ fn websocket_diagnostic(error: &CodexWebSocketExchangeError) -> ProviderDiagnost
CodexWebSocketExchangeError::ConnectTimeout { timeout } => {
format!("OpenAI WebSocket connection timed out after {timeout:?}")
}
CodexWebSocketExchangeError::FastPathTimeout { timeout } => {
format!("OpenAI WebSocket fast-path budget expired after {timeout:?}")
}
CodexWebSocketExchangeError::OriginCircuitOpen => {
"OpenAI WebSocket origin circuit is open".to_owned()
}
@@ -1176,7 +1173,6 @@ pub(super) fn websocket_send_state(error: &CodexWebSocketExchangeError) -> Upstr
CodexWebSocketExchangeError::InvalidRequest(_)
| CodexWebSocketExchangeError::Connect(_)
| CodexWebSocketExchangeError::ConnectTimeout { .. }
| CodexWebSocketExchangeError::FastPathTimeout { .. }
| CodexWebSocketExchangeError::OriginCircuitOpen
| CodexWebSocketExchangeError::OriginHalfOpenBusy
| CodexWebSocketExchangeError::SharedConnectFailed
@@ -1204,7 +1200,6 @@ pub(super) fn websocket_error_kind(error: &CodexWebSocketExchangeError) -> Provi
| CodexWebSocketExchangeError::InvalidSse(_)
| CodexWebSocketExchangeError::UnexpectedBinaryEvent => ProviderErrorKind::Protocol,
CodexWebSocketExchangeError::ConnectTimeout { .. }
| CodexWebSocketExchangeError::FastPathTimeout { .. }
| CodexWebSocketExchangeError::SendTimeout { .. }
| CodexWebSocketExchangeError::ReceiveIdleTimeout { .. }
| CodexWebSocketExchangeError::InitialEventTimeout { .. } => ProviderErrorKind::Timeout,
backend/crates/providers/openai/src/transport/client.rs
-3
View File
@@ -854,9 +854,6 @@ pub(super) fn http_fallback_decision(
error: &CodexWebSocketExchangeError,
) -> CodexTransportDecision {
match error.classified() {
CodexWebSocketExchangeError::FastPathTimeout { .. } => {
CodexTransportDecision::Http2WebSocketBudgetExhausted
}
CodexWebSocketExchangeError::OriginCircuitOpen
| CodexWebSocketExchangeError::OriginHalfOpenBusy => {
CodexTransportDecision::Http2BreakerOpen
backend/crates/providers/openai/src/transport/client_sse.rs
+28 -4
View File
@@ -38,9 +38,9 @@ use crate::transport::{
websocket::{
CodexWebSocketConnection, CodexWebSocketExchangeError, CodexWebSocketPool,
CodexWebSocketPoolKey, CodexWebSocketStreamingExchange, DEFAULT_INITIAL_EVENT_TIMEOUT,
WEBSOCKET_FAST_PATH_BUDGET, WebSocketOriginBreaker, WebSocketPoolDecision,
execute_prepared_response_create_request_stream, post_send_ambiguous,
prepare_response_create_request_with_pool, websocket_audit_dir,
WEBSOCKET_FAST_PATH_BUDGET, WebSocketFastPath, WebSocketOriginBreaker,
WebSocketPoolDecision, execute_prepared_response_create_request_stream,
post_send_ambiguous, prepare_response_create_request_with_pool, websocket_audit_dir,
write_websocket_audit_artifact_from_env,
},
};
@@ -306,7 +306,31 @@ impl CodexBackendClient {
)
.await;
let prepared = match prepared {
Ok(prepared) => prepared,
Ok(WebSocketFastPath::Ready(prepared)) => prepared,
Ok(WebSocketFastPath::Missed) => {
let decision = CodexTransportDecision::Http2WebSocketBudgetExhausted;
let wait_ms = elapsed_duration_millis(prepare_started_at.elapsed());
tracing::info!(
request_id = %context.request_id,
account_id = pool_account_id.or(context.account_id).unwrap_or_default(),
transport_requirement = requirement.as_str(),
transport_decision = decision.as_str(),
transport_decision_wait_ms = wait_ms,
websocket_preconnect_continues = pool_log_context.is_some(),
"WebSocket fast path missed; using same-account HTTP"
);
return Ok(PreparedResponseTransport {
requirement,
route: PreparedResponseRoute::Http,
metrics: CodexTransportMetrics {
decision: Some(decision),
ws_connect_ms: None,
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()
backend/crates/providers/openai/src/transport/websocket/coordinator.rs
+82 -54
View File
@@ -29,6 +29,14 @@ use super::{
pub(crate) const WEBSOCKET_FAST_PATH_BUDGET: Duration = Duration::from_millis(800);
/// WebSocket 快路径的控制流结果；预算未命中不代表连接失败。
pub(crate) enum WebSocketFastPath<T> {
/// 在前台预算内获得结果。
Ready(T),
/// 前台停止等待；池化 opening 可以继续在后台完成。
Missed,
}
/// 尚未发送 `response.create` 的 WebSocket。只有该类型可以安全切换到 HTTP。
pub(crate) struct PreparedWebSocket {
connection: PooledWebSocketConnection,
@@ -104,7 +112,7 @@ pub(crate) async fn prepare_response_create_request_with_pool(
fast_path_budget: Option<Duration>,
require_pool: bool,
fallback_initial_event_timeout: Option<Duration>,
) -> Result<PreparedWebSocket, CodexWebSocketExchangeError> {
) -> Result<WebSocketFastPath<PreparedWebSocket>, CodexWebSocketExchangeError> {
let decision_started_at = Instant::now();
let Some((pool, key)) = pool else {
if require_pool || !request.continuation().permits_fresh_connection() {
@@ -142,7 +150,7 @@ pub(crate) async fn prepare_response_create_request_with_pool(
PreviousResponseUnavailableReason::LatestResponseMismatch,
));
}
Ok(PreparedWebSocket {
Ok(WebSocketFastPath::Ready(PreparedWebSocket {
connection: *connection,
binding: PoolBinding::Pooled {
lease: Box::new(lease),
@@ -152,7 +160,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(),
})
}))
}
WebSocketPoolAcquire::Connect(connect_lease) => {
if !request.continuation().permits_fresh_connection() {
@@ -178,14 +186,14 @@ pub(crate) async fn prepare_response_create_request_with_pool(
PreviousResponseUnavailableReason::ConnectionBusy,
));
}
let outcome = wait_for_shared_connect(waiter, fast_path_budget).await?;
match outcome {
WebSocketPoolConnectOutcome::Ready => Err(continuation_unavailable(
PreviousResponseUnavailableReason::ConnectionBusy,
)),
WebSocketPoolConnectOutcome::Failed | WebSocketPoolConnectOutcome::Pending => {
Err(CodexWebSocketExchangeError::SharedConnectFailed)
}
match wait_for_shared_connect(waiter, fast_path_budget).await {
WebSocketFastPath::Missed => Ok(WebSocketFastPath::Missed),
WebSocketFastPath::Ready(WebSocketPoolConnectOutcome::Ready) => Err(
continuation_unavailable(PreviousResponseUnavailableReason::ConnectionBusy),
),
WebSocketFastPath::Ready(
WebSocketPoolConnectOutcome::Failed | WebSocketPoolConnectOutcome::Pending,
) => Err(CodexWebSocketExchangeError::SharedConnectFailed),
}
}
WebSocketPoolAcquire::ContinuationLost(observation) => Err(continuation_unavailable(
@@ -203,7 +211,7 @@ async fn prepare_unpooled_websocket(
fast_path_budget: Option<Duration>,
initial_event_timeout: Option<Duration>,
decision_started_at: Instant,
) -> Result<PreparedWebSocket, CodexWebSocketExchangeError> {
) -> Result<WebSocketFastPath<PreparedWebSocket>, CodexWebSocketExchangeError> {
let permit = acquire_breaker_permit(breaker, origin_key)?;
let context = pump_log_context_from_connection(request.connection());
let connected = connect_with_budget(
@@ -213,14 +221,22 @@ async fn prepare_unpooled_websocket(
context,
)
.await;
let connected = match connected {
Ok(WebSocketFastPath::Ready(connection)) => Ok(connection),
Ok(WebSocketFastPath::Missed) => {
permit.fast_timeout();
return Ok(WebSocketFastPath::Missed);
}
Err(error) => Err(error),
};
let (connection, connect_elapsed) = finish_breaker_attempt(permit, connected)?;
Ok(PreparedWebSocket {
Ok(WebSocketFastPath::Ready(PreparedWebSocket {
connection,
binding: PoolBinding::Unpooled,
connect_elapsed: Some(connect_elapsed),
decision_wait_elapsed: decision_started_at.elapsed(),
initial_event_timeout,
})
}))
}
async fn prepare_pooled_websocket(
@@ -231,7 +247,7 @@ async fn prepare_pooled_websocket(
origin_key: &str,
fast_path_budget: Option<Duration>,
decision_started_at: Instant,
) -> Result<PreparedWebSocket, CodexWebSocketExchangeError> {
) -> Result<WebSocketFastPath<PreparedWebSocket>, CodexWebSocketExchangeError> {
let permit = match acquire_breaker_permit(breaker, origin_key) {
Ok(permit) => permit,
Err(error) => {
@@ -245,18 +261,20 @@ async fn prepare_pooled_websocket(
connect_lease,
permit,
);
let handoff = waiter.wait(fast_path_budget).await?;
Ok(PreparedWebSocket {
connection: *handoff.connection,
binding: PoolBinding::Pooled {
lease: Box::new(handoff.lease),
reused: false,
decision: WebSocketPoolDecision::new(),
},
connect_elapsed: Some(handoff.connect_elapsed),
decision_wait_elapsed: decision_started_at.elapsed(),
initial_event_timeout: pool.initial_event_timeout(),
})
match waiter.wait(fast_path_budget).await? {
WebSocketFastPath::Ready(handoff) => Ok(WebSocketFastPath::Ready(PreparedWebSocket {
connection: *handoff.connection,
binding: PoolBinding::Pooled {
lease: Box::new(handoff.lease),
reused: false,
decision: WebSocketPoolDecision::new(),
},
connect_elapsed: Some(handoff.connect_elapsed),
decision_wait_elapsed: decision_started_at.elapsed(),
initial_event_timeout: pool.initial_event_timeout(),
})),
WebSocketFastPath::Missed => Ok(WebSocketFastPath::Missed),
}
}
struct PooledWebSocketConnectWaiter {
@@ -275,7 +293,7 @@ impl PooledWebSocketConnectWaiter {
async fn wait(
&mut self,
fast_path_budget: Option<Duration>,
) -> Result<PooledWebSocketHandoff, CodexWebSocketExchangeError> {
) -> Result<WebSocketFastPath<PooledWebSocketHandoff>, CodexWebSocketExchangeError> {
let received = match fast_path_budget {
Some(budget) => {
let remaining = budget.saturating_sub(self.started_at.elapsed());
@@ -283,15 +301,15 @@ impl PooledWebSocketConnectWaiter {
Ok(received) => received,
Err(_) => {
self.fast_path_reporter.missed();
return Err(CodexWebSocketExchangeError::FastPathTimeout {
timeout: budget,
});
return Ok(WebSocketFastPath::Missed);
}
}
}
None => (&mut self.receiver).await,
};
received.map_err(|_| CodexWebSocketExchangeError::SharedConnectFailed)?
received
.map_err(|_| CodexWebSocketExchangeError::SharedConnectFailed)?
.map(WebSocketFastPath::Ready)
}
}
@@ -327,7 +345,7 @@ fn start_pooled_websocket_connect(
);
return;
}
result = connect_with_budget(&connection, keepalive, None, context) => result,
result = connect_websocket_connection(&connection, keepalive, context) => result,
};
match finish_breaker_attempt(permit, connected) {
Ok((connection, connect_elapsed)) => {
@@ -401,8 +419,7 @@ fn opening_failure_reason(error: &CodexWebSocketExchangeError) -> &'static str {
CodexWebSocketExchangeError::InvalidRequest(_) => "request_build_failure",
CodexWebSocketExchangeError::Connect(error)
| CodexWebSocketExchangeError::Transport(error) => transport_metric_reason(error),
CodexWebSocketExchangeError::ConnectTimeout { .. }
| CodexWebSocketExchangeError::FastPathTimeout { .. } => "connect_timeout",
CodexWebSocketExchangeError::ConnectTimeout { .. } => "connect_timeout",
CodexWebSocketExchangeError::Upstream(_) => "upgrade_rejected",
CodexWebSocketExchangeError::OriginCircuitOpen
| CodexWebSocketExchangeError::OriginHalfOpenBusy => "connect_suppressed",
@@ -432,18 +449,32 @@ async fn connect_with_budget(
keepalive: PumpKeepalive,
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),
}
}
async fn connect_websocket_connection(
connection: &CodexWebSocketConnection,
keepalive: PumpKeepalive,
context: PumpLogContext,
) -> Result<(PooledWebSocketConnection, Duration), CodexWebSocketExchangeError> {
let started_at = Instant::now();
let connected = match fast_path_budget {
Some(budget) => timeout(
budget,
connect_pumped_websocket(connection, keepalive, context),
)
.await
.map_err(|_| CodexWebSocketExchangeError::FastPathTimeout { timeout: budget })?,
None => connect_pumped_websocket(connection, keepalive, context).await,
}?;
let (websocket, response) = connected;
let (websocket, response) = connect_pumped_websocket(connection, keepalive, context).await?;
Ok((
PooledWebSocketConnection {
websocket,
@@ -477,10 +508,6 @@ fn finish_breaker_attempt(
permit.succeed();
Ok(connection)
}
Err(error @ CodexWebSocketExchangeError::FastPathTimeout { .. }) => {
permit.fast_timeout();
Err(error)
}
Err(CodexWebSocketExchangeError::Upstream(upstream)) if upstream.status_code < 500 => {
// 账号或请求级 opening 响应证明 origin 可达，不得污染 transport 熔断器。
permit.succeed();
@@ -496,15 +523,16 @@ fn finish_breaker_attempt(
async fn wait_for_shared_connect(
waiter: pool::WebSocketPoolConnectWaiter,
fast_path_budget: Option<Duration>,
) -> Result<WebSocketPoolConnectOutcome, CodexWebSocketExchangeError> {
) -> WebSocketFastPath<WebSocketPoolConnectOutcome> {
match fast_path_budget {
Some(budget) => {
let remaining = waiter.remaining_budget(budget);
timeout(remaining, waiter.wait())
.await
.map_err(|_| CodexWebSocketExchangeError::FastPathTimeout { timeout: budget })
match timeout(remaining, waiter.wait()).await {
Ok(outcome) => WebSocketFastPath::Ready(outcome),
Err(_) => WebSocketFastPath::Missed,
}
}
None => Ok(waiter.wait().await),
None => WebSocketFastPath::Ready(waiter.wait().await),
}
}
backend/crates/providers/openai/src/transport/websocket/error.rs
-7
View File
@@ -31,12 +31,6 @@ pub enum CodexWebSocketExchangeError {
/// 建连超时时长。
timeout: Duration,
},
/// 普通请求的 WebSocket 快路径预算耗尽，payload 尚未发送。
#[error("websocket fast-path connect budget exhausted after {timeout:?}")]
FastPathTimeout {
/// 快路径等待时长。
timeout: Duration,
},
/// origin WebSocket 冷建连熔断中。
#[error("websocket origin circuit is open")]
OriginCircuitOpen,
@@ -341,7 +335,6 @@ impl CodexWebSocketExchangeError {
self,
Self::Connect(_)
| Self::ConnectTimeout { .. }
| Self::FastPathTimeout { .. }
| Self::OriginCircuitOpen
| Self::OriginHalfOpenBusy
| Self::SharedConnectFailed
backend/crates/providers/openai/src/transport/websocket/mod.rs
+1 -1
View File
@@ -37,7 +37,7 @@ pub use self::{
};
pub(crate) use self::{
coordinator::{
PreparedWebSocket, WEBSOCKET_FAST_PATH_BUDGET,
PreparedWebSocket, WEBSOCKET_FAST_PATH_BUDGET, WebSocketFastPath,
execute_prepared_response_create_request_stream, post_send_ambiguous,
prepare_response_create_request_with_pool,
},
backend/crates/providers/openai/tests/provider/contract.rs
+145
View File
@@ -459,6 +459,7 @@ fn contract_account_scope() -> Arc<FrozenAccountScope> {
"acct_usage_limit_request_path",
"acct_websocket_close",
"acct_websocket_cooldown",
"acct_websocket_fast_path",
"acct_websocket_busy_replay",
"acct_websocket_turn_state",
]
@@ -1489,6 +1490,150 @@ async fn websocket_close_after_delivery_preserves_details_and_forces_a_fresh_cli
server.await.expect("WebSocket server");
}
#[tokio::test]
async fn websocket_fast_path_miss_uses_http_and_keeps_background_preconnect() {
const ACCOUNT_ID: &str = "acct_websocket_fast_path";
const CONVERSATION_ID: &str = "conversation-websocket-fast-path";
const SESSION_ID: &str = "websocket-fast-path-session";
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, ACCOUNT_ID).await;
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind upstream listener");
let base_url = format!(
"http://{}",
listener.local_addr().expect("listener address")
);
let (preconnect_ready_tx, preconnect_ready_rx) = oneshot::channel();
let server = tokio::spawn(async move {
let (websocket_stream, _) = listener
.accept()
.await
.expect("accept background WebSocket opening");
let websocket = tokio::spawn(async move {
tokio::time::sleep(Duration::from_secs(1)).await;
let mut websocket = accept_codex_test_websocket(websocket_stream).await;
assert!(
timeout(Duration::from_millis(100), websocket.next())
.await
.is_err(),
"background preconnect must not send the HTTP-fallback request payload"
);
preconnect_ready_tx
.send(())
.expect("signal background WebSocket readiness");
let request = websocket
.next()
.await
.expect("next request should reuse the background WebSocket")
.expect("valid reused WebSocket request");
let payload: Value =
serde_json::from_str(request.to_text().expect("reused WebSocket request text"))
.expect("reused WebSocket request JSON");
assert_eq!(payload.get("type"), Some(&json!("response.create")));
websocket
.send(Message::Text(
json!({
"type": "response.completed",
"response": {
"id": "resp_background_preconnect",
"model": "gpt-5.4",
"status": "completed",
"output": [],
"usage": {
"input_tokens": 1,
"output_tokens": 1,
"total_tokens": 2
}
}
})
.to_string()
.into(),
))
.await
.expect("complete reused WebSocket request");
});
let (mut http, _) = listener
.accept()
.await
.expect("accept HTTP fallback request");
let request = capture_http_request(&mut http).await;
assert!(String::from_utf8_lossy(&request).starts_with("POST /codex/responses"));
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 HTTP fallback response");
websocket.await.expect("background WebSocket server");
});
let provider = provider_with_base_url(&store, base_url);
let operation = |thread_id| {
Operation::Generate(generate_with_persisted_session_context(
ACCOUNT_ID,
CONVERSATION_ID,
SESSION_ID,
thread_id,
))
};
let mut first = provider
.execute(
planned_request("openai", operation("thread-http-fallback")),
context("req_fast_path_http", CancellationToken::new()),
)
.await
.expect("prepare first provider stream");
let mut first_transport = None;
let mut first_decision = None;
while let Some(event) = first.next().await {
let event = event.expect("fast-path miss must not surface as a provider error");
if let Some(observation) = event.response_observation() {
first_transport = Some(observation.transport().as_str().to_owned());
if let Some(provider_metadata) = observation.provider_metadata() {
let metadata: Value = serde_json::from_str(provider_metadata.as_json())
.expect("OpenAI provider metadata JSON");
first_decision = metadata
.get("transportDecision")
.and_then(Value::as_str)
.map(ToOwned::to_owned);
}
}
}
assert_eq!(first_transport.as_deref(), Some("http_sse"));
assert_eq!(first_decision.as_deref(), Some("http2_ws_budget_exhausted"));
preconnect_ready_rx
.await
.expect("background WebSocket should become ready");
tokio::time::sleep(Duration::from_millis(20)).await;
let mut second = provider
.execute(
planned_request("openai", operation("thread-websocket-reuse")),
context("req_background_ws_reuse", CancellationToken::new()),
)
.await
.expect("prepare second provider stream");
let mut second_transport = None;
let mut second_pool = None;
while let Some(event) = second.next().await {
let event = event.expect("background WebSocket reuse should complete");
if let Some(observation) = event.response_observation() {
second_transport = Some(observation.transport().as_str().to_owned());
second_pool = Some(observation.websocket_pool());
}
}
assert_eq!(second_transport.as_deref(), Some("websocket"));
assert_eq!(second_pool, Some(Some(WebSocketPoolKind::Reuse)));
server.await.expect("HTTP and WebSocket upstream server");
}
#[tokio::test]
async fn ambiguous_websocket_close_probes_a_fresh_websocket_for_the_next_new_chain() {
const ACCOUNT_ID: &str = "acct_websocket_close";