From 3bf7030fba6273553cfbd021fd7d3390b81d0308 Mon Sep 17 00:00:00 2001 From: Zyy <39424122+zyycn@users.noreply.github.com> Date: Sun, 27 Sep 2026 19:51:41 +0800 Subject: [PATCH] fix(openai): widen and trace stream replay grace (#304) Co-authored-by: Zyy Co-authored-by: Codex --- .../openai/src/provider/execution.rs | 4 +- .../providers/openai/src/provider/failure.rs | 62 +++- .../providers/openai/src/provider/mod.rs | 7 +- .../openai/tests/provider/contract/mod.rs | 9 +- .../tests/provider/contract/precommit.rs | 317 ++++++++++++++++++ 5 files changed, 372 insertions(+), 27 deletions(-) create mode 100644 backend/crates/providers/openai/tests/provider/contract/precommit.rs diff --git a/backend/crates/providers/openai/src/provider/execution.rs b/backend/crates/providers/openai/src/provider/execution.rs index 25b17f34f..e91986e86 100644 --- a/backend/crates/providers/openai/src/provider/execution.rs +++ b/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; diff --git a/backend/crates/providers/openai/src/provider/failure.rs b/backend/crates/providers/openai/src/provider/failure.rs index b2ff52d40..eda5225e1 100644 --- a/backend/crates/providers/openai/src/provider/failure.rs +++ b/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) { /// 提交边界前的上游事件预取。 /// -/// 原始 chunk 计数而不是重编码后的 event 大小。时间与 64 KiB 共同限定无感换号 +/// 原始 chunk 计数而不是重编码后的 event 大小。时间与字节阈值共同限定无感换号 /// 窗口；任一边界到达都会提交已缓存 wire，不能因网关私有资源规则伪造上游协议 /// 失败。一旦提交，后续事件不再具备无痕重放资格。 pub(super) struct PreCommitClientEvents { pending: Vec, prefetched_bytes: usize, - replay_grace_deadline: Option, + replay_grace_started_at: Option, 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) -> Vec { 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 { - self.replay_grace_deadline + pub(super) fn replay_grace_deadline(&self) -> Option { + 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 { + pub(super) fn commit_pending(&mut self, reason: PreCommitReleaseReason) -> Vec { + // 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) } } diff --git a/backend/crates/providers/openai/src/provider/mod.rs b/backend/crates/providers/openai/src/provider/mod.rs index cbeef99c7..19352c686 100644 --- a/backend/crates/providers/openai/src/provider/mod.rs +++ b/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); diff --git a/backend/crates/providers/openai/tests/provider/contract/mod.rs b/backend/crates/providers/openai/tests/provider/contract/mod.rs index 77f161743..8b4d90366 100644 --- a/backend/crates/providers/openai/tests/provider/contract/mod.rs +++ b/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) diff --git a/backend/crates/providers/openai/tests/provider/contract/precommit.rs b/backend/crates/providers/openai/tests/provider/contract/precommit.rs new file mode 100644 index 000000000..7392b9137 --- /dev/null +++ b/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 { + 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, 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(); +}