mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
fix: add resilient Responses WebSocket transport
Expose the official Responses WebSocket endpoint so Codex can retry interrupted turns without proxy-side response splicing. Preserve upstream frames under backpressure, report pump exit reasons, organize Responses handlers by transport, and poll fingerprint updates daily. Keep the warning policy in the backend manifest instead of a repository-wide Cargo config. Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
@@ -1,2 +0,0 @@
|
||||
[build]
|
||||
warnings = "deny"
|
||||
+1
-2
@@ -49,6 +49,5 @@ Thumbs.db
|
||||
/.reasonix/
|
||||
/.worktrees/
|
||||
/worktrees/
|
||||
/.cargo/*
|
||||
!/.cargo/config.toml
|
||||
/.cargo/
|
||||
*.tmp
|
||||
|
||||
@@ -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` | 模型列表 |
|
||||
|
||||
Generated
+2
@@ -364,6 +364,8 @@ dependencies = [
|
||||
"futures",
|
||||
"hex",
|
||||
"hmac",
|
||||
"hyper",
|
||||
"hyper-util",
|
||||
"indexmap",
|
||||
"ipnet",
|
||||
"percent-encoding",
|
||||
|
||||
+3
-2
@@ -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]
|
||||
|
||||
@@ -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<String> {
|
||||
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::<Value>(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<Value>) {
|
||||
let kind = NO_AVAILABLE_ACCOUNTS_ERROR;
|
||||
(
|
||||
|
||||
@@ -5,4 +5,3 @@ pub mod errors;
|
||||
pub mod models;
|
||||
pub mod responses;
|
||||
pub mod router;
|
||||
pub mod sse;
|
||||
|
||||
@@ -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<String, Value>) -> &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<String, Value>,
|
||||
headers: &HeaderMap,
|
||||
forced_subagent: Option<&str>,
|
||||
client_ip: Option<String>,
|
||||
client_api_key_id: String,
|
||||
) -> Result<PreparedResponsesRequest, ResponsesRequestValidationError> {
|
||||
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<String, ResponsesModelValidationError> {
|
||||
) -> Result<String, ResponsesRequestValidationError> {
|
||||
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,
|
||||
}
|
||||
@@ -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 {
|
||||
@@ -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<TokioIo<Upgraded>>;
|
||||
|
||||
/// 将已认证的 `GET /v1/responses` 请求升级为官方 Responses WebSocket 协议。
|
||||
pub async fn responses_websocket(State(state): State<AppState>, 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::<RequestId>()
|
||||
.cloned()
|
||||
.unwrap_or_else(RequestId::generate);
|
||||
let client_ip = request
|
||||
.extensions()
|
||||
.get::<crate::api::middleware::request_id::ClientIp>()
|
||||
.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<String>,
|
||||
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<Result<bytes::Bytes, crate::dispatch::errors::ResponseDispatchStreamError>>),
|
||||
Downstream(Option<Result<Message, WebSocketError>>),
|
||||
}
|
||||
|
||||
async fn forward_sse_events(
|
||||
socket: &mut ResponsesWebSocket,
|
||||
events: Vec<SseEvent>,
|
||||
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::<Map<String, Value>>();
|
||||
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<Map<String, Value>, ResponseCreateFrameError> {
|
||||
let Value::Object(mut body) = serde_json::from_str::<Value>(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;
|
||||
}
|
||||
@@ -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<AppState> {
|
||||
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))
|
||||
|
||||
@@ -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<String>, Option<String>, Option<String>);
|
||||
|
||||
/// Appcast 指纹更新检查器。
|
||||
|
||||
@@ -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<u8>,
|
||||
}
|
||||
|
||||
impl SseEventDecoder {
|
||||
/// 追加一个字节块并返回其中已经完整的事件。
|
||||
pub fn push(&mut self, chunk: &[u8]) -> Result<Vec<SseEvent>, 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<Vec<SseEvent>, 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)
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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<String>,
|
||||
initial_event_timeout: Option<Duration>,
|
||||
) -> 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<WebSocketStreamPoolReturn>,
|
||||
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"
|
||||
);
|
||||
|
||||
@@ -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<MaybeTlsStream<TcpStream>>;
|
||||
@@ -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<String> },
|
||||
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<String> {
|
||||
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<Message, tungstenite::Error>,
|
||||
exit_reason: Option<PumpExitReason>,
|
||||
}
|
||||
|
||||
/// 由后台 pump 任务托管的 WebSocket 连接句柄。
|
||||
pub(crate) struct PumpedWebSocket {
|
||||
connection_id: Uuid,
|
||||
tx_command: mpsc::Sender<PumpCommand>,
|
||||
rx_message: mpsc::Receiver<Result<Message, tungstenite::Error>>,
|
||||
closed: Arc<AtomicBool>,
|
||||
exit_reason: Arc<Mutex<Option<PumpExitReason>>>,
|
||||
pump: Option<JoinHandle<()>>,
|
||||
}
|
||||
|
||||
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::<PumpCommand>(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<PumpExitReason> {
|
||||
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<PumpCommand>,
|
||||
tx_message: mpsc::Sender<Result<Message, tungstenite::Error>>,
|
||||
keepalive: PumpKeepalive,
|
||||
closed: Arc<AtomicBool>,
|
||||
exit_reason: Arc<Mutex<Option<PumpExitReason>>>,
|
||||
) {
|
||||
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<PendingInbound> = 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<Result<Message, tungstenite::Error>>,
|
||||
pending_inbound: &mut Option<PendingInbound>,
|
||||
connection_id: Uuid,
|
||||
backpressure_events: &mut u64,
|
||||
inbound: PendingInbound,
|
||||
) -> Option<PumpExitReason> {
|
||||
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<PumpExitReason> {
|
||||
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<Duration>,
|
||||
liveness_timeout: Option<Duration>,
|
||||
check_liveness: bool,
|
||||
) -> Option<PumpExitReason> {
|
||||
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!` 分支实际禁用。
|
||||
|
||||
@@ -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,
|
||||
@@ -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::<Value>(&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<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>> {
|
||||
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<tokio::net::TcpStream>,
|
||||
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<S>(websocket: &mut WebSocketStream<S>) -> Vec<Value>
|
||||
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::<Value>(&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()
|
||||
}
|
||||
@@ -76,7 +76,7 @@ fn replay_snapshot(turn_input: Vec<Value>, turn_output: Vec<Value>) -> ResponseR
|
||||
}
|
||||
}
|
||||
|
||||
async fn accept_async(stream: TcpStream) -> Result<WebSocketStream<TcpStream>, WsError> {
|
||||
pub(crate) async fn accept_async(stream: TcpStream) -> Result<WebSocketStream<TcpStream>, 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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -36,7 +36,6 @@ deploy/.env
|
||||
worktrees/
|
||||
.codex-ws-audit/
|
||||
.agents/
|
||||
.cargo/*
|
||||
!.cargo/config.toml
|
||||
.cargo/
|
||||
.reasonix/
|
||||
*.tmp
|
||||
|
||||
+14
-9
@@ -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 和运行信息 |
|
||||
|
||||
Reference in New Issue
Block a user