diff --git a/backend/crates/gateway-api/Cargo.toml b/backend/crates/gateway-api/Cargo.toml index a9ed466f2..ee2de5ddc 100644 --- a/backend/crates/gateway-api/Cargo.toml +++ b/backend/crates/gateway-api/Cargo.toml @@ -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"] } diff --git a/backend/crates/gateway-api/src/openai/responses/websocket.rs b/backend/crates/gateway-api/src/openai/responses/websocket.rs deleted file mode 100644 index c494e5991..000000000 --- a/backend/crates/gateway-api/src/openai/responses/websocket.rs +++ /dev/null @@ -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, - connect_info: Option>>, - 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, - user_agent: Option, - 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, - user_agent: Option, - request_headers: OpenAiRequestHeaders, - lifecycle: Arc, - _connection_guard: Box, -} - -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, - provider_state: Option, -} - -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) { - 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, -) { - 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, 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 { - decode_response_create_inner(payload, request_headers) -} - -fn decode_response_create_inner( - payload: &str, - request_headers: &OpenAiRequestHeaders, -) -> Result { - let Value::Object(mut body) = serde_json::from_str::(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) -> 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; -} diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/connection.rs b/backend/crates/gateway-api/src/openai/responses/websocket/connection.rs new file mode 100644 index 000000000..ece0f2206 --- /dev/null +++ b/backend/crates/gateway-api/src/openai/responses/websocket/connection.rs @@ -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>, + phase: FramePhase, +} + +impl WriteContext { + pub(super) fn request(request_id: &Arc, 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>, +} + +#[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) { + 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, + opened_at: Instant, + expired: Arc, + commands: Option>, + incoming: mpsc::Receiver, + pump_task: Option>, + stats: Arc, + config: ConnectionConfig, + exit_reason: Option, +} + +impl ResponsesWebSocketConnection { + pub(super) fn new( + socket: WebSocket, + connection_id: String, + cancellation: CancellationToken, + ) -> Self { + spawn_connection( + socket, + Arc::::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 { + 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>, + ) { + 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( + socket: S, + connection_id: Arc, + cancellation: CancellationToken, + config: ConnectionConfig, +) -> ResponsesWebSocketConnection +where + S: Stream> + Sink + 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( + mut socket: S, + mut commands: mpsc::Receiver, + incoming: mpsc::Sender, + connection_id: Arc, + cancellation: CancellationToken, + opened_at: Instant, + expired: Arc, + stats: Arc, + config: ConnectionConfig, +) where + S: Stream> + Sink + 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, + 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, + }) +} diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/forward.rs b/backend/crates/gateway-api/src/openai/responses/websocket/forward.rs new file mode 100644 index 000000000..0b6bdd1c3 --- /dev/null +++ b/backend/crates/gateway-api/src/openai/responses/websocket/forward.rs @@ -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, + provider_state: Option, +} + +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) { + 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::::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, +) { + 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, +) -> 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, 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, +) -> 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, +) -> 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, +) -> 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, + 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, + request_id: &Arc, + 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, + 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, + 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" + ); +} diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs b/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs new file mode 100644 index 000000000..9d0750338 --- /dev/null +++ b/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs @@ -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, + connect_info: Option>>, + 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, + user_agent: Option, + 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, + user_agent: Option, + request_headers: OpenAiRequestHeaders, + lifecycle: Arc, + connection_guard: Box, +} + +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::::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; + } +} diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/protocol.rs b/backend/crates/gateway-api/src/openai/responses/websocket/protocol.rs new file mode 100644 index 000000000..0f4d12a99 --- /dev/null +++ b/backend/crates/gateway-api/src/openai/responses/websocket/protocol.rs @@ -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 { + decode_response_create_inner(payload, request_headers) +} + +fn decode_response_create_inner( + payload: &str, + request_headers: &OpenAiRequestHeaders, +) -> Result { + let Value::Object(mut body) = serde_json::from_str::(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, + ) +} diff --git a/backend/crates/gateway-api/tests/architecture.rs b/backend/crates/gateway-api/tests/architecture.rs index aadd1eb4d..9091d0ab1 100644 --- a/backend/crates/gateway-api/tests/architecture.rs +++ b/backend/crates/gateway-api/tests/architecture.rs @@ -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(); diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket/connection.rs b/backend/crates/gateway-api/tests/openai/responses/websocket/connection.rs new file mode 100644 index 000000000..e1b22780d --- /dev/null +++ b/backend/crates/gateway-api/tests/openai/responses/websocket/connection.rs @@ -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, +} + +impl ExecutionSession for CrossingLimitSession { + fn next_event(&mut self) -> BoxFuture<'_, Result, 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, EngineError>> { + Box::pin(async { Err(EngineError::InvalidDeliveryState) }) + } + + fn response_headers(&self) -> &[ProviderResponseHeader] { + &[] + } + + fn commit_downstream(&mut self, _: Option) -> 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) -> BoxFuture<'static, ()> { + Box::pin(async move { + self.trace.finalized.store(true, Ordering::Release); + }) + } +} + +struct CrossingLimitExecution { + client: AuthenticatedClient, + trace: Arc, +} + +impl ExecutionService for CrossingLimitExecution { + fn authenticate( + &self, + plaintext: &str, + ) -> Result { + (plaintext == "sk_ws_crossing_limit") + .then(|| self.client.clone()) + .ok_or(ClientAuthenticationError::InvalidKey) + } + + fn public_models(&self, _: &AuthenticatedClient) -> Vec { + 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> { + 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> { + 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>, + written: UnboundedSender, + stall_writes: bool, + dropped: Arc, +} + +impl Stream for TestSocket { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + self.incoming.poll_recv(context) + } +} + +impl Sink for TestSocket { + type Error = TestSocketError; + + fn poll_ready(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + 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> { + Poll::Ready(Ok(())) + } + + fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } +} + +impl Drop for TestSocket { + fn drop(&mut self) { + self.dropped.store(true, Ordering::Release); + } +} + +struct PumpHarness { + connection: ResponsesWebSocketConnection, + incoming: UnboundedSender>, + written: UnboundedReceiver, + dropped: Arc, +} + +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::::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::::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::(&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::(&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::(&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::(&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(); +} diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket.rs b/backend/crates/gateway-api/tests/openai/responses/websocket/mod.rs similarity index 99% rename from backend/crates/gateway-api/tests/openai/responses/websocket.rs rename to backend/crates/gateway-api/tests/openai/responses/websocket/mod.rs index c8b7fa3dd..eadd401fa 100644 --- a/backend/crates/gateway-api/tests/openai/responses/websocket.rs +++ b/backend/crates/gateway-api/tests/openai/responses/websocket/mod.rs @@ -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> { + self.trace.starts.fetch_add(1, Ordering::AcqRel); Box::pin(async move { Ok(StartedExecution { request_id: gateway_core::engine::ModelRequestId::new("req_ws_atomic") diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket/protocol.rs b/backend/crates/gateway-api/tests/openai/responses/websocket/protocol.rs new file mode 100644 index 000000000..05b76e066 --- /dev/null +++ b/backend/crates/gateway-api/tests/openai/responses/websocket/protocol.rs @@ -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") + ); +}