fix(openai): widen and trace stream replay grace (#304)

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>

backend/crates/providers/openai/src/provider/execution.rs

@@ -837,7 +837,7 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
             .with_requested_service_tier(request.service_tier())
             .with_request_tool_pricing(upstream_model.as_str(), request.tools())
             .with_raw_sse_passthrough();
-        let mut pre_commit_events = PreCommitClientEvents::new();
+        let mut pre_commit_events = PreCommitClientEvents::new(trace);
         loop {
             let Some(stream_deadline) = remaining(context.deadline()) else {
                 if allows_account_state_mutation {
@@ -882,7 +882,7 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
             let next = match next {
                 Ok(PreCommitPoll::Upstream(next)) => next,
                 Ok(PreCommitPoll::GraceElapsed) => {
-                    for event in pre_commit_events.commit_pending() {
+                    for event in pre_commit_events.commit_pending(PreCommitReleaseReason::GraceTimeout) {
                         yield event;
                     }
                     continue;

backend/crates/providers/openai/src/provider/failure.rs

@@ -1,6 +1,7 @@
 //! OpenAI 上游失败分类、恢复决策与稳定错误投影。
 
 use super::*;
+use gateway_core::diagnostics::TraceContext;
 
 /// OpenAI 失败对 Smart 账号分数的结构化 reason 闭集。
 ///
@@ -83,23 +84,35 @@ pub(super) async fn wait_for_replay_grace(deadline: Option<Instant>) {
 
 /// 提交边界前的上游事件预取。
 ///
-/// 原始 chunk 计数而不是重编码后的 event 大小。时间与 64 KiB 共同限定无感换号
+/// 原始 chunk 计数而不是重编码后的 event 大小。时间与字节阈值共同限定无感换号
 /// 窗口；任一边界到达都会提交已缓存 wire，不能因网关私有资源规则伪造上游协议
 /// 失败。一旦提交，后续事件不再具备无痕重放资格。
 pub(super) struct PreCommitClientEvents {
     pending: Vec<ProviderEvent>,
     prefetched_bytes: usize,
-    replay_grace_deadline: Option<Instant>,
+    replay_grace_started_at: Option<Instant>,
     committed: bool,
+    trace: TraceContext,
+}
+
+#[derive(Clone, Copy, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub(super) enum PreCommitReleaseReason {
+    SemanticOutput,
+    Terminal,
+    ByteLimit,
+    GraceTimeout,
+    Eof,
 }
 
 impl PreCommitClientEvents {
-    pub(super) const fn new() -> Self {
+    pub(super) const fn new(trace: TraceContext) -> Self {
         Self {
             pending: Vec::new(),
             prefetched_bytes: 0,
-            replay_grace_deadline: None,
+            replay_grace_started_at: None,
             committed: false,
+            trace,
         }
     }
 
@@ -120,14 +133,17 @@ impl PreCommitClientEvents {
         }
         let starts_replay_grace = incoming.iter().any(ProviderEvent::has_client_event);
         self.pending.extend(incoming);
-        if timing_signals.semantic_output
-            || completed
-            || self.prefetched_bytes > MAX_STREAM_PREFETCH_BYTES
-        {
-            return self.commit_pending();
+        if timing_signals.semantic_output {
+            return self.commit_pending(PreCommitReleaseReason::SemanticOutput);
         }
-        if starts_replay_grace && self.replay_grace_deadline.is_none() {
-            self.replay_grace_deadline = Instant::now().checked_add(STREAM_REPLAY_GRACE);
+        if completed {
+            return self.commit_pending(PreCommitReleaseReason::Terminal);
+        }
+        if self.prefetched_bytes > MAX_STREAM_PREFETCH_BYTES {
+            return self.commit_pending(PreCommitReleaseReason::ByteLimit);
+        }
+        if starts_replay_grace && self.replay_grace_started_at.is_none() {
+            self.replay_grace_started_at = Some(Instant::now());
         }
         Vec::new()
     }
@@ -142,7 +158,7 @@ impl PreCommitClientEvents {
         if self.committed {
             events
         } else {
-            self.commit_pending()
+            self.commit_pending(PreCommitReleaseReason::Eof)
         }
     }
 
@@ -151,28 +167,38 @@ impl PreCommitClientEvents {
             return incoming;
         }
         self.pending.extend(incoming);
-        self.commit_pending()
+        // 此入口只在同批次已出现语义输出、随后失败时释放原始事件。
+        self.commit_pending(PreCommitReleaseReason::SemanticOutput)
     }
 
     pub(super) fn take_for_failure(&mut self, incoming: Vec<ProviderEvent>) -> Vec<ProviderEvent> {
         self.pending.extend(incoming);
         self.prefetched_bytes = 0;
-        self.replay_grace_deadline = None;
+        self.replay_grace_started_at = None;
         std::mem::take(&mut self.pending)
     }
 
-    pub(super) const fn replay_grace_deadline(&self) -> Option<Instant> {
-        self.replay_grace_deadline
+    pub(super) fn replay_grace_deadline(&self) -> Option<Instant> {
+        self.replay_grace_started_at
+            .and_then(|started| started.checked_add(STREAM_REPLAY_GRACE))
     }
 
     pub(super) const fn is_committed(&self) -> bool {
         self.committed
     }
 
-    pub(super) fn commit_pending(&mut self) -> Vec<ProviderEvent> {
+    pub(super) fn commit_pending(&mut self, reason: PreCommitReleaseReason) -> Vec<ProviderEvent> {
+        // Provider 只记录释放缓存的原因；实际下游提交仍由 Core 记录和裁决。
+        if self.trace.is_enabled() {
+            self.trace.record("provider.precommit.released", json!({
+                "reason": reason,
+                "prefetchedBytes": self.prefetched_bytes,
+                "waitMs": self.replay_grace_started_at.map_or(0, |started| started.elapsed().as_millis()),
+            }));
+        }
         self.committed = true;
         self.prefetched_bytes = 0;
-        self.replay_grace_deadline = None;
+        self.replay_grace_started_at = None;
         std::mem::take(&mut self.pending)
     }
 }

backend/crates/providers/openai/src/provider/mod.rs

@@ -120,12 +120,13 @@ 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;
-/// 提交边界前最多保留 64 KiB 原始上游 chunk；达到阈值后结束无感换号窗口，
+/// 提交边界前预取 128 KiB 原始上游 chunk；容纳携带配置回显的前导事件，
+/// 超过阈值后结束无感换号窗口（最后一个 chunk 可越过阈值），
 /// 但不会把上游数据改写成协议失败。
-const MAX_STREAM_PREFETCH_BYTES: usize = 64 * 1024;
+const MAX_STREAM_PREFETCH_BYTES: usize = 128 * 1024;
 /// 短暂保留 response.created 等结构事件，让随后到达的明确拒绝可以无感换号；
 /// 到期即放行，避免模型长时间思考时让客户端一直收不到首事件。
-const STREAM_REPLAY_GRACE: Duration = Duration::from_millis(1_200);
+const STREAM_REPLAY_GRACE: Duration = Duration::from_millis(2_500);
 // 额度拒绝后先给上游额度结算留出时间，再以受限时长同步 usage 快照。
 const QUOTA_FAILURE_REFRESH_DELAY: Duration = Duration::from_secs(2);
 const QUOTA_FAILURE_REFRESH_TIMEOUT: Duration = Duration::from_secs(5);

backend/crates/providers/openai/tests/provider/contract/mod.rs

@@ -1,5 +1,6 @@
 mod account_isolation;
 mod capacity;
+mod precommit;
 mod response_interrupt;
 
 use std::collections::{BTreeMap, BTreeSet};
@@ -8559,7 +8560,7 @@ async fn ordinary_request_should_bound_structural_event_replay_grace() {
         .await
         .expect("prepare provider stream");
 
-    let first_event = timeout(Duration::from_secs(2), async {
+    let first_event = timeout(Duration::from_secs(4), async {
         loop {
             let event = stream
                 .next()
@@ -9073,10 +9074,10 @@ async fn exact_websocket_busy_then_replay_scope_relaxation_is_rejected_before_se
 }
 
 #[tokio::test]
-async fn continuation_prefetch_over_64_kib_should_commit_wire_without_protocol_failure() {
+async fn continuation_prefetch_over_128_kib_should_commit_wire_without_protocol_failure() {
     let store = Arc::new(MemoryAccountStore::default());
     create_account(&store, "acct_prefetch_limit").await;
-    let padding = "x".repeat(64 * 1024);
+    let padding = "x".repeat(128 * 1024);
     let body = format!(
         "event: response.created\ndata: {}\n\n",
         json!({
@@ -9089,7 +9090,7 @@ async fn continuation_prefetch_over_64_kib_should_commit_wire_without_protocol_f
             }
         })
     );
-    assert!(body.len() > 64 * 1024);
+    assert!(body.len() > 128 * 1024);
     let (base_url, release, _first_chunk_sent, server) =
         paused_chunked_sse_server(body, String::new()).await;
     let mut stream = provider_with_base_url(&store, base_url)

backend/crates/providers/openai/tests/provider/contract/precommit.rs

@@ -0,0 +1,317 @@
+use gateway_core::diagnostics::TraceContext;
+use gateway_core::engine::provider::ProviderStream;
+use gateway_core::error::ProviderError;
+use gateway_core::event::ProviderEvent;
+use gateway_protocol::openai::sse::encode_sse_event;
+
+use super::*;
+
+fn traced_context(trace: &TraceContext, cancellation: CancellationToken) -> AttemptContext {
+    AttemptContext::new(
+        RequestAttemptContext::new(
+            ModelRequestId::new("req_precommit").unwrap(),
+            ClientApiKeyId::new("key_openai_contract").unwrap(),
+        )
+        .with_trace(trace.clone()),
+        NonZeroU32::new(1).unwrap(),
+        SystemTime::now() + Duration::from_secs(30),
+        account_policy(),
+        AccountAttemptContext::new(BTreeSet::new(), None, None)
+            .with_account_scope(contract_account_scope()),
+        None,
+        cancellation,
+    )
+}
+
+fn structural_event(kind: &str, bytes: usize) -> Value {
+    let mut event = json!({
+        "type": kind,
+        "response": {
+            "id": "resp_precommit", "model": "gpt-5.4", "status": "in_progress",
+            "instructions": "", "tools": [], "output": []
+        }
+    });
+    event["response"]["instructions"] = json!("x".repeat(bytes - event.to_string().len()));
+    assert_eq!(event.to_string().len(), bytes);
+    event
+}
+
+fn overload() -> Value {
+    json!({
+        "type": "error",
+        "error": {"type": "service_unavailable_error", "code": "server_is_overloaded", "message": "busy"}
+    })
+}
+
+fn sse(event: &Value) -> String {
+    encode_sse_event(event["type"].as_str().unwrap(), &event.to_string())
+}
+
+fn releases(trace: &TraceContext) -> Vec<Value> {
+    trace.snapshot().unwrap()["events"]
+        .as_array()
+        .unwrap()
+        .iter()
+        .filter(|event| event["stage"] == "provider.precommit.released")
+        .cloned()
+        .collect()
+}
+
+async fn next_client_event(stream: &mut ProviderStream) -> ProviderEvent {
+    loop {
+        let event = stream.next().await.unwrap().unwrap();
+        if event.has_client_event() {
+            return event;
+        }
+    }
+}
+
+async fn stream_failure(stream: &mut ProviderStream) -> (Vec<ProviderEvent>, ProviderError) {
+    let mut events = Vec::new();
+    loop {
+        match stream.next().await.expect("upstream failure") {
+            Ok(event) => events.push(event),
+            Err(error) => return (events, error),
+        }
+    }
+}
+
+#[tokio::test]
+async fn large_structural_events_preserve_overload_replay_past_the_old_grace_on_http_and_websocket()
+{
+    for websocket in [false, true] {
+        let store = Arc::new(MemoryAccountStore::default());
+        create_account(&store, "acct_provider_contract").await;
+        // #259 仅保留了事件长度；用合成配置回显复现尺寸，不依赖现场私有正文。
+        let created = structural_event("response.created", 38_781);
+        let progress = structural_event("response.in_progress", 38_785);
+        let failure = overload();
+        let expected = vec![created.clone(), progress.clone(), failure.clone()];
+        let (base_url, release, server) = if websocket {
+            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+            let base_url = format!("http://{}", listener.local_addr().unwrap());
+            let (release, released) = oneshot::channel();
+            let server = tokio::spawn(async move {
+                let (socket, _) = listener.accept().await.unwrap();
+                let mut ws = accept_codex_test_websocket(socket).await;
+                ws.next().await.unwrap().unwrap();
+                for event in [created, progress] {
+                    ws.send(Message::Text(event.to_string().into()))
+                        .await
+                        .unwrap();
+                }
+                released.await.unwrap();
+                ws.send(Message::Text(failure.to_string().into()))
+                    .await
+                    .unwrap();
+            });
+            (base_url, release, server)
+        } else {
+            let (base_url, release, _, server) =
+                paused_chunked_sse_server(sse(&created) + &sse(&progress), sse(&failure)).await;
+            (base_url, release, server)
+        };
+        let trace = TraceContext::new("req_precommit");
+        let operation = if websocket {
+            generate_operation()
+        } else {
+            http_generate_operation()
+        };
+        let mut stream = provider_with_base_url(&store, base_url)
+            .execute(
+                planned_request("openai", operation),
+                traced_context(&trace, CancellationToken::new()),
+            )
+            .await
+            .unwrap();
+
+        assert!(
+            timeout(Duration::from_millis(1_500), next_client_event(&mut stream))
+                .await
+                .is_err()
+        );
+        release.send(()).unwrap();
+        let (events, mut error) = timeout(Duration::from_secs(1), stream_failure(&mut stream))
+            .await
+            .unwrap();
+        assert!(events.iter().all(|event| !event.has_client_event()));
+        assert!(error.replay_is_safe());
+        assert!(matches!(
+            error.pre_delivery_retry(),
+            Some(PreDeliveryRetry::SameAccountTransientRetry { .. })
+        ));
+        let atomic = error.take_atomic_client_events();
+        let actual: Vec<_> = atomic
+            .iter()
+            .filter_map(|event| event.wire_event().map(|wire| wire.data().clone()))
+            .collect();
+        assert_eq!(actual, expected);
+        assert!(
+            releases(&trace).is_empty(),
+            "discardable failure must not be reported as a release"
+        );
+        server.await.unwrap();
+    }
+}
+
+#[tokio::test]
+async fn later_structural_events_do_not_extend_grace_and_late_overload_is_not_replayable() {
+    let store = Arc::new(MemoryAccountStore::default());
+    create_account(&store, "acct_provider_contract").await;
+    let created = sse(&structural_event("response.created", 38_781));
+    let progress = sse(&structural_event("response.in_progress", 38_785));
+    let expected_bytes = created.len() + progress.len();
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let base_url = format!("http://{}", listener.local_addr().unwrap());
+    let (release, released) = oneshot::channel();
+    let server = tokio::spawn(async move {
+        let (mut socket, _) = listener.accept().await.unwrap();
+        read_http_request(&mut socket).await;
+        socket.write_all(b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n").await.unwrap();
+        write_http_chunk(&mut socket, &created).await;
+        tokio::time::sleep(Duration::from_millis(1_500)).await;
+        write_http_chunk(&mut socket, &progress).await;
+        released.await.unwrap();
+        write_http_chunk(&mut socket, &sse(&overload())).await;
+        socket.write_all(b"0\r\n\r\n").await.unwrap();
+    });
+    let trace = TraceContext::new("req_precommit");
+    let mut stream = provider_with_base_url(&store, base_url)
+        .execute(
+            planned_request("openai", http_generate_operation()),
+            traced_context(&trace, CancellationToken::new()),
+        )
+        .await
+        .unwrap();
+    let first = timeout(Duration::from_millis(3_500), next_client_event(&mut stream))
+        .await
+        .unwrap();
+    assert_eq!(
+        first.wire_event().unwrap().event_type(),
+        Some("response.created")
+    );
+    release.send(()).unwrap();
+    let (events, error) = timeout(Duration::from_secs(1), stream_failure(&mut stream))
+        .await
+        .unwrap();
+    assert!(!error.replay_is_safe());
+    assert!(error.pre_delivery_retry().is_none());
+    assert!(events.iter().any(|event| {
+        event
+            .wire_event()
+            .is_some_and(|wire| wire.data() == &overload())
+    }));
+    let releases = releases(&trace);
+    assert_eq!(releases.len(), 1);
+    assert_eq!(releases[0]["attemptIndex"], 1);
+    assert_eq!(releases[0]["data"]["reason"], "grace_timeout");
+    assert_eq!(releases[0]["data"]["prefetchedBytes"], expected_bytes);
+    assert!((2_500..3_500).contains(&releases[0]["data"]["waitMs"].as_u64().unwrap()));
+    server.await.unwrap();
+}
+
+#[tokio::test]
+async fn immediate_release_preserves_wire_and_records_the_boundary_once() {
+    let semantic = json!({"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello"});
+    let terminal = json!({"type": "response.completed", "response": {"id": "resp_precommit", "model": "gpt-5.4", "status": "completed", "output": []}});
+    let tool = json!({"type": "response.web_search_call.in_progress", "item_id": "search_precommit", "output_index": 0});
+    for (reason, body) in [
+        (
+            "semantic_output",
+            sse(&structural_event("response.created", 512)) + &sse(&semantic),
+        ),
+        (
+            "semantic_output",
+            sse(&structural_event("response.created", 512)) + &sse(&tool),
+        ),
+        (
+            "terminal",
+            sse(&structural_event("response.created", 512)) + &sse(&terminal),
+        ),
+        (
+            "byte_limit",
+            sse(&structural_event("response.created", 129 * 1024)),
+        ),
+        ("eof", sse(&structural_event("response.created", 512))),
+    ] {
+        let store = Arc::new(MemoryAccountStore::default());
+        create_account(&store, "acct_provider_contract").await;
+        let (base_url, release, _, server) =
+            paused_chunked_sse_server(body.clone(), String::new()).await;
+        // EOF 是另一种边界；其余场景必须在上游结束前立即释放。
+        let mut release = Some(release);
+        if reason == "eof" {
+            release.take().unwrap().send(()).unwrap();
+        }
+        let trace = TraceContext::new("req_precommit");
+        let mut stream = provider_with_base_url(&store, base_url)
+            .execute(
+                planned_request("openai", http_generate_operation()),
+                traced_context(&trace, CancellationToken::new()),
+            )
+            .await
+            .unwrap();
+        let first = timeout(Duration::from_secs(1), next_client_event(&mut stream))
+            .await
+            .unwrap_or_else(|error| panic!("{reason} did not release immediately: {error}"));
+        if let Some(release) = release {
+            release.send(()).unwrap();
+        }
+        let mut wire = Vec::new();
+        wire.extend_from_slice(first.wire_event().unwrap().raw_sse_frame().unwrap());
+        while let Some(event) = stream.next().await {
+            if let Some(frame) = event
+                .unwrap()
+                .wire_event()
+                .and_then(|wire| wire.raw_sse_frame())
+            {
+                wire.extend_from_slice(frame);
+            }
+        }
+        assert_eq!(wire, body.as_bytes());
+        let releases = releases(&trace);
+        assert_eq!(releases.len(), 1);
+        assert_eq!(releases[0]["data"]["reason"], reason);
+        if reason == "byte_limit" {
+            assert!(releases[0]["data"]["prefetchedBytes"].as_u64().unwrap() > 128 * 1024);
+        } else {
+            assert_eq!(releases[0]["data"]["prefetchedBytes"], body.len());
+        }
+        server.await.unwrap();
+    }
+}
+
+#[tokio::test]
+async fn cancelling_buffered_structural_events_does_not_release_or_offer_replay() {
+    let store = Arc::new(MemoryAccountStore::default());
+    create_account(&store, "acct_provider_contract").await;
+    let (base_url, release, _, server) = paused_chunked_sse_server(
+        sse(&structural_event("response.created", 38_781)),
+        String::new(),
+    )
+    .await;
+    let trace = TraceContext::new("req_precommit");
+    let cancellation = CancellationToken::new();
+    let mut stream = provider_with_base_url(&store, base_url)
+        .execute(
+            planned_request("openai", http_generate_operation()),
+            traced_context(&trace, cancellation.clone()),
+        )
+        .await
+        .unwrap();
+    assert!(
+        timeout(Duration::from_millis(100), next_client_event(&mut stream))
+            .await
+            .is_err()
+    );
+    cancellation.cancel();
+    let (events, error) = timeout(Duration::from_secs(1), stream_failure(&mut stream))
+        .await
+        .unwrap();
+    assert_eq!(error.kind(), ProviderErrorKind::Cancelled);
+    assert!(!error.replay_is_safe());
+    assert!(events.iter().all(|event| !event.has_client_event()));
+    assert!(releases(&trace).is_empty());
+    release.send(()).unwrap();
+    server.await.unwrap();
+}