fix: align upstream request and account lifecycle

- use OpenAI fork WebSocket transport and aligned request headers

- persist refresh/account failure states and dedupe account imports by identity

- document real-chain and risk-control audit results

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-06-20 15:33:42 +08:00
co-authored by Codex
parent b8c14f5dd1
commit 17de9ffe57
59 changed files with 7937 additions and 1256 deletions
Generated
+50 -8
View File
@@ -373,7 +373,6 @@ dependencies = [
"chrono",
"codex-proxy-core",
"codex-proxy-platform",
"flate2",
"futures",
"indexmap",
"reqwest",
@@ -390,6 +389,7 @@ dependencies = [
"tokio-rustls",
"tokio-tungstenite",
"tracing",
"tungstenite",
"uuid",
"wiremock",
]
@@ -433,7 +433,6 @@ dependencies = [
"chrono",
"config",
"dirs",
"flate2",
"hex",
"hmac",
"rand 0.10.1",
@@ -947,6 +946,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c"
dependencies = [
"crc32fast",
"libz-sys",
"miniz_oxide",
]
@@ -1208,6 +1208,30 @@ dependencies = [
"hashbrown 0.16.1",
]
[[package]]
name = "headers"
version = "0.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b3314d5adb5d94bcdf56771f2e50dbbc80bb4bdf88967526706205ac9eff24eb"
dependencies = [
"base64",
"bytes",
"headers-core",
"http",
"httpdate",
"mime",
"sha1 0.10.6",
]
[[package]]
name = "headers-core"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "54b4a22553d4242c49fddb9ba998a99962b5cc6f22cb5a3482bec22522403ce4"
dependencies = [
"http",
]
[[package]]
name = "heck"
version = "0.5.0"
@@ -1595,6 +1619,17 @@ dependencies = [
"vcpkg",
]
[[package]]
name = "libz-sys"
version = "1.1.29"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "85bc9657773828b90eeb625adff10eeac83cc21bbfd8e23a03eaa8a33c9e28d9"
dependencies = [
"cc",
"pkg-config",
"vcpkg",
]
[[package]]
name = "linux-raw-sys"
version = "0.12.1"
@@ -2968,9 +3003,8 @@ dependencies = [
[[package]]
name = "tokio-tungstenite"
version = "0.29.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8f72a05e828585856dacd553fba484c242c46e391fb0e58917c942ee9202915c"
version = "0.28.0"
source = "git+https://github.com/openai-oss-forks/tokio-tungstenite?rev=132f5b39c862e3a970f731d709608b3e6276d5f6#132f5b39c862e3a970f731d709608b3e6276d5f6"
dependencies = [
"futures-util",
"log",
@@ -3175,12 +3209,13 @@ checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b"
[[package]]
name = "tungstenite"
version = "0.29.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6c01152af293afb9c7c2a57e4b559c5620b421f6d133261c60dd2d0cdb38e6b8"
version = "0.27.0"
source = "git+https://github.com/openai-oss-forks/tungstenite-rs?rev=9200079d3b54a1ff51072e24d81fd354f085156f#9200079d3b54a1ff51072e24d81fd354f085156f"
dependencies = [
"bytes",
"data-encoding",
"flate2",
"headers",
"http",
"httparse",
"log",
@@ -3189,6 +3224,7 @@ dependencies = [
"rustls-pki-types",
"sha1 0.10.6",
"thiserror 2.0.18",
"utf-8",
]
[[package]]
@@ -3276,6 +3312,12 @@ dependencies = [
"serde",
]
[[package]]
name = "utf-8"
version = "0.7.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9"
[[package]]
name = "utf8_iter"
version = "1.0.4"
+6 -2
View File
@@ -20,7 +20,6 @@ bytes = "1.11.1"
chrono = { version = "0.4.45", features = ["serde", "clock"] }
config = "0.15.23"
dirs = "5.0.1"
flate2 = "1.1.9"
futures = "0.3.32"
hex = "0.4.3"
hmac = "0.13.0"
@@ -39,15 +38,20 @@ sqlx = { version = "0.9.0", features = ["runtime-tokio", "tls-rustls", "sqlite",
thiserror = "2.0.18"
tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread", "signal", "fs", "time", "sync", "net", "test-util"] }
tokio-rustls = { version = "0.26.4", default-features = false, features = ["ring", "tls12"] }
tokio-tungstenite = { version = "0.29.0", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
tokio-tungstenite = { version = "0.28.0", features = ["proxy", "rustls-tls-native-roots"] }
tower = { version = "0.5.3", features = ["util"] }
tower-http = { version = "0.6.11", features = ["trace", "cors", "request-id", "timeout"] }
tracing = "0.1.44"
tracing-appender = "0.2.5"
tracing-subscriber = { version = "0.3.23", features = ["env-filter", "json"] }
tungstenite = { version = "0.27.0", features = ["deflate", "proxy"] }
uuid = { version = "1.23.3", features = ["v4", "v7", "serde"] }
zeroize = "1.8.2"
[patch.crates-io]
tokio-tungstenite = { git = "https://github.com/openai-oss-forks/tokio-tungstenite", rev = "132f5b39c862e3a970f731d709608b3e6276d5f6" }
tungstenite = { git = "https://github.com/openai-oss-forks/tungstenite-rs", rev = "9200079d3b54a1ff51072e24d81fd354f085156f" }
[workspace.lints.rust]
future_incompatible = { level = "warn", priority = -1 }
nonstandard_style = { level = "deny", priority = -1 }
+1 -1
View File
@@ -12,7 +12,6 @@ bytes.workspace = true
codex-proxy-core = { path = "../core" }
codex-proxy-platform = { path = "../platform" }
chrono.workspace = true
flate2.workspace = true
futures.workspace = true
indexmap.workspace = true
reqwest.workspace = true
@@ -28,6 +27,7 @@ tokio.workspace = true
tokio-rustls.workspace = true
tokio-tungstenite.workspace = true
tracing.workspace = true
tungstenite.workspace = true
uuid.workspace = true
[dev-dependencies]
+20 -15
View File
@@ -114,6 +114,8 @@ pub enum CodexClientError {
body: String,
/// 推导出的重试秒数。
retry_after_seconds: Option<u64>,
/// 上游透传的 `set-cookie` 列表。
set_cookie_headers: Vec<String>,
},
}
@@ -290,6 +292,8 @@ pub struct CodexRequestContext<'a> {
pub struct CodexBackendResponse {
/// 完整 SSE 文本。
pub body: String,
/// 实际使用的上游传输。
pub transport: CodexBackendTransport,
/// 从 SSE 中提取出的最终 usage。
pub usage: Option<TokenUsage>,
/// 响应头里的最新 turn state。
@@ -476,6 +480,7 @@ impl CodexBackendClient {
retry_after_seconds: retry_after_seconds
.or_else(|| retry_after_seconds_from_body(&body)),
body,
set_cookie_headers,
});
}
@@ -483,6 +488,7 @@ impl CodexBackendClient {
let usage = extract_sse_usage(&body).map_err(CodexClientError::InvalidSse)?;
Ok(CodexBackendResponse {
body,
transport: CodexBackendTransport::HttpSse,
usage,
turn_state,
set_cookie_headers,
@@ -516,6 +522,7 @@ impl CodexBackendClient {
retry_after_seconds: retry_after_seconds
.or_else(|| retry_after_seconds_from_body(&body)),
body,
set_cookie_headers,
});
}
@@ -539,7 +546,7 @@ impl CodexBackendClient {
let prepared = CodexWebSocketConnection::responses_create_request(
&self.base_url,
&generate_key(),
websocket_header_pairs(&headers, context),
websocket_header_pairs(&headers),
upstream_request,
)
.map_err(CodexClientError::WebSocketEncode)?;
@@ -562,6 +569,7 @@ impl CodexBackendClient {
Ok(CodexBackendResponse {
body: exchange.body,
transport: CodexBackendTransport::WebSocket,
usage: exchange.usage,
turn_state: exchange.turn_state,
set_cookie_headers: exchange.set_cookie_headers,
@@ -578,7 +586,7 @@ impl CodexBackendClient {
let prepared = CodexWebSocketConnection::responses_create_request(
&self.base_url,
&generate_key(),
websocket_header_pairs(&headers, context),
websocket_header_pairs(&headers),
upstream_request,
)
.map_err(CodexClientError::WebSocketEncode)?;
@@ -658,6 +666,7 @@ impl CodexBackendClient {
retry_after_seconds: retry_after_seconds
.or_else(|| retry_after_seconds_from_body(&body)),
body,
set_cookie_headers,
});
}
@@ -669,6 +678,7 @@ impl CodexBackendClient {
"Compact response is not valid JSON: {}",
truncate_for_error(&body)
),
set_cookie_headers: set_cookie_headers.clone(),
})?;
Ok(CodexCompactResponse {
@@ -699,6 +709,7 @@ impl CodexBackendClient {
retry_after_seconds: retry_after_seconds
.or_else(|| retry_after_seconds_from_body(&body)),
body,
set_cookie_headers: Vec::new(),
});
}
@@ -715,6 +726,7 @@ impl CodexBackendClient {
body: last_invalid_body
.map(|body| format!("invalid usage response: {}", truncate_for_error(&body)))
.unwrap_or_else(|| "usage endpoint is unavailable".to_string()),
set_cookie_headers: Vec::new(),
})
}
@@ -749,6 +761,7 @@ impl CodexBackendClient {
status: StatusCode::BAD_GATEWAY,
retry_after_seconds: None,
body: "backend model catalog is unavailable".to_string(),
set_cookie_headers: Vec::new(),
})
}
@@ -1201,24 +1214,14 @@ fn header_pairs(headers: &HeaderMap) -> Vec<(String, String)> {
.collect()
}
fn websocket_header_pairs(
headers: &HeaderMap,
context: CodexRequestContext<'_>,
) -> Vec<(String, String)> {
let mut pairs = header_pairs(headers)
fn websocket_header_pairs(headers: &HeaderMap) -> Vec<(String, String)> {
let pairs = header_pairs(headers)
.into_iter()
.filter(|(name, _)| {
!name.eq_ignore_ascii_case("content-type")
&& !name.eq_ignore_ascii_case("accept")
&& !name.eq_ignore_ascii_case("session_id")
!name.eq_ignore_ascii_case("content-type") && !name.eq_ignore_ascii_case("accept")
})
.collect::<Vec<_>>();
if let Some(session_id) = context.session_id {
pairs.push(("session-id".to_string(), session_id.to_string()));
pairs.push(("thread-id".to_string(), session_id.to_string()));
}
pairs
}
@@ -1230,10 +1233,12 @@ fn websocket_exchange_error_to_client_error(
status_code,
retry_after_seconds,
body,
set_cookie_headers,
} => CodexClientError::Upstream {
status: StatusCode::from_u16(status_code).unwrap_or(StatusCode::BAD_GATEWAY),
body,
retry_after_seconds,
set_cookie_headers,
},
error => CodexClientError::WebSocket(error),
}
+12 -1
View File
@@ -160,6 +160,17 @@ impl FingerprintRepository {
build_number: &str,
chromium_version: Option<&str>,
) -> Result<(), sqlx::Error> {
let resolved_chromium_version = match chromium_version {
Some(version) => version.to_string(),
None => sqlx::query_scalar::<_, String>(
"select chromium_version from fingerprints where id = ?",
)
.bind(AUTO_UPDATED_FINGERPRINT_ID)
.fetch_optional(&self.pool)
.await?
.unwrap_or_else(|| DEFAULT_AUTO_UPDATE_CHROMIUM_VERSION.to_string()),
};
sqlx::query(
"insert into fingerprints (id, app_version, build_number, platform, arch, chromium_version, user_agent_template, source, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?) \
on conflict(id) do update set app_version = excluded.app_version, build_number = excluded.build_number, platform = excluded.platform, arch = excluded.arch, chromium_version = excluded.chromium_version, user_agent_template = excluded.user_agent_template, source = excluded.source, created_at = excluded.created_at",
@@ -169,7 +180,7 @@ impl FingerprintRepository {
.bind(build_number)
.bind(AUTO_UPDATE_PLATFORM)
.bind(AUTO_UPDATE_ARCH)
.bind(chromium_version.unwrap_or(DEFAULT_AUTO_UPDATE_CHROMIUM_VERSION))
.bind(resolved_chromium_version)
.bind(AUTO_UPDATE_USER_AGENT_TEMPLATE)
.bind(AUTO_UPDATE_SOURCE)
.bind(Utc::now().to_rfc3339())
+51 -222
View File
@@ -1,6 +1,6 @@
//! Codex WebSocket 连接建立。
use std::{io, pin::Pin, sync::Arc, time::Duration, time::Instant};
use std::{pin::Pin, sync::Arc, time::Duration, time::Instant};
use bytes::Bytes;
use codex_proxy_core::protocol::codex::{
@@ -19,31 +19,21 @@ use codex_proxy_core::protocol::codex::{
},
};
use futures::{channel::mpsc, SinkExt, Stream, StreamExt};
use rustls_pki_types::ServerName;
use serde_json::Value;
use thiserror::Error;
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpStream,
sync::Mutex,
time::timeout,
};
use tokio_rustls::TlsConnector;
use tokio_tungstenite::{
tungstenite::{
self,
error::{CapacityError, ProtocolError, TlsError, UrlError},
handshake::derive_accept_key,
http::{HeaderName, HeaderValue, Response as WsResponse, StatusCode as WsStatusCode, Uri},
protocol::Role,
Message,
},
MaybeTlsStream, WebSocketStream,
use tokio::{sync::Mutex, time::timeout};
use tokio_tungstenite::{connect_async_tls_with_config, Connector};
use tungstenite::{
self,
extensions::{compression::deflate::DeflateConfig, ExtensionsConfig},
handshake::client::Request as WsRequest,
http::{HeaderMap, Response as WsResponse},
protocol::WebSocketConfig,
Message,
};
use crate::codex::client::{maybe_build_rustls_client_config_with_custom_ca, native_root_store};
use crate::codex::client::maybe_build_rustls_client_config_with_custom_ca;
use super::deflate::{permessage_deflate_enabled, PerMessageDeflateStream};
use super::pool::{
CodexWebSocketConnectionMetadata, CodexWebSocketPool, CodexWebSocketPoolKey, CodexWsStream,
PooledWebSocketConnection, WebSocketPoolAcquire,
@@ -52,7 +42,6 @@ use super::pool::{
const REDACTED_HEADER_VALUE: &str = "<redacted>";
const CODEX_RESPONSES_PATH: &str = "/codex/responses";
const WEBSOCKET_EXTENSIONS: &str = "permessage-deflate; client_max_window_bits";
const MAX_OPENING_RESPONSE_BYTES: usize = 64 * 1024;
const WEBSOCKET_RECEIVE_IDLE_TIMEOUT: Duration = Duration::from_secs(20);
/// WebSocket 连接适配器。
@@ -133,6 +122,8 @@ pub enum CodexWebSocketExchangeError {
retry_after_seconds: Option<u64>,
/// 原始错误帧。
body: String,
/// 上游透传的 `set-cookie` 列表。
set_cookie_headers: Vec<String>,
},
/// 上游返回 `response.incomplete`。
#[error("Incomplete response returned, reason: {reason}")]
@@ -479,28 +470,39 @@ fn websocket_host_header(endpoint: &str) -> Option<String> {
async fn connect_websocket(
connection: &CodexWebSocketConnection,
) -> Result<(CodexWsStream, WsResponse<Option<Vec<u8>>>), CodexWebSocketExchangeError> {
let uri = connection
.endpoint()
.parse::<Uri>()
.map_err(tungstenite::Error::from)?;
let mut stream = connect_stream(&uri).await?;
let request_bytes = opening_request_bytes(connection);
stream
.write_all(&request_bytes)
.await
.map_err(tungstenite::Error::Io)?;
stream.flush().await.map_err(tungstenite::Error::Io)?;
let (response, buffered) = read_opening_response(&mut stream).await?;
if response.status() != WsStatusCode::SWITCHING_PROTOCOLS {
return Err(websocket_opening_error(response));
let request = websocket_handshake_request(connection)?;
let connector = maybe_build_rustls_client_config_with_custom_ca()
.map_err(|error| tungstenite::Error::Io(std::io::Error::other(error)))?
.map(Connector::Rustls);
match connect_async_tls_with_config(request, Some(websocket_config()), false, connector).await {
Ok((websocket, response)) => Ok((websocket, response)),
Err(tungstenite::Error::Http(response)) => Err(websocket_opening_error(*response)),
Err(error) => Err(error.into()),
}
verify_accept_key(connection, &response)?;
let permessage_deflate = response_accepts_permessage_deflate(&response);
let stream = PerMessageDeflateStream::new(stream, permessage_deflate, buffered);
let websocket =
WebSocketStream::from_partially_read(stream, Vec::new(), Role::Client, None).await;
Ok((websocket, response))
}
fn websocket_handshake_request(
connection: &CodexWebSocketConnection,
) -> Result<WsRequest, tungstenite::http::Error> {
let mut builder = WsRequest::builder()
.method("GET")
.uri(connection.endpoint());
for (name, value) in connection.headers() {
if name.eq_ignore_ascii_case("sec-websocket-extensions") {
continue;
}
builder = builder.header(name.as_str(), value.as_str());
}
builder.body(())
}
fn websocket_config() -> WebSocketConfig {
let mut extensions = ExtensionsConfig::default();
extensions.permessage_deflate = Some(DeflateConfig::default());
let mut config = WebSocketConfig::default();
config.extensions = extensions;
config
}
fn websocket_opening_error(response: WsResponse<Option<Vec<u8>>>) -> CodexWebSocketExchangeError {
@@ -521,55 +523,10 @@ fn websocket_opening_error(response: WsResponse<Option<Vec<u8>>>) -> CodexWebSoc
status_code,
retry_after_seconds,
body,
set_cookie_headers: websocket_set_cookie_headers(response.headers()),
}
}
async fn connect_stream(uri: &Uri) -> Result<MaybeTlsStream<TcpStream>, tungstenite::Error> {
let host = uri.host().ok_or(UrlError::NoHostName)?;
if host.is_empty() {
return Err(UrlError::EmptyHostName.into());
}
let port = uri
.port_u16()
.or_else(|| match uri.scheme_str() {
Some("wss") => Some(443),
Some("ws") => Some(80),
_ => None,
})
.ok_or(UrlError::UnsupportedUrlScheme)?;
let socket = TcpStream::connect((host, port)).await?;
match uri.scheme_str() {
Some("ws") => Ok(MaybeTlsStream::Plain(socket)),
Some("wss") => connect_tls_stream(socket, host).await,
_ => Err(UrlError::UnsupportedUrlScheme.into()),
}
}
async fn connect_tls_stream(
socket: TcpStream,
host: &str,
) -> Result<MaybeTlsStream<TcpStream>, tungstenite::Error> {
let server_name =
ServerName::try_from(host.to_string()).map_err(|_| TlsError::InvalidDnsName)?;
let config = if let Some(config) = maybe_build_rustls_client_config_with_custom_ca()
.map_err(|error| tungstenite::Error::Io(io::Error::other(error)))?
{
config
} else {
Arc::new(
rustls::ClientConfig::builder()
.with_root_certificates(native_root_store().map_err(tungstenite::Error::Io)?)
.with_no_client_auth(),
)
};
let stream = TlsConnector::from(config)
.connect(server_name, socket)
.await
.map_err(tungstenite::Error::Io)?;
Ok(MaybeTlsStream::Rustls(stream))
}
fn opening_request_bytes(connection: &CodexWebSocketConnection) -> Vec<u8> {
let mut bytes = Vec::new();
bytes.extend_from_slice(request_line_for_endpoint(connection.endpoint()).as_bytes());
@@ -584,130 +541,6 @@ fn opening_request_bytes(connection: &CodexWebSocketConnection) -> Vec<u8> {
bytes
}
async fn read_opening_response(
stream: &mut MaybeTlsStream<TcpStream>,
) -> Result<(WsResponse<Option<Vec<u8>>>, Vec<u8>), tungstenite::Error> {
let mut bytes = Vec::new();
let mut chunk = [0_u8; 1024];
let header_end = loop {
let read = stream.read(&mut chunk).await?;
if read == 0 {
return Err(ProtocolError::HandshakeIncomplete.into());
}
bytes.extend_from_slice(&chunk[..read]);
if bytes.len() > MAX_OPENING_RESPONSE_BYTES {
return Err(CapacityError::MessageTooLong {
size: bytes.len(),
max_size: MAX_OPENING_RESPONSE_BYTES,
}
.into());
}
if let Some(index) = find_header_end(&bytes) {
break index;
}
};
let body_start = header_end + 4;
let mut body = bytes[body_start..].to_vec();
let mut response = parse_response_head(&bytes[..header_end])?;
if response.status() != WsStatusCode::SWITCHING_PROTOCOLS {
read_response_body(stream, response.headers(), &mut body).await?;
*response.body_mut() = Some(body);
return Ok((response, Vec::new()));
}
Ok((response, body))
}
fn find_header_end(bytes: &[u8]) -> Option<usize> {
bytes.windows(4).position(|window| window == b"\r\n\r\n")
}
fn parse_response_head(head: &[u8]) -> Result<WsResponse<Option<Vec<u8>>>, tungstenite::Error> {
let head = std::str::from_utf8(head)?;
let mut lines = head.split("\r\n");
let status_line = lines
.next()
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "empty response"))?;
let status = status_line
.split_whitespace()
.nth(1)
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "missing response status"))?
.parse::<u16>()
.map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
let mut response = WsResponse::new(None);
*response.status_mut() = WsStatusCode::from_u16(status)?;
for line in lines {
let Some((name, value)) = line.split_once(':') else {
continue;
};
response.headers_mut().append(
HeaderName::from_bytes(name.trim().as_bytes())?,
HeaderValue::from_bytes(value.trim_start().as_bytes())?,
);
}
Ok(response)
}
async fn read_response_body(
stream: &mut MaybeTlsStream<TcpStream>,
headers: &tokio_tungstenite::tungstenite::http::HeaderMap,
body: &mut Vec<u8>,
) -> Result<(), tungstenite::Error> {
let Some(content_length) = headers
.get("content-length")
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse::<usize>().ok())
else {
return Ok(());
};
while body.len() < content_length {
let remaining = content_length - body.len();
let mut chunk = vec![0_u8; remaining.min(1024)];
let read = stream.read(&mut chunk).await?;
if read == 0 {
break;
}
body.extend_from_slice(&chunk[..read]);
}
body.truncate(content_length);
Ok(())
}
fn verify_accept_key(
connection: &CodexWebSocketConnection,
response: &WsResponse<Option<Vec<u8>>>,
) -> Result<(), tungstenite::Error> {
let key = connection
.headers()
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("sec-websocket-key"))
.map(|(_, value)| value.as_str())
.ok_or(ProtocolError::MissingSecWebSocketKey)?;
let expected = derive_accept_key(key.as_bytes());
let actual = response
.headers()
.get("sec-websocket-accept")
.ok_or(ProtocolError::SecWebSocketAcceptKeyMismatch)?;
if actual == expected.as_str() {
Ok(())
} else {
Err(ProtocolError::SecWebSocketAcceptKeyMismatch.into())
}
}
fn response_accepts_permessage_deflate(response: &WsResponse<Option<Vec<u8>>>) -> bool {
response
.headers()
.get_all("sec-websocket-extensions")
.iter()
.filter_map(|value| value.to_str().ok())
.any(permessage_deflate_enabled)
}
async fn collect_websocket_response(
mut websocket: CodexWsStream,
mut metadata: CodexWebSocketConnectionMetadata,
@@ -742,6 +575,7 @@ async fn collect_websocket_response(
status_code: classified.status_code,
retry_after_seconds,
body: raw,
set_cookie_headers: Vec::new(),
});
}
if let Some(reason) = websocket_incomplete_response_reason(&raw) {
@@ -887,6 +721,7 @@ async fn forward_websocket_response_stream(
status_code: classified.status_code,
retry_after_seconds,
body: raw,
set_cookie_headers: Vec::new(),
}));
return;
}
@@ -990,18 +825,14 @@ fn websocket_event_type(raw: &str) -> Option<String> {
})
}
fn websocket_response_turn_state(
headers: &tokio_tungstenite::tungstenite::http::HeaderMap,
) -> Option<String> {
fn websocket_response_turn_state(headers: &HeaderMap) -> Option<String> {
headers
.get("x-codex-turn-state")
.and_then(|value| value.to_str().ok())
.map(ToString::to_string)
}
fn websocket_set_cookie_headers(
headers: &tokio_tungstenite::tungstenite::http::HeaderMap,
) -> Vec<String> {
fn websocket_set_cookie_headers(headers: &HeaderMap) -> Vec<String> {
headers
.get_all("set-cookie")
.iter()
@@ -1009,9 +840,7 @@ fn websocket_set_cookie_headers(
.collect()
}
fn websocket_rate_limit_headers(
headers: &tokio_tungstenite::tungstenite::http::HeaderMap,
) -> Vec<(String, String)> {
fn websocket_rate_limit_headers(headers: &HeaderMap) -> Vec<(String, String)> {
headers
.iter()
.filter(|(name, _)| is_rate_limit_header(name.as_str()))
@@ -1,309 +0,0 @@
//! WebSocket permessage-deflate 处理。
use std::{
collections::VecDeque,
io::{self, Read},
pin::Pin,
task::{Context, Poll},
};
use flate2::read::DeflateDecoder;
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
const PERMESSAGE_DEFLATE_TAIL: [u8; 4] = [0x00, 0x00, 0xff, 0xff];
const OPCODE_CONTINUATION: u8 = 0x0;
const OPCODE_TEXT: u8 = 0x1;
const OPCODE_BINARY: u8 = 0x2;
const SERVER_FRAME_CHUNK_BYTES: usize = 8192;
/// Async stream wrapper that inflates server permessage-deflate data frames
/// before tungstenite sees them.
pub(crate) struct PerMessageDeflateStream<S> {
inner: S,
enabled: bool,
raw_input: Vec<u8>,
decoded_input: VecDeque<u8>,
compressed_message: Option<CompressedMessage>,
}
struct CompressedMessage {
opcode: u8,
payload: Vec<u8>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct FrameParts {
fin: bool,
rsv1: bool,
opcode: u8,
masked: bool,
payload_offset: usize,
payload_len: usize,
frame_len: usize,
}
impl<S> PerMessageDeflateStream<S> {
pub(crate) fn new(inner: S, enabled: bool, preloaded: Vec<u8>) -> Self {
Self {
inner,
enabled,
raw_input: preloaded,
decoded_input: VecDeque::new(),
compressed_message: None,
}
}
}
impl<S> AsyncRead for PerMessageDeflateStream<S>
where
S: AsyncRead + AsyncWrite + Unpin,
{
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buffer: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
let this = self.get_mut();
if !this.enabled {
return Pin::new(&mut this.inner).poll_read(cx, buffer);
}
loop {
if copy_decoded_input(&mut this.decoded_input, buffer) {
return Poll::Ready(Ok(()));
}
match this.rewrite_next_frame() {
Ok(true) => continue,
Ok(false) => {}
Err(error) => return Poll::Ready(Err(error)),
}
let mut chunk = [0_u8; SERVER_FRAME_CHUNK_BYTES];
let mut read_buffer = ReadBuf::new(&mut chunk);
match Pin::new(&mut this.inner).poll_read(cx, &mut read_buffer) {
Poll::Pending => return Poll::Pending,
Poll::Ready(Err(error)) => return Poll::Ready(Err(error)),
Poll::Ready(Ok(())) => {
let filled = read_buffer.filled();
if filled.is_empty() {
return Poll::Ready(Ok(()));
}
this.raw_input.extend_from_slice(filled);
}
}
}
}
}
impl<S> AsyncWrite for PerMessageDeflateStream<S>
where
S: AsyncRead + AsyncWrite + Unpin,
{
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
data: &[u8],
) -> Poll<io::Result<usize>> {
Pin::new(&mut self.get_mut().inner).poll_write(cx, data)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.get_mut().inner).poll_flush(cx)
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.get_mut().inner).poll_shutdown(cx)
}
}
impl<S> PerMessageDeflateStream<S> {
fn rewrite_next_frame(&mut self) -> io::Result<bool> {
let Some(frame) = parse_frame_parts(&self.raw_input)? else {
return Ok(false);
};
let raw_frame = self.raw_input.drain(..frame.frame_len).collect::<Vec<_>>();
let payload = &raw_frame[frame.payload_offset..frame.payload_offset + frame.payload_len];
if frame.masked {
self.decoded_input.extend(raw_frame);
return Ok(true);
}
if frame.rsv1 && is_initial_data_frame(frame.opcode) {
if frame.fin {
let inflated = inflate_permessage_deflate_message(payload)?;
self.decoded_input
.extend(encode_server_frame(frame.opcode, false, &inflated));
} else {
self.compressed_message = Some(CompressedMessage {
opcode: frame.opcode,
payload: payload.to_vec(),
});
}
return Ok(true);
}
if frame.opcode == OPCODE_CONTINUATION {
if let Some(message) = self.compressed_message.as_mut() {
message.payload.extend_from_slice(payload);
if frame.fin {
let message = self.compressed_message.take().ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidData, "missing compressed message")
})?;
let inflated = inflate_permessage_deflate_message(&message.payload)?;
self.decoded_input.extend(encode_server_frame(
message.opcode,
false,
&inflated,
));
}
return Ok(true);
}
}
self.decoded_input.extend(raw_frame);
Ok(true)
}
}
fn copy_decoded_input(decoded_input: &mut VecDeque<u8>, buffer: &mut ReadBuf<'_>) -> bool {
if decoded_input.is_empty() || buffer.remaining() == 0 {
return false;
}
let copy_len = decoded_input.len().min(buffer.remaining());
let bytes = decoded_input.drain(..copy_len).collect::<Vec<_>>();
buffer.put_slice(&bytes);
true
}
/// 判断扩展响应是否启用了 permessage-deflate。
pub fn permessage_deflate_enabled(extension_header: &str) -> bool {
extension_header
.split(',')
.any(|value| value.trim().starts_with("permessage-deflate"))
}
/// 解压单条 permessage-deflate 消息 payload。
pub fn inflate_permessage_deflate_message(payload: &[u8]) -> io::Result<Vec<u8>> {
let mut compressed = Vec::with_capacity(payload.len() + PERMESSAGE_DEFLATE_TAIL.len());
compressed.extend_from_slice(payload);
compressed.extend_from_slice(&PERMESSAGE_DEFLATE_TAIL);
let mut decoder = DeflateDecoder::new(compressed.as_slice());
let mut inflated = Vec::new();
decoder.read_to_end(&mut inflated)?;
if inflated.is_empty() {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"permessage-deflate frame produced no output",
));
}
Ok(inflated)
}
/// 将启用 RSV1 的压缩服务端 data frame 重写为普通未压缩 frame。
pub fn rewrite_permessage_deflate_server_frame(frame: &[u8]) -> io::Result<Option<Vec<u8>>> {
let Some(parts) = parse_frame_parts(frame)? else {
return Ok(None);
};
if parts.frame_len != frame.len() || parts.masked || !parts.rsv1 || !parts.fin {
return Ok(None);
}
if !is_data_frame(parts.opcode) {
return Ok(None);
}
let payload = &frame[parts.payload_offset..parts.payload_offset + parts.payload_len];
let inflated = inflate_permessage_deflate_message(payload)?;
Ok(Some(encode_server_frame(parts.opcode, false, &inflated)))
}
fn parse_frame_parts(bytes: &[u8]) -> io::Result<Option<FrameParts>> {
if bytes.len() < 2 {
return Ok(None);
}
let first = bytes[0];
let second = bytes[1];
let mut offset = 2;
let payload_len = match second & 0x7f {
len @ 0..=125 => usize::from(len),
126 => {
if bytes.len() < offset + 2 {
return Ok(None);
}
let len = u16::from_be_bytes([bytes[offset], bytes[offset + 1]]);
offset += 2;
usize::from(len)
}
127 => {
if bytes.len() < offset + 8 {
return Ok(None);
}
let len = u64::from_be_bytes([
bytes[offset],
bytes[offset + 1],
bytes[offset + 2],
bytes[offset + 3],
bytes[offset + 4],
bytes[offset + 5],
bytes[offset + 6],
bytes[offset + 7],
]);
offset += 8;
usize::try_from(len).map_err(|_| {
io::Error::new(io::ErrorKind::InvalidData, "websocket frame is too large")
})?
}
_ => unreachable!("websocket length marker is masked to 7 bits"),
};
let masked = second & 0x80 != 0;
if masked {
offset += 4;
}
let frame_len = offset
.checked_add(payload_len)
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "websocket frame overflow"))?;
if bytes.len() < frame_len {
return Ok(None);
}
Ok(Some(FrameParts {
fin: first & 0x80 != 0,
rsv1: first & 0x40 != 0,
opcode: first & 0x0f,
masked,
payload_offset: offset,
payload_len,
frame_len,
}))
}
fn encode_server_frame(opcode: u8, rsv1: bool, payload: &[u8]) -> Vec<u8> {
let mut frame = Vec::new();
let rsv1_bit = if rsv1 { 0x40 } else { 0 };
frame.push(0x80 | rsv1_bit | opcode);
match payload.len() {
len @ 0..=125 => frame.push(len as u8),
len @ 126..=65_535 => {
frame.push(126);
frame.extend_from_slice(&(len as u16).to_be_bytes());
}
len => {
frame.push(127);
frame.extend_from_slice(&(len as u64).to_be_bytes());
}
}
frame.extend_from_slice(payload);
frame
}
fn is_data_frame(opcode: u8) -> bool {
matches!(opcode, OPCODE_TEXT | OPCODE_BINARY | OPCODE_CONTINUATION)
}
fn is_initial_data_frame(opcode: u8) -> bool {
matches!(opcode, OPCODE_TEXT | OPCODE_BINARY)
}
@@ -1,6 +1,5 @@
//! Codex WebSocket 适配器。
pub mod connect;
pub mod deflate;
pub mod opening;
pub mod pool;
+1 -3
View File
@@ -10,8 +10,6 @@ use futures::{SinkExt, StreamExt};
use tokio::{net::TcpStream, sync::Mutex, time::timeout};
use tokio_tungstenite::{tungstenite::Message, MaybeTlsStream, WebSocketStream};
use super::deflate::PerMessageDeflateStream;
const DEFAULT_MAX_PER_ACCOUNT: usize = 8;
const DEFAULT_MAX_AGE: Duration = Duration::from_secs(55 * 60);
const DEFAULT_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(25);
@@ -426,7 +424,7 @@ struct WebSocketPoolState {
shutting_down: bool,
}
pub(crate) type CodexWsStream = WebSocketStream<PerMessageDeflateStream<MaybeTlsStream<TcpStream>>>;
pub(crate) type CodexWsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
#[derive(Clone)]
pub(crate) struct CodexWebSocketConnectionMetadata {
+24 -1
View File
@@ -126,7 +126,7 @@ impl TokenRefresher for OpenAiOAuthClient {
])
.send()
.await
.map_err(|_| RefreshFailure::Transport)?;
.map_err(|error| refresh_transport_failure(&error))?;
let status = response.status();
let body = response
.text()
@@ -188,3 +188,26 @@ fn classify_refresh_failure(status: StatusCode, body: &str) -> RefreshFailure {
}
RefreshFailure::Transport
}
fn refresh_transport_failure(error: &reqwest::Error) -> RefreshFailure {
if is_safe_to_retry_refresh_transport(error) {
RefreshFailure::RetryableTransport
} else {
RefreshFailure::Transport
}
}
fn is_safe_to_retry_refresh_transport(error: &reqwest::Error) -> bool {
let message = error.to_string().to_ascii_lowercase();
message.contains("econnrefused")
|| message.contains("could not resolve proxy")
|| message.contains("could not resolve host")
|| message.contains("curl exited with code 5")
|| message.contains("curl exited with code 6")
|| message.contains("curl exited with code 7")
|| message.contains("curl exited with code 35")
|| message.contains("dns error")
|| message.contains("connection refused")
|| message.contains("network is unreachable")
|| message.contains("tls handshake")
}
+237
View File
@@ -43,6 +43,7 @@ select
access_token_cipher,
refresh_token_cipher,
access_token_expires_at,
next_refresh_at,
status,
added_at,
updated_at
@@ -60,6 +61,7 @@ select
access_token_cipher,
refresh_token_cipher,
access_token_expires_at,
next_refresh_at,
status,
added_at,
updated_at
@@ -80,6 +82,7 @@ select
access_token_cipher,
refresh_token_cipher,
access_token_expires_at,
next_refresh_at,
status,
added_at,
updated_at
@@ -100,6 +103,7 @@ select
access_token_cipher,
refresh_token_cipher,
access_token_expires_at,
next_refresh_at,
status,
added_at,
updated_at
@@ -152,6 +156,7 @@ select
a.access_token_cipher,
a.refresh_token_cipher,
a.access_token_expires_at,
a.next_refresh_at,
a.status,
a.quota_limit_reached,
a.quota_verify_required,
@@ -191,6 +196,7 @@ select
a.access_token_cipher,
a.refresh_token_cipher,
a.access_token_expires_at,
a.next_refresh_at,
a.status,
a.quota_limit_reached,
a.quota_verify_required,
@@ -395,6 +401,17 @@ set
updated_at = ?
where id = ?";
const APPLY_QUOTA_SNAPSHOT_SQL: &str = r"
update accounts
set
quota_json = ?,
quota_fetched_at = ?,
quota_limit_reached = ?,
quota_verify_required = 0,
quota_cooldown_until = ?,
updated_at = ?
where id = ?";
const SELECT_RATE_LIMIT_WINDOW_SQL: &str = r"
select
window_reset_at,
@@ -466,6 +483,7 @@ set
access_token_cipher = ?,
refresh_token_cipher = ?,
access_token_expires_at = ?,
next_refresh_at = ?,
status = ?,
updated_at = ?
where id = ?";
@@ -479,10 +497,68 @@ set
plan_type = ?,
access_token_cipher = ?,
access_token_expires_at = ?,
next_refresh_at = ?,
status = ?,
updated_at = ?
where id = ?";
const SET_NEXT_REFRESH_AT_SQL: &str = r"
update accounts
set
next_refresh_at = ?,
updated_at = ?
where id = ?";
const UPDATE_IMPORTED_ACCOUNT_WITH_REFRESH_SQL: &str = r"
update accounts
set
email = ?,
account_id = ?,
user_id = ?,
label = ?,
plan_type = ?,
access_token_cipher = ?,
refresh_token_cipher = ?,
access_token_expires_at = ?,
status = ?,
quota_json = coalesce(?, quota_json),
quota_fetched_at = case when ? is null then quota_fetched_at else ? end,
quota_limit_reached = 0,
quota_cooldown_until = null,
quota_verify_required = ?,
updated_at = ?
where id = ?";
const UPDATE_IMPORTED_ACCOUNT_PRESERVING_REFRESH_SQL: &str = r"
update accounts
set
email = ?,
account_id = ?,
user_id = ?,
label = ?,
plan_type = ?,
access_token_cipher = ?,
access_token_expires_at = ?,
status = ?,
quota_json = coalesce(?, quota_json),
quota_fetched_at = case when ? is null then quota_fetched_at else ? end,
quota_limit_reached = 0,
quota_cooldown_until = null,
quota_verify_required = ?,
updated_at = ?
where id = ?";
const APPLY_IMPORTED_QUOTA_STATE_SQL: &str = r"
update accounts
set
quota_json = coalesce(?, quota_json),
quota_fetched_at = case when ? is null then quota_fetched_at else ? end,
quota_limit_reached = 0,
quota_cooldown_until = null,
quota_verify_required = ?,
updated_at = ?
where id = ?";
/// SQLite 账号仓储错误。
#[derive(Debug, Error)]
pub enum SqliteAccountStoreError {
@@ -548,10 +624,25 @@ pub struct AccountClaimsUpdate {
pub refresh_token: Option<SecretString>,
/// access token 过期时间。
pub access_token_expires_at: Option<DateTime<Utc>>,
/// 下一次允许刷新 token 的时间。
pub next_refresh_at: Option<DateTime<Utc>>,
/// 新状态。
pub status: AccountStatus,
}
/// native 导入时用于覆盖已有账号的数据。
#[derive(Debug)]
pub struct ImportedAccountUpdate {
/// 账号基础信息和 token。
pub account: NewAccount,
/// 可选的配额快照 JSON。
pub quota_json: Option<String>,
/// 配额快照拉取时间。
pub quota_fetched_at: Option<DateTime<Utc>>,
/// 是否需要重新校验配额。
pub quota_verify_required: bool,
}
/// 已存储账号数据。
#[derive(Debug, Clone)]
pub struct StoredAccount {
@@ -573,6 +664,8 @@ pub struct StoredAccount {
pub refresh_token: Option<SecretString>,
/// access token 过期时间。
pub access_token_expires_at: Option<DateTime<Utc>>,
/// 下一次允许刷新 token 的时间。
pub next_refresh_at: Option<DateTime<Utc>>,
/// 账号状态。
pub status: AccountStatus,
/// 创建时间。
@@ -858,6 +951,7 @@ impl SqliteAccountStore {
update.access_token_expires_at,
)?;
let status = status_to_db(update.status);
let next_refresh_at = update.next_refresh_at.map(|value| value.to_rfc3339());
let result = if let Some(refresh_token_cipher) = refresh_token_cipher {
sqlx::query(UPDATE_ACCOUNT_CLAIMS_WITH_REFRESH_SQL)
@@ -868,6 +962,7 @@ impl SqliteAccountStore {
.bind(access_token_cipher)
.bind(refresh_token_cipher)
.bind(access_token_expires_at)
.bind(next_refresh_at.as_deref())
.bind(status)
.bind(updated_at)
.bind(id)
@@ -881,6 +976,7 @@ impl SqliteAccountStore {
.bind(update.plan_type)
.bind(access_token_cipher)
.bind(access_token_expires_at)
.bind(next_refresh_at.as_deref())
.bind(status)
.bind(updated_at)
.bind(id)
@@ -890,6 +986,88 @@ impl SqliteAccountStore {
Ok(result.rows_affected() > 0)
}
/// 用 native 导入数据覆盖已有账号,并清理本地配额冷却状态。
pub async fn update_from_import(
&self,
update: ImportedAccountUpdate,
) -> SqliteAccountStoreResult<bool> {
let account = update.account;
let TokenWrite {
access_token_cipher,
refresh_token_cipher,
access_token_expires_at,
updated_at,
} = self.prepare_token_write(
&account.access_token,
account.refresh_token.as_ref(),
account.access_token_expires_at,
)?;
let status = status_to_db(account.status);
let quota_fetched_at = update.quota_fetched_at.map(|value| value.to_rfc3339());
let quota_verify_required = if update.quota_verify_required { 1 } else { 0 };
let result = if let Some(refresh_token_cipher) = refresh_token_cipher {
sqlx::query(UPDATE_IMPORTED_ACCOUNT_WITH_REFRESH_SQL)
.bind(account.email)
.bind(account.account_id)
.bind(account.user_id)
.bind(account.label)
.bind(account.plan_type)
.bind(access_token_cipher)
.bind(refresh_token_cipher)
.bind(access_token_expires_at)
.bind(status)
.bind(update.quota_json.as_deref())
.bind(quota_fetched_at.as_deref())
.bind(quota_fetched_at.as_deref())
.bind(quota_verify_required)
.bind(updated_at)
.bind(account.id)
.execute(&self.pool)
.await?
} else {
sqlx::query(UPDATE_IMPORTED_ACCOUNT_PRESERVING_REFRESH_SQL)
.bind(account.email)
.bind(account.account_id)
.bind(account.user_id)
.bind(account.label)
.bind(account.plan_type)
.bind(access_token_cipher)
.bind(access_token_expires_at)
.bind(status)
.bind(update.quota_json.as_deref())
.bind(quota_fetched_at.as_deref())
.bind(quota_fetched_at.as_deref())
.bind(quota_verify_required)
.bind(updated_at)
.bind(account.id)
.execute(&self.pool)
.await?
};
Ok(result.rows_affected() > 0)
}
/// 应用 native 导入中的配额快照,并清理本地配额冷却状态。
pub async fn apply_import_quota_state(
&self,
account_id: &str,
quota_json: Option<&str>,
quota_fetched_at: Option<DateTime<Utc>>,
quota_verify_required: bool,
) -> SqliteAccountStoreResult<bool> {
let quota_fetched_at = quota_fetched_at.map(|value| value.to_rfc3339());
let result = sqlx::query(APPLY_IMPORTED_QUOTA_STATE_SQL)
.bind(quota_json)
.bind(quota_fetched_at.as_deref())
.bind(quota_fetched_at.as_deref())
.bind(if quota_verify_required { 1 } else { 0 })
.bind(Utc::now().to_rfc3339())
.bind(account_id)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
/// 按创建时间倒序分页列出账号。
pub async fn list(
&self,
@@ -986,6 +1164,22 @@ impl SqliteAccountStore {
Ok(result.rows_affected() > 0)
}
/// 更新下一次允许刷新 token 的时间。
pub async fn set_next_refresh_at(
&self,
id: &str,
next_refresh_at: Option<DateTime<Utc>>,
) -> SqliteAccountStoreResult<bool> {
let next_refresh_at = next_refresh_at.map(|value| value.to_rfc3339());
let result = sqlx::query(SET_NEXT_REFRESH_AT_SQL)
.bind(next_refresh_at.as_deref())
.bind(Utc::now().to_rfc3339())
.bind(id)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
/// 更新账号标签。
pub async fn set_label(
&self,
@@ -1173,6 +1367,27 @@ impl SqliteAccountStore {
Ok(result.rows_affected() > 0)
}
/// 应用已经验证过的账号配额快照。
pub async fn apply_quota_snapshot(
&self,
account_id: &str,
quota_json: &str,
limit_reached: bool,
cooldown_until: Option<DateTime<Utc>>,
) -> SqliteAccountStoreResult<bool> {
let now = Utc::now().to_rfc3339();
let result = sqlx::query(APPLY_QUOTA_SNAPSHOT_SQL)
.bind(quota_json)
.bind(&now)
.bind(if limit_reached { 1 } else { 0 })
.bind(cooldown_until.map(|value| value.to_rfc3339()))
.bind(now)
.bind(account_id)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
/// 标记账号进入配额冷却期。
pub async fn mark_quota_limited_until(
&self,
@@ -1285,6 +1500,9 @@ impl SqliteAccountStore {
access_token_expires_at: parse_optional_rfc3339(
row.get::<Option<String>, _>("access_token_expires_at"),
)?,
next_refresh_at: parse_optional_rfc3339(
row.get::<Option<String>, _>("next_refresh_at"),
)?,
status: status_from_db(&row.get::<String, _>("status"))?,
added_at: parse_rfc3339(&row.get::<String, _>("added_at"))?,
updated_at: parse_rfc3339(&row.get::<String, _>("updated_at"))?,
@@ -1397,6 +1615,24 @@ impl AccountStore for SqliteAccountStore {
.map_err(map_account_store_error)
}
async fn apply_quota_snapshot(
&self,
account_id: &str,
quota_json: &str,
limit_reached: bool,
cooldown_until: Option<DateTime<Utc>>,
) -> AccountStoreResult<bool> {
SqliteAccountStore::apply_quota_snapshot(
self,
account_id,
quota_json,
limit_reached,
cooldown_until,
)
.await
.map_err(map_account_store_error)
}
async fn sync_rate_limit_window(
&self,
account_id: &str,
@@ -1458,6 +1694,7 @@ fn pool_account_from_row(
access_token_expires_at: parse_optional_rfc3339(
row.get::<Option<String>, _>("access_token_expires_at"),
)?,
next_refresh_at: parse_optional_rfc3339(row.get::<Option<String>, _>("next_refresh_at"))?,
status: status_from_db(&row.get::<String, _>("status"))?,
quota_limit_reached: row.get::<i64, _>("quota_limit_reached") != 0,
quota_verify_required: row.get::<i64, _>("quota_verify_required") != 0,
+27 -52
View File
@@ -1,5 +1,4 @@
use std::{
io::Write,
process::Command,
sync::{
atomic::{AtomicUsize, Ordering},
@@ -12,7 +11,6 @@ use codex_proxy_adapters::codex::{
client,
websocket::pool::{CodexWebSocketPool, CodexWebSocketPoolConfig},
};
use flate2::{write::DeflateEncoder, Compression};
use futures::{SinkExt, StreamExt};
use serde_json::json;
use tokio::{
@@ -21,12 +19,11 @@ use tokio::{
time::timeout,
};
use tokio_tungstenite::{
accept_async, accept_hdr_async,
accept_hdr_async_with_config,
tungstenite::{
handshake::{
derive_accept_key,
server::{Request as WsRequest, Response as WsResponse},
},
extensions::{compression::deflate::DeflateConfig, ExtensionsConfig},
handshake::server::{Request as WsRequest, Response as WsResponse},
protocol::WebSocketConfig,
Message,
},
};
@@ -42,44 +39,31 @@ mod codex_websocket;
#[path = "codex/codex_websocket_pool.rs"]
mod codex_websocket_pool;
fn compressed_permessage_deflate_payload(payload: &[u8]) -> Vec<u8> {
let mut encoder = DeflateEncoder::new(Vec::new(), Compression::fast());
encoder.write_all(payload).expect("payload should compress");
encoder.finish().expect("compressed payload")
fn websocket_accept_config() -> WebSocketConfig {
let mut extensions = ExtensionsConfig::default();
extensions.permessage_deflate = Some(DeflateConfig::default());
let mut config = WebSocketConfig::default();
config.extensions = extensions;
config
}
fn server_websocket_frame(opcode: u8, rsv1: bool, payload: &[u8]) -> Vec<u8> {
let mut frame = Vec::new();
frame.push(0x80 | if rsv1 { 0x40 } else { 0 } | opcode);
match payload.len() {
len @ 0..=125 => frame.push(len as u8),
len @ 126..=65_535 => {
frame.push(126);
frame.extend_from_slice(&(len as u16).to_be_bytes());
}
len => {
frame.push(127);
frame.extend_from_slice(&(len as u64).to_be_bytes());
}
}
frame.extend_from_slice(payload);
frame
}
fn server_websocket_payload(frame: &[u8]) -> &[u8] {
let len_marker = frame[1] & 0x7f;
let (offset, len) = match len_marker {
len @ 0..=125 => (2, usize::from(len)),
126 => (4, usize::from(u16::from_be_bytes([frame[2], frame[3]]))),
127 => {
let len = u64::from_be_bytes([
frame[2], frame[3], frame[4], frame[5], frame[6], frame[7], frame[8], frame[9],
]) as usize;
(10, len)
}
_ => unreachable!("websocket length marker is masked to 7 bits"),
};
&frame[offset..offset + len]
async fn accept_codex_test_websocket(
stream: TcpStream,
) -> tokio_tungstenite::WebSocketStream<TcpStream> {
accept_hdr_async_with_config(
stream,
|_request: &WsRequest, mut response: WsResponse| {
response.headers_mut().insert(
"sec-websocket-extensions",
"permessage-deflate".parse().unwrap(),
);
Ok(response)
},
Some(websocket_accept_config()),
)
.await
.unwrap()
}
fn read_header_names(request: &str) -> Vec<String> {
@@ -123,15 +107,6 @@ async fn read_http_request(stream: &mut TcpStream) -> String {
String::from_utf8(request).unwrap()
}
fn websocket_opening_header(request: &str, name: &str) -> Option<String> {
request.lines().find_map(|line| {
let (header_name, value) = line.split_once(':')?;
header_name
.eq_ignore_ascii_case(name)
.then(|| value.trim().to_string())
})
}
async fn write_empty_http_response(stream: &mut TcpStream) {
stream
.write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 0\r\n\r\n")
@@ -39,7 +39,7 @@ async fn fingerprint_repository_should_upsert_auto_update_record() {
assert_eq!(count.0, 1);
assert_eq!(stored.app_version, "26.800.2");
assert_eq!(stored.build_number, "6002");
assert_eq!(stored.chromium_version, "146");
assert_eq!(stored.chromium_version, "147");
}
#[tokio::test]
+58 -15
View File
@@ -91,6 +91,10 @@ fn websocket_connection_should_build_standard_opening_headers_around_business_he
"https://chatgpt.com/backend-api",
"test-websocket-key",
vec![
(
"chatgpt-account-id".to_string(),
"chatgpt-account".to_string(),
),
(
"authorization".to_string(),
"Bearer access-token".to_string(),
@@ -117,6 +121,7 @@ fn websocket_connection_should_build_standard_opening_headers_around_business_he
"Upgrade",
"Sec-WebSocket-Version",
"Sec-WebSocket-Key",
"chatgpt-account-id",
"authorization",
"user-agent",
"openai-beta",
@@ -135,6 +140,7 @@ fn websocket_connection_should_build_standard_opening_headers_around_business_he
("Upgrade", "websocket"),
("Sec-WebSocket-Version", "13"),
("Sec-WebSocket-Key", "test-websocket-key"),
("chatgpt-account-id", "<redacted>"),
("authorization", "<redacted>"),
("user-agent", "Codex Desktop/test"),
("openai-beta", "responses_websockets=2026-02-06"),
@@ -156,8 +162,13 @@ async fn websocket_execute_response_create_request_should_capture_handshake_head
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket =
accept_hdr_async(stream, |_request: &WsRequest, mut response: WsResponse| {
let mut websocket = accept_hdr_async_with_config(
stream,
|_request: &WsRequest, mut response: WsResponse| {
response.headers_mut().insert(
"sec-websocket-extensions",
"permessage-deflate".parse().unwrap(),
);
response
.headers_mut()
.insert("x-codex-turn-state", "turn-from-handshake".parse().unwrap());
@@ -171,9 +182,11 @@ async fn websocket_execute_response_create_request_should_capture_handshake_head
.headers_mut()
.insert("x-ratelimit-remaining-requests", "41".parse().unwrap());
Ok(response)
})
.await
.unwrap();
},
Some(websocket_accept_config()),
)
.await
.unwrap();
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
@@ -239,8 +252,14 @@ async fn codex_backend_client_websocket_should_forward_security_chain_headers_an
let headers_for_server = Arc::clone(&received_headers);
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket =
accept_hdr_async(stream, move |request: &WsRequest, response: WsResponse| {
let mut websocket = accept_hdr_async_with_config(
stream,
move |request: &WsRequest, response: WsResponse| {
let mut response = response;
response.headers_mut().insert(
"sec-websocket-extensions",
"permessage-deflate".parse().unwrap(),
);
let headers = request
.headers()
.iter()
@@ -253,9 +272,11 @@ async fn codex_backend_client_websocket_should_forward_security_chain_headers_an
.collect::<Vec<_>>();
*headers_for_server.lock().unwrap() = headers;
Ok(response)
})
.await
.unwrap();
},
Some(websocket_accept_config()),
)
.await
.unwrap();
let message = websocket.next().await.unwrap().unwrap();
let payload = serde_json::from_str::<serde_json::Value>(&message.into_text().unwrap())
.expect("client payload should be json");
@@ -328,18 +349,40 @@ async fn codex_backend_client_websocket_should_forward_security_chain_headers_an
assert!(headers
.iter()
.any(|(name, value)| name == "x-client-request-id" && value == "cp_derived"));
assert!(headers
.iter()
.any(|(name, value)| name == "x-codex-installation-id" && value == "install-123"));
assert!(headers
.iter()
.any(|(name, value)| name == "x-openai-internal-codex-residency" && value == "us"));
assert!(headers
.iter()
.any(|(name, value)| name == "x-codex-turn-state" && value == "turn-state"));
assert!(headers.iter().any(|(name, value)| {
name == "x-codex-turn-metadata" && value == "{\"thread_source\":\"subagent\"}"
}));
assert!(headers
.iter()
.any(|(name, value)| name == "x-codex-beta-features" && value == "feature-a"));
assert!(headers.iter().any(|(name, value)| {
name == "x-responsesapi-include-timing-metrics" && value == "true"
}));
assert!(headers
.iter()
.any(|(name, value)| name == "version" && value == "26.318.11754"));
assert!(headers
.iter()
.any(|(name, value)| name == "x-codex-parent-thread-id" && value == "parent-456"));
assert!(headers
.iter()
.any(|(name, value)| name == "x-openai-subagent" && value == "review"));
assert!(headers
.iter()
.any(|(name, value)| name == "session-id" && value == "cp_derived"));
assert!(headers
.iter()
.any(|(name, value)| name == "thread-id" && value == "cp_derived"));
.any(|(name, value)| name == "session_id" && value == "cp_derived"));
assert!(headers.iter().all(|(name, _)| name != "content-type"));
assert!(headers.iter().all(|(name, _)| name != "accept"));
assert!(headers.iter().all(|(name, _)| name != "session_id"));
assert!(headers.iter().all(|(name, _)| name != "session-id"));
assert!(headers.iter().all(|(name, _)| name != "thread-id"));
}
#[tokio::test]
+72 -93
View File
@@ -92,73 +92,42 @@ fn websocket_responses_endpoint_should_convert_http_base_url_to_ws_endpoint() {
);
}
#[test]
fn websocket_deflate_should_rewrite_compressed_server_text_frame() {
let payload = json!({
"type": "response.completed",
"response": {
"id": "resp_deflate",
"object": "response"
}
})
.to_string();
let compressed = compressed_permessage_deflate_payload(payload.as_bytes());
let frame = server_websocket_frame(0x1, true, &compressed);
let rewritten =
codex_proxy_adapters::codex::websocket::deflate::rewrite_permessage_deflate_server_frame(
&frame,
)
.expect("compressed frame should rewrite")
.expect("compressed frame should produce rewritten frame");
assert_eq!(rewritten[0] & 0x40, 0);
assert_eq!(rewritten[0] & 0x0f, 0x1);
assert_eq!(server_websocket_payload(&rewritten), payload.as_bytes());
}
#[test]
fn websocket_deflate_should_leave_uncompressed_server_frame_unchanged() {
let frame = server_websocket_frame(0x1, false, b"plain");
let rewritten =
codex_proxy_adapters::codex::websocket::deflate::rewrite_permessage_deflate_server_frame(
&frame,
)
.expect("plain frame should parse");
assert_eq!(rewritten, None);
}
#[tokio::test]
async fn codex_backend_client_should_decode_live_permessage_deflate_websocket_frame() {
async fn codex_backend_client_should_decode_permessage_deflate_context_takeover_frames() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut stream).await;
let websocket_key = websocket_opening_header(&request, "sec-websocket-key")
.expect("opening request should include sec-websocket-key");
let accept_key = derive_accept_key(websocket_key.as_bytes());
stream
.write_all(
format!(
"HTTP/1.1 101 Switching Protocols\r\n\
Upgrade: websocket\r\n\
Connection: Upgrade\r\n\
Sec-WebSocket-Accept: {accept_key}\r\n\
Sec-WebSocket-Extensions: permessage-deflate\r\n\
\r\n"
)
.as_bytes(),
)
.await
.unwrap();
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_hdr_async_with_config(
stream,
|request: &WsRequest, response: WsResponse| {
let extensions = request
.headers()
.get("sec-websocket-extensions")
.and_then(|value| value.to_str().ok())
.unwrap_or_default();
assert!(extensions.contains("permessage-deflate"));
let mut response = response;
response.headers_mut().insert(
"sec-websocket-extensions",
"permessage-deflate".parse().unwrap(),
);
Ok(response)
},
Some(websocket_accept_config()),
)
.await
.unwrap();
let delta = json!({
"type": "response.output_text.delta",
"delta": "hello from websocket"
})
.to_string();
let completed = json!({
"type": "response.completed",
"response": {
"id": "resp_live_deflate",
"id": "resp_6f8d0c2b5a4e4a0d9c1b7e3f2a8d5c6b",
"object": "response",
"usage": {
"input_tokens": 3,
@@ -168,12 +137,14 @@ async fn codex_backend_client_should_decode_live_permessage_deflate_websocket_fr
}
})
.to_string();
let compressed = compressed_permessage_deflate_payload(completed.as_bytes());
stream
.write_all(&server_websocket_frame(0x1, true, &compressed))
.await
.unwrap();
stream.flush().await.unwrap();
let Some(Ok(Message::Text(_payload))) = websocket.next().await else {
panic!("client should send response.create payload");
};
for payload in [delta, completed] {
websocket.send(Message::Text(payload.into())).await.unwrap();
}
});
let backend = client::CodexBackendClient::new(
client::build_reqwest_client(false).unwrap(),
@@ -197,7 +168,10 @@ async fn codex_backend_client_should_decode_live_permessage_deflate_websocket_fr
.expect("deflated websocket response should decode");
server.await.unwrap();
assert!(response.body.contains("resp_live_deflate"));
assert!(response.body.contains("hello from websocket"));
assert!(response
.body
.contains("resp_6f8d0c2b5a4e4a0d9c1b7e3f2a8d5c6b"));
}
#[test]
@@ -316,7 +290,6 @@ fn websocket_connection_should_prepare_capture_payload_with_canonical_field_orde
"thread_id": "capture-thread-secret",
"safe": "capture",
}));
request.generate = Some(false);
let prepared =
codex_proxy_adapters::codex::websocket::connect::CodexWebSocketConnection::responses_create_request(
@@ -334,15 +307,11 @@ fn websocket_connection_should_prepare_capture_payload_with_canonical_field_orde
"\"model\":\"gpt-5.5\"",
"\"instructions\":\"private capture instructions\"",
"\"input\":",
"\"tools\":[]",
"\"tool_choice\":\"auto\"",
"\"parallel_tool_calls\":true",
"\"reasoning\":null",
"\"store\":false",
"\"stream\":true",
"\"include\":[]",
"\"tool_choice\":\"auto\"",
"\"parallel_tool_calls\":true",
"\"prompt_cache_key\":\"session-1\"",
"\"generate\":false",
"\"client_metadata\":",
],
);
@@ -354,7 +323,7 @@ async fn websocket_execute_response_create_request_should_collect_completed_sse(
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let message = websocket.next().await.unwrap().unwrap();
let payload = serde_json::from_str::<serde_json::Value>(&message.into_text().unwrap())
.expect("client payload should be json");
@@ -413,7 +382,7 @@ async fn websocket_execute_response_create_request_should_surface_response_faile
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
@@ -457,6 +426,7 @@ async fn websocket_execute_response_create_request_should_surface_response_faile
status_code,
retry_after_seconds,
body,
..
} = error
else {
panic!("expected upstream websocket error");
@@ -500,6 +470,7 @@ async fn websocket_execute_response_create_request_should_preserve_opening_error
status_code,
retry_after_seconds,
body,
..
} = error
else {
panic!("expected upstream opening error");
@@ -515,7 +486,7 @@ async fn websocket_execute_response_create_request_should_reject_binary_event()
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Binary(b"unexpected-binary".to_vec().into()))
@@ -544,7 +515,7 @@ async fn websocket_execute_response_create_request_should_surface_wrapped_error_
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
@@ -578,6 +549,7 @@ async fn websocket_execute_response_create_request_should_surface_wrapped_error_
status_code,
retry_after_seconds,
body,
..
} = error
else {
panic!("expected wrapped upstream error");
@@ -593,7 +565,7 @@ async fn websocket_execute_response_create_request_should_surface_connection_lim
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
@@ -640,7 +612,7 @@ async fn websocket_execute_response_create_request_should_skip_invalid_stream_ev
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
@@ -700,7 +672,7 @@ async fn websocket_execute_response_create_request_should_capture_internal_metad
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
@@ -781,7 +753,7 @@ async fn websocket_execute_response_create_request_should_surface_incomplete_res
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
@@ -836,7 +808,7 @@ async fn websocket_execute_response_create_request_should_reject_invalid_complet
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
@@ -895,7 +867,7 @@ async fn websocket_execute_response_create_request_should_ignore_completed_witho
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
@@ -942,7 +914,7 @@ async fn websocket_execute_response_create_request_should_ignore_success_status_
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
@@ -993,7 +965,7 @@ async fn websocket_execute_response_create_request_should_timeout_when_upstream_
let (request_seen_tx, request_seen_rx) = tokio::sync::oneshot::channel();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
request_seen_tx.send(()).unwrap();
futures::future::pending::<()>().await;
@@ -1030,7 +1002,7 @@ async fn websocket_execute_response_create_request_should_reply_to_server_ping_b
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Ping("codex-ping".as_bytes().to_vec().into()))
@@ -1094,7 +1066,7 @@ async fn codex_backend_client_stream_should_reject_binary_websocket_event() {
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Binary(b"unexpected-binary".to_vec().into()))
@@ -1144,7 +1116,7 @@ async fn codex_backend_client_stream_should_error_when_websocket_closes_before_t
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
@@ -1212,8 +1184,13 @@ async fn codex_backend_client_should_use_websocket_when_previous_response_id_is_
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket =
accept_hdr_async(stream, |_request: &WsRequest, mut response: WsResponse| {
let mut websocket = accept_hdr_async_with_config(
stream,
|_request: &WsRequest, mut response: WsResponse| {
response.headers_mut().insert(
"sec-websocket-extensions",
"permessage-deflate".parse().unwrap(),
);
response
.headers_mut()
.insert("x-codex-turn-state", "turn-ws-client".parse().unwrap());
@@ -1227,9 +1204,11 @@ async fn codex_backend_client_should_use_websocket_when_previous_response_id_is_
.headers_mut()
.insert("x-ratelimit-remaining-requests", "17".parse().unwrap());
Ok(response)
})
.await
.unwrap();
},
Some(websocket_accept_config()),
)
.await
.unwrap();
let message = websocket.next().await.unwrap().unwrap();
let payload = serde_json::from_str::<serde_json::Value>(&message.into_text().unwrap())
.expect("client payload should be json");
@@ -27,7 +27,7 @@ async fn codex_backend_client_should_reuse_pooled_websocket_for_same_account_and
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
for response_id in ["resp_pool_first", "resp_pool_second"] {
let _message = websocket.next().await.unwrap().unwrap();
websocket
@@ -101,7 +101,7 @@ async fn websocket_pool_should_bypass_busy_key_with_one_shot_connections() {
let (release_first_tx, release_first_rx) = tokio::sync::oneshot::channel();
let server = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
let mut first_websocket = accept_async(first_stream).await.unwrap();
let mut first_websocket = accept_codex_test_websocket(first_stream).await;
let _first_message = first_websocket.next().await.unwrap().unwrap();
first_websocket
.send(Message::Text(
@@ -116,7 +116,7 @@ async fn websocket_pool_should_bypass_busy_key_with_one_shot_connections() {
.unwrap();
let (second_stream, _) = listener.accept().await.unwrap();
let mut second_websocket = accept_async(second_stream).await.unwrap();
let mut second_websocket = accept_codex_test_websocket(second_stream).await;
let _second_message = second_websocket.next().await.unwrap().unwrap();
second_websocket
.send(Message::Text(
@@ -127,7 +127,7 @@ async fn websocket_pool_should_bypass_busy_key_with_one_shot_connections() {
second_websocket.close(None).await.unwrap();
let (third_stream, _) = listener.accept().await.unwrap();
let mut third_websocket = accept_async(third_stream).await.unwrap();
let mut third_websocket = accept_codex_test_websocket(third_stream).await;
let _third_message = third_websocket.next().await.unwrap().unwrap();
third_websocket
.send(Message::Text(
@@ -204,7 +204,7 @@ async fn websocket_pool_should_bypass_new_keys_after_account_cap() {
for response_id in ["resp_cap_first", "resp_cap_second", "resp_cap_third"] {
let (stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
@@ -264,7 +264,7 @@ async fn codex_backend_client_should_ping_idle_pooled_websocket_during_maintenan
let ping_count_for_server = Arc::clone(&ping_count);
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _first_message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
@@ -368,7 +368,7 @@ async fn websocket_pool_should_evict_idle_connection_when_ping_times_out() {
let server = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let mut first_websocket = accept_async(first_stream).await.unwrap();
let mut first_websocket = accept_codex_test_websocket(first_stream).await;
let _first_message = first_websocket.next().await.unwrap().unwrap();
first_websocket
.send(Message::Text(
@@ -378,7 +378,7 @@ async fn websocket_pool_should_evict_idle_connection_when_ping_times_out() {
.unwrap();
let (second_stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let mut second_websocket = accept_async(second_stream).await.unwrap();
let mut second_websocket = accept_codex_test_websocket(second_stream).await;
let _second_message = second_websocket.next().await.unwrap().unwrap();
second_websocket
.send(Message::Text(
@@ -434,7 +434,7 @@ async fn websocket_pool_should_gc_expired_idle_connections() {
let server = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let mut first_websocket = accept_async(first_stream).await.unwrap();
let mut first_websocket = accept_codex_test_websocket(first_stream).await;
let _first_message = first_websocket.next().await.unwrap().unwrap();
first_websocket
.send(Message::Text(
@@ -451,7 +451,7 @@ async fn websocket_pool_should_gc_expired_idle_connections() {
let (second_stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let mut second_websocket = accept_async(second_stream).await.unwrap();
let mut second_websocket = accept_codex_test_websocket(second_stream).await;
let _second_message = second_websocket.next().await.unwrap().unwrap();
second_websocket
.send(Message::Text(
@@ -507,7 +507,7 @@ async fn codex_backend_client_should_ping_idle_pooled_websocket_from_background_
let ping_count_for_server = Arc::clone(&ping_count);
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_async(stream).await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _first_message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
@@ -616,7 +616,7 @@ async fn codex_backend_client_should_close_idle_pooled_websocket_when_account_is
let server = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let mut first_websocket = accept_async(first_stream).await.unwrap();
let mut first_websocket = accept_codex_test_websocket(first_stream).await;
let _first_message = first_websocket.next().await.unwrap().unwrap();
first_websocket
.send(Message::Text(
@@ -646,7 +646,7 @@ async fn codex_backend_client_should_close_idle_pooled_websocket_when_account_is
let (second_stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let mut second_websocket = accept_async(second_stream).await.unwrap();
let mut second_websocket = accept_codex_test_websocket(second_stream).await;
let _second_message = second_websocket.next().await.unwrap().unwrap();
second_websocket
.send(Message::Text(
@@ -718,7 +718,7 @@ async fn codex_backend_client_should_stop_reusing_pooled_websockets_after_shutdo
let server = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let mut first_websocket = accept_async(first_stream).await.unwrap();
let mut first_websocket = accept_codex_test_websocket(first_stream).await;
let _first_message = first_websocket.next().await.unwrap().unwrap();
first_websocket
.send(Message::Text(
@@ -748,7 +748,7 @@ async fn codex_backend_client_should_stop_reusing_pooled_websockets_after_shutdo
let (second_stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let mut second_websocket = accept_async(second_stream).await.unwrap();
let mut second_websocket = accept_codex_test_websocket(second_stream).await;
let _second_message = second_websocket.next().await.unwrap().unwrap();
second_websocket
.send(Message::Text(
@@ -820,7 +820,7 @@ async fn codex_backend_client_should_close_idle_pooled_websocket_after_liveness_
let server = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let mut first_websocket = accept_async(first_stream).await.unwrap();
let mut first_websocket = accept_codex_test_websocket(first_stream).await;
let _first_message = first_websocket.next().await.unwrap().unwrap();
first_websocket
.send(Message::Text(
@@ -850,7 +850,7 @@ async fn codex_backend_client_should_close_idle_pooled_websocket_after_liveness_
let (second_stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let mut second_websocket = accept_async(second_stream).await.unwrap();
let mut second_websocket = accept_codex_test_websocket(second_stream).await;
let _second_message = second_websocket.next().await.unwrap().unwrap();
second_websocket
.send(Message::Text(
@@ -923,7 +923,7 @@ async fn codex_backend_client_should_discard_pooled_websocket_after_upstream_err
let server = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let mut first_websocket = accept_async(first_stream).await.unwrap();
let mut first_websocket = accept_codex_test_websocket(first_stream).await;
let _first_message = first_websocket.next().await.unwrap().unwrap();
first_websocket
.send(Message::Text(
@@ -946,7 +946,7 @@ async fn codex_backend_client_should_discard_pooled_websocket_after_upstream_err
let (second_stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let mut second_websocket = accept_async(second_stream).await.unwrap();
let mut second_websocket = accept_codex_test_websocket(second_stream).await;
let _second_message = second_websocket.next().await.unwrap().unwrap();
second_websocket
.send(Message::Text(
+74
View File
@@ -7,6 +7,8 @@ use tokio::sync::RwLock;
const PATH_BLOCK_THRESHOLD: u32 = 3;
const PATH_BLOCK_STALE_AFTER: Duration = Duration::hours(1);
const CHALLENGE_BACKOFF_SECONDS: [i64; 4] = [10, 30, 90, 120];
const CHALLENGE_STALE_AFTER: Duration = Duration::hours(1);
#[derive(Debug, Clone, Copy)]
struct PathBlockState {
@@ -14,6 +16,25 @@ struct PathBlockState {
last_at: DateTime<Utc>,
}
#[derive(Debug, Clone, Copy)]
struct ChallengeCooldownState {
challenge_count: u32,
updated_at: DateTime<Utc>,
}
/// Cloudflare challenge cooldown state after recording one challenge.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CloudflareChallengeCooldown {
/// Current non-stale consecutive challenge count.
pub challenge_count: u32,
/// Backoff delay selected for this challenge.
pub delay_seconds: i64,
/// Time until which this account should be skipped.
pub cooldown_until: DateTime<Utc>,
/// Time at which this challenge was recorded.
pub updated_at: DateTime<Utc>,
}
/// Tracks per-account Cloudflare path-block failures.
#[derive(Debug, Clone, Default)]
pub struct CloudflarePathBlockTracker {
@@ -65,3 +86,56 @@ impl CloudflarePathBlockTracker {
self.count(account_id, now).await >= PATH_BLOCK_THRESHOLD
}
}
/// Tracks per-account Cloudflare challenge cooldown escalation.
#[derive(Debug, Clone, Default)]
pub struct CloudflareChallengeCooldownTracker {
states: Arc<RwLock<HashMap<String, ChallengeCooldownState>>>,
}
impl CloudflareChallengeCooldownTracker {
/// Creates an empty challenge cooldown tracker.
pub fn new() -> Self {
Self::default()
}
/// Records one challenge and returns the current non-stale cooldown state.
pub async fn record_challenge(
&self,
account_id: &str,
now: DateTime<Utc>,
) -> CloudflareChallengeCooldown {
let mut states = self.states.write().await;
let challenge_count = states
.get(account_id)
.filter(|state| now.signed_duration_since(state.updated_at) <= CHALLENGE_STALE_AFTER)
.map(|state| state.challenge_count.saturating_add(1))
.unwrap_or(1);
states.insert(
account_id.to_string(),
ChallengeCooldownState {
challenge_count,
updated_at: now,
},
);
let delay_seconds = challenge_delay_seconds(challenge_count);
CloudflareChallengeCooldown {
challenge_count,
delay_seconds,
cooldown_until: now + Duration::seconds(delay_seconds),
updated_at: now,
}
}
/// Clears tracked challenge cooldown state for an account.
pub async fn reset(&self, account_id: &str) {
self.states.write().await.remove(account_id);
}
}
fn challenge_delay_seconds(challenge_count: u32) -> i64 {
let index = challenge_count
.saturating_sub(1)
.min((CHALLENGE_BACKOFF_SECONDS.len() - 1) as u32) as usize;
CHALLENGE_BACKOFF_SECONDS[index]
}
+3
View File
@@ -41,6 +41,8 @@ pub struct Account {
pub refresh_token: Option<String>,
/// 访问令牌过期时间。
pub access_token_expires_at: Option<DateTime<Utc>>,
/// 下一次允许刷新 token 的时间。
pub next_refresh_at: Option<DateTime<Utc>>,
/// 当前账号状态。
pub status: AccountStatus,
/// 是否已经触发配额封顶。
@@ -104,6 +106,7 @@ impl Account {
access_token: format!("token-{id}"),
refresh_token: Some(format!("refresh-{id}")),
access_token_expires_at: Some(Utc::now() + Duration::hours(1)),
next_refresh_at: None,
status,
quota_limit_reached: false,
quota_verify_required: false,
+30
View File
@@ -189,6 +189,11 @@ impl AccountPool {
self.accounts.insert(account.id.clone(), account);
}
/// 获取账号池中的账号快照。
pub fn get(&self, account_id: &str) -> Option<Account> {
self.accounts.get(account_id).cloned()
}
/// 移除账号及其在途槽位。
pub fn remove(&mut self, account_id: &str) -> bool {
let removed = self.accounts.remove(account_id).is_some();
@@ -292,6 +297,31 @@ impl AccountPool {
true
}
/// 应用已经验证过的配额状态。
pub fn apply_quota_state(
&mut self,
account_id: &str,
limit_reached: bool,
cooldown_until: Option<DateTime<Utc>>,
) -> bool {
let Some(account) = self.accounts.get_mut(account_id) else {
return false;
};
account.quota_verify_required = false;
account.quota_limit_reached = limit_reached;
account.quota_cooldown_until = limit_reached.then_some(cooldown_until).flatten();
if let Some(cooldown_until) = account.quota_cooldown_until {
account.window_reset_at = account
.window_reset_at
.filter(|existing| *existing > cooldown_until)
.or(Some(cooldown_until));
}
if limit_reached {
self.slots.remove(account_id);
}
true
}
/// 同步上游返回的限流窗口信息。
pub fn sync_rate_limit_window(
&mut self,
+12
View File
@@ -76,6 +76,18 @@ pub trait AccountStore: Send + Sync + 'static {
Ok(false)
}
/// 应用已经验证过的账号配额快照。
async fn apply_quota_snapshot(
&self,
account_id: &str,
quota_json: &str,
limit_reached: bool,
cooldown_until: Option<DateTime<Utc>>,
) -> AccountStoreResult<bool> {
let _ = (limit_reached, cooldown_until);
self.update_quota_json(account_id, quota_json).await
}
/// 同步账号当前 rate-limit 统计窗口。
async fn sync_rate_limit_window(
&self,
+13 -5
View File
@@ -225,8 +225,11 @@ pub enum RefreshFailure {
/// 账号被显式禁用。
#[error("account is disabled")]
Disabled,
/// 刷新请求在传输层失败。
#[error("refresh transport failed")]
/// 刷新请求在到达服务端前失败,可安全复用当前 refresh token 重试。
#[error("refresh transport failed before server processing")]
RetryableTransport,
/// 刷新请求在传输层失败,refresh token 可能已经被服务端消费。
#[error("refresh transport failed after possible server processing")]
Transport,
}
@@ -236,8 +239,11 @@ pub enum RefreshError {
/// 并发限制信号量已关闭。
#[error("refresh task semaphore closed")]
ConcurrencyClosed,
/// 刷新请求在传输层失败。
#[error("refresh transport failed")]
/// 刷新请求在到达服务端前失败,可安全复用当前 refresh token 重试。
#[error("refresh transport failed before server processing")]
RetryableTransport,
/// 刷新请求在传输层失败,refresh token 可能已经被服务端消费。
#[error("refresh transport failed after possible server processing")]
Transport,
}
@@ -295,6 +301,7 @@ where
match self.client.refresh(refresh_token).await {
Ok(token_pair) => Ok(apply_token_pair(account, token_pair)),
Err(RefreshFailure::RetryableTransport) => Err(RefreshError::RetryableTransport),
Err(RefreshFailure::Transport) => Err(RefreshError::Transport),
Err(error) => Ok(apply_refresh_failure(account, error)),
}
@@ -352,10 +359,11 @@ pub fn apply_token_pair(account: &Account, token_pair: TokenPair) -> Account {
pub fn apply_refresh_failure(account: &Account, failure: RefreshFailure) -> Account {
let mut updated = account.clone();
updated.status = match failure {
RefreshFailure::InvalidGrant => AccountStatus::Expired,
RefreshFailure::InvalidGrant => AccountStatus::Disabled,
RefreshFailure::QuotaExhausted => AccountStatus::QuotaExhausted,
RefreshFailure::Banned => AccountStatus::Banned,
RefreshFailure::Disabled => AccountStatus::Disabled,
RefreshFailure::RetryableTransport => AccountStatus::Active,
RefreshFailure::Transport => AccountStatus::Active,
};
updated
+74 -77
View File
@@ -1192,13 +1192,14 @@ pub fn websocket_response_create_payload(request: &CodexResponsesRequest) -> Val
Value::String("response.create".to_string()),
);
insert_value(&mut body, "model", Value::String(request.model.clone()));
if !request.instructions.is_empty() {
insert_value(
&mut body,
"instructions",
Value::String(request.instructions.clone()),
);
}
insert_value(
&mut body,
"instructions",
Value::String(request.instructions.clone()),
);
insert_value(&mut body, "input", Value::Array(request.input.clone()));
insert_value(&mut body, "store", Value::Bool(request.store));
insert_value(&mut body, "stream", Value::Bool(request.stream));
if let Some(previous_response_id) = &request.previous_response_id {
insert_value(
&mut body,
@@ -1206,12 +1207,12 @@ pub fn websocket_response_create_payload(request: &CodexResponsesRequest) -> Val
Value::String(previous_response_id.clone()),
);
}
insert_value(&mut body, "input", Value::Array(request.input.clone()));
insert_value(
&mut body,
"tools",
Value::Array(request.tools.clone().unwrap_or_default()),
);
if let Some(reasoning) = &request.reasoning {
insert_value(&mut body, "reasoning", reasoning.clone());
}
if let Some(tools) = non_empty_tools(request) {
insert_value(&mut body, "tools", Value::Array(tools.to_vec()));
}
insert_value(
&mut body,
"tool_choice",
@@ -1225,25 +1226,9 @@ pub fn websocket_response_create_payload(request: &CodexResponsesRequest) -> Val
"parallel_tool_calls",
Value::Bool(request.parallel_tool_calls.unwrap_or(true)),
);
insert_value(
&mut body,
"reasoning",
request.reasoning.clone().unwrap_or(Value::Null),
);
insert_value(&mut body, "store", Value::Bool(request.store));
insert_value(&mut body, "stream", Value::Bool(request.stream));
insert_value(
&mut body,
"include",
request
.include
.clone()
.unwrap_or_default()
.into_iter()
.map(Value::String)
.collect::<Vec<_>>()
.into(),
);
if let Some(text) = &request.text {
insert_value(&mut body, "text", text.clone());
}
if let Some(service_tier) = &request.service_tier {
insert_value(
&mut body,
@@ -1258,11 +1243,17 @@ pub fn websocket_response_create_payload(request: &CodexResponsesRequest) -> Val
Value::String(prompt_cache_key.clone()),
);
}
if let Some(text) = &request.text {
insert_value(&mut body, "text", text.clone());
}
if let Some(generate) = request.generate {
insert_value(&mut body, "generate", Value::Bool(generate));
if let Some(include) = non_empty_include(request) {
insert_value(
&mut body,
"include",
include
.iter()
.cloned()
.map(Value::String)
.collect::<Vec<_>>()
.into(),
);
}
if let Some(client_metadata) = &request.client_metadata {
insert_value(&mut body, "client_metadata", client_metadata.clone());
@@ -1288,14 +1279,19 @@ impl Serialize for OrderedResponseCreatePayload<'_> {
let mut map = serializer.serialize_map(Some(websocket_payload_keys(request).len()))?;
map.serialize_entry("type", "response.create")?;
map.serialize_entry("model", &request.model)?;
if !request.instructions.is_empty() {
map.serialize_entry("instructions", &request.instructions)?;
}
map.serialize_entry("instructions", &request.instructions)?;
map.serialize_entry("input", &request.input)?;
map.serialize_entry("store", &request.store)?;
map.serialize_entry("stream", &request.stream)?;
if let Some(previous_response_id) = &request.previous_response_id {
map.serialize_entry("previous_response_id", previous_response_id)?;
}
map.serialize_entry("input", &request.input)?;
map.serialize_entry("tools", request.tools.as_deref().unwrap_or(&[]))?;
if let Some(reasoning) = &request.reasoning {
map.serialize_entry("reasoning", reasoning)?;
}
if let Some(tools) = non_empty_tools(request) {
map.serialize_entry("tools", tools)?;
}
if let Some(tool_choice) = &request.tool_choice {
map.serialize_entry("tool_choice", tool_choice)?;
} else {
@@ -1305,25 +1301,17 @@ impl Serialize for OrderedResponseCreatePayload<'_> {
"parallel_tool_calls",
&request.parallel_tool_calls.unwrap_or(true),
)?;
if let Some(reasoning) = &request.reasoning {
map.serialize_entry("reasoning", reasoning)?;
} else {
map.serialize_entry("reasoning", &Value::Null)?;
if let Some(text) = &request.text {
map.serialize_entry("text", text)?;
}
map.serialize_entry("store", &request.store)?;
map.serialize_entry("stream", &request.stream)?;
map.serialize_entry("include", request.include.as_deref().unwrap_or(&[]))?;
if let Some(service_tier) = &request.service_tier {
map.serialize_entry("service_tier", service_tier)?;
}
if let Some(prompt_cache_key) = &request.prompt_cache_key {
map.serialize_entry("prompt_cache_key", prompt_cache_key)?;
}
if let Some(text) = &request.text {
map.serialize_entry("text", text)?;
}
if let Some(generate) = request.generate {
map.serialize_entry("generate", &generate)?;
if let Some(include) = non_empty_include(request) {
map.serialize_entry("include", include)?;
}
if let Some(client_metadata) = &request.client_metadata {
map.serialize_entry("client_metadata", client_metadata)?;
@@ -1337,38 +1325,36 @@ fn insert_value(body: &mut Map<String, Value>, key: &str, value: Value) {
}
fn websocket_payload_keys(request: &CodexResponsesRequest) -> Vec<String> {
let mut keys = vec!["type".to_string(), "model".to_string()];
if !request.instructions.is_empty() {
keys.push("instructions".to_string());
}
let mut keys = vec![
"type".to_string(),
"model".to_string(),
"instructions".to_string(),
"input".to_string(),
"store".to_string(),
"stream".to_string(),
];
if request.previous_response_id.is_some() {
keys.push("previous_response_id".to_string());
}
keys.extend(
[
"input",
"tools",
"tool_choice",
"parallel_tool_calls",
"reasoning",
"store",
"stream",
"include",
]
.into_iter()
.map(ToString::to_string),
);
if request.reasoning.is_some() {
keys.push("reasoning".to_string());
}
if non_empty_tools(request).is_some() {
keys.push("tools".to_string());
}
keys.push("tool_choice".to_string());
keys.push("parallel_tool_calls".to_string());
if request.text.is_some() {
keys.push("text".to_string());
}
if request.service_tier.is_some() {
keys.push("service_tier".to_string());
}
if request.prompt_cache_key.is_some() {
keys.push("prompt_cache_key".to_string());
}
if request.text.is_some() {
keys.push("text".to_string());
}
if request.generate.is_some() {
keys.push("generate".to_string());
if non_empty_include(request).is_some() {
keys.push("include".to_string());
}
if request.client_metadata.is_some() {
keys.push("client_metadata".to_string());
@@ -1376,6 +1362,17 @@ fn websocket_payload_keys(request: &CodexResponsesRequest) -> Vec<String> {
keys
}
fn non_empty_tools(request: &CodexResponsesRequest) -> Option<&[Value]> {
request.tools.as_deref().filter(|tools| !tools.is_empty())
}
fn non_empty_include(request: &CodexResponsesRequest) -> Option<&[String]> {
request
.include
.as_deref()
.filter(|include| !include.is_empty())
}
fn redact_payload_body(body: Value) -> Value {
let Value::Object(body) = body else {
return body;
+440
View File
@@ -80,6 +80,378 @@ pub enum ChatTranslationError {
EmptyMessages,
}
/// Chat Completions 流式转换失败。
#[derive(Debug, Error, PartialEq, Eq)]
pub enum ChatStreamTranslationError {
/// 上游 SSE 无法解析。
#[error("invalid upstream SSE response: {0}")]
InvalidSse(#[from] SseError),
/// 上游返回错误事件。
#[error("{0}")]
Upstream(String),
}
/// 增量转换 Codex Responses SSE 为 OpenAI Chat Completions chunk SSE。
#[derive(Debug)]
pub struct ChatCompletionStreamTranslator {
pending: String,
id: String,
created: i64,
model: String,
include_reasoning: bool,
tuple_schema: Option<Value>,
tuple_text_buffer: Option<String>,
has_tool_calls: bool,
has_content: bool,
function_call_items: BTreeMap<String, FunctionCallInfo>,
tool_call_indices: BTreeMap<String, usize>,
call_ids_with_deltas: BTreeSet<String>,
next_tool_call_index: usize,
closed: bool,
}
impl ChatCompletionStreamTranslator {
/// 构造新的 chat stream 转换器。
pub fn new(
model: impl Into<String>,
include_reasoning: bool,
tuple_schema: Option<Value>,
) -> Self {
let tuple_text_buffer = tuple_schema.as_ref().map(|_| String::new());
Self {
pending: String::new(),
id: chat_completion_stream_id(),
created: Utc::now().timestamp(),
model: model.into(),
include_reasoning,
tuple_schema,
tuple_text_buffer,
has_tool_calls: false,
has_content: false,
function_call_items: BTreeMap::new(),
tool_call_indices: BTreeMap::new(),
call_ids_with_deltas: BTreeSet::new(),
next_tool_call_index: 0,
closed: false,
}
}
/// TS 版本会先发 assistant role chunk。
pub fn initial_frame(&self) -> String {
self.frame(json!({
"choices": [{
"delta": {"role": "assistant"},
"finish_reason": Value::Null,
"index": 0,
}],
}))
}
/// 推入一段上游 SSE 文本,返回已完成事件对应的下游 SSE 文本。
pub fn push_str(&mut self, chunk: &str) -> Result<String, ChatStreamTranslationError> {
if self.closed {
return Ok(String::new());
}
self.pending.push_str(chunk);
let mut output = String::new();
while let Some((event_end, separator_len)) = sse_event_separator(&self.pending) {
let event_frame = self.pending[..event_end].to_string();
self.pending.drain(..event_end + separator_len);
if event_frame.trim().is_empty() {
continue;
}
let events = parse_sse_events(&format!("{event_frame}\n\n"))?;
for event in events {
self.push_event(&event, &mut output)?;
}
}
Ok(output)
}
fn push_event(
&mut self,
event: &crate::protocol::codex::sse::SseEvent,
output: &mut String,
) -> Result<(), ChatStreamTranslationError> {
let Ok(value) = serde_json::from_str::<Value>(&event.data) else {
return Ok(());
};
let event_type = event
.event
.as_deref()
.or_else(|| value.get("type").and_then(Value::as_str));
if matches!(event_type, Some("error" | "response.failed")) {
return Err(ChatStreamTranslationError::Upstream(
codex_stream_error_message(&value),
));
}
match event_type {
Some("response.output_text.delta") => {
self.push_text_delta(&value, output);
}
Some("response.reasoning_summary_text.delta") if self.include_reasoning => {
if let Some(delta) = value.get("delta").and_then(Value::as_str) {
self.has_content = true;
output.push_str(&self.frame(json!({
"choices": [{
"delta": {"reasoning_content": delta},
"finish_reason": Value::Null,
"index": 0,
}],
})));
}
}
Some("response.output_item.added") => {
self.push_function_call_start(&value, output);
}
Some("response.function_call_arguments.delta") => {
self.push_function_call_delta(&value, output);
}
Some("response.function_call_arguments.done") => {
self.push_function_call_done(&value, output);
}
Some("response.output_item.done") => {
self.push_image_generation_done(&value, output);
}
Some("response.completed") => {
self.push_completed(&value, output);
}
_ => {}
}
Ok(())
}
fn push_text_delta(&mut self, value: &Value, output: &mut String) {
let Some(delta) = value.get("delta").and_then(Value::as_str) else {
return;
};
self.has_content = true;
if let Some(buffer) = &mut self.tuple_text_buffer {
buffer.push_str(delta);
return;
}
output.push_str(&self.frame(json!({
"choices": [{
"delta": {"content": delta},
"finish_reason": Value::Null,
"index": 0,
}],
})));
}
fn push_function_call_start(&mut self, value: &Value, output: &mut String) {
let Some((item_id, call_id, name)) = function_call_item(value.get("item")) else {
return;
};
self.has_tool_calls = true;
self.has_content = true;
let index = self.next_tool_call_index;
self.next_tool_call_index += 1;
self.function_call_items.insert(
item_id,
FunctionCallInfo {
call_id: call_id.clone(),
name: name.clone(),
},
);
self.tool_call_indices.insert(call_id.clone(), index);
output.push_str(&self.frame(json!({
"choices": [{
"delta": {
"tool_calls": [{
"function": {"arguments": "", "name": name},
"id": call_id,
"index": index,
"type": "function",
}],
},
"finish_reason": Value::Null,
"index": 0,
}],
})));
}
fn push_function_call_delta(&mut self, value: &Value, output: &mut String) {
let Some((call_id, delta)) = self.function_call_delta(value) else {
return;
};
self.call_ids_with_deltas.insert(call_id.clone());
let index = self.tool_call_indices.get(&call_id).copied().unwrap_or(0);
output.push_str(&self.frame(json!({
"choices": [{
"delta": {
"tool_calls": [{
"function": {"arguments": delta},
"index": index,
}],
},
"finish_reason": Value::Null,
"index": 0,
}],
})));
}
fn push_function_call_done(&mut self, value: &Value, output: &mut String) {
let Some((call_id, arguments)) = self.function_call_done(value) else {
return;
};
if self.call_ids_with_deltas.contains(&call_id) {
return;
}
let index = self.tool_call_indices.get(&call_id).copied().unwrap_or(0);
output.push_str(&self.frame(json!({
"choices": [{
"delta": {
"tool_calls": [{
"function": {"arguments": arguments},
"index": index,
}],
},
"finish_reason": Value::Null,
"index": 0,
}],
})));
}
fn push_image_generation_done(&mut self, value: &Value, output: &mut String) {
let Some(item) = value.get("item") else {
return;
};
if item.get("type").and_then(Value::as_str) != Some("image_generation_call") {
return;
}
self.has_tool_calls = true;
self.has_content = true;
let index = self.next_tool_call_index;
self.next_tool_call_index += 1;
let id = item.get("id").and_then(Value::as_str).unwrap_or_default();
let mut arguments = json!({
"result": item.get("result").and_then(Value::as_str).unwrap_or_default(),
});
if let Some(revised_prompt) = item.get("revised_prompt") {
arguments["revised_prompt"] = revised_prompt.clone();
}
output.push_str(&self.frame(json!({
"choices": [{
"delta": {
"tool_calls": [{
"function": {"arguments": "", "name": "image_generation"},
"id": id,
"index": index,
"type": "function",
}],
},
"finish_reason": Value::Null,
"index": 0,
}],
})));
output.push_str(&self.frame(json!({
"choices": [{
"delta": {
"tool_calls": [{
"function": {"arguments": arguments.to_string()},
"index": index,
}],
},
"finish_reason": Value::Null,
"index": 0,
}],
})));
}
fn push_completed(&mut self, value: &Value, output: &mut String) {
self.flush_tuple_text(output);
if !self.has_content {
output.push_str(&self.frame(json!({
"choices": [{
"delta": {"content": "[Error] Codex returned an empty response. Please retry."},
"finish_reason": Value::Null,
"index": 0,
}],
})));
}
let response = value.get("response");
output.push_str(&self.frame(json!({
"choices": [{
"delta": {},
"finish_reason": if self.has_tool_calls { "tool_calls" } else { "stop" },
"index": 0,
}],
"usage": openai_stream_usage(response.and_then(|response| response.get("usage"))),
})));
output.push_str("data: [DONE]\n\n");
self.closed = true;
}
fn flush_tuple_text(&mut self, output: &mut String) {
let Some(buffer) = self.tuple_text_buffer.take() else {
return;
};
if buffer.is_empty() {
return;
}
let content = self
.tuple_schema
.as_ref()
.and_then(|tuple_schema| reconvert_tuple_text(&buffer, tuple_schema))
.unwrap_or(buffer);
output.push_str(&self.frame(json!({
"choices": [{
"delta": {"content": content},
"finish_reason": Value::Null,
"index": 0,
}],
})));
}
fn function_call_delta(&self, value: &Value) -> Option<(String, String)> {
let event_id = value
.get("call_id")
.or_else(|| value.get("item_id"))
.and_then(Value::as_str)?;
let delta = value.get("delta").and_then(Value::as_str)?;
let call_id = self
.function_call_items
.get(event_id)
.map(|info| info.call_id.as_str())
.unwrap_or(event_id)
.to_string();
Some((call_id, delta.to_string()))
}
fn function_call_done(&self, value: &Value) -> Option<(String, String)> {
let event_id = value
.get("call_id")
.or_else(|| value.get("item_id"))
.and_then(Value::as_str)?;
let call_id = self
.function_call_items
.get(event_id)
.map(|info| info.call_id.as_str())
.unwrap_or(event_id)
.to_string();
let arguments = value
.get("arguments")
.and_then(Value::as_str)
.unwrap_or_default()
.to_string();
Some((call_id, arguments))
}
fn frame(&self, mut chunk: Value) -> String {
chunk["id"] = Value::String(self.id.clone());
chunk["object"] = Value::String("chat.completion.chunk".to_string());
chunk["created"] = json!(self.created);
chunk["model"] = Value::String(self.model.clone());
format!("data: {chunk}\n\n")
}
}
/// 将 OpenAI Chat Completions 请求纯转换为 Codex Responses 请求。
pub fn translate_chat_to_codex(
request: ChatCompletionRequest,
@@ -618,3 +990,71 @@ fn openai_usage(usage: Option<&Value>) -> Value {
}
openai_usage
}
fn openai_stream_usage(usage: Option<&Value>) -> Value {
let Some(usage) = usage else {
return Value::Null;
};
let prompt_tokens = usage
.get("input_tokens")
.and_then(Value::as_u64)
.unwrap_or_default();
let completion_tokens = usage
.get("output_tokens")
.and_then(Value::as_u64)
.unwrap_or_default();
let mut openai_usage = json!({
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens.saturating_add(completion_tokens),
});
if let Some(cached_tokens) = usage
.get("input_tokens_details")
.and_then(|details| details.get("cached_tokens"))
.and_then(Value::as_u64)
{
openai_usage["prompt_tokens_details"] = json!({
"cached_tokens": cached_tokens,
});
}
if let Some(reasoning_tokens) = usage
.get("output_tokens_details")
.and_then(|details| details.get("reasoning_tokens"))
.and_then(Value::as_u64)
{
openai_usage["completion_tokens_details"] = json!({
"reasoning_tokens": reasoning_tokens,
});
}
openai_usage
}
fn sse_event_separator(input: &str) -> Option<(usize, usize)> {
let lf = input.find("\n\n").map(|index| (index, 2));
let crlf = input.find("\r\n\r\n").map(|index| (index, 4));
match (lf, crlf) {
(Some(left), Some(right)) => Some(if left.0 <= right.0 { left } else { right }),
(Some(found), None) | (None, Some(found)) => Some(found),
(None, None) => None,
}
}
fn chat_completion_stream_id() -> String {
let suffix = Uuid::new_v4()
.simple()
.to_string()
.chars()
.take(24)
.collect::<String>();
format!("chatcmpl-{suffix}")
}
fn codex_stream_error_message(value: &Value) -> String {
value
.pointer("/error/message")
.or_else(|| value.pointer("/response/error/message"))
.and_then(Value::as_str)
.filter(|message| !message.trim().is_empty())
.unwrap_or("Upstream Codex response failed")
.to_string()
}
+63 -6
View File
@@ -110,6 +110,7 @@ fn default_responses_stream() -> bool {
/// 将 OpenAI Responses 请求转换为 Codex Responses 请求。
pub fn translate_response_to_codex(request: OpenAiResponsesRequest) -> CodexResponsesRequest {
let prepared_text = prepare_text_format(request.text, true);
let client_metadata = sanitize_client_metadata(request.client_metadata);
let mut codex_request = CodexResponsesRequest::new_http_sse(
request.model,
request.instructions.unwrap_or_default(),
@@ -117,12 +118,32 @@ pub fn translate_response_to_codex(request: OpenAiResponsesRequest) -> CodexResp
);
codex_request.previous_response_id = request.previous_response_id;
codex_request.turn_state = request.turn_state;
codex_request.turn_metadata = non_empty_string(request.turn_metadata);
codex_request.beta_features = non_empty_string(request.beta_features);
codex_request.include_timing_metrics = non_empty_string(request.include_timing_metrics);
codex_request.turn_metadata = first_request_string(
request.turn_metadata,
client_metadata.as_ref(),
"x-codex-turn-metadata",
);
codex_request.beta_features = first_request_string(
request.beta_features,
client_metadata.as_ref(),
"x-codex-beta-features",
);
codex_request.include_timing_metrics = first_request_string(
request.include_timing_metrics,
client_metadata.as_ref(),
"x-responsesapi-include-timing-metrics",
);
codex_request.version = non_empty_string(request.version);
codex_request.codex_window_id = non_empty_string(request.codex_window_id);
codex_request.parent_thread_id = non_empty_string(request.parent_thread_id);
codex_request.codex_window_id = first_request_string(
request.codex_window_id,
client_metadata.as_ref(),
"x-codex-window-id",
);
codex_request.parent_thread_id = first_request_string(
request.parent_thread_id,
client_metadata.as_ref(),
"x-codex-parent-thread-id",
);
codex_request.reasoning = responses_reasoning(request.reasoning);
codex_request.tools = non_empty_array(request.tools);
codex_request.tool_choice = request.tool_choice;
@@ -134,7 +155,7 @@ pub fn translate_response_to_codex(request: OpenAiResponsesRequest) -> CodexResp
codex_request.prompt_cache_key = non_empty_string(request.prompt_cache_key);
codex_request.explicit_prompt_cache_key = codex_request.prompt_cache_key.is_some();
codex_request.include = string_array(request.include);
codex_request.client_metadata = sanitize_client_metadata(request.client_metadata);
codex_request.client_metadata = client_metadata;
ensure_reasoning_include(&mut codex_request);
match request.use_websocket {
Some(true) => codex_request.use_websocket = true,
@@ -214,6 +235,7 @@ fn sanitize_responses_input(input: Value) -> Vec<Value> {
match input {
Value::Array(items) => sanitize_codex_input_items(items),
Value::Null => Vec::new(),
Value::String(text) => vec![responses_input_text_message(text)],
value => vec![value],
}
}
@@ -222,6 +244,9 @@ fn sanitize_codex_input_items(input: Vec<Value>) -> Vec<Value> {
input
.into_iter()
.filter_map(|item| {
if let Value::String(text) = item {
return Some(responses_input_text_message(text));
}
let Value::Object(object) = item else {
return Some(item);
};
@@ -234,6 +259,19 @@ fn sanitize_codex_input_items(input: Vec<Value>) -> Vec<Value> {
.collect()
}
fn responses_input_text_message(text: String) -> Value {
json!({
"type": "message",
"role": "user",
"content": [
{
"type": "input_text",
"text": text
}
]
})
}
fn sanitize_reasoning_item(item: &Map<String, Value>) -> Option<Value> {
let id = non_empty_str(item.get("id"))?;
let summary = sanitize_summary(item.get("summary"))?;
@@ -375,6 +413,25 @@ fn sanitize_client_metadata(client_metadata: Option<Value>) -> Option<Value> {
(!metadata.is_empty()).then_some(Value::Object(metadata))
}
fn first_request_string(
direct: Option<String>,
client_metadata: Option<&Value>,
metadata_key: &str,
) -> Option<String> {
non_empty_string(direct).or_else(|| metadata_string(client_metadata, metadata_key))
}
fn metadata_string(client_metadata: Option<&Value>, key: &str) -> Option<String> {
client_metadata?
.as_object()?
.get(key)?
.as_str()
.and_then(|value| {
let trimmed = value.trim();
(!trimmed.is_empty()).then(|| trimmed.to_string())
})
}
fn non_empty_array(value: Option<Value>) -> Option<Vec<Value>> {
let Value::Array(values) = value? else {
return None;
+26
View File
@@ -1,5 +1,6 @@
//! 配额检查策略。
use chrono::{DateTime, Utc};
use serde_json::{json, Map, Value};
/// 判断配额百分比是否达到阈值。
@@ -7,6 +8,31 @@ pub fn quota_reached(used_percent: f64, threshold: f64) -> bool {
used_percent >= threshold
}
/// 判断配额快照的主 rate-limit 是否已经触顶。
pub fn quota_snapshot_limit_reached(quota: &Value) -> bool {
quota
.pointer("/rate_limit/limit_reached")
.and_then(Value::as_bool)
.unwrap_or(false)
}
/// 从配额快照读取主 rate-limit 的 UTC 重置时间。
pub fn quota_snapshot_reset_at(quota: &Value) -> Option<DateTime<Utc>> {
let reset_at = quota
.pointer("/rate_limit/reset_at")
.and_then(Value::as_i64)
.filter(|value| *value > 0)?;
DateTime::<Utc>::from_timestamp(reset_at, 0)
}
/// 从配额快照读取主 rate-limit 窗口大小。
pub fn quota_snapshot_limit_window_seconds(quota: &Value) -> Option<u64> {
quota
.pointer("/rate_limit/limit_window_seconds")
.and_then(Value::as_u64)
.filter(|value| *value > 0)
}
/// 从 Codex usage 响应中提取持久化配额快照。
pub fn quota_from_usage(usage: &Value) -> Value {
let additional = usage
+54
View File
@@ -0,0 +1,54 @@
use chrono::{Duration, TimeZone, Utc};
use codex_proxy_core::accounts::cloudflare::CloudflareChallengeCooldownTracker;
#[tokio::test]
async fn challenge_cooldown_should_escalate_and_cap_delay() {
let tracker = CloudflareChallengeCooldownTracker::new();
let now = Utc.with_ymd_and_hms(2026, 6, 20, 10, 0, 0).unwrap();
let first = tracker.record_challenge("acct_a", now).await;
let second = tracker
.record_challenge("acct_a", now + Duration::seconds(1))
.await;
let third = tracker
.record_challenge("acct_a", now + Duration::seconds(2))
.await;
let fourth = tracker
.record_challenge("acct_a", now + Duration::seconds(3))
.await;
let fifth = tracker
.record_challenge("acct_a", now + Duration::seconds(4))
.await;
assert_eq!(
(
first.delay_seconds,
second.delay_seconds,
third.delay_seconds,
fourth.delay_seconds,
fifth.delay_seconds,
),
(10, 30, 90, 120, 120)
);
}
#[tokio::test]
async fn challenge_cooldown_should_reset_after_stale_window() {
let tracker = CloudflareChallengeCooldownTracker::new();
let now = Utc.with_ymd_and_hms(2026, 6, 20, 10, 0, 0).unwrap();
let first = tracker.record_challenge("acct_a", now).await;
let after_stale = tracker
.record_challenge("acct_a", now + Duration::hours(2))
.await;
assert_eq!(
(
first.challenge_count,
first.delay_seconds,
after_stale.challenge_count,
after_stale.delay_seconds,
),
(1, 10, 1, 10)
);
}
@@ -35,19 +35,18 @@ fn codex_websocket_payload_audit_snapshot_should_redact_user_content() {
"type",
"model",
"instructions",
"previous_response_id",
"input",
"store",
"stream",
"previous_response_id",
"reasoning",
"tools",
"tool_choice",
"parallel_tool_calls",
"reasoning",
"store",
"stream",
"include",
"text",
"service_tier",
"prompt_cache_key",
"text",
"generate",
"include",
"client_metadata",
]
);
@@ -92,20 +91,32 @@ fn codex_websocket_response_create_payload_text_should_preserve_canonical_field_
"\"model\":\"gpt-5.5\"",
"\"instructions\":\"private capture instructions\"",
"\"input\":",
"\"tools\":[]",
"\"tool_choice\":\"auto\"",
"\"parallel_tool_calls\":true",
"\"reasoning\":null",
"\"store\":false",
"\"stream\":true",
"\"include\":[]",
"\"tool_choice\":\"auto\"",
"\"parallel_tool_calls\":true",
"\"prompt_cache_key\":\"session-1\"",
"\"generate\":false",
"\"client_metadata\":",
],
);
}
#[test]
fn codex_websocket_response_create_payload_text_should_include_empty_instructions() {
let request = CodexResponsesRequest::new_http_sse("gpt-5.5", "", Vec::new());
let payload =
codex_proxy_core::protocol::codex::websocket::websocket_response_create_payload_text(
&request,
)
.expect("payload should serialize");
let value: serde_json::Value = serde_json::from_str(&payload).unwrap();
let snapshot = websocket_payload_audit_snapshot(&request);
assert_eq!(value["instructions"], "");
assert_eq!(snapshot.top_level_keys[2], "instructions");
}
#[test]
fn codex_websocket_audit_artifact_should_record_transport_opening_and_payload() {
let mut request = CodexResponsesRequest::new_http_sse(
@@ -69,3 +69,42 @@ fn chat_completion_from_codex_sse_should_convert_completed_response() {
assert_eq!(response["choices"][0]["message"]["content"], "hello");
}
#[test]
fn chat_completion_stream_translator_should_emit_openai_chunks() {
let mut translator =
codex_proxy_core::protocol::openai::chat::ChatCompletionStreamTranslator::new(
"gpt-5.5", false, None,
);
let first = concat!(
"event: response.output_text.delta\n",
"data: {\"delta\":\"he",
);
let rest = concat!(
"llo\"}\n\n",
"event: response.completed\n",
"data: {\"response\":{\"id\":\"resp_1\",\"usage\":{\"input_tokens\":2,\"output_tokens\":3}}}\n\n",
);
let pending_output = translator
.push_str(first)
.expect("partial event should be buffered");
assert!(pending_output.is_empty());
let output = format!(
"{}{}",
translator.initial_frame(),
translator
.push_str(rest)
.expect("stream conversion should succeed")
);
assert_substrings_appear_in_order(
&output,
&[
"\"delta\":{\"role\":\"assistant\"}",
"\"delta\":{\"content\":\"hello\"}",
"\"finish_reason\":\"stop\"",
"\"usage\":{\"completion_tokens\":3,\"prompt_tokens\":2,\"total_tokens\":5}",
"data: [DONE]",
],
);
}
@@ -105,7 +105,17 @@ fn openai_response_request_should_translate_to_codex_request() {
let codex = translate_response_to_codex(request);
assert_eq!(codex.model, "gpt-5.5");
assert_eq!(codex.input, vec![json!("hello")]);
assert_eq!(
codex.input,
vec![json!({
"type": "message",
"role": "user",
"content": [{
"type": "input_text",
"text": "hello"
}]
})]
);
assert_eq!(codex.previous_response_id.as_deref(), Some("resp_previous"));
assert_eq!(codex.turn_state.as_deref(), Some("turn_previous"));
assert!(!codex.use_websocket);
@@ -236,6 +246,38 @@ fn openai_response_request_should_forward_parity_fields_to_codex() {
assert!(codex.force_http_sse);
}
#[test]
fn openai_response_request_should_fallback_context_fields_to_client_metadata() {
let request = serde_json::from_value::<OpenAiResponsesRequest>(json!({
"model": "gpt-5.5",
"stream": false,
"input": [],
"client_metadata": {
"x-codex-turn-metadata": " meta-from-metadata ",
"x-codex-beta-features": " beta-from-metadata ",
"x-responsesapi-include-timing-metrics": " true ",
"x-codex-window-id": " window-from-metadata ",
"x-codex-parent-thread-id": " parent-from-metadata "
},
"betaFeatures": "beta-direct"
}))
.expect("responses request should deserialize");
let codex = translate_response_to_codex(request);
assert_eq!(codex.turn_metadata.as_deref(), Some("meta-from-metadata"));
assert_eq!(codex.beta_features.as_deref(), Some("beta-direct"));
assert_eq!(codex.include_timing_metrics.as_deref(), Some("true"));
assert_eq!(
codex.codex_window_id.as_deref(),
Some("window-from-metadata")
);
assert_eq!(
codex.parent_thread_id.as_deref(),
Some("parent-from-metadata")
);
}
#[test]
fn codex_responses_model_options_should_apply_suffix_defaults_and_include_reasoning() {
let mut request = CodexResponsesRequest::new_http_sse("gpt-5.5-high-fast", "", Vec::new());
-1
View File
@@ -13,7 +13,6 @@ base64.workspace = true
chrono.workspace = true
config.workspace = true
dirs.workspace = true
flate2.workspace = true
hex.workspace = true
hmac.workspace = true
rand.workspace = true
+1
View File
@@ -51,6 +51,7 @@ create table if not exists accounts (
access_token_cipher text not null,
refresh_token_cipher text,
access_token_expires_at text,
next_refresh_at text,
status text not null check (status in ('active', 'expired', 'quota_exhausted', 'refreshing', 'disabled', 'banned')),
quota_json text,
quota_fetched_at text,
+31
View File
@@ -22,6 +22,37 @@ pub async fn connect_sqlite(database_url: &str) -> Result<SqlitePool, sqlx::Erro
async fn initialize_schema(pool: &SqlitePool) -> Result<(), sqlx::Error> {
let schema = include_str!("schema.sql");
sqlx::raw_sql(schema).execute(pool).await?;
apply_lightweight_migrations(pool).await?;
Ok(())
}
async fn apply_lightweight_migrations(pool: &SqlitePool) -> Result<(), sqlx::Error> {
add_accounts_next_refresh_at_if_missing(pool).await?;
Ok(())
}
async fn add_accounts_next_refresh_at_if_missing(pool: &SqlitePool) -> Result<(), sqlx::Error> {
let table_exists: Option<(String,)> =
sqlx::query_as("select name from sqlite_master where type = 'table' and name = ?")
.bind("accounts")
.fetch_optional(pool)
.await?;
if table_exists.is_none() {
return Ok(());
}
let column_exists: Option<(String,)> =
sqlx::query_as("select name from pragma_table_info('accounts') where name = ?")
.bind("next_refresh_at")
.fetch_optional(pool)
.await?;
if column_exists.is_some() {
return Ok(());
}
sqlx::query("alter table accounts add column next_refresh_at text")
.execute(pool)
.await?;
Ok(())
}
File diff suppressed because it is too large Load Diff
+26 -6
View File
@@ -11,7 +11,10 @@ use codex_proxy_core::{
model::AccountStatus,
ports::{AccountStore, AccountStoreError},
},
serving::quota::quota_from_usage,
serving::quota::{
quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds,
quota_snapshot_reset_at,
},
};
use thiserror::Error;
use tokio::time::{interval, sleep, Instant};
@@ -105,7 +108,7 @@ impl QuotaRefreshTask {
candidates = summary.candidates,
skipped_recent = summary.skipped_recent,
failed = summary.failed,
"没有需要刷新的 quota 锁定账号"
"没有需要刷新的 quota 锁定或待验证账号"
);
}
Err(error) => {
@@ -124,7 +127,7 @@ impl QuotaRefreshTask {
SchedulerHandle::new(shutdown_tx)
}
/// 执行一次配额锁定账号刷新。
/// 执行一次配额锁定或待验证账号刷新。
pub async fn refresh_locked_accounts_once(
&self,
) -> QuotaRefreshTaskResult<QuotaRefreshSummary> {
@@ -151,7 +154,8 @@ impl QuotaRefreshTask {
let mut candidates = accounts
.into_iter()
.filter(|account| {
account.status == AccountStatus::Active && account.quota_limit_reached
account.status == AccountStatus::Active
&& (account.quota_limit_reached || account.quota_verify_required)
})
.peekable();
@@ -188,15 +192,31 @@ impl QuotaRefreshTask {
{
Ok(raw) => {
let quota = quota_from_usage(&raw);
let limit_reached = quota_snapshot_limit_reached(&quota);
let reset_at = quota_snapshot_reset_at(&quota);
if self
.store
.update_quota_json(&account.id, &quota.to_string())
.apply_quota_snapshot(
&account.id,
&quota.to_string(),
limit_reached,
limit_reached.then_some(reset_at).flatten(),
)
.await?
{
summary.refreshed += 1;
} else {
summary.missing += 1;
}
if let Some(reset_at) = reset_at {
self.store
.sync_rate_limit_window(
&account.id,
reset_at,
quota_snapshot_limit_window_seconds(&quota),
)
.await?;
}
}
Err(error) => {
warn!(
@@ -222,7 +242,7 @@ impl QuotaRefreshTask {
pub struct QuotaRefreshSummary {
/// 扫描账号数。
pub scanned: usize,
/// 配额锁定候选账号数。
/// 配额锁定或待验证候选账号数。
pub candidates: usize,
/// 成功刷新账号数。
pub refreshed: usize,
+161 -74
View File
@@ -57,8 +57,20 @@ where
lease_owner: String,
retry_delays: Vec<Duration>,
in_flight: Arc<Mutex<HashSet<String>>>,
recovery_not_before: Arc<Mutex<HashMap<String, DateTime<Utc>>>>,
timers: Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
timers: Arc<Mutex<HashMap<String, ScheduledRefreshTimer>>>,
}
struct ScheduledRefreshTimer {
scheduled_at: DateTime<Utc>,
trigger: RefreshTrigger,
handle: JoinHandle<()>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ScheduleTimerResult {
Scheduled,
Replaced,
Unchanged,
}
impl<C> Clone for TokenRefreshTask<C>
@@ -75,7 +87,6 @@ where
lease_owner: self.lease_owner.clone(),
retry_delays: self.retry_delays.clone(),
in_flight: self.in_flight.clone(),
recovery_not_before: self.recovery_not_before.clone(),
timers: self.timers.clone(),
}
}
@@ -96,7 +107,6 @@ where
lease_owner: refresh_lease_owner(),
retry_delays: default_refresh_retry_delays(),
in_flight: Arc::new(Mutex::new(HashSet::new())),
recovery_not_before: Arc::new(Mutex::new(HashMap::new())),
timers: Arc::new(Mutex::new(HashMap::new())),
}
}
@@ -117,7 +127,6 @@ where
lease_owner: refresh_lease_owner(),
retry_delays: default_refresh_retry_delays(),
in_flight: Arc::new(Mutex::new(HashSet::new())),
recovery_not_before: Arc::new(Mutex::new(HashMap::new())),
timers: Arc::new(Mutex::new(HashMap::new())),
}
}
@@ -222,7 +231,12 @@ where
let timer = match account.status {
AccountStatus::Refreshing => Some((Duration::ZERO, RefreshTrigger::Unauthorized)),
AccountStatus::Expired => {
let delay = Duration::from_secs(30 + expired_recovery_index * 2);
let delay = match account.next_refresh_at {
Some(next_refresh_at) if next_refresh_at > now => {
duration_until(now, next_refresh_at)
}
_ => Duration::from_secs(30 + expired_recovery_index * 2),
};
expired_recovery_index += 1;
Some((delay, RefreshTrigger::Unauthorized))
}
@@ -252,11 +266,18 @@ where
}
let scheduled_at = scheduled_at_from_delay(now, delay);
if self
self.persist_next_refresh_at_if_changed(&account, Some(scheduled_at))
.await?;
let schedule_result = self
.schedule_account_timer(account.id.clone(), delay, trigger, scheduled_at)
.await
{
summary.replaced += 1;
.await;
match schedule_result {
ScheduleTimerResult::Unchanged => {
summary.skipped += 1;
continue;
}
ScheduleTimerResult::Replaced => summary.replaced += 1,
ScheduleTimerResult::Scheduled => {}
}
if account.status == AccountStatus::Expired {
@@ -292,7 +313,7 @@ where
};
for account in accounts {
if self.is_recovery_delayed(&account.id, now).await {
if account.next_refresh_at.is_some_and(|value| value > now) {
summary.skipped += 1;
continue;
}
@@ -343,11 +364,11 @@ where
}
fn timer_delay_for_account(&self, account: &Account, now: DateTime<Utc>) -> Option<Duration> {
let expires_at = account
.access_token_expires_at
.or_else(|| jwt_expiration(&account.access_token))?;
let margin_seconds = self.refresh_margin_seconds.min(i64::MAX as u64) as i64;
let refresh_at = expires_at - chrono::Duration::seconds(margin_seconds);
if let Some(next_refresh_at) = account.next_refresh_at.filter(|value| *value > now) {
return Some(duration_until(now, next_refresh_at));
}
let refresh_at = self.computed_refresh_at(account)?;
if refresh_at <= now {
return Some(Duration::ZERO);
@@ -356,14 +377,42 @@ where
(refresh_at - now).to_std().ok()
}
fn computed_refresh_at(&self, account: &Account) -> Option<DateTime<Utc>> {
let expires_at = account
.access_token_expires_at
.or_else(|| jwt_expiration(&account.access_token))?;
let margin_seconds = self.refresh_margin_seconds.min(i64::MAX as u64) as i64;
Some(expires_at - chrono::Duration::seconds(margin_seconds))
}
async fn persist_next_refresh_at_if_changed(
&self,
account: &Account,
next_refresh_at: Option<DateTime<Utc>>,
) -> TokenRefreshTaskResult<()> {
if same_refresh_time(account.next_refresh_at, next_refresh_at) {
return Ok(());
}
persist_next_refresh_at(&self.store, &account.id, next_refresh_at).await
}
fn schedule_account_timer(
&self,
account_id: String,
delay: Duration,
trigger: RefreshTrigger,
scheduled_at: DateTime<Utc>,
) -> Pin<Box<dyn Future<Output = bool> + Send + '_>> {
) -> Pin<Box<dyn Future<Output = ScheduleTimerResult> + Send + '_>> {
Box::pin(async move {
{
let timers = self.timers.lock().await;
if timers.get(&account_id).is_some_and(|timer| {
timer.scheduled_at == scheduled_at && timer.trigger == trigger
}) {
return ScheduleTimerResult::Unchanged;
}
}
let replaced = self.clear_scheduled_timer(&account_id).await;
let task = self.clone();
let timer_account_id = account_id.clone();
@@ -384,14 +433,25 @@ where
}
});
self.timers.lock().await.insert(account_id, handle);
replaced
self.timers.lock().await.insert(
account_id,
ScheduledRefreshTimer {
scheduled_at,
trigger,
handle,
},
);
if replaced {
ScheduleTimerResult::Replaced
} else {
ScheduleTimerResult::Scheduled
}
})
}
async fn clear_scheduled_timer(&self, account_id: &str) -> bool {
if let Some(handle) = self.timers.lock().await.remove(account_id) {
handle.abort();
if let Some(timer) = self.timers.lock().await.remove(account_id) {
timer.handle.abort();
true
} else {
false
@@ -400,8 +460,8 @@ where
async fn clear_scheduled_timers(&self) {
let mut timers = self.timers.lock().await;
for (_, handle) in timers.drain() {
handle.abort();
for (_, timer) in timers.drain() {
timer.handle.abort();
}
}
@@ -411,16 +471,15 @@ where
trigger: RefreshTrigger,
now: DateTime<Utc>,
) -> TokenRefreshTaskResult<TokenRefreshOutcome> {
if self.is_recovery_delayed(account_id, now).await {
return Ok(TokenRefreshOutcome::Skipped);
}
let Some(account) = self.store.get(account_id).await? else {
return Err(TokenRefreshTaskError::AccountNotFound(
account_id.to_string(),
));
};
let account = stored_account_to_refresh_account(account);
if account.next_refresh_at.is_some_and(|value| value > now) {
return Ok(TokenRefreshOutcome::Skipped);
}
if !self.try_mark_in_flight(&account.id).await {
return Ok(TokenRefreshOutcome::Skipped);
@@ -502,31 +561,33 @@ where
trigger: RefreshTrigger,
now: DateTime<Utc>,
) -> TokenRefreshTaskResult<TokenRefreshOutcome> {
let attempt = self.prepare_refresh_attempt(account).await?;
if attempt.refresh_token.is_none() {
persist_status(&self.store, &account.id, AccountStatus::Expired).await?;
self.clear_recovery_delay(&account.id).await;
return Ok(TokenRefreshOutcome::StatusUpdated);
}
let max_attempts = self.retry_delays.len() + 1;
let mut permanent_failures = 0;
for attempt_index in 0..max_attempts {
let Some(attempt) = self.prepare_refresh_attempt(account, now).await? else {
return Ok(TokenRefreshOutcome::Skipped);
};
if attempt.refresh_token.is_none() {
persist_status(&self.store, &attempt.id, AccountStatus::Expired).await?;
persist_next_refresh_at(&self.store, &attempt.id, None).await?;
return Ok(TokenRefreshOutcome::StatusUpdated);
}
match self
.scheduler
.refresh_account_at(&attempt, trigger, now)
.await
{
Ok(mut updated) if account_requires_token_write(account, &updated) => {
Ok(mut updated) if account_requires_token_write(&attempt, &updated) => {
updated.access_token_expires_at =
jwt_expiration(&updated.access_token).or(updated.access_token_expires_at);
updated.next_refresh_at = self.computed_refresh_at(&updated);
persist_token_update(&self.store, &updated).await?;
self.clear_recovery_delay(&account.id).await;
return Ok(TokenRefreshOutcome::Refreshed(Box::new(updated)));
}
Ok(updated) if account.status != updated.status => {
Ok(updated) if attempt.status != updated.status => {
if is_permanent_refresh_status(updated.status) {
permanent_failures += 1;
if permanent_failures < PERMANENT_FAILURE_CONFIRMATION_THRESHOLD
@@ -536,17 +597,21 @@ where
}
}
persist_status(&self.store, &account.id, updated.status).await?;
self.clear_recovery_delay(&account.id).await;
persist_status(&self.store, &attempt.id, updated.status).await?;
persist_next_refresh_at(&self.store, &attempt.id, None).await?;
return Ok(TokenRefreshOutcome::StatusUpdated);
}
Ok(_) => return Ok(TokenRefreshOutcome::Skipped),
Err(RefreshError::Transport) if self.sleep_before_retry(attempt_index).await => {
Err(RefreshError::RetryableTransport)
if self.sleep_before_retry(attempt_index).await =>
{
continue;
}
Err(RefreshError::Transport) => {
persist_status(&self.store, &account.id, AccountStatus::Active).await?;
self.schedule_recovery(&account.id, now).await;
Err(RefreshError::RetryableTransport) | Err(RefreshError::Transport) => {
persist_status(&self.store, &attempt.id, AccountStatus::Active).await?;
let next_refresh_at = now + chrono::Duration::seconds(RECOVERY_DELAY_SECONDS);
persist_next_refresh_at(&self.store, &attempt.id, Some(next_refresh_at))
.await?;
return Ok(TokenRefreshOutcome::Failed);
}
Err(error) => return Err(TokenRefreshTaskError::Refresh(error)),
@@ -556,14 +621,30 @@ where
Ok(TokenRefreshOutcome::Skipped)
}
async fn prepare_refresh_attempt(&self, account: &Account) -> TokenRefreshTaskResult<Account> {
if account.status != AccountStatus::Refreshing {
persist_status(&self.store, &account.id, AccountStatus::Refreshing).await?;
async fn prepare_refresh_attempt(
&self,
account: &Account,
now: DateTime<Utc>,
) -> TokenRefreshTaskResult<Option<Account>> {
let Some(stored) = self.store.get(&account.id).await? else {
return Err(TokenRefreshTaskError::AccountNotFound(account.id.clone()));
};
let mut attempt = stored_account_to_refresh_account(stored);
if matches!(
attempt.status,
AccountStatus::Disabled | AccountStatus::Banned
) || attempt.next_refresh_at.is_some_and(|value| value > now)
|| attempt.refresh_token != account.refresh_token
{
return Ok(None);
}
if attempt.status != AccountStatus::Refreshing {
persist_status(&self.store, &attempt.id, AccountStatus::Refreshing).await?;
}
let mut attempt = account.clone();
attempt.status = AccountStatus::Active;
Ok(attempt)
Ok(Some(attempt))
}
async fn sleep_before_retry(&self, attempt_index: usize) -> bool {
@@ -585,31 +666,6 @@ where
self.in_flight.lock().await.remove(account_id);
}
async fn is_recovery_delayed(&self, account_id: &str, now: DateTime<Utc>) -> bool {
let mut recovery_not_before = self.recovery_not_before.lock().await;
let Some(not_before) = recovery_not_before.get(account_id).copied() else {
return false;
};
if now < not_before {
return true;
}
recovery_not_before.remove(account_id);
false
}
async fn schedule_recovery(&self, account_id: &str, now: DateTime<Utc>) {
self.recovery_not_before.lock().await.insert(
account_id.to_string(),
now + chrono::Duration::seconds(RECOVERY_DELAY_SECONDS),
);
}
async fn clear_recovery_delay(&self, account_id: &str) {
self.recovery_not_before.lock().await.remove(account_id);
}
async fn try_acquire_refresh_lease(
&self,
account_id: &str,
@@ -728,6 +784,7 @@ async fn persist_token_update(
.map(|token| SecretString::new(token.into())),
access_token_expires_at: jwt_expiration(&account.access_token)
.or(account.access_token_expires_at),
next_refresh_at: account.next_refresh_at,
status: account.status,
},
)
@@ -754,6 +811,23 @@ async fn persist_status(
}
}
async fn persist_next_refresh_at(
store: &SqliteAccountStore,
account_id: &str,
next_refresh_at: Option<DateTime<Utc>>,
) -> TokenRefreshTaskResult<()> {
let updated = store
.set_next_refresh_at(account_id, next_refresh_at)
.await?;
if updated {
Ok(())
} else {
Err(TokenRefreshTaskError::AccountNotFound(
account_id.to_string(),
))
}
}
fn account_requires_token_write(previous: &Account, updated: &Account) -> bool {
previous.access_token != updated.access_token
|| previous.refresh_token != updated.refresh_token
@@ -793,6 +867,18 @@ fn scheduled_at_from_delay(now: DateTime<Utc>, delay: Duration) -> DateTime<Utc>
}
}
fn duration_until(now: DateTime<Utc>, scheduled_at: DateTime<Utc>) -> Duration {
(scheduled_at - now).to_std().unwrap_or(Duration::ZERO)
}
fn same_refresh_time(left: Option<DateTime<Utc>>, right: Option<DateTime<Utc>>) -> bool {
match (left, right) {
(Some(left), Some(right)) => left.timestamp() == right.timestamp(),
(None, None) => true,
_ => false,
}
}
fn stored_account_to_refresh_account(stored: StoredAccount) -> Account {
Account {
id: stored.id,
@@ -806,6 +892,7 @@ fn stored_account_to_refresh_account(stored: StoredAccount) -> Account {
.refresh_token
.map(|token| token.expose_secret().to_string()),
access_token_expires_at: stored.access_token_expires_at,
next_refresh_at: stored.next_refresh_at,
status: stored.status,
quota_limit_reached: false,
quota_verify_required: false,
+34 -1
View File
@@ -11,7 +11,7 @@ use async_trait::async_trait;
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
use chrono::{Duration, TimeZone, Utc};
use codex_proxy_adapters::sqlite::{
accounts::{NewAccount, SqliteAccountStore},
accounts::{AccountClaimsUpdate, NewAccount, SqliteAccountStore},
refresh_leases::SqliteRefreshLeaseStore,
};
use codex_proxy_core::{
@@ -148,6 +148,39 @@ impl TokenRefresher for SequenceTokenRefresher {
}
}
#[derive(Clone)]
struct RefreshTokenRotatingRefresher {
store: SqliteAccountStore,
account_id: String,
calls: Arc<AtomicUsize>,
access_token: String,
}
#[async_trait]
impl TokenRefresher for RefreshTokenRotatingRefresher {
async fn refresh(&self, _refresh_token: &str) -> Result<TokenPair, RefreshFailure> {
self.calls.fetch_add(1, Ordering::SeqCst);
self.store
.update_from_claims(
&self.account_id,
AccountClaimsUpdate {
email: Some("rotated@example.com".to_string()),
account_id: Some("chatgpt-rotated".to_string()),
user_id: Some("user-rotated".to_string()),
plan_type: Some("plus".to_string()),
access_token: SecretString::new(self.access_token.clone().into()),
refresh_token: Some(SecretString::new("refresh-rotated".to_string().into())),
access_token_expires_at: Some(Utc::now() + Duration::hours(1)),
next_refresh_at: None,
status: AccountStatus::Active,
},
)
.await
.expect("rotated token should persist");
Err(RefreshFailure::RetryableTransport)
}
}
fn test_jwt(exp: i64) -> String {
let header = json!({"alg": "none", "typ": "JWT"});
let payload = json!({ "exp": exp });
@@ -127,7 +127,7 @@ async fn token_refresh_task_should_recover_refreshing_account_after_restart() {
}
#[tokio::test]
async fn token_refresh_task_should_restore_active_status_after_transport_failure() {
async fn token_refresh_task_should_not_retry_ambiguous_transport_failure() {
let dir = tempfile::tempdir().expect("temp dir");
let db = dir.path().join("token-refresh-transport-failure.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
@@ -180,8 +180,9 @@ async fn token_refresh_task_should_restore_active_status_after_transport_failure
let observed_statuses = observed_statuses.lock().await.clone();
assert_eq!(summary.failed, 1);
assert_eq!(observed_statuses, [AccountStatus::Refreshing; 5]);
assert_eq!(observed_statuses, [AccountStatus::Refreshing]);
assert_eq!(stored.status, AccountStatus::Active);
assert!(stored.next_refresh_at.is_some_and(|next| next > now));
assert_eq!(
stored.access_token.expose_secret(),
old_access_token.as_str()
@@ -217,7 +218,7 @@ async fn token_refresh_task_should_delay_recovery_after_retry_exhaustion() {
.await
.expect("account should be inserted");
let mut responses = Vec::new();
responses.extend((0..5).map(|_| Err(RefreshFailure::Transport)));
responses.extend((0..5).map(|_| Err(RefreshFailure::RetryableTransport)));
responses.push(Ok(TokenPair {
access_token: new_access_token.clone(),
refresh_token: None,
@@ -270,7 +271,75 @@ async fn token_refresh_task_should_delay_recovery_after_retry_exhaustion() {
}
#[tokio::test]
async fn token_refresh_task_should_confirm_invalid_grant_before_expiring_account() {
async fn token_refresh_task_should_not_reuse_stale_refresh_token_after_retryable_failure() {
let dir = tempfile::tempdir().expect("temp dir");
let db = dir.path().join("token-refresh-stale-rt.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
.await
.expect("sqlite pool");
let store = SqliteAccountStore::new(pool, SecretBox::new([31u8; 32]));
let now = Utc.with_ymd_and_hms(2026, 6, 19, 16, 0, 0).unwrap();
let old_access_token = test_jwt((now + Duration::seconds(30)).timestamp());
let rotated_access_token = test_jwt((now + Duration::hours(1)).timestamp());
store
.insert(NewAccount {
id: "acct-refresh-stale-rt".to_string(),
email: Some("stale-rt@example.com".to_string()),
account_id: Some("chatgpt-stale-rt".to_string()),
user_id: Some("user-stale-rt".to_string()),
label: None,
plan_type: Some("plus".to_string()),
access_token: SecretString::new(old_access_token.into()),
refresh_token: Some(SecretString::new("refresh-stale-old".to_string().into())),
access_token_expires_at: Some(now + Duration::seconds(30)),
status: AccountStatus::Active,
})
.await
.expect("account should be inserted");
let calls = Arc::new(AtomicUsize::new(0));
let refresher = RefreshTokenRotatingRefresher {
store: store.clone(),
account_id: "acct-refresh-stale-rt".to_string(),
calls: calls.clone(),
access_token: rotated_access_token.clone(),
};
let task = codex_proxy_runtime::tasks::token_refresh::TokenRefreshTask::new(
store.clone(),
RefreshPolicy {
refresh_margin_seconds: 300,
refresh_concurrency: 1,
},
refresher,
)
.with_retry_delays(vec![StdDuration::ZERO; 4]);
let summary = task
.refresh_due_accounts_once_at(now)
.await
.expect("stale refresh token should skip retry");
let stored = store
.get("acct-refresh-stale-rt")
.await
.expect("account should load")
.expect("account should exist");
assert_eq!(summary.skipped, 1);
assert_eq!(calls.load(Ordering::SeqCst), 1);
assert_eq!(
stored.access_token.expose_secret(),
rotated_access_token.as_str()
);
assert_eq!(
stored
.refresh_token
.as_ref()
.map(|token| token.expose_secret()),
Some("refresh-rotated")
);
}
#[tokio::test]
async fn token_refresh_task_should_confirm_invalid_grant_before_disabling_account() {
let dir = tempfile::tempdir().expect("temp dir");
let db = dir.path().join("token-refresh-permanent-failure.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
@@ -329,7 +398,7 @@ async fn token_refresh_task_should_confirm_invalid_grant_before_expiring_account
assert_eq!(summary.status_updated, 1);
assert_eq!(observed_statuses, [AccountStatus::Refreshing; 2]);
assert_eq!(stored.status, AccountStatus::Expired);
assert_eq!(stored.status, AccountStatus::Disabled);
assert_eq!(
stored.access_token.expose_secret(),
old_access_token.as_str()
@@ -208,7 +208,7 @@ async fn token_refresh_task_should_fire_per_account_timer_at_refresh_time() {
.await
.expect("sqlite pool");
let store = SqliteAccountStore::new(pool, SecretBox::new([21u8; 32]));
let now = Utc.with_ymd_and_hms(2026, 6, 19, 17, 0, 0).unwrap();
let now = Utc::now() + Duration::minutes(5);
let old_access_token = test_jwt((now + Duration::seconds(1)).timestamp());
let new_access_token = test_jwt((now + Duration::hours(1)).timestamp());
store
@@ -107,7 +107,7 @@ async fn token_refresh_task_should_persist_success_after_transient_transport_ret
account_id: "acct-refresh-transient-success".to_string(),
observed_statuses: observed_statuses.clone(),
responses: Arc::new(Mutex::new(VecDeque::from(vec![
Err(RefreshFailure::Transport),
Err(RefreshFailure::RetryableTransport),
Ok(TokenPair {
access_token: new_access_token.clone(),
refresh_token: None,
+39 -4
View File
@@ -10,8 +10,16 @@ use codex_proxy_core::{
use codex_proxy_platform::config::WebSocketPoolConfig;
use futures::{SinkExt, StreamExt};
use serde_json::json;
use tokio::net::TcpListener;
use tokio_tungstenite::{accept_async, tungstenite::Message};
use tokio::net::{TcpListener, TcpStream};
use tokio_tungstenite::{
accept_hdr_async_with_config,
tungstenite::{
extensions::{compression::deflate::DeflateConfig, ExtensionsConfig},
handshake::server::{Request as WsRequest, Response as WsResponse},
protocol::WebSocketConfig,
Message,
},
};
#[tokio::test]
async fn codex_backend_client_should_apply_configured_websocket_pool() {
@@ -22,7 +30,7 @@ async fn codex_backend_client_should_apply_configured_websocket_pool() {
let server = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let mut first_websocket = accept_async(first_stream).await.unwrap();
let mut first_websocket = accept_runtime_test_websocket(first_stream).await;
let _first_message = first_websocket.next().await.unwrap().unwrap();
first_websocket
.send(Message::Text(
@@ -45,7 +53,7 @@ async fn codex_backend_client_should_apply_configured_websocket_pool() {
accepted = listener.accept() => {
let (second_stream, _) = accepted.unwrap();
accepted_connections_for_server.fetch_add(1, Ordering::SeqCst);
let mut second_websocket = accept_async(second_stream).await.unwrap();
let mut second_websocket = accept_runtime_test_websocket(second_stream).await;
let _second_message = second_websocket.next().await.unwrap().unwrap();
second_websocket
.send(Message::Text(
@@ -93,6 +101,33 @@ async fn codex_backend_client_should_apply_configured_websocket_pool() {
assert_eq!(accepted_connections.load(Ordering::SeqCst), 1);
}
fn websocket_accept_config() -> WebSocketConfig {
let mut extensions = ExtensionsConfig::default();
extensions.permessage_deflate = Some(DeflateConfig::default());
let mut config = WebSocketConfig::default();
config.extensions = extensions;
config
}
async fn accept_runtime_test_websocket(
stream: TcpStream,
) -> tokio_tungstenite::WebSocketStream<TcpStream> {
accept_hdr_async_with_config(
stream,
|_request: &WsRequest, mut response: WsResponse| {
response.headers_mut().insert(
"sec-websocket-extensions",
"permessage-deflate".parse().unwrap(),
);
Ok(response)
},
Some(websocket_accept_config()),
)
.await
.unwrap()
}
fn request_context<'a>(
request_id: &'a str,
account_id: Option<&'a str>,
+1
View File
@@ -12,6 +12,7 @@ codex-proxy-assets = { path = "../assets" }
codex-proxy-core = { path = "../core" }
codex-proxy-platform = { path = "../platform" }
codex-proxy-runtime = { path = "../runtime" }
futures.workspace = true
serde.workspace = true
serde_json.workspace = true
tokio.workspace = true
@@ -22,7 +22,7 @@ use crate::{
struct AccountImportData {
imported: u32,
skipped: u32,
source_format: &'static str,
source_format: String,
}
/// `POST /api/admin/accounts/import`
@@ -42,7 +42,7 @@ pub async fn import_accounts(
AccountImportData {
imported: imported.imported,
skipped: imported.skipped,
source_format: "native",
source_format: imported.source_format.to_string(),
},
request_id,
),
+216 -5
View File
@@ -1,15 +1,29 @@
//! OpenAI 聊天处理器。
use axum::{
body::Bytes,
body::{Body, Bytes},
extract::State,
http::{HeaderMap, StatusCode},
http::{
header::{CACHE_CONTROL, CONTENT_TYPE},
HeaderMap, StatusCode,
},
response::{IntoResponse, Response},
Extension, Json,
};
use codex_proxy_core::protocol::openai::chat::{translate_chat_to_codex, ChatCompletionRequest};
use codex_proxy_runtime::{services::ChatDispatchError, state::AppState};
use codex_proxy_core::{
models::catalog::ModelCatalog,
protocol::openai::chat::{
translate_chat_to_codex, ChatCompletionRequest, ChatCompletionStreamTranslator,
},
serving::responses::apply_response_model_options,
};
use codex_proxy_runtime::{
services::{ChatDispatchError, ResponseDispatchError, ResponseDispatchStream},
state::AppState,
};
use futures::{stream as futures_stream, StreamExt};
use serde_json::{json, Value};
use std::convert::Infallible;
use crate::middleware::request_id::RequestId;
@@ -40,7 +54,10 @@ pub async fn chat_completions(
if !catalog.is_recognized_model_name(&model) {
return model_not_found_response().into_response();
}
let codex_request = match translate_chat_to_codex(chat_request) {
let parsed_model = catalog.parse_model_name(&model);
let display_model = ModelCatalog::build_display_model_name(&parsed_model);
let stream = chat_request.stream;
let mut codex_request = match translate_chat_to_codex(chat_request) {
Ok(request) => request,
Err(_) => {
return openai_error_response(
@@ -52,6 +69,40 @@ pub async fn chat_completions(
.into_response();
}
};
apply_response_model_options(
&mut codex_request,
&parsed_model,
state.services.models.config(),
);
let include_reasoning = codex_request
.reasoning
.as_ref()
.and_then(|reasoning| reasoning.get("effort"))
.and_then(Value::as_str)
.is_some_and(|effort| !effort.trim().is_empty());
let tuple_schema = codex_request.tuple_schema.clone();
if stream {
return match state
.services
.responses
.stream(
request_id.as_str(),
"/v1/chat/completions",
codex_request,
&model,
)
.await
{
Ok(stream) => live_chat_event_stream_response(
stream,
&display_model,
include_reasoning,
tuple_schema,
),
Err(error) => response_dispatch_chat_stream_error_response(error),
};
}
match state
.services
@@ -112,6 +163,31 @@ pub async fn chat_completions(
"upstream_error",
)
.into_response(),
Err(ChatDispatchError::Disabled {
count,
upstream_error,
}) => openai_error_response(
StatusCode::UNAUTHORIZED,
&format!(
"All accounts exhausted ({count} disabled). Codex upstream error: {upstream_error}"
),
"server_error",
"upstream_error",
)
.into_response(),
Err(ChatDispatchError::Banned {
count,
upstream_error,
status_code,
}) => openai_error_response(
StatusCode::from_u16(status_code).unwrap_or(StatusCode::FORBIDDEN),
&format!(
"All accounts exhausted ({count} banned). Codex upstream error: {upstream_error}"
),
"server_error",
"upstream_error",
)
.into_response(),
Err(ChatDispatchError::CloudflareChallenge {
count,
upstream_error,
@@ -195,3 +271,138 @@ fn openai_error_response(
})),
)
}
fn live_chat_event_stream_response(
stream: ResponseDispatchStream,
model: &str,
include_reasoning: bool,
tuple_schema: Option<Value>,
) -> Response {
let mut translator =
ChatCompletionStreamTranslator::new(model.to_string(), include_reasoning, tuple_schema);
let initial_frame = translator.initial_frame();
let body_stream =
futures_stream::once(async move { Ok::<Bytes, Infallible>(Bytes::from(initial_frame)) })
.chain(stream.body.map(move |result| {
let body = match result {
Ok(bytes) => {
let text = String::from_utf8_lossy(&bytes);
translator
.push_str(&text)
.unwrap_or_else(|error| chat_stream_error_sse_frame(&error.to_string()))
}
Err(error) => chat_stream_error_sse_frame(&error.to_string()),
};
Ok::<Bytes, Infallible>(Bytes::from(body))
}));
Response::builder()
.status(StatusCode::OK)
.header(CONTENT_TYPE, "text/event-stream")
.header(CACHE_CONTROL, "no-cache")
.header("connection", "keep-alive")
.header("x-accel-buffering", "no")
.body(Body::from_stream(body_stream))
.unwrap_or_else(|_| {
openai_error_response(
StatusCode::INTERNAL_SERVER_ERROR,
"Failed to build stream response",
"server_error",
"stream_response_error",
)
.into_response()
})
}
fn response_dispatch_chat_stream_error_response(error: ResponseDispatchError) -> Response {
let message = match error {
ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore => {
"No active upstream account is available".to_string()
}
ResponseDispatchError::Upstream(_) => "Upstream Codex request failed".to_string(),
ResponseDispatchError::QuotaExhausted {
count,
upstream_error,
} => format!(
"All accounts exhausted ({count} quota-exhausted). Codex upstream error: {upstream_error}"
),
ResponseDispatchError::RateLimited {
count,
upstream_error,
} => format!(
"All accounts exhausted ({count} rate-limited). Codex upstream error: {upstream_error}"
),
ResponseDispatchError::Expired {
count,
upstream_error,
} => format!(
"All accounts exhausted ({count} expired). Codex upstream error: {upstream_error}"
),
ResponseDispatchError::Disabled {
count,
upstream_error,
} => format!(
"All accounts exhausted ({count} disabled). Codex upstream error: {upstream_error}"
),
ResponseDispatchError::Banned {
count,
upstream_error,
..
} => format!(
"All accounts exhausted ({count} banned). Codex upstream error: {upstream_error}"
),
ResponseDispatchError::CloudflareChallenge {
count,
upstream_error,
} => format!(
"All accounts exhausted ({count} cloudflare-challenge). Codex upstream error: {upstream_error}"
),
ResponseDispatchError::CloudflarePathBlocked {
count,
upstream_error,
} => format!(
"All accounts exhausted ({count} cloudflare-path-block). Codex upstream error: {upstream_error}"
),
ResponseDispatchError::ModelUnsupported {
count,
upstream_error,
} => format!(
"All accounts exhausted ({count} model-unsupported). Codex upstream error: {upstream_error}"
),
ResponseDispatchError::InvalidSse(_)
| ResponseDispatchError::MissingCompleted
| ResponseDispatchError::EmptyUpstreamResponse => "Invalid upstream Codex response".to_string(),
ResponseDispatchError::Failed(_) => "Upstream Codex response failed".to_string(),
};
chat_stream_error_response(&message)
}
fn chat_stream_error_response(message: &str) -> Response {
Response::builder()
.status(StatusCode::OK)
.header(CONTENT_TYPE, "text/event-stream")
.header(CACHE_CONTROL, "no-cache")
.header("connection", "keep-alive")
.body(Body::from(chat_stream_error_sse_frame(message)))
.unwrap_or_else(|_| {
openai_error_response(
StatusCode::INTERNAL_SERVER_ERROR,
"Failed to build stream response",
"server_error",
"stream_response_error",
)
.into_response()
})
}
fn chat_stream_error_sse_frame(message: &str) -> String {
format!(
"data: {}\n\n",
json!({
"error": {
"message": message,
"type": "stream_error",
}
})
)
}
+78 -4
View File
@@ -37,7 +37,7 @@ pub async fn responses(
headers: HeaderMap,
body: Bytes,
) -> Response {
handle_responses(state, request_id, headers, body, None).await
handle_responses(state, request_id, headers, body, "/v1/responses", None).await
}
/// `POST /v1/responses/review`
@@ -47,7 +47,15 @@ pub async fn review_responses(
headers: HeaderMap,
body: Bytes,
) -> Response {
handle_responses(state, request_id, headers, body, Some("review")).await
handle_responses(
state,
request_id,
headers,
body,
"/v1/responses/review",
Some("review"),
)
.await
}
/// `POST /v1/responses/compact`
@@ -65,6 +73,7 @@ async fn handle_responses(
request_id: RequestId,
headers: HeaderMap,
body: Bytes,
route: &str,
forced_subagent: Option<&str>,
) -> Response {
if !authorize_client_api_key(&state, &headers).await {
@@ -107,7 +116,7 @@ async fn handle_responses(
return match state
.services
.responses
.stream(request_id.as_str(), codex_request, &model)
.stream(request_id.as_str(), route, codex_request, &model)
.await
{
Ok(stream) => live_event_stream_response(stream),
@@ -118,7 +127,7 @@ async fn handle_responses(
match state
.services
.responses
.complete(request_id.as_str(), codex_request, &model)
.complete(request_id.as_str(), route, codex_request, &model)
.await
{
Ok(body) => (StatusCode::OK, Json(body)).into_response(),
@@ -168,6 +177,31 @@ async fn handle_responses(
"upstream_error",
)
.into_response(),
Err(ResponseDispatchError::Disabled {
count,
upstream_error,
}) => openai_error_response(
StatusCode::UNAUTHORIZED,
&format!(
"All accounts exhausted ({count} disabled). Codex upstream error: {upstream_error}"
),
"server_error",
"upstream_error",
)
.into_response(),
Err(ResponseDispatchError::Banned {
count,
upstream_error,
status_code,
}) => openai_error_response(
StatusCode::from_u16(status_code).unwrap_or(StatusCode::FORBIDDEN),
&format!(
"All accounts exhausted ({count} banned). Codex upstream error: {upstream_error}"
),
"server_error",
"upstream_error",
)
.into_response(),
Err(ResponseDispatchError::CloudflareChallenge {
count,
upstream_error,
@@ -449,6 +483,27 @@ fn response_dispatch_compact_error_response(error: ResponseDispatchError) -> Res
),
)
.into_response(),
ResponseDispatchError::Disabled {
count,
upstream_error,
} => responses_dispatch_error_response(
StatusCode::UNAUTHORIZED,
&format!(
"All accounts exhausted ({count} disabled). Codex upstream error: {upstream_error}"
),
)
.into_response(),
ResponseDispatchError::Banned {
count,
upstream_error,
status_code,
} => responses_dispatch_error_response(
StatusCode::from_u16(status_code).unwrap_or(StatusCode::FORBIDDEN),
&format!(
"All accounts exhausted ({count} banned). Codex upstream error: {upstream_error}"
),
)
.into_response(),
ResponseDispatchError::CloudflareChallenge {
count,
upstream_error,
@@ -566,6 +621,25 @@ fn response_dispatch_stream_error_response(error: ResponseDispatchError) -> Resp
"All accounts exhausted ({count} expired). Codex upstream error: {upstream_error}"
),
),
ResponseDispatchError::Disabled {
count,
upstream_error,
} => (
StatusCode::UNAUTHORIZED,
format!(
"All accounts exhausted ({count} disabled). Codex upstream error: {upstream_error}"
),
),
ResponseDispatchError::Banned {
count,
upstream_error,
status_code,
} => (
StatusCode::from_u16(status_code).unwrap_or(StatusCode::FORBIDDEN),
format!(
"All accounts exhausted ({count} banned). Codex upstream error: {upstream_error}"
),
),
ResponseDispatchError::CloudflareChallenge {
count,
upstream_error,
@@ -124,6 +124,280 @@ async fn admin_accounts_import_should_store_native_account_tokens() {
assert_eq!(stored.status, AccountStatus::Disabled);
}
#[tokio::test]
async fn admin_accounts_import_should_update_existing_sub2api_account_and_clear_stale_quota_lock() {
let (app, state, pool, _dir, secret_box) =
admin_accounts_test_app("admin-accounts-import-update-existing.sqlite", 116).await;
let new_access_token = test_jwt(
"chatgpt-import-update",
Some("user-import-update"),
Some("new@example.com"),
Some("plus"),
);
seed_encrypted_account(
&pool,
secret_box.clone(),
NewAccount {
id: "acct_import_update".to_string(),
email: Some("old@example.com".to_string()),
account_id: Some("chatgpt-import-update".to_string()),
user_id: Some("user-import-update".to_string()),
label: Some("old-label".to_string()),
plan_type: Some("free".to_string()),
access_token: SecretString::new("old-access".to_string().into()),
refresh_token: Some(SecretString::new("old-refresh".to_string().into())),
access_token_expires_at: None,
status: AccountStatus::Expired,
},
)
.await;
sqlx::query(
"update accounts set quota_limit_reached = 1, quota_cooldown_until = ?, quota_verify_required = 1 where id = ?",
)
.bind("2026-06-25T00:00:00Z")
.bind("acct_import_update")
.execute(&pool)
.await
.unwrap();
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/api/admin/accounts/import")
.header("content-type", "application/json")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_accounts_import_update")
.body(Body::from(
json!({
"accounts": [{
"id": "acct_import_update",
"email": "new@example.com",
"accountId": "chatgpt-import-update",
"userId": "user-import-update",
"label": "new-label",
"planType": "plus",
"token": new_access_token,
"refreshToken": "new-refresh",
"status": "active",
"cachedQuota": {
"rate_limit": {
"allowed": true,
"limit_reached": false,
"used_percent": 5
}
},
"quotaFetchedAt": "2026-06-19T14:00:00Z",
"quotaVerifyRequired": false,
"proxyApiKey": "exported-key-prefix",
"usage": {"requestCount": 9}
}]
})
.to_string(),
))
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let body = response_json(response).await;
let stored = SqliteAccountStore::new(pool.clone(), secret_box)
.get("acct_import_update")
.await
.unwrap()
.unwrap();
let quota_state: (i64, Option<String>, i64, Option<String>) = sqlx::query_as(
"select quota_limit_reached, quota_cooldown_until, quota_verify_required, quota_json from accounts where id = ?",
)
.bind("acct_import_update")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(status, StatusCode::OK);
assert_eq!(body["requestId"], "req_accounts_import_update");
assert_eq!(body["data"]["imported"], 1);
assert_eq!(body["data"]["skipped"], 0);
assert_eq!(body["data"]["sourceFormat"], "sub2api");
assert_eq!(stored.email.as_deref(), Some("new@example.com"));
assert_eq!(stored.label.as_deref(), Some("new-label"));
assert_eq!(stored.plan_type.as_deref(), Some("plus"));
assert_eq!(stored.access_token.expose_secret(), new_access_token);
assert_eq!(stored.refresh_token.unwrap().expose_secret(), "new-refresh");
assert_eq!(stored.status, AccountStatus::Active);
assert_eq!(quota_state.0, 0);
assert!(quota_state.1.is_none());
assert_eq!(quota_state.2, 0);
assert!(quota_state.3.is_some());
let acquired = state
.services
.account_pool
.acquire("gpt-5.5", Utc::now())
.await
.unwrap()
.account;
assert_eq!(acquired.id, "acct_import_update");
assert_eq!(acquired.access_token, new_access_token);
}
#[tokio::test]
async fn admin_accounts_import_should_update_existing_identity_when_import_id_changed() {
let (app, _state, pool, _dir, secret_box) =
admin_accounts_test_app("admin-accounts-import-update-existing-identity.sqlite", 126).await;
let new_access_token = test_jwt(
"chatgpt-import-identity",
Some("user-import-identity"),
Some("identity-new@example.com"),
Some("plus"),
);
seed_encrypted_account(
&pool,
secret_box.clone(),
NewAccount {
id: "acct_existing_identity".to_string(),
email: Some("identity-old@example.com".to_string()),
account_id: Some("chatgpt-import-identity".to_string()),
user_id: Some("user-import-identity".to_string()),
label: Some("old-label".to_string()),
plan_type: Some("free".to_string()),
access_token: SecretString::new("old-access".to_string().into()),
refresh_token: Some(SecretString::new("old-refresh".to_string().into())),
access_token_expires_at: None,
status: AccountStatus::Disabled,
},
)
.await;
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/api/admin/accounts/import")
.header("content-type", "application/json")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_accounts_import_update_identity")
.body(Body::from(
json!({
"accounts": [{
"id": "acct_changed_export_id",
"email": "identity-new@example.com",
"accountId": "chatgpt-import-identity",
"userId": "user-import-identity",
"label": "new-label",
"planType": "plus",
"token": new_access_token,
"refreshToken": "new-refresh",
"status": "active",
"cachedQuota": {
"rate_limit": {
"allowed": true,
"limit_reached": false,
"used_percent": 6
}
}
}]
})
.to_string(),
))
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let body = response_json(response).await;
let store = SqliteAccountStore::new(pool.clone(), secret_box);
let stored = store.get("acct_existing_identity").await.unwrap().unwrap();
let changed_id = store.get("acct_changed_export_id").await.unwrap();
let account_count: i64 = sqlx::query_scalar("select count(*) from accounts")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(status, StatusCode::OK);
assert_eq!(body["requestId"], "req_accounts_import_update_identity");
assert_eq!(body["data"]["imported"], 1);
assert_eq!(body["data"]["skipped"], 0);
assert_eq!(body["data"]["sourceFormat"], "sub2api");
assert_eq!(account_count, 1);
assert!(changed_id.is_none());
assert_eq!(stored.email.as_deref(), Some("identity-new@example.com"));
assert_eq!(stored.label.as_deref(), Some("new-label"));
assert_eq!(stored.plan_type.as_deref(), Some("plus"));
assert_eq!(stored.access_token.expose_secret(), new_access_token);
assert_eq!(stored.refresh_token.unwrap().expose_secret(), "new-refresh");
assert_eq!(stored.status, AccountStatus::Active);
}
#[tokio::test]
async fn admin_accounts_import_should_expire_sub2api_account_when_token_is_expired() {
let (app, state, pool, _dir, secret_box) =
admin_accounts_test_app("admin-accounts-import-expired-sub2api.sqlite", 117).await;
let expired_access_token = test_jwt_with_exp(
Some("chatgpt-expired-import"),
Some("user-expired-import"),
Some("expired@example.com"),
Some("free"),
1_600_000_000,
);
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/api/admin/accounts/import")
.header("content-type", "application/json")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_accounts_import_expired_sub2api")
.body(Body::from(
json!({
"accounts": [{
"id": "acct_expired_import",
"email": "expired@example.com",
"accountId": "chatgpt-expired-import",
"userId": "user-expired-import",
"label": "expired-label",
"planType": "free",
"token": expired_access_token,
"refreshToken": "expired-refresh",
"status": "active",
"cachedQuota": {
"rate_limit": {
"allowed": true,
"limit_reached": false,
"used_percent": 5
}
}
}]
})
.to_string(),
))
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let body = response_json(response).await;
let stored = SqliteAccountStore::new(pool, secret_box)
.get("acct_expired_import")
.await
.unwrap()
.unwrap();
let auth_status = state.services.admin_accounts.auth_status().await.unwrap();
assert_eq!(status, StatusCode::OK);
assert_eq!(body["data"]["sourceFormat"], "sub2api");
assert_eq!(stored.status, AccountStatus::Expired);
assert_eq!(auth_status.pool.active, 0);
assert_eq!(auth_status.pool.expired, 1);
assert!(auth_status.user.is_none());
assert!(state
.services
.account_pool
.acquire("gpt-5.5", Utc::now())
.await
.is_none());
}
#[tokio::test]
async fn admin_accounts_import_should_require_admin_session_cookie() {
let dir = tempfile::tempdir().unwrap();
@@ -573,7 +573,7 @@ async fn admin_account_refresh_should_update_tokens_and_runtime_pool_without_ret
}
#[tokio::test]
async fn admin_account_refresh_should_mark_invalid_refresh_token_as_expired() {
async fn admin_account_refresh_should_disable_invalid_refresh_token() {
let dir = tempfile::tempdir().unwrap();
let db = dir
.path()
@@ -630,13 +630,13 @@ async fn admin_account_refresh_should_mark_invalid_refresh_token_as_expired() {
assert_eq!(response.status(), StatusCode::OK);
let body = response_json(response).await;
assert_eq!(body["data"]["result"], "dead");
assert_eq!(body["data"]["status"], "expired");
assert_eq!(body["data"]["status"], "disabled");
let stored_status: (String,) = sqlx::query_as("select status from accounts where id = ?")
.bind("acct_refresh_invalid_route")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(stored_status.0, "expired");
assert_eq!(stored_status.0, "disabled");
assert!(state
.services
.account_pool
@@ -531,5 +531,5 @@ async fn admin_accounts_health_check_should_refresh_oauth_without_touching_codex
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(dead_status.0, "expired");
assert_eq!(dead_status.0, "disabled");
}
@@ -20,12 +20,9 @@
"content": "continue"
}
],
"tools": [],
"tool_choice": "auto",
"parallel_tool_calls": true,
"reasoning": null,
"store": false,
"stream": true,
"include": [],
"previous_response_id": "resp_implicit_resume_first"
"previous_response_id": "resp_implicit_resume_first",
"tool_choice": "auto",
"parallel_tool_calls": true
}
+73 -5
View File
@@ -44,11 +44,14 @@ use tokio::{
time::timeout,
};
use tokio_tungstenite::{
accept_async, accept_hdr_async,
accept_hdr_async_with_config,
tungstenite::{
handshake::server::{Request as WsRequest, Response as WsResponse},
Message,
extensions::{compression::deflate::DeflateConfig, ExtensionsConfig},
handshake::server::{Callback, Request as WsRequest, Response as WsResponse},
protocol::WebSocketConfig,
Error as WsError, Message,
},
WebSocketStream,
};
use tower::util::ServiceExt;
use wiremock::{
@@ -58,6 +61,40 @@ use wiremock::{
const TEST_INSTALLATION_ID: &str = "b4f9d503-07b1-457b-a0da-87e6836b1c43";
fn websocket_accept_config() -> WebSocketConfig {
let mut extensions = ExtensionsConfig::default();
extensions.permessage_deflate = Some(DeflateConfig::default());
let mut config = WebSocketConfig::default();
config.extensions = extensions;
config
}
async fn accept_async(stream: TcpStream) -> Result<WebSocketStream<TcpStream>, WsError> {
accept_hdr_async(stream, |_request: &WsRequest, response: WsResponse| {
Ok(response)
})
.await
}
async fn accept_hdr_async<C>(
stream: TcpStream,
callback: C,
) -> Result<WebSocketStream<TcpStream>, WsError>
where
C: Callback + Unpin,
{
accept_hdr_async_with_config(
stream,
move |request: &WsRequest, response: WsResponse| {
let response = callback.on_request(request, response)?;
Ok(response)
},
Some(websocket_accept_config()),
)
.await
}
const CHAT_SUCCESS_SSE: &str = concat!(include_str!("fixtures/chat/success.sse"), "\n");
const RESPONSES_SUCCESS_SSE: &str = concat!(
include_str!("fixtures/responses/http_sse/success.sse"),
@@ -1188,19 +1225,49 @@ async fn response_text(response: axum::response::Response) -> String {
String::from_utf8(bytes.to_vec()).unwrap()
}
async fn wait_for_session_affinity_turn_state(
pool: &SqlitePool,
response_id: &str,
) -> Option<String> {
timeout(StdDuration::from_secs(1), async {
loop {
if let Some(turn_state) = sqlx::query_scalar::<_, Option<String>>(
"select turn_state from session_affinities where response_id = ?",
)
.bind(response_id)
.fetch_optional(pool)
.await
.unwrap()
.flatten()
{
return Some(turn_state);
}
tokio::time::sleep(StdDuration::from_millis(10)).await;
}
})
.await
.ok()
.flatten()
}
struct ResponseEventLog {
level: String,
request_id: Option<String>,
account_id: Option<String>,
route: Option<String>,
status_code: Option<i64>,
metadata_json: String,
}
async fn latest_response_event_log(pool: &SqlitePool) -> ResponseEventLog {
latest_event_log(pool, "v1.response").await
}
async fn latest_event_log(pool: &SqlitePool, kind: &str) -> ResponseEventLog {
let page = SqliteEventLogStore::new(pool.clone())
.list(
EventLogFilter {
kind: Some("v1.response".to_string()),
kind: Some(kind.to_string()),
..EventLogFilter::default()
},
None,
@@ -1212,11 +1279,12 @@ async fn latest_response_event_log(pool: &SqlitePool) -> ResponseEventLog {
.items
.into_iter()
.next()
.expect("expected a v1.response event log");
.unwrap_or_else(|| panic!("expected a {kind} event log"));
ResponseEventLog {
level: event_level_name(event.level).to_string(),
request_id: event.request_id,
account_id: event.account_id,
route: event.route,
status_code: event.status_code,
metadata_json: event.metadata.to_string(),
}
@@ -1,5 +1,15 @@
use super::*;
fn assert_substrings_appear_in_order(haystack: &str, needles: &[&str]) {
let mut cursor = 0;
for needle in needles {
let Some(offset) = haystack[cursor..].find(needle) else {
panic!("expected substring {needle:?} after byte {cursor} in:\n{haystack}");
};
cursor += offset + needle.len();
}
}
#[tokio::test]
async fn chat_completions_should_dispatch_to_codex_and_return_openai_response() {
let server = MockServer::start().await;
@@ -15,7 +25,7 @@ async fn chat_completions_should_dispatch_to_codex_and_return_openai_response()
.mount(&server)
.await;
let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await;
let (app, api_key, pool, _dir) = test_app_with_account_pool_and_logging(server.uri()).await;
let response = app
.oneshot(
Request::builder()
@@ -23,9 +33,10 @@ async fn chat_completions_should_dispatch_to_codex_and_return_openai_response()
.uri("/v1/chat/completions")
.header("authorization", format!("Bearer {api_key}"))
.header("content-type", "application/json")
.header("x-request-id", "req_chat_nonstream_log")
.body(Body::from(
json!({
"model": "gpt-5.5",
"model": "gpt-5.5-high-fast",
"messages": [
{"role": "system", "content": "You are concise."},
{"role": "user", "content": "Say hello"}
@@ -42,10 +53,22 @@ async fn chat_completions_should_dispatch_to_codex_and_return_openai_response()
assert_eq!(status, StatusCode::OK);
assert_eq!(body["object"], "chat.completion");
assert_eq!(body["model"], "gpt-5.5");
assert_eq!(body["model"], "gpt-5.5-high-fast");
assert_eq!(body["choices"][0]["message"]["content"], "hello");
assert_eq!(body["usage"]["prompt_tokens"], 9);
assert_eq!(body["usage"]["completion_tokens"], 3);
let requests = server.received_requests().await.unwrap();
let upstream_body: Value = serde_json::from_slice(&requests[0].body).unwrap();
assert_eq!(upstream_body["model"], "gpt-5.5");
assert_eq!(upstream_body["service_tier"], "priority");
assert_eq!(
upstream_body["reasoning"],
json!({"summary": "auto", "effort": "high"})
);
assert_eq!(
upstream_body["include"],
json!(["reasoning.encrypted_content"])
);
let usage: (i64, i64, i64) = sqlx::query_as(
"select request_count, input_tokens, output_tokens from account_usage where account_id = ?",
)
@@ -54,6 +77,103 @@ async fn chat_completions_should_dispatch_to_codex_and_return_openai_response()
.await
.unwrap();
assert_eq!(usage, (1, 9, 3));
let event = latest_event_log(&pool, "v1.chat").await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert_eq!(event.request_id.as_deref(), Some("req_chat_nonstream_log"));
assert_eq!(event.account_id.as_deref(), Some("acct_chat"));
assert_eq!(event.route.as_deref(), Some("/v1/chat/completions"));
assert_eq!(event.status_code, Some(200));
assert_eq!(metadata["route"], "/v1/chat/completions");
assert_eq!(metadata["apiKind"], "chat");
assert_eq!(metadata["stream"], false);
assert_eq!(metadata["transport"], "http_sse");
assert_eq!(metadata["usage"]["inputTokens"], 9);
assert_eq!(metadata["usage"]["outputTokens"], 3);
}
#[tokio::test]
async fn chat_completions_stream_should_translate_codex_sse_to_openai_chunks() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.and(header("authorization", "Bearer access-secret"))
.and(header("chatgpt-account-id", "chatgpt-account"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(concat!(
"event: response.reasoning_summary_text.delta\n",
"data: {\"delta\":\"I considered the context.\"}\n\n",
"event: response.output_text.delta\n",
"data: {\"delta\":\"hello\"}\n\n",
"event: response.completed\n",
"data: {\"response\":{\"id\":\"resp_1\",\"output_text\":\"hello\",\"usage\":{\"input_tokens\":9,\"output_tokens\":3}}}\n\n",
)),
)
.mount(&server)
.await;
let (app, api_key, pool, _dir) = test_app_with_account_pool_and_logging(server.uri()).await;
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header("authorization", format!("Bearer {api_key}"))
.header("content-type", "application/json")
.header("x-request-id", "req_chat_stream_log")
.body(Body::from(
json!({
"model": "gpt-5.5-high-fast",
"stream": true,
"messages": [
{"role": "user", "content": "Say hello"}
]
})
.to_string(),
))
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let content_type = response
.headers()
.get("content-type")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string();
let body = String::from_utf8(
to_bytes(response.into_body(), usize::MAX)
.await
.unwrap()
.to_vec(),
)
.unwrap();
assert_eq!(status, StatusCode::OK);
assert!(content_type.starts_with("text/event-stream"));
assert!(body.contains("\"object\":\"chat.completion.chunk\""));
assert!(body.contains("\"model\":\"gpt-5.5-high-fast\""));
assert_substrings_appear_in_order(
&body,
&[
"\"delta\":{\"role\":\"assistant\"}",
"\"delta\":{\"reasoning_content\":\"I considered the context.\"}",
"\"delta\":{\"content\":\"hello\"}",
"\"finish_reason\":\"stop\"",
"\"usage\":{\"completion_tokens\":3,\"prompt_tokens\":9,\"total_tokens\":12}",
"data: [DONE]",
],
);
let event = latest_event_log(&pool, "v1.chat").await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert_eq!(event.request_id.as_deref(), Some("req_chat_stream_log"));
assert_eq!(event.route.as_deref(), Some("/v1/chat/completions"));
assert_eq!(metadata["route"], "/v1/chat/completions");
assert_eq!(metadata["apiKind"], "chat");
assert_eq!(metadata["stream"], true);
assert_eq!(metadata["transport"], "http_sse");
}
#[tokio::test]
@@ -80,6 +200,7 @@ async fn chat_completions_should_forward_runtime_installation_id_to_codex() {
.uri("/v1/chat/completions")
.header("authorization", format!("Bearer {api_key}"))
.header("content-type", "application/json")
.header("x-request-id", "req_chat_429_exhausted")
.body(Body::from(
json!({
"model": "gpt-5.5",
@@ -222,7 +343,7 @@ async fn chat_completions_should_return_rate_limit_error_when_429_fallback_is_ex
.mount(&server)
.await;
let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await;
let (app, api_key, pool, _dir) = test_app_with_account_pool_and_logging(server.uri()).await;
let response = app
.oneshot(
Request::builder()
@@ -258,6 +379,26 @@ async fn chat_completions_should_return_rate_limit_error_when_429_fallback_is_ex
assert_eq!(body["error"]["code"], "upstream_error");
assert_eq!(quota_state.0, 1);
assert!(quota_state.1.is_some());
let event = latest_event_log(&pool, "v1.chat").await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert_eq!(event.level, "error");
assert!(event
.request_id
.as_deref()
.is_some_and(|request_id| request_id.starts_with("req_")));
assert_eq!(event.account_id.as_deref(), Some("acct_chat"));
assert_eq!(event.route.as_deref(), Some("/v1/chat/completions"));
assert_eq!(event.status_code, Some(429));
assert_eq!(metadata["route"], "/v1/chat/completions");
assert_eq!(metadata["apiKind"], "chat");
assert_eq!(metadata["stream"], false);
assert_eq!(metadata["transport"], "http_sse");
assert_eq!(metadata["failed"], true);
assert_eq!(metadata["failureClass"], "rate_limited");
assert_eq!(metadata["exhaustedCount"], 1);
assert!(metadata["upstreamError"]
.as_str()
.is_some_and(|error| error.contains("rate limited")));
}
#[tokio::test]
@@ -49,6 +49,50 @@ async fn responses_review_route_should_force_review_subagent_upstream() {
);
}
#[tokio::test]
async fn responses_review_route_should_record_review_route_in_event_log() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.and(header("x-openai-subagent", "review"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(RESPONSES_SUCCESS_SSE),
)
.expect(1)
.mount(&server)
.await;
let (app, api_key, pool, _dir) = test_app_with_account_pool_and_logging(server.uri()).await;
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/responses/review")
.header("authorization", format!("Bearer {api_key}"))
.header("content-type", "application/json")
.header("x-request-id", "req_review_route_log")
.body(Body::from(
json!({
"model": "gpt-5.5",
"input": [],
"stream": false,
"use_websocket": false
})
.to_string(),
))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let event = latest_response_event_log(&pool).await;
assert_eq!(event.request_id.as_deref(), Some("req_review_route_log"));
assert_eq!(event.route.as_deref(), Some("/v1/responses/review"));
}
#[tokio::test]
async fn responses_compact_should_post_json_to_codex_compact_upstream() {
let server = MockServer::start().await;
@@ -59,20 +103,25 @@ async fn responses_compact_should_post_json_to_codex_compact_upstream() {
.and(header("content-type", "application/json"))
.and(header("openai-beta", "responses_websockets=2026-02-06"))
.and(header("x-openai-internal-codex-residency", "us"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"output": [
{
"type": "message",
"role": "assistant",
"content": [{"type": "output_text", "text": "compacted"}]
}
]
})))
.respond_with(
ResponseTemplate::new(200)
.insert_header("x-ratelimit-limit-requests", "33")
.set_body_json(json!({
"output": [
{
"type": "message",
"role": "assistant",
"content": [{"type": "output_text", "text": "compacted"}]
}
],
"usage": {"input_tokens": 11, "output_tokens": 4}
})),
)
.expect(1)
.mount(&server)
.await;
let (app, api_key, _dir) = test_app_with_account(server.uri()).await;
let (app, api_key, pool, _dir) = test_app_with_account_pool_and_logging(server.uri()).await;
let response = app
.oneshot(
Request::builder()
@@ -80,6 +129,7 @@ async fn responses_compact_should_post_json_to_codex_compact_upstream() {
.uri("/v1/responses/compact")
.header("authorization", format!("Bearer {api_key}"))
.header("content-type", "application/json")
.header("x-request-id", "req_compact_success_log")
.body(Body::from(
json!({
"model": "gpt-5.5-fast",
@@ -161,6 +211,19 @@ async fn responses_compact_should_post_json_to_codex_compact_upstream() {
upstream_body["input"][2]["encrypted_content"],
"enc_compact"
);
let event = latest_response_event_log(&pool).await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert_eq!(event.request_id.as_deref(), Some("req_compact_success_log"));
assert_eq!(event.account_id.as_deref(), Some("acct_chat"));
assert_eq!(event.route.as_deref(), Some("/v1/responses/compact"));
assert_eq!(event.status_code, Some(200));
assert_eq!(metadata["route"], "/v1/responses/compact");
assert_eq!(metadata["apiKind"], "responses");
assert_eq!(metadata["stream"], false);
assert_eq!(metadata["compact"], true);
assert_eq!(metadata["usage"]["inputTokens"], 11);
assert_eq!(metadata["usage"]["outputTokens"], 4);
assert_rate_limit_header(&metadata, "x-ratelimit-limit-requests", "33");
}
#[tokio::test]
@@ -179,7 +242,7 @@ async fn responses_compact_should_return_rate_limit_error_when_fallback_is_exhau
.mount(&server)
.await;
let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(server.uri()).await;
let (app, api_key, pool, _dir) = test_app_with_account_pool_and_logging(server.uri()).await;
let response = app
.oneshot(
Request::builder()
@@ -187,6 +250,7 @@ async fn responses_compact_should_return_rate_limit_error_when_fallback_is_exhau
.uri("/v1/responses/compact")
.header("authorization", format!("Bearer {api_key}"))
.header("content-type", "application/json")
.header("x-request-id", "req_compact_rate_limited_log")
.body(Body::from(
json!({
"model": "gpt-5.5",
@@ -208,6 +272,27 @@ async fn responses_compact_should_return_rate_limit_error_when_fallback_is_exhau
assert_eq!(body["error"]["code"], "rate_limit_exceeded");
assert!(message.contains("All accounts exhausted (1 rate-limited)"));
assert!(message.contains("compact quota reached"));
let event = latest_response_event_log(&pool).await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert_eq!(
event.request_id.as_deref(),
Some("req_compact_rate_limited_log")
);
assert_eq!(event.account_id.as_deref(), Some("acct_chat"));
assert_eq!(event.route.as_deref(), Some("/v1/responses/compact"));
assert_eq!(event.status_code, Some(429));
assert_eq!(event.level, "error");
assert_eq!(metadata["route"], "/v1/responses/compact");
assert_eq!(metadata["apiKind"], "responses");
assert_eq!(metadata["stream"], false);
assert_eq!(metadata["compact"], true);
assert_eq!(metadata["transport"], "http");
assert_eq!(metadata["failed"], true);
assert_eq!(metadata["failureClass"], "rate_limited");
assert_eq!(metadata["exhaustedCount"], 1);
assert!(metadata["upstreamError"]
.as_str()
.is_some_and(|value| value.contains("compact quota reached")));
}
#[tokio::test]
@@ -206,7 +206,7 @@ async fn responses_should_use_websocket_upstream_by_default_while_serving_sse()
assert!(body.contains("\"id\":\"resp_ws_default\""));
assert_eq!(captured.payload["type"], "response.create");
assert_eq!(captured.payload["model"], "gpt-5.5");
assert_eq!(captured.payload["generate"], false);
assert!(captured.payload.get("generate").is_none());
assert!(captured.payload.get("previous_response_id").is_none());
assert!(captured.payload["prompt_cache_key"]
.as_str()
@@ -217,6 +217,44 @@ async fn responses_should_use_websocket_upstream_by_default_while_serving_sse()
.is_some_and(|value| value.starts_with("cp_") && value.ends_with(":0")));
}
#[tokio::test]
async fn responses_non_stream_should_record_websocket_transport_metadata() {
let (base_url, upstream) =
spawn_single_websocket_completed_upstream("resp_ws_non_stream_log").await;
let (app, api_key, pool, _dir) = test_app_with_account_pool_and_logging(base_url).await;
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/responses")
.header("authorization", format!("Bearer {api_key}"))
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "gpt-5.5",
"input": [],
"stream": false,
"generate": false
})
.to_string(),
))
.unwrap(),
)
.await
.unwrap();
let body = response_json(response).await;
let captured = upstream.await.unwrap();
assert_eq!(body["id"], "resp_ws_non_stream_log");
assert_eq!(captured.payload["type"], "response.create");
let event = latest_response_event_log(&pool).await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert_eq!(metadata["stream"], false);
assert_eq!(metadata["transport"], "websocket");
}
#[tokio::test]
async fn responses_should_ignore_camel_case_use_websocket_field() {
let (base_url, upstream) =
@@ -396,12 +434,9 @@ async fn responses_should_forward_parity_fields_context_headers_and_account_scop
.header("authorization", format!("Bearer {api_key}"))
.header("content-type", "application/json")
.header("x-codex-turn-state", "turn-header")
.header("x-codex-turn-metadata", "meta-header")
.header("x-codex-beta-features", "beta-header")
.header("x-responsesapi-include-timing-metrics", "false")
.header("version", "header-version")
.header("x-codex-window-id", "window-header")
.header("x-codex-parent-thread-id", "parent-header")
.header("x-openai-subagent", "review")
.body(Body::from(
json!({
@@ -410,14 +445,17 @@ async fn responses_should_forward_parity_fields_context_headers_and_account_scop
"use_websocket": false,
"input": [],
"prompt_cache_key": "pcache",
"client_metadata": {"safe": "yes", "drop": 42},
"client_metadata": {
"safe": "yes",
"drop": 42,
"x-codex-turn-metadata": "meta-metadata",
"x-codex-window-id": "window-metadata",
"x-codex-parent-thread-id": "parent-metadata"
},
"turnState": "turn-body",
"turnMetadata": "meta-body",
"betaFeatures": "beta-body",
"includeTimingMetrics": "true",
"version": "2026-06-12",
"codexWindowId": "window-body",
"parentThreadId": "parent-body",
"use_websocket": false
})
.to_string(),
@@ -437,7 +475,8 @@ async fn responses_should_forward_parity_fields_context_headers_and_account_scop
.find(|request| request.url.path() == "/codex/responses")
.expect("responses upstream request should be sent");
let upstream_body: Value = serde_json::from_slice(&upstream.body).unwrap();
let identity = build_conversation_identity(Some("pcache"), Some("window-body"), "acct_chat");
let identity =
build_conversation_identity(Some("pcache"), Some("window-metadata"), "acct_chat");
let conversation_id = identity
.conversation_id
.as_deref()
@@ -479,16 +518,19 @@ async fn responses_should_forward_parity_fields_context_headers_and_account_scop
);
assert_eq!(
upstream_body["client_metadata"]["x-codex-turn-metadata"],
"meta-body"
"meta-metadata"
);
assert_eq!(
upstream_body["client_metadata"]["x-codex-parent-thread-id"],
"parent-body"
"parent-metadata"
);
assert_eq!(upstream_header("session_id"), Some(conversation_id));
assert_eq!(upstream_header("x-codex-window-id"), Some(window_id));
assert_eq!(upstream_header("x-codex-turn-state"), Some("turn-body"));
assert_eq!(upstream_header("x-codex-turn-metadata"), Some("meta-body"));
assert_eq!(
upstream_header("x-codex-turn-metadata"),
Some("meta-metadata")
);
assert_eq!(upstream_header("x-codex-beta-features"), Some("beta-body"));
assert_eq!(
upstream_header("x-responsesapi-include-timing-metrics"),
@@ -497,7 +539,7 @@ async fn responses_should_forward_parity_fields_context_headers_and_account_scop
assert_eq!(upstream_header("version"), Some("2026-06-12"));
assert_eq!(
upstream_header("x-codex-parent-thread-id"),
Some("parent-body")
Some("parent-metadata")
);
assert_eq!(upstream_header("x-openai-subagent"), Some("review"));
for local_field in [
@@ -2015,7 +2015,7 @@ async fn responses_should_return_rate_limit_error_when_429_fallback_is_exhausted
.mount(&server)
.await;
let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await;
let (app, api_key, pool, _dir) = test_app_with_account_pool_and_logging(server.uri()).await;
let response = app
.oneshot(
Request::builder()
@@ -2023,6 +2023,7 @@ async fn responses_should_return_rate_limit_error_when_429_fallback_is_exhausted
.uri("/v1/responses")
.header("authorization", format!("Bearer {api_key}"))
.header("content-type", "application/json")
.header("x-request-id", "req_responses_429_exhausted")
.body(Body::from(
json!({
"model": "gpt-5.5",
@@ -2053,6 +2054,23 @@ async fn responses_should_return_rate_limit_error_when_429_fallback_is_exhausted
assert_eq!(body["error"]["code"], "upstream_error");
assert_eq!(quota_state.0, 1);
assert!(quota_state.1.is_some());
let event = latest_response_event_log(&pool).await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert_eq!(event.level, "error");
assert_eq!(
event.request_id.as_deref(),
Some("req_responses_429_exhausted")
);
assert_eq!(event.account_id.as_deref(), Some("acct_chat"));
assert_eq!(event.route.as_deref(), Some("/v1/responses"));
assert_eq!(event.status_code, Some(429));
assert_eq!(metadata["stream"], false);
assert_eq!(metadata["transport"], "http_sse");
assert_eq!(metadata["failureClass"], "rate_limited");
assert_eq!(metadata["exhaustedCount"], 1);
assert!(metadata["upstreamError"]
.as_str()
.is_some_and(|error| error.contains("rate limited")));
}
#[tokio::test]
@@ -2068,7 +2086,7 @@ async fn responses_stream_should_return_rate_limit_error_when_429_fallback_is_ex
.mount(&server)
.await;
let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await;
let (app, api_key, pool, _dir) = test_app_with_account_pool_and_logging(server.uri()).await;
let response = app
.oneshot(
Request::builder()
@@ -2096,6 +2114,12 @@ async fn responses_stream_should_return_rate_limit_error_when_429_fallback_is_ex
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string();
let response_request_id = response
.headers()
.get("x-request-id")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string();
let body = response_text(response).await;
let quota_state: (i64, Option<String>) = sqlx::query_as(
"select quota_limit_reached, quota_cooldown_until from accounts where id = ?",
@@ -2113,6 +2137,24 @@ async fn responses_stream_should_return_rate_limit_error_when_429_fallback_is_ex
assert!(body.contains("rate limited"));
assert_eq!(quota_state.0, 1);
assert!(quota_state.1.is_some());
assert!(!response_request_id.is_empty());
let event = latest_response_event_log(&pool).await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert_eq!(event.level, "error");
assert_eq!(
event.request_id.as_deref(),
Some(response_request_id.as_str())
);
assert_eq!(event.account_id.as_deref(), Some("acct_chat"));
assert_eq!(event.route.as_deref(), Some("/v1/responses"));
assert_eq!(event.status_code, Some(429));
assert_eq!(metadata["stream"], true);
assert_eq!(metadata["transport"], "http_sse");
assert_eq!(metadata["failureClass"], "rate_limited");
assert_eq!(metadata["exhaustedCount"], 1);
assert!(metadata["upstreamError"]
.as_str()
.is_some_and(|error| error.contains("rate limited")));
}
#[tokio::test]
@@ -462,7 +462,7 @@ async fn responses_websocket_stream_should_record_metadata_turn_state_for_contin
let second_headers = second_headers.lock().unwrap().clone();
(first_payload, second_payload, second_headers)
});
let (app, api_key, _pool, _dir) = test_app_with_account_pool_config(base_url, |config| {
let (app, api_key, pool, _dir) = test_app_with_account_pool_config(base_url, |config| {
config.ws_pool.enabled = false;
})
.await;
@@ -492,6 +492,10 @@ async fn responses_websocket_stream_should_record_metadata_turn_state_for_contin
let first_body = response_text(first_response).await;
assert!(first_body.contains("\"id\":\"resp_metadata_turn_state\""));
assert!(!first_body.contains("response.metadata"));
assert_eq!(
wait_for_session_affinity_turn_state(&pool, "resp_metadata_turn_state").await,
Some("turn-from-metadata".to_string())
);
let second_response = app
.oneshot(
@@ -457,6 +457,7 @@ async fn responses_stream_should_record_event_log_after_completed_stream() {
assert_eq!(event.account_id.as_deref(), Some("acct_chat"));
assert_eq!(event.status_code, Some(200));
assert_eq!(metadata["stream"], true);
assert_eq!(metadata["transport"], "http_sse");
assert_eq!(metadata["completed"], true);
assert_eq!(metadata["responseId"], "resp_stream_usage");
assert_eq!(metadata["usage"]["inputTokens"], 3);
+847
View File
File diff suppressed because one or more lines are too long
File diff suppressed because it is too large Load Diff
+260
View File
@@ -0,0 +1,260 @@
# 风控链路审计
审计日期:2026-06-20
## 范围
本次只审计 OpenAI / Codex 上游真实请求链路相关的风控实现,不纳入通用 proxy、VPN 或其它 provider。
参照对象:
- TS 仓库:`/home/zyy/桌面/Codes/codex-proxy`
- 重点版本:`v2.0.78-beta.15c49d6`
- 后续风控增强参照:当前 TS HEAD `v2.0.80`
- Rust 仓库:当前工作区 `/home/zyy/Codes/codex-proxy-rs`
## 对齐原则
本轮严格以 TS 实现和 OpenAI 原版实现为依据,不再从单次抓包或经验规则直接推导代码行为。
- TS 当前 HEAD 是 `v2.0.80`,已经包含 `v2.0.78-beta.15c49d6` 之后的 Cloudflare cooldown、上游 session identity 隔离、reasoning replay 等风控增强;没有特别说明时,优先对齐当前 TS HEAD。
- `docs/openai-res.txt` 是真实链路验收证据,不是直接改代码的唯一来源。抓包和 TS 行为冲突时,先记录差异,再用 TS / OpenAI 原版 / 新真实链路三方确认。
- TS 没有实现的增强,不写成“已对齐”。例如 active quota refresher 的 30 分钟节流在 TS 中是进程内 `Map`,不是持久化字段;如果后续要做跨重启持久化节流,应标为 Rust 侧增强,而不是 TS 对齐项。
- 请求头、payload、账号状态流转都先查 TS 源码和调用点;缺少明确来源的字段不直接加入真实请求链路。
相关历史:
- `981f94c fix(auth): add cascading ban defense on session affinity degradation`
- `c559141 fix(auth): add active quota refresher and drift defense verification`
- `3178d21 fix(auth): address PR #624 review findings`
- `d19ed47 fix(auth): resolve TypeScript null-narrowing in drift-defense loop`
- `abc5b98 fix(auth): only strip session on ban, not quota rotation`
- `d4ec5a7 Merge pull request #624 from icebear0828/fix/cascading-ban-defense`
- 后续增强:`b079af1` Cloudflare challenge cooldown、`bc4b13b` 上游 session identity 隔离、`9eb8612` reasoning replay cache。
## 结论
当前 Rust 已经覆盖了不少后续 TS 增强,尤其是上游 identity 隔离、reasoning replay、WebSocket required/fallback 安全、Cloudflare path-block 处理、禁用/封禁状态持久化、`next_refresh_at` 持久化。
本次已补齐 TS PR #624 中与请求链路相关的 P0 缺口,并补齐后续 TS 风控增强里的 refresh token 安全重试与 Cloudflare challenge cooldown:
1. `quota_verify_required` 已接入请求发送前 `/usage` 漂移验证,chat / responses / stream / compact 都会在真实业务请求前验证 dirty 账号。
2. 后台 `QuotaRefreshTask` 已同时扫描 `quota_limit_reached` 和 `quota_verify_required`。
3. session affinity 降级时,只有 preferred 账号为 `Banned` / `Disabled` 才会在发送前剥离历史;quota rotation 不预剥离。
4. refresh token transport error 已拆成可安全重试的 pre-flight 失败和不可安全重试的 ambiguous 失败;每次真正刷新前都会重读 DB refresh token,避免复用 stale RT。
5. 401/token invalid 与 refresh `invalid_grant` 已拆成两阶段:请求链路普通 token invalid 先落 `Expired`,允许 RT 自救;一旦 refresh token 本身确认 `invalid_grant` / invalid / expired,则落 `Disabled` 并停止后续调度。TS 当前把 refresh `invalid_grant` 也落 `Expired`,但 Rust 有持久化周期调度,照搬会反复刷新坏 RT。
6. Cloudflare challenge cooldown 已按 10/30/90/120 秒递增,1 小时无 challenge 后重置,并继续持久化 `cloudflare_cooldown_until`。
仍需继续做真实链路层面的验证:HTTP SSE header、TLS 指纹、真实上游是否持续返回可复用的 `cf_clearance`,以及 OpenAI 原版客户端后续风控差异。cookie 自动捕获范围和捕获时机已先按 TS 对齐。
## PR #624 四个修复逐项对照
这组修复都在 TS `v2.0.78-beta.15c49d6` 之前合入,属于当前必须对齐的风控基线。
| TS 修复 | TS 行为 | Rust 当前状态 | 判断 |
| --- | --- | --- | --- |
| `fix(auth): add active quota refresher and drift defense verification` | `ActiveQuotaRefresher` 每 15 分钟扫描 active 账号,命中 `limit_reached` 或 `quotaVerifyRequired` 都主动拉 `/usage`;请求链路遇 dirty 账号先做 `/usage` 校验,仍限流则排除并重新 acquire,最多 5 次 | Rust 已在 `QuotaRefreshTask` 扫描 locked/dirty;`verify_acquired_quota_if_required` 已接入 chat / responses / stream / compact;成功验证会通过 `apply_quota_snapshot` 同步 DB 和内存池 | 已对齐 |
| `fix(auth): add cascading ban defense on session affinity degradation` | 当 `previous_response_id` 指向的 preferred 账号已经 `banned/disabled`,但本次 acquire 降级到其它账号时,发送前剥离 `previous_response_id` 和 `turnState`,避免封禁关联扩散 | Rust 已在 Responses complete/stream 发送前检查 preferred 账号状态;仅 `Banned` / `Disabled` 且换号时剥离 `previous_response_id` / `turn_state` 并忘记 affinity;不额外清理 `turn_metadata` / `prompt_cache_key` | 已对齐 |
| `fix(auth): only strip session on ban, not quota rotation` | 只在 `banned/disabled` 这种 ban-risk 状态下预剥离;quota rotation 不预剥离,交给 `previous_response_not_found` 的普通 retry | Rust 的预剥离 guard 只匹配 `AccountStatus::Banned | AccountStatus::Disabled`;dirty quota / rate-limit / quota exhausted 只走排除重调度或原有 history recovery | 已对齐 |
| `fix(auth): resolve TypeScript null-narrowing in drift-defense loop` | 这是 TS 类型收窄修复,保证 drift-defense loop 中 acquire 变量在重新赋值后类型安全 | Rust 不存在 TS null-narrowing 问题;业务语义已由 `QuotaVerificationDecision` 显式区分 ready、retry、max attempts | 类型问题不适用,业务已覆盖 |
## 覆盖矩阵
| 能力 | TS 参照 | Rust 当前状态 | 判断 |
| --- | --- | --- | --- |
| Active quota refresher | `src/auth/active-quota-refresher.ts:63-99` 同时扫 `limit_reached` 和 `quotaVerifyRequired` | `crates/runtime/src/tasks/quota_refresh.rs` 同时扫 `quota_limit_reached || quota_verify_required`,成功后应用完整 quota state | 覆盖 |
| 请求前 quota 漂移验证 | `src/routes/shared/proxy-handler.ts:92-139`,最多 5 次 `/usage` 校验并排除仍限流账号 | `crates/runtime/src/services.rs` 的 chat / responses / stream / compact acquire 后先调用 `verify_acquired_quota_if_required`;仍限流则 release + exclude + reacquire | 覆盖 |
| quota 冷却后 dirty 标记 | `account-registry.ts` 离线窗口 rollover 后 dirty | `crates/core/src/accounts/pool.rs:670-689` 冷却或窗口过期后置 `quota_verify_required=true` | 部分覆盖 |
| ban-risk session 剥离 | `src/routes/shared/proxy-retry-recovery.ts:110-152` 仅 preferred 为 `banned/disabled` 且换号时剥离 | `ResponseDispatchService::apply_cascading_ban_defense` 仅在 preferred 为 `Banned/Disabled` 且换号时剥离;quota rotation 不剥离 | 覆盖 |
| same-account history recovery | TS 对 `previous_response_not_found` / unanswered function call strip 后同账号重试 | `crates/runtime/src/services.rs:2040-2048`、`2580-2624`、`3293-3300` 可错误后 strip/retry | 覆盖 |
| 上游 session identity 隔离 | `src/proxy/codex-api.ts:108-139`、`328-358`、`384-413` | `crates/core/src/gateway/conversation.rs:27-46`、`145-168`;`crates/runtime/src/services.rs:2863-2885`、`2901-2923`;`crates/adapters/src/codex/client.rs:1064-1110` | 覆盖 |
| reasoning replay cache | TS `9eb8612` | `crates/core/src/serving/reasoning_replay.rs:8-10`、`44-108`、`173-238`;记录/驱逐在 `crates/runtime/src/services.rs:3812-3895` | 覆盖 |
| WebSocket previous_response 安全 | TS `PreviousResponseWebSocketError` 禁止降级 | `crates/core/src/serving/responses.rs:21-39` + `crates/adapters/src/codex/client.rs:390-448`,`previous_response_id` 为 WebSocketRequired | 覆盖 |
| WS 请求头和 deflate | TS 使用 Node `ws` 完整握手;项目只传业务 headers,`ws@8.19.0` 客户端默认 `perMessageDeflate=true` | Rust 使用 OpenAI fork tungstenite;`crates/adapters/src/codex/websocket/connect.rs:377-399` 构造 opening 审计快照,`497-503` 启用 permessage-deflate | 覆盖,仍需真实抓包持续比对 |
| Cloudflare challenge cooldown | TS `src/auth/cf-challenge-cooldown.ts:1-56` 递增 10/30/90/120 秒 | `CloudflareChallengeCooldownTracker` 按 10/30/90/120 秒递增并 1 小时 stale reset;runtime 持久化 `cloudflare_cooldown_until`,成功响应后清理 cooldown state | 覆盖 |
| Cloudflare path-block | TS 空 body 404 清 cookie,累计 3 次 disable | `crates/runtime/src/services.rs:3095-3100` 识别空 404;`878-899` 清 cookie / 计数 / disable;`crates/core/src/accounts/cloudflare.rs:8-66` 3 次阈值、1 小时 stale | 覆盖 |
| cookie 自动捕获 | TS `cookie-jar.ts` 只允许自动捕获 `cf_clearance`;`codex-api.ts` 在 HTTP/compact transport response 后、状态码判断前 capture,WS 升级响应 header 也会回传 | Rust `SqliteCookieStore::capture_set_cookie` 仍只持久化 `cf_clearance`;HTTP/compact/stream/WS 的成功响应和 upstream error 都会在 runtime 消费 `set_cookie_headers` | 已对齐 |
| refresh disabled/banned skip | TS refresh scheduler 跳过 disabled/banned | `crates/runtime/src/tasks/token_refresh.rs:220-229` 跳过 disabled/banned | 覆盖 |
| `next_refresh_at` 持久化 | TS 持久化下次刷新日期 | `crates/runtime/src/tasks/token_refresh.rs:315-319`、`388-397`、`513-541` | 覆盖 |
| 永久刷新失败确认 | TS 永久失败 2 次确认 | `crates/runtime/src/tasks/token_refresh.rs:44`、`589-600` | 覆盖 |
| token invalid 状态映射 | TS 请求链路 401 非 deactivated 落 `expired`;401 deactivated 落 `banned` 但返回仍是 401;refresh `invalid_grant` 连续确认后落 `expired` | Rust 请求链路仍将 `token_revoked` / `token_invalid` 落 `Expired`;refresh token 确认 `InvalidGrant` 后落 `Disabled`,避免持久化周期调度反复刷新坏 RT;`account_deactivated` / `banned` 仍为 `Banned`,并保留上游触发状态码返回客户端 | 有证据偏离 TS |
| refresh token 安全重试 | TS `src/tls/direct-fallback.ts:52-75` 只重试安全 pre-flight 错误 | `RefreshFailure::RetryableTransport` 才重试;`RefreshFailure::Transport` 直接进入 recovery delay,不复用 RT 重试 | 覆盖 |
| stale RT / 跨进程刷新防护 | TS refresh 前比对磁盘 RT,发现不同则同步不消费旧 RT | `TokenRefreshTask::prepare_refresh_attempt` 每次真正刷新前重读 DB;refresh token 与扫描快照不一致时跳过,不消费旧 RT | 覆盖 |
| stream 失败可观测 | TS `recordStreamCloseEvent` 会把 stream 异常携带 requestId、path、model、account 写入排查日志 | Rust 不迁移 TS 旧日志体系;Responses stream 启动阶段终态失败写入现有 `event_logs`,包含 `stream=true`、`transport`、`failureClass`、`exhaustedCount`、`upstreamError`;live stream completed/failed 继续走已有事件 | 覆盖 |
## 关键发现
### 1. quota 漂移防御已补齐
TS 在请求链路里对 `quotaVerifyRequired` 做前置 `/usage` 校验:
- 最多 `MAX_VERIFY_ATTEMPTS = 5`,避免单请求放大 `/usage`。
- 校验后仍 `limit_reached` 的账号会 release,加入 `verifiedExcludeIds`,重新 acquire。
- 校验失败不会清 dirty flag,交给下一次请求或后台 refresher。
Rust 当前具备状态字段和发送前验证:
- `accounts.quota_verify_required` 已在 schema 中持久化:`crates/platform/src/storage/schema.sql:59`
- 账号模型有字段:`crates/core/src/accounts/model.rs:51`
- SQLite 可读写该字段:`crates/adapters/src/sqlite/accounts.rs:162`、`517`、`536`、`547`
- quota 冷却过期后会设置 dirty:`crates/core/src/accounts/pool.rs:685-689`
- `crates/core/src/serving/quota.rs` 提供主 rate-limit 的触顶、reset_at、窗口解析。
- `crates/adapters/src/sqlite/accounts.rs` 新增 `apply_quota_snapshot`,成功验证会清 `quota_verify_required`,并按 `/usage` 结果同步 `quota_limit_reached` / `quota_cooldown_until`。
- `crates/core/src/accounts/pool.rs` 新增 `apply_quota_state`,请求前验证会同步内存账号池。
- `crates/runtime/src/services.rs` 的 `verify_acquired_quota_if_required` 在业务请求发送前处理 dirty 账号,最多 5 次;仍限流则 release、exclude、reacquire;`/usage` 拉取失败会保留 dirty 并继续原链路。
- `crates/runtime/src/tasks/quota_refresh.rs` 后台候选已扩展为 `quota_limit_reached || quota_verify_required`。
TS 复核点:
- TS `ActiveQuotaRefresher` 每 15 分钟 jitter 调度,per-account 最小 30 分钟刷新间隔保存在进程内 `lastRefreshedAt: Map<string, number>`。
- TS 没有把 active quota refresher 的最小刷新间隔持久化到账号数据。
- TS 的 token refresh scheduler 另有持久化 `next_refresh_at` 语义;它约束 access token 自动刷新,不等同于 active quota refresher 的 `/usage` 主动校验节流。
注意:后台 `QuotaRefreshTask` 当前仍只持有 SQLite store,不直接刷新运行时内存池。真实发送链路已经通过请求前验证同步 DB 和内存;如果需要后台刷新后立即让内存池解锁,需要后续把 `RuntimeAccountPoolService` 注入 task。
### 2. cascading-ban defense 已补齐
TS 的语义很明确:
- 如果 preferred account 因 session affinity 指向旧账号,但 acquire 实际换到新账号;
- 且 preferred account 状态是 `banned` 或 `disabled`;
- 才剥离 `previous_response_id` / `turnState`,避免把封禁风险关联到新账号;
- quota rotation 不剥离,由普通 `previous_response_not_found` retry 处理。
Rust 目前已有并已补齐:
- `previous_response_id` 查 preferred account:`crates/runtime/src/services.rs:1849-1858`
- 错误后 history recovery:`crates/runtime/src/services.rs:1860-1875`
- `previous_response_not_found` / `invalid_encrypted_content` 后 recovery:`crates/runtime/src/services.rs:2040-2048`、`2580-2624`
- 实际剥离函数:`crates/runtime/src/services.rs` 的 `strip_request_history`
- 新增 `ResponseDispatchService::apply_cascading_ban_defense`,在 Responses complete/stream acquire 后、发送前执行。
- 仅当 preferred 账号状态为 `Banned` 或 `Disabled`,且 acquired 账号不同,才剥离 `previous_response_id` / `turn_state` 并忘记 affinity;`turn_metadata` 和 `prompt_cache_key` 保留,和 TS `applyCascadingBanDefense` / `applyProxyRetryRecoveryDecision` 一致。
- `QuotaExhausted` / rate-limit / dirty quota rotation 不做预剥离,继续保留现有错误后 history recovery。
- 普通 history recovery 的隐式续链回滚也按 TS `restoreImplicitResumeRequestState` + `applyProxyRetryRecoveryDecision` 对齐:恢复原请求后只清 `previous_response_id` / `turn_state`,不清 `turn_metadata`。
### 3. refresh token 重试策略已收窄
Rust 当前安全点:
- disabled/banned 不参与调度:`crates/runtime/src/tasks/token_refresh.rs:220-229`
- `next_refresh_at` 未来时间会跳过:`crates/runtime/src/tasks/token_refresh.rs:315-319`
- token 更新后会持久化下一次刷新时间:`crates/runtime/src/tasks/token_refresh.rs:581-586`
- 永久失败需达到阈值 2:`crates/runtime/src/tasks/token_refresh.rs:44`、`589-600`
- 请求链路 `token_revoked` / `token_invalid` 先归类为 `expired`;refresh token 自身确认 `invalid_grant` / `invalid_token` / `access_denied` / `refresh_token_expired` 后归类为 `disabled`;`refresh_token_reused` / `account has been deactivated` / `banned` 归类为 `banned`。
- refresh lease 防止同一时刻多进程刷新同一账号:`crates/runtime/src/tasks/token_refresh.rs`
- `RefreshFailure::RetryableTransport` 表示 DNS、连接拒绝、网络不可达、TLS handshake 等 pre-flight 错误,允许同 RT 重试。
- `RefreshFailure::Transport` 表示读响应失败、body parse 失败、未知 HTTP transport 等 ambiguous 状态,不再复用同 RT 重试,直接进入 recovery delay。
- `OpenAiOAuthClient::refresh` 只在 `.send()` 的明确 pre-flight 错误上返回 `RetryableTransport`;读 body 和解析 token 失败都返回普通 `Transport`。
- `TokenRefreshTask::prepare_refresh_attempt` 每次真正调用 refresher 前重读 DB refresh token;如果 DB token 与本次扫描快照不同,直接跳过,避免消费 stale RT。
### 4. Cloudflare challenge cooldown 已覆盖
Rust 已有并已补齐:
- 403 且包含 challenge 信号时识别为 Cloudflare challenge:`crates/runtime/src/services.rs:3087-3093`
- challenge 后删除账号 cookie 并设置 cooldown:`CloudflareRecovery::apply_challenge`
- 账号池 acquire 会跳过 cooldown 中账号:`crates/core/src/accounts/pool.rs:535-549`
- cooldown 持久化:`crates/adapters/src/sqlite/accounts.rs:1375-1389`
- `CloudflareChallengeCooldownTracker` 记录 per-account count 和 updated_at。
- backoff 与 TS 对齐:10 / 30 / 90 / 120 秒,超过 120 秒后封顶。
- 1 小时无 challenge 后重置计数。
- 成功响应后通过 `reset_account_recovery` 清理 challenge cooldown 和 path-block state。
### 5. Cloudflare path-block 基本覆盖
Rust 当前实现与 TS 后续增强基本对齐:
- 空 body 404 识别为 path-block:`crates/runtime/src/services.rs:3095-3100`
- path-block 后清 cookie、计数、三次后 disabled:`crates/runtime/src/services.rs:878-899`
- 阈值 3、1 小时 stale:`crates/core/src/accounts/cloudflare.rs:8-66`
- 成功请求后 reset path-block:`crates/runtime/src/services.rs:2225`、`2369`、`3713-3714`
- cookie 只捕获 `cf_clearance`,请求时按 domain/path 匹配:`crates/adapters/src/sqlite/cookies.rs:11-13`、`95-147`
- HTTP/SSE、compact、stream 和 WebSocket 升级响应带回的 `Set-Cookie` 已在 runtime 层统一 capture;HTTP/compact 非成功状态也会通过 `CodexClientError::Upstream.set_cookie_headers` 传回后 capture,时机与 TS “收到 transport response 先 capture,再分类状态码”一致。
这个方向可以保留。后续真实链路只观察 OpenAI 是否稳定返回可复用的 `cf_clearance`;不因一次响应出现其它 cookie 就直接扩大自动捕获范围。
### 6. 上游身份隔离已覆盖
TS 后续 `bc4b13b` 的核心是:客户端传来的 `prompt_cache_key` / window id 不能跨账号原样复用,必须加 account scope hash。
Rust 当前已实现同等机制:
- `build_conversation_identity` 将 conversation/window 变成 `cp_` / `cw_` hash:`crates/core/src/gateway/conversation.rs:27-46`、`145-168`
- Responses 非流式和流式请求都构造 account-scoped identity:`crates/runtime/src/services.rs:2863-2885`、`2901-2923`
- body 中 `prompt_cache_key` 会替换成 account-scoped `session_id`:`crates/adapters/src/codex/client.rs:1064-1073`
- `client_metadata[x-codex-window-id]` 会写 account-scoped window id:`crates/adapters/src/codex/client.rs:1076-1110`
- OpenAI Responses 翻译层已按 TS `firstRequestString` 从 `client_metadata[x-codex-window-id]` 兜底读取客户端 window id,因此 metadata-only 请求也会进入 account-scoped identity。
- WS header 当前按 TS 的统一 headers 投影:过滤 HTTP/SSE 专用 `content-type`、`accept`,保留 `session_id`、fingerprint 默认头、cookie、`x-openai-internal-codex-residency`、`x-codex-installation-id`、`x-codex-turn-state`、`x-codex-turn-metadata`、`x-codex-beta-features`、`x-responsesapi-include-timing-metrics`、`version`、`x-codex-parent-thread-id`、`x-openai-subagent` 等条件头,见 `crates/adapters/src/codex/client.rs:1192-1228`
这个能力当前应保留,不要回退到原始客户端 ID。
### 7. reasoning replay 已覆盖
Rust 已有:
- TTL 55 分钟、最多 512 条、单条 256KB、总 4MB:`crates/core/src/serving/reasoning_replay.rs:8-10`
- lookup 按 `account_id + conversation_id + variant_hash + response_id` 隔离:`crates/core/src/serving/reasoning_replay.rs:88-108`
- 只保存 reasoning encrypted_content 和 function_call replay 项:`crates/core/src/serving/reasoning_replay.rs:173-238`
- completed 后记录:`crates/runtime/src/services.rs:3812-3873`
- `invalid_encrypted_content` 后驱逐:`crates/runtime/src/services.rs:3876-3905`
这个与 TS `9eb8612` 的方向一致。
### 8. WebSocket 和 fallback 安全已基本覆盖
Rust 当前:
- 使用 OpenAI fork 的 `tokio-tungstenite` / `tungstenite-rs`:`Cargo.toml:52-53`
- WS opening header 顺序明确,包含 `permessage-deflate; client_max_window_bits`:`crates/adapters/src/codex/websocket/connect.rs:377-399`
- fork tungstenite 配置 deflate:`crates/adapters/src/codex/websocket/connect.rs:497-503`
- receive idle timeout 20 秒:`crates/adapters/src/codex/websocket/connect.rs:44`、`769-775`
- WS pool ping 间隔 25 秒、timeout 5 秒:`crates/adapters/src/codex/websocket/pool.rs:16-17`
- `previous_response_id` 强制 WebSocket,不允许 HTTP SSE fallback:`crates/core/src/serving/responses.rs:21-39`
TS 对照:
- `src/proxy/ws-transport.ts` 的 `buildWsConstructorOpts` 只设置 `{ headers }` 和可选 agent,没有手写 `Sec-WebSocket-Extensions`。
- `ws@8.19.0` 客户端默认 `perMessageDeflate=true`,真实 client offer 由库生成。
- `ws@8.19.0` 的 `WebSocketServer` 默认 `perMessageDeflate=false`,旧 TS 测试 helper 默认值不能作为真实客户端不开 deflate 的证据。
仍需注意:
- HTTP SSE 仍是 `reqwest` + `rustls`,不是 TS 的 curl/libcurl/native transport 指纹。
- TLS 指纹不能仅靠 header 判断,必须用真实抓包或服务端反馈继续验证。
- 不应在文档或日志中暴露 token、账号邮箱、refresh token。
### 9. 请求头现状
WS 链路当前先与 TS 发送语义保持一致:
- opening 固定头:`Host`、`Connection`、`Upgrade`、`Sec-WebSocket-Version`、`Sec-WebSocket-Key`
- 业务/指纹头:来自统一 Codex header map,过滤 `content-type`、`accept` 后保留;包含默认 fingerprint headers、`authorization`、`chatgpt-account-id`、`originator`、`openai-beta`、`x-openai-internal-codex-residency`、`x-client-request-id`、`x-codex-installation-id`、`session_id`、Codex context 条件头
- 扩展头:`sec-websocket-extensions: permessage-deflate; client_max_window_bits`
HTTP SSE 链路当前会带:
- base headers:`authorization`、`chatgpt-account-id`、`originator`、`user-agent`、`sec-ch-ua` 和 fingerprint 默认头。
- response headers:`content-type`、`accept: text/event-stream`、`openai-beta`、`x-openai-internal-codex-residency`、`x-client-request-id`、`x-codex-installation-id`、`session_id`、`x-codex-window-id`、`x-codex-turn-state`、`x-codex-turn-metadata` 等,见 `crates/adapters/src/codex/client.rs:788-835`。
注意:TS WebSocket/HTTP SSE 当前都使用 `session_id`;真实 WS 抓包和 OpenAI 原版 Rust 使用 `session-id` / `thread-id`。本轮先保持 TS,后续用真实链路决定是否需要有依据地偏离 TS。
TS `config/fingerprint.yaml` 是 HTTP header 指纹基线,不是 TLS 指纹;Rust 的默认 `Fingerprint` 已覆盖同一组默认 headers、header order,并以 TS `config/default.yaml` 的 `darwin/arm64` UA 为默认基线。本轮复核确认 `auth_domains` / `auth_domain_exclusions` 只在 TS config schema/loading 中出现,未找到发送链路使用点,暂不作为 Rust 必须新增的运行时逻辑。
TS 的 Codex context 字段读取规则是“请求体直接字段优先,`client_metadata` 兜底”。Rust 已补齐该规则,覆盖 `x-codex-turn-metadata`、`x-codex-beta-features`、`x-responsesapi-include-timing-metrics`、`x-codex-window-id`、`x-codex-parent-thread-id`;`turnState` 和 `version` 仍保持直接字段/header 来源。WS header 投影已回到从统一 header map 派生,避免手写白名单漏掉 `x-codex-turn-state` 等 TS 条件头。
WS payload 当前也按 TS `createResponseViaWebSocket` 对齐:`instructions` 即使为空也发送;`tool_choice` / `parallel_tool_calls` 使用 TS 默认值;`reasoning`、`tools`、`include` 仅在 TS 会赋值时发送,不再保留旧 Rust 快照里的 `reasoning:null`、`tools:[]`、`include:[]`。
## 后续推进顺序
建议按这个顺序修:
1. P1:继续用真实链路比 HTTP SSE header、TLS 指纹,以及真实上游 cookie 返回情况。
2. P1:持续跟踪 OpenAI 原版客户端和 TS 后续风控增强,发现差异先记录再修复。
## 不建议做的事
- 不要为了“测试覆盖”先新增大量迁移兼容测试。优先用真实链路验证关键风控行为,再补最小固定测试。
- 不要把 quota rotation 也纳入预剥离。TS 修复明确要求“仅封禁/禁用风险才剥离”,quota rotation 由普通 history recovery 处理。
- 不要把旧项目所有行为都默认照搬。遇到差异先记录事实,再用真实链路和原版请求对比判断。
- 不要扩大到非 OpenAI provider 的兼容层。
+487
View File
@@ -0,0 +1,487 @@
# 上游请求链路审计
日期:2026-06-20
状态:已完成审计,并按确认方向推进实现代码。
## 目标
本轮只审计 OpenAI/Codex 上游请求链路:
- `/backend-api/codex/responses` HTTP/SSE
- `/backend-api/codex/responses` WebSocket
- `/backend-api/codex/responses/compact`
- `/backend-api/codex/models`
- `/backend-api/wham/usage` 这类账号/额度辅助入口
不把项目扩展成通用 proxy、VPN 或多 provider 网关;Anthropic、Gemini、OpenRouter、Ollama 等旧项目内容不进入本轮设计。
## 结论
1. WebSocket 连接应使用 OpenAI 原版锁定的 `openai-oss-forks/tokio-tungstenite` 完整握手,不再手写实际 opening request。
2. WebSocket 帧压缩不应长期维护自写 `PerMessageDeflateStream`,已切到 fork 内建 `permessage-deflate` 协商和解压/压缩状态。
3. `opening_request_text()` 继续作为审计快照保留;真实发送字节由 fork 根据 request headers 和 `WebSocketConfig` 生成。
4. HTTP/SSE header 不应按 WebSocket 抓包删到最小集合。原版和旧 TS 都会带 Codex compatibility headers,例如 `x-codex-turn-state`、`x-codex-window-id`、`x-codex-parent-thread-id`、`x-openai-subagent`、`x-codex-installation-id`。
5. 默认指纹先按 TS `config/default.yaml` 与 `config/fingerprint.yaml` 对齐;真实 Linux 抓包和 macOS 解包信息只作为后续验证证据,不直接覆盖 TS 基线。
6. TLS 层当前 Rust 和原版 Rust 都走 rustls/native roots/custom CA 思路,不能从 HTTP 抓包直接证明 TLS 指纹完全一致。先对齐原版 Rust 依赖和请求头,再用真实链路继续验证。
7. Chat 请求的模型后缀、默认 reasoning、默认 service tier、`fast -> priority` 和展示模型名先以 TS 当前实现为准,不从抓包或经验补规则。
8. 账号导入时的 `chatgpt-account-id` 补齐先按 TS:导入/持久化从 JWT claim 提取,发送链路只消费已解析出的 account id;不能在请求头层随意发空值或测试占位值。
## 基线规则
本轮不再“猜着写”。所有请求链路变更必须能对应到以下至少一个来源:
- TS 当前实现:`/home/zyy/桌面/Codes/codex-proxy`
- OpenAI 原版 Rust:`/tmp/openai-codex-src/codex-rs`
- 新的真实链路证据,并且已和 TS / 原版对比过
`docs/openai-res.txt` 的真实抓包用于验收和发现差异,但不能单独反推出全局白名单。比如该抓包的 WS opening 没有 `x-codex-installation-id`,而 TS 当前 HTTP/WS 统一 header map 会发送它;本轮先保持 TS 行为,再用新抓包验证。
## 证据索引
### 真实抓包
`docs/openai-res.txt` 中 WebSocket opening header 的顺序和集合是:
```text
Host
Connection
Upgrade
Sec-WebSocket-Version
Sec-WebSocket-Key
chatgpt-account-id
authorization
user-agent
originator
openai-beta
x-client-request-id
session-id
thread-id
x-codex-window-id
x-codex-turn-metadata
sec-websocket-extensions
```
其中 `sec-websocket-extensions` 为:
```text
permessage-deflate; client_max_window_bits
```
抓包中 `/backend-api/wham/usage` 只有较少 HTTP header:`user-agent`、`authorization`、`chatgpt-account-id`、`accept`、`cookie`。这说明不同上游入口的 header 集合不同,不能把 WS opening 的最小集合直接套到所有 HTTP 请求。
### 当前 Rust 实现
调整前 workspace 依赖:
- `tokio-tungstenite = 0.29.0`,未启用 fork 的 `deflate` 能力。
- `reqwest = 0.12.28`,使用 rustls native roots,启用 gzip/brotli/zstd/deflate/http2。
- `rustls = 0.23.36`,`tokio-rustls = 0.26.4`。
当前 WebSocket opening 审计快照仍按真实抓包顺序生成:
- `crates/adapters/src/codex/websocket/connect.rs`
- `CodexWebSocketConnection::responses`
- `opening_request_bytes`
当前 WS business headers:
- `chatgpt-account-id`
- `authorization`
- `user-agent`
- `originator`
- `openai-beta`
- `x-client-request-id`
- `session-id`
- `thread-id`
- `x-codex-window-id`
- `x-codex-turn-metadata`
- `x-openai-subagent`,仅在 request metadata 提供合法 subagent 时发送
调整前仍有自写 deflate 层:
- `crates/adapters/src/codex/websocket/deflate.rs`
- `PerMessageDeflateStream`
- `flate2::Decompress`
- 支持 server context takeover 后,真实 WS stream 已跑通
这说明“旧修复可用”,但维护面仍然偏大:自写代码需要自己处理 frame parsing、RSV1、continuation、context takeover、错误传播和后续 tungstenite 行为变化。当前实现已删除这层自写协议代码,并改用 fork 的 `connect_async_tls_with_config` 完成握手、扩展协商和 response 校验。
### OpenAI 原版 Rust
原版路径:
- `/tmp/openai-codex-src/codex-rs/Cargo.toml`
- `/tmp/openai-codex-src/codex-rs/codex-api/src/endpoint/responses_websocket.rs`
- `/tmp/openai-codex-src/codex-rs/core/src/client.rs`
- `/tmp/openai-codex-src/codex-rs/core/src/responses_metadata.rs`
原版锁定 fork:
```toml
tokio-tungstenite = { version = "0.28.0", features = ["proxy", "rustls-tls-native-roots"] }
tungstenite = { version = "0.27.0", features = ["deflate", "proxy"] }
[patch.crates-io]
tokio-tungstenite = { git = "https://github.com/openai-oss-forks/tokio-tungstenite", rev = "132f5b39c862e3a970f731d709608b3e6276d5f6" }
tungstenite = { git = "https://github.com/openai-oss-forks/tungstenite-rs", rev = "9200079d3b54a1ff51072e24d81fd354f085156f" }
```
原版 WebSocket 连接:
- 使用 `connect_async_tls_with_config`
- `WebSocketConfig.extensions.permessage_deflate = Some(DeflateConfig::default())`
- `false` 表示不禁用 Nagle
- 自定义 CA 存在时给 WebSocket 显式 rustls connector,避免 WS 和 HTTP 的 CA 策略不一致
原版 Responses header 语义:
- `x-client-request-id` 使用 thread id
- `session-id` / `thread-id` 来自 session metadata
- `x-codex-window-id`、`x-codex-turn-metadata`、`x-codex-parent-thread-id` 属于 compatibility headers
- `x-codex-turn-metadata` 的 canonical 来源是 `client_metadata["x-codex-turn-metadata"]`,HTTP/ws 直接 header 是兼容投影
- `x-openai-subagent` 支持 review、compact、memory、collab 等内部语义
### 旧 TS 项目
旧项目路径:
- `/home/zyy/桌面/Codes/codex-proxy/src/proxy/codex-api.ts`
- `/home/zyy/桌面/Codes/codex-proxy/src/proxy/ws-transport.ts`
- `/home/zyy/桌面/Codes/codex-proxy/src/fingerprint/manager.ts`
- `/home/zyy/桌面/Codes/codex-proxy/config/fingerprint.yaml`
- `/home/zyy/桌面/Codes/codex-proxy/src/proxy/cookie-jar.ts`
有价值的差异线索:
- `src/fingerprint/manager.ts:91-113` 的 `buildHeaders` 明确只做这些事:写 `Authorization`,用传入 `accountId` 或 JWT claim 补 `ChatGPT-Account-Id`,写 `originator`,合并 `buildRawDefaultHeaders`,最后按 `fingerprint.header_order` 排序。
- `buildRawDefaultHeaders` 只展开 `User-Agent`、`sec-ch-ua`,再合并 `fingerprint.yaml` 的 `default_headers`。
- `src/proxy/codex-api.ts:createResponseViaHttp` 和 `createResponseViaWebSocket` 在统一 base headers 上追加 `OpenAI-Beta`、`x-openai-internal-codex-residency`、`x-client-request-id`、`x-codex-installation-id`、session/window/context headers。
- `ws-transport.ts` 使用 Node `ws` 完整握手,业务 headers 通过构造参数传入;没有手写 raw opening,也没有手写 `Sec-WebSocket-Extensions`。
- `buildWsConstructorOpts` 只传 `{ headers }` 和可选 agent;`ws@8.19.0` 客户端默认 `perMessageDeflate: true`,因此扩展 offer 由库生成。
- `ws@8.19.0` 的 `WebSocketServer` 默认 `perMessageDeflate: false`。旧 TS 测试 helper 若未显式打开 server deflate,不代表真实客户端不会发送扩展 offer。
- WebSocket 升级响应头通过 `upgrade` 事件采集后转回 SSE response headers。
- WebSocket 打开后发送 `response.create` JSON 文本帧,并用 25s ping 保活。
- `buildConversationIdentity` 用账号作用域派生 `cp_...` conversation id 和 `cw_...` window id。
- HTTP 与 WS 都会补 `OpenAI-Beta: responses_websockets=2026-02-06`。
- HTTP/SSE 会带 `Accept: text/event-stream`,compact 不带。
- `x-codex-installation-id` 同时出现在 header 和 `client_metadata`。
- `x-codex-turn-metadata`、`x-codex-beta-features`、`x-responsesapi-include-timing-metrics`、`Version`、`x-codex-parent-thread-id` 都是条件性透传。
- cookie 自动捕获只允许 `cf_clearance`,不自动持久化 `__cf_bm`。
- `CodexApi.captureCookies()` 的调用点在 transport response 返回后立即执行:HTTP SSE 和 compact 都在状态码判断前 capture;WebSocket 升级响应 header 由 `ws-transport.ts` 采集后随 response 返回;warmup 路径也会 capture,但当前账号导入链路禁用 warmup,普通 `getUsage()` 不做 capture。
- `fingerprint.yaml` 是 HTTP header 指纹基线,不是 TLS/JA3 指纹;其中定义 UA 模板、认证域白名单、header 顺序和默认浏览器态 header。
- `auth_domains` / `auth_domain_exclusions` 在当前 TS HEAD 只进入 schema,未找到请求发送链路使用点。
- `src/auth/account-registry.ts:addAccount` 和 `src/auth/account-persistence.ts:normalizeAccountEntries` 都从 JWT 中回填 `accountId`、`userId`、`planType`、`email`。因此 Rust 导入阶段刷新并解析 claims 后写入 `account_id`,是和 TS 一致的方向;请求发送阶段不需要引入 `proxyApiKey`、`proxy_api_key` 等 sub2api 私有字段。
不应迁移的内容:
- 多 provider 适配逻辑。
- 通用 proxy/VPN 语义。
- 旧迁移兼容测试。
- 与 OpenAI/Codex 无关的历史抽象。
## 三方对比
### WebSocket opening
| 项目 | 生成方式 | 风险 |
| --- | --- | --- |
| 真实抓包 | Codex Desktop 发送的最小 opening header | 当前最可信的行为目标 |
| 当前 Rust | fork `tokio-tungstenite` 完整握手 | 与原版一致;header 顺序由库生成 |
| OpenAI 原版 Rust | fork `tokio-tungstenite` 自动握手 | deflate 协商和状态可靠 |
| 旧 TS | Node `ws` transport | 同样由库生成 opening 并处理 permessage-deflate |
判断:TS 版和 OpenAI 原版都不是手写 raw opening。为了拿到 fork 内部 negotiated extension state,Rust 当前也应使用 fork 完整握手;审计快照只用于排查差异。
### WebSocket permessage-deflate
| 方案 | 优点 | 风险 |
| --- | --- | --- |
| 当前自写 `PerMessageDeflateStream` | 已真实跑通;不改依赖 | 协议维护面大,容易漏 frame/extension 边界 |
| OpenAI fork `tungstenite` | 与原版同源;支持协商和上下文 | 引入 git patch;版本从当前 0.29 回到原版 0.27/0.28 组合 |
| 旧 TS `ws@8.19.0` | 客户端默认开启 `perMessageDeflate`,无需项目手写扩展头 | 测试 server 默认关闭 deflate,不能用测试 helper 默认值反推真实客户端 opening |
判断:已切 fork,删除自写 deflate。运行时风险小于继续维护自写协议层。
### HTTP/SSE header
当前 Rust HTTP/SSE 已覆盖主要 Codex 头:
- `authorization`
- `chatgpt-account-id`
- `originator`
- `user-agent`
- `content-type`
- `accept: text/event-stream`
- `openai-beta`
- `x-openai-internal-codex-residency`
- `x-client-request-id`
- `x-codex-installation-id`
- `session_id`
- `x-codex-window-id`
- `x-codex-turn-state`
- `x-codex-turn-metadata`
- `x-codex-beta-features`
- `x-responsesapi-include-timing-metrics`
- `version`
- `x-codex-parent-thread-id`
- `x-openai-subagent`
- `cookie`
注意:
- `session_id` 是 TS 当前 WebSocket/HTTP SSE 都会发送的兼容字段;`v2.0.78-beta.15c49d6` 和当前 `dev/v2.0.80` 都没有把它转换为 `session-id` / `thread-id`。
- WS opening 真实抓包和 OpenAI 原版 Rust 用的是 `session-id` / `thread-id`,不是 `session_id`。这是后续真实链路验证项,本轮不在 Rust 侧猜测改写。
- 当前 WS opening 按 TS `buildHeaders` + `applyCodexContextHeaders` 语义从统一 header map 投影:过滤 HTTP/SSE 专用的 `content-type`、`accept`,保留 `session_id`、fingerprint 默认头、cookie、`x-openai-internal-codex-residency`、`x-codex-installation-id` 和各类 Codex context 条件头。
- `docs/openai-res.txt` 的一次真实 WS opening 抓包是最小集合,但 TS 当前实现会发送更多兼容/条件头;本轮基线先与 TS 保持一致,不用手写白名单猜测。
- `x-codex-turn-metadata`、`x-codex-beta-features`、`x-responsesapi-include-timing-metrics`、`x-codex-window-id`、`x-codex-parent-thread-id` 的 canonical 来源允许是 `client_metadata`;Rust 已按 TS `firstRequestString` 语义补齐“直接字段优先,metadata 兜底”。这会影响 account-scoped window id、HTTP header 投影和 WS business header 投影。
- `x-openai-subagent` 在 TS WebSocket 和 HTTP 路径都会从 `client_metadata` 规范化后写入 header;OpenAI 原版 Responses stream 也会写该 header。Rust 已补齐 WS header 投影。
### Chat 模型选项
TS 参照:
- `src/translation/openai-to-codex.ts` 在 Chat 翻译阶段解析模型名后缀,并按“显式请求字段 > 模型后缀 > 配置默认值”的顺序设置 `reasoning.effort` 和 `service_tier`。
- `src/routes/chat.ts` 在翻译后判断 `codexRequest.reasoning?.effort`,因此后缀或配置默认 reasoning 也会影响 Chat stream 输出格式。
- `src/routes/chat.ts` 使用 `buildDisplayModelName(parseModelName(req.model))` 作为 OpenAI Chat 响应中的展示模型名。
- `src/routes/shared/proxy-request-preparation.ts` 在带 reasoning 且 include 为空时补 `["reasoning.encrypted_content"]`。
- `src/proxy/codex-api.ts` 的 `normalizeServiceTierForUpstream` 在真正发 HTTP/WS 前把 `fast` 映射为 `priority`。
Rust 对齐结果:
- `/v1/chat/completions` handler 在 `translate_chat_to_codex` 后立即调用 `apply_response_model_options`,让 suffix/config/default 对 `include_reasoning` 和 Chat stream 可见。
- `ChatDispatchService::complete` 也调用同一个函数兜底,避免非 handler 调用绕过模型选项。
- Chat stream 响应模型名改为 `ModelCatalog::build_display_model_name(parsed_model)`,与 TS display model 规则一致。
- `apply_response_model_options` 继续复用 Responses 路径的 `fast -> priority` 和 reasoning include 规则,避免 Chat/Responses 分叉。
### WebSocket payload
TS `src/proxy/codex-api.ts:createResponseViaWebSocket` 的 `wsRequest` 构造顺序是:
```text
type
model
instructions
input
store
stream
previous_response_id?
reasoning?
tools?
tool_choice
parallel_tool_calls
text?
service_tier?
prompt_cache_key?
include?
client_metadata
```
对齐判断:
- `instructions` 使用 TS 语义 `request.instructions ?? ""`,空字符串也必须发送;真实上游会在缺失该字段时返回 400。
- `tool_choice` 和 `parallel_tool_calls` 即使客户端未传也按 TS 补默认值:`auto`、`true`。
- `reasoning`、`tools`、`include` 不再发送旧 Rust 快照里的 `null` / 空数组;TS 只在有值或非空数组时赋值。
- `generate` 不属于 TS 当前 `WsCreateRequest`,WebSocket payload 不按 Rust 旧测试快照额外发送。
- `docs/openai-res.txt` 中原版抓包的 payload 字段集合可作为后续原版对照证据,但本轮先以 TS 当前实现为 Rust 基线。
## 指纹审计
TS 当前 `config/fingerprint.yaml`:
```yaml
user_agent_template: "Codex Desktop/{version} ({platform}; {arch})"
auth_domains: ["chatgpt.com", "*.chatgpt.com", "openai.com", "*.openai.com"]
auth_domain_exclusions: ["ab.chatgpt.com"]
default_headers:
Accept-Encoding: "gzip, deflate, br, zstd"
Accept-Language: "en-US,en;q=0.9"
sec-ch-ua-mobile: "?0"
sec-ch-ua-platform: '"macOS"'
sec-fetch-site: "same-origin"
sec-fetch-mode: "cors"
sec-fetch-dest: "empty"
```
对应 TS 发送逻辑:
- `buildHeaders` / `buildHeadersWithContentType` 负责展开 UA、补 `sec-ch-ua`、合并默认 headers 并按 `header_order` 排序。
- `sec-ch-ua` 来自 `client.chromium_version`,不是写在 `fingerprint.yaml` 里。
- `auth_domains` / `auth_domain_exclusions` 当前只在 config schema/loading 里出现,没有发现实际发送链路使用点,不能作为 Rust 侧必须新增的运行时逻辑。
对 Rust 的影响判断:
- 这个 YAML 可以作为“默认 HTTP header 指纹基线”的证据,不应被理解成 TLS/JA3/ALPN 指纹配置。
- Rust 侧当前不需要引入同名 YAML 运行时配置;先保持 `Fingerprint::default_codex_desktop()` 这类强类型默认值,减少外部配置漂移面。
- `header_order` 在 TS 配置里是原始大小写,Rust 内部统一用小写 key 排序;这不影响实际 HTTP header 语义,但真实抓包排查时要按 wire 输出核对。
- `auth_domains` / `auth_domain_exclusions` 如果后续要启用,应该只用于 cookie/auth header 作用域校验,不能自动扩大成通用 proxy 域名白名单。
TS 默认指纹基线:
```text
originator: Codex Desktop
app_version: 26.519.81530
build_number: 3178
platform: darwin
arch: arm64
chromium_version: 146
user-agent: Codex Desktop/{version} ({platform}; {arch})
sec-ch-ua-platform: "macOS"
```
真实 Linux 抓包中的 UA:
```text
Codex Desktop/0.140.0 (Debian 13.0.0; x86_64) gnome-terminal (Codex Desktop; 26.609.71450)
```
macOS Codex Desktop 解包与探测来源:
- `/home/zyy/桌面/Codes/codex-desktop-linux/codex-app/resources/app.asar`
- `/tmp/codex-desktop-fingerprint/fingerprint-report.md`
确认到的 macOS arm64 UA:
```text
Codex Desktop/26.609.71450 (Macintosh; Intel Mac OS X; arm64)
```
风险:
- TS 默认指纹、真实 Linux 抓包和 macOS 解包信息不一致,不能在没有 TS 证据时自行混合。
- HTTP 默认 header 里仍有 macOS `sec-ch-ua-platform`,而真实 WS opening 不带 `sec-ch-*`。
- 如果上游把 UA、account id、cookie、TLS、IP、timing 关联起来,指纹漂移会增加风控风险。
当前对齐原则:
1. 先以 TS `config/default.yaml` + `config/fingerprint.yaml` 为 Rust 默认基线:版本 `26.519.81530`、构建 `3178`、平台 `darwin`、架构 `arm64`、Chromium `146`。
2. 后续仍应支持从真实导入/配置中持久化当前使用的 Codex Desktop 指纹。
3. WS opening 继续不带 `sec-ch-*`。
4. HTTP/SSE 是否保留 `sec-ch-*` 需要用真实抓包继续确认;在没有 HTTP `/codex/responses` 完整抓包前,不做大删减。
Rust 当前状态:
- `crates/core/src/gateway/fingerprint.rs` 已硬编码同一组默认 headers 和 header 顺序。
- 默认 UA 已对齐 TS 展开结果:`Codex Desktop/26.519.81530 (darwin; arm64)`。
- auto-update 没有 extracted Chromium 时保留已有 Chromium;首次无历史值时使用 TS 默认 `146`,不再回退到旧 `142`。
- `crates/adapters/src/codex/client.rs` 的 HTTP/SSE 路径会发送 `sec-ch-ua` 和 `fingerprint.default_headers`;WS opening 路径按 TS 当前实现从统一 Codex header map 投影,过滤 HTTP/SSE 专用 `content-type`、`accept`,保留 `session_id`,不按抓包最小集合手写白名单。
## TLS 与网络安全面
当前 Rust HTTP:
- `reqwest`
- rustls native roots
- 自定义 CA 支持
- `.no_proxy()`
- 可配置 `force_http11`
- 默认保留 HTTP/2
当前 Rust WS:
- fork `connect_async_tls_with_config`
- rustls native roots/custom CA
- 不禁用 Nagle
- opening response 和 permessage-deflate 协商由 fork 校验
OpenAI 原版 WS:
- fork `tokio-tungstenite`
- rustls native roots
- 自定义 CA connector
- 不禁用 Nagle
- 内建 permessage-deflate
审计判断:
- 目前与原版 Rust 的 WS/TLS 策略方向一致。
- `.no_proxy()` 符合本项目“不做 proxy/VPN”的定位。
- 仅凭 `docs/openai-res.txt` 不能判断 TLS ClientHello/JA3/ALPN 是否完全一致。
- 如果后续真实链路仍被风控,需要单独抓 TLS 指纹和 ALPN,而不是继续猜 header。
## Cookie 审计
当前 Rust 与旧 TS 的关键策略一致:
- 自动捕获只持久化 `cf_clearance`。
- 不自动持久化 `__cf_bm`。
- 请求时按 domain/path/expiry 过滤 cookie。
- Cloudflare challenge 或 path-block 时会清理账号 cookie 并进入冷却/禁用路径。
- HTTP/SSE、compact、stream 和 WebSocket 升级成功响应都会把上游 `Set-Cookie` 交给 runtime capture。
- HTTP/SSE 和 compact 的非 2xx 响应会通过 `CodexClientError::Upstream.set_cookie_headers` 带回 cookie,再由 runtime capture;这对齐 TS “收到 transport response 先 capture,再按状态码分类”的时机。
- WebSocket 错误帧本身没有 HTTP response headers,按空 cookie 列表处理;WebSocket opening 失败如果带 `Set-Cookie`,会从握手响应头传回。
这个方向合理。`__cf_bm` 与 IP、UA、TLS、时间强绑定,重放可能比不带 cookie 更差。
## 使用 OpenAI fork 的风险
收益:
- 与 OpenAI 原版 Codex Rust 同源。
- 解决真实 WS `permessage-deflate` 的长期维护问题。
- fork 已有 extension negotiation、context takeover、压缩/解压测试。
- 避免手写 opening 后无法访问 fork 私有 negotiated extensions 状态的问题。
风险:
- 引入 git patch,依赖供应链需要固定 rev。
- 当前项目 `tokio-tungstenite` 是 0.29,原版组合是 `tokio-tungstenite 0.28` + `tungstenite 0.27`,存在 API 细节回退风险。
- `[patch.crates-io]` 会影响整个 workspace 内所有 `tungstenite` / `tokio-tungstenite` 使用点。
- fork 不是 crates.io 正式发布版本,后续升级需要跟 OpenAI 原版同步审计。
风险控制:
- 固定到原版相同 rev。
- 只改 WS transport 相关代码,不顺手改账号、导入、日志。
- 保留 WS opening 审计快照测试。
- 新增最小的 deflate context takeover 行为测试,替代自写 deflate helper 测试。
- 跑真实 `gpt-5.5` HTTP/SSE + WS stream 复测。
## 本轮实现项
1. Cargo 依赖切到 OpenAI 原版 fork。
2. `connect.rs` 改用 fork `connect_async_tls_with_config` 完整握手。
3. 通过 `WebSocketConfig.extensions.permessage_deflate = Some(DeflateConfig::default())` 让 fork 生成 offer 并校验上游 agreed extensions。
4. 删除 `crates/adapters/src/codex/websocket/deflate.rs` 和相关 type alias。
5. 保留 WS opening 审计快照测试。
6. 适配器层 deflate 行为由 fork 内建握手覆盖,不再维护自写压缩帧。
7. 默认 Codex Desktop 指纹先对齐 TS 配置:`26.519.81530 / 3178 / darwin / arm64 / Chromium 146`。
8. OpenAI Responses 翻译层补齐 TS `firstRequestString` 语义:直接字段为空时从 string-only `client_metadata` 读取 Codex context 字段。
9. WebSocket header 投影回到从统一 Codex header map 派生,过滤 HTTP/SSE 专用头,保留 TS `createResponseViaWebSocket` 会发送的 `session_id`、`x-codex-turn-state`、`x-codex-installation-id`、`x-openai-internal-codex-residency`、`x-openai-subagent` 等条件头。
10. `/v1/chat/completions` 的 `stream:true` 复用 Responses stream 调度链路,并按 TS `streamCodexToOpenAI` 语义输出 OpenAI `chat.completion.chunk` SSE。
11. WebSocket `response.create` payload 按 TS 字段顺序和可选字段规则序列化:保留空 `instructions`,省略空 `tools` / `include` 和 `reasoning:null`。
12. `/v1/chat/completions` 的模型后缀、reasoning、service tier 和 display model 已按 TS Chat 路径对齐;定向测试固定了 `gpt-5.5-high-fast` 上游发送为 `model=gpt-5.5`、`service_tier=priority`、`reasoning.effort=high`,并验证 stream reasoning delta 会转为 OpenAI `reasoning_content`。
13. Chat 非流和 compact 成功路径已接入结构化事件日志,复用现有 `record_response_event`,按 route 区分 `v1.chat` / `v1.response` 并记录 accountId、usage、rate-limit headers;chat 非流失败耗尽和 compact 终态失败路径补充 `level=error`、`failureClass`、`exhaustedCount`、`upstreamError`。
14. Responses stream 启动阶段终态失败已接入结构化事件日志:账号耗尽、dirty quota 校验达到最大次数、模型不支持重试耗尽、首帧前空响应 / 缺 completion / SSE 解析失败都会记录 `stream=true`、`transport`、`failureClass`,与 TS `recordStreamCloseEvent` 的排查目标保持一致,但落到本项目现有 `event_logs`。
15. cookie capture 时机按 TS 补齐:业务响应成功、stream 启动成功、compact 成功、HTTP/compact upstream error、WebSocket opening error 都会消费上游 `Set-Cookie`,底层仍只允许自动持久化 `cf_clearance`。
## 验证顺序
1. `cargo fmt`。
2. `cargo check -p codex-proxy-runtime -p codex-proxy-server`。
3. `cargo test -p codex-proxy-adapters --test codex`。
4. `cargo test -p codex-proxy-server --test openai_chat_upstream openai_chat_routes::`。
5. `cargo test -p codex-proxy-core --test protocol protocol_openai_chat:: -- --nocapture`。
6. 重启服务,用自然请求文本和 `gpt-5.5` 跑真实 HTTP/SSE、WS 和 chat stream。
已完成的真实验证:
- `/v1/chat/completions` + `stream:true`:请求 `9805212d-401f-4a10-a8f1-3d5bc422299e` 返回 HTTP 200、`text/event-stream`、OpenAI `chat.completion.chunk`、final usage 和 `data: [DONE]`,无 error SSE。
- `/v1/chat/completions` + `stream:true` 日志复测:请求 `dc22ea52-9f74-428f-8fe5-3b1a3def0114` 返回 HTTP 200、`text/event-stream`、OpenAI `chat.completion.chunk`、final usage 和 `data: [DONE]`,管理端事件为 `kind=v1.chat`、`route=/v1/chat/completions`、`metadata.apiKind=chat`、`metadata.completed=true`。
真实请求验证时避免 `test`、`probe`、`real_chain` 这类显眼测试语义,请求 ID 使用 UUID 形态,请求文本使用正常用户表达。
## 后续仍需审计
- HTTP `/codex/responses` 的真实抓包 header,需要和当前 Rust 的 HTTP/SSE 逐项比对。
- 原版 attestation header 何时生成;当前项目是否需要支持,不能盲目伪造。
- `x-openai-internal-codex-residency` 在当前真实账号链路中的必要性。
- `version` header 的真实来源和是否仍被上游使用。
- TLS ClientHello/ALPN 是否需要与原版二进制进一步对齐。
- 账号导入时是否能补齐真实 `account_id`、指纹、installation id、cookie 状态。
## 当前建议
- WS opening/deflate 使用 OpenAI fork 完整握手。
- HTTP/SSE 保留 Codex compatibility headers,不按 WS 抓包全删。
- 后续真实链路继续以 `gpt-5.5`、自然请求文本、真实账号状态为准。