starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-06 16:06:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(openai): harden downstream websocket delivery

Browse Source 
Move downstream socket ownership into a bounded pump, preserve official Codex control-frame behavior, and bound downstream writes.

Enforce the official connection lifetime after active terminal delivery, add structured transport diagnostics, and mirror connection/protocol coverage under the test tree without adding stream_id.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-08-31 03:27:11 +08:00
co-authored by Codex
parent 6b2d335bc5
commit 75e22d7a5b
10 changed files with 2129 additions and 703 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-api/Cargo.toml
+2 -2
View File
@@ -24,7 +24,7 @@ secrecy = "0.10"
serde = { version = "1.0.228", features = ["derive"] }
serde_json = { version = "1.0.150", features = ["arbitrary_precision", "preserve_order"] }
thiserror = "2.0.18"
tokio = { version = "1.52.3", default-features = false, features = ["macros", "time"] }
tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt", "sync", "time"] }
tower-http = { version = "0.7.0", features = ["cors", "fs", "request-id", "timeout", "trace"] }
tracing = "0.1.44"
url = "2.5.8"
@@ -33,7 +33,7 @@ uuid = { version = "1.23.3", features = ["v7"] }
[dev-dependencies]
async-trait = "0.1.89"
axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio"] }
tokio = { version = "1.52.3", default-features = false, features = ["macros", "net", "rt-multi-thread"] }
tokio = { version = "1.52.3", default-features = false, features = ["macros", "net", "rt-multi-thread", "test-util"] }
tokio-tungstenite = "0.29.0"
tower = { version = "0.5.3", features = ["util"] }
backend/crates/gateway-api/src/openai/responses/websocket.rs
-699
View File
@@ -1,699 +0,0 @@
//! OpenAI Responses WebSocket adapter。
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use axum::{
extract::{
Extension, State,
connect_info::ConnectInfo,
ws::{CloseFrame, Message, WebSocket, WebSocketUpgrade, close_code},
},
http::{HeaderMap, HeaderName, StatusCode},
response::{IntoResponse, Response},
};
use futures::StreamExt;
use gateway_core::engine::execution::{AuthenticatedClient, ClientTransport, StartedExecution};
use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError};
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::event::ProviderResponseHeader;
use gateway_core::lifecycle::{ConnectionGuard, ConnectionLifecycle};
use gateway_core::operation::ProviderSessionState;
use serde_json::{Map, Value, json};
use thiserror::Error;
use crate::{
ApiState,
openai::{
auth::{authenticate_client, authentication_error_response},
error::{gateway_error_contract, gateway_error_from_engine, runtime_unavailable_response},
service::OpenAiService,
},
};
use super::{
DecodedResponsesRequest, OpenAiResponsesEncoder, ProtocolErrorBody, RequestDecodeError,
http::{PendingExecution, request_client_context},
request::OpenAiRequestHeaders,
};
const ACTIVE_RESPONSE_VIOLATION: &str = "Only one response.create may be active per connection";
const TEXT_FRAMES_ONLY: &str = "Responses WebSocket accepts text frames only";
/// 将已认证的 `GET /v1/responses` 升级为 Responses WebSocket。
pub(crate) async fn responses_websocket(
State(state): State<ApiState>,
connect_info: Option<Extension<ConnectInfo<SocketAddr>>>,
headers: HeaderMap,
websocket: WebSocketUpgrade,
) -> Response {
let service = state.openai().clone();
let client = match authenticate_client(&service, &headers) {
Ok(client) => client,
Err(error) => return authentication_error_response(error),
};
let (client_ip, user_agent) = request_client_context(
&headers,
connect_info.map(|Extension(ConnectInfo(address))| address),
);
ResponsesWebSocketAdapter::new(service).upgrade_with_client_context(
websocket,
client,
client_ip,
user_agent,
OpenAiRequestHeaders::from_headers(&headers),
)
}
/// 已鉴权 Responses WebSocket 升级边界。
#[derive(Clone)]
pub(crate) struct ResponsesWebSocketAdapter {
service: OpenAiService,
}
impl ResponsesWebSocketAdapter {
/// 绑定应用提供的唯一 OpenAI 客户端服务端口。
#[must_use]
pub const fn new(service: OpenAiService) -> Self {
Self { service }
}
fn upgrade_with_client_context(
&self,
websocket: WebSocketUpgrade,
client: AuthenticatedClient,
client_ip: Option<IpAddr>,
user_agent: Option<String>,
request_headers: OpenAiRequestHeaders,
) -> Response {
let connection_guard = match self.service.try_register_connection() {
Ok(guard) => guard,
Err(_) => return runtime_unavailable_response().into_response(),
};
let session = ResponsesWebSocketSession {
service: self.service.clone(),
client,
connection_id: self.service.next_request_id().replacen("req_", "ws_", 1),
client_ip,
user_agent,
request_headers,
lifecycle: self.service.lifecycle(),
_connection_guard: connection_guard,
};
websocket
// 覆盖 axum/tungstenite 的私有 64 MiB message 与 16 MiB frame 默认值。
// Responses JSON 的协议可接受性由上游决定，代理不另设 wire 长度上限。
.max_message_size(usize::MAX)
.max_frame_size(usize::MAX)
.on_upgrade(move |socket| async move {
serve_responses_websocket(socket, session).await;
})
}
}
struct ResponsesWebSocketSession {
service: OpenAiService,
client: AuthenticatedClient,
connection_id: String,
client_ip: Option<IpAddr>,
user_agent: Option<String>,
request_headers: OpenAiRequestHeaders,
lifecycle: Arc<dyn ConnectionLifecycle>,
_connection_guard: Box<dyn ConnectionGuard>,
}
async fn serve_responses_websocket(mut socket: WebSocket, session: ResponsesWebSocketSession) {
tracing::info!(
websocket_connection_id = %session.connection_id,
"Responses WebSocket connected"
);
let mut request_count = 0_u64;
let mut replay = ConnectionReplaySnapshot::default();
let cancellation = session.lifecycle.cancellation();
loop {
let message = tokio::select! {
() = cancellation.cancelled() => break,
message = socket.next() => message,
};
let Some(message) = message else {
break;
};
let payload = match message {
Ok(Message::Text(payload)) => payload,
Ok(Message::Ping(_) | Message::Pong(_)) => continue,
Ok(Message::Close(_)) => break,
Ok(Message::Binary(_)) => {
close_for_protocol_violation(&mut socket, TEXT_FRAMES_ONLY).await;
break;
}
Err(error) => {
tracing::info!(
websocket_connection_id = %session.connection_id,
error = %error,
"Responses WebSocket receive failed"
);
break;
}
};
request_count = request_count.saturating_add(1);
let correlation_id = session.service.next_request_id();
let decoded =
match decode_response_create_with_context(payload.as_str(), &session.request_headers) {
Ok(decoded) => {
decoded.with_client_context(session.client_ip, session.user_agent.clone())
}
Err(error) => {
tracing::info!(
websocket_connection_id = %session.connection_id,
request_id = %correlation_id,
error = %error,
"Responses WebSocket request rejected"
);
if send_protocol_error(
&mut socket,
StatusCode::BAD_REQUEST,
error.protocol_body(),
&correlation_id,
)
.await
== ForwardOutcome::Disconnect
{
break;
}
continue;
}
};
let decoded = replay.prepare(decoded);
let started = match session
.service
.start_response(
session.client.clone(),
decoded,
ClientTransport::WebSocket,
"/v1/responses",
)
.await
{
Ok(started) => started,
Err(error) => {
if send_gateway_error(&mut socket, &error, &correlation_id).await
== ForwardOutcome::Disconnect
{
break;
}
continue;
}
};
if forward_execution(&mut socket, started, &mut replay).await == ForwardOutcome::Disconnect
{
break;
}
}
tracing::info!(
websocket_connection_id = %session.connection_id,
request_count,
"Responses WebSocket disconnected"
);
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum ForwardOutcome {
Continue,
Disconnect,
}
#[derive(Default)]
struct ConnectionReplaySnapshot {
last_response_id: Option<String>,
provider_state: Option<ProviderSessionState>,
}
impl ConnectionReplaySnapshot {
fn prepare(&self, request: DecodedResponsesRequest) -> DecodedResponsesRequest {
match (
request.metadata().continuation().previous_response_id(),
self.last_response_id.as_deref(),
self.provider_state.as_ref(),
) {
(Some(previous), Some(last), Some(state)) if previous == last => {
request.with_provider_session_state(state.clone())
}
_ => request,
}
}
fn commit(&mut self, response_id: String, provider_state: Option<ProviderSessionState>) {
self.last_response_id = Some(response_id);
self.provider_state = provider_state;
}
}
async fn forward_execution(
socket: &mut WebSocket,
started: StartedExecution,
replay: &mut ConnectionReplaySnapshot,
) -> ForwardOutcome {
let request_id = started.request_id.to_string();
let streaming = started.stream;
let mut execution = PendingExecution::new(started.session);
if !streaming {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"WebSocket execution was not initialized as a stream",
);
return send_gateway_error(socket, &error, &request_id).await;
}
let first = loop {
match next_active_input(socket, &mut execution).await {
ActiveInput::Event(Ok(Some(event))) => break event,
ActiveInput::Event(Ok(None)) => {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway response ended before its first event",
);
return send_gateway_error(socket, &error, &request_id).await;
}
ActiveInput::Event(Err(error)) => {
let error = gateway_error_from_engine(&error);
return send_gateway_error(socket, &error, &request_id).await;
}
ActiveInput::Control => continue,
ActiveInput::Disconnect => return ForwardOutcome::Disconnect,
ActiveInput::ProtocolViolation => {
close_for_protocol_violation(socket, ACTIVE_RESPONSE_VIOLATION).await;
return ForwardOutcome::Disconnect;
}
}
};
let requirement = first.commit_requirement();
let mut first = first.into_provider_events();
if requirement != CommitRequirement::CommitBeforeDelivery {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway first event did not require commit",
);
return send_gateway_error(socket, &error, &request_id).await;
}
let mut encoder = OpenAiResponsesEncoder::new();
let mut provider_state = None;
let mut first_messages = Vec::new();
for event in &mut first {
if let Some(update) = event.take_session_update() {
provider_state = Some(update);
}
first_messages.extend(encoder.push_websocket(event));
}
if first_messages.is_empty() {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway commit batch encoded no output",
);
return send_gateway_error(socket, &error, &request_id).await;
}
let Some(response_session) = execution.session_mut() else {
return ForwardOutcome::Disconnect;
};
let response_headers = response_session.response_headers().to_vec();
if let Err(error) = response_session.commit_downstream(None).await {
let error = gateway_error_from_engine(&error);
return send_gateway_error(socket, &error, &request_id).await;
}
if encoder.is_completed() {
if let Err(outcome) = confirm_completed_execution(socket, &mut execution, &request_id).await
{
return outcome;
}
commit_connection_replay(replay, &encoder, provider_state.take());
if !send_text(
socket,
response_metadata_event(&request_id, &response_headers),
)
.await
|| !send_messages(socket, first_messages).await
{
return ForwardOutcome::Disconnect;
}
execution.disarm();
return ForwardOutcome::Continue;
}
if !send_text(
socket,
response_metadata_event(&request_id, &response_headers),
)
.await
|| !send_messages(socket, first_messages).await
{
return ForwardOutcome::Disconnect;
}
loop {
match next_active_input(socket, &mut execution).await {
ActiveInput::Event(Ok(Some(delivery))) => {
let requirement = delivery.commit_requirement();
let mut events = delivery.into_provider_events();
if requirement != CommitRequirement::AlreadyCommitted {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway requested another downstream commit",
);
return send_gateway_error(socket, &error, &request_id).await;
}
let mut messages = Vec::new();
for event in &mut events {
if let Some(update) = event.take_session_update() {
provider_state = Some(update);
}
messages.extend(encoder.push_websocket(event));
}
if encoder.is_completed() {
if let Err(outcome) =
confirm_completed_execution(socket, &mut execution, &request_id).await
{
return outcome;
}
commit_connection_replay(replay, &encoder, provider_state.take());
if !send_messages(socket, messages).await {
return ForwardOutcome::Disconnect;
}
execution.disarm();
return ForwardOutcome::Continue;
}
if !send_messages(socket, messages).await {
return ForwardOutcome::Disconnect;
}
}
ActiveInput::Event(Ok(None)) => {
if execution
.session_mut()
.is_some_and(|session| session.is_finalized())
{
commit_connection_replay(replay, &encoder, provider_state.take());
execution.disarm();
return ForwardOutcome::Continue;
}
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway response ended without finalizing the execution",
);
return send_gateway_error(socket, &error, &request_id).await;
}
ActiveInput::Event(Err(error)) => {
if encoder.has_wire_failure() {
execution.disarm();
return ForwardOutcome::Continue;
}
let error = gateway_error_from_engine(&error);
return send_gateway_error(socket, &error, &request_id).await;
}
ActiveInput::Control => {}
ActiveInput::Disconnect => return ForwardOutcome::Disconnect,
ActiveInput::ProtocolViolation => {
close_for_protocol_violation(socket, ACTIVE_RESPONSE_VIOLATION).await;
return ForwardOutcome::Disconnect;
}
}
}
}
fn commit_connection_replay(
replay: &mut ConnectionReplaySnapshot,
encoder: &OpenAiResponsesEncoder,
provider_state: Option<ProviderSessionState>,
) {
if let Some(response_id) = encoder.response_id() {
replay.commit(response_id.to_owned(), provider_state);
}
}
async fn confirm_completed_execution(
socket: &mut WebSocket,
execution: &mut PendingExecution,
request_id: &str,
) -> Result<(), ForwardOutcome> {
loop {
match next_active_input(socket, execution).await {
ActiveInput::Event(Ok(None))
if execution
.session_mut()
.is_some_and(|session| session.is_finalized()) =>
{
return Ok(());
}
ActiveInput::Event(Ok(None)) => {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway response was not finalized after its terminal event",
);
return Err(send_gateway_error(socket, &error, request_id).await);
}
ActiveInput::Event(Ok(Some(_))) => {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway response continued after its terminal event",
);
return Err(send_gateway_error(socket, &error, request_id).await);
}
ActiveInput::Event(Err(error)) => {
let error = gateway_error_from_engine(&error);
return Err(send_gateway_error(socket, &error, request_id).await);
}
ActiveInput::Control => {}
ActiveInput::Disconnect => return Err(ForwardOutcome::Disconnect),
ActiveInput::ProtocolViolation => {
close_for_protocol_violation(socket, ACTIVE_RESPONSE_VIOLATION).await;
return Err(ForwardOutcome::Disconnect);
}
}
}
}
enum ActiveInput {
Event(Result<Option<CoordinatedEvent>, EngineError>),
Control,
Disconnect,
ProtocolViolation,
}
async fn next_active_input(
socket: &mut WebSocket,
execution: &mut PendingExecution,
) -> ActiveInput {
let Some(session) = execution.session_mut() else {
return ActiveInput::Disconnect;
};
tokio::select! {
event = session.next_event() => ActiveInput::Event(event),
message = socket.next() => match message {
Some(Ok(Message::Ping(_) | Message::Pong(_))) => ActiveInput::Control,
Some(Ok(Message::Close(_))) | None | Some(Err(_)) => ActiveInput::Disconnect,
Some(Ok(Message::Text(_) | Message::Binary(_))) => ActiveInput::ProtocolViolation,
},
}
}
/// 使用连接级请求头和 Provider 上下文解码官方 `response.create` 文本帧。
///
/// 缺省 `stream` 等价于 WebSocket 固有的流式语义；显式 `false` 会被拒绝。
///
/// # Errors
///
/// 帧不是合法 JSON object、消息类型错误、显式关闭 stream，或 Responses 请求
/// 无法映射到 canonical operation 时返回不包含正文内容的稳定错误。
pub fn decode_response_create_with_context(
payload: &str,
request_headers: &OpenAiRequestHeaders,
) -> Result<DecodedResponsesRequest, ResponseCreateFrameError> {
decode_response_create_inner(payload, request_headers)
}
fn decode_response_create_inner(
payload: &str,
request_headers: &OpenAiRequestHeaders,
) -> Result<DecodedResponsesRequest, ResponseCreateFrameError> {
let Value::Object(mut body) = serde_json::from_str::<Value>(payload)
.map_err(|_| ResponseCreateFrameError::InvalidJson)?
else {
return Err(ResponseCreateFrameError::ExpectedObject);
};
match body.remove("type") {
Some(Value::String(message_type)) if message_type == "response.create" => {}
_ => return Err(ResponseCreateFrameError::UnsupportedType),
}
if matches!(body.get("stream"), Some(value) if value.as_bool() != Some(true)) {
return Err(ResponseCreateFrameError::StreamingRequired);
}
super::request::decode_request_object(body, false, request_headers)
.map_err(ResponseCreateFrameError::Request)
}
/// `response.create` 帧的稳定安全错误。
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum ResponseCreateFrameError {
/// 文本不是合法 JSON。
#[error("response.create frame must be valid JSON")]
InvalidJson,
/// 顶层不是 object。
#[error("response.create frame must be a JSON object")]
ExpectedObject,
/// `type` 缺失或不是 `response.create`。
#[error("unsupported Responses WebSocket message type")]
UnsupportedType,
/// WebSocket 请求显式声明 `stream=false`。
#[error("Responses WebSocket requests require stream=true")]
StreamingRequired,
/// 内层 Responses 请求无法映射到 canonical operation。
#[error(transparent)]
Request(RequestDecodeError),
}
impl ResponseCreateFrameError {
fn protocol_body(&self) -> ProtocolErrorBody {
match self {
Self::Request(error) => error.protocol_body(),
Self::InvalidJson => RequestDecodeError::MalformedJson.protocol_body(),
Self::ExpectedObject => RequestDecodeError::ExpectedObject.protocol_body(),
Self::UnsupportedType => RequestDecodeError::InvalidValue {
field: "type".to_owned(),
}
.protocol_body(),
Self::StreamingRequired => RequestDecodeError::InvalidValue {
field: "stream".to_owned(),
}
.protocol_body(),
}
}
}
fn response_metadata_event(
request_id: &str,
response_headers: &[ProviderResponseHeader],
) -> String {
let connection_options = super::response_connection_options(response_headers);
let mut headers = Map::new();
for header in response_headers {
if !super::response_header_is_forwardable(header.name(), &connection_options) {
continue;
}
let Ok(name) = HeaderName::from_bytes(header.name().as_bytes()) else {
continue;
};
let Ok(value) = std::str::from_utf8(header.value()) else {
continue;
};
// 官方 response.metadata 使用 string map，无法表达同名多值；后值保持既有覆盖语义。
headers.insert(name.as_str().to_owned(), Value::String(value.to_owned()));
}
headers
.entry("x-request-id".to_owned())
.or_insert_with(|| Value::String(request_id.to_owned()));
json!({
"type": "response.metadata",
"headers": headers,
})
.to_string()
}
async fn send_gateway_error(
socket: &mut WebSocket,
error: &GatewayError,
request_id: &str,
) -> ForwardOutcome {
let (status, default_type, default_code) = gateway_error_contract(error.kind());
if send_error_event(
socket,
status,
error.client_error_type().unwrap_or(default_type),
error.client_error_code().unwrap_or(default_code),
error.client_message(),
None,
request_id,
)
.await
{
ForwardOutcome::Continue
} else {
ForwardOutcome::Disconnect
}
}
async fn send_protocol_error(
socket: &mut WebSocket,
status: StatusCode,
body: ProtocolErrorBody,
request_id: &str,
) -> ForwardOutcome {
let error = body.error;
if send_error_event(
socket,
status,
error.kind,
error.code,
&error.message,
error.param.as_deref(),
request_id,
)
.await
{
ForwardOutcome::Continue
} else {
ForwardOutcome::Disconnect
}
}
async fn send_error_event(
socket: &mut WebSocket,
status: StatusCode,
error_type: &str,
code: &str,
message: &str,
param: Option<&str>,
request_id: &str,
) -> bool {
let mut error = Map::new();
error.insert("type".to_owned(), Value::String(error_type.to_owned()));
error.insert("code".to_owned(), Value::String(code.to_owned()));
error.insert("message".to_owned(), Value::String(message.to_owned()));
if let Some(param) = param {
error.insert("param".to_owned(), Value::String(param.to_owned()));
}
send_text(
socket,
json!({
"type": "error",
"status": status.as_u16(),
"error": error,
"request_id": request_id,
})
.to_string(),
)
.await
}
async fn send_messages(socket: &mut WebSocket, messages: Vec<String>) -> bool {
for message in messages {
if !send_text(socket, message).await {
return false;
}
}
true
}
async fn send_text(socket: &mut WebSocket, payload: String) -> bool {
if let Err(error) = socket.send(Message::Text(payload.into())).await {
tracing::info!(error = %error, "Responses WebSocket send failed");
return false;
}
true
}
async fn close_for_protocol_violation(socket: &mut WebSocket, reason: &'static str) {
let _ = socket
.send(Message::Close(Some(CloseFrame {
code: close_code::POLICY,
reason: reason.into(),
})))
.await;
}
backend/crates/gateway-api/src/openai/responses/websocket/connection.rs
+550
View File
@@ -0,0 +1,550 @@
//! 下游客户端 WebSocket 的单 owner pump 与有界收发边界。
use std::{
fmt,
sync::{
Arc,
atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
},
time::Duration,
};
use axum::extract::ws::{CloseFrame, Message, WebSocket, close_code};
use futures::{Sink, SinkExt, Stream, StreamExt};
use gateway_core::engine::CancellationToken;
use thiserror::Error;
use tokio::{
sync::{mpsc, oneshot},
task::JoinHandle,
time::{Instant, timeout},
};
const OUTBOUND_COMMAND_BUFFER: usize = 32;
const INBOUND_EVENT_BUFFER: usize = 32;
const DOWNSTREAM_WRITE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
const CONNECTION_MAX_AGE: Duration = Duration::from_secs(60 * 60);
#[derive(Clone, Copy)]
pub(super) struct ConnectionConfig {
pub(super) write_timeout: Duration,
pub(super) max_age: Duration,
}
impl ConnectionConfig {
pub(super) const PRODUCTION: Self = Self {
write_timeout: DOWNSTREAM_WRITE_TIMEOUT,
max_age: CONNECTION_MAX_AGE,
};
}
/// 业务层可观察的客户端输入；Ping/Pong 始终由 pump 消费。
pub(super) enum ConnectionEvent {
Text(String),
Binary,
Expired,
Exited(PumpExitReason),
}
/// 下游写入阶段；名称刻意使用 write，而不是暗示客户端已消费的 delivery。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum FramePhase {
Metadata,
First,
Data,
Terminal,
FirstAndTerminal,
Error,
ConnectionLimit,
Close,
}
impl FramePhase {
const fn as_str(self) -> &'static str {
match self {
Self::Metadata => "metadata",
Self::First => "first",
Self::Data => "data",
Self::Terminal => "terminal",
Self::FirstAndTerminal => "first_and_terminal",
Self::Error => "error",
Self::ConnectionLimit => "connection_limit",
Self::Close => "close",
}
}
const fn is_milestone(self) -> bool {
matches!(
self,
Self::First
| Self::Terminal
| Self::FirstAndTerminal
| Self::Error
| Self::ConnectionLimit
| Self::Close
)
}
}
#[derive(Clone)]
pub(super) struct WriteContext {
request_id: Option<Arc<str>>,
phase: FramePhase,
}
impl WriteContext {
pub(super) fn request(request_id: &Arc<str>, phase: FramePhase) -> Self {
Self {
request_id: Some(Arc::clone(request_id)),
phase,
}
}
pub(super) const fn connection(phase: FramePhase) -> Self {
Self {
request_id: None,
phase,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum PumpExitReason {
ClientClose,
PeerEof,
ReadError,
WriteError,
WriteTimeout,
LifecycleShutdown,
ConnectionMaxAge,
CoordinatorDropped,
InboundOverload,
ServerClose,
PumpStopped,
}
impl PumpExitReason {
pub(super) const fn as_str(self) -> &'static str {
match self {
Self::ClientClose => "client_close",
Self::PeerEof => "peer_eof",
Self::ReadError => "read_error",
Self::WriteError => "write_error",
Self::WriteTimeout => "write_timeout",
Self::LifecycleShutdown => "lifecycle_shutdown",
Self::ConnectionMaxAge => "connection_max_age",
Self::CoordinatorDropped => "coordinator_dropped",
Self::InboundOverload => "inbound_overload",
Self::ServerClose => "server_close",
Self::PumpStopped => "pump_stopped",
}
}
}
#[derive(Debug, Error)]
pub(super) enum ConnectionWriteError {
#[error("downstream WebSocket pump is closed")]
Closed,
#[error("downstream WebSocket write timed out after {timeout:?}")]
Timeout { timeout: Duration },
#[error("downstream WebSocket transport write failed: {message}")]
Transport { message: String },
}
struct ConnectionCommand {
message: Message,
context: WriteContext,
acknowledged: oneshot::Sender<Result<(), ConnectionWriteError>>,
}
#[derive(Default)]
struct ConnectionStats {
command_queue_high_water: AtomicUsize,
command_backpressure_count: AtomicU64,
ping_received_count: AtomicU64,
pong_written_count: AtomicU64,
}
impl ConnectionStats {
fn observe_command_queue(&self, sender: &mpsc::Sender<ConnectionCommand>) {
let queued = OUTBOUND_COMMAND_BUFFER
.saturating_sub(sender.capacity())
.saturating_add(1)
.min(OUTBOUND_COMMAND_BUFFER);
self.command_queue_high_water
.fetch_max(queued, Ordering::Relaxed);
if sender.capacity() == 0 {
self.command_backpressure_count
.fetch_add(1, Ordering::Relaxed);
}
}
}
pub(super) struct ResponsesWebSocketConnection {
connection_id: Arc<str>,
opened_at: Instant,
expired: Arc<AtomicBool>,
commands: Option<mpsc::Sender<ConnectionCommand>>,
incoming: mpsc::Receiver<ConnectionEvent>,
pump_task: Option<JoinHandle<()>>,
stats: Arc<ConnectionStats>,
config: ConnectionConfig,
exit_reason: Option<PumpExitReason>,
}
impl ResponsesWebSocketConnection {
pub(super) fn new(
socket: WebSocket,
connection_id: String,
cancellation: CancellationToken,
) -> Self {
spawn_connection(
socket,
Arc::<str>::from(connection_id),
cancellation,
ConnectionConfig::PRODUCTION,
)
}
pub(super) fn id(&self) -> &str {
&self.connection_id
}
pub(super) fn is_expired(&self) -> bool {
self.expired.load(Ordering::Acquire)
}
pub(super) fn age(&self) -> Duration {
self.opened_at.elapsed()
}
pub(super) async fn next_event(&mut self) -> Option<ConnectionEvent> {
let event = self.incoming.recv().await;
if let Some(ConnectionEvent::Exited(reason)) = event.as_ref() {
self.exit_reason.get_or_insert(*reason);
} else if event.is_none() {
self.exit_reason.get_or_insert(PumpExitReason::PumpStopped);
}
event
}
pub(super) async fn send_text(
&mut self,
payload: String,
context: WriteContext,
) -> Result<(), ConnectionWriteError> {
self.send(Message::Text(payload.into()), context).await
}
pub(super) async fn close_policy(
&mut self,
reason: &'static str,
request_id: Option<&Arc<str>>,
) {
let context = request_id.map_or_else(
|| WriteContext::connection(FramePhase::Close),
|request_id| WriteContext::request(request_id, FramePhase::Close),
);
self.close(
CloseFrame {
code: close_code::POLICY,
reason: reason.into(),
},
context,
PumpExitReason::ServerClose,
)
.await;
}
pub(super) async fn close_for_connection_limit(&mut self, reason: &'static str) {
self.close(
CloseFrame {
code: close_code::NORMAL,
reason: reason.into(),
},
WriteContext::connection(FramePhase::Close),
PumpExitReason::ConnectionMaxAge,
)
.await;
}
async fn close(
&mut self,
frame: CloseFrame,
context: WriteContext,
exit_reason: PumpExitReason,
) {
if self
.send(Message::Close(Some(frame)), context)
.await
.is_ok()
{
self.exit_reason.get_or_insert(exit_reason);
}
}
async fn send(
&mut self,
message: Message,
context: WriteContext,
) -> Result<(), ConnectionWriteError> {
let Some(commands) = self.commands.as_ref().cloned() else {
return Err(ConnectionWriteError::Closed);
};
self.stats.observe_command_queue(&commands);
let (acknowledged, acknowledgement) = oneshot::channel();
let started_at = Instant::now();
let request_id = context.request_id.clone();
let phase = context.phase;
let command = ConnectionCommand {
message,
context,
acknowledged,
};
let write = async move {
commands
.send(command)
.await
.map_err(|_| ConnectionWriteError::Closed)?;
acknowledgement
.await
.map_err(|_| ConnectionWriteError::Closed)?
};
let result = match timeout(self.config.write_timeout, write).await {
Ok(result) => result,
Err(_) => Err(ConnectionWriteError::Timeout {
timeout: self.config.write_timeout,
}),
};
let duration = started_at.elapsed();
match &result {
Ok(()) if phase.is_milestone() => tracing::info!(
websocket_connection_id = %self.connection_id,
request_id = request_id.as_deref().unwrap_or(""),
frame_phase = phase.as_str(),
write_duration_ms = duration.as_millis(),
"Responses WebSocket frame write succeeded"
),
Ok(()) => tracing::debug!(
websocket_connection_id = %self.connection_id,
request_id = request_id.as_deref().unwrap_or(""),
frame_phase = phase.as_str(),
write_duration_ms = duration.as_millis(),
"Responses WebSocket frame write succeeded"
),
Err(error) => tracing::info!(
websocket_connection_id = %self.connection_id,
request_id = request_id.as_deref().unwrap_or(""),
frame_phase = phase.as_str(),
write_duration_ms = duration.as_millis(),
error = %error,
"Responses WebSocket frame write failed"
),
}
if let Err(error) = &result {
let reason = match error {
ConnectionWriteError::Timeout { .. } => PumpExitReason::WriteTimeout,
ConnectionWriteError::Transport { .. } => PumpExitReason::WriteError,
ConnectionWriteError::Closed => PumpExitReason::PumpStopped,
};
self.terminate(reason);
}
result
}
fn terminate(&mut self, reason: PumpExitReason) {
self.exit_reason.get_or_insert(reason);
self.commands.take();
if let Some(task) = self.pump_task.take() {
task.abort();
}
}
pub(super) fn log_summary(&self, request_count: u64) {
tracing::info!(
websocket_connection_id = %self.connection_id,
request_count,
connection_age_ms = self.age().as_millis(),
pump_exit_reason = self
.exit_reason
.map_or("unknown", PumpExitReason::as_str),
command_queue_high_water = self
.stats
.command_queue_high_water
.load(Ordering::Relaxed),
command_backpressure_count = self
.stats
.command_backpressure_count
.load(Ordering::Relaxed),
ping_received_count = self.stats.ping_received_count.load(Ordering::Relaxed),
pong_written_count = self.stats.pong_written_count.load(Ordering::Relaxed),
"Responses WebSocket disconnected"
);
}
}
impl Drop for ResponsesWebSocketConnection {
fn drop(&mut self) {
self.commands.take();
if let Some(task) = self.pump_task.take() {
task.abort();
}
}
}
pub(super) fn spawn_connection<S, E>(
socket: S,
connection_id: Arc<str>,
cancellation: CancellationToken,
config: ConnectionConfig,
) -> ResponsesWebSocketConnection
where
S: Stream<Item = Result<Message, E>> + Sink<Message, Error = E> + Unpin + Send + 'static,
E: fmt::Display + Send + 'static,
{
let opened_at = Instant::now();
let expired = Arc::new(AtomicBool::new(false));
let stats = Arc::new(ConnectionStats::default());
let (command_tx, command_rx) = mpsc::channel(OUTBOUND_COMMAND_BUFFER);
let (incoming_tx, incoming_rx) = mpsc::channel(INBOUND_EVENT_BUFFER);
let pump_task = tokio::spawn(run_pump(
socket,
command_rx,
incoming_tx,
Arc::clone(&connection_id),
cancellation,
opened_at,
Arc::clone(&expired),
Arc::clone(&stats),
config,
));
ResponsesWebSocketConnection {
connection_id,
opened_at,
expired,
commands: Some(command_tx),
incoming: incoming_rx,
pump_task: Some(pump_task),
stats,
config,
exit_reason: None,
}
}
#[expect(clippy::too_many_arguments)]
async fn run_pump<S, E>(
mut socket: S,
mut commands: mpsc::Receiver<ConnectionCommand>,
incoming: mpsc::Sender<ConnectionEvent>,
connection_id: Arc<str>,
cancellation: CancellationToken,
opened_at: Instant,
expired: Arc<AtomicBool>,
stats: Arc<ConnectionStats>,
config: ConnectionConfig,
) where
S: Stream<Item = Result<Message, E>> + Sink<Message, Error = E> + Unpin,
E: fmt::Display,
{
let deadline = tokio::time::sleep_until(opened_at + config.max_age);
tokio::pin!(deadline);
let mut deadline_elapsed = false;
let reason = loop {
tokio::select! {
biased;
() = cancellation.cancelled() => break PumpExitReason::LifecycleShutdown,
() = &mut deadline, if !deadline_elapsed => {
deadline_elapsed = true;
expired.store(true, Ordering::Release);
if let Err(reason) = emit_incoming(&incoming, ConnectionEvent::Expired) {
break reason;
}
}
command = commands.recv() => {
let Some(command) = command else {
break PumpExitReason::CoordinatorDropped;
};
let ConnectionCommand {
message,
context,
acknowledged,
} = command;
let closing = matches!(message, Message::Close(_));
let result = socket.send(message).await.map_err(|error| {
let message = error.to_string();
tracing::debug!(
websocket_connection_id = %connection_id,
request_id = context.request_id.as_deref().unwrap_or(""),
frame_phase = context.phase.as_str(),
error = %message,
"Responses WebSocket pump transport write failed"
);
ConnectionWriteError::Transport { message }
});
let failed = result.is_err();
let _ = acknowledged.send(result);
if failed {
break PumpExitReason::WriteError;
}
if closing {
break PumpExitReason::ServerClose;
}
}
message = socket.next() => match message {
Some(Ok(Message::Text(payload))) => {
if let Err(reason) = emit_incoming(&incoming, ConnectionEvent::Text(payload.to_string())) {
break reason;
}
}
Some(Ok(Message::Binary(_))) => {
if let Err(reason) = emit_incoming(&incoming, ConnectionEvent::Binary) {
break reason;
}
}
Some(Ok(Message::Ping(payload))) => {
stats.ping_received_count.fetch_add(1, Ordering::Relaxed);
match timeout(config.write_timeout, socket.send(Message::Pong(payload))).await {
Ok(Ok(())) => {
stats.pong_written_count.fetch_add(1, Ordering::Relaxed);
}
Ok(Err(error)) => {
tracing::info!(
websocket_connection_id = %connection_id,
error = %error,
"Responses WebSocket Pong write failed"
);
break PumpExitReason::WriteError;
}
Err(_) => break PumpExitReason::WriteTimeout,
}
}
Some(Ok(Message::Pong(_))) => {}
Some(Ok(Message::Close(_))) => break PumpExitReason::ClientClose,
Some(Err(error)) => {
tracing::info!(
websocket_connection_id = %connection_id,
error = %error,
"Responses WebSocket receive failed"
);
break PumpExitReason::ReadError;
}
None => break PumpExitReason::PeerEof,
},
}
};
let _ = incoming.try_send(ConnectionEvent::Exited(reason));
tracing::debug!(
websocket_connection_id = %connection_id,
pump_exit_reason = reason.as_str(),
"Responses WebSocket pump exited"
);
}
fn emit_incoming(
incoming: &mpsc::Sender<ConnectionEvent>,
event: ConnectionEvent,
) -> Result<(), PumpExitReason> {
incoming.try_send(event).map_err(|error| match error {
mpsc::error::TrySendError::Full(_) => PumpExitReason::InboundOverload,
mpsc::error::TrySendError::Closed(_) => PumpExitReason::CoordinatorDropped,
})
}
backend/crates/gateway-api/src/openai/responses/websocket/forward.rs
+482
View File
@@ -0,0 +1,482 @@
//! 已启动 Responses execution 到客户端 WebSocket wire event 的串行转发。
use std::sync::Arc;
use axum::http::StatusCode;
use gateway_core::engine::execution::StartedExecution;
use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError};
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::operation::ProviderSessionState;
use tokio::time::Instant;
use crate::openai::error::{gateway_error_contract, gateway_error_from_engine};
use super::{
super::{DecodedResponsesRequest, OpenAiResponsesEncoder, PendingExecution, ProtocolErrorBody},
connection::{ConnectionEvent, FramePhase, ResponsesWebSocketConnection, WriteContext},
protocol::{error_event, response_metadata_event},
};
const ACTIVE_RESPONSE_VIOLATION: &str = "Only one response.create may be active per connection";
#[derive(Clone, Copy, PartialEq, Eq)]
pub(super) enum ForwardOutcome {
Continue,
Disconnect,
}
#[derive(Default)]
pub(super) struct ConnectionReplaySnapshot {
last_response_id: Option<String>,
provider_state: Option<ProviderSessionState>,
}
impl ConnectionReplaySnapshot {
pub(super) fn prepare(&self, request: DecodedResponsesRequest) -> DecodedResponsesRequest {
match (
request.metadata().continuation().previous_response_id(),
self.last_response_id.as_deref(),
self.provider_state.as_ref(),
) {
(Some(previous), Some(last), Some(state)) if previous == last => {
request.with_provider_session_state(state.clone())
}
_ => request,
}
}
fn commit(&mut self, response_id: String, provider_state: Option<ProviderSessionState>) {
self.last_response_id = Some(response_id);
self.provider_state = provider_state;
}
}
pub(super) async fn forward_execution(
connection: &mut ResponsesWebSocketConnection,
started: StartedExecution,
replay: &mut ConnectionReplaySnapshot,
) -> ForwardOutcome {
let request_id = Arc::<str>::from(started.request_id.to_string());
let streaming = started.stream;
let mut execution = PendingExecution::new(started.session);
if !streaming {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"WebSocket execution was not initialized as a stream",
);
return send_gateway_error(connection, &error, &request_id).await;
}
let first = loop {
match next_active_input(connection, &mut execution).await {
ActiveInput::Event(Ok(Some(event))) => break event,
ActiveInput::Event(Ok(None)) => {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway response ended before its first event",
);
return send_gateway_error(connection, &error, &request_id).await;
}
ActiveInput::Event(Err(error)) => {
let error = gateway_error_from_engine(&error);
return send_gateway_error(connection, &error, &request_id).await;
}
ActiveInput::Control => continue,
ActiveInput::Disconnect => return ForwardOutcome::Disconnect,
ActiveInput::ProtocolViolation => {
connection
.close_policy(ACTIVE_RESPONSE_VIOLATION, Some(&request_id))
.await;
return ForwardOutcome::Disconnect;
}
}
};
let requirement = first.commit_requirement();
let mut first = first.into_provider_events();
if requirement != CommitRequirement::CommitBeforeDelivery {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway first event did not require commit",
);
return send_gateway_error(connection, &error, &request_id).await;
}
let mut encoder = OpenAiResponsesEncoder::new();
let mut provider_state = None;
let mut first_messages = Vec::new();
for event in &mut first {
if let Some(update) = event.take_session_update() {
provider_state = Some(update);
}
first_messages.extend(encoder.push_websocket(event));
}
if first_messages.is_empty() {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway commit batch encoded no output",
);
return send_gateway_error(connection, &error, &request_id).await;
}
let Some(response_session) = execution.session_mut() else {
return ForwardOutcome::Disconnect;
};
let response_headers = response_session.response_headers().to_vec();
if let Err(error) = response_session.commit_downstream(None).await {
let error = gateway_error_from_engine(&error);
return send_gateway_error(connection, &error, &request_id).await;
}
let mut first_frame_written = false;
if encoder.is_completed() {
if let Err(outcome) =
confirm_completed_execution(connection, &mut execution, &request_id).await
{
return outcome;
}
let provider_terminal_at = Instant::now();
commit_connection_replay(replay, &encoder, provider_state.take());
let writes_succeeded = send_metadata(
connection,
&request_id,
response_metadata_event(&request_id, &response_headers),
)
.await
&& send_messages(
connection,
first_messages,
&request_id,
&mut first_frame_written,
true,
)
.await;
if !writes_succeeded {
log_terminal_write_failure(connection, &request_id, provider_terminal_at.elapsed());
return ForwardOutcome::Disconnect;
}
log_terminal_write_success(connection, &request_id, provider_terminal_at.elapsed());
execution.disarm();
return ForwardOutcome::Continue;
}
if !send_metadata(
connection,
&request_id,
response_metadata_event(&request_id, &response_headers),
)
.await
|| !send_messages(
connection,
first_messages,
&request_id,
&mut first_frame_written,
encoder.has_wire_failure(),
)
.await
{
return ForwardOutcome::Disconnect;
}
loop {
match next_active_input(connection, &mut execution).await {
ActiveInput::Event(Ok(Some(delivery))) => {
let requirement = delivery.commit_requirement();
let mut events = delivery.into_provider_events();
if requirement != CommitRequirement::AlreadyCommitted {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway requested another downstream commit",
);
return send_gateway_error(connection, &error, &request_id).await;
}
let mut messages = Vec::new();
for event in &mut events {
if let Some(update) = event.take_session_update() {
provider_state = Some(update);
}
messages.extend(encoder.push_websocket(event));
}
if encoder.is_completed() {
if let Err(outcome) =
confirm_completed_execution(connection, &mut execution, &request_id).await
{
return outcome;
}
let provider_terminal_at = Instant::now();
commit_connection_replay(replay, &encoder, provider_state.take());
if !send_messages(
connection,
messages,
&request_id,
&mut first_frame_written,
true,
)
.await
{
log_terminal_write_failure(
connection,
&request_id,
provider_terminal_at.elapsed(),
);
return ForwardOutcome::Disconnect;
}
log_terminal_write_success(
connection,
&request_id,
provider_terminal_at.elapsed(),
);
execution.disarm();
return ForwardOutcome::Continue;
}
let terminal_batch = encoder.has_wire_failure();
if !send_messages(
connection,
messages,
&request_id,
&mut first_frame_written,
terminal_batch,
)
.await
{
return ForwardOutcome::Disconnect;
}
}
ActiveInput::Event(Ok(None)) => {
if execution
.session_mut()
.is_some_and(|session| session.is_finalized())
{
commit_connection_replay(replay, &encoder, provider_state.take());
execution.disarm();
return ForwardOutcome::Continue;
}
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway response ended without finalizing the execution",
);
return send_gateway_error(connection, &error, &request_id).await;
}
ActiveInput::Event(Err(error)) => {
if encoder.has_wire_failure() {
execution.disarm();
return ForwardOutcome::Continue;
}
let error = gateway_error_from_engine(&error);
return send_gateway_error(connection, &error, &request_id).await;
}
ActiveInput::Control => {}
ActiveInput::Disconnect => return ForwardOutcome::Disconnect,
ActiveInput::ProtocolViolation => {
connection
.close_policy(ACTIVE_RESPONSE_VIOLATION, Some(&request_id))
.await;
return ForwardOutcome::Disconnect;
}
}
}
}
fn commit_connection_replay(
replay: &mut ConnectionReplaySnapshot,
encoder: &OpenAiResponsesEncoder,
provider_state: Option<ProviderSessionState>,
) {
if let Some(response_id) = encoder.response_id() {
replay.commit(response_id.to_owned(), provider_state);
}
}
async fn confirm_completed_execution(
connection: &mut ResponsesWebSocketConnection,
execution: &mut PendingExecution,
request_id: &Arc<str>,
) -> Result<(), ForwardOutcome> {
loop {
match next_active_input(connection, execution).await {
ActiveInput::Event(Ok(None))
if execution
.session_mut()
.is_some_and(|session| session.is_finalized()) =>
{
return Ok(());
}
ActiveInput::Event(Ok(None)) => {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway response was not finalized after its terminal event",
);
return Err(send_gateway_error(connection, &error, request_id).await);
}
ActiveInput::Event(Ok(Some(_))) => {
let error = GatewayError::new(
GatewayErrorKind::Internal,
"gateway response continued after its terminal event",
);
return Err(send_gateway_error(connection, &error, request_id).await);
}
ActiveInput::Event(Err(error)) => {
let error = gateway_error_from_engine(&error);
return Err(send_gateway_error(connection, &error, request_id).await);
}
ActiveInput::Control => {}
ActiveInput::Disconnect => return Err(ForwardOutcome::Disconnect),
ActiveInput::ProtocolViolation => {
connection
.close_policy(ACTIVE_RESPONSE_VIOLATION, Some(request_id))
.await;
return Err(ForwardOutcome::Disconnect);
}
}
}
}
enum ActiveInput {
Event(Result<Option<CoordinatedEvent>, EngineError>),
Control,
Disconnect,
ProtocolViolation,
}
async fn next_active_input(
connection: &mut ResponsesWebSocketConnection,
execution: &mut PendingExecution,
) -> ActiveInput {
let Some(session) = execution.session_mut() else {
return ActiveInput::Disconnect;
};
tokio::select! {
event = session.next_event() => ActiveInput::Event(event),
message = connection.next_event() => match message {
Some(ConnectionEvent::Expired) => ActiveInput::Control,
Some(ConnectionEvent::Text(_) | ConnectionEvent::Binary) => {
ActiveInput::ProtocolViolation
}
Some(ConnectionEvent::Exited(_)) | None => ActiveInput::Disconnect,
},
}
}
pub(super) async fn send_gateway_error(
connection: &mut ResponsesWebSocketConnection,
error: &GatewayError,
request_id: &Arc<str>,
) -> ForwardOutcome {
let (status, default_type, default_code) = gateway_error_contract(error.kind());
send_error(
connection,
status,
error.client_error_type().unwrap_or(default_type),
error.client_error_code().unwrap_or(default_code),
error.client_message(),
None,
request_id,
)
.await
}
pub(super) async fn send_protocol_error(
connection: &mut ResponsesWebSocketConnection,
status: StatusCode,
body: ProtocolErrorBody,
request_id: &Arc<str>,
) -> ForwardOutcome {
let error = body.error;
send_error(
connection,
status,
error.kind,
error.code,
&error.message,
error.param.as_deref(),
request_id,
)
.await
}
async fn send_error(
connection: &mut ResponsesWebSocketConnection,
status: StatusCode,
error_type: &str,
code: &str,
message: &str,
param: Option<&str>,
request_id: &Arc<str>,
) -> ForwardOutcome {
if connection
.send_text(
error_event(status, error_type, code, message, param, Some(request_id)),
WriteContext::request(request_id, FramePhase::Error),
)
.await
.is_ok()
{
ForwardOutcome::Continue
} else {
ForwardOutcome::Disconnect
}
}
async fn send_metadata(
connection: &mut ResponsesWebSocketConnection,
request_id: &Arc<str>,
message: String,
) -> bool {
connection
.send_text(
message,
WriteContext::request(request_id, FramePhase::Metadata),
)
.await
.is_ok()
}
async fn send_messages(
connection: &mut ResponsesWebSocketConnection,
messages: Vec<String>,
request_id: &Arc<str>,
first_frame_written: &mut bool,
terminal_batch: bool,
) -> bool {
let message_count = messages.len();
for (index, message) in messages.into_iter().enumerate() {
let is_terminal = terminal_batch && index + 1 == message_count;
let phase = match (*first_frame_written, is_terminal) {
(false, true) => FramePhase::FirstAndTerminal,
(false, false) => FramePhase::First,
(true, true) => FramePhase::Terminal,
(true, false) => FramePhase::Data,
};
if connection
.send_text(message, WriteContext::request(request_id, phase))
.await
.is_err()
{
return false;
}
*first_frame_written = true;
}
true
}
fn log_terminal_write_success(
connection: &ResponsesWebSocketConnection,
request_id: &Arc<str>,
provider_terminal_to_write: std::time::Duration,
) {
tracing::info!(
websocket_connection_id = connection.id(),
request_id = %request_id,
provider_terminal_to_terminal_write_ms = provider_terminal_to_write.as_millis(),
terminal_frame_written = true,
"Provider execution terminal was written to the WebSocket transport"
);
}
fn log_terminal_write_failure(
connection: &ResponsesWebSocketConnection,
request_id: &Arc<str>,
provider_terminal_to_write: std::time::Duration,
) {
tracing::info!(
websocket_connection_id = connection.id(),
request_id = %request_id,
provider_terminal_to_terminal_write_ms = provider_terminal_to_write.as_millis(),
provider_succeeded_but_terminal_write_failed = true,
"Provider execution succeeded before the terminal WebSocket frame write failed"
);
}
backend/crates/gateway-api/src/openai/responses/websocket/mod.rs
+248
View File
@@ -0,0 +1,248 @@
//! OpenAI Responses WebSocket 的升级、连接与串行 session 编排。
mod connection;
mod forward;
mod protocol;
use std::{
net::{IpAddr, SocketAddr},
sync::Arc,
};
use axum::{
extract::{
Extension, State,
connect_info::ConnectInfo,
ws::{WebSocket, WebSocketUpgrade},
},
http::{HeaderMap, StatusCode},
response::{IntoResponse, Response},
};
use gateway_core::{
engine::execution::{AuthenticatedClient, ClientTransport},
lifecycle::{ConnectionGuard, ConnectionLifecycle},
};
use crate::{
ApiState,
openai::{
auth::{authenticate_client, authentication_error_response},
error::runtime_unavailable_response,
service::OpenAiService,
},
};
use super::{http::request_client_context, request::OpenAiRequestHeaders};
use connection::{ConnectionEvent, FramePhase, ResponsesWebSocketConnection, WriteContext};
use forward::{
ConnectionReplaySnapshot, ForwardOutcome, forward_execution, send_gateway_error,
send_protocol_error,
};
use protocol::connection_limit_event;
pub use protocol::{ResponseCreateFrameError, decode_response_create_with_context};
const TEXT_FRAMES_ONLY: &str = "Responses WebSocket accepts text frames only";
const CONNECTION_LIMIT_CLOSE_REASON: &str = "Responses websocket connection limit reached";
/// 将已认证的 `GET /v1/responses` 升级为 Responses WebSocket。
pub(crate) async fn responses_websocket(
State(state): State<ApiState>,
connect_info: Option<Extension<ConnectInfo<SocketAddr>>>,
headers: HeaderMap,
websocket: WebSocketUpgrade,
) -> Response {
let service = state.openai().clone();
let client = match authenticate_client(&service, &headers) {
Ok(client) => client,
Err(error) => return authentication_error_response(error),
};
let (client_ip, user_agent) = request_client_context(
&headers,
connect_info.map(|Extension(ConnectInfo(address))| address),
);
ResponsesWebSocketAdapter::new(service).upgrade_with_client_context(
websocket,
client,
client_ip,
user_agent,
OpenAiRequestHeaders::from_headers(&headers),
)
}
/// 已鉴权 Responses WebSocket 升级边界。
#[derive(Clone)]
pub(crate) struct ResponsesWebSocketAdapter {
service: OpenAiService,
}
impl ResponsesWebSocketAdapter {
/// 绑定应用提供的唯一 OpenAI 客户端服务端口。
#[must_use]
pub const fn new(service: OpenAiService) -> Self {
Self { service }
}
fn upgrade_with_client_context(
&self,
websocket: WebSocketUpgrade,
client: AuthenticatedClient,
client_ip: Option<IpAddr>,
user_agent: Option<String>,
request_headers: OpenAiRequestHeaders,
) -> Response {
let connection_guard = match self.service.try_register_connection() {
Ok(guard) => guard,
Err(_) => return runtime_unavailable_response().into_response(),
};
let session = ResponsesWebSocketSession {
service: self.service.clone(),
client,
connection_id: self.service.next_request_id().replacen("req_", "ws_", 1),
client_ip,
user_agent,
request_headers,
lifecycle: self.service.lifecycle(),
connection_guard,
};
websocket
// 覆盖 axum/tungstenite 的私有 64 MiB message 与 16 MiB frame 默认值。
// Responses JSON 的协议可接受性由上游决定，代理不另设 wire 长度上限。
.max_message_size(usize::MAX)
.max_frame_size(usize::MAX)
.on_upgrade(move |socket| async move {
serve_responses_websocket(socket, session).await;
})
}
}
struct ResponsesWebSocketSession {
service: OpenAiService,
client: AuthenticatedClient,
connection_id: String,
client_ip: Option<IpAddr>,
user_agent: Option<String>,
request_headers: OpenAiRequestHeaders,
lifecycle: Arc<dyn ConnectionLifecycle>,
connection_guard: Box<dyn ConnectionGuard>,
}
async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocketSession) {
let ResponsesWebSocketSession {
service,
client,
connection_id,
client_ip,
user_agent,
request_headers,
lifecycle,
connection_guard,
} = session;
tracing::info!(
websocket_connection_id = %connection_id,
"Responses WebSocket connected"
);
let _connection_guard = connection_guard;
let cancellation = lifecycle.cancellation();
let mut connection = ResponsesWebSocketConnection::new(socket, connection_id, cancellation);
let mut request_count = 0_u64;
let mut replay = ConnectionReplaySnapshot::default();
loop {
let Some(event) = connection.next_event().await else {
break;
};
let payload = match event {
ConnectionEvent::Text(payload) => payload,
ConnectionEvent::Binary => {
connection.close_policy(TEXT_FRAMES_ONLY, None).await;
break;
}
ConnectionEvent::Expired => {
expire_connection(&mut connection).await;
break;
}
ConnectionEvent::Exited(_) => break,
};
if connection.is_expired() {
expire_connection(&mut connection).await;
break;
}
request_count = request_count.saturating_add(1);
let correlation_id = Arc::<str>::from(service.next_request_id());
let decoded = match decode_response_create_with_context(&payload, &request_headers) {
Ok(decoded) => decoded.with_client_context(client_ip, user_agent.clone()),
Err(error) => {
tracing::info!(
websocket_connection_id = connection.id(),
request_id = %correlation_id,
error = %error,
"Responses WebSocket request rejected"
);
if send_protocol_error(
&mut connection,
StatusCode::BAD_REQUEST,
error.protocol_body(),
&correlation_id,
)
.await
== ForwardOutcome::Disconnect
{
break;
}
continue;
}
};
let decoded = replay.prepare(decoded);
// deadline 与 Text 可能同时就绪；在任何上游执行开始前再次封住该竞争窗口。
if connection.is_expired() {
expire_connection(&mut connection).await;
break;
}
let started = match service
.start_response(
client.clone(),
decoded,
ClientTransport::WebSocket,
"/v1/responses",
)
.await
{
Ok(started) => started,
Err(error) => {
if send_gateway_error(&mut connection, &error, &correlation_id).await
== ForwardOutcome::Disconnect
{
break;
}
continue;
}
};
if forward_execution(&mut connection, started, &mut replay).await
== ForwardOutcome::Disconnect
{
break;
}
if connection.is_expired() {
expire_connection(&mut connection).await;
break;
}
}
connection.log_summary(request_count);
}
async fn expire_connection(connection: &mut ResponsesWebSocketConnection) {
if connection
.send_text(
connection_limit_event(),
WriteContext::connection(FramePhase::ConnectionLimit),
)
.await
.is_ok()
{
connection
.close_for_connection_limit(CONNECTION_LIMIT_CLOSE_REASON)
.await;
}
}
backend/crates/gateway-api/src/openai/responses/websocket/protocol.rs
+151
View File
@@ -0,0 +1,151 @@
//! Responses WebSocket 入站请求与下行事件的纯协议映射。
use axum::http::{HeaderName, StatusCode};
use gateway_core::event::ProviderResponseHeader;
use serde_json::{Map, Value, json};
use thiserror::Error;
use super::super::{
DecodedResponsesRequest, ProtocolErrorBody, RequestDecodeError, request::OpenAiRequestHeaders,
};
/// 使用连接级请求头和 Provider 上下文解码官方 `response.create` 文本帧。
///
/// 缺省 `stream` 等价于 WebSocket 固有的流式语义；显式 `false` 会被拒绝。
///
/// # Errors
///
/// 帧不是合法 JSON object、消息类型错误、显式关闭 stream，或 Responses 请求
/// 无法映射到 canonical operation 时返回不包含正文内容的稳定错误。
pub fn decode_response_create_with_context(
payload: &str,
request_headers: &OpenAiRequestHeaders,
) -> Result<DecodedResponsesRequest, ResponseCreateFrameError> {
decode_response_create_inner(payload, request_headers)
}
fn decode_response_create_inner(
payload: &str,
request_headers: &OpenAiRequestHeaders,
) -> Result<DecodedResponsesRequest, ResponseCreateFrameError> {
let Value::Object(mut body) = serde_json::from_str::<Value>(payload)
.map_err(|_| ResponseCreateFrameError::InvalidJson)?
else {
return Err(ResponseCreateFrameError::ExpectedObject);
};
match body.remove("type") {
Some(Value::String(message_type)) if message_type == "response.create" => {}
_ => return Err(ResponseCreateFrameError::UnsupportedType),
}
if matches!(body.get("stream"), Some(value) if value.as_bool() != Some(true)) {
return Err(ResponseCreateFrameError::StreamingRequired);
}
super::super::request::decode_request_object(body, false, request_headers)
.map_err(ResponseCreateFrameError::Request)
}
/// `response.create` 帧的稳定安全错误。
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum ResponseCreateFrameError {
/// 文本不是合法 JSON。
#[error("response.create frame must be valid JSON")]
InvalidJson,
/// 顶层不是 object。
#[error("response.create frame must be a JSON object")]
ExpectedObject,
/// `type` 缺失或不是 `response.create`。
#[error("unsupported Responses WebSocket message type")]
UnsupportedType,
/// WebSocket 请求显式声明 `stream=false`。
#[error("Responses WebSocket requests require stream=true")]
StreamingRequired,
/// 内层 Responses 请求无法映射到 canonical operation。
#[error(transparent)]
Request(RequestDecodeError),
}
impl ResponseCreateFrameError {
pub(super) fn protocol_body(&self) -> ProtocolErrorBody {
match self {
Self::Request(error) => error.protocol_body(),
Self::InvalidJson => RequestDecodeError::MalformedJson.protocol_body(),
Self::ExpectedObject => RequestDecodeError::ExpectedObject.protocol_body(),
Self::UnsupportedType => RequestDecodeError::InvalidValue {
field: "type".to_owned(),
}
.protocol_body(),
Self::StreamingRequired => RequestDecodeError::InvalidValue {
field: "stream".to_owned(),
}
.protocol_body(),
}
}
}
pub(super) fn response_metadata_event(
request_id: &str,
response_headers: &[ProviderResponseHeader],
) -> String {
let connection_options = super::super::response_connection_options(response_headers);
let mut headers = Map::new();
for header in response_headers {
if !super::super::response_header_is_forwardable(header.name(), &connection_options) {
continue;
}
let Ok(name) = HeaderName::from_bytes(header.name().as_bytes()) else {
continue;
};
let Ok(value) = std::str::from_utf8(header.value()) else {
continue;
};
// 官方 response.metadata 使用 string map，无法表达同名多值；后值保持既有覆盖语义。
headers.insert(name.as_str().to_owned(), Value::String(value.to_owned()));
}
headers
.entry("x-request-id".to_owned())
.or_insert_with(|| Value::String(request_id.to_owned()));
json!({
"type": "response.metadata",
"headers": headers,
})
.to_string()
}
pub(super) fn error_event(
status: StatusCode,
error_type: &str,
code: &str,
message: &str,
param: Option<&str>,
request_id: Option<&str>,
) -> String {
let mut error = Map::new();
error.insert("type".to_owned(), Value::String(error_type.to_owned()));
error.insert("code".to_owned(), Value::String(code.to_owned()));
error.insert("message".to_owned(), Value::String(message.to_owned()));
if let Some(param) = param {
error.insert("param".to_owned(), Value::String(param.to_owned()));
}
let mut event = Map::new();
event.insert("type".to_owned(), Value::String("error".to_owned()));
event.insert("status".to_owned(), Value::Number(status.as_u16().into()));
event.insert("error".to_owned(), Value::Object(error));
if let Some(request_id) = request_id {
event.insert(
"request_id".to_owned(),
Value::String(request_id.to_owned()),
);
}
Value::Object(event).to_string()
}
pub(super) fn connection_limit_event() -> String {
error_event(
StatusCode::BAD_REQUEST,
"invalid_request_error",
"websocket_connection_limit_reached",
"Responses websocket connection limit reached (60 minutes). Create a new websocket connection to continue.",
None,
None,
)
}
backend/crates/gateway-api/tests/architecture.rs
+7 -2
View File
@@ -62,7 +62,10 @@ fn source_tree_should_match_frozen_machine_manifest() {
"src/openai/responses/mod.rs",
"src/openai/responses/request.rs",
"src/openai/responses/response.rs",
"src/openai/responses/websocket.rs",
"src/openai/responses/websocket/connection.rs",
"src/openai/responses/websocket/forward.rs",
"src/openai/responses/websocket/mod.rs",
"src/openai/responses/websocket/protocol.rs",
"src/openai/router.rs",
"src/openai/service.rs",
];
@@ -99,7 +102,9 @@ fn test_tree_should_match_frozen_rust_mirror() {
"tests/openai/models.rs",
"tests/openai/responses/http.rs",
"tests/openai/responses/mod.rs",
"tests/openai/responses/websocket.rs",
"tests/openai/responses/websocket/connection.rs",
"tests/openai/responses/websocket/mod.rs",
"tests/openai/responses/websocket/protocol.rs",
"tests/openai/router.rs",
];
expected.sort_unstable();
backend/crates/gateway-api/tests/openai/responses/websocket/connection.rs
+651
View File
@@ -0,0 +1,651 @@
use std::{
fmt,
pin::Pin,
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
task::{Context, Poll},
time::{Duration, SystemTime},
};
use axum::{
extract::ws::Message,
http::{StatusCode, header::AUTHORIZATION},
};
use futures::{Sink, SinkExt, Stream, StreamExt, future::BoxFuture};
use gateway_core::{
engine::{
CancellationToken, CommitRequirement, CoordinatedEvent, EngineError,
execution::{
AuthenticatedClient, ClientAuthenticationError, ExecutionService, ExecutionSession,
StartExecution, StartProviderExecution, StartedExecution,
},
},
error::GatewayError,
event::{GatewayEvent, ProtocolWireEvent, ProviderEvent, ProviderResponseHeader, ResponseMeta},
routing::PublicModelId,
};
use serde_json::{Value, json};
use tokio::sync::mpsc::{
UnboundedReceiver, UnboundedSender, error::TryRecvError, unbounded_channel,
};
use tokio_tungstenite::{
connect_async,
tungstenite::{
Message as ClientMessage, client::IntoClientRequest, protocol::frame::coding::CloseCode,
},
};
use super::{AtomicFailureExecution, AtomicFailureTrace};
use crate::openai::{api_router, authenticated_client};
#[derive(Default)]
struct CrossingLimitTrace {
release_terminal: tokio::sync::Notify,
started_emitted: AtomicBool,
terminal_emitted: AtomicBool,
finalized: AtomicBool,
}
struct CrossingLimitSession {
trace: Arc<CrossingLimitTrace>,
}
impl ExecutionSession for CrossingLimitSession {
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>> {
Box::pin(async move {
if !self.trace.started_emitted.swap(true, Ordering::AcqRel) {
return Ok(Some(crossing_limit_started_batch()));
}
if !self.trace.terminal_emitted.load(Ordering::Acquire) {
self.trace.release_terminal.notified().await;
if !self.trace.terminal_emitted.swap(true, Ordering::AcqRel) {
return Ok(Some(crossing_limit_terminal_batch()));
}
}
self.trace.finalized.store(true, Ordering::Release);
Ok(None)
})
}
fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result<Vec<ProviderEvent>, EngineError>> {
Box::pin(async { Err(EngineError::InvalidDeliveryState) })
}
fn response_headers(&self) -> &[ProviderResponseHeader] {
&[]
}
fn commit_downstream(&mut self, _: Option<u16>) -> BoxFuture<'_, Result<(), EngineError>> {
Box::pin(async { Ok(()) })
}
fn record_client_status(&mut self, _: u16) -> BoxFuture<'_, Result<(), EngineError>> {
Box::pin(async { Ok(()) })
}
fn is_finalized(&self) -> bool {
self.trace.finalized.load(Ordering::Acquire)
}
fn cancel(&self) {}
fn detach_finalize(self: Box<Self>) -> BoxFuture<'static, ()> {
Box::pin(async move {
self.trace.finalized.store(true, Ordering::Release);
})
}
}
struct CrossingLimitExecution {
client: AuthenticatedClient,
trace: Arc<CrossingLimitTrace>,
}
impl ExecutionService for CrossingLimitExecution {
fn authenticate(
&self,
plaintext: &str,
) -> Result<AuthenticatedClient, ClientAuthenticationError> {
(plaintext == "sk_ws_crossing_limit")
.then(|| self.client.clone())
.ok_or(ClientAuthenticationError::InvalidKey)
}
fn public_models(&self, _: &AuthenticatedClient) -> Vec<PublicModelId> {
vec![PublicModelId::new("model-a").expect("public model")]
}
fn contains_public_model(&self, _: &AuthenticatedClient, model: &PublicModelId) -> bool {
model.as_str() == "model-a"
}
fn start(
&self,
request: StartExecution,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
Box::pin(async move {
Ok(StartedExecution {
request_id: gateway_core::engine::ModelRequestId::new("req_ws_crossing_limit")
.expect("request id"),
created_at: SystemTime::now(),
stream: request.metadata.stream,
session: Box::new(CrossingLimitSession {
trace: Arc::clone(&self.trace),
}),
})
})
}
fn start_provider_endpoint(
&self,
_: StartProviderExecution,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
Box::pin(async { unreachable!("WebSocket test does not execute provider endpoints") })
}
}
fn crossing_limit_started_batch() -> CoordinatedEvent {
let started = ProviderEvent::canonical_with_wire(
vec![GatewayEvent::Started(ResponseMeta::new(
"resp_ws_crossing_limit",
"model-a",
))],
ProtocolWireEvent::json(
"openai",
Some("response.created".to_owned()),
json!({
"type": "response.created",
"response": {
"id": "resp_ws_crossing_limit",
"model": "model-a",
"status": "in_progress"
}
}),
)
.expect("created wire"),
);
CoordinatedEvent::try_batch(vec![started], CommitRequirement::CommitBeforeDelivery)
.expect("started WebSocket batch")
}
fn crossing_limit_terminal_batch() -> CoordinatedEvent {
let completed = ProviderEvent::wire(
ProtocolWireEvent::json(
"openai",
Some("response.completed".to_owned()),
json!({
"type": "response.completed",
"response": {
"id": "resp_ws_crossing_limit",
"model": "model-a",
"status": "completed"
}
}),
)
.expect("completed wire"),
);
CoordinatedEvent::try_batch(vec![completed], CommitRequirement::AlreadyCommitted)
.expect("terminal WebSocket batch")
}
// 连接 runtime 没有正式公共 API；镜像测试直接编译该私有模块，以测试 socket pump，
// 不把测试钩子暴露给 gateway-api 的消费者。
#[expect(
dead_code,
reason = "the mirrored test only exercises pump-facing seams"
)]
#[path = "../../../../src/openai/responses/websocket/connection.rs"]
mod connection_under_test;
use connection_under_test::{
ConnectionConfig, ConnectionEvent, ConnectionWriteError, FramePhase,
ResponsesWebSocketConnection, WriteContext, spawn_connection,
};
#[derive(Debug)]
struct TestSocketError;
impl fmt::Display for TestSocketError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("test socket error")
}
}
struct TestSocket {
incoming: UnboundedReceiver<Result<Message, TestSocketError>>,
written: UnboundedSender<Message>,
stall_writes: bool,
dropped: Arc<AtomicBool>,
}
impl Stream for TestSocket {
type Item = Result<Message, TestSocketError>;
fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.incoming.poll_recv(context)
}
}
impl Sink<Message> for TestSocket {
type Error = TestSocketError;
fn poll_ready(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
if self.stall_writes {
Poll::Pending
} else {
Poll::Ready(Ok(()))
}
}
fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
self.written.send(item).map_err(|_| TestSocketError)
}
fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
}
impl Drop for TestSocket {
fn drop(&mut self) {
self.dropped.store(true, Ordering::Release);
}
}
struct PumpHarness {
connection: ResponsesWebSocketConnection,
incoming: UnboundedSender<Result<Message, TestSocketError>>,
written: UnboundedReceiver<Message>,
dropped: Arc<AtomicBool>,
}
fn test_connection(stall_writes: bool) -> PumpHarness {
let (incoming_tx, incoming_rx) = unbounded_channel();
let (written_tx, written_rx) = unbounded_channel();
let dropped = Arc::new(AtomicBool::new(false));
let socket = TestSocket {
incoming: incoming_rx,
written: written_tx,
stall_writes,
dropped: Arc::clone(&dropped),
};
let connection = spawn_connection(
socket,
Arc::from("ws_test"),
CancellationToken::new(),
ConnectionConfig::PRODUCTION,
);
PumpHarness {
connection,
incoming: incoming_tx,
written: written_rx,
dropped,
}
}
#[tokio::test]
async fn pump_replies_to_ping_with_the_same_pong_payload() {
let PumpHarness {
connection,
incoming,
mut written,
..
} = test_connection(false);
incoming
.send(Ok(Message::Ping(vec![1, 2, 3].into())))
.expect("send Ping to test socket");
let message = written.recv().await.expect("pump must write Pong");
let Message::Pong(payload) = message else {
panic!("pump wrote a non-Pong control frame");
};
assert_eq!(payload.as_ref(), &[1, 2, 3]);
drop(connection);
}
#[tokio::test(start_paused = true)]
async fn pump_consumes_pong_without_forwarding_a_business_event() {
let PumpHarness {
mut connection,
incoming,
..
} = test_connection(false);
incoming
.send(Ok(Message::Pong(vec![4, 5, 6].into())))
.expect("send Pong to test socket");
let result = tokio::time::timeout(Duration::from_secs(1), connection.next_event()).await;
assert!(result.is_err(), "Pong unexpectedly reached the coordinator");
}
#[tokio::test(start_paused = true)]
async fn pump_does_not_emit_an_active_ping() {
let PumpHarness {
connection,
incoming: _incoming,
mut written,
..
} = test_connection(false);
tokio::time::advance(Duration::from_secs(5 * 60)).await;
tokio::task::yield_now().await;
assert!(matches!(written.try_recv(), Err(TryRecvError::Empty)));
drop(connection);
}
#[tokio::test]
async fn outbound_commands_are_written_in_acknowledged_order() {
let PumpHarness {
mut connection,
incoming: _incoming,
mut written,
..
} = test_connection(false);
let request_id = Arc::<str>::from("req_test");
connection
.send_text(
"first".to_owned(),
WriteContext::request(&request_id, FramePhase::First),
)
.await
.expect("first write acknowledgement");
connection
.send_text(
"second".to_owned(),
WriteContext::request(&request_id, FramePhase::Terminal),
)
.await
.expect("second write acknowledgement");
let first = written.recv().await.expect("first wire message");
let second = written.recv().await.expect("second wire message");
assert!(
matches!(first, Message::Text(payload) if payload.as_str() == "first")
&& matches!(second, Message::Text(payload) if payload.as_str() == "second")
);
}
#[tokio::test(start_paused = true)]
async fn stalled_write_times_out_at_three_hundred_seconds_and_aborts_the_pump() {
let PumpHarness {
mut connection,
incoming: _incoming,
written: _written,
dropped,
} = test_connection(true);
let request_id = Arc::<str>::from("req_test");
let error = connection
.send_text(
"payload".to_owned(),
WriteContext::request(&request_id, FramePhase::First),
)
.await
.expect_err("stalled write must time out");
tokio::task::yield_now().await;
assert!(
matches!(error, ConnectionWriteError::Timeout { timeout } if timeout == Duration::from_secs(300))
&& dropped.load(Ordering::Acquire)
);
}
#[tokio::test(start_paused = true)]
async fn pump_marks_the_connection_expired_at_sixty_minutes() {
let PumpHarness {
mut connection,
incoming: _incoming,
written: _written,
..
} = test_connection(false);
let event = connection
.next_event()
.await
.expect("deadline must emit an event");
assert!(matches!(event, ConnectionEvent::Expired) && connection.is_expired());
}
#[tokio::test]
async fn dropping_the_connection_aborts_and_drops_the_socket_owner() {
let PumpHarness {
connection,
incoming: _incoming,
written: _written,
dropped,
} = test_connection(false);
drop(connection);
tokio::task::yield_now().await;
assert!(dropped.load(Ordering::Acquire));
}
#[tokio::test(start_paused = true)]
async fn idle_connection_reaches_the_official_limit_without_starting_an_execution() {
let trace = Arc::new(AtomicFailureTrace::default());
let execution = Arc::new(AtomicFailureExecution {
client: authenticated_client("sk_ws_atomic"),
trace: Arc::clone(&trace),
response_headers: Vec::new(),
});
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind WebSocket test server");
let address = listener.local_addr().expect("test server address");
let server = tokio::spawn(async move {
axum::serve(listener, app)
.await
.expect("serve WebSocket test router");
});
let mut request = format!("ws://{address}/v1/responses")
.into_client_request()
.expect("WebSocket request");
request.headers_mut().insert(
AUTHORIZATION,
"Bearer sk_ws_atomic".parse().expect("authorization"),
);
let (mut socket, response) = connect_async(request).await.expect("upgrade WebSocket");
assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS);
tokio::time::advance(Duration::from_secs(60 * 60)).await;
tokio::task::yield_now().await;
let message = tokio::time::timeout(Duration::from_secs(1), socket.next())
.await
.expect("connection limit response timeout")
.expect("connection remains available for the limit error")
.expect("valid WebSocket frame");
let ClientMessage::Text(text) = message else {
panic!("connection limit must be a text error event");
};
let value = serde_json::from_str::<Value>(&text).expect("connection limit JSON");
assert_eq!(
value,
json!({
"type": "error",
"status": 400,
"error": {
"type": "invalid_request_error",
"code": "websocket_connection_limit_reached",
"message": "Responses websocket connection limit reached (60 minutes). Create a new websocket connection to continue."
}
})
);
assert_eq!(trace.starts.load(Ordering::Acquire), 0);
socket.close(None).await.expect("close WebSocket");
server.abort();
}
#[tokio::test]
async fn active_response_crossing_the_limit_writes_terminal_before_the_limit_error() {
let trace = Arc::new(CrossingLimitTrace::default());
let execution = Arc::new(CrossingLimitExecution {
client: authenticated_client("sk_ws_crossing_limit"),
trace: Arc::clone(&trace),
});
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind WebSocket test server");
let address = listener.local_addr().expect("test server address");
let server = tokio::spawn(async move {
axum::serve(listener, app)
.await
.expect("serve WebSocket test router");
});
let mut request = format!("ws://{address}/v1/responses")
.into_client_request()
.expect("WebSocket request");
request.headers_mut().insert(
AUTHORIZATION,
"Bearer sk_ws_crossing_limit"
.parse()
.expect("authorization"),
);
let (mut socket, response) = connect_async(request).await.expect("upgrade WebSocket");
assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS);
socket
.send(ClientMessage::Text(
json!({
"type": "response.create",
"model": "model-a",
"input": "hello"
})
.to_string()
.into(),
))
.await
.expect("send response.create");
let mut initial_types = Vec::new();
loop {
tokio::task::yield_now().await;
let message = tokio::time::timeout(Duration::from_secs(1), socket.next())
.await
.unwrap_or_else(|_| panic!("initial response timeout after {initial_types:?}"))
.expect("WebSocket remains open")
.expect("valid WebSocket frame");
let ClientMessage::Text(text) = message else {
panic!("unexpected initial WebSocket frame after {initial_types:?}: {message:?}");
};
let value = serde_json::from_str::<Value>(&text).expect("initial response JSON");
let event_type = value.get("type").and_then(Value::as_str);
initial_types.extend(event_type.map(str::to_owned));
if event_type == Some("response.created") {
break;
}
}
tokio::time::pause();
tokio::time::advance(Duration::from_secs(60 * 60)).await;
tokio::task::yield_now().await;
tokio::time::resume();
trace.release_terminal.notify_one();
let mut event_types = Vec::new();
loop {
tokio::task::yield_now().await;
let message = tokio::time::timeout(Duration::from_secs(1), socket.next())
.await
.expect("post-limit response timeout")
.expect("WebSocket remains open through the limit error")
.expect("valid WebSocket frame");
let ClientMessage::Text(text) = message else {
continue;
};
let value = serde_json::from_str::<Value>(&text).expect("post-limit response JSON");
let Some(event_type) = value.get("type").and_then(Value::as_str) else {
continue;
};
event_types.push(event_type.to_owned());
if event_type == "error" {
break;
}
}
assert_eq!(event_types, ["response.completed", "error"]);
assert!(trace.finalized.load(Ordering::Acquire));
server.abort();
}
#[tokio::test]
async fn second_text_frame_during_an_active_response_keeps_the_single_response_contract() {
let trace = Arc::new(CrossingLimitTrace::default());
let execution = Arc::new(CrossingLimitExecution {
client: authenticated_client("sk_ws_crossing_limit"),
trace,
});
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind WebSocket test server");
let address = listener.local_addr().expect("test server address");
let server = tokio::spawn(async move {
axum::serve(listener, app)
.await
.expect("serve WebSocket test router");
});
let mut request = format!("ws://{address}/v1/responses")
.into_client_request()
.expect("WebSocket request");
request.headers_mut().insert(
AUTHORIZATION,
"Bearer sk_ws_crossing_limit"
.parse()
.expect("authorization"),
);
let (mut socket, _) = connect_async(request).await.expect("upgrade WebSocket");
let response_create = ClientMessage::Text(
json!({
"type": "response.create",
"model": "model-a",
"input": "hello"
})
.to_string()
.into(),
);
socket
.send(response_create.clone())
.await
.expect("send first response.create");
loop {
let message = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.expect("initial response timeout")
.expect("WebSocket remains open")
.expect("valid WebSocket frame");
let ClientMessage::Text(text) = message else {
continue;
};
let value = serde_json::from_str::<Value>(&text).expect("initial response JSON");
if value.get("type").and_then(Value::as_str) == Some("response.created") {
break;
}
}
socket
.send(response_create)
.await
.expect("send overlapping response.create");
let message = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.expect("policy close timeout")
.expect("WebSocket returns a policy close")
.expect("valid WebSocket frame");
let ClientMessage::Close(Some(frame)) = message else {
panic!("overlapping response.create must close the connection");
};
assert_eq!(frame.code, CloseCode::Policy);
server.abort();
}
backend/crates/gateway-api/tests/openai/responses/websocket.rs → backend/crates/gateway-api/tests/openai/responses/websocket/mod.rs
+5
View File
@@ -1,3 +1,6 @@
mod connection;
mod protocol;
use std::sync::{
Arc,
atomic::{AtomicBool, AtomicUsize, Ordering},
@@ -171,6 +174,7 @@ fn response_create_should_reject_non_boolean_stream_without_disclosing_body_valu
#[derive(Default)]
struct AtomicFailureTrace {
starts: AtomicUsize,
next_calls: AtomicUsize,
committed: AtomicBool,
finalized: AtomicBool,
@@ -258,6 +262,7 @@ impl ExecutionService for AtomicFailureExecution {
&self,
request: StartExecution,
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
self.trace.starts.fetch_add(1, Ordering::AcqRel);
Box::pin(async move {
Ok(StartedExecution {
request_id: gateway_core::engine::ModelRequestId::new("req_ws_atomic")
backend/crates/gateway-api/tests/openai/responses/websocket/protocol.rs
+33
View File
@@ -0,0 +1,33 @@
use gateway_core::operation::Operation;
use serde_json::json;
use super::decode_response_create;
#[test]
fn latest_official_codex_response_create_fixture_decodes_unchanged() {
// Audited against openai/codex main 94cbbddafc1776d5e377bca1b05932c697e82238.
let decoded = decode_response_create(
&json!({
"type": "response.create",
"model": "smart-code",
"input": "hello",
"store": false
})
.to_string(),
)
.expect("latest official Codex response.create fixture must decode");
let Operation::Generate(request) = decoded.operation() else {
panic!("Responses must map to Generate");
};
assert_eq!(
request.protocol_payload().body(),
json!({
"model": "smart-code",
"input": "hello",
"store": false
})
.as_object()
.expect("fixture body")
);
}