diff --git a/.cargo/config.toml b/.cargo/config.toml deleted file mode 100644 index 5a9134515..000000000 --- a/.cargo/config.toml +++ /dev/null @@ -1,2 +0,0 @@ -[build] -warnings = "deny" diff --git a/.gitignore b/.gitignore index 3b25d6ca9..3a54f3979 100644 --- a/.gitignore +++ b/.gitignore @@ -49,6 +49,5 @@ Thumbs.db /.reasonix/ /.worktrees/ /worktrees/ -/.cargo/* -!/.cargo/config.toml +/.cargo/ *.tmp diff --git a/README.md b/README.md index f1d24f6dc..49fc63149 100644 --- a/README.md +++ b/README.md @@ -23,8 +23,8 @@ Codex Proxy RS 为多个 Codex 账号提供统一的 Responses API。客户端 ## 功能 -- 兼容 `/v1/responses`、SSE 流式响应、Review、Compact 和模型查询。 -- 支持 HTTP/SSE 与 WebSocket 上游连接。 +- 兼容 `/v1/responses`、SSE、官方 Responses WebSocket、Review、Compact 和模型查询。 +- 客户端与上游均可使用 HTTP/SSE 或 WebSocket。 - 提供 `smart`、额度重置优先、轮询和粘性四种账号调度策略。 - 单个账号发生额度、认证、模型或传输故障时，继续使用账号池中的其他可用账号。 - 使用会话亲和和受限历史重放处理 `previous_response_not_found` 与跨账号续接。 @@ -130,11 +130,28 @@ curl http://127.0.0.1:8080/v1/responses \ }' ``` +官方 Codex 使用自定义 provider 时，可在 `~/.codex/config.toml` 中启用 Responses WebSocket： + +```toml +model_provider = "codex-proxy-rs" + +[model_providers.codex-proxy-rs] +name = "Codex Proxy RS" +base_url = "https://your-domain.example/v1" +env_key = "CODEX_PROXY_API_KEY" +wire_api = "responses" +requires_openai_auth = true +supports_websockets = true +``` + +`CODEX_PROXY_API_KEY` 的值为管理端创建的 `sk_...` 客户端 Key。HTTPS 反向代理必须转发 WebSocket 的 `Upgrade` 和 `Connection` 请求头。 + 可用接口： | 路由 | 说明 | | --- | --- | | `POST /v1/responses` | Responses 非流式或 SSE 流式请求 | +| `GET /v1/responses` | 官方 Responses WebSocket | | `POST /v1/responses/review` | Review 请求 | | `POST /v1/responses/compact` | Compact 请求 | | `GET /v1/models` | 模型列表 | diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 8a186a0c9..371cf1410 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -364,6 +364,8 @@ dependencies = [ "futures", "hex", "hmac", + "hyper", + "hyper-util", "indexmap", "ipnet", "percent-encoding", diff --git a/backend/Cargo.toml b/backend/Cargo.toml index 4a5473e72..c5a9e08d3 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -18,6 +18,8 @@ flate2 = "1.1.9" futures = "0.3.32" hex = "0.4.3" hmac = "0.13.0" +hyper = { version = "1.10.1", features = ["http1", "server"] } +hyper-util = { version = "0.1.20", features = ["tokio"] } indexmap = { version = "2.0", features = ["serde"] } ipnet = "2.12.0" percent-encoding = "2.3.2" @@ -62,8 +64,7 @@ tokio-tungstenite = { git = "https://github.com/openai-oss-forks/tokio-tungsteni tungstenite = { git = "https://github.com/openai-oss-forks/tungstenite-rs", rev = "4fffad30fe373adbdcffab9545e9e9bf4f2fc19f" } [lints.rust] -future_incompatible = { level = "warn", priority = -1 } -nonstandard_style = { level = "deny", priority = -1 } +warnings = "deny" unsafe_code = "forbid" [lints.clippy] diff --git a/backend/src/api/client/errors.rs b/backend/src/api/client/errors.rs index 687045e83..ddb8ab70a 100644 --- a/backend/src/api/client/errors.rs +++ b/backend/src/api/client/errors.rs @@ -138,6 +138,79 @@ pub fn responses_stream_dispatch_failed_sse_event(error: &ResponseDispatchError) responses_failed_sse_event(kind, &http_error.message) } +pub(super) fn responses_websocket_dispatch_error_event( + error: &ResponseDispatchError, + request_id: &str, +) -> String { + if let Some(event) = upstream_client_websocket_error_event(error, request_id) { + return event; + } + let http_error = + response_dispatch_http_error(error, ResponseDispatchMessageStyle::ResponsesStream); + let kind = response_dispatch_error_kind(error); + responses_websocket_error_event( + http_error.status, + kind.error_type, + kind.code, + &http_error.message, + request_id, + ) +} + +pub(super) fn responses_websocket_error_event( + status: StatusCode, + error_type: &str, + code: &str, + message: &str, + request_id: &str, +) -> String { + json!({ + "type": "error", + "status": status.as_u16(), + "error": { + "type": error_type, + "code": code, + "message": message, + }, + "headers": { + "x-request-id": request_id, + }, + }) + .to_string() +} + +fn upstream_client_websocket_error_event( + error: &ResponseDispatchError, + request_id: &str, +) -> Option { + let ResponseDispatchError::Upstream(CodexClientError::Upstream { status, body, .. }) = error + else { + return None; + }; + if !status.is_client_error() { + return None; + } + let Value::Object(mut payload) = serde_json::from_str::(body).ok()? else { + return None; + }; + if !payload.get("error").is_some_and(Value::is_object) { + return None; + } + payload.insert("type".to_string(), Value::String("error".to_string())); + payload.insert("status".to_string(), Value::Number(status.as_u16().into())); + let headers = payload + .entry("headers".to_string()) + .or_insert_with(|| Value::Object(serde_json::Map::new())); + if !headers.is_object() { + *headers = Value::Object(serde_json::Map::new()); + } + headers.as_object_mut()?.insert( + "x-request-id".to_string(), + Value::String(request_id.to_string()), + ); + Some(Value::Object(payload).to_string()) +} + fn responses_no_available_accounts_response() -> (StatusCode, Json) { let kind = NO_AVAILABLE_ACCOUNTS_ERROR; ( diff --git a/backend/src/api/client/mod.rs b/backend/src/api/client/mod.rs index ea188f667..cfcb349a6 100644 --- a/backend/src/api/client/mod.rs +++ b/backend/src/api/client/mod.rs @@ -5,4 +5,3 @@ pub mod errors; pub mod models; pub mod responses; pub mod router; -pub mod sse; diff --git a/backend/src/api/client/responses.rs b/backend/src/api/client/responses/mod.rs similarity index 89% rename from backend/src/api/client/responses.rs rename to backend/src/api/client/responses/mod.rs index d7e069caf..756317a9f 100644 --- a/backend/src/api/client/responses.rs +++ b/backend/src/api/client/responses/mod.rs @@ -1,6 +1,11 @@ //! OpenAI Responses API 透明代理入口：解析原始 JSON body，提取调度元数据， //! 做最小 patch 后交给调度层，请求语义原样透传上游。 +mod sse; +mod websocket; + +pub use websocket::responses_websocket; + use axum::{ Extension, Json, body::{Body, Bytes}, @@ -20,6 +25,7 @@ use crate::{ }, }; +use self::sse::{done_sse_frame, event_stream_response as sse_event_stream_response}; use super::{ errors::{ invalid_responses_request_response, missing_client_api_key_response, @@ -27,7 +33,6 @@ use super::{ responses_dispatch_error_response_ref, responses_stream_dispatch_failed_sse_event, }, models::model_catalog_for_state, - sse::{done_sse_frame, event_stream_response as sse_event_stream_response}, }; const OPENAI_SUBAGENT_HEADER: &str = "x-openai-subagent"; @@ -322,29 +327,39 @@ async fn handle_responses( let Some(body) = parse_responses_body(&body) else { return invalid_responses_request_response().into_response(); }; - let model = match validated_responses_model(&state, request_model(&body)).await { - Ok(model) => model, - Err(ResponsesModelValidationError::InvalidRequest) => { + let prepared = match prepare_responses_request( + &state, + body, + &headers, + forced_subagent, + client_ip, + client_api_key_id, + ) + .await + { + Ok(prepared) => prepared, + Err(ResponsesRequestValidationError::InvalidRequest) => { return invalid_responses_request_response().into_response(); } - Err(ResponsesModelValidationError::ModelNotFound) => { + Err(ResponsesRequestValidationError::ModelNotFound) => { return model_not_found_response().into_response(); } }; - let mut codex_request = build_codex_request(body, &headers, forced_subagent); - codex_request.client_ip = client_ip; - codex_request.client_api_key_id = Some(client_api_key_id); + let PreparedResponsesRequest { + mut request, + requested_model, + } = prepared; // 客户端未显式指定 transport 时默认偏好 WebSocket。 - if !codex_request.force_http_sse { - codex_request.use_websocket = true; + if !request.force_http_sse { + request.use_websocket = true; } - let stream = codex_request.stream(); + let stream = request.stream(); if stream { return match state .services .responses - .stream(request_id.as_str(), route, codex_request, &model) + .stream(request_id.as_str(), route, request, &requested_model) .await { Ok(stream) => live_event_stream_response(stream), @@ -355,7 +370,7 @@ async fn handle_responses( match state .services .responses - .complete(request_id.as_str(), route, codex_request, &model) + .complete(request_id.as_str(), route, request, &requested_model) .await { Ok(result) => apply_safe_response_headers( @@ -384,10 +399,10 @@ async fn handle_compact_responses( }; let model = match validated_responses_model(&state, request_model(&body)).await { Ok(model) => model, - Err(ResponsesModelValidationError::InvalidRequest) => { + Err(ResponsesRequestValidationError::InvalidRequest) => { return invalid_responses_request_response().into_response(); } - Err(ResponsesModelValidationError::ModelNotFound) => { + Err(ResponsesRequestValidationError::ModelNotFound) => { return model_not_found_response().into_response(); } }; @@ -423,23 +438,47 @@ fn request_model(body: &Map) -> &str { .unwrap_or_default() } +pub(super) struct PreparedResponsesRequest { + pub(super) request: CodexResponsesRequest, + pub(super) requested_model: String, +} + +pub(super) async fn prepare_responses_request( + state: &AppState, + body: Map, + headers: &HeaderMap, + forced_subagent: Option<&str>, + client_ip: Option, + client_api_key_id: String, +) -> Result { + let requested_model = validated_responses_model(state, request_model(&body)).await?; + let mut request = build_codex_request(body, headers, forced_subagent); + request.client_ip = client_ip; + request.client_api_key_id = Some(client_api_key_id); + Ok(PreparedResponsesRequest { + request, + requested_model, + }) +} + async fn validated_responses_model( state: &AppState, raw_model: &str, -) -> Result { +) -> Result { let model = raw_model.trim(); if model.is_empty() { - return Err(ResponsesModelValidationError::InvalidRequest); + return Err(ResponsesRequestValidationError::InvalidRequest); } let catalog = model_catalog_for_state(state).await; if catalog.is_recognized_model_name(model) { Ok(model.to_string()) } else { - Err(ResponsesModelValidationError::ModelNotFound) + Err(ResponsesRequestValidationError::ModelNotFound) } } -enum ResponsesModelValidationError { +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum ResponsesRequestValidationError { InvalidRequest, ModelNotFound, } diff --git a/backend/src/api/client/sse.rs b/backend/src/api/client/responses/sse.rs similarity index 94% rename from backend/src/api/client/sse.rs rename to backend/src/api/client/responses/sse.rs index f64cae753..fd73c0aed 100644 --- a/backend/src/api/client/sse.rs +++ b/backend/src/api/client/responses/sse.rs @@ -13,7 +13,7 @@ use axum::{ use crate::upstream::openai::protocol::sse::DONE_SSE_FRAME; -use super::errors::openai_error_response; +use crate::api::client::errors::openai_error_response; /// 编码 OpenAI 流结束标记。 pub fn done_sse_frame() -> &'static str { diff --git a/backend/src/api/client/responses/websocket.rs b/backend/src/api/client/responses/websocket.rs new file mode 100644 index 000000000..c8f32c305 --- /dev/null +++ b/backend/src/api/client/responses/websocket.rs @@ -0,0 +1,506 @@ +//! 官方 Responses WebSocket 入站协议。 + +use axum::{ + body::Body, + extract::{Request, State}, + http::{HeaderMap, StatusCode}, + response::{IntoResponse, Response}, +}; +use futures::{SinkExt, StreamExt}; +use hyper::upgrade::Upgraded; +use hyper_util::rt::TokioIo; +use serde_json::{Map, Value, json}; +use thiserror::Error; +use tokio_tungstenite::{ + WebSocketStream, + tungstenite::{ + Error as WebSocketError, Message, + handshake::server::create_response_with_body, + protocol::{CloseFrame, Role, WebSocketConfig, frame::coding::CloseCode}, + }, +}; + +use crate::{ + api::{AppState, middleware::request_id::RequestId}, + dispatch::service::ResponseDispatchStream, + upstream::openai::protocol::{ + sse::{SseEvent, SseEventDecoder}, + websocket::{is_terminal_websocket_event, websocket_event_type}, + }, +}; + +use super::{PreparedResponsesRequest, ResponsesRequestValidationError, prepare_responses_request}; +use crate::api::client::{ + auth::authorized_client_api_key_id, + errors::{ + missing_client_api_key_response, openai_error_response, + responses_websocket_dispatch_error_event, responses_websocket_error_event, + }, + router::MAX_CLIENT_REQUEST_BODY_BYTES, +}; + +const RESPONSES_ROUTE: &str = "/v1/responses"; +const INVALID_REQUEST_MESSAGE: &str = "Invalid responses request"; +const MODEL_NOT_FOUND_MESSAGE: &str = "Model not found"; +const INVALID_STREAM_MESSAGE: &str = "Responses WebSocket requests require stream=true"; +const INVALID_UPSTREAM_EVENT_MESSAGE: &str = "Invalid upstream Responses event"; +const STREAM_DISCONNECTED_MESSAGE: &str = "Upstream stream closed before response.completed"; + +type ResponsesWebSocket = WebSocketStream>; + +/// 将已认证的 `GET /v1/responses` 请求升级为官方 Responses WebSocket 协议。 +pub async fn responses_websocket(State(state): State, mut request: Request) -> Response { + let headers = request.headers().clone(); + let Some(client_api_key_id) = authorized_client_api_key_id(&state, &headers).await else { + return missing_client_api_key_response().into_response(); + }; + let connection_request_id = request + .extensions() + .get::() + .cloned() + .unwrap_or_else(RequestId::generate); + let client_ip = request + .extensions() + .get::() + .map(|client_ip| client_ip.as_str().to_string()); + let connection_id = connection_request_id.as_str().to_string(); + let response = match create_response_with_body(&request, Body::empty) { + Ok(response) => response, + Err(error) => { + tracing::info!( + websocket_connection_id = %connection_id, + error = %error, + "invalid Responses WebSocket upgrade request" + ); + return openai_error_response( + StatusCode::BAD_REQUEST, + "Invalid Responses WebSocket upgrade request", + "invalid_request_error", + "invalid_websocket_upgrade", + ) + .into_response(); + } + }; + let upgrade = hyper::upgrade::on(&mut request); + tokio::spawn(async move { + match upgrade.await { + Ok(upgraded) => { + let socket = WebSocketStream::from_raw_socket( + TokioIo::new(upgraded), + Role::Server, + Some(downstream_websocket_config()), + ) + .await; + serve_responses_websocket( + socket, + ResponsesWebSocketSession { + state, + headers, + client_ip, + client_api_key_id, + connection_id, + }, + ) + .await; + } + Err(error) => { + tracing::warn!( + websocket_connection_id = %connection_id, + error = %error, + "Responses WebSocket upgrade failed" + ); + } + } + }); + response +} + +fn downstream_websocket_config() -> WebSocketConfig { + let mut config = WebSocketConfig::default(); + config.max_message_size = Some(MAX_CLIENT_REQUEST_BODY_BYTES); + config.max_frame_size = Some(MAX_CLIENT_REQUEST_BODY_BYTES); + config +} + +struct ResponsesWebSocketSession { + state: AppState, + headers: HeaderMap, + client_ip: Option, + client_api_key_id: String, + connection_id: String, +} + +async fn serve_responses_websocket( + mut socket: ResponsesWebSocket, + session: ResponsesWebSocketSession, +) { + tracing::info!( + websocket_connection_id = %session.connection_id, + "Responses WebSocket connected" + ); + let mut request_count = 0u64; + + while let Some(message) = socket.next().await { + let payload = match message { + Ok(Message::Text(payload)) => payload, + Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_)) => continue, + Ok(Message::Close(_)) => break, + Ok(Message::Binary(_)) => { + close_for_protocol_violation( + &mut socket, + "Responses WebSocket accepts 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 request_id = RequestId::generate(); + let body = match parse_response_create(payload.as_str()) { + Ok(body) => body, + Err(error) => { + tracing::info!( + websocket_connection_id = %session.connection_id, + request_id = %request_id.as_str(), + error = %error, + "Responses WebSocket request rejected" + ); + if !send_error( + &mut socket, + StatusCode::BAD_REQUEST, + "invalid_request_error", + "invalid_request", + &error.to_string(), + request_id.as_str(), + ) + .await + { + break; + } + continue; + } + }; + let prepared = match prepare_responses_request( + &session.state, + body, + &session.headers, + None, + session.client_ip.clone(), + session.client_api_key_id.clone(), + ) + .await + { + Ok(prepared) => prepared, + Err(error) => { + let (status, code, message) = match error { + ResponsesRequestValidationError::InvalidRequest => ( + StatusCode::BAD_REQUEST, + "invalid_request", + INVALID_REQUEST_MESSAGE, + ), + ResponsesRequestValidationError::ModelNotFound => ( + StatusCode::NOT_FOUND, + "model_not_found", + MODEL_NOT_FOUND_MESSAGE, + ), + }; + if !send_error( + &mut socket, + status, + "invalid_request_error", + code, + message, + request_id.as_str(), + ) + .await + { + break; + } + continue; + } + }; + + if !dispatch_response_create(&mut socket, &session, request_id, prepared).await { + break; + } + } + + tracing::info!( + websocket_connection_id = %session.connection_id, + request_count, + "Responses WebSocket disconnected" + ); +} + +async fn dispatch_response_create( + socket: &mut ResponsesWebSocket, + session: &ResponsesWebSocketSession, + request_id: RequestId, + prepared: PreparedResponsesRequest, +) -> bool { + let PreparedResponsesRequest { + mut request, + requested_model, + } = prepared; + if !request.force_http_sse { + request.use_websocket = true; + } + let stream = session + .state + .services + .responses + .stream( + request_id.as_str(), + RESPONSES_ROUTE, + request, + &requested_model, + ) + .await; + let stream = match stream { + Ok(stream) => stream, + Err(error) => { + let event = responses_websocket_dispatch_error_event(&error, request_id.as_str()); + return send_text(socket, event).await; + } + }; + forward_dispatch_stream(socket, stream, request_id.as_str()).await +} + +async fn forward_dispatch_stream( + socket: &mut ResponsesWebSocket, + stream: ResponseDispatchStream, + request_id: &str, +) -> bool { + let ResponseDispatchStream { + body, + response_headers, + } = stream; + if !send_text( + socket, + response_metadata_event(request_id, &response_headers), + ) + .await + { + return false; + } + + let mut body = body; + let mut decoder = SseEventDecoder::default(); + let mut terminal_sent = false; + loop { + let next = if terminal_sent { + ActiveResponseInput::Upstream(body.next().await) + } else { + tokio::select! { + next = body.next() => ActiveResponseInput::Upstream(next), + message = socket.next() => ActiveResponseInput::Downstream(message), + } + }; + match next { + ActiveResponseInput::Upstream(Some(Ok(chunk))) => { + let events = match decoder.push(&chunk) { + Ok(events) => events, + Err(error) => { + tracing::warn!(request_id, error = %error, "failed to decode Responses SSE for WebSocket forwarding"); + let sent = send_error( + socket, + StatusCode::BAD_GATEWAY, + "server_error", + "invalid_upstream_response", + INVALID_UPSTREAM_EVENT_MESSAGE, + request_id, + ) + .await; + drop(body); + return sent; + } + }; + if !forward_sse_events(socket, events, request_id, &mut terminal_sent).await { + return false; + } + } + ActiveResponseInput::Upstream(Some(Err(error))) => { + tracing::warn!(request_id, error = %error, "Responses dispatch stream failed"); + if terminal_sent { + return true; + } + return send_error( + socket, + StatusCode::BAD_GATEWAY, + "server_error", + "upstream_error", + STREAM_DISCONNECTED_MESSAGE, + request_id, + ) + .await; + } + ActiveResponseInput::Upstream(None) => { + let events = match decoder.finish() { + Ok(events) => events, + Err(error) => { + tracing::warn!(request_id, error = %error, "failed to finish Responses SSE decoding"); + return send_error( + socket, + StatusCode::BAD_GATEWAY, + "server_error", + "invalid_upstream_response", + INVALID_UPSTREAM_EVENT_MESSAGE, + request_id, + ) + .await; + } + }; + if !forward_sse_events(socket, events, request_id, &mut terminal_sent).await { + return false; + } + if terminal_sent { + return true; + } + return send_error( + socket, + StatusCode::BAD_GATEWAY, + "server_error", + "stream_disconnected", + STREAM_DISCONNECTED_MESSAGE, + request_id, + ) + .await; + } + ActiveResponseInput::Downstream(Some(Ok( + Message::Ping(_) | Message::Pong(_) | Message::Frame(_), + ))) => {} + ActiveResponseInput::Downstream(Some(Ok(Message::Close(_))) | None) => return false, + ActiveResponseInput::Downstream(Some(Ok(Message::Text(_) | Message::Binary(_)))) => { + close_for_protocol_violation( + socket, + "Only one response.create may be active per connection", + ) + .await; + return false; + } + ActiveResponseInput::Downstream(Some(Err(error))) => { + tracing::info!(request_id, error = %error, "Responses WebSocket client disconnected during response"); + return false; + } + } + } +} + +enum ActiveResponseInput { + Upstream(Option>), + Downstream(Option>), +} + +async fn forward_sse_events( + socket: &mut ResponsesWebSocket, + events: Vec, + request_id: &str, + terminal_sent: &mut bool, +) -> bool { + for event in events { + let Some(event_type) = websocket_event_type(&event.data) else { + tracing::warn!( + request_id, + "Responses SSE event does not contain a valid JSON type" + ); + let _ = send_error( + socket, + StatusCode::BAD_GATEWAY, + "server_error", + "invalid_upstream_response", + INVALID_UPSTREAM_EVENT_MESSAGE, + request_id, + ) + .await; + return false; + }; + if !send_text(socket, event.data).await { + return false; + } + *terminal_sent |= is_terminal_websocket_event(&event_type); + } + true +} + +fn response_metadata_event(request_id: &str, response_headers: &[(String, String)]) -> String { + let mut headers = response_headers + .iter() + .map(|(name, value)| (name.clone(), Value::String(value.clone()))) + .collect::>(); + headers.insert( + "x-request-id".to_string(), + Value::String(request_id.to_string()), + ); + json!({ + "type": "response.metadata", + "headers": headers, + }) + .to_string() +} + +fn parse_response_create(payload: &str) -> Result, ResponseCreateFrameError> { + let Value::Object(mut body) = serde_json::from_str::(payload)? else { + return Err(ResponseCreateFrameError::ExpectedObject); + }; + match body.remove("type") { + Some(Value::String(message_type)) if message_type == "response.create" => {} + _ => return Err(ResponseCreateFrameError::UnsupportedType), + } + if body.get("stream").and_then(Value::as_bool) != Some(true) { + return Err(ResponseCreateFrameError::StreamingRequired); + } + Ok(body) +} + +#[derive(Debug, Error)] +enum ResponseCreateFrameError { + #[error("Invalid response.create JSON: {0}")] + InvalidJson(#[from] serde_json::Error), + #[error("A response.create payload must be a JSON object")] + ExpectedObject, + #[error("Unsupported Responses WebSocket message type")] + UnsupportedType, + #[error("{INVALID_STREAM_MESSAGE}")] + StreamingRequired, +} + +async fn send_error( + socket: &mut ResponsesWebSocket, + status: StatusCode, + error_type: &str, + code: &str, + message: &str, + request_id: &str, +) -> bool { + send_text( + socket, + responses_websocket_error_event(status, error_type, code, message, request_id), + ) + .await +} + +async fn send_text(socket: &mut ResponsesWebSocket, 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 ResponsesWebSocket, reason: &'static str) { + let _ = socket + .send(Message::Close(Some(CloseFrame { + code: CloseCode::Policy, + reason: reason.into(), + }))) + .await; +} diff --git a/backend/src/api/client/router.rs b/backend/src/api/client/router.rs index a74ca88e7..c6ff352e6 100644 --- a/backend/src/api/client/router.rs +++ b/backend/src/api/client/router.rs @@ -10,7 +10,7 @@ use crate::api::AppState; use super::{ models::{model_catalog, model_detail, model_info, models}, - responses::{compact_responses, responses, review_responses}, + responses::{compact_responses, responses, responses_websocket, review_responses}, }; pub const MAX_CLIENT_REQUEST_BODY_BYTES: usize = 16 * 1024 * 1024; @@ -18,7 +18,7 @@ pub const MAX_CLIENT_REQUEST_BODY_BYTES: usize = 16 * 1024 * 1024; /// 构造 OpenAI 兼容 API 路由。 pub fn router() -> Router { Router::new() - .route("/v1/responses", post(responses)) + .route("/v1/responses", get(responses_websocket).post(responses)) .route("/v1/responses/review", post(review_responses)) .route("/v1/responses/compact", post(compact_responses)) .route("/v1/models", get(models)) diff --git a/backend/src/upstream/openai/fingerprint/updater.rs b/backend/src/upstream/openai/fingerprint/updater.rs index 15076565d..d05ba261e 100644 --- a/backend/src/upstream/openai/fingerprint/updater.rs +++ b/backend/src/upstream/openai/fingerprint/updater.rs @@ -12,7 +12,7 @@ use super::{ }; const APPCAST_TIMEOUT: Duration = Duration::from_secs(30); -pub const APPCAST_POLL_INTERVAL: Duration = Duration::from_hours(72); +pub const APPCAST_POLL_INTERVAL: Duration = Duration::from_hours(24); type AppcastFields = (Option, Option, Option); /// Appcast 指纹更新检查器。 diff --git a/backend/src/upstream/openai/protocol/sse.rs b/backend/src/upstream/openai/protocol/sse.rs index 6d1d10fec..53bb17464 100644 --- a/backend/src/upstream/openai/protocol/sse.rs +++ b/backend/src/upstream/openai/protocol/sse.rs @@ -22,6 +22,50 @@ pub const MAX_SSE_EVENT_BUFFER_BYTES: usize = 64 * 1024 * 1024; /// SSE 流结束标记帧。 pub const DONE_SSE_FRAME: &str = "data: [DONE]\n\n"; +/// 用于处理任意分块边界的增量 SSE 解码器。 +#[derive(Debug, Default)] +pub struct SseEventDecoder { + pending: Vec, +} + +impl SseEventDecoder { + /// 追加一个字节块并返回其中已经完整的事件。 + pub fn push(&mut self, chunk: &[u8]) -> Result, SseError> { + self.pending.extend_from_slice(chunk); + let mut events = Vec::new(); + let mut consumed = 0usize; + + while let Some(frame_len) = sse_frame_end(&self.pending[consumed..]) { + let end = consumed.saturating_add(frame_len); + let frame = std::str::from_utf8(&self.pending[consumed..end]) + .map_err(|error| SseError::ParseError(error.to_string()))?; + events.extend(parse_sse_events(frame)?); + consumed = end; + } + + if consumed != 0 { + self.pending.drain(..consumed); + } + if self.pending.len() > MAX_SSE_EVENT_BUFFER_BYTES { + return Err(SseError::BufferExceeded { + max_bytes: MAX_SSE_EVENT_BUFFER_BYTES, + }); + } + Ok(events) + } + + /// 流结束时解析尚未带空行分隔符的最后一帧。 + pub fn finish(&mut self) -> Result, SseError> { + if self.pending.is_empty() { + return Ok(Vec::new()); + } + let pending = std::mem::take(&mut self.pending); + let frame = std::str::from_utf8(&pending) + .map_err(|error| SseError::ParseError(error.to_string()))?; + parse_sse_events(frame) + } +} + /// 编码 OpenAI Responses `response.failed` SSE 事件。 pub fn response_failed_sse_event(error_type: &str, code: &str, message: &str) -> String { response_failed_sse_event_with_id(None, error_type, code, message) diff --git a/backend/src/upstream/openai/transport/client_sse.rs b/backend/src/upstream/openai/transport/client_sse.rs index 24a42312f..e9e111cbb 100644 --- a/backend/src/upstream/openai/transport/client_sse.rs +++ b/backend/src/upstream/openai/transport/client_sse.rs @@ -368,6 +368,13 @@ impl CodexBackendClient { } } .map_err(websocket_exchange_error_to_client_error)?; + tracing::info!( + request_id = %context.request_id, + account_id = pool_account_id.or(context.account_id).unwrap_or_default(), + websocket_connection_id = %exchange.websocket_connection_id, + ws_pool = exchange.pool_decision.map_or("unpooled", WebSocketPoolDecision::kind), + "websocket response stream established" + ); log_websocket_pool_decision( context.request_id, pool_account_id.or(context.account_id), diff --git a/backend/src/upstream/openai/transport/websocket.rs b/backend/src/upstream/openai/transport/websocket.rs index d453d7fff..b69d69161 100644 --- a/backend/src/upstream/openai/transport/websocket.rs +++ b/backend/src/upstream/openai/transport/websocket.rs @@ -44,7 +44,7 @@ use super::websocket_pool::{ PooledWebSocketConnection, WebSocketContinuationState, WebSocketPoolAcquire, WebSocketPoolDecision, WebSocketPoolLease, }; -use super::websocket_pump::{PumpKeepalive, PumpedWebSocket, RawWsStream}; +use super::websocket_pump::{PumpExitReason, PumpKeepalive, PumpedWebSocket, RawWsStream}; use super::{client::CodexResponseMetadata, diagnostics::CodexUpstreamDiagnostics, response_meta}; use uuid::Uuid; diff --git a/backend/src/upstream/openai/transport/websocket_frames.rs b/backend/src/upstream/openai/transport/websocket_frames.rs index 93d2ff412..f9221149e 100644 --- a/backend/src/upstream/openai/transport/websocket_frames.rs +++ b/backend/src/upstream/openai/transport/websocket_frames.rs @@ -27,6 +27,8 @@ pub struct CodexWebSocketExchange { /// Responses WebSocket live SSE exchange result. pub struct CodexWebSocketStreamingExchange { + /// 关联请求日志与底层 pump 生命周期日志的连接标识。 + pub(crate) websocket_connection_id: Uuid, /// Live SSE bytes converted from WebSocket events. pub body: CodexWebSocketSseStream, /// 上游 metadata 帧中的最新 turn state。 @@ -494,6 +496,7 @@ pub(super) fn stream_websocket_response( prefetched_frames: Vec, initial_event_timeout: Option, ) -> CodexWebSocketStreamingExchange { + let websocket_connection_id = websocket.connection_id(); let response_metadata = metadata.clone(); let rate_limit_header_updates = Arc::new(Mutex::new(Vec::new())); let rate_limit_header_updates_for_task = Arc::clone(&rate_limit_header_updates); @@ -519,6 +522,7 @@ pub(super) fn stream_websocket_response( }); CodexWebSocketStreamingExchange { + websocket_connection_id, body: Box::pin(body), turn_state: response_metadata.turn_state, set_cookie_headers: response_metadata.set_cookie_headers, @@ -789,8 +793,21 @@ async fn discard_stream_websocket( pool_return: Option, reason: StreamWebSocketDiscardReason, ) { + let websocket_connection_id = websocket.connection_id(); + let pump_exit = websocket.exit_reason(); + let pump_exit_reason = pump_exit + .as_ref() + .map(PumpExitReason::as_str) + .unwrap_or("running"); + let pump_exit_detail = pump_exit + .as_ref() + .and_then(PumpExitReason::detail) + .unwrap_or_default(); tracing::info!( + websocket_connection_id = %websocket_connection_id, reason = reason.as_str(), + pump_exit_reason, + pump_exit_detail, pooled = pool_return.is_some(), "discarding Responses WebSocket stream" ); diff --git a/backend/src/upstream/openai/transport/websocket_pump.rs b/backend/src/upstream/openai/transport/websocket_pump.rs index a6a2a2a81..563c422e9 100644 --- a/backend/src/upstream/openai/transport/websocket_pump.rs +++ b/backend/src/upstream/openai/transport/websocket_pump.rs @@ -11,9 +11,10 @@ //! 收发都通过 channel 与 pump 任务交互： //! - 发送：`send` 走 command channel，等待 pump 回执。 //! - 接收：`next` 从 message channel 取出 pump 转发的入站帧（`Ping`/`Pong` 已被 pump 吞掉）。 +//! - 入站缓冲满时暂停读取 socket，并继续处理关闭命令与保活；消费恢复后按原顺序继续转发。 use std::sync::{ - Arc, + Arc, Mutex, atomic::{AtomicBool, Ordering}, }; use std::time::Duration; @@ -26,6 +27,7 @@ use tokio::{ time::{Instant, MissedTickBehavior}, }; use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, tungstenite::Message}; +use uuid::Uuid; /// 底层 tungstenite WebSocket 流。 pub(crate) type RawWsStream = WebSocketStream>; @@ -60,17 +62,89 @@ enum PumpCommand { const PUMP_MESSAGE_BUFFER: usize = 64; +#[derive(Debug, Clone)] +pub(crate) enum PumpExitReason { + CommandChannelClosed, + LocalClose, + OutboundTransportError { message: String }, + UpstreamCloseFrame { frame: Option }, + UpstreamEof, + InboundTransportError { message: String }, + LivenessTimeout { timeout: Duration }, + KeepaliveTransportError { message: String }, + MessageReceiverClosed, +} + +impl PumpExitReason { + pub(crate) fn as_str(&self) -> &'static str { + match self { + Self::CommandChannelClosed => "command_channel_closed", + Self::LocalClose => "local_close", + Self::OutboundTransportError { .. } => "outbound_transport_error", + Self::UpstreamCloseFrame { .. } => "upstream_close_frame", + Self::UpstreamEof => "upstream_eof", + Self::InboundTransportError { .. } => "inbound_transport_error", + Self::LivenessTimeout { .. } => "liveness_timeout", + Self::KeepaliveTransportError { .. } => "keepalive_transport_error", + Self::MessageReceiverClosed => "message_receiver_closed", + } + } + + pub(crate) fn detail(&self) -> Option { + match self { + Self::OutboundTransportError { message } + | Self::InboundTransportError { message } + | Self::KeepaliveTransportError { message } => Some(message.clone()), + Self::UpstreamCloseFrame { frame } => frame.clone(), + Self::LivenessTimeout { timeout } => Some(format!("{timeout:?}")), + Self::CommandChannelClosed + | Self::LocalClose + | Self::UpstreamEof + | Self::MessageReceiverClosed => None, + } + } + + fn is_unexpected(&self) -> bool { + matches!( + self, + Self::OutboundTransportError { .. } + | Self::UpstreamEof + | Self::InboundTransportError { .. } + | Self::LivenessTimeout { .. } + | Self::KeepaliveTransportError { .. } + ) + } + + fn should_send_close(&self) -> bool { + matches!( + self, + Self::CommandChannelClosed + | Self::UpstreamCloseFrame { .. } + | Self::LivenessTimeout { .. } + | Self::MessageReceiverClosed + ) + } +} + +struct PendingInbound { + item: Result, + exit_reason: Option, +} + /// 由后台 pump 任务托管的 WebSocket 连接句柄。 pub(crate) struct PumpedWebSocket { + connection_id: Uuid, tx_command: mpsc::Sender, rx_message: mpsc::Receiver>, closed: Arc, + exit_reason: Arc>>, pump: Option>, } impl std::fmt::Debug for PumpedWebSocket { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("PumpedWebSocket") + .field("connection_id", &self.connection_id) .field("closed", &self.closed.load(Ordering::Acquire)) .finish() } @@ -81,16 +155,29 @@ impl PumpedWebSocket { pub(crate) fn new(inner: RawWsStream, keepalive: PumpKeepalive) -> Self { let (tx_command, rx_command) = mpsc::channel::(PUMP_COMMAND_BUFFER); let (tx_message, rx_message) = mpsc::channel(PUMP_MESSAGE_BUFFER); + let connection_id = Uuid::new_v4(); let closed = Arc::new(AtomicBool::new(false)); let closed_for_task = Arc::clone(&closed); + let exit_reason = Arc::new(Mutex::new(None)); + let exit_reason_for_task = Arc::clone(&exit_reason); let pump = tokio::spawn(async move { - pump_loop(inner, rx_command, tx_message, keepalive).await; - closed_for_task.store(true, Ordering::Release); + pump_loop( + connection_id, + inner, + rx_command, + tx_message, + keepalive, + closed_for_task, + exit_reason_for_task, + ) + .await; }); Self { + connection_id, tx_command, rx_message, closed, + exit_reason, pump: Some(pump), } } @@ -123,9 +210,22 @@ impl PumpedWebSocket { self.closed.load(Ordering::Acquire) || self.tx_command.is_closed() } + pub(crate) fn connection_id(&self) -> Uuid { + self.connection_id + } + + pub(crate) fn exit_reason(&self) -> Option { + self.exit_reason + .lock() + .expect("WebSocket pump exit reason lock poisoned") + .clone() + } + /// 主动关闭连接（best-effort 发送 Close 帧）。 pub(crate) async fn close(&self) { - let _ = self.send(Message::Close(None)).await; + if !self.is_closed() { + let _ = self.send(Message::Close(None)).await; + } } } @@ -138,10 +238,13 @@ impl Drop for PumpedWebSocket { } async fn pump_loop( + connection_id: Uuid, mut inner: RawWsStream, mut rx_command: mpsc::Receiver, tx_message: mpsc::Sender>, keepalive: PumpKeepalive, + closed: Arc, + exit_reason: Arc>>, ) { let mut last_activity = Instant::now(); let ping_interval = keepalive.ping_interval.filter(|d| !d.is_zero()); @@ -162,29 +265,65 @@ async fn pump_loop( ticker }); - loop { + let mut pending_inbound: Option = None; + let mut backpressure_events = 0_u64; + + let reason = 'pump: loop { + if pending_inbound.is_some() { + tokio::select! { + permit = tx_message.reserve() => { + let Ok(permit) = permit else { + break 'pump PumpExitReason::MessageReceiverClosed; + }; + let pending = pending_inbound + .take() + .expect("pending inbound frame must exist"); + permit.send(pending.item); + if let Some(reason) = pending.exit_reason { + break 'pump reason; + } + } + command = rx_command.recv() => { + let Some(command) = command else { + break 'pump PumpExitReason::CommandChannelClosed; + }; + if let Some(reason) = handle_command(&mut inner, command).await { + break 'pump reason; + } + } + _ = tick(&mut ticker) => { + if let Some(reason) = handle_tick( + &mut inner, + last_activity, + ping_interval, + liveness_timeout, + false, + ).await { + break 'pump reason; + } + } + } + continue; + } + tokio::select! { command = rx_command.recv() => { - match command { - // 所有句柄已丢弃：连接不再被使用，退出 pump。 - None => break, - Some(PumpCommand::Send { message, ack }) => { - let result = inner.send(message).await; - let should_break = result.is_err(); - let _ = ack.send(result); - if should_break { - break; - } - } + let Some(command) = command else { + break 'pump PumpExitReason::CommandChannelClosed; + }; + if let Some(reason) = handle_command(&mut inner, command).await { + break 'pump reason; } } message = inner.next() => { match message { - None => break, + None => break 'pump PumpExitReason::UpstreamEof, Some(Ok(Message::Ping(payload))) => { last_activity = Instant::now(); - if inner.send(Message::Pong(payload)).await.is_err() { - break; + if let Err(error) = inner.send(Message::Pong(payload)).await { + break 'pump PumpExitReason::KeepaliveTransportError { + message: error.to_string(), + }; } } Some(Ok(Message::Pong(_))) => { @@ -192,38 +331,158 @@ async fn pump_loop( } Some(Ok(message)) => { last_activity = Instant::now(); - let is_close = matches!(message, Message::Close(_)); - if tx_message.try_send(Ok(message)).is_err() { - break; - } - if is_close { - break; + let terminal_reason = match &message { + Message::Close(frame) => Some(PumpExitReason::UpstreamCloseFrame { + frame: frame.as_ref().map(|frame| format!("{frame:?}")), + }), + _ => None, + }; + if let Some(reason) = enqueue_inbound( + &tx_message, + &mut pending_inbound, + connection_id, + &mut backpressure_events, + PendingInbound { + item: Ok(message), + exit_reason: terminal_reason, + }, + ) { + break 'pump reason; } } Some(Err(err)) => { - let _ = tx_message.try_send(Err(err)); - break; + let terminal_reason = PumpExitReason::InboundTransportError { + message: err.to_string(), + }; + if let Some(reason) = enqueue_inbound( + &tx_message, + &mut pending_inbound, + connection_id, + &mut backpressure_events, + PendingInbound { + item: Err(err), + exit_reason: Some(terminal_reason), + }, + ) { + break 'pump reason; + } } } } _ = tick(&mut ticker) => { - // 先检查 liveness：长时间无入站活动则判定失活退出。 - if let Some(timeout) = liveness_timeout - && last_activity.elapsed() >= timeout { - break; - } - // 再按需主动 ping 保活（穿透 NAT/中间盒空闲计时器）。 - if ping_interval.is_some() - && inner.send(Message::Ping(Vec::new().into())).await.is_err() - { - break; + if let Some(reason) = handle_tick( + &mut inner, + last_activity, + ping_interval, + liveness_timeout, + true, + ).await { + break 'pump reason; } } } - } + }; - // best-effort：尝试通知对端关闭，随后 drop 底层流释放 socket。 - let _ = inner.send(Message::Close(None)).await; + closed.store(true, Ordering::Release); + *exit_reason + .lock() + .expect("WebSocket pump exit reason lock poisoned") = Some(reason.clone()); + log_pump_exit(connection_id, &reason, backpressure_events); + + if reason.should_send_close() { + let _ = inner.send(Message::Close(None)).await; + } +} + +fn enqueue_inbound( + tx_message: &mpsc::Sender>, + pending_inbound: &mut Option, + connection_id: Uuid, + backpressure_events: &mut u64, + inbound: PendingInbound, +) -> Option { + let PendingInbound { item, exit_reason } = inbound; + match tx_message.try_send(item) { + Ok(()) => exit_reason, + Err(mpsc::error::TrySendError::Full(item)) => { + *backpressure_events += 1; + if *backpressure_events == 1 { + tracing::info!( + websocket_connection_id = %connection_id, + buffer_capacity = PUMP_MESSAGE_BUFFER, + "Responses WebSocket pump applying inbound backpressure" + ); + } + *pending_inbound = Some(PendingInbound { item, exit_reason }); + None + } + Err(mpsc::error::TrySendError::Closed(_)) => Some(PumpExitReason::MessageReceiverClosed), + } +} + +async fn handle_command(inner: &mut RawWsStream, command: PumpCommand) -> Option { + let PumpCommand::Send { message, ack } = command; + let is_close = matches!(message, Message::Close(_)); + let result = inner.send(message).await; + let transport_error = result.as_ref().err().map(ToString::to_string); + let _ = ack.send(result); + match transport_error { + Some(message) => Some(PumpExitReason::OutboundTransportError { message }), + None if is_close => Some(PumpExitReason::LocalClose), + None => None, + } +} + +async fn handle_tick( + inner: &mut RawWsStream, + last_activity: Instant, + ping_interval: Option, + liveness_timeout: Option, + check_liveness: bool, +) -> Option { + if check_liveness + && let Some(timeout) = liveness_timeout + && last_activity.elapsed() >= timeout + { + return Some(PumpExitReason::LivenessTimeout { timeout }); + } + if ping_interval.is_some() + && let Err(error) = inner.send(Message::Ping(Vec::new().into())).await + { + return Some(PumpExitReason::KeepaliveTransportError { + message: error.to_string(), + }); + } + None +} + +fn log_pump_exit(connection_id: Uuid, reason: &PumpExitReason, backpressure_events: u64) { + let detail = reason.detail().unwrap_or_default(); + if reason.is_unexpected() { + tracing::warn!( + websocket_connection_id = %connection_id, + pump_exit_reason = reason.as_str(), + pump_exit_detail = detail, + backpressure_events, + "Responses WebSocket pump stopped unexpectedly" + ); + } else if matches!(reason, PumpExitReason::UpstreamCloseFrame { .. }) { + tracing::info!( + websocket_connection_id = %connection_id, + pump_exit_reason = reason.as_str(), + pump_exit_detail = detail, + backpressure_events, + "Responses WebSocket pump received close frame" + ); + } else { + tracing::debug!( + websocket_connection_id = %connection_id, + pump_exit_reason = reason.as_str(), + pump_exit_detail = detail, + backpressure_events, + "Responses WebSocket pump stopped" + ); + } } /// 等待 ping ticker；`None` 时永远挂起，让 `select!` 分支实际禁用。 diff --git a/backend/tests/api/client/responses.rs b/backend/tests/api/client/responses/mod.rs similarity index 98% rename from backend/tests/api/client/responses.rs rename to backend/tests/api/client/responses/mod.rs index 6af645d16..4dbddb2a2 100644 --- a/backend/tests/api/client/responses.rs +++ b/backend/tests/api/client/responses/mod.rs @@ -13,6 +13,8 @@ use crate::support::{ storage::{background_task_stores, create_test_redis, init_test_db, test_database_url}, }; +mod websocket; + #[tokio::test] async fn responses_route_should_reject_missing_client_api_key() { let (app, _key, _dir) = test_app_with_client_api_key().await; @@ -151,7 +153,7 @@ async fn responses_route_should_default_omitted_stream_to_sse() { assert!(body.ends_with("data: [DONE]\n\n")); } -async fn test_app_with_client_api_key() -> ( +pub(super) async fn test_app_with_client_api_key() -> ( axum::Router, String, crate::support::storage::TestDatabaseGuard, diff --git a/backend/tests/api/client/responses/websocket.rs b/backend/tests/api/client/responses/websocket.rs new file mode 100644 index 000000000..c6d470596 --- /dev/null +++ b/backend/tests/api/client/responses/websocket.rs @@ -0,0 +1,267 @@ +use std::{net::SocketAddr, time::Duration}; + +use axum::Router; +use futures::{SinkExt, StreamExt}; +use serde_json::{Value, json}; +use tokio::{net::TcpListener, task::JoinHandle, time::timeout}; +use tokio_tungstenite::{ + WebSocketStream, connect_async, + tungstenite::{ + Error as WebSocketError, Message, + client::IntoClientRequest, + http::{HeaderValue, header::AUTHORIZATION}, + }, +}; + +use crate::dispatch::service::accept_async as accept_upstream_websocket; + +struct TestServer { + address: SocketAddr, + task: JoinHandle<()>, +} + +impl Drop for TestServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +#[tokio::test] +async fn responses_websocket_should_reject_missing_client_api_key_during_handshake() { + let (app, _api_key, _dir) = super::test_app_with_client_api_key().await; + let server = spawn_app(app).await; + let request = format!("ws://{}/v1/responses", server.address) + .into_client_request() + .unwrap(); + + let error = connect_async(request).await.unwrap_err(); + + let WebSocketError::Http(response) = error else { + panic!("expected HTTP handshake rejection"); + }; + assert_eq!(response.status(), 401); +} + +#[tokio::test] +async fn responses_websocket_should_return_official_error_frame_when_accounts_are_unavailable() { + let (app, api_key, _dir) = super::test_app_with_client_api_key().await; + let server = spawn_app(app).await; + let mut websocket = connect_responses_websocket(server.address, &api_key).await; + + websocket + .send(Message::Text(response_create_payload("no accounts").into())) + .await + .unwrap(); + let events = receive_response(&mut websocket).await; + + assert_eq!(events.last().unwrap()["type"], "error"); + assert_eq!(events.last().unwrap()["status"], 503); + assert_eq!( + events.last().unwrap()["error"]["code"], + "no_available_accounts" + ); +} + +#[tokio::test] +async fn responses_websocket_should_forward_multiple_response_create_requests_on_one_connection() { + let upstream_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let upstream_base_url = format!("http://{}", upstream_listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (stream, _) = upstream_listener.accept().await.unwrap(); + let mut websocket = accept_upstream_websocket(stream).await.unwrap(); + let mut requests = Vec::new(); + for index in 1..=2 { + let request = websocket.next().await.unwrap().unwrap(); + requests.push(serde_json::from_str::(&request.into_text().unwrap()).unwrap()); + send_upstream_response(&mut websocket, &format!("resp_{index}"), "ok").await; + } + requests + }); + let (app, api_key, _dir) = + crate::dispatch::service::test_app_with_account(upstream_base_url).await; + let server = spawn_app(app).await; + let mut websocket = connect_responses_websocket(server.address, &api_key).await; + + websocket + .send(Message::Text(response_create_payload("first").into())) + .await + .unwrap(); + let first = receive_response(&mut websocket).await; + websocket + .send(Message::Text(response_create_payload("second").into())) + .await + .unwrap(); + let second = receive_response(&mut websocket).await; + let upstream_requests = upstream.await.unwrap(); + + assert_eq!(response_event_types(&first), response_event_types(&second)); + assert_eq!(first.first().unwrap()["type"], "response.metadata"); + assert_eq!(first.last().unwrap()["type"], "response.completed"); + assert_eq!(upstream_requests.len(), 2); + assert_eq!(upstream_requests[0]["type"], "response.create"); + assert_eq!(upstream_requests[1]["type"], "response.create"); +} + +#[tokio::test] +async fn responses_websocket_should_allow_official_retry_after_midstream_upstream_disconnect() { + let upstream_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let upstream_base_url = format!("http://{}", upstream_listener.local_addr().unwrap()); + let upstream = tokio::spawn(async move { + let (first_stream, _) = upstream_listener.accept().await.unwrap(); + let mut first = accept_upstream_websocket(first_stream).await.unwrap(); + let _request = first.next().await.unwrap().unwrap(); + first + .send(Message::Text( + json!({ + "type": "response.output_text.delta", + "delta": "partial", + }) + .to_string() + .into(), + )) + .await + .unwrap(); + drop(first); + + let (second_stream, _) = upstream_listener.accept().await.unwrap(); + let mut second = accept_upstream_websocket(second_stream).await.unwrap(); + let _request = second.next().await.unwrap().unwrap(); + send_upstream_response(&mut second, "resp_retry", "recovered").await; + }); + let (app, api_key, _dir) = + crate::dispatch::service::test_app_with_account(upstream_base_url).await; + let server = spawn_app(app).await; + + let mut first = connect_responses_websocket(server.address, &api_key).await; + first + .send(Message::Text(response_create_payload("retry me").into())) + .await + .unwrap(); + let failed = receive_response(&mut first).await; + drop(first); + + let mut second = connect_responses_websocket(server.address, &api_key).await; + second + .send(Message::Text(response_create_payload("retry me").into())) + .await + .unwrap(); + let recovered = receive_response(&mut second).await; + upstream.await.unwrap(); + + assert!(response_event_types(&failed).contains(&"response.output_text.delta")); + assert_eq!(failed.last().unwrap()["type"], "response.failed"); + assert_eq!( + failed.last().unwrap()["response"]["error"]["code"], + "stream_disconnected" + ); + assert_eq!(recovered.last().unwrap()["type"], "response.completed"); +} + +async fn spawn_app(app: Router) -> TestServer { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let task = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + TestServer { address, task } +} + +async fn connect_responses_websocket( + address: SocketAddr, + api_key: &str, +) -> WebSocketStream> { + let mut request = format!("ws://{address}/v1/responses") + .into_client_request() + .unwrap(); + request.headers_mut().insert( + AUTHORIZATION, + HeaderValue::from_str(&format!("Bearer {api_key}")).unwrap(), + ); + connect_async(request).await.unwrap().0 +} + +fn response_create_payload(content: &str) -> String { + json!({ + "type": "response.create", + "model": "gpt-5.5", + "instructions": "test", + "input": [{"role": "user", "content": content}], + "tools": [], + "tool_choice": "auto", + "parallel_tool_calls": true, + "reasoning": {"effort": "medium"}, + "prompt_cache_key": "downstream-websocket-test", + "store": false, + "stream": true, + "include": [], + }) + .to_string() +} + +async fn send_upstream_response( + websocket: &mut WebSocketStream, + response_id: &str, + text: &str, +) { + for event in [ + json!({ + "type": "response.created", + "response": {"id": response_id, "status": "in_progress"}, + }), + json!({ + "type": "response.output_text.delta", + "delta": text, + }), + json!({ + "type": "response.completed", + "response": { + "id": response_id, + "object": "response", + "status": "completed", + "usage": { + "input_tokens": 3, + "output_tokens": 1, + "total_tokens": 4, + }, + }, + }), + ] { + websocket + .send(Message::Text(event.to_string().into())) + .await + .unwrap(); + } +} + +async fn receive_response(websocket: &mut WebSocketStream) -> Vec +where + S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, +{ + let mut events = Vec::new(); + loop { + let message = timeout(Duration::from_secs(5), websocket.next()) + .await + .expect("timed out waiting for Responses WebSocket event") + .expect("Responses WebSocket closed before a terminal event") + .expect("failed to receive Responses WebSocket event"); + let Message::Text(payload) = message else { + continue; + }; + let event = serde_json::from_str::(&payload).unwrap(); + let terminal = matches!( + event.get("type").and_then(Value::as_str), + Some("response.completed" | "response.incomplete" | "response.failed" | "error") + ); + events.push(event); + if terminal { + return events; + } + } +} + +fn response_event_types(events: &[Value]) -> Vec<&str> { + events + .iter() + .filter_map(|event| event.get("type").and_then(Value::as_str)) + .collect() +} diff --git a/backend/tests/dispatch/service.rs b/backend/tests/dispatch/service.rs index 3e8570209..9c13e8719 100644 --- a/backend/tests/dispatch/service.rs +++ b/backend/tests/dispatch/service.rs @@ -76,7 +76,7 @@ fn replay_snapshot(turn_input: Vec, turn_output: Vec) -> ResponseR } } -async fn accept_async(stream: TcpStream) -> Result, WsError> { +pub(crate) async fn accept_async(stream: TcpStream) -> Result, WsError> { accept_hdr_async(stream, |_request, _response| {}).await } @@ -268,7 +268,7 @@ async fn test_app_state_with_storage( AppState::from(&services) } -async fn test_app_with_account( +pub(crate) async fn test_app_with_account( base_url: String, ) -> ( axum::Router, diff --git a/backend/tests/upstream/openai/fingerprint/mod.rs b/backend/tests/upstream/openai/fingerprint/mod.rs index ddc5275f8..03e6f41d2 100644 --- a/backend/tests/upstream/openai/fingerprint/mod.rs +++ b/backend/tests/upstream/openai/fingerprint/mod.rs @@ -1,9 +1,16 @@ -use std::num::NonZeroU32; +use std::{num::NonZeroU32, time::Duration}; -use codex_proxy_rs::upstream::openai::fingerprint::{PgFingerprintStore, UpdateChecker}; +use codex_proxy_rs::upstream::openai::fingerprint::{ + APPCAST_POLL_INTERVAL, PgFingerprintStore, UpdateChecker, +}; use crate::support::storage::init_test_db; +#[test] +fn fingerprint_appcast_poll_interval_should_be_one_day() { + assert_eq!(APPCAST_POLL_INTERVAL, Duration::from_hours(24)); +} + #[tokio::test] async fn fingerprint_store_should_update_current_record() { let (pool, _guard) = init_test_db("fingerprint-store").await; diff --git a/backend/tests/upstream/openai/protocol/mod.rs b/backend/tests/upstream/openai/protocol/mod.rs index 985e993ce..523c3e0c7 100644 --- a/backend/tests/upstream/openai/protocol/mod.rs +++ b/backend/tests/upstream/openai/protocol/mod.rs @@ -10,7 +10,7 @@ use codex_proxy_rs::upstream::openai::protocol::responses::{ transport_for_request, }; use codex_proxy_rs::upstream::openai::protocol::sse::{ - parse_sse_events, response_failed_sse_event, + SseEventDecoder, parse_sse_events, response_failed_sse_event, }; use codex_proxy_rs::upstream::openai::protocol::websocket::{ OpeningAuditHeader, OpeningAuditSnapshot, classify_websocket_error_frame, diff --git a/backend/tests/upstream/openai/protocol/sse.rs b/backend/tests/upstream/openai/protocol/sse.rs index b16e81e6c..b747d04ff 100644 --- a/backend/tests/upstream/openai/protocol/sse.rs +++ b/backend/tests/upstream/openai/protocol/sse.rs @@ -72,3 +72,51 @@ fn first_output_detection_should_match_crlf_frames() { assert!(response_body_has_first_output(body)); } + +#[test] +fn incremental_decoder_should_wait_for_complete_frame() { + let mut decoder = SseEventDecoder::default(); + + let events = decoder + .push(b"event: response.created\ndata: {\"type\":\"response.created\"}") + .unwrap(); + + assert!(events.is_empty()); +} + +#[test] +fn incremental_decoder_should_decode_events_across_chunk_boundaries() { + let mut decoder = SseEventDecoder::default(); + decoder + .push(b"event: response.created\ndata: {\"type\":\"response.") + .unwrap(); + + let events = decoder + .push( + b"created\"}\n\nevent: response.completed\ndata: {\"type\":\"response.completed\"}\n\n", + ) + .unwrap(); + + assert_eq!(events.len(), 2); +} + +#[test] +fn incremental_decoder_should_ignore_done_frame() { + let mut decoder = SseEventDecoder::default(); + + let events = decoder.push(b"data: [DONE]\n\n").unwrap(); + + assert!(events.is_empty()); +} + +#[test] +fn incremental_decoder_finish_should_decode_unterminated_final_frame() { + let mut decoder = SseEventDecoder::default(); + decoder + .push(b"event: response.completed\ndata: {\"type\":\"response.completed\"}") + .unwrap(); + + let events = decoder.finish().unwrap(); + + assert_eq!(events.len(), 1); +} diff --git a/backend/tests/upstream/openai/transport/websocket.rs b/backend/tests/upstream/openai/transport/websocket.rs index 81fe4c5c0..6b2f85cdc 100644 --- a/backend/tests/upstream/openai/transport/websocket.rs +++ b/backend/tests/upstream/openai/transport/websocket.rs @@ -1317,6 +1317,156 @@ async fn codex_backend_client_stream_should_error_when_websocket_closes_before_t ); } +#[tokio::test] +async fn codex_backend_client_stream_should_preserve_burst_during_downstream_backpressure() { + const BURST_FRAMES: usize = 256; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_codex_test_websocket(stream).await; + let _message = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.output_text.delta", + "delta": "initial;" + }) + .to_string() + .into(), + )) + .await + .unwrap(); + + for index in 0..BURST_FRAMES { + let frame = json!({ + "type": "response.output_text.delta", + "delta": format!("chunk-{index};") + }) + .to_string(); + if websocket.send(Message::Text(frame.into())).await.is_err() { + return; + } + } + let _ = websocket + .send(Message::Text( + completed_websocket_response("resp_backpressure", 3, 257).into(), + )) + .await; + let _ = websocket.close(None).await; + }); + let backend = CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + format!("http://{addr}"), + crate::support::fingerprint::runtime_test_fingerprint(), + ); + let request = pooled_websocket_request("conversation-backpressure"); + + let mut response = backend + .create_response_stream( + &request, + request_context("req_stream_backpressure", Some("chatgpt-account")), + ) + .await + .expect("websocket stream should open after the first output frame"); + + tokio::time::sleep(Duration::from_millis(100)).await; + let body = timeout(Duration::from_secs(5), async { + let mut body = String::new(); + while let Some(chunk) = response.body.next().await { + let chunk = chunk.expect("backpressured websocket frame should remain valid"); + body.push_str(std::str::from_utf8(&chunk).unwrap()); + } + body + }) + .await + .expect("backpressured websocket stream should finish"); + server.await.unwrap(); + + assert_eq!( + body.matches("event: response.output_text.delta").count(), + BURST_FRAMES + 1 + ); + assert!(body.contains("initial;")); + assert!(body.contains("chunk-0;")); + assert!(body.contains("chunk-255;")); + assert!(body.contains("resp_backpressure")); +} + +#[tokio::test] +async fn codex_backend_client_stream_should_cancel_while_inbound_is_backpressured() { + const BURST_FRAMES: usize = 256; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (burst_sent_tx, burst_sent_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_codex_test_websocket(stream).await; + let _message = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.output_text.delta", + "delta": "initial;" + }) + .to_string() + .into(), + )) + .await + .unwrap(); + for index in 0..BURST_FRAMES { + websocket + .send(Message::Text( + json!({ + "type": "response.output_text.delta", + "delta": format!("chunk-{index};") + }) + .to_string() + .into(), + )) + .await + .unwrap(); + } + burst_sent_tx.send(()).unwrap(); + + timeout(Duration::from_secs(1), async { + loop { + match websocket.next().await { + Some(Ok(Message::Close(_))) => return, + Some(Ok(_)) => continue, + Some(Err(error)) => panic!("websocket close should be graceful: {error}"), + None => panic!("websocket ended without a close frame"), + } + } + }) + .await + .expect("client cancellation should reach a backpressured pump"); + }); + let backend = CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + format!("http://{addr}"), + crate::support::fingerprint::runtime_test_fingerprint(), + ); + let request = pooled_websocket_request("conversation-backpressure-cancel"); + + let response = backend + .create_response_stream( + &request, + request_context("req_stream_backpressure_cancel", Some("chatgpt-account")), + ) + .await + .expect("websocket stream should open after the first output frame"); + burst_sent_rx + .await + .expect("server should finish the inbound burst"); + tokio::time::sleep(Duration::from_millis(100)).await; + + drop(response.body); + server.await.unwrap(); +} + #[tokio::test(start_paused = true)] async fn codex_backend_client_stream_should_wait_for_terminal_after_active_websocket_gap() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); diff --git a/deploy/Dockerfile b/deploy/Dockerfile index b1506c02f..e79ca05a3 100644 --- a/deploy/Dockerfile +++ b/deploy/Dockerfile @@ -21,7 +21,6 @@ ENV CPR_VERSION=${CPR_VERSION} ENV CPR_GIT_SHA=${CPR_GIT_SHA} ENV CPR_BUILD_TIME=${CPR_BUILD_TIME} ENV CPR_BUILD_TYPE=${CPR_BUILD_TYPE} -COPY .cargo/config.toml /app/.cargo/config.toml COPY backend/Cargo.toml backend/Cargo.lock ./ COPY backend/build.rs ./build.rs COPY backend/src ./src diff --git a/deploy/Dockerfile.dockerignore b/deploy/Dockerfile.dockerignore index 9df521485..8f04c4391 100644 --- a/deploy/Dockerfile.dockerignore +++ b/deploy/Dockerfile.dockerignore @@ -36,7 +36,6 @@ deploy/.env worktrees/ .codex-ws-audit/ .agents/ -.cargo/* -!.cargo/config.toml +.cargo/ .reasonix/ *.tmp diff --git a/docs/architecture.md b/docs/architecture.md index c55ff4692..fd75cdf5e 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -6,7 +6,7 @@ Codex Proxy RS 是一个单进程网关。Rust 进程同时提供： -- OpenAI Responses 兼容接口：`/v1/*` +- OpenAI Responses 兼容接口：`/v1/*`，包括 HTTP/SSE 和官方 Responses WebSocket - 管理端接口：`/api/admin/*` - Vue 管理端静态资源 - PostgreSQL、Redis 健康检查：`/healthz` @@ -294,10 +294,12 @@ api/ ├── client/ │ ├── router.rs │ ├── auth.rs -│ ├── responses.rs +│ ├── responses/ +│ │ ├── mod.rs +│ │ ├── sse.rs +│ │ └── websocket.rs │ ├── models.rs -│ ├── errors.rs -│ └── sse.rs +│ └── errors.rs └── admin/ ├── router.rs ├── response.rs @@ -315,7 +317,7 @@ api/ - `router.rs`：组合客户端 API、管理端 API、SPA 静态资源和 `/healthz`。 - `middleware/request_id.rs`：接收或生成请求 ID，并写入响应头。 - `middleware/trace.rs`：HTTP 访问日志。 -- `client`：`/v1` 路由、Bearer API Key 鉴权、Responses 请求解析和 SSE 输出。 +- `client`：`/v1` 路由、Bearer API Key 鉴权和 Responses 入站协议。`responses/mod.rs` 负责共享请求构造，`sse.rs` 与 `websocket.rs` 只处理各自的下游传输。 - `admin`：`/api/admin` 路由、会话或管理员 API Key 鉴权、统一响应结构。所有管理端响应带 `Cache-Control: no-store`。 - `assets.rs`：静态文件服务和 Vue Router history fallback；未知 `/api`、`/v1` 路径不会回退到 SPA。 @@ -370,11 +372,11 @@ bootstrap/ ## 代理请求链路 -一次 `POST /v1/responses` 请求经过以下步骤： +一次 `POST /v1/responses` 请求或 WebSocket `response.create` 消息经过以下步骤： 1. 请求 ID 中间件建立全链路标识，访问日志记录基础 HTTP 信息。 -2. `client::auth` 从 `Authorization: Bearer sk_...` 取出完整 Key，并在 PostgreSQL 校验是否存在且启用。 -3. API 层解析请求，提取客户端 IP、User-Agent 和 transport-only 参数。 +2. `client::auth` 从 `Authorization: Bearer sk_...` 取出完整 Key，并在 PostgreSQL 校验是否存在且启用。WebSocket 在 HTTP upgrade 前完成鉴权。 +3. API 层解析 HTTP JSON 或官方 `response.create` 文本帧，提取客户端 IP、User-Agent 和 transport-only 参数。 4. `models` 解析模型别名；`dispatch` 建立会话变体标识和 previous response 恢复计划。 5. `AccountAttemptLedger` 按当前调度策略冻结完整候选顺序。会话亲和账号可以排在前面，但不会绕过可用性检查。 6. 账号池租用并发槽位，必要时先刷新 quota；请求间隔在发送上游前执行。 @@ -425,7 +427,9 @@ replay 按轮保存增量输入和输出，不保存每一轮累计的完整输 超过任一边界时不再生成新的 replay 节点，正常代理响应不因此失败。 -流式请求只有在响应尚未提交给客户端时才能透明恢复或换号。首段预取会识别终止事件和账号级失败；一旦真实输出已经发给客户端，后续断流只能以 SSE 失败结束，不能在另一账号上重放后拼接到已有输出。 +流式请求只有在响应尚未提交给客户端时才能透明恢复或换号。首段预取会识别终止事件和账号级失败；一旦真实输出已经发给客户端，代理会用 `response.failed` 结束当前 SSE 或 WebSocket response，不能在另一账号上重放后拼接到已有输出。 + +Responses 协议没有按事件序号续传未完成 response 的请求字段。官方 Codex 在 turn 层持有完整会话历史和工具执行状态，收到可重试的断流后会建立新请求重新采样。代理保留这一责任边界：新请求仍进入账号池调度和会话恢复，但不会伪造同一个 response 的续传。 ## 账号作用域身份 @@ -549,6 +553,7 @@ frontend/src/ | --- | --- | | `GET /healthz` | PostgreSQL 和 Redis 健康检查，成功返回 `204` | | `POST /v1/responses` | Responses 非流式或 SSE 流式入口 | +| `GET /v1/responses` | 官方 Responses WebSocket upgrade 入口 | | `POST /v1/responses/review` | Review 请求入口 | | `POST /v1/responses/compact` | Compact 请求入口 | | `GET /v1/models*` | 模型列表、详情、catalog 和运行信息 |