perf(responses): reduce WebSocket copies and trace completion latency

Move owned request bodies back into WebSocket requests and forward validated
UTF-8 frames without copying their payloads. Record quota, account feedback,
and cyber policy durations while preserving terminal delivery ordering.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-10-24 21:05:22 +08:00
co-authored by Codex
parent 2eed9b20ed
commit 6648198df4
4 changed files with 30 additions and 10 deletions
@@ -5,7 +5,7 @@ use super::{
pump::run_pump,
state::{ConnectionCommand, ConnectionStats},
};
use axum::extract::ws::{CloseFrame, Message, WebSocket, close_code};
use axum::extract::ws::{CloseFrame, Message, Utf8Bytes, WebSocket, close_code};
use futures::{Sink, Stream};
use gateway_core::{
engine::middleware::{FrozenMiddlewarePlan, MiddlewareHeader},
@@ -142,7 +142,7 @@ impl ResponsesWebSocketConnection {
/// pump 已关闭、写入超时或 transport 失败时返回稳定错误
pub async fn send_text(
&mut self,
payload: String,
payload: impl Into<Utf8Bytes>,
context: WriteContext,
) -> Result<WriteOutcome, ConnectionWriteError> {
self.send(Message::Text(payload.into()), context).await
@@ -3,7 +3,7 @@
use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use axum::http::StatusCode;
use axum::{extract::ws::Utf8Bytes, http::StatusCode};
use bytes::Bytes;
use futures::future::BoxFuture;
use gateway_core::engine::CommitRequirement;
@@ -268,7 +268,7 @@ pub(super) async fn forward_response(
return fail_body(connection, body, MiddlewareError::InvalidState, &request_id).await;
}
let terminal = frame.terminal();
let Ok(message) = String::from_utf8(frame.into_bytes().to_vec()) else {
let Ok(message) = Utf8Bytes::try_from(frame.into_bytes()) else {
return fail_body(connection, body, MiddlewareError::InvalidState, &request_id).await;
};
let phase = match (first_frame_written, terminal) {
@@ -1152,26 +1152,42 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
session_transport_recovery.websocket_succeeded(key);
}
if allows_account_state_mutation && (completed || terminal_failure.is_some()) {
let started = Instant::now();
synchronize_passive_quota(
&quota,
&active_account,
passive_quota_observation.rate_limits(),
)
.await;
trace.record("provider.quota_sync", json!({
"durationMs": started.elapsed().as_millis(),
}));
}
if let Some((failure, _)) = terminal_failure.as_ref() {
let started = Instant::now();
apply_failure(&failure_context, &active_account, failure)
.await;
trace.record("provider.account_failure", json!({
"durationMs": started.elapsed().as_millis(),
}));
}
attach_openai_session_update(&mut events, &mut session_capture);
if allows_account_state_mutation && completed && terminal_failure.is_none() {
// 完成事实交给下游前更新账号健康,绑定已经在发送前完成
let started = Instant::now();
selector
.record_success(&active_account)
.await;
trace.record("provider.account_success", json!({
"durationMs": started.elapsed().as_millis(),
}));
let started = Instant::now();
selector
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
.await;
trace.record("provider.cyber_policy_success", json!({
"durationMs": started.elapsed().as_millis(),
}));
}
if (rate_limits_changed
|| response_model_changed
@@ -318,12 +318,16 @@ impl CodexBackendClient {
Some(&headers),
);
}
*websocket_request.body_mut() = body.as_object().cloned().ok_or_else(|| {
CodexClientError::Privacy(gateway_core::settings::privacy::PrivacyError {
rule_index: 0,
reason: "WS 请求正文必须是对象",
})
})?;
// 正文已归当前请求所有,移回对象避免再次复制完整对话
let serde_json::Value::Object(body) = body else {
return Err(CodexClientError::Privacy(
gateway_core::settings::privacy::PrivacyError {
rule_index: 0,
reason: "WS 请求正文必须是对象",
},
));
};
*websocket_request.body_mut() = body;
let mut websocket_create = CodexWebSocketConnection::responses_create_request_for_path(
&self.base_url,
self.protocol.responses_path(),