mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 16:06:08 +00:00
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:
Generated
+50
-8
@@ -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
@@ -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 }
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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]
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
|
||||
+2373
-261
File diff suppressed because it is too large
Load Diff
@@ -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("a);
|
||||
let reset_at = quota_snapshot_reset_at("a);
|
||||
if self
|
||||
.store
|
||||
.update_quota_json(&account.id, "a.to_string())
|
||||
.apply_quota_snapshot(
|
||||
&account.id,
|
||||
"a.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("a),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
Err(error) => {
|
||||
warn!(
|
||||
@@ -222,7 +242,7 @@ impl QuotaRefreshTask {
|
||||
pub struct QuotaRefreshSummary {
|
||||
/// 扫描账号数。
|
||||
pub scanned: usize,
|
||||
/// 配额锁定候选账号数。
|
||||
/// 配额锁定或待验证候选账号数。
|
||||
pub candidates: usize,
|
||||
/// 成功刷新账号数。
|
||||
pub refreshed: usize,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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>,
|
||||
|
||||
@@ -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,
|
||||
),
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
})
|
||||
)
|
||||
}
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
|
||||
+3
-6
@@ -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
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because it is too large
Load Diff
@@ -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 的兼容层。
|
||||
@@ -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`、自然请求文本、真实账号状态为准。
|
||||
Reference in New Issue
Block a user