From 49cee2397f423772f78e88891a0756f7ff47a1ed Mon Sep 17 00:00:00 2001 From: Zyy Date: Wed, 4 Nov 2026 12:17:26 +0800 Subject: [PATCH] fix(openai): avoid repeated replay errors after large SSE requests Keep oversized Codex threads on SSE until a full request falls below 13.5 MiB. After delivering a response that requires full replay, close an idle Codex socket so the next step can send its history directly. Co-authored-by: Codex --- .../src/openai/responses/validation.rs | 5 ++ .../responses/websocket/connection/handle.rs | 24 ++++++ .../responses/websocket/connection/pump.rs | 5 +- .../src/openai/responses/websocket/forward.rs | 84 ++++++++++-------- .../src/openai/responses/websocket/mod.rs | 16 +++- .../websocket/connection/middleware.rs | 4 +- .../responses/websocket/connection/mod.rs | 60 +++++++++++-- backend/crates/gateway-core/src/operation.rs | 16 ++++ .../openai/src/provider/execution.rs | 8 ++ .../providers/openai/src/provider/mod.rs | 20 +++-- .../providers/openai/src/session_transport.rs | 85 +++++++++++++------ .../openai/tests/provider/contract/mod.rs | 13 ++- docs/api.md | 13 ++- docs/architecture.md | 3 +- 14 files changed, 264 insertions(+), 92 deletions(-) diff --git a/backend/crates/gateway-api/src/openai/responses/validation.rs b/backend/crates/gateway-api/src/openai/responses/validation.rs index 9ca525085..6b4a9e4cf 100644 --- a/backend/crates/gateway-api/src/openai/responses/validation.rs +++ b/backend/crates/gateway-api/src/openai/responses/validation.rs @@ -275,6 +275,11 @@ impl ResponsesDeliveryValidator { } } + pub(super) fn completed_response_id(&self) -> Option<&str> { + self.response_id() + .filter(|_| self.state.terminal == Some(TerminalKind::Completed)) + } + pub(super) fn finish_websocket_delivery(&self) -> Result<(), ResponseValidationError> { if !self.active || self.state.terminal.is_some() { Ok(()) diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/connection/handle.rs b/backend/crates/gateway-api/src/openai/responses/websocket/connection/handle.rs index 93866c499..43420199a 100644 --- a/backend/crates/gateway-api/src/openai/responses/websocket/connection/handle.rs +++ b/backend/crates/gateway-api/src/openai/responses/websocket/connection/handle.rs @@ -36,6 +36,7 @@ pub struct ResponsesWebSocketConnection { commands: Option>, incoming: mpsc::Receiver, deferred: VecDeque, + input_gate: Arc>, exited: oneshot::Receiver, pump_task: Option>, stats: Arc, @@ -172,6 +173,26 @@ impl ResponsesWebSocketConnection { .await; } + pub(in super::super) async fn close_for_replay_if_idle(&mut self) -> bool { + // 原子划分入站帧与正常关闭，正在处理或已排队的帧都先按原合同执行 + let Ok(_guard) = self.input_gate.clone().try_lock_owned() else { + return false; + }; + if !self.deferred.is_empty() || !self.incoming.is_empty() { + return false; + } + self.close( + CloseFrame { + code: close_code::NORMAL, + reason: "Reconnect with full input.".into(), + }, + WriteContext::connection(FramePhase::Close), + PumpExitReason::ServerClose, + ) + .await; + true + } + pub(in super::super) async fn close_for_connection_limit(&mut self, reason: &'static str) { self.close( CloseFrame { @@ -344,11 +365,13 @@ where let extensions = crate::middleware::current() .map(|context| context.extensions) .unwrap_or_default(); + let input_gate = Arc::new(tokio::sync::Mutex::new(())); let pump_task = tokio::spawn(run_pump( socket, activity.clone(), command_rx, incoming_tx, + Arc::clone(&input_gate), exit_tx, Arc::clone(&connection_id), cancellation.clone(), @@ -368,6 +391,7 @@ where commands: Some(command_tx), incoming: incoming_rx, deferred: VecDeque::new(), + input_gate, exited: exit_rx, pump_task: Some(pump_task), stats, diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/connection/pump.rs b/backend/crates/gateway-api/src/openai/responses/websocket/connection/pump.rs index 19c3abc8b..422a2637e 100644 --- a/backend/crates/gateway-api/src/openai/responses/websocket/connection/pump.rs +++ b/backend/crates/gateway-api/src/openai/responses/websocket/connection/pump.rs @@ -29,6 +29,7 @@ pub(super) async fn run_pump( activity: watch::Sender, commands: mpsc::Receiver, incoming: mpsc::Sender, + input_gate: Arc>, exited: oneshot::Sender, connection_id: Arc, cancellation: CancellationToken, @@ -63,7 +64,7 @@ pub(super) async fn run_pump( // 心跳独立于业务写入等待，但共享唯一 writer，避免长写入或排队阻止失联检测 let reason = tokio::select! { biased; - reason = read_connection(&mut reader, &incoming, sender.clone(), &activity, &context) => reason, + reason = read_connection(&mut reader, &incoming, &input_gate, sender.clone(), &activity, &context) => reason, reason = write_connection(writer.clone(), commands, &context) => reason, reason = heartbeat_connection(writer.clone(), received_activity, &context) => reason, }; @@ -81,6 +82,7 @@ pub(super) async fn run_pump( async fn read_connection( socket: &mut S, incoming: &mpsc::Sender, + input_gate: &tokio::sync::Mutex<()>, sender: Arc, activity: &watch::Sender, context: &PumpContext, @@ -107,6 +109,7 @@ where } } message = socket.next() => { + let _guard = input_gate.lock().await; if matches!(&message, Some(Ok(_))) { // 链路活动以真实入站为准，插件丢弃消息不应抹掉对端已经响应的事实 activity.send_replace(Instant::now()); diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/forward.rs b/backend/crates/gateway-api/src/openai/responses/websocket/forward.rs index 0647f4355..987c0ad7d 100644 --- a/backend/crates/gateway-api/src/openai/responses/websocket/forward.rs +++ b/backend/crates/gateway-api/src/openai/responses/websocket/forward.rs @@ -23,7 +23,9 @@ use crate::openai::responses::validation::{ResponseValidationFacts, ResponsesDel use super::{ super::{DecodedResponsesRequest, OpenAiResponsesEncoder, ProtocolErrorBody}, - connection::{ConnectionEvent, FramePhase, ResponsesWebSocketConnection, WriteContext}, + connection::{ + ConnectionEvent, FramePhase, ResponsesWebSocketConnection, WriteContext, WriteOutcome, + }, protocol::{ error_event, initial_engine_error_event, is_response_create, response_metadata_event, }, @@ -32,6 +34,7 @@ use super::{ #[derive(Clone, Copy, PartialEq, Eq)] pub(super) enum ForwardOutcome { Continue, + ReconnectIfIdle, Disconnect, } @@ -55,9 +58,25 @@ impl ConnectionReplaySnapshot { } } - fn commit(&mut self, response_id: String, provider_state: Option) { - self.last_response_id = Some(response_id); - self.provider_state = provider_state; + fn commit(&mut self, capture: &ReplayCaptureHandle, delivered: Option<&str>) -> ForwardOutcome { + let mut capture = capture + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let Some(id) = capture.response_id.take() else { + return ForwardOutcome::Continue; + }; + self.provider_state = capture.provider_state.take(); + let reconnect = delivered == Some(&id) + && self + .provider_state + .as_ref() + .is_some_and(ProviderSessionState::prefers_websocket_reconnect); + self.last_response_id = Some(id); + if reconnect { + ForwardOutcome::ReconnectIfIdle + } else { + ForwardOutcome::Continue + } } } @@ -199,7 +218,7 @@ pub(super) async fn forward_response( let mut first_frame_written = false; let mut current = Some(first); - let mut terminal_seen = false; + let mut terminal_write = None; let mut first_protocol_validated = true; loop { let frame = match current.take() { @@ -211,12 +230,14 @@ pub(super) async fn forward_response( response_control, &mut interrupt, &request_id, - terminal_seen, + terminal_write.is_some(), ) .await { BodyInput::Frame(Ok(Some(frame))) => frame, - BodyInput::Frame(Ok(None)) if terminal_seen && body.is_finalized() => { + BodyInput::Frame(Ok(None)) + if terminal_write.is_some() && body.is_finalized() => + { if delivery_validator.finish_websocket_delivery().is_err() { return fail_body( connection, @@ -226,16 +247,12 @@ pub(super) async fn forward_response( ) .await; } - { - let captured = capture - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if let Some(response_id) = captured.response_id.clone() { - replay.commit(response_id, captured.provider_state.clone()); - } - } + let delivered = delivery_validator + .completed_response_id() + .filter(|_| terminal_write == Some(WriteOutcome::Written)); + let outcome = replay.commit(&capture, delivered); body.close().await; - return ForwardOutcome::Continue; + return outcome; } BodyInput::Frame(Ok(None)) => { return fail_body( @@ -256,7 +273,7 @@ pub(super) async fn forward_response( } } }; - if terminal_seen || validate_frame(&frame).is_err() { + if terminal_write.is_some() || validate_frame(&frame).is_err() { return fail_body(connection, body, MiddlewareError::InvalidState, &request_id).await; } if first_protocol_validated { @@ -278,30 +295,27 @@ pub(super) async fn forward_response( (true, false) => FramePhase::Data, }; let created_id = super::interrupt::ResponseInterrupt::created_id(&message); - match connection + let Ok(write) = connection .send_text(message, WriteContext::request(&request_id, phase)) .await - { - Ok(super::connection::WriteOutcome::Written) => { - first_frame_written = true; - if let Some(id) = created_id { - interrupt.observe_created(id, response_control); - tokio::select! { - _ = connection.wait_for_exit() => { - detach_body(body); - return ForwardOutcome::Disconnect; - } - () = interrupt.forward_pending(response_control) => {} + else { + detach_body(body); + return ForwardOutcome::Disconnect; + }; + if write == WriteOutcome::Written { + first_frame_written = true; + if let Some(id) = created_id { + interrupt.observe_created(id, response_control); + tokio::select! { + _ = connection.wait_for_exit() => { + detach_body(body); + return ForwardOutcome::Disconnect; } + () = interrupt.forward_pending(response_control) => {} } } - Ok(super::connection::WriteOutcome::Suppressed) => {} - Err(_) => { - detach_body(body); - return ForwardOutcome::Disconnect; - } } - terminal_seen = terminal; + terminal_write = terminal.then_some(write); } } diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs b/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs index 735b05464..9a067c806 100644 --- a/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs +++ b/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs @@ -179,6 +179,12 @@ async fn serve_responses_websocket(socket: Socket, session: ResponsesWebSocketSe crate::middleware::websocket::CancelOnDrop(Some(cancellation.clone())); let request_headers = request_headers.with_downstream_websocket_connection_id(connection_id.clone()); + // Codex Core 观察正常关闭后会重建完整输入；原始帧由中间件改写时无法证明同一终态已交付 + let reconnect_for_replay = crate::openai::auth::identify_codex_client(&raw_headers) + .is_some_and(|client| client.kind() == gateway_core::policy::CodexClientKind::Cli) + && !middleware + .as_ref() + .is_some_and(FrozenMiddlewarePlan::has_websocket); let mut connection = ResponsesWebSocketConnection::new( socket, connection_id, @@ -401,7 +407,7 @@ async fn serve_responses_websocket(socket: Socket, session: ResponsesWebSocketSe continue; } }; - if forward_response( + let outcome = forward_response( &mut connection, response, request_id, @@ -410,8 +416,12 @@ async fn serve_responses_websocket(socket: Socket, session: ResponsesWebSocketSe validation, &response_control, ) - .await - == ForwardOutcome::Disconnect + .await; + if outcome == ForwardOutcome::Disconnect + || (outcome == ForwardOutcome::ReconnectIfIdle + && reconnect_for_replay + && !connection.is_expired() + && connection.close_for_replay_if_idle().await) { break; } diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket/connection/middleware.rs b/backend/crates/gateway-api/tests/openai/responses/websocket/connection/middleware.rs index 4719750b9..10c3595d4 100644 --- a/backend/crates/gateway-api/tests/openai/responses/websocket/connection/middleware.rs +++ b/backend/crates/gateway-api/tests/openai/responses/websocket/connection/middleware.rs @@ -356,7 +356,7 @@ async fn rewrite_upgrade_and_control_active_response_keep_http_context_until_con ), ); let (trace, mut socket, server) = - start_active_response_with_middleware(Some(frozen), "/custom/websocket").await; + start_active_response_with(Some(frozen), "/custom/websocket", "codex-tui/0.166.0").await; // 升级后的默认会话仍持有入口上下文；101 正文结束不再提前取消它 assert!( !plan @@ -441,7 +441,7 @@ async fn input_plugin_runs_while_a_real_socket_write_is_blocked() { ), ); let (trace, mut socket, server) = - start_active_response_with_middleware(Some(frozen), "/custom/websocket").await; + start_active_response_with(Some(frozen), "/custom/websocket", "codex-tui/0.166.0").await; trace.release_terminal.notify_one(); // 仅窥视帧头以确认写入开始，不消费 64 MiB 输出，保留真实 TCP 接收窗口背压 let MaybeTlsStream::Plain(transport) = socket.get_ref() else { diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket/connection/mod.rs b/backend/crates/gateway-api/tests/openai/responses/websocket/connection/mod.rs index 8246e4fbd..1fb6c8b96 100644 --- a/backend/crates/gateway-api/tests/openai/responses/websocket/connection/mod.rs +++ b/backend/crates/gateway-api/tests/openai/responses/websocket/connection/mod.rs @@ -49,6 +49,7 @@ use crate::openai::{api_router, authenticated_client}; #[derive(Default)] struct CrossingLimitTrace { + reconnect_for_replay: AtomicBool, release_terminal: tokio::sync::Notify, starts: AtomicUsize, inputs: Mutex>, @@ -74,7 +75,9 @@ impl ExecutionSession for CrossingLimitSession { if !self.terminal_emitted { self.trace.release_terminal.notified().await; self.terminal_emitted = true; - return Ok(Some(crossing_limit_terminal_batch())); + return Ok(Some(crossing_limit_terminal_batch( + self.trace.reconnect_for_replay.load(Ordering::Acquire), + ))); } self.finalized = true; self.trace.finalized.store(true, Ordering::Release); @@ -197,8 +200,8 @@ fn crossing_limit_started_batch() -> CoordinatedEvent { .expect("started WebSocket batch") } -fn crossing_limit_terminal_batch() -> CoordinatedEvent { - let completed = ProviderEvent::wire( +fn crossing_limit_terminal_batch(reconnect: bool) -> CoordinatedEvent { + let mut completed = ProviderEvent::wire( ProtocolWireEvent::json( "openai", Some("response.completed".to_owned()), @@ -207,12 +210,20 @@ fn crossing_limit_terminal_batch() -> CoordinatedEvent { "response": { "id": "resp_ws_crossing_limit", "model": "model-a", + "output": [], "status": "completed" } }), ) .expect("completed wire"), ); + if reconnect { + completed.attach_session_update( + gateway_core::operation::ProviderSessionState::new("openai", serde_json::Map::new()) + .unwrap() + .with_websocket_reconnect(), + ); + } CoordinatedEvent::try_batch(vec![completed], CommitRequirement::AlreadyCommitted) .expect("terminal WebSocket batch") } @@ -816,12 +827,13 @@ async fn start_active_response() -> ( TestClientSocket, tokio::task::JoinHandle<()>, ) { - start_active_response_with_middleware(None, "/v1/responses").await + start_active_response_with(None, "/v1/responses", "integration-test").await } -async fn start_active_response_with_middleware( +async fn start_active_response_with( plan: Option, path: &str, + user_agent: &str, ) -> ( Arc, TestClientSocket, @@ -872,6 +884,9 @@ async fn start_active_response_with_middleware( .parse() .expect("authorization"), ); + request + .headers_mut() + .insert("user-agent", user_agent.parse().unwrap()); let (mut socket, _) = connect_async(request).await.expect("upgrade WebSocket"); let response_create = ClientMessage::Text( json!({ @@ -905,8 +920,19 @@ async fn start_active_response_with_middleware( } #[tokio::test] -async fn requests_received_during_an_active_response_should_execute_in_order() { - let (trace, mut socket, server) = start_active_response().await; +async fn replay_reconnect_waits_for_queued_requests_and_writes_the_last_terminal() { + assert_replay_queue("codex-tui/0.166.0", true).await; +} + +#[tokio::test] +async fn third_party_client_keeps_its_socket_after_replay_required_responses() { + assert_replay_queue("integration-test", false).await; +} + +async fn assert_replay_queue(user_agent: &str, expects_close: bool) { + let (trace, mut socket, server) = + start_active_response_with(None, "/v1/responses", user_agent).await; + trace.reconnect_for_replay.store(true, Ordering::Release); for input in ["second", "third"] { socket .send(ClientMessage::Text( @@ -964,7 +990,25 @@ async fn requests_received_during_an_active_response_should_execute_in_order() { vec![json!("hello"), json!("second"), json!("third")] ); assert!(!trace.cancelled.load(Ordering::Acquire)); - socket.close(None).await.expect("close WebSocket"); + if !expects_close { + socket + .send(ClientMessage::Ping(vec![9].into())) + .await + .unwrap(); + } + let next = tokio::time::timeout(Duration::from_secs(2), socket.next()) + .await + .expect("close or pong") + .unwrap() + .unwrap(); + if expects_close { + assert!( + matches!(next, ClientMessage::Close(Some(frame)) if frame.code == CloseCode::Normal) + ); + } else { + assert!(matches!(next, ClientMessage::Pong(_))); + socket.close(None).await.unwrap(); + } server.abort(); } diff --git a/backend/crates/gateway-core/src/operation.rs b/backend/crates/gateway-core/src/operation.rs index bce0b6ca5..dbe9e895b 100644 --- a/backend/crates/gateway-core/src/operation.rs +++ b/backend/crates/gateway-core/src/operation.rs @@ -123,6 +123,9 @@ pub struct ProviderSessionState { payload: Map, #[serde(default, skip_serializing_if = "Option::is_none")] extension_owner: Option>, + // 仅供本次终态交付后处置客户端连接，不持久化或传给外部适配器 + #[serde(skip)] + reconnect_websocket: bool, } /// 受管扩展会话的宿主归属，不能把插件私有状态交给原生 Provider 解码 @@ -155,6 +158,7 @@ impl ProviderSessionState { provider, payload, extension_owner: None, + reconnect_websocket: false, }) } @@ -168,6 +172,18 @@ impl ProviderSessionState { &self.payload } + /// 提示能重连重放的客户端在本次成功交付后重建连接 + #[must_use] + pub const fn with_websocket_reconnect(mut self) -> Self { + self.reconnect_websocket = true; + self + } + + #[must_use] + pub const fn prefers_websocket_reconnect(&self) -> bool { + self.reconnect_websocket + } + #[must_use] pub fn with_extension_owner(mut self, owner: ExtensionSessionOwner) -> Self { self.extension_owner = Some(Box::new(owner)); diff --git a/backend/crates/providers/openai/src/provider/execution.rs b/backend/crates/providers/openai/src/provider/execution.rs index 28ba51fc6..68d7ec298 100644 --- a/backend/crates/providers/openai/src/provider/execution.rs +++ b/backend/crates/providers/openai/src/provider/execution.rs @@ -448,6 +448,7 @@ pub(super) struct OpenAiSessionCapture { pub(super) client_turn_id: Option, pub(super) reasoning_effort_override: Option, pub(super) response_store: bool, + pub(super) oversized_http: bool, pub(super) continuation_scope: Option, } @@ -514,6 +515,13 @@ pub(super) fn attach_openai_session_update( let Ok(update) = encode_openai_session_capture(&capture) else { return; }; + let update = if capture.oversized_http + && capture.continuation_scope == Some(OpenAiContinuationScope::ReplayRequired) + { + update.with_websocket_reconnect() + } else { + update + }; events[terminal_index].attach_session_update(update); } diff --git a/backend/crates/providers/openai/src/provider/mod.rs b/backend/crates/providers/openai/src/provider/mod.rs index 72e4cb7f5..8dee125ca 100644 --- a/backend/crates/providers/openai/src/provider/mod.rs +++ b/backend/crates/providers/openai/src/provider/mod.rs @@ -66,7 +66,7 @@ use crate::credential::{ derive_previous_response_id_hash, derive_turn_alias, }; use crate::request_identity::{account_session_with_headers, turn_id_with_headers}; -use crate::session_transport::CodexSessionTransportRecovery; +use crate::session_transport::{CodexSessionTransportRecovery, WS_PAYLOAD_LIMIT}; use crate::transport::canonical::{ CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalOutcome, }; @@ -124,8 +124,6 @@ const PROVIDER_NAME: &str = "openai"; const HTTP_SSE_TRANSPORT: &str = "http_sse"; const HTTP_JSON_TRANSPORT: &str = "http_json"; const WEBSOCKET_TRANSPORT: &str = "websocket"; -// 在已观测到的 Codex OAuth 上游 16 MiB 附近消息边界前留出传输 metadata 余量 -const WEBSOCKET_HTTP_FALLBACK_THRESHOLD_BYTES: usize = 15 * 1024 * 1024; const MAX_COOKIE_HEADER_BYTES: usize = 16 * 1024; // 额度拒绝后的 usage 补查自身执行时长上限；结算等待的随机延迟见 jitter 模块 const QUOTA_FAILURE_REFRESH_TIMEOUT: Duration = Duration::from_secs(5); @@ -1056,6 +1054,7 @@ impl CodexProvider { } else { requested_transport }; + let mut oversized_http = false; if codex_oauth && transport == CodexProviderTransport::PreferWebSocket && (matches!( @@ -1070,12 +1069,12 @@ impl CodexProvider { websocket_response_create_payload_len(&upstream_request).map_err(|_| { provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent) })?; - if payload_bytes >= WEBSOCKET_HTTP_FALLBACK_THRESHOLD_BYTES { - if requirement == TransportRequirement::ExactWebSocketContinuation { + if requirement == TransportRequirement::ExactWebSocketContinuation { + if payload_bytes >= WS_PAYLOAD_LIMIT { tracing::info!( request_id = context.request_id().as_str(), payload_bytes, - threshold_bytes = WEBSOCKET_HTTP_FALLBACK_THRESHOLD_BYTES, + threshold_bytes = WS_PAYLOAD_LIMIT, continuation_recovery_action = "client_replay_required", "OpenAI oversized WebSocket continuation requires full client input before send" ); @@ -1083,12 +1082,16 @@ impl CodexProvider { "transport_payload_too_large", )); } + } else if self + .session_transport_recovery + .http_for_full_payload(session_transport_key.as_ref(), payload_bytes) + { transport = CodexProviderTransport::HttpOnly; + oversized_http = true; tracing::info!( request_id = context.request_id().as_str(), payload_bytes, - threshold_bytes = WEBSOCKET_HTTP_FALLBACK_THRESHOLD_BYTES, - "OpenAI oversized new-chain request selected HTTP before WebSocket send" + "OpenAI full request selected HTTP by WebSocket payload size" ); } } @@ -1147,6 +1150,7 @@ impl CodexProvider { ) .reasoning_effort_override, response_store, + oversized_http, continuation_scope: None, }); let allows_account_state_mutation = lease.allows_account_state_mutation(); diff --git a/backend/crates/providers/openai/src/session_transport.rs b/backend/crates/providers/openai/src/session_transport.rs index 9bdf23ed1..e39d0c3fd 100644 --- a/backend/crates/providers/openai/src/session_transport.rs +++ b/backend/crates/providers/openai/src/session_transport.rs @@ -1,4 +1,4 @@ -//! 根会话的连续 WS 失败预算；耗尽后保持 HTTP，空闲过期后释放状态 +//! 按客户端线程维护 WS 失败预算及大包 HTTP 回退区间，空闲过期后释放状态 use std::collections::HashMap; use std::sync::{Arc, Mutex, MutexGuard}; @@ -7,6 +7,11 @@ use std::time::Duration; use gateway_core::provider_ports::ProviderSessionAffinityKey; use tokio::time::Instant; +// 在实测的 16 MiB 上游 WS 边界前留半 MiB，再用 2 MiB 回差避免频繁切换 +const MIB: usize = 1024 * 1024; +pub(crate) const WS_PAYLOAD_LIMIT: usize = 15 * MIB + MIB / 2; +const WS_PAYLOAD_RESUME: usize = WS_PAYLOAD_LIMIT - 2 * MIB; + const SESSION_IDLE_RETENTION: Duration = Duration::from_secs(8 * 60 * 60); const MAX_SESSION_TRANSPORTS: usize = 16_384; @@ -17,6 +22,7 @@ pub(crate) struct CodexSessionTransportRecovery { struct SessionTransport { state: TransportState, + oversized: bool, last_used: Instant, } @@ -46,39 +52,70 @@ impl CodexSessionTransportRecovery { key: &ProviderSessionAffinityKey, max_retries: u32, ) -> bool { - self.update(key, |state| match state { - TransportState::WebSocket { failures } if failures < max_retries => { - TransportState::WebSocket { - failures: failures + 1, + self.update(key, |session| { + session.state = match session.state { + TransportState::WebSocket { failures } if failures < max_retries => { + TransportState::WebSocket { + failures: failures + 1, + } } - } - TransportState::WebSocket { .. } | TransportState::Http => TransportState::Http, + _ => TransportState::Http, + }; + matches!(session.state, TransportState::Http) }) } pub(crate) fn disable_websocket(&self, key: &ProviderSessionAffinityKey) { - self.update(key, |_| TransportState::Http); + self.update(key, |session| session.state = TransportState::Http); } pub(crate) fn websocket_succeeded(&self, key: &ProviderSessionAffinityKey) { - self.lock().remove(key); + // 预热或增量成功不能证明下一份完整请求能放进 WS + self.update(key, |session| { + session.state = TransportState::WebSocket { failures: 0 } + }); } - fn update( + /// 只用完整新链的实际帧长度更新大小偏好，重连沿用同一线程的判断 + pub(crate) fn http_for_full_payload( + &self, + key: Option<&ProviderSessionAffinityKey>, + bytes: usize, + ) -> bool { + let Some(key) = key else { + return bytes >= WS_PAYLOAD_LIMIT; + }; + self.update(key, |session| { + if bytes >= WS_PAYLOAD_LIMIT { + session.oversized = true; + } else if bytes < WS_PAYLOAD_RESUME { + session.oversized = false; + } + session.oversized + }) + } + + fn update( &self, key: &ProviderSessionAffinityKey, - update: impl FnOnce(TransportState) -> TransportState, - ) -> bool { + update: impl FnOnce(&mut SessionTransport) -> T, + ) -> T { let now = Instant::now(); let mut sessions = self.lock(); - let previous = sessions + let mut session = sessions .remove(key) - .filter(|session| { - now.saturating_duration_since(session.last_used) < SESSION_IDLE_RETENTION - }) - .map_or(TransportState::WebSocket { failures: 0 }, |session| { - session.state + .filter(|s| now.saturating_duration_since(s.last_used) < SESSION_IDLE_RETENTION) + .unwrap_or(SessionTransport { + state: TransportState::WebSocket { failures: 0 }, + oversized: false, + last_used: now, }); + let result = update(&mut session); + if !session.oversized && matches!(session.state, TransportState::WebSocket { failures: 0 }) + { + return result; + } + session.last_used = now; if sessions.len() >= MAX_SESSION_TRANSPORTS { sessions.retain(|_, session| { now.saturating_duration_since(session.last_used) < SESSION_IDLE_RETENTION @@ -92,16 +129,8 @@ impl CodexSessionTransportRecovery { sessions.remove(&oldest); } } - let state = update(previous); - let uses_http = matches!(state, TransportState::Http); - sessions.insert( - key.clone(), - SessionTransport { - state, - last_used: now, - }, - ); - uses_http + sessions.insert(key.clone(), session); + result } fn lock(&self) -> MutexGuard<'_, HashMap> { diff --git a/backend/crates/providers/openai/tests/provider/contract/mod.rs b/backend/crates/providers/openai/tests/provider/contract/mod.rs index 43ef8d999..c9ee5b3d9 100644 --- a/backend/crates/providers/openai/tests/provider/contract/mod.rs +++ b/backend/crates/providers/openai/tests/provider/contract/mod.rs @@ -5108,7 +5108,8 @@ async fn oversized_http_new_chain_uses_http_without_disabling_session_websocket( #[tokio::test] async fn oversized_http_threshold_counts_utf8_and_json_escaping() { - const THRESHOLD: usize = 15 * 1024 * 1024; + const THRESHOLD: usize = 15 * 1024 * 1024 + 512 * 1024; + const RESUME: usize = THRESHOLD - 2 * 1024 * 1024; let store = Arc::new(MemoryAccountStore::default()); create_account(&store, "acct_provider_contract").await; let server = MockServer::start().await; @@ -5159,7 +5160,14 @@ async fn oversized_http_threshold_counts_utf8_and_json_escaping() { frame["session_id"] = body["session_id"].clone(); let baseline_len = serde_json::to_vec(&frame).unwrap().len(); - for (payload_len, expected) in [(THRESHOLD - 1, "websocket"), (THRESHOLD, "http_sse")] { + for (payload_len, expected) in [ + (THRESHOLD - 1, "websocket"), + (THRESHOLD, "http_sse"), + (THRESHOLD - 1, "http_sse"), + (RESUME, "http_sse"), + (RESUME - 1, "websocket"), + (THRESHOLD - 1, "websocket"), + ] { let text = format!("{seed_text}{}", "x".repeat(payload_len - baseline_len)); body.get_mut("input").unwrap()[0]["content"][0]["text"] = json!(text); let stream = Arc::clone(&provider) @@ -5507,6 +5515,7 @@ async fn oversized_websocket_http_response_continues_with_client_tool_history_re } drop(first); let state = state.expect("HTTP response publishes replay requirement"); + assert!(state.prefers_websocket_reconnect()); assert_eq!(state.payload()["continuation_scope"], "replay_required"); let tool_output = json!({"type": "function_call_output", "call_id": "call_large", "output": "done"}); diff --git a/docs/api.md b/docs/api.md index d3d3913d1..cc967a4cd 100644 --- a/docs/api.md +++ b/docs/api.md @@ -475,13 +475,13 @@ Responses WebSocket 接受文本消息，`response.create` 在同一连接串行 网关分别维护两段链路的心跳：上游 WS 定时发送 Ping 保活并回复 Pong，下游 WS 独立探活；Ping/Pong 不跨段转发，也不转换为 SSE 事件。 下游探活包含消息接收中的字节进度，完整消息收齐后才进入中间件与请求处理。 -客户端断开时取消当前执行，并释放仍归该客户端所有的上游连接。`response.completed` 只结束本轮响应，保留客户端 WS 等待下一轮请求 +客户端断开时取消当前执行，并释放仍归该客户端所有的上游连接。`response.completed` 结束本轮响应。通常保留客户端 WS 等待下一轮请求，大包回退后的 Codex 重连规则见下文 | 入站 → 上游 | 正常回答结束后的连接处理 | | --- | --- | | SSE → SSE | 结束本次响应流，底层 HTTP 连接由 HTTP 连接池管理 | | SSE → WS | 结束 SSE，上游 WS 归池，按无客户端绑定的空闲期限回收 | -| WS → SSE | 上游本次响应结束，下游 WS 保留 | +| WS → SSE | 上游本次响应结束，下游 WS 通常保留；大包回退可触发正常重连 | | WS → WS | 保留当前绑定的上游 WS，支持下一轮续接 | 活跃客户端 WS 的绑定不受普通空闲回收期限约束；上游关闭、总寿命到期或连接被回收仍可能使连接内续接失效， @@ -520,7 +520,8 @@ OAuth 账号默认 `prefer_websocket`，客户端使用 HTTP/SSE 时仍可能选 客户端可跳过预热继续正式请求，外部插件 adapter 按各自的协议能力处理 Codex OAuth backend 的候选上游为 WS 时，无 `previous_response_id` 的普通新链若规范化 -`response.create` 达到 15 MiB，发送前选择 HTTP/SSE；HTTP 和 WS 入站均适用，下游交付协议不变。 +`response.create` 达到 15.5 MiB，发送前选择 HTTP/SSE；HTTP 和 WS 入站均适用。能关联到客户端线程时， +后续完整新链沿用 HTTP/SSE，直到帧长低于 13.5 MiB 再尝试 WS。预热和增量帧不改变此大小判断。 该阈值为已观察到的上游消息大小边界预留余量，不是网关输入长度上限或 OpenAI 公布的统一限制； API Key 账号、小请求、显式 warmup、未知外部 previous ID 和持久化续接不因该大小策略改用 HTTP @@ -528,7 +529,11 @@ OAuth 的大 connection-local WS 续接，以及 HTTP `store=false` 成功后标 的 WS 增量，均在发送前返回 `status: 400 / previous_response_not_found`。客户端应清除旧 ID， 携带完整历史及工具调用／输出重试。官方 Codex 支持重连 WS 后完整重发，也可能按自身重试预算 切到 HTTP；其他客户端需要自行实现此合同。代理不缓存 transcript，不把增量输入当作独立新链， -连接内续接失去原连接时仍需客户端恢复；HTTP 链后续步骤可能增加一次恢复信号、重连及全量上传 +连接内续接失去原连接时仍需客户端恢复 + +因大小走 HTTP/SSE 且需要完整重放时，网关向可识别的 Codex CLI/Core 写出成功终态后， +下游空闲则正常关闭 WS。客户端观察到关闭后，下一步在新连接上发送完整输入。 +已接收的后续请求按序处理；启用 WS 帧中间件时保留连接。若增量先到，仍按上述 400 合同恢复 已建立模型执行的 Responses、Images 和 Search HTTP 响应按以下规则返回关联 ID： `x-gateway-request-id` 为模型执行 ID；`x-request-id` 保留有效上游值，只有上游 diff --git a/docs/architecture.md b/docs/architecture.md index 45d3a6251..516c45154 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -621,7 +621,8 @@ client，OIDC 的 JWKS 缓存与单飞归属对应出口状态。自动刷新提 客户端配置示例维护在 [部署文档](../deploy/README.md#客户端配置) 客户端到代理与代理到上游的传输选择相互独立。客户端的 `supports_websockets = false` -不禁止 Provider 使用上游 WebSocket。响应终态前的 Close 1000 仍视为失败， +不禁止 Provider 使用上游 WebSocket。OpenAI Provider 按线程独立维护大包回退与 WS 失败预算， +以终态会话更新提示下游重连；API 负责确认成功写出和入站队列空闲。响应终态前的 Close 1000 仍视为失败， 后续恢复请求成功也不改写原失败请求的结果 ### 错误与诊断三层边界