mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
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:
@@ -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(
|
||||
"a,
|
||||
&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(),
|
||||
|
||||
Reference in New Issue
Block a user