From 47efce7a9049d375f9f9cd1ccfd547f142c82db5 Mon Sep 17 00:00:00 2001 From: zyycn Date: Sun, 26 Jul 2026 20:05:55 +0800 Subject: [PATCH] perf(openai): cut websocket hot-path parsing and copies Parse each WebSocket frame once in the exchange reducer and reuse the parsed JSON across rate-limit, metadata, turn-state, and forwarding extraction (previously up to eight full parses per frame). Serialize the response.create frame by borrowing the request body instead of deep-copying it (byte-identical output, including a client-supplied type key), skip audit artifact construction entirely when no audit directory is configured, and take pending turn-state updates instead of cloning them on every chunk. Co-authored-by: Claude --- .../crates/providers/openai/src/provider.rs | 5 +- .../openai/src/transport/client_sse.rs | 20 +++-- .../src/transport/protocol/websocket.rs | 83 +++++++++++-------- .../openai/src/transport/websocket/audit.rs | 23 +++-- .../transport/websocket/exchange/reducer.rs | 30 +++---- .../openai/src/transport/websocket/mod.rs | 2 +- .../tests/transport/protocol/websocket.rs | 41 ++++++--- 7 files changed, 127 insertions(+), 77 deletions(-) diff --git a/backend/crates/providers/openai/src/provider.rs b/backend/crates/providers/openai/src/provider.rs index 6f4746772..9134f5397 100644 --- a/backend/crates/providers/openai/src/provider.rs +++ b/backend/crates/providers/openai/src/provider.rs @@ -879,7 +879,8 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { }; if let (Some(capture), Some(updates)) = (session_capture.as_mut(), turn_state_updates.as_ref()) - && let Some(turn_state) = updates.lock().await.clone() + // take：仅在上游写入过新 turn_state 时更新 capture，避免逐 chunk 克隆。 + && let Some(turn_state) = updates.lock().await.take() { capture.turn_state = Some(turn_state); } @@ -996,7 +997,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { } if let (Some(capture), Some(updates)) = (session_capture.as_mut(), turn_state_updates.as_ref()) - && let Some(turn_state) = updates.lock().await.clone() + && let Some(turn_state) = updates.lock().await.take() { capture.turn_state = Some(turn_state); } diff --git a/backend/crates/providers/openai/src/transport/client_sse.rs b/backend/crates/providers/openai/src/transport/client_sse.rs index 4f54f4745..f3888445d 100644 --- a/backend/crates/providers/openai/src/transport/client_sse.rs +++ b/backend/crates/providers/openai/src/transport/client_sse.rs @@ -31,7 +31,7 @@ use crate::transport::{ CodexWebSocketConnection, CodexWebSocketPool, CodexWebSocketPoolKey, DEFAULT_INITIAL_EVENT_TIMEOUT, WEBSOCKET_FAST_PATH_BUDGET, WebSocketOriginBreaker, WebSocketPoolDecision, execute_prepared_response_create_request_stream, - post_send_ambiguous, prepare_response_create_request_with_pool, + post_send_ambiguous, prepare_response_create_request_with_pool, websocket_audit_dir, write_websocket_audit_artifact_from_env, }, }; @@ -176,13 +176,17 @@ impl CodexBackendClient { &websocket_request, ) .map_err(CodexClientError::WebSocketEncode)?; - let artifact = websocket_audit_artifact_from_attempt( - &websocket_request, - websocket_create.connection().opening_audit_snapshot(), - websocket_payload_audit_snapshot(&websocket_request), - ); - if let Err(error) = write_websocket_audit_artifact_from_env(&artifact).await { - tracing::warn!(error = %error, "Failed to write Codex WebSocket audit artifact"); + // 审计未启用时跳过 artifact 构造：payload 快照会深拷贝整个请求 body， + // 且位于首字节前的关键路径上。 + if websocket_audit_dir().is_some() { + let artifact = websocket_audit_artifact_from_attempt( + &websocket_request, + websocket_create.connection().opening_audit_snapshot(), + websocket_payload_audit_snapshot(&websocket_request), + ); + if let Err(error) = write_websocket_audit_artifact_from_env(&artifact).await { + tracing::warn!(error = %error, "Failed to write Codex WebSocket audit artifact"); + } } let connection_profile = websocket_connection_profile(&headers); let pool_key = diff --git a/backend/crates/providers/openai/src/transport/protocol/websocket.rs b/backend/crates/providers/openai/src/transport/protocol/websocket.rs index a91b173c2..c628c137a 100644 --- a/backend/crates/providers/openai/src/transport/protocol/websocket.rs +++ b/backend/crates/providers/openai/src/transport/protocol/websocket.rs @@ -50,43 +50,32 @@ pub struct WebSocketAuditArtifact { /// 将一条公开 WebSocket JSON 事件编码为 SSE 帧。 pub fn websocket_event_to_sse_frame(raw: &str) -> Option { - let event = websocket_event_type(raw)?; - if is_internal_websocket_event(&event) { - return None; - } - if websocket_event_should_skip(raw) { - return None; - } - Some(encode_sse_event(&event, raw)) + let value = serde_json::from_str::(raw).ok()?; + websocket_event_frame(&value, raw) } -pub(crate) fn websocket_event_type(raw: &str) -> Option { - serde_json::from_str::(raw).ok().and_then(|value| { - value - .get("type") - .and_then(Value::as_str) - .map(ToString::to_string) - }) +/// 同上，复用已解析的事件 JSON；`data` 段逐字节内嵌 `raw` 原文。 +/// +/// 仅剥离传输层内部帧：`codex.rate_limits` 与承载 turn_state 的 +/// `response.metadata` 不下发客户端；业务事件一律原样转发，不做 schema 校验。 +pub(crate) fn websocket_event_frame(value: &Value, raw: &str) -> Option { + let event = websocket_event_type(value)?; + if is_internal_websocket_event(event) || event == "response.metadata" { + return None; + } + Some(encode_sse_event(event, raw)) +} + +pub(crate) fn websocket_event_type(value: &Value) -> Option<&str> { + value.get("type").and_then(Value::as_str) } fn is_internal_websocket_event(event: &str) -> bool { event == "codex.rate_limits" } -/// 判断 WebSocket 事件是否应在转发前剥离。 -/// -/// 仅剥离传输层内部帧：`response.metadata` 承载 turn_state，由上层提取转存到 -/// 会话状态，不下发客户端。业务事件一律原样转发，不做 schema 校验。 -/// （`codex.rate_limits` 内部事件由 `is_internal_websocket_event` 单独剥离。） -fn websocket_event_should_skip(raw: &str) -> bool { - websocket_event_type(raw).as_deref() == Some("response.metadata") -} - /// 提取 `response.metadata` 帧中的字符串响应头。 -pub fn websocket_metadata_headers(raw: &str) -> Vec<(String, String)> { - let Ok(value) = serde_json::from_str::(raw) else { - return Vec::new(); - }; +pub fn websocket_metadata_headers(value: &Value) -> Vec<(String, String)> { if value.get("type").and_then(Value::as_str) != Some("response.metadata") { return Vec::new(); } @@ -100,9 +89,8 @@ pub fn websocket_metadata_headers(raw: &str) -> Vec<(String, String)> { } /// 从 `response.metadata` 帧中提取 `x-codex-turn-state`。 -pub fn websocket_metadata_turn_state(raw: &str) -> Option { - let value = serde_json::from_str::(raw).ok()?; - websocket_metadata_headers(raw) +pub fn websocket_metadata_turn_state(value: &Value) -> Option { + websocket_metadata_headers(value) .into_iter() .find_map(|(name, value)| { name.eq_ignore_ascii_case("x-codex-turn-state") @@ -128,8 +116,7 @@ fn json_value_as_string(value: &Value) -> Option { /// 从 `response.completed` 旁路提取 response ID，供连接池记录续接能力。 /// /// 该值不参与客户端 wire 的可交付性判断；无法读取时只是不记录连接内续接状态。 -pub fn websocket_response_completed_id(raw: &str) -> Option { - let value = serde_json::from_str::(raw).ok()?; +pub fn websocket_response_completed_id(value: &Value) -> Option { if value.get("type").and_then(Value::as_str) != Some("response.completed") { return None; } @@ -179,11 +166,39 @@ pub fn websocket_response_create_payload(request: &CodexResponsesRequest) -> Val Value::Object(payload) } +/// 借用原始 body 序列化 `response.create` 帧，不复制字段。 +/// +/// 输出与序列化 [`websocket_response_create_payload`] 的合并 Map 逐字节一致： +/// body 自带 `type` 键时以 body 值置于首位，其余字段按插入顺序原样透传。 +struct ResponseCreateFrame<'a> { + body: &'a Map, +} + +impl serde::Serialize for ResponseCreateFrame<'_> { + fn serialize(&self, serializer: S) -> Result { + use serde::ser::SerializeMap as _; + let mut map = serializer.serialize_map(Some(self.body.len() + 1))?; + match self.body.get("type") { + Some(body_type) => map.serialize_entry("type", body_type)?, + None => map.serialize_entry("type", "response.create")?, + } + for (key, value) in self.body { + if key == "type" { + continue; + } + map.serialize_entry(key, value)?; + } + map.end() + } +} + /// 生成 Responses WebSocket `response.create` 文本帧内容。 pub fn websocket_response_create_payload_text( request: &CodexResponsesRequest, ) -> Result { - serde_json::to_string(&websocket_response_create_payload(request)) + serde_json::to_string(&ResponseCreateFrame { + body: request.body(), + }) } fn websocket_payload_keys(request: &CodexResponsesRequest) -> Vec { diff --git a/backend/crates/providers/openai/src/transport/websocket/audit.rs b/backend/crates/providers/openai/src/transport/websocket/audit.rs index 7af156f0f..46bf3ecbe 100644 --- a/backend/crates/providers/openai/src/transport/websocket/audit.rs +++ b/backend/crates/providers/openai/src/transport/websocket/audit.rs @@ -3,6 +3,7 @@ use std::{ io, path::{Path, PathBuf}, + sync::OnceLock, }; use chrono::Utc; @@ -19,6 +20,19 @@ const REDACTED_HEADER_VALUE: &str = ""; /// WebSocket audit artifact 输出目录环境变量。 pub const WS_AUDIT_DIR_ENV: &str = "CODEX_PROXY_WS_AUDIT_DIR"; +static WS_AUDIT_DIR: OnceLock> = OnceLock::new(); + +/// 返回进程级审计输出目录；未配置时调用方应整体跳过 artifact 构造。 +pub fn websocket_audit_dir() -> Option<&'static Path> { + WS_AUDIT_DIR + .get_or_init(|| { + std::env::var_os(WS_AUDIT_DIR_ENV) + .filter(|value| !value.is_empty()) + .map(PathBuf::from) + }) + .as_deref() +} + /// 显式写入 WebSocket audit artifact。 pub async fn write_websocket_audit_artifact_for_dir( dir: Option<&Path>, @@ -39,14 +53,7 @@ pub async fn write_websocket_audit_artifact_for_dir( pub async fn write_websocket_audit_artifact_from_env( artifact: &WebSocketAuditArtifact, ) -> io::Result> { - let Some(dir) = std::env::var_os(WS_AUDIT_DIR_ENV) - .filter(|value| !value.is_empty()) - .map(PathBuf::from) - else { - return Ok(None); - }; - - write_websocket_audit_artifact_for_dir(Some(&dir), artifact).await + write_websocket_audit_artifact_for_dir(websocket_audit_dir(), artifact).await } impl CodexWebSocketConnection { diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/reducer.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/reducer.rs index 09b2c2945..0c5150c3f 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/reducer.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/reducer.rs @@ -1,9 +1,10 @@ //! WebSocket aggregate/stream 共用事件归约器。 use gateway_protocol::openai::events; +use serde_json::Value; use crate::transport::protocol::websocket::{ - websocket_event_to_sse_frame, websocket_event_type, websocket_metadata_headers, + websocket_event_frame, websocket_event_type, websocket_metadata_headers, websocket_metadata_turn_state, websocket_response_completed_id, }; use crate::transport::response_meta; @@ -33,40 +34,41 @@ pub(super) fn reduce_websocket_event( metadata: &mut CodexWebSocketConnectionMetadata, continuation: &mut WebSocketContinuationState, ) -> Result { - if let Some(headers) = websocket_rate_limit_event_headers(raw) { + // 每帧只解析一次 JSON，后续提取全部复用同一 Value； + // 不可解析的帧不承载可路由的事件类型，忽略。 + let Ok(value) = serde_json::from_str::(raw) else { + return Ok(ExchangeAction::Ignore); + }; + if let Some(parsed) = events::parse_rate_limits_event(&value) { + let headers = events::rate_limits_to_header_pairs(&parsed); metadata.rate_limit_headers.extend(headers.iter().cloned()); return Ok(ExchangeAction::RateLimits(headers)); } response_meta::merge_response_metadata( &mut metadata.response_metadata, - websocket_metadata_headers(raw), + websocket_metadata_headers(&value), ); - if let Some(turn_state) = websocket_metadata_turn_state(raw) { + if let Some(turn_state) = websocket_metadata_turn_state(&value) { metadata.turn_state = Some(turn_state.clone()); return Ok(ExchangeAction::TurnState(turn_state)); } - let event = websocket_event_type(raw); - if event.as_deref() == Some("response.completed") - && let Some(response_id) = websocket_response_completed_id(raw) + let event = websocket_event_type(&value); + if event == Some("response.completed") + && let Some(response_id) = websocket_response_completed_id(&value) { continuation.record_completed(response_id); } - let terminal = match event.as_deref() { + let terminal = match event { Some("response.completed") => Some(WebSocketTerminalKind::Completed), Some("response.incomplete") => Some(WebSocketTerminalKind::Incomplete), Some("response.failed" | "error") => Some(WebSocketTerminalKind::Failed), _ => None, }; - Ok(match websocket_event_to_sse_frame(raw) { + Ok(match websocket_event_frame(&value, raw) { Some(frame) => ExchangeAction::Forward { frame, terminal }, None => ExchangeAction::Ignore, }) } - -fn websocket_rate_limit_event_headers(raw: &str) -> Option> { - events::parse_rate_limits_event_raw(raw) - .map(|parsed| events::rate_limits_to_header_pairs(&parsed)) -} diff --git a/backend/crates/providers/openai/src/transport/websocket/mod.rs b/backend/crates/providers/openai/src/transport/websocket/mod.rs index 6d1d892c3..e291cca82 100644 --- a/backend/crates/providers/openai/src/transport/websocket/mod.rs +++ b/backend/crates/providers/openai/src/transport/websocket/mod.rs @@ -12,7 +12,7 @@ mod pump; pub use self::{ audit::{ - WS_AUDIT_DIR_ENV, write_websocket_audit_artifact_for_dir, + WS_AUDIT_DIR_ENV, websocket_audit_dir, write_websocket_audit_artifact_for_dir, write_websocket_audit_artifact_from_env, }, breaker::{ diff --git a/backend/crates/providers/openai/tests/transport/protocol/websocket.rs b/backend/crates/providers/openai/tests/transport/protocol/websocket.rs index dd861f3d1..46f691fd5 100644 --- a/backend/crates/providers/openai/tests/transport/protocol/websocket.rs +++ b/backend/crates/providers/openai/tests/transport/protocol/websocket.rs @@ -90,6 +90,32 @@ fn websocket_response_create_payload_should_keep_explicit_empty_instructions() { assert_eq!(value["instructions"], ""); } +#[test] +fn websocket_response_create_payload_text_should_match_merged_map_serialization() { + // body 自带 type 键与未知字段：文本帧必须与「合并 Map 再序列化」逐字节一致。 + let request = CodexResponsesRequest::from_body( + json!({ + "model": "gpt-test", + "type": "client-supplied", + "x_unknown": {"nested": [1, 2.5, "3"]}, + "input": [] + }) + .as_object() + .expect("request object") + .clone(), + ); + + let merged = serde_json::to_string( + &provider_openai::transport::protocol::websocket::websocket_response_create_payload( + &request, + ), + ) + .expect("merged payload"); + let streamed = websocket_response_create_payload_text(&request).expect("streamed payload"); + + assert_eq!(streamed, merged); +} + #[test] fn required_websocket_audit_should_forbid_http_fallback() { let request = CodexResponsesRequest::from_body( @@ -138,8 +164,7 @@ fn websocket_metadata_turn_state_should_accept_case_insensitive_header() { let event = json!({ "type": "response.metadata", "headers": {"X-Codex-Turn-State": ["turn-from-metadata"]} - }) - .to_string(); + }); assert_eq!( websocket_metadata_turn_state(&event).as_deref(), @@ -155,16 +180,14 @@ fn websocket_completed_id_should_read_the_id_without_validating_the_response_sha "id": "resp_valid", "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2} } - }) - .to_string(); + }); let invalid_usage = json!({ "type": "response.completed", "response": { "id": "resp_invalid", "usage": {"input_tokens": "bad", "output_tokens": 1, "total_tokens": 1} } - }) - .to_string(); + }); assert_eq!( websocket_response_completed_id(&valid), @@ -226,16 +249,14 @@ fn websocket_completed_id_should_leave_unreadable_completion_untracked() { "response": { "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2} } - }) - .to_string(); + }); let incomplete_usage = json!({ "type": "response.completed", "response": { "id": "resp_incomplete_usage", "usage": {"input_tokens": 1, "output_tokens": 1} } - }) - .to_string(); + }); assert_eq!(websocket_response_completed_id(&missing_id), None); assert_eq!(