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 <noreply@openai.com>
2026-11-06 08:06:19 +00:00 · 2026-11-04 12:17:26 +08:00
parent 1915806bd1
commit 49cee2397f
14 changed files with 264 additions and 92 deletions
@@ -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(())
@@ -36,6 +36,7 @@ pub struct ResponsesWebSocketConnection {
    commands: Option<mpsc::Sender<ConnectionCommand>>,
    incoming: mpsc::Receiver<PendingConnectionEvent>,
    deferred: VecDeque<PendingConnectionEvent>,
+    input_gate: Arc<tokio::sync::Mutex<()>>,
    exited: oneshot::Receiver<PumpExitReason>,
    pump_task: Option<JoinHandle<()>>,
    stats: Arc<ConnectionStats>,
@@ -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,
@@ -29,6 +29,7 @@ pub(super) async fn run_pump<S, E>(
    activity: watch::Sender<Instant>,
    commands: mpsc::Receiver<ConnectionCommand>,
    incoming: mpsc::Sender<PendingConnectionEvent>,
+    input_gate: Arc<tokio::sync::Mutex<()>>,
    exited: oneshot::Sender<PumpExitReason>,
    connection_id: Arc<str>,
    cancellation: CancellationToken,
@@ -63,7 +64,7 @@ pub(super) async fn run_pump<S, E>(
    // 心跳独立于业务写入等待，但共享唯一 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<S, E>(
 async fn read_connection<S, E>(
    socket: &mut S,
    incoming: &mpsc::Sender<PendingConnectionEvent>,
+    input_gate: &tokio::sync::Mutex<()>,
    sender: Arc<dyn contract::Sender>,
    activity: &watch::Sender<Instant>,
    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());
@@ -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<ProviderSessionState>) {
-        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);
    }
 }

@@ -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;
        }
@@ -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 {
@@ -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<Vec<Value>>,
@@ -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<gateway_core::engine::middleware::FrozenMiddlewarePlan>,
    path: &str,
+    user_agent: &str,
 ) -> (
    Arc<CrossingLimitTrace>,
    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();
 }

@@ -123,6 +123,9 @@ pub struct ProviderSessionState {
    payload: Map<String, Value>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    extension_owner: Option<Box<ExtensionSessionOwner>>,
+    // 仅供本次终态交付后处置客户端连接，不持久化或传给外部适配器
+    #[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));
@@ -448,6 +448,7 @@ pub(super) struct OpenAiSessionCapture {
    pub(super) client_turn_id: Option<String>,
    pub(super) reasoning_effort_override: Option<String>,
    pub(super) response_store: bool,
+    pub(super) oversized_http: bool,
    pub(super) continuation_scope: Option<OpenAiContinuationScope>,
 }

@@ -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);
 }

@@ -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();
@@ -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<T>(
        &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<ProviderSessionAffinityKey, SessionTransport>> {
@@ -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"});
@@ -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` 保留有效上游值，只有上游
@@ -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 仍视为失败，
 后续恢复请求成功也不改写原失败请求的结果

 ### 错误与诊断三层边界