mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 00:05:49 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
53739ebe5c | ||
|
|
7f420257a6 | ||
|
|
4f85209dc9 | ||
|
|
cd89e2e8fc | ||
|
|
0cb5a29699 | ||
|
|
61788804b4 | ||
|
|
2751ea6aa4 | ||
|
|
58fe66789c | ||
|
|
573b6b26db | ||
|
|
d425428ebe | ||
|
|
b26d3b386a | ||
|
|
cf0b14c5fd | ||
|
|
056561de41 | ||
|
|
f2215df72e | ||
|
|
72f8957163 | ||
|
|
73f0c15790 | ||
|
|
0558693949 | ||
|
|
2a9cb02cc9 | ||
|
|
5ce36931a8 | ||
|
|
0d826d9eba | ||
|
|
304242957a | ||
|
|
ef37ae8d9e | ||
|
|
fc13d4f2b6 | ||
|
|
78f9c0c800 | ||
|
|
a3bc7bc445 | ||
|
|
0132114285 | ||
|
|
5970efbb7c | ||
|
|
7bf28b571d | ||
|
|
bc3bd6bc43 | ||
|
|
9acfab0a86 | ||
|
|
3e5a84d2a0 | ||
|
|
92948f51a9 | ||
|
|
776637fccb | ||
|
|
82acb08d26 | ||
|
|
7744e1b0b4 | ||
|
|
0b8a1add18 | ||
|
|
2dd60c08b6 | ||
|
|
cce9880937 | ||
|
|
b8ea1cd076 | ||
|
|
4117840f75 | ||
|
|
f642764f34 | ||
|
|
482e192e2d | ||
|
|
ee51a07de4 | ||
|
|
bf23658aaf | ||
|
|
bc4728332e | ||
|
|
2cb5cda077 | ||
|
|
9d90dedb81 | ||
|
|
6074cff68c | ||
|
|
6ab8e7d0a6 | ||
|
|
17f9ef8915 | ||
|
|
15e0a28bd0 |
@@ -55,6 +55,9 @@ jobs:
|
||||
target: runtime
|
||||
load: true
|
||||
tags: ${{ env.CPR_IMAGE }}
|
||||
no-cache-filters: |
|
||||
runtime-base
|
||||
runtime
|
||||
cache-from: type=gha,scope=runtime-ci
|
||||
cache-to: type=gha,mode=max,scope=runtime-ci
|
||||
|
||||
@@ -76,7 +79,7 @@ jobs:
|
||||
${{ hashFiles('trivy-results.sarif') != '' &&
|
||||
(github.event_name != 'pull_request' ||
|
||||
github.event.pull_request.head.repo.full_name == github.repository) }}
|
||||
uses: github/codeql-action/upload-sarif@5595ccaf912efad79be6eef63a5619ff05969be3 # v4.37.6
|
||||
uses: github/codeql-action/upload-sarif@ff2f1c621b7f889edc0d3c761ac2e6a3f8cdb0dd # v4.37.7
|
||||
with:
|
||||
sarif_file: trivy-results.sarif
|
||||
category: runtime-container
|
||||
|
||||
@@ -468,6 +468,9 @@ jobs:
|
||||
push: true
|
||||
build-args: |
|
||||
CPR_PREBUILT_DIR=dist/docker
|
||||
no-cache-filters: |
|
||||
runtime-base
|
||||
runtime-prebuilt
|
||||
cache-from: type=registry,ref=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:buildcache-${{ matrix.os }}-${{ matrix.arch }}
|
||||
cache-to: type=registry,ref=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:buildcache-${{ matrix.os }}-${{ matrix.arch }},mode=min
|
||||
provenance: mode=max
|
||||
|
||||
Generated
+44
-36
@@ -120,9 +120,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "async-trait"
|
||||
version = "0.1.91"
|
||||
version = "0.1.92"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ae36dc4177970ef04fde5178d3e2429882def40e57a451f919c098f72baa6cec"
|
||||
checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
@@ -451,9 +451,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-types"
|
||||
version = "1.6.1"
|
||||
version = "1.6.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d6dc683efb34b9e755675b37fedbe0103141e5b6df7bdc9eb6967756a8c167d8"
|
||||
checksum = "fce83ce9abbb198d25bc7131e468d0f9fe1257125e58c39f3f9fc9f5098c9647"
|
||||
dependencies = [
|
||||
"base64-simd",
|
||||
"bytes",
|
||||
@@ -801,9 +801,9 @@ checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c"
|
||||
|
||||
[[package]]
|
||||
name = "cookie"
|
||||
version = "0.18.1"
|
||||
version = "0.18.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4ddef33a339a91ea89fb53151bd0a4689cfce27055c291dfa69945475d22c747"
|
||||
checksum = "1a373e3602691c3cdea496d2f0ee5935151e6168fe87739483c463db1b2f2f87"
|
||||
dependencies = [
|
||||
"percent-encoding",
|
||||
"time",
|
||||
@@ -1246,7 +1246,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.52.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1384,9 +1384,9 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
|
||||
|
||||
[[package]]
|
||||
name = "futures"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a88cf1f829d945f548cf8fec32c61b1f202b6d93b45848602fc02af4b12ad218"
|
||||
checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3"
|
||||
dependencies = [
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
@@ -1399,9 +1399,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "futures-channel"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "262590f4fe6afeb0bc83be1daa64e52657fe185690a958af7f3ad0e92085c5ae"
|
||||
checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"futures-sink",
|
||||
@@ -1409,15 +1409,15 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "futures-core"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7"
|
||||
checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e"
|
||||
|
||||
[[package]]
|
||||
name = "futures-executor"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6754879cc9f2c66f88c6e5c35344bb0bdb0708b0352b1201815667c7eabc7458"
|
||||
checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"futures-task",
|
||||
@@ -1437,32 +1437,32 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "futures-io"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4577ecaa3c4f96589d473f679a71b596316f6641bc350038b962a5daf0085d7a"
|
||||
checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed"
|
||||
|
||||
[[package]]
|
||||
name = "futures-macro"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2d6d3cde68c518367be28956066ddfef33813991b77a55005a69dae04bf3b10b"
|
||||
checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.119",
|
||||
"syn 3.0.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures-sink"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e34418ac499d6305c2fb5ad0ed2f6ac998c5f8ca209b4510f7f94242c647e307"
|
||||
checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d"
|
||||
|
||||
[[package]]
|
||||
name = "futures-task"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109"
|
||||
checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd"
|
||||
|
||||
[[package]]
|
||||
name = "futures-timer"
|
||||
@@ -1472,9 +1472,9 @@ checksum = "af43fadb8a98512d547e37b4e92e0ced13e205c061b87b4623eff01d918d6968"
|
||||
|
||||
[[package]]
|
||||
name = "futures-util"
|
||||
version = "0.3.33"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa"
|
||||
checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc"
|
||||
dependencies = [
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
@@ -1686,9 +1686,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "h2"
|
||||
version = "0.4.15"
|
||||
version = "0.4.16"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155"
|
||||
checksum = "a9f37a958b41b3b19ee2707c06439c0e9e547e847223eb791ecb0cb821c65e27"
|
||||
dependencies = [
|
||||
"atomic-waker",
|
||||
"bytes",
|
||||
@@ -2581,6 +2581,7 @@ dependencies = [
|
||||
"cookie",
|
||||
"crypto_box",
|
||||
"ed25519-dalek",
|
||||
"flate2",
|
||||
"futures",
|
||||
"gateway-admin",
|
||||
"gateway-core",
|
||||
@@ -2608,6 +2609,7 @@ dependencies = [
|
||||
"tracing-subscriber",
|
||||
"tungstenite 0.27.0",
|
||||
"url",
|
||||
"urlencoding",
|
||||
"uuid",
|
||||
"wiremock",
|
||||
"zstd",
|
||||
@@ -2712,7 +2714,7 @@ dependencies = [
|
||||
"once_cell",
|
||||
"socket2",
|
||||
"tracing",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.52.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2982,7 +2984,7 @@ dependencies = [
|
||||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.52.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3623,26 +3625,26 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
|
||||
dependencies = [
|
||||
"fastrand",
|
||||
"getrandom 0.3.4",
|
||||
"getrandom 0.4.3",
|
||||
"once_cell",
|
||||
"rustix",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.52.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror"
|
||||
version = "2.0.19"
|
||||
version = "2.0.20"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9"
|
||||
checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f"
|
||||
dependencies = [
|
||||
"thiserror-impl",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror-impl"
|
||||
version = "2.0.19"
|
||||
version = "2.0.20"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd"
|
||||
checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
@@ -4078,6 +4080,12 @@ dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "urlencoding"
|
||||
version = "2.1.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "daf8dba3b7eb870caf1ddeed7bc9d2a049f3cfdfae7cb521b087cc33ae4c49da"
|
||||
|
||||
[[package]]
|
||||
name = "utf-8"
|
||||
version = "0.7.6"
|
||||
|
||||
@@ -27,8 +27,9 @@ impl LoadableConfig for GatewayConfig {
|
||||
return Err(ConfigError::InvalidField("schema_version"));
|
||||
}
|
||||
self.host.resolve_and_validate(source_dir)?;
|
||||
let runtime_data_dir = self.host.runtime_data_dir().to_path_buf();
|
||||
self.store
|
||||
.resolve_and_validate(source_dir)
|
||||
.resolve_and_validate(&runtime_data_dir)
|
||||
.map_err(|_| ConfigError::InvalidField("store"))?;
|
||||
self.admin
|
||||
.resolve_and_validate(source_dir)
|
||||
@@ -37,7 +38,7 @@ impl LoadableConfig for GatewayConfig {
|
||||
.resolve_and_validate(source_dir)
|
||||
.map_err(|_| ConfigError::InvalidField("api"))?;
|
||||
self.openai
|
||||
.resolve_and_validate(source_dir)
|
||||
.resolve_and_validate(&runtime_data_dir)
|
||||
.map_err(|_| ConfigError::InvalidField("openai"))?;
|
||||
self.xai
|
||||
.resolve_and_validate(source_dir)
|
||||
|
||||
@@ -18,10 +18,18 @@ fn config_loader_should_load_complete_terminal_example() {
|
||||
fn config_loader_should_resolve_paths_relative_to_config_file() {
|
||||
let (config, _directory) = parse_config(&valid_config()).expect("resolved config");
|
||||
let debug = format!("{config:?}");
|
||||
assert!(debug.contains(".runtime/data"));
|
||||
assert!(debug.contains(".runtime/logs"));
|
||||
assert!(debug.contains("frontend/dist"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_loader_should_reject_missing_runtime_data_dir() {
|
||||
let config = valid_config().replace(" runtime_data_dir: '../.runtime/data'\n", "");
|
||||
|
||||
assert!(parse_config(&config).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_loader_should_inject_connection_passwords_into_urls() {
|
||||
parse_config(&valid_config()).expect("Store validates password injection into both URLs");
|
||||
@@ -135,7 +143,10 @@ fn config_loader_should_reject_removed_fingerprint_section() {
|
||||
|
||||
#[test]
|
||||
fn config_loader_should_reject_invalid_codex_cli_version() {
|
||||
assert_rejected(valid_config().replace("codex_version: '0.146.0'", "codex_version: 'latest'"));
|
||||
assert_rejected(valid_config().replace(
|
||||
"codex_version: '0.147.0-alpha.6.6'",
|
||||
"codex_version: 'latest'",
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -154,7 +165,7 @@ fn config_loader_should_reject_zero_server_port() {
|
||||
|
||||
#[test]
|
||||
fn config_loader_should_reject_invalid_desktop_profile_fields() {
|
||||
assert_rejected(valid_config().replace("desktop_build: '6119'", "desktop_build: 'build'"));
|
||||
assert_rejected(valid_config().replace("desktop_build: '6415'", "desktop_build: 'build'"));
|
||||
}
|
||||
|
||||
fn assert_rejected(config: String) {
|
||||
|
||||
@@ -570,13 +570,34 @@ pub enum QuotaLocalUsageAttribution {
|
||||
Unavailable,
|
||||
}
|
||||
|
||||
/// Provider quota bucket 中窗口的官方位置语义。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ProviderQuotaWindowRole {
|
||||
Primary,
|
||||
Secondary,
|
||||
Monthly,
|
||||
}
|
||||
|
||||
impl ProviderQuotaWindowRole {
|
||||
#[must_use]
|
||||
pub const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Primary => "primary",
|
||||
Self::Secondary => "secondary",
|
||||
Self::Monthly => "monthly",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 一个 Provider quota 窗口的公共投影。
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct ProviderQuotaWindow {
|
||||
pub key: String,
|
||||
pub group: String,
|
||||
pub label: String,
|
||||
pub source: Option<String>,
|
||||
pub limit_id: Option<String>,
|
||||
pub limit_name: Option<String>,
|
||||
pub role: Option<ProviderQuotaWindowRole>,
|
||||
pub local_usage_attribution: QuotaLocalUsageAttribution,
|
||||
pub window_seconds: Option<u64>,
|
||||
pub used_percent: Option<f64>,
|
||||
@@ -608,8 +629,9 @@ pub struct ProviderQuota {
|
||||
impl ProviderQuota {
|
||||
/// 返回账号展开面板使用的当前额度窗口用量。
|
||||
///
|
||||
/// 只选择具备完整时间边界且可按账号归属的窗口;优先级与 Dashboard
|
||||
/// 代表性额度窗口一致,同一优先级保持 Provider 投影顺序。
|
||||
/// Provider 已提供账号级本地用量时,固定窗口和无固定重置点的滚动窗口
|
||||
/// 都可作为代表用量;优先级与 Dashboard 代表性额度窗口一致,
|
||||
/// 同一优先级保持 Provider 投影顺序。
|
||||
#[must_use]
|
||||
pub fn representative_window_usage(&self) -> Option<&AccountUsage> {
|
||||
self.windows
|
||||
@@ -618,7 +640,6 @@ impl ProviderQuota {
|
||||
.filter(|(_, window)| {
|
||||
window.local_usage_attribution == QuotaLocalUsageAttribution::AccountWide
|
||||
&& window.window_seconds.is_some()
|
||||
&& window.reset_at.is_some()
|
||||
&& window.local_usage.is_some()
|
||||
})
|
||||
.min_by_key(|(index, window)| (quota_usage_priority(window), *index))
|
||||
@@ -631,25 +652,74 @@ impl ProviderQuota {
|
||||
/// 这里只解释跨 Provider 共享的窗口语义,绝不读取 Provider 私有 JSON。
|
||||
#[must_use]
|
||||
pub fn representative_used_percent(&self) -> Option<f64> {
|
||||
if self.limit_reached {
|
||||
return Some(100.0);
|
||||
}
|
||||
self.representative_used_window()
|
||||
.map(|(_, used_percent)| used_percent)
|
||||
}
|
||||
|
||||
/// 将已确认的账号级额度耗尽事实投影到展示窗口。
|
||||
///
|
||||
/// Provider 已指出具体触顶窗口时只归一化这些窗口;否则归一化 Dashboard
|
||||
/// 同样会选择的代表窗口,避免把多个独立额度窗口全部伪造成已用尽。
|
||||
pub fn apply_limit_reached_display(&mut self) {
|
||||
if !self.limit_reached {
|
||||
return;
|
||||
}
|
||||
let reached_windows = self
|
||||
.windows
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(index, window)| window.limit_reached.then_some(index))
|
||||
.collect::<Vec<_>>();
|
||||
if !reached_windows.is_empty() {
|
||||
for index in reached_windows {
|
||||
self.windows[index].used_percent = Some(100.0);
|
||||
}
|
||||
return;
|
||||
}
|
||||
let representative = self
|
||||
.representative_used_window()
|
||||
.map(|(index, _)| index)
|
||||
.or_else(|| {
|
||||
self.windows
|
||||
.iter()
|
||||
.enumerate()
|
||||
.min_by_key(|(index, window)| (quota_usage_priority(window), *index))
|
||||
.map(|(index, _)| index)
|
||||
});
|
||||
if let Some(index) = representative {
|
||||
self.windows[index].used_percent = Some(100.0);
|
||||
self.windows[index].limit_reached = true;
|
||||
}
|
||||
}
|
||||
|
||||
fn representative_used_window(&self) -> Option<(usize, f64)> {
|
||||
self.windows
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|window| {
|
||||
let (index, window) = window;
|
||||
let used_percent = window
|
||||
.used_percent
|
||||
.filter(|value| value.is_finite())
|
||||
.map(|value| value.clamp(0.0, 100.0))?;
|
||||
Some((quota_usage_priority(window), used_percent))
|
||||
Some((index, quota_usage_priority(window), used_percent))
|
||||
})
|
||||
.fold(None::<(u8, f64)>, |selected, candidate| match selected {
|
||||
Some(current)
|
||||
if current.0 < candidate.0
|
||||
|| (current.0 == candidate.0 && current.1 >= candidate.1) =>
|
||||
{
|
||||
.fold(
|
||||
None::<(usize, u8, f64)>,
|
||||
|selected, candidate| match selected {
|
||||
Some(current)
|
||||
}
|
||||
_ => Some(candidate),
|
||||
})
|
||||
.map(|(_, used_percent)| used_percent)
|
||||
if current.1 < candidate.1
|
||||
|| (current.1 == candidate.1 && current.2 >= candidate.2) =>
|
||||
{
|
||||
Some(current)
|
||||
}
|
||||
_ => Some(candidate),
|
||||
},
|
||||
)
|
||||
.map(|(index, _, used_percent)| (index, used_percent))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -31,23 +31,51 @@ pub enum ProviderAdminErrorKind {
|
||||
Internal,
|
||||
}
|
||||
|
||||
/// 不泄漏 OAuth material 与 Provider 响应体的管理错误。
|
||||
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
|
||||
/// 不携带 OAuth 请求材料的管理错误。
|
||||
///
|
||||
/// `message` 用于当前认证管理请求展示;上游没有标准 `error.message` 时,
|
||||
/// Provider 可以按其官方契约回退为完整的非成功响应正文。
|
||||
#[derive(Clone, PartialEq, Eq, thiserror::Error)]
|
||||
#[error("provider admin operation failed: {kind:?}")]
|
||||
pub struct ProviderAdminError {
|
||||
kind: ProviderAdminErrorKind,
|
||||
message: Option<String>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for ProviderAdminError {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter
|
||||
.debug_struct("ProviderAdminError")
|
||||
.field("kind", &self.kind)
|
||||
.field("message", &self.message.as_ref().map(|_| "<redacted>"))
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderAdminError {
|
||||
#[must_use]
|
||||
pub const fn new(kind: ProviderAdminErrorKind) -> Self {
|
||||
Self { kind }
|
||||
Self {
|
||||
kind,
|
||||
message: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_message(mut self, message: impl Into<String>) -> Self {
|
||||
self.message = Some(message.into());
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn kind(&self) -> ProviderAdminErrorKind {
|
||||
self.kind
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn message(&self) -> Option<&str> {
|
||||
self.message.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
/// 一个具体 Provider 对管理控制面提供的解析、验证、上游交互与运行时资源回收能力。
|
||||
|
||||
@@ -565,11 +565,7 @@ impl AccountsService for DefaultAccountsService {
|
||||
let upstream_body = error
|
||||
.upstream_response()
|
||||
.map(|response| String::from_utf8_lossy(response.body()).into_owned());
|
||||
let message = upstream_body
|
||||
.as_deref()
|
||||
.filter(|body| !body.is_empty())
|
||||
.unwrap_or_else(|| error.client_message())
|
||||
.to_owned();
|
||||
let message = error.client_message().to_owned();
|
||||
vec![AccountConnectionTestEvent::Failed {
|
||||
message,
|
||||
provider_error_code: error.client_error_code().map(ToOwned::to_owned),
|
||||
|
||||
@@ -59,7 +59,11 @@ fn map_provider_error(
|
||||
ProviderAdminErrorKind::Unavailable => AdminErrorKind::Unavailable,
|
||||
ProviderAdminErrorKind::Internal => AdminErrorKind::Internal,
|
||||
};
|
||||
AdminError::new(kind, format!("{resource} operation failed"))
|
||||
let message = error
|
||||
.message()
|
||||
.map(ToOwned::to_owned)
|
||||
.unwrap_or_else(|| format!("{resource} operation failed"));
|
||||
AdminError::new(kind, message)
|
||||
}
|
||||
|
||||
async fn publish_committed(
|
||||
|
||||
@@ -6,6 +6,7 @@ use gateway_core::engine::credential::OpaqueProviderData;
|
||||
use gateway_admin::{
|
||||
model::{
|
||||
PageSize, Revision,
|
||||
accounts::AccountUsage,
|
||||
auth::LoginCommand,
|
||||
client_keys::ClientKeyPageSize,
|
||||
observability::{DecimalAmount, RequestOutcome, TimeRange},
|
||||
@@ -135,6 +136,97 @@ fn representative_quota_should_prefer_short_window_and_highest_usage() {
|
||||
assert_eq!(quota.representative_used_percent(), Some(45.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn representative_window_usage_should_accept_account_wide_rolling_window_without_reset() {
|
||||
let usage = AccountUsage {
|
||||
account_id: "acct_xai".to_owned(),
|
||||
request_count: 5,
|
||||
success_count: 5,
|
||||
input_tokens: Some(100),
|
||||
output_tokens: Some(20),
|
||||
cached_tokens: Some(80),
|
||||
cache_write_tokens: Some(0),
|
||||
reasoning_tokens: Some(0),
|
||||
image_input_tokens: Some(0),
|
||||
image_output_tokens: Some(0),
|
||||
image_request_count: 0,
|
||||
image_request_failed_count: 0,
|
||||
total_tokens: Some(120),
|
||||
cost_coverage: Default::default(),
|
||||
costs: Vec::new(),
|
||||
last_used_at: Some(Utc::now()),
|
||||
request_buckets: Vec::new(),
|
||||
models: Vec::new(),
|
||||
};
|
||||
let quota = ProviderQuota {
|
||||
observed_at: None,
|
||||
refresh_token_expires_at: None,
|
||||
windows: vec![ProviderQuotaWindow {
|
||||
key: "free-rolling-24h".to_owned(),
|
||||
group: "shortTerm".to_owned(),
|
||||
label: "日限额".to_owned(),
|
||||
limit_id: None,
|
||||
limit_name: None,
|
||||
role: None,
|
||||
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
|
||||
window_seconds: Some(86_400),
|
||||
used_percent: None,
|
||||
reset_at: None,
|
||||
limit_reached: false,
|
||||
local_usage: Some(usage.clone()),
|
||||
provider_data: None,
|
||||
}],
|
||||
limit_reached: false,
|
||||
provider_data: None,
|
||||
};
|
||||
|
||||
assert_eq!(quota.representative_window_usage(), Some(&usage));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exhausted_quota_should_project_full_usage_to_only_the_representative_window() {
|
||||
let mut quota = ProviderQuota {
|
||||
observed_at: None,
|
||||
refresh_token_expires_at: None,
|
||||
windows: vec![
|
||||
quota_window("monthly", Some(2_592_000), Some(99.0)),
|
||||
quota_window("shortTerm", Some(604_800), Some(80.0)),
|
||||
quota_window("shortTerm", Some(18_000), Some(30.0)),
|
||||
quota_window("shortTerm", Some(14_400), Some(95.0)),
|
||||
],
|
||||
limit_reached: true,
|
||||
provider_data: None,
|
||||
};
|
||||
|
||||
assert_eq!(quota.representative_used_percent(), Some(100.0));
|
||||
quota.apply_limit_reached_display();
|
||||
|
||||
assert_eq!(quota.windows[0].used_percent, Some(99.0));
|
||||
assert_eq!(quota.windows[1].used_percent, Some(80.0));
|
||||
assert_eq!(quota.windows[2].used_percent, Some(30.0));
|
||||
assert_eq!(quota.windows[3].used_percent, Some(100.0));
|
||||
assert!(quota.windows[3].limit_reached);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exhausted_quota_should_preserve_the_provider_identified_reached_window() {
|
||||
let mut reached = quota_window("monthly", Some(2_592_000), Some(98.0));
|
||||
reached.limit_reached = true;
|
||||
let mut quota = ProviderQuota {
|
||||
observed_at: None,
|
||||
refresh_token_expires_at: None,
|
||||
windows: vec![reached, quota_window("shortTerm", Some(18_000), Some(95.0))],
|
||||
limit_reached: true,
|
||||
provider_data: None,
|
||||
};
|
||||
|
||||
quota.apply_limit_reached_display();
|
||||
|
||||
assert_eq!(quota.windows[0].used_percent, Some(100.0));
|
||||
assert_eq!(quota.windows[1].used_percent, Some(95.0));
|
||||
assert!(!quota.windows[1].limit_reached);
|
||||
}
|
||||
|
||||
fn quota_window(
|
||||
group: &str,
|
||||
window_seconds: Option<u64>,
|
||||
@@ -144,7 +236,9 @@ fn quota_window(
|
||||
key: group.to_owned(),
|
||||
group: group.to_owned(),
|
||||
label: group.to_owned(),
|
||||
source: None,
|
||||
limit_id: None,
|
||||
limit_name: None,
|
||||
role: None,
|
||||
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
|
||||
window_seconds,
|
||||
used_percent,
|
||||
|
||||
@@ -60,7 +60,7 @@ pub(super) type EventLog = Arc<Mutex<Vec<&'static str>>>;
|
||||
pub(super) struct FakeProviderAdmin {
|
||||
kind: ProviderKind,
|
||||
events: EventLog,
|
||||
failure: Mutex<Option<ProviderAdminErrorKind>>,
|
||||
failure: Mutex<Option<ProviderAdminError>>,
|
||||
quota_failure: Mutex<Option<ProviderAdminErrorKind>>,
|
||||
pending: Arc<Mutex<Option<PendingAuthorizationMutation>>>,
|
||||
retry_authorization_after_abort: Mutex<bool>,
|
||||
@@ -89,7 +89,12 @@ impl FakeProviderAdmin {
|
||||
}
|
||||
|
||||
pub(super) fn fail_next(&self, kind: ProviderAdminErrorKind) {
|
||||
*self.failure.lock().expect("provider failure") = Some(kind);
|
||||
*self.failure.lock().expect("provider failure") = Some(ProviderAdminError::new(kind));
|
||||
}
|
||||
|
||||
fn fail_next_with_message(&self, kind: ProviderAdminErrorKind, message: &str) {
|
||||
*self.failure.lock().expect("provider failure") =
|
||||
Some(ProviderAdminError::new(kind).with_message(message));
|
||||
}
|
||||
|
||||
pub(super) fn fail_next_quota(&self, kind: ProviderAdminErrorKind) {
|
||||
@@ -145,7 +150,7 @@ impl FakeProviderAdmin {
|
||||
|
||||
fn require_available(&self) -> Result<(), ProviderAdminError> {
|
||||
match self.failure.lock().expect("provider failure").take() {
|
||||
Some(kind) => Err(ProviderAdminError::new(kind)),
|
||||
Some(error) => Err(error),
|
||||
None => Ok(()),
|
||||
}
|
||||
}
|
||||
@@ -1257,7 +1262,9 @@ async fn accounts_list_should_degrade_quota_failure_to_empty_window_without_drop
|
||||
key: "primary".to_owned(),
|
||||
group: "shortTerm".to_owned(),
|
||||
label: "5小时限额".to_owned(),
|
||||
source: None,
|
||||
limit_id: None,
|
||||
limit_name: None,
|
||||
role: None,
|
||||
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
|
||||
window_seconds: Some(5 * 60 * 60),
|
||||
used_percent: Some(97.0),
|
||||
@@ -1386,7 +1393,9 @@ async fn accounts_list_should_not_derive_rate_limited_from_provider_quota_view()
|
||||
key: "primary".to_owned(),
|
||||
group: "shortTerm".to_owned(),
|
||||
label: "5小时限额".to_owned(),
|
||||
source: None,
|
||||
limit_id: None,
|
||||
limit_name: None,
|
||||
role: None,
|
||||
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
|
||||
window_seconds: Some(5 * 60 * 60),
|
||||
used_percent: Some(100.0),
|
||||
@@ -1431,7 +1440,9 @@ async fn accounts_list_should_not_derive_exhaustion_from_provider_quota_view() {
|
||||
key: "primary".to_owned(),
|
||||
group: "shortTerm".to_owned(),
|
||||
label: "5小时限额".to_owned(),
|
||||
source: None,
|
||||
limit_id: None,
|
||||
limit_name: None,
|
||||
role: None,
|
||||
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
|
||||
window_seconds: Some(5 * 60 * 60),
|
||||
used_percent: Some(100.0),
|
||||
@@ -1477,7 +1488,9 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
|
||||
key: "primary".to_owned(),
|
||||
group: "shortTerm".to_owned(),
|
||||
label: "5小时限额".to_owned(),
|
||||
source: None,
|
||||
limit_id: None,
|
||||
limit_name: None,
|
||||
role: None,
|
||||
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
|
||||
window_seconds: Some(5 * 60 * 60),
|
||||
used_percent: Some(97.0),
|
||||
@@ -1542,7 +1555,9 @@ async fn accounts_list_should_not_attach_account_usage_to_model_specific_quota_w
|
||||
key: "core-primary".to_owned(),
|
||||
group: "weekly".to_owned(),
|
||||
label: "周额度".to_owned(),
|
||||
source: Some("codex".to_owned()),
|
||||
limit_id: Some("codex".to_owned()),
|
||||
limit_name: None,
|
||||
role: None,
|
||||
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
|
||||
window_seconds: Some(7 * 24 * 60 * 60),
|
||||
used_percent: Some(1.0),
|
||||
@@ -1554,8 +1569,10 @@ async fn accounts_list_should_not_attach_account_usage_to_model_specific_quota_w
|
||||
ProviderQuotaWindow {
|
||||
key: "codex-bengalfox-primary".to_owned(),
|
||||
group: "weekly".to_owned(),
|
||||
label: "GPT-5.3-Codex-Spark · 周额度".to_owned(),
|
||||
source: Some("codex_bengalfox".to_owned()),
|
||||
label: "周额度".to_owned(),
|
||||
limit_id: Some("codex_bengalfox".to_owned()),
|
||||
limit_name: Some("GPT-5.3-Codex-Spark".to_owned()),
|
||||
role: None,
|
||||
local_usage_attribution: QuotaLocalUsageAttribution::Unavailable,
|
||||
window_seconds: Some(7 * 24 * 60 * 60),
|
||||
used_percent: Some(0.0),
|
||||
@@ -1650,6 +1667,32 @@ async fn accounts_refresh_provider_failure_should_not_call_store_commit() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_refresh_should_preserve_the_provider_failure_message() {
|
||||
let events = events();
|
||||
let provider = FakeProviderAdmin::new("openai", events.clone());
|
||||
let upstream_message = "Your refresh token has already been used.";
|
||||
provider.fail_next_with_message(ProviderAdminErrorKind::Conflict, upstream_message);
|
||||
let store = FakeAccountStore::new("openai", events.clone());
|
||||
let services = accounts_service(provider, store).await;
|
||||
|
||||
let error = services
|
||||
.accounts()
|
||||
.refresh(
|
||||
&context("refresh-provider-message"),
|
||||
ProviderAccountId::new("acct_test").expect("account ID"),
|
||||
)
|
||||
.await
|
||||
.expect_err("Provider preparation must fail");
|
||||
|
||||
assert_eq!(error.kind(), gateway_admin::model::AdminErrorKind::Conflict);
|
||||
assert_eq!(error.to_string(), upstream_message);
|
||||
assert_eq!(
|
||||
recorded(&events),
|
||||
["store.load_account", "provider.prepare_refresh"]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_refresh_store_failure_should_drop_guard_after_commit_attempt() {
|
||||
let events = events();
|
||||
|
||||
@@ -305,8 +305,12 @@ pub struct AccountQuotaView {
|
||||
pub struct AccountQuotaWindowView {
|
||||
pub key: String,
|
||||
pub group: String,
|
||||
pub limit_id: Option<String>,
|
||||
pub limit_name: Option<String>,
|
||||
pub role: Option<String>,
|
||||
pub window_seconds: Option<u64>,
|
||||
pub label_display: String,
|
||||
pub window_label_display: String,
|
||||
pub used_percent: Option<f64>,
|
||||
pub used_percent_display: String,
|
||||
pub limit_reached: bool,
|
||||
@@ -1574,10 +1578,11 @@ fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
|
||||
}
|
||||
|
||||
fn account_quota_view(
|
||||
quota: ProviderQuota,
|
||||
mut quota: ProviderQuota,
|
||||
rate_limited_until: Option<DateTime<Utc>>,
|
||||
now: DateTime<Utc>,
|
||||
) -> (AccountQuotaView, Option<String>) {
|
||||
quota.apply_limit_reached_display();
|
||||
let refresh_token_expires_at = quota
|
||||
.refresh_token_expires_at
|
||||
.map(|value| china_rfc3339(&value));
|
||||
@@ -1602,7 +1607,9 @@ fn quota_window_view(window: ProviderQuotaWindow) -> AccountQuotaWindowView {
|
||||
key,
|
||||
group,
|
||||
label,
|
||||
source: _,
|
||||
limit_id,
|
||||
limit_name,
|
||||
role,
|
||||
local_usage_attribution: _,
|
||||
window_seconds,
|
||||
used_percent,
|
||||
@@ -1611,10 +1618,17 @@ fn quota_window_view(window: ProviderQuotaWindow) -> AccountQuotaWindowView {
|
||||
local_usage,
|
||||
provider_data: _,
|
||||
} = window;
|
||||
let label_display = limit_name
|
||||
.as_ref()
|
||||
.map_or_else(|| label.clone(), |name| format!("{name} · {label}"));
|
||||
AccountQuotaWindowView {
|
||||
label_display: label,
|
||||
label_display,
|
||||
window_label_display: label,
|
||||
key,
|
||||
group,
|
||||
limit_id,
|
||||
limit_name,
|
||||
role: role.map(|value| value.as_str().to_owned()),
|
||||
window_seconds,
|
||||
used_percent,
|
||||
used_percent_display: used_percent
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//! `/healthz` 对 Core、Store 与 Host worker 健康事实的聚合。
|
||||
//! `/healthz` 对 Core、Store 与 Host 关键 worker 健康事实的聚合。
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
@@ -6,8 +6,10 @@ use std::time::Duration;
|
||||
use axum::{extract::State, http::StatusCode};
|
||||
use futures::future::join_all;
|
||||
use gateway_core::health::{
|
||||
HealthProbe, HealthState, WorkerHealthSnapshot, WorkerHealthSource, WorkerRuntimeState,
|
||||
HealthProbe, HealthState, WorkerHealthKey, WorkerHealthSnapshot, WorkerHealthSource,
|
||||
WorkerRuntimeState,
|
||||
};
|
||||
use gateway_core::task::WorkerKind;
|
||||
|
||||
use crate::ApiState;
|
||||
|
||||
@@ -32,7 +34,13 @@ impl HealthStatus {
|
||||
}
|
||||
|
||||
pub(crate) async fn healthy(&self) -> bool {
|
||||
if !self.workers.snapshot().iter().all(worker_is_healthy) {
|
||||
if !self
|
||||
.workers
|
||||
.snapshot()
|
||||
.iter()
|
||||
.filter(|worker| worker_affects_healthz(worker))
|
||||
.all(worker_is_healthy)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
let checks = self.probes.iter().map(|probe| probe.check());
|
||||
@@ -54,6 +62,16 @@ pub(crate) async fn healthz(State(state): State<ApiState>) -> StatusCode {
|
||||
}
|
||||
}
|
||||
|
||||
fn worker_affects_healthz(worker: &WorkerHealthSnapshot) -> bool {
|
||||
let WorkerHealthKey::Task(id) = &worker.key else {
|
||||
return false;
|
||||
};
|
||||
!matches!(
|
||||
id.kind(),
|
||||
WorkerKind::OAuthRefresh | WorkerKind::QuotaCatalogHealth
|
||||
)
|
||||
}
|
||||
|
||||
fn worker_is_healthy(worker: &WorkerHealthSnapshot) -> bool {
|
||||
match worker.state {
|
||||
WorkerRuntimeState::Disabled | WorkerRuntimeState::Standby => true,
|
||||
|
||||
@@ -73,7 +73,6 @@ pub(super) fn response_header_is_forwardable(name: &str, connection_options: &[S
|
||||
| "x-openai-organization"
|
||||
| "x-openai-project"
|
||||
| "x-codex-installation-id"
|
||||
| "x-codex-turn-state"
|
||||
| "x-codex-turn-metadata"
|
||||
)
|
||||
}
|
||||
|
||||
@@ -11,7 +11,9 @@ use gateway_core::{
|
||||
StartedExecution,
|
||||
},
|
||||
error::GatewayError,
|
||||
health::{WorkerHealthKey, WorkerHealthSnapshot, WorkerHealthSource, WorkerRuntimeState},
|
||||
routing::PublicModelId,
|
||||
task::{WorkerId, WorkerKind},
|
||||
};
|
||||
use tower::ServiceExt;
|
||||
|
||||
@@ -31,6 +33,75 @@ async fn healthz_should_return_no_content_when_all_inputs_are_healthy() {
|
||||
assert_eq!(response.status(), StatusCode::NO_CONTENT);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn healthz_should_ignore_provider_maintenance_worker_failure() {
|
||||
let worker_health = StaticWorkerHealth(vec![worker_snapshot(
|
||||
WorkerKind::QuotaCatalogHealth,
|
||||
WorkerRuntimeState::BackingOff,
|
||||
)]);
|
||||
let response = crate::openai::api_router_with_worker_health(
|
||||
Arc::new(UnusedExecution),
|
||||
Arc::new(worker_health),
|
||||
)
|
||||
.await
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.uri("/healthz")
|
||||
.body(Body::empty())
|
||||
.expect("health request"),
|
||||
)
|
||||
.await
|
||||
.expect("health response");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::NO_CONTENT);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn healthz_should_reject_critical_worker_failure() {
|
||||
let worker_health = StaticWorkerHealth(vec![worker_snapshot(
|
||||
WorkerKind::RuntimeSnapshotReconciliation,
|
||||
WorkerRuntimeState::BackingOff,
|
||||
)]);
|
||||
let response = crate::openai::api_router_with_worker_health(
|
||||
Arc::new(UnusedExecution),
|
||||
Arc::new(worker_health),
|
||||
)
|
||||
.await
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.uri("/healthz")
|
||||
.body(Body::empty())
|
||||
.expect("health request"),
|
||||
)
|
||||
.await
|
||||
.expect("health response");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct StaticWorkerHealth(Vec<WorkerHealthSnapshot>);
|
||||
|
||||
impl WorkerHealthSource for StaticWorkerHealth {
|
||||
fn snapshot(&self) -> Vec<WorkerHealthSnapshot> {
|
||||
self.0.clone()
|
||||
}
|
||||
}
|
||||
|
||||
fn worker_snapshot(kind: WorkerKind, state: WorkerRuntimeState) -> WorkerHealthSnapshot {
|
||||
let id = WorkerId::try_new(kind, "health-test").expect("worker ID");
|
||||
WorkerHealthSnapshot {
|
||||
key: WorkerHealthKey::Task(id),
|
||||
state,
|
||||
consecutive_failures: 1,
|
||||
completed_cycles: 0,
|
||||
last_fencing_token: None,
|
||||
last_success_at: None,
|
||||
last_failure_at: None,
|
||||
last_error: Some("test worker failure".to_owned()),
|
||||
}
|
||||
}
|
||||
|
||||
struct UnusedExecution;
|
||||
|
||||
impl ExecutionService for UnusedExecution {
|
||||
|
||||
@@ -44,9 +44,29 @@ pub(super) async fn api_router(execution: Arc<dyn ExecutionService>) -> axum::Ro
|
||||
api_router_with_origins(execution, Vec::new()).await
|
||||
}
|
||||
|
||||
pub(super) async fn api_router_with_worker_health(
|
||||
execution: Arc<dyn ExecutionService>,
|
||||
worker_health: Arc<dyn WorkerHealthSource>,
|
||||
) -> axum::Router {
|
||||
api_router_with_origins_and_worker_health(execution, Vec::new(), worker_health).await
|
||||
}
|
||||
|
||||
pub(super) async fn api_router_with_origins(
|
||||
execution: Arc<dyn ExecutionService>,
|
||||
cors_allowed_origins: Vec<String>,
|
||||
) -> axum::Router {
|
||||
api_router_with_origins_and_worker_health(
|
||||
execution,
|
||||
cors_allowed_origins,
|
||||
Arc::new(EmptyWorkerHealth),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn api_router_with_origins_and_worker_health(
|
||||
execution: Arc<dyn ExecutionService>,
|
||||
cors_allowed_origins: Vec<String>,
|
||||
worker_health: Arc<dyn WorkerHealthSource>,
|
||||
) -> axum::Router {
|
||||
let admin = crate::admin::AdminTestFixture::new().await;
|
||||
gateway_api::initialize(
|
||||
@@ -59,7 +79,7 @@ pub(super) async fn api_router_with_origins(
|
||||
execution,
|
||||
admin.services,
|
||||
Vec::new(),
|
||||
Arc::new(EmptyWorkerHealth),
|
||||
worker_health,
|
||||
Arc::new(TestLifecycle::default()),
|
||||
)
|
||||
.expect("API bundle")
|
||||
|
||||
@@ -788,6 +788,10 @@ async fn streaming_response_should_append_ordinary_headers_and_skip_unrepresenta
|
||||
let trace = Arc::new(Trace::default());
|
||||
let headers = vec![
|
||||
ProviderResponseHeader::new("x-models-etag", Bytes::from_static(b"models-v2")),
|
||||
ProviderResponseHeader::new(
|
||||
"x-codex-turn-state",
|
||||
Bytes::from_static(b"turn-state-from-upstream"),
|
||||
),
|
||||
ProviderResponseHeader::new("x-future-multi", Bytes::from_static(b"first")),
|
||||
ProviderResponseHeader::new("x-future-multi", Bytes::from_static(b"second")),
|
||||
ProviderResponseHeader::new("x-future-bytes", Bytes::from_static(b"\xffopaque")),
|
||||
@@ -834,6 +838,13 @@ async fn streaming_response_should_append_ordinary_headers_and_skip_unrepresenta
|
||||
.map(HeaderValue::as_bytes),
|
||||
Some(b"\xffopaque".as_slice())
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get("x-codex-turn-state")
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("turn-state-from-upstream")
|
||||
);
|
||||
assert!(response.headers().get("authorization").is_none());
|
||||
assert!(response.headers().get("connection").is_none());
|
||||
assert!(response.headers().get("x-hop-secret").is_none());
|
||||
|
||||
@@ -411,6 +411,10 @@ async fn websocket_response_metadata_should_preserve_ordinary_headers_without_bl
|
||||
ProviderResponseHeader::new("x-future-header", Bytes::from_static(b"future-value")),
|
||||
ProviderResponseHeader::new("x-future-multi", Bytes::from_static(b"first")),
|
||||
ProviderResponseHeader::new("x-future-multi", Bytes::from_static(b"second")),
|
||||
ProviderResponseHeader::new(
|
||||
"x-codex-turn-state",
|
||||
Bytes::from_static(b"turn-state-from-upstream"),
|
||||
),
|
||||
ProviderResponseHeader::new("x-future-bytes", Bytes::from_static(b"\xffopaque")),
|
||||
ProviderResponseHeader::new("bad\0name", Bytes::from_static(b"unrepresentable")),
|
||||
ProviderResponseHeader::new(
|
||||
@@ -483,6 +487,10 @@ async fn websocket_response_metadata_should_preserve_ordinary_headers_without_bl
|
||||
// 官方 metadata 是 string map,同名多值只能保持既有的后值覆盖语义。
|
||||
assert_eq!(headers.get("x-future-multi"), Some(&json!("second")));
|
||||
assert_eq!(headers.get("x-request-id"), Some(&json!("req_upstream")));
|
||||
assert_eq!(
|
||||
headers.get("x-codex-turn-state"),
|
||||
Some(&json!("turn-state-from-upstream"))
|
||||
);
|
||||
for omitted in [
|
||||
"x-future-bytes",
|
||||
"bad\0name",
|
||||
|
||||
@@ -1162,10 +1162,11 @@ where
|
||||
true
|
||||
}
|
||||
|
||||
/// 原生续写的原账号在本地调度阶段已不可用时,改用完整输入重放。
|
||||
/// 原生续写的原账号在本地调度阶段已不可用时,交给 Provider 执行跨账号恢复。
|
||||
///
|
||||
/// 这不是强制 Smart:下一次 attempt 仍使用请求配置的调度策略。只有网关持有
|
||||
/// 对应 Provider 的可携带会话状态时,才可安全移除 native handle 并换号。
|
||||
/// 对应 Provider 的会话状态时才进入恢复;是否保留 native handle、执行 probe
|
||||
/// 或使用完整 transcript,由 Provider 自己的协议边界决定。
|
||||
fn prepare_unavailable_native_continuation_replay(&mut self, error: &ProviderError) -> bool {
|
||||
if self.account_selection.required_account().is_some()
|
||||
|| self.continuation_attempt != ContinuationAttempt::Native
|
||||
|
||||
@@ -946,7 +946,28 @@ impl fmt::Debug for LoadedCredential {
|
||||
}
|
||||
}
|
||||
|
||||
/// 刷新成功后的完整 CAS 写回。
|
||||
/// 与 credential revision CAS 同事务提交的账号错误事实。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CredentialStateWrite {
|
||||
pub credential_state: CredentialState,
|
||||
pub observed_at: SystemTime,
|
||||
pub error_reason: Option<AccountErrorReason>,
|
||||
pub message: Option<String>,
|
||||
}
|
||||
|
||||
/// [`CredentialCasUpdate`] 跨 store 边界的命名字段。
|
||||
pub struct CredentialCasUpdateParts {
|
||||
pub account_id: ProviderAccountId,
|
||||
pub expected_revision: CredentialRevision,
|
||||
pub profile: ProviderAccountUpdate,
|
||||
pub credential: PlaintextCredential,
|
||||
pub has_refresh_token: bool,
|
||||
pub access_token_expires_at: Option<SystemTime>,
|
||||
pub next_refresh_at: Option<SystemTime>,
|
||||
pub account_state: Option<Box<CredentialStateWrite>>,
|
||||
}
|
||||
|
||||
/// 刷新后的完整 CAS 写回。
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct CredentialCasUpdate {
|
||||
account_id: ProviderAccountId,
|
||||
@@ -956,6 +977,7 @@ pub struct CredentialCasUpdate {
|
||||
has_refresh_token: bool,
|
||||
access_token_expires_at: Option<SystemTime>,
|
||||
next_refresh_at: Option<SystemTime>,
|
||||
account_state: Option<Box<CredentialStateWrite>>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for CredentialCasUpdate {
|
||||
@@ -969,6 +991,7 @@ impl fmt::Debug for CredentialCasUpdate {
|
||||
.field("has_refresh_token", &self.has_refresh_token)
|
||||
.field("access_token_expires_at", &self.access_token_expires_at)
|
||||
.field("next_refresh_at", &self.next_refresh_at)
|
||||
.field("account_state", &self.account_state)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -1002,9 +1025,34 @@ impl CredentialCasUpdate {
|
||||
has_refresh_token,
|
||||
access_token_expires_at,
|
||||
next_refresh_at,
|
||||
account_state: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// 将刷新调度与账号错误事实放入同一个 revision CAS。
|
||||
#[must_use]
|
||||
pub fn with_account_state(
|
||||
mut self,
|
||||
credential_state: CredentialState,
|
||||
observed_at: SystemTime,
|
||||
error_reason: Option<AccountErrorReason>,
|
||||
message: Option<String>,
|
||||
) -> Self {
|
||||
let message = message.filter(|value| !value.trim().is_empty());
|
||||
let error_reason = if credential_state == CredentialState::Ready {
|
||||
message.as_ref().and(error_reason)
|
||||
} else {
|
||||
error_reason.or_else(|| credential_state.error_reason())
|
||||
};
|
||||
self.account_state = Some(Box::new(CredentialStateWrite {
|
||||
credential_state,
|
||||
observed_at,
|
||||
error_reason,
|
||||
message,
|
||||
}));
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn account_id(&self) -> &ProviderAccountId {
|
||||
&self.account_id
|
||||
@@ -1041,26 +1089,17 @@ impl CredentialCasUpdate {
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn into_parts(
|
||||
self,
|
||||
) -> (
|
||||
ProviderAccountId,
|
||||
CredentialRevision,
|
||||
ProviderAccountUpdate,
|
||||
PlaintextCredential,
|
||||
bool,
|
||||
Option<SystemTime>,
|
||||
Option<SystemTime>,
|
||||
) {
|
||||
(
|
||||
self.account_id,
|
||||
self.expected_revision,
|
||||
self.profile,
|
||||
self.credential,
|
||||
self.has_refresh_token,
|
||||
self.access_token_expires_at,
|
||||
self.next_refresh_at,
|
||||
)
|
||||
pub fn into_parts(self) -> CredentialCasUpdateParts {
|
||||
CredentialCasUpdateParts {
|
||||
account_id: self.account_id,
|
||||
expected_revision: self.expected_revision,
|
||||
profile: self.profile,
|
||||
credential: self.credential,
|
||||
has_refresh_token: self.has_refresh_token,
|
||||
access_token_expires_at: self.access_token_expires_at,
|
||||
next_refresh_at: self.next_refresh_at,
|
||||
account_state: self.account_state,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1102,6 +1141,14 @@ pub enum QuotaWriteOutcome {
|
||||
Conflict,
|
||||
}
|
||||
|
||||
/// 只推进 Provider quota 文档的最后成功查询时间。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct QuotaObservationTouch {
|
||||
pub account_id: ProviderAccountId,
|
||||
pub expected_revision: CredentialRevision,
|
||||
pub observed_at: SystemTime,
|
||||
}
|
||||
|
||||
/// 不改 Provider 原始 quota JSON 的额度访问事实写入。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct QuotaAccessChange {
|
||||
@@ -1117,9 +1164,13 @@ pub struct AccountStateChange {
|
||||
pub expected_revision: CredentialRevision,
|
||||
pub credential_state: CredentialState,
|
||||
pub observed_at: SystemTime,
|
||||
/// 受控错误原因;恢复 Ready 时必须清空。
|
||||
/// 受控错误原因;无失败事实的 Ready 写入必须清空。
|
||||
///
|
||||
/// Ready 凭据可以在 AT 过期后继续 RT 恢复,此时允许保留最近一次
|
||||
/// 刷新失败原因,供错误状态投影展示。
|
||||
pub error_reason: Option<AccountErrorReason>,
|
||||
/// 已脱敏的上游诊断信息;恢复 Ready 时必须清空。
|
||||
/// 供管理端展示的错误消息;结构化上游失败应保留原始 message,不能写入整个正文。
|
||||
/// 刷新成功或其他无失败事实的 Ready 写入必须清空。
|
||||
pub message: Option<String>,
|
||||
}
|
||||
|
||||
@@ -1170,6 +1221,11 @@ pub trait ProviderAccountStore: Send + Sync {
|
||||
observation: QuotaObservation,
|
||||
) -> Result<QuotaWriteOutcome, StoreError>;
|
||||
|
||||
async fn touch_quota_observation(
|
||||
&self,
|
||||
touch: QuotaObservationTouch,
|
||||
) -> Result<QuotaWriteOutcome, StoreError>;
|
||||
|
||||
async fn apply_quota_access(
|
||||
&self,
|
||||
change: QuotaAccessChange,
|
||||
@@ -1588,14 +1644,26 @@ impl AccountSelector {
|
||||
|
||||
let candidate = match context.policy.strategy() {
|
||||
RotationStrategy::QuotaResetPriority => {
|
||||
eligible.sort_by_key(|candidate| {
|
||||
let default_concurrency = context.policy.max_concurrent_per_account();
|
||||
eligible.sort_by(|left, right| {
|
||||
(
|
||||
candidate.signals.quota_reset_at.is_none(),
|
||||
candidate.signals.quota_reset_at,
|
||||
candidate.signals.in_flight,
|
||||
candidate.signals.last_started_at,
|
||||
candidate.account.id().clone(),
|
||||
left.signals.quota_reset_at.is_none(),
|
||||
left.signals.quota_reset_at,
|
||||
)
|
||||
.cmp(&(
|
||||
right.signals.quota_reset_at.is_none(),
|
||||
right.signals.quota_reset_at,
|
||||
))
|
||||
.then_with(|| {
|
||||
capacity_utilization(left, default_concurrency)
|
||||
.total_cmp(&capacity_utilization(right, default_concurrency))
|
||||
})
|
||||
.then_with(|| {
|
||||
left.signals
|
||||
.last_started_at
|
||||
.cmp(&right.signals.last_started_at)
|
||||
})
|
||||
.then_with(|| left.account.id().cmp(right.account.id()))
|
||||
});
|
||||
eligible.first().copied()?
|
||||
}
|
||||
@@ -1604,9 +1672,11 @@ impl AccountSelector {
|
||||
let index = context.round_robin_cursor as usize % eligible.len();
|
||||
eligible.get(index).copied()?
|
||||
}
|
||||
RotationStrategy::Smart => {
|
||||
select_smart_candidate(&eligible, context.round_robin_cursor)?
|
||||
}
|
||||
RotationStrategy::Smart => select_smart_candidate(
|
||||
&eligible,
|
||||
context.policy.max_concurrent_per_account(),
|
||||
context.round_robin_cursor,
|
||||
)?,
|
||||
RotationStrategy::Sticky => {
|
||||
eligible.sort_by_key(|candidate| {
|
||||
(
|
||||
@@ -1676,8 +1746,18 @@ const SMART_QUOTA_WEIGHT: f64 = 0.8;
|
||||
const SMART_FAILURE_WEIGHT: f64 = 1.0;
|
||||
const SMART_LATENCY_WEIGHT: f64 = 0.5;
|
||||
|
||||
fn capacity_utilization(candidate: &AccountCandidate, default_concurrency: NonZeroU32) -> f64 {
|
||||
f64::from(candidate.signals.in_flight)
|
||||
/ f64::from(
|
||||
candidate
|
||||
.account
|
||||
.effective_concurrency(default_concurrency)
|
||||
.get(),
|
||||
)
|
||||
}
|
||||
|
||||
struct SmartNormalization {
|
||||
max_in_flight: u32,
|
||||
max_capacity_utilization: f64,
|
||||
min_quota: Option<u64>,
|
||||
max_quota: Option<u64>,
|
||||
min_latency_ms: Option<u64>,
|
||||
@@ -1685,11 +1765,11 @@ struct SmartNormalization {
|
||||
}
|
||||
|
||||
impl SmartNormalization {
|
||||
fn from_candidates(candidates: &[&AccountCandidate]) -> Self {
|
||||
let max_in_flight = candidates
|
||||
fn from_candidates(candidates: &[&AccountCandidate], default_concurrency: NonZeroU32) -> Self {
|
||||
let max_capacity_utilization = candidates
|
||||
.iter()
|
||||
.map(|candidate| candidate.signals.in_flight)
|
||||
.max()
|
||||
.map(|candidate| capacity_utilization(candidate, default_concurrency))
|
||||
.reduce(f64::max)
|
||||
.unwrap_or_default();
|
||||
let min_quota = candidates
|
||||
.iter()
|
||||
@@ -1710,7 +1790,7 @@ impl SmartNormalization {
|
||||
.filter(|latency| *latency > 0)
|
||||
.max();
|
||||
Self {
|
||||
max_in_flight,
|
||||
max_capacity_utilization,
|
||||
min_quota,
|
||||
max_quota,
|
||||
min_latency_ms,
|
||||
@@ -1721,12 +1801,18 @@ impl SmartNormalization {
|
||||
|
||||
fn select_smart_candidate<'a>(
|
||||
candidates: &[&'a AccountCandidate],
|
||||
default_concurrency: NonZeroU32,
|
||||
cursor: u64,
|
||||
) -> Option<&'a AccountCandidate> {
|
||||
let normalization = SmartNormalization::from_candidates(candidates);
|
||||
let normalization = SmartNormalization::from_candidates(candidates, default_concurrency);
|
||||
let mut ranked = candidates
|
||||
.iter()
|
||||
.map(|candidate| (*candidate, smart_score(candidate, &normalization)))
|
||||
.map(|candidate| {
|
||||
(
|
||||
*candidate,
|
||||
smart_score(candidate, default_concurrency, &normalization),
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
ranked.sort_by(|(left, left_score), (right, right_score)| {
|
||||
right_score
|
||||
@@ -1742,11 +1828,14 @@ fn select_smart_candidate<'a>(
|
||||
Some(ranked[index].0)
|
||||
}
|
||||
|
||||
fn smart_score(candidate: &AccountCandidate, normalization: &SmartNormalization) -> f64 {
|
||||
let load = lower_is_better(
|
||||
u64::from(candidate.signals.in_flight),
|
||||
0,
|
||||
u64::from(normalization.max_in_flight),
|
||||
fn smart_score(
|
||||
candidate: &AccountCandidate,
|
||||
default_concurrency: NonZeroU32,
|
||||
normalization: &SmartNormalization,
|
||||
) -> f64 {
|
||||
let load = lower_ratio_is_better(
|
||||
capacity_utilization(candidate, default_concurrency),
|
||||
normalization.max_capacity_utilization,
|
||||
);
|
||||
let quota = candidate.signals.quota_remaining_rank.map_or(0.5, |quota| {
|
||||
higher_is_better(
|
||||
@@ -1781,6 +1870,13 @@ fn smart_score(candidate: &AccountCandidate, normalization: &SmartNormalization)
|
||||
+ SMART_LATENCY_WEIGHT * latency
|
||||
}
|
||||
|
||||
fn lower_ratio_is_better(value: f64, maximum: f64) -> f64 {
|
||||
if maximum <= 0.0 {
|
||||
return 1.0;
|
||||
}
|
||||
1.0 - (value / maximum).clamp(0.0, 1.0)
|
||||
}
|
||||
|
||||
fn lower_is_better(value: u64, minimum: u64, maximum: u64) -> f64 {
|
||||
if maximum <= minimum {
|
||||
return 1.0;
|
||||
|
||||
@@ -362,9 +362,14 @@ pub enum ContinuationAttempt {
|
||||
None,
|
||||
/// 使用 Store 解析出的原生 handle 与账号绑定。
|
||||
Native,
|
||||
/// 在原账号上用完整连接内 transcript 重放。
|
||||
/// 在原账号上执行 Provider 定义的恢复。
|
||||
///
|
||||
/// 只有持有可携带 transcript 的 Provider 才能清除 native handle 后重放;
|
||||
/// 只持有 opaque state 的 Provider 必须保留 continuation 依赖或返回客户端错误。
|
||||
ReplayOwner,
|
||||
/// transcript 已可携带,允许选择其他账号。
|
||||
/// 允许选择其他账号并执行 Provider 定义的恢复。
|
||||
///
|
||||
/// 该枚举本身不证明客户端 input 完整,也不授权 Provider 清除 previous-response。
|
||||
ReplayAny,
|
||||
}
|
||||
|
||||
|
||||
@@ -407,6 +407,84 @@ pub trait ProviderCatalogCachePort: Send + Sync {
|
||||
) -> BoxFuture<'a, Result<Option<OpaqueProviderData>, ProviderStoreError>>;
|
||||
}
|
||||
|
||||
/// Provider 从官方制品核验出的可重建请求画像。
|
||||
///
|
||||
/// Core 只用单调制品序号约束覆盖顺序;具体版本字段由对应 Provider 放在
|
||||
/// `profile` 中解释。每个 Provider 在 Store 中只保留一份最新画像。
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct ProviderArtifactProfile {
|
||||
provider_kind: ProviderKind,
|
||||
artifact_sequence: u64,
|
||||
verified_at: SystemTime,
|
||||
profile: OpaqueProviderData,
|
||||
}
|
||||
|
||||
impl ProviderArtifactProfile {
|
||||
#[must_use]
|
||||
pub const fn new(
|
||||
provider_kind: ProviderKind,
|
||||
artifact_sequence: u64,
|
||||
verified_at: SystemTime,
|
||||
profile: OpaqueProviderData,
|
||||
) -> Self {
|
||||
Self {
|
||||
provider_kind,
|
||||
artifact_sequence,
|
||||
verified_at,
|
||||
profile,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn provider_kind(&self) -> &ProviderKind {
|
||||
&self.provider_kind
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn artifact_sequence(&self) -> u64 {
|
||||
self.artifact_sequence
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn verified_at(&self) -> SystemTime {
|
||||
self.verified_at
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn profile(&self) -> &OpaqueProviderData {
|
||||
&self.profile
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for ProviderArtifactProfile {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("ProviderArtifactProfile")
|
||||
.field("provider_kind", &self.provider_kind)
|
||||
.field("artifact_sequence", &self.artifact_sequence)
|
||||
.field("verified_at", &self.verified_at)
|
||||
.field("profile", &"[PROVIDER-OWNED]")
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
pub trait ProviderArtifactProfileCachePort: Send + Sync {
|
||||
/// 覆盖同一 Provider 的固定 cache key。
|
||||
///
|
||||
/// 返回 `false` 表示 Store 已持有更高的制品序号;相同序号但内容不同必须返回
|
||||
/// [`ProviderStoreErrorKind::Conflict`],不能静默改写已核验画像。
|
||||
fn replace_if_newer(
|
||||
&self,
|
||||
profile: ProviderArtifactProfile,
|
||||
ttl: Duration,
|
||||
) -> BoxFuture<'_, Result<bool, ProviderStoreError>>;
|
||||
|
||||
fn read<'a>(
|
||||
&'a self,
|
||||
provider_kind: &'a ProviderKind,
|
||||
) -> BoxFuture<'a, Result<Option<ProviderArtifactProfile>, ProviderStoreError>>;
|
||||
}
|
||||
|
||||
/// Redis 中可重建的账号状态投影。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderCredentialState {
|
||||
@@ -946,6 +1024,7 @@ pub struct ProviderStorePorts {
|
||||
session_exclusions: Arc<dyn ProviderSessionExclusionPort>,
|
||||
account_feedback: Arc<AccountFeedbackStats>,
|
||||
catalog_cache: Arc<dyn ProviderCatalogCachePort>,
|
||||
artifact_profiles: Arc<dyn ProviderArtifactProfileCachePort>,
|
||||
credential_state: Arc<dyn ProviderCredentialStatePort>,
|
||||
cooldowns: Arc<dyn ProviderCooldownPort>,
|
||||
runtime_policy: Arc<dyn ProviderRuntimePolicyPort>,
|
||||
@@ -962,6 +1041,7 @@ impl ProviderStorePorts {
|
||||
session_affinity: Arc<dyn ProviderSessionAffinityPort>,
|
||||
session_exclusions: Arc<dyn ProviderSessionExclusionPort>,
|
||||
catalog_cache: Arc<dyn ProviderCatalogCachePort>,
|
||||
artifact_profiles: Arc<dyn ProviderArtifactProfileCachePort>,
|
||||
credential_state: Arc<dyn ProviderCredentialStatePort>,
|
||||
cooldowns: Arc<dyn ProviderCooldownPort>,
|
||||
runtime_policy: Arc<dyn ProviderRuntimePolicyPort>,
|
||||
@@ -974,6 +1054,7 @@ impl ProviderStorePorts {
|
||||
session_exclusions,
|
||||
account_feedback: Arc::new(AccountFeedbackStats::default()),
|
||||
catalog_cache,
|
||||
artifact_profiles,
|
||||
credential_state,
|
||||
cooldowns,
|
||||
runtime_policy,
|
||||
@@ -1011,6 +1092,11 @@ impl ProviderStorePorts {
|
||||
Arc::clone(&self.catalog_cache)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn artifact_profiles(&self) -> Arc<dyn ProviderArtifactProfileCachePort> {
|
||||
Arc::clone(&self.artifact_profiles)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn credential_state(&self) -> Arc<dyn ProviderCredentialStatePort> {
|
||||
Arc::clone(&self.credential_state)
|
||||
|
||||
@@ -493,6 +493,31 @@ fn bare_atomic_response_failed(response_id: &str) -> ProviderError {
|
||||
.with_atomic_client_events(vec![failed])
|
||||
}
|
||||
|
||||
fn atomic_history_unavailable(response_id: &str) -> ProviderError {
|
||||
let failed = ProviderEvent::wire(
|
||||
ProtocolWireEvent::json(
|
||||
"openai",
|
||||
Some("response.failed".to_owned()),
|
||||
json!({
|
||||
"type": "response.failed",
|
||||
"response": {
|
||||
"id": response_id,
|
||||
"status": "failed",
|
||||
"error": {
|
||||
"code": "previous_response_not_found",
|
||||
"message": "Previous response was not found. Retrying the full request."
|
||||
}
|
||||
}
|
||||
}),
|
||||
)
|
||||
.expect("history unavailable wire"),
|
||||
);
|
||||
ProviderError::new(ProviderErrorKind::InvalidRequest, UpstreamSendState::Sent)
|
||||
.with_status(400)
|
||||
.with_continuation_failure(ContinuationFailure::HistoryUnavailable)
|
||||
.with_atomic_client_events(vec![failed])
|
||||
}
|
||||
|
||||
fn image_stream(image_output_tokens: Option<u64>) -> Vec<Result<GatewayEvent, ProviderError>> {
|
||||
let mut usage = Usage::new();
|
||||
usage.input_tokens = Some(12);
|
||||
@@ -1481,7 +1506,7 @@ fn streaming_native_continuation_delivers_started_event_before_later_output() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn native_continuation_replays_owner_before_safely_switching_account() {
|
||||
fn native_continuation_uses_provider_defined_owner_then_cross_account_recovery() {
|
||||
let Operation::Generate(generate) = generate_operation() else {
|
||||
panic!("generate operation");
|
||||
};
|
||||
@@ -1536,7 +1561,7 @@ fn native_continuation_replays_owner_before_safely_switching_account() {
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.expect("start execution");
|
||||
block_on(session.collect_uncommitted()).expect("portable replay succeeds");
|
||||
block_on(session.collect_uncommitted()).expect("provider-defined recovery succeeds");
|
||||
block_on(session.commit_downstream(Some(200))).expect("commit response");
|
||||
|
||||
let contexts = provider.contexts.lock().expect("contexts lock");
|
||||
@@ -1561,7 +1586,7 @@ fn native_continuation_replays_owner_before_safely_switching_account() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unavailable_native_continuation_replays_with_the_configured_account_policy() {
|
||||
fn unavailable_native_continuation_uses_provider_recovery_with_the_configured_account_policy() {
|
||||
let Operation::Generate(generate) = generate_operation() else {
|
||||
panic!("generate operation");
|
||||
};
|
||||
@@ -1602,7 +1627,7 @@ fn unavailable_native_continuation_replays_with_the_configured_account_policy()
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.expect("start execution");
|
||||
block_on(session.collect_uncommitted()).expect("portable replay succeeds");
|
||||
block_on(session.collect_uncommitted()).expect("provider-defined recovery succeeds");
|
||||
block_on(session.commit_downstream(Some(200))).expect("commit response");
|
||||
|
||||
let contexts = provider.contexts.lock().expect("contexts lock");
|
||||
@@ -1625,6 +1650,76 @@ fn unavailable_native_continuation_replays_with_the_configured_account_policy()
|
||||
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn continuation_history_failure_without_replay_proof_is_delivered_without_hidden_retry() {
|
||||
let Operation::Generate(generate) = generate_operation() else {
|
||||
panic!("generate operation");
|
||||
};
|
||||
let operation = Operation::Generate(generate.with_provider_session_state(
|
||||
ProviderSessionState::new("openai", Map::new()).expect("provider session state"),
|
||||
));
|
||||
let route_plan = plan(&operation);
|
||||
let (coordinator, _, provider) = coordinator(vec![
|
||||
Script::ObservedStream {
|
||||
account_id: "acct_one",
|
||||
items: vec![Err(atomic_history_unavailable("response-history-missing"))],
|
||||
},
|
||||
Script::Stream {
|
||||
account_id: "acct_two",
|
||||
items: complete_stream(None),
|
||||
},
|
||||
]);
|
||||
let continuation = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("previous-response"),
|
||||
PreviousResponseId::new("upstream-response"),
|
||||
ClientApiKeyId::new("key_client_1").expect("client key"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
ProviderAccountId::new("acct_one").expect("account"),
|
||||
);
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation,
|
||||
route_plan,
|
||||
None,
|
||||
Some(ContinuationBinding::Pinned(continuation)),
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.expect("start execution");
|
||||
|
||||
let terminal = block_on(session.next_event())
|
||||
.expect("history failure remains deliverable")
|
||||
.expect("terminal history failure batch");
|
||||
assert_eq!(
|
||||
terminal.commit_requirement(),
|
||||
CommitRequirement::CommitBeforeDelivery
|
||||
);
|
||||
let events = terminal.into_provider_events();
|
||||
let failure = events
|
||||
.iter()
|
||||
.filter_map(ProviderEvent::wire_event)
|
||||
.find(|wire| wire.event_type() == Some("response.failed"))
|
||||
.expect("raw response.failed event");
|
||||
assert_eq!(
|
||||
failure
|
||||
.data()
|
||||
.pointer("/response/error/code")
|
||||
.and_then(Value::as_str),
|
||||
Some("previous_response_not_found")
|
||||
);
|
||||
block_on(session.commit_downstream(Some(200))).expect("commit terminal history failure");
|
||||
assert!(matches!(
|
||||
block_on(session.next_event()),
|
||||
Err(EngineError::Provider(_))
|
||||
));
|
||||
|
||||
let contexts = provider.contexts.lock().expect("contexts lock");
|
||||
assert_eq!(contexts.len(), 1);
|
||||
assert_eq!(
|
||||
contexts[0].continuation_attempt(),
|
||||
ContinuationAttempt::Native
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn diagnostic_account_reaches_provider_and_matching_metadata_succeeds() {
|
||||
let operation = operation(RetrySafety::Idempotent);
|
||||
|
||||
@@ -51,6 +51,15 @@ fn candidate(id: &str, in_flight: u32, remaining: Option<u64>) -> AccountCandida
|
||||
}
|
||||
}
|
||||
|
||||
fn candidate_with_concurrency(id: &str, in_flight: u32, concurrency: u32) -> AccountCandidate {
|
||||
let mut candidate = candidate(id, in_flight, Some(100));
|
||||
candidate.account = candidate.account.with_scheduling(
|
||||
Some(AccountConcurrencyLimit::new(concurrency).expect("concurrency override")),
|
||||
AccountWeight::DEFAULT,
|
||||
);
|
||||
candidate
|
||||
}
|
||||
|
||||
fn weighted_candidate(id: &str, weight: u16, in_flight: u32) -> AccountCandidate {
|
||||
let mut candidate = candidate(id, in_flight, None);
|
||||
candidate.account = candidate.account.with_scheduling(
|
||||
@@ -348,6 +357,20 @@ fn smart_selector_should_prefer_lower_inflight_count() {
|
||||
assert_eq!(selected.candidate().account.id().as_str(), "acct_idle");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn smart_selector_should_prefer_lower_capacity_utilization_over_lower_inflight_count() {
|
||||
let candidates = [
|
||||
candidate_with_concurrency("acct_smaller", 1, 2),
|
||||
candidate_with_concurrency("acct_larger", 2, 10),
|
||||
];
|
||||
|
||||
let selected = AccountSelector
|
||||
.select(&candidates, &context(RotationStrategy::Smart))
|
||||
.expect("candidate available");
|
||||
|
||||
assert_eq!(selected.candidate().account.id().as_str(), "acct_larger");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_should_prioritize_the_highest_weight_for_every_rotation_strategy() {
|
||||
let candidates = vec![
|
||||
@@ -622,6 +645,22 @@ fn quota_reset_selector_should_prefer_known_earliest_window() {
|
||||
assert_eq!(selected.candidate().account.id().as_str(), "acct_earlier");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quota_reset_selector_should_use_capacity_utilization_as_load_tiebreaker() {
|
||||
let reset_at = SystemTime::now() + Duration::from_secs(60);
|
||||
let mut smaller = candidate_with_concurrency("acct_smaller", 1, 2);
|
||||
smaller.signals.quota_reset_at = Some(reset_at);
|
||||
let mut larger = candidate_with_concurrency("acct_larger", 2, 10);
|
||||
larger.signals.quota_reset_at = Some(reset_at);
|
||||
let candidates = [smaller, larger];
|
||||
|
||||
let selected = AccountSelector
|
||||
.select(&candidates, &context(RotationStrategy::QuotaResetPriority))
|
||||
.expect("candidate available");
|
||||
|
||||
assert_eq!(selected.candidate().account.id().as_str(), "acct_larger");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn smart_selector_should_prefer_known_quota_over_unknown_after_load_is_equal() {
|
||||
let candidates = vec![
|
||||
|
||||
@@ -41,6 +41,7 @@ pub fn load_config<T: LoadableConfig>() -> Result<T, ConfigError> {
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct HostConfig {
|
||||
pub listen: ListenConfig,
|
||||
pub runtime_data_dir: PathBuf,
|
||||
pub logging: LoggingConfig,
|
||||
#[serde(default)]
|
||||
pub system_update: SystemUpdateConfig,
|
||||
@@ -76,11 +77,21 @@ impl HostConfig {
|
||||
"host.worker_shutdown_timeout_seconds",
|
||||
));
|
||||
}
|
||||
if self.runtime_data_dir.as_os_str().is_empty() {
|
||||
return Err(ConfigError::InvalidField("host.runtime_data_dir"));
|
||||
}
|
||||
resolve_relative_path(source_dir, &mut self.runtime_data_dir);
|
||||
self.logging.resolve_and_validate(source_dir)?;
|
||||
self.system_update.resolve_and_validate(source_dir)?;
|
||||
self.system_update
|
||||
.resolve_and_validate(source_dir, &self.runtime_data_dir)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn runtime_data_dir(&self) -> &Path {
|
||||
&self.runtime_data_dir
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn drain_timeout(&self) -> Duration {
|
||||
Duration::from_secs(self.drain_timeout_seconds)
|
||||
|
||||
@@ -81,10 +81,10 @@ impl Default for SystemUpdateConfig {
|
||||
});
|
||||
let update_state_file = environment_value("CPR_UPDATE_STATE_FILE")
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|| PathBuf::from("/app/.runtime/data/update-state.json"));
|
||||
.unwrap_or_else(|| PathBuf::from("update-state.json"));
|
||||
let update_lock_file = environment_value("CPR_UPDATE_LOCK_FILE")
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|| update_state_file.with_extension("lock"));
|
||||
.unwrap_or_else(|| update_state_file.with_file_name("update.lock"));
|
||||
let update_temp_dir = environment_value("CPR_UPDATE_TEMP_DIR")
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|| state::default_temp_dir(&update_state_file));
|
||||
@@ -120,21 +120,28 @@ impl Default for SystemUpdateConfig {
|
||||
}
|
||||
|
||||
impl SystemUpdateConfig {
|
||||
pub(crate) fn resolve_and_validate(&mut self, source_dir: &Path) -> Result<(), ConfigError> {
|
||||
for path in [
|
||||
self.executable_path.as_mut(),
|
||||
Some(&mut self.web_dist_dir),
|
||||
Some(&mut self.update_state_file),
|
||||
Some(&mut self.update_lock_file),
|
||||
Some(&mut self.update_temp_dir),
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
pub(crate) fn resolve_and_validate(
|
||||
&mut self,
|
||||
source_dir: &Path,
|
||||
runtime_data_dir: &Path,
|
||||
) -> Result<(), ConfigError> {
|
||||
for path in [self.executable_path.as_mut(), Some(&mut self.web_dist_dir)]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
{
|
||||
if path.is_relative() {
|
||||
*path = source_dir.join(&*path);
|
||||
}
|
||||
}
|
||||
for path in [
|
||||
&mut self.update_state_file,
|
||||
&mut self.update_lock_file,
|
||||
&mut self.update_temp_dir,
|
||||
] {
|
||||
if path.is_relative() {
|
||||
*path = runtime_data_dir.join(&*path);
|
||||
}
|
||||
}
|
||||
if self.version.trim().is_empty() {
|
||||
return Err(ConfigError::InvalidField("host.system_update.version"));
|
||||
}
|
||||
|
||||
@@ -34,6 +34,40 @@ fn host_config_resolves_only_host_owned_relative_paths() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn host_config_resolves_runtime_data_dir_relative_to_configuration() {
|
||||
let mut config = valid_config();
|
||||
config
|
||||
.resolve_and_validate(std::path::Path::new("/srv/gateway"))
|
||||
.expect("valid host config");
|
||||
|
||||
assert_eq!(
|
||||
config.runtime_data_dir(),
|
||||
PathBuf::from("/srv/gateway/runtime-data")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn host_config_derives_update_paths_from_runtime_data_dir() {
|
||||
let mut config = valid_config();
|
||||
config
|
||||
.resolve_and_validate(std::path::Path::new("/srv/gateway"))
|
||||
.expect("valid host config");
|
||||
|
||||
assert_eq!(
|
||||
(
|
||||
config.system_update.update_state_file,
|
||||
config.system_update.update_lock_file,
|
||||
config.system_update.update_temp_dir,
|
||||
),
|
||||
(
|
||||
PathBuf::from("/srv/gateway/runtime-data/update-state.json"),
|
||||
PathBuf::from("/srv/gateway/runtime-data/update.lock"),
|
||||
PathBuf::from("/srv/gateway/runtime-data/update-tmp"),
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn host_config_rejects_zero_drain_window() {
|
||||
let mut config = valid_config();
|
||||
@@ -48,9 +82,9 @@ fn host_config_rejects_zero_drain_window() {
|
||||
|
||||
fn valid_config() -> HostConfig {
|
||||
let system_update = SystemUpdateConfig {
|
||||
update_state_file: PathBuf::from(".runtime/update-state.json"),
|
||||
update_lock_file: PathBuf::from(".runtime/update-state.lock"),
|
||||
update_temp_dir: PathBuf::from(".runtime/update-tmp"),
|
||||
update_state_file: PathBuf::from("update-state.json"),
|
||||
update_lock_file: PathBuf::from("update.lock"),
|
||||
update_temp_dir: PathBuf::from("update-tmp"),
|
||||
..SystemUpdateConfig::default()
|
||||
};
|
||||
HostConfig {
|
||||
@@ -58,6 +92,7 @@ fn valid_config() -> HostConfig {
|
||||
host: "127.0.0.1".to_owned(),
|
||||
port: 8080,
|
||||
},
|
||||
runtime_data_dir: PathBuf::from("runtime-data"),
|
||||
logging: LoggingConfig {
|
||||
level: "info".to_owned(),
|
||||
stdout: true,
|
||||
|
||||
@@ -23,6 +23,7 @@ fn logging_requires_at_least_one_sink() {
|
||||
host: "127.0.0.1".to_owned(),
|
||||
port: 8080,
|
||||
},
|
||||
runtime_data_dir: PathBuf::from("/tmp/runtime-data"),
|
||||
logging: LoggingConfig {
|
||||
level: "info".to_owned(),
|
||||
stdout: false,
|
||||
@@ -98,6 +99,7 @@ fn write_oauth_recovery_log(directory: PathBuf) {
|
||||
host: "127.0.0.1".to_owned(),
|
||||
port: 8080,
|
||||
},
|
||||
runtime_data_dir: PathBuf::from("/tmp/runtime-data"),
|
||||
logging: LoggingConfig {
|
||||
level: "off".to_owned(),
|
||||
stdout: false,
|
||||
|
||||
@@ -208,11 +208,24 @@ pub struct RateLimitWindow {
|
||||
pub reset_at: Option<i64>,
|
||||
}
|
||||
|
||||
/// 限流项键的来源。
|
||||
///
|
||||
/// HTTP headers 和事件里的 `metered_limit_name` 提供稳定 `limit_id`;
|
||||
/// WebSocket `additional_rate_limits` 在部分协议版本里只以可读名称为 map key,
|
||||
/// 此时必须由持久化层用既有名称元数据解析,不能把名称误当成稳定 ID。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum RateLimitKeySource {
|
||||
LimitId,
|
||||
LimitName,
|
||||
}
|
||||
|
||||
/// 单个计量项的限流信息。
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct RateLimitDetails {
|
||||
/// 稳定的计量项 ID。
|
||||
pub limit_id: String,
|
||||
/// [`limit_id`](Self::limit_id) 是稳定 ID,还是名称生成的临时解析键。
|
||||
pub key_source: RateLimitKeySource,
|
||||
/// 上游提供的可读名称。
|
||||
pub limit_name: Option<String>,
|
||||
/// 当前请求是否被允许。
|
||||
@@ -341,20 +354,40 @@ pub fn parse_rate_limits_event(value: &Value) -> Option<ParsedRateLimits> {
|
||||
return None;
|
||||
}
|
||||
|
||||
let limit_id = rate_limit_name(value)
|
||||
.map(normalize_limit_id)
|
||||
let explicit_limit_id = top_level_rate_limit_id(value).map(normalize_limit_id);
|
||||
let explicit_limit_name = non_empty_string(value, "limit_name");
|
||||
let default_limit_id = explicit_limit_id
|
||||
.clone()
|
||||
.unwrap_or_else(|| "codex".to_string());
|
||||
let limit_name = value
|
||||
.get("metered_limit_name")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(|_| value.get("limit_name"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|name| !name.is_empty())
|
||||
.map(ToString::to_string);
|
||||
let details = value
|
||||
.get("rate_limits")
|
||||
.and_then(|details| parse_details_from_object(details, &limit_id, limit_name));
|
||||
let default_details = value.get("rate_limits").and_then(|details| {
|
||||
parse_details_from_object(
|
||||
details,
|
||||
&default_limit_id,
|
||||
explicit_limit_name,
|
||||
RateLimitKeySource::LimitId,
|
||||
)
|
||||
});
|
||||
let mut limits = parse_additional_rate_limits(value);
|
||||
let matching_additional_ids = default_details
|
||||
.as_ref()
|
||||
.map(|default| {
|
||||
limits
|
||||
.iter()
|
||||
.filter(|(_, additional)| same_rate_limit_facts(default, additional))
|
||||
.map(|(limit_id, _)| limit_id.clone())
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
// 当前 WebSocket 协议会把活动具名桶同时镜像到顶层 `rate_limits`,但不附
|
||||
// `metered_limit_name`。只要 additional 中存在相同事实,顶层就不是 core 桶。
|
||||
let default_is_additional_alias =
|
||||
explicit_limit_id.is_none() && !matching_additional_ids.is_empty();
|
||||
if let Some(details) = default_details
|
||||
&& !default_is_additional_alias
|
||||
{
|
||||
limits.insert(default_limit_id.clone(), details);
|
||||
}
|
||||
let credits = value.get("credits").and_then(parse_credits_from_object);
|
||||
let plan_type = value
|
||||
.get("plan_type")
|
||||
@@ -363,17 +396,18 @@ pub fn parse_rate_limits_event(value: &Value) -> Option<ParsedRateLimits> {
|
||||
.filter(|plan| !plan.is_empty())
|
||||
.map(ToString::to_string);
|
||||
|
||||
if details.is_none() && credits.is_none() && plan_type.is_none() {
|
||||
if limits.is_empty() && credits.is_none() && plan_type.is_none() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut limits = BTreeMap::new();
|
||||
if let Some(details) = details {
|
||||
limits.insert(limit_id.clone(), details);
|
||||
}
|
||||
let active_limit = explicit_limit_id.or_else(|| match matching_additional_ids.as_slice() {
|
||||
[limit_id] => Some(limit_id.clone()),
|
||||
[] if !limits.is_empty() => Some(default_limit_id),
|
||||
_ => None,
|
||||
});
|
||||
Some(ParsedRateLimits {
|
||||
limits,
|
||||
active_limit: Some(limit_id),
|
||||
active_limit,
|
||||
credits,
|
||||
plan_type,
|
||||
promo_message: value
|
||||
@@ -597,6 +631,7 @@ fn parse_details_from_lookup(
|
||||
}
|
||||
Some(RateLimitDetails {
|
||||
limit_id: limit_id.to_string(),
|
||||
key_source: RateLimitKeySource::LimitId,
|
||||
limit_name,
|
||||
allowed,
|
||||
limit_reached,
|
||||
@@ -687,6 +722,7 @@ fn parse_details_from_object(
|
||||
value: &Value,
|
||||
limit_id: &str,
|
||||
limit_name: Option<String>,
|
||||
key_source: RateLimitKeySource,
|
||||
) -> Option<RateLimitDetails> {
|
||||
let primary = value.get("primary").and_then(parse_window_from_object);
|
||||
let secondary = value.get("secondary").and_then(parse_window_from_object);
|
||||
@@ -697,6 +733,7 @@ fn parse_details_from_object(
|
||||
}
|
||||
Some(RateLimitDetails {
|
||||
limit_id: limit_id.to_string(),
|
||||
key_source,
|
||||
limit_name,
|
||||
allowed,
|
||||
limit_reached,
|
||||
@@ -752,11 +789,59 @@ fn value_as_positive_i64(value: &Value) -> Option<i64> {
|
||||
.filter(|value| *value > 0)
|
||||
}
|
||||
|
||||
fn rate_limit_name(value: &Value) -> Option<&str> {
|
||||
fn rate_limit_id(value: &Value) -> Option<&str> {
|
||||
value
|
||||
.get("metered_limit_name")
|
||||
.or_else(|| value.get("limit_name"))
|
||||
.or_else(|| value.get("limit_id"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn top_level_rate_limit_id(value: &Value) -> Option<&str> {
|
||||
rate_limit_id(value).or_else(|| {
|
||||
value
|
||||
.get("limit_name")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
})
|
||||
}
|
||||
|
||||
fn non_empty_string(value: &Value, field: &str) -> Option<String> {
|
||||
value
|
||||
.get(field)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToString::to_string)
|
||||
}
|
||||
|
||||
fn parse_additional_rate_limits(value: &Value) -> BTreeMap<String, RateLimitDetails> {
|
||||
value
|
||||
.get("additional_rate_limits")
|
||||
.and_then(Value::as_object)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(|(map_key, value)| {
|
||||
let explicit_limit_id = rate_limit_id(value).map(normalize_limit_id);
|
||||
let (limit_id, key_source) = explicit_limit_id.map_or_else(
|
||||
|| (normalize_limit_id(map_key), RateLimitKeySource::LimitName),
|
||||
|limit_id| (limit_id, RateLimitKeySource::LimitId),
|
||||
);
|
||||
let limit_name = non_empty_string(value, "limit_name")
|
||||
.or_else(|| Some(map_key.trim().to_owned()).filter(|name| !name.is_empty()));
|
||||
parse_details_from_object(value, &limit_id, limit_name, key_source)
|
||||
.map(|details| (limit_id, details))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn same_rate_limit_facts(left: &RateLimitDetails, right: &RateLimitDetails) -> bool {
|
||||
left.allowed == right.allowed
|
||||
&& left.limit_reached == right.limit_reached
|
||||
&& left.primary == right.primary
|
||||
&& left.secondary == right.secondary
|
||||
}
|
||||
|
||||
fn normalize_limit_id(value: impl AsRef<str>) -> String {
|
||||
|
||||
@@ -70,7 +70,7 @@ pub const MAX_SSE_EVENT_BUFFER_BYTES: usize = 64 * 1024 * 1024;
|
||||
pub const DONE_SSE_FRAME: &str = "data: [DONE]\n\n";
|
||||
|
||||
/// 用于处理任意分块边界的增量 SSE 解码器。
|
||||
#[derive(Debug, Default)]
|
||||
#[derive(Debug)]
|
||||
pub struct SseEventDecoder {
|
||||
pending: Vec<u8>,
|
||||
/// `pending` 中已确认不含帧分隔符的前缀长度。新数据到达时从该边界
|
||||
@@ -79,6 +79,19 @@ pub struct SseEventDecoder {
|
||||
/// 超大未完成帧已开始直接交付;在遇到帧分隔符前只保留跨 chunk 扫描尾部,
|
||||
/// 不再尝试解析该帧。
|
||||
opaque_frame: bool,
|
||||
/// 只有整个 SSE 流的第一条物理行允许携带 UTF-8 BOM。
|
||||
stream_start: bool,
|
||||
}
|
||||
|
||||
impl Default for SseEventDecoder {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
pending: Vec::new(),
|
||||
scanned: 0,
|
||||
opaque_frame: false,
|
||||
stream_start: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SseEventDecoder {
|
||||
@@ -115,7 +128,8 @@ impl SseEventDecoder {
|
||||
while let Some(end) = self.next_frame_end(consumed) {
|
||||
let frame = std::str::from_utf8(&self.pending[consumed..end])
|
||||
.map_err(|error| SseError::ParseError(error.to_string()))?;
|
||||
events.extend(parse_sse_events(frame)?);
|
||||
events.extend(parse_sse_events_inner(frame, self.stream_start)?);
|
||||
self.stream_start = false;
|
||||
consumed = end;
|
||||
}
|
||||
|
||||
@@ -149,9 +163,10 @@ impl SseEventDecoder {
|
||||
} else {
|
||||
std::str::from_utf8(&raw)
|
||||
.ok()
|
||||
.and_then(|frame| parse_sse_events(frame).ok())
|
||||
.and_then(|frame| parse_sse_events_inner(frame, self.stream_start).ok())
|
||||
.unwrap_or_default()
|
||||
};
|
||||
self.stream_start = false;
|
||||
frames.push(SseFrame::new(raw, events));
|
||||
consumed = end;
|
||||
}
|
||||
@@ -173,7 +188,9 @@ impl SseEventDecoder {
|
||||
let pending = std::mem::take(&mut self.pending);
|
||||
let frame = std::str::from_utf8(&pending)
|
||||
.map_err(|error| SseError::ParseError(error.to_string()))?;
|
||||
parse_sse_events(frame)
|
||||
let events = parse_sse_events_inner(frame, self.stream_start)?;
|
||||
self.stream_start = false;
|
||||
Ok(events)
|
||||
}
|
||||
|
||||
/// 在流结束时返回未带空行分隔符的最后一帧及其旁路解析结果。
|
||||
@@ -191,9 +208,10 @@ impl SseEventDecoder {
|
||||
} else {
|
||||
std::str::from_utf8(&raw)
|
||||
.ok()
|
||||
.and_then(|frame| parse_sse_events(frame).ok())
|
||||
.and_then(|frame| parse_sse_events_inner(frame, self.stream_start).ok())
|
||||
.unwrap_or_default()
|
||||
};
|
||||
self.stream_start = false;
|
||||
self.opaque_frame = false;
|
||||
vec![SseFrame::new(raw, events)]
|
||||
}
|
||||
@@ -207,6 +225,7 @@ impl SseEventDecoder {
|
||||
let remaining = self.pending.split_off(end);
|
||||
let raw = std::mem::replace(&mut self.pending, remaining);
|
||||
frames.push(SseFrame::new(raw, Vec::new()));
|
||||
self.stream_start = false;
|
||||
self.opaque_frame = false;
|
||||
self.scanned = 0;
|
||||
true
|
||||
@@ -226,6 +245,7 @@ impl SseEventDecoder {
|
||||
let raw = std::mem::replace(&mut self.pending, tail);
|
||||
self.scanned = 0;
|
||||
frames.push(SseFrame::new(raw, Vec::new()));
|
||||
self.stream_start = false;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -351,12 +371,21 @@ impl EventBuilder {
|
||||
|
||||
/// 解析 SSE 事件流。
|
||||
pub fn parse_sse_events(input: &str) -> Result<Vec<SseEvent>, SseError> {
|
||||
parse_sse_events_inner(input, true)
|
||||
}
|
||||
|
||||
fn parse_sse_events_inner(input: &str, strip_leading_bom: bool) -> Result<Vec<SseEvent>, SseError> {
|
||||
let mut events = Vec::new();
|
||||
let mut builder = EventBuilder::default();
|
||||
let mut saw_sse_syntax = false;
|
||||
let mut event_buffer_bytes = 0usize;
|
||||
|
||||
for raw_line in input.lines() {
|
||||
for (index, raw_line) in input.lines().enumerate() {
|
||||
let raw_line = if strip_leading_bom && index == 0 {
|
||||
raw_line.strip_prefix('\u{feff}').unwrap_or(raw_line)
|
||||
} else {
|
||||
raw_line
|
||||
};
|
||||
let line = raw_line.strip_suffix('\r').unwrap_or(raw_line);
|
||||
if line.is_empty() {
|
||||
event_buffer_bytes = 0;
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use gateway_protocol::openai::events::{
|
||||
RateLimitWindow, TokenUsage, extract_sse_usage, extract_usage, is_codex_quota_header_name,
|
||||
is_rate_limit_header_name, parse_rate_limit_headers, parse_rate_limits_event,
|
||||
rate_limits_to_header_pairs, retry_after_seconds_from_body,
|
||||
RateLimitKeySource, RateLimitWindow, TokenUsage, extract_sse_usage, extract_usage,
|
||||
is_codex_quota_header_name, is_rate_limit_header_name, parse_rate_limit_headers,
|
||||
parse_rate_limits_event, rate_limits_to_header_pairs, retry_after_seconds_from_body,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
@@ -483,6 +483,58 @@ fn parse_rate_limits_event_should_extract_websocket_metered_limit() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_rate_limits_event_should_preserve_name_keyed_additional_limit() {
|
||||
let event = json!({
|
||||
"type": "codex.rate_limits",
|
||||
"plan_type": "pro",
|
||||
"rate_limits": {
|
||||
"allowed": true,
|
||||
"limit_reached": false,
|
||||
"primary": {
|
||||
"used_percent": 0,
|
||||
"window_minutes": 10080,
|
||||
"reset_after_seconds": 354916,
|
||||
"reset_at": 1787396838
|
||||
},
|
||||
"secondary": null
|
||||
},
|
||||
"additional_rate_limits": {
|
||||
"GPT-5.3-Codex-Spark": {
|
||||
"allowed": true,
|
||||
"limit_reached": false,
|
||||
"primary": {
|
||||
"used_percent": 0,
|
||||
"window_minutes": 10080,
|
||||
"reset_after_seconds": 354916,
|
||||
"reset_at": 1787396838
|
||||
},
|
||||
"secondary": null
|
||||
}
|
||||
},
|
||||
"credits": null
|
||||
});
|
||||
|
||||
let parsed = parse_rate_limits_event(&event).expect("rate-limit event");
|
||||
let details = parsed
|
||||
.limits
|
||||
.get("gpt_5.3_codex_spark")
|
||||
.expect("name-keyed additional limit");
|
||||
|
||||
assert_eq!(parsed.active_limit.as_deref(), Some("gpt_5.3_codex_spark"));
|
||||
assert_eq!(details.key_source, RateLimitKeySource::LimitName);
|
||||
assert_eq!(details.limit_name.as_deref(), Some("GPT-5.3-Codex-Spark"));
|
||||
assert_eq!(
|
||||
details.primary,
|
||||
Some(RateLimitWindow {
|
||||
used_percent: 0.0,
|
||||
window_minutes: Some(10080),
|
||||
reset_at: Some(1787396838),
|
||||
})
|
||||
);
|
||||
assert!(!parsed.limits.contains_key("codex"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_rate_limits_event_should_reject_unrelated_event_type() {
|
||||
for event in [
|
||||
|
||||
@@ -31,6 +31,50 @@ fn parser_should_preserve_multiline_data_id_and_retry() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parser_should_ignore_a_leading_utf8_bom() {
|
||||
let input = concat!(
|
||||
"\u{feff}event: response.created\n",
|
||||
"data: {\"type\":\"response.created\"}\n\n",
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
parse_sse_events(input).expect("SSE with BOM"),
|
||||
vec![SseEvent {
|
||||
event: Some("response.created".to_owned()),
|
||||
data: "{\"type\":\"response.created\"}".to_owned(),
|
||||
id: None,
|
||||
retry: None,
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn incremental_decoder_should_ignore_bom_only_at_the_stream_start() {
|
||||
let mut decoder = SseEventDecoder::default();
|
||||
let first = decoder
|
||||
.push(
|
||||
concat!(
|
||||
"\u{feff}event: response.created\n",
|
||||
"data: {\"type\":\"response.created\"}\n\n",
|
||||
)
|
||||
.as_bytes(),
|
||||
)
|
||||
.expect("first SSE frame");
|
||||
let second = decoder
|
||||
.push(
|
||||
concat!(
|
||||
"\u{feff}event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\"}\n\n",
|
||||
)
|
||||
.as_bytes(),
|
||||
)
|
||||
.expect("second SSE frame");
|
||||
|
||||
assert_eq!(first[0].event.as_deref(), Some("response.created"));
|
||||
assert_eq!(second[0].event, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parser_should_ignore_done_control_frame() {
|
||||
assert!(
|
||||
|
||||
@@ -50,7 +50,7 @@ impl StoreBundle {
|
||||
pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
const REDIS_NAMESPACE: &str = "codex-proxy-rs";
|
||||
|
||||
config.resolve_and_validate(std::path::Path::new("."))?;
|
||||
config.validate_resolved()?;
|
||||
let pool = postgres::connect_and_migrate(&config.database_url()?, config.pool).await?;
|
||||
let redis_client = ::redis::Client::open(config.redis_url()?)
|
||||
.map_err(|_| redis_unavailable("create Redis client"))?;
|
||||
@@ -79,6 +79,10 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
redis_connection.clone(),
|
||||
REDIS_NAMESPACE,
|
||||
)?);
|
||||
let artifact_profiles = Arc::new(redis::RedisProviderArtifactProfileRepository::new(
|
||||
redis_connection.clone(),
|
||||
REDIS_NAMESPACE,
|
||||
)?);
|
||||
let runtime_policy = Arc::new(postgres::PgRuntimeSettingsRepository::new(pool.clone()));
|
||||
let oauth_pending = Arc::new(redis::RedisOAuthPendingFlowRepository::new(
|
||||
redis_connection.clone(),
|
||||
@@ -167,6 +171,7 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
REDIS_NAMESPACE,
|
||||
)?),
|
||||
credential_state.clone(),
|
||||
artifact_profiles,
|
||||
credential_state,
|
||||
cooldowns,
|
||||
runtime_policy,
|
||||
@@ -204,7 +209,7 @@ pub(crate) fn backup_ports(
|
||||
config: &StoreConfig,
|
||||
) -> StoreResult<BackupStorePorts> {
|
||||
let staging = Arc::new(backup::staging::StagingArea::open(
|
||||
std::path::PathBuf::from(".runtime/data/backup-staging"),
|
||||
config.backup_staging_dir().to_path_buf(),
|
||||
backup::staging::DEFAULT_MAX_ARCHIVE_BYTES,
|
||||
)?);
|
||||
let repository = Arc::new(postgres::PgBackupRepository::new(pool));
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
//! Store 启动配置、环境变量解析与校验。
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use super::*;
|
||||
|
||||
pub(crate) const DATABASE_URL_ENV: &str = "CPR_DATABASE_URL";
|
||||
@@ -15,6 +17,8 @@ pub struct StoreConfig {
|
||||
pub(crate) redis: StoreConnectionConfig,
|
||||
#[serde(default)]
|
||||
pub(crate) pool: StorePoolConfig,
|
||||
#[serde(skip)]
|
||||
backup_staging_dir: PathBuf,
|
||||
}
|
||||
|
||||
/// PostgreSQL 连接池预算;acquire 超时决定池耗尽时快速失败而非排队积压。
|
||||
@@ -47,7 +51,18 @@ impl StorePoolConfig {
|
||||
}
|
||||
|
||||
impl StoreConfig {
|
||||
pub fn resolve_and_validate(&mut self, _source_dir: &std::path::Path) -> StoreResult<()> {
|
||||
pub fn resolve_and_validate(&mut self, runtime_data_dir: &Path) -> StoreResult<()> {
|
||||
if runtime_data_dir.as_os_str().is_empty() {
|
||||
return Err(StoreError::InvalidData {
|
||||
entity: "store config",
|
||||
message: "runtime_data_dir must not be empty".to_owned(),
|
||||
});
|
||||
}
|
||||
self.backup_staging_dir = runtime_data_dir.join("backup-staging");
|
||||
self.validate_resolved()
|
||||
}
|
||||
|
||||
pub(crate) fn validate_resolved(&mut self) -> StoreResult<()> {
|
||||
if let Some(url) = optional_environment_value(DATABASE_URL_ENV)? {
|
||||
self.database.url = url;
|
||||
}
|
||||
@@ -63,6 +78,12 @@ impl StoreConfig {
|
||||
self.database.validate("database")?;
|
||||
self.redis.validate("redis")?;
|
||||
self.pool.validate()?;
|
||||
if self.backup_staging_dir.as_os_str().is_empty() {
|
||||
return Err(StoreError::InvalidData {
|
||||
entity: "store config",
|
||||
message: "runtime_data_dir was not resolved".to_owned(),
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -73,6 +94,12 @@ impl StoreConfig {
|
||||
pub(crate) fn redis_url(&self) -> StoreResult<String> {
|
||||
self.redis.connection_url("redis")
|
||||
}
|
||||
|
||||
/// 返回由统一运行数据根目录派生的备份暂存目录。
|
||||
#[must_use]
|
||||
pub fn backup_staging_dir(&self) -> &Path {
|
||||
&self.backup_staging_dir
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn optional_environment_value(name: &'static str) -> StoreResult<Option<String>> {
|
||||
|
||||
@@ -52,9 +52,12 @@ pub(crate) async fn provider_account_usage(
|
||||
}
|
||||
statement.push(
|
||||
" group by pa.id, pa.provider_kind, pa.authentication_kind, pa.name, pa.email,
|
||||
pa.plan_type
|
||||
order by max(mr.started_at) desc nulls last, pa.name, pa.id limit ",
|
||||
pa.plan_type",
|
||||
);
|
||||
if query.account_ids.is_none() {
|
||||
statement.push(" having count(mr.id) > 0");
|
||||
}
|
||||
statement.push(" order by max(mr.started_at) desc nulls last, pa.name, pa.id limit ");
|
||||
statement.push_bind(i64::from(query.limit));
|
||||
let rows = statement
|
||||
.build()
|
||||
|
||||
@@ -103,7 +103,7 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
&self,
|
||||
update: CredentialCasUpdate,
|
||||
) -> Result<CredentialCasOutcome, CoreStoreError> {
|
||||
let (
|
||||
let CredentialCasUpdateParts {
|
||||
account_id,
|
||||
expected_revision,
|
||||
profile,
|
||||
@@ -111,7 +111,8 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
has_refresh_token,
|
||||
access_token_expires_at,
|
||||
next_refresh_at,
|
||||
) = update.into_parts();
|
||||
account_state,
|
||||
} = update.into_parts();
|
||||
if profile.account_id != account_id {
|
||||
return Err(CoreStoreError::new(CoreStoreErrorKind::InvalidData));
|
||||
}
|
||||
@@ -122,13 +123,38 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
CREDENTIALS_MAX_BYTES,
|
||||
)
|
||||
.map_err(core_store_error)?;
|
||||
let credential_state = account_state
|
||||
.as_ref()
|
||||
.map(|state| state.credential_state.as_str());
|
||||
let credential_observed_at = account_state
|
||||
.as_ref()
|
||||
.map(|state| DateTime::<Utc>::from(state.observed_at));
|
||||
let error_reason = account_state
|
||||
.as_ref()
|
||||
.and_then(|state| state.error_reason.map(AccountErrorReason::as_str));
|
||||
let message = account_state
|
||||
.as_ref()
|
||||
.and_then(|state| state.message.as_deref());
|
||||
let next = sqlx::query_scalar::<_, i64>(
|
||||
"update provider_accounts
|
||||
set name = $3, email = $4, plan_type = $5,
|
||||
provider_credentials_json = $6,
|
||||
credential_revision = credential_revision + 1,
|
||||
has_refresh_token = $7, access_token_expires_at = $8,
|
||||
next_refresh_at = $9, updated_at = now()
|
||||
next_refresh_at = $9,
|
||||
credential_state = case
|
||||
when $10::text is not null and enabled and upstream_user_id is not null
|
||||
then $10 else credential_state end,
|
||||
credential_observed_at = case
|
||||
when $10::text is not null and enabled and upstream_user_id is not null
|
||||
then $11 else credential_observed_at end,
|
||||
last_error_reason = case
|
||||
when $10::text is not null and enabled and upstream_user_id is not null
|
||||
then $12 else last_error_reason end,
|
||||
last_error_message = case
|
||||
when $10::text is not null and enabled and upstream_user_id is not null
|
||||
then $13 else last_error_message end,
|
||||
updated_at = greatest(now(), coalesce($11, now()))
|
||||
where id = $1 and credential_revision = $2
|
||||
returning credential_revision",
|
||||
)
|
||||
@@ -141,6 +167,10 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
.bind(has_refresh_token)
|
||||
.bind(access_token_expires_at.map(DateTime::<Utc>::from))
|
||||
.bind(next_refresh_at.map(DateTime::<Utc>::from))
|
||||
.bind(credential_state)
|
||||
.bind(credential_observed_at)
|
||||
.bind(error_reason)
|
||||
.bind(message)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::Unavailable))?;
|
||||
@@ -280,6 +310,25 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
})
|
||||
}
|
||||
|
||||
async fn touch_quota_observation(
|
||||
&self,
|
||||
touch: QuotaObservationTouch,
|
||||
) -> Result<QuotaWriteOutcome, CoreStoreError> {
|
||||
let updated = self
|
||||
.touch_provider_quota_observation(
|
||||
touch.account_id.as_str(),
|
||||
Revision::new(touch.expected_revision.get()).map_err(core_store_error)?,
|
||||
DateTime::<Utc>::from(touch.observed_at),
|
||||
)
|
||||
.await
|
||||
.map_err(core_store_error)?;
|
||||
Ok(if updated {
|
||||
QuotaWriteOutcome::Updated
|
||||
} else {
|
||||
QuotaWriteOutcome::Conflict
|
||||
})
|
||||
}
|
||||
|
||||
async fn apply_state_change(&self, change: AccountStateChange) -> Result<(), CoreStoreError> {
|
||||
let updated = self
|
||||
.apply_provider_account_state(ProviderAccountStateUpdate {
|
||||
|
||||
@@ -38,12 +38,13 @@ use sqlx::{PgPool, Postgres, Row, Transaction};
|
||||
|
||||
use gateway_core::engine::credential::{
|
||||
AccountConcurrencyLimit, AccountErrorReason, AccountStateChange, AccountWeight,
|
||||
CredentialCasOutcome, CredentialCasUpdate, CredentialRevision as CoreCredentialRevision,
|
||||
CredentialState, LoadedCredential, NewProviderAccount as CoreNewProviderAccount,
|
||||
OpaqueProviderData, PlaintextCredential, ProviderAccount as CoreProviderAccount,
|
||||
ProviderAccountId as CoreProviderAccountId, ProviderAccountIdentity, ProviderAccountStore,
|
||||
CredentialCasOutcome, CredentialCasUpdate, CredentialCasUpdateParts,
|
||||
CredentialRevision as CoreCredentialRevision, CredentialState, LoadedCredential,
|
||||
NewProviderAccount as CoreNewProviderAccount, OpaqueProviderData, PlaintextCredential,
|
||||
ProviderAccount as CoreProviderAccount, ProviderAccountId as CoreProviderAccountId,
|
||||
ProviderAccountIdentity, ProviderAccountStore,
|
||||
ProviderAccountUpdate as CoreProviderAccountUpdate, QuotaAccessChange, QuotaAccessState,
|
||||
QuotaEvidence, QuotaObservation, QuotaState, QuotaWriteOutcome,
|
||||
QuotaEvidence, QuotaObservation, QuotaObservationTouch, QuotaState, QuotaWriteOutcome,
|
||||
};
|
||||
use gateway_core::error::{StoreError as CoreStoreError, StoreErrorKind as CoreStoreErrorKind};
|
||||
use gateway_core::routing::{AccountGroupId, ProviderKind};
|
||||
|
||||
@@ -29,6 +29,12 @@ pub trait ProviderAccountRepository: Send + Sync {
|
||||
observed_at: DateTime<Utc>,
|
||||
state: QuotaState,
|
||||
) -> StoreResult<bool>;
|
||||
async fn touch_provider_quota_observation(
|
||||
&self,
|
||||
account_id: &str,
|
||||
expected_revision: Revision,
|
||||
observed_at: DateTime<Utc>,
|
||||
) -> StoreResult<bool>;
|
||||
async fn apply_provider_quota_access(
|
||||
&self,
|
||||
account_id: &str,
|
||||
@@ -237,13 +243,11 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
else credential_observed_at
|
||||
end,
|
||||
last_error_reason = case
|
||||
when enabled and upstream_user_id is not null then
|
||||
case when $3 = 'ready' then null else $5 end
|
||||
when enabled and upstream_user_id is not null then $5
|
||||
else last_error_reason
|
||||
end,
|
||||
last_error_message = case
|
||||
when enabled and upstream_user_id is not null then
|
||||
case when $3 = 'ready' then null else $6 end
|
||||
when enabled and upstream_user_id is not null then $6
|
||||
else last_error_message
|
||||
end,
|
||||
updated_at = case when enabled then greatest(now(), $4) else updated_at end
|
||||
@@ -354,6 +358,29 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
Ok(result.rows_affected() == 1)
|
||||
}
|
||||
|
||||
async fn touch_provider_quota_observation(
|
||||
&self,
|
||||
account_id: &str,
|
||||
expected_revision: Revision,
|
||||
observed_at: DateTime<Utc>,
|
||||
) -> StoreResult<bool> {
|
||||
require_nonempty(ENTITY, "account_id", account_id)?;
|
||||
let result = sqlx::query(
|
||||
"update provider_accounts
|
||||
set quota_observed_at = $3, updated_at = greatest(now(), $3)
|
||||
where id = $1 and credential_revision = $2
|
||||
and provider_quota_json is not null
|
||||
and (quota_observed_at is null or quota_observed_at <= $3)",
|
||||
)
|
||||
.bind(account_id)
|
||||
.bind(to_i64(expected_revision.get())?)
|
||||
.bind(observed_at)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("touch provider quota observation"))?;
|
||||
Ok(result.rows_affected() == 1)
|
||||
}
|
||||
|
||||
async fn delete_provider_account(&self, id: &str) -> StoreResult<bool> {
|
||||
require_nonempty(ENTITY, "id", id)?;
|
||||
let result = sqlx::query("delete from provider_accounts where id = $1 and not enabled")
|
||||
|
||||
@@ -0,0 +1,175 @@
|
||||
//! Provider 官方制品画像的单键 Redis cache。
|
||||
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use gateway_core::engine::credential::OpaqueProviderData;
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderArtifactProfile, ProviderArtifactProfileCachePort, ProviderStoreError,
|
||||
ProviderStoreErrorKind,
|
||||
};
|
||||
use gateway_core::routing::ProviderKind;
|
||||
use redis::{Script, aio::ConnectionManager};
|
||||
|
||||
use crate::StoreResult;
|
||||
|
||||
use super::{MAX_REDIS_EXACT_INTEGER, namespace};
|
||||
|
||||
const MAX_PROFILE_BYTES: usize = 16 * 1024;
|
||||
|
||||
const REPLACE_SCRIPT: &str = r#"
|
||||
local current_sequence = redis.call('HGET', KEYS[1], 'artifact_sequence')
|
||||
if current_sequence then
|
||||
local current_number = tonumber(current_sequence)
|
||||
local candidate_number = tonumber(ARGV[1])
|
||||
if current_number > candidate_number then return 0 end
|
||||
if current_number == candidate_number then
|
||||
local current_profile = redis.call('HGET', KEYS[1], 'profile')
|
||||
if current_profile ~= ARGV[3] then return -1 end
|
||||
end
|
||||
end
|
||||
redis.call('HSET', KEYS[1],
|
||||
'artifact_sequence', ARGV[1],
|
||||
'verified_at_ms', ARGV[2],
|
||||
'profile', ARGV[3])
|
||||
redis.call('PEXPIRE', KEYS[1], ARGV[4])
|
||||
return 1
|
||||
"#;
|
||||
|
||||
/// 每个 Provider 只占一个覆盖写 key;TTL 负责进程长期停机后的最终清理。
|
||||
#[derive(Clone)]
|
||||
pub struct RedisProviderArtifactProfileRepository {
|
||||
connection: ConnectionManager,
|
||||
namespace: String,
|
||||
}
|
||||
|
||||
impl RedisProviderArtifactProfileRepository {
|
||||
pub fn new(connection: ConnectionManager, key_namespace: &str) -> StoreResult<Self> {
|
||||
Ok(Self {
|
||||
connection,
|
||||
namespace: namespace(key_namespace)?,
|
||||
})
|
||||
}
|
||||
|
||||
fn key(&self, provider_kind: &ProviderKind) -> String {
|
||||
format!(
|
||||
"{}:provider:{}:artifact-profile:v1",
|
||||
self.namespace,
|
||||
provider_kind.as_str()
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderArtifactProfileCachePort for RedisProviderArtifactProfileRepository {
|
||||
fn replace_if_newer(
|
||||
&self,
|
||||
profile: ProviderArtifactProfile,
|
||||
ttl: Duration,
|
||||
) -> futures::future::BoxFuture<'_, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async move {
|
||||
let artifact_sequence = profile.artifact_sequence();
|
||||
if artifact_sequence == 0 || artifact_sequence > MAX_REDIS_EXACT_INTEGER {
|
||||
return Err(provider_invalid("validate artifact profile sequence"));
|
||||
}
|
||||
let ttl_ms = exact_positive_millis(ttl, "validate artifact profile TTL")?;
|
||||
let verified_at_ms = profile
|
||||
.verified_at()
|
||||
.duration_since(SystemTime::UNIX_EPOCH)
|
||||
.ok()
|
||||
.and_then(|duration| u64::try_from(duration.as_millis()).ok())
|
||||
.filter(|millis| *millis <= MAX_REDIS_EXACT_INTEGER)
|
||||
.ok_or_else(|| provider_invalid("validate artifact profile timestamp"))?;
|
||||
let payload = serde_json::to_vec(&serde_json::Value::Object(
|
||||
profile.profile().expose_to_provider().clone(),
|
||||
))
|
||||
.map_err(|_| provider_invalid("encode artifact profile"))?;
|
||||
if payload.len() > MAX_PROFILE_BYTES {
|
||||
return Err(provider_invalid("validate artifact profile size"));
|
||||
}
|
||||
|
||||
let mut connection = self.connection.clone();
|
||||
let outcome = Script::new(REPLACE_SCRIPT)
|
||||
.key(self.key(profile.provider_kind()))
|
||||
.arg(artifact_sequence)
|
||||
.arg(verified_at_ms)
|
||||
.arg(payload)
|
||||
.arg(ttl_ms)
|
||||
.invoke_async::<i64>(&mut connection)
|
||||
.await
|
||||
.map_err(|_| provider_unavailable("replace artifact profile"))?;
|
||||
match outcome {
|
||||
1 => Ok(true),
|
||||
0 => Ok(false),
|
||||
-1 => Err(ProviderStoreError::new(
|
||||
ProviderStoreErrorKind::Conflict,
|
||||
"replace artifact profile",
|
||||
)),
|
||||
_ => Err(provider_invalid("decode artifact profile write outcome")),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn read<'a>(
|
||||
&'a self,
|
||||
provider_kind: &'a ProviderKind,
|
||||
) -> futures::future::BoxFuture<'a, Result<Option<ProviderArtifactProfile>, ProviderStoreError>>
|
||||
{
|
||||
Box::pin(async move {
|
||||
let mut connection = self.connection.clone();
|
||||
let values: (Option<u64>, Option<u64>, Option<Vec<u8>>) = redis::cmd("HMGET")
|
||||
.arg(self.key(provider_kind))
|
||||
.arg("artifact_sequence")
|
||||
.arg("verified_at_ms")
|
||||
.arg("profile")
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.map_err(|_| provider_unavailable("read artifact profile"))?;
|
||||
let (artifact_sequence, verified_at_ms, payload) = values;
|
||||
if artifact_sequence.is_none() && verified_at_ms.is_none() && payload.is_none() {
|
||||
return Ok(None);
|
||||
}
|
||||
let (Some(artifact_sequence), Some(verified_at_ms), Some(payload)) =
|
||||
(artifact_sequence, verified_at_ms, payload)
|
||||
else {
|
||||
return Err(provider_invalid("decode artifact profile"));
|
||||
};
|
||||
if artifact_sequence == 0 || artifact_sequence > MAX_REDIS_EXACT_INTEGER {
|
||||
return Err(provider_invalid("decode artifact profile sequence"));
|
||||
}
|
||||
if payload.len() > MAX_PROFILE_BYTES {
|
||||
return Err(provider_invalid("decode artifact profile size"));
|
||||
}
|
||||
let serde_json::Value::Object(fields) = serde_json::from_slice(&payload)
|
||||
.map_err(|_| provider_invalid("decode artifact profile"))?
|
||||
else {
|
||||
return Err(provider_invalid("decode artifact profile"));
|
||||
};
|
||||
let verified_at = SystemTime::UNIX_EPOCH
|
||||
.checked_add(Duration::from_millis(verified_at_ms))
|
||||
.ok_or_else(|| provider_invalid("decode artifact profile timestamp"))?;
|
||||
Ok(Some(ProviderArtifactProfile::new(
|
||||
provider_kind.clone(),
|
||||
artifact_sequence,
|
||||
verified_at,
|
||||
OpaqueProviderData::new(fields),
|
||||
)))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn exact_positive_millis(
|
||||
duration: Duration,
|
||||
operation: &'static str,
|
||||
) -> Result<u64, ProviderStoreError> {
|
||||
u64::try_from(duration.as_millis())
|
||||
.ok()
|
||||
.filter(|millis| *millis > 0 && *millis <= MAX_REDIS_EXACT_INTEGER)
|
||||
.ok_or_else(|| provider_invalid(operation))
|
||||
}
|
||||
|
||||
fn provider_unavailable(operation: &'static str) -> ProviderStoreError {
|
||||
ProviderStoreError::new(ProviderStoreErrorKind::Unavailable, operation)
|
||||
}
|
||||
|
||||
fn provider_invalid(operation: &'static str) -> ProviderStoreError {
|
||||
ProviderStoreError::new(ProviderStoreErrorKind::InvalidData, operation)
|
||||
}
|
||||
@@ -7,6 +7,7 @@ use sha2::{Digest, Sha256};
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
mod artifact_profile;
|
||||
mod client_admission;
|
||||
mod coordination_buffer;
|
||||
mod credential_cooldown;
|
||||
@@ -20,6 +21,7 @@ mod provider_session_exclusion;
|
||||
mod runtime_change;
|
||||
pub(crate) mod worker_lease;
|
||||
|
||||
pub use artifact_profile::*;
|
||||
pub use client_admission::*;
|
||||
pub use coordination_buffer::*;
|
||||
pub use credential_cooldown::*;
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
use gateway_store::StoreConfig;
|
||||
|
||||
const PASSWORD: &str = "111111111111111111111111111111111111111111111111";
|
||||
|
||||
#[test]
|
||||
fn store_config_should_derive_backup_staging_from_runtime_data_dir() {
|
||||
let root = tempfile::tempdir().expect("runtime data root");
|
||||
let mut config = valid_config();
|
||||
|
||||
config
|
||||
.resolve_and_validate(root.path())
|
||||
.expect("valid Store configuration");
|
||||
|
||||
assert_eq!(
|
||||
config.backup_staging_dir(),
|
||||
root.path().join("backup-staging")
|
||||
);
|
||||
}
|
||||
|
||||
fn valid_config() -> StoreConfig {
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"database": {
|
||||
"url": "postgres://codex_proxy@127.0.0.1:5432/codex_proxy",
|
||||
"password": PASSWORD,
|
||||
},
|
||||
"redis": {
|
||||
"url": "redis://127.0.0.1:6379/",
|
||||
"password": PASSWORD,
|
||||
},
|
||||
}))
|
||||
.expect("test Store configuration")
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
mod config;
|
||||
mod postgres;
|
||||
mod redis;
|
||||
mod s3;
|
||||
|
||||
@@ -730,8 +730,18 @@ async fn dashboard_summary_totals_include_history_outside_selected_range() {
|
||||
.await
|
||||
.expect("dashboard summary");
|
||||
assert_eq!(dashboard.requests.request_count, 2);
|
||||
assert_eq!(dashboard.account_usage[0].request_count, 0);
|
||||
assert_eq!(dashboard.account_usage[0].request_buckets.len(), 2);
|
||||
assert!(dashboard.account_usage.is_empty());
|
||||
|
||||
let explicit_account_usage = repository
|
||||
.provider_account_usage(
|
||||
ProviderAccountUsageQuery::for_accounts(range, vec!["acct_observe".to_owned()])
|
||||
.expect("account usage query"),
|
||||
)
|
||||
.await
|
||||
.expect("provider account usage");
|
||||
assert_eq!(explicit_account_usage.len(), 1);
|
||||
assert_eq!(explicit_account_usage[0].request_count, 0);
|
||||
|
||||
assert_eq!(
|
||||
(
|
||||
dashboard.totals.request_count,
|
||||
|
||||
@@ -18,8 +18,8 @@ use gateway_core::engine::credential::{
|
||||
AccountErrorReason, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
|
||||
CredentialRevision, CredentialState, OpaqueProviderData, PlaintextCredential,
|
||||
ProviderAccountId, ProviderAccountIdentity, ProviderAccountStore, ProviderAccountUpdate,
|
||||
QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaObservation, QuotaState,
|
||||
QuotaWriteOutcome,
|
||||
QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaObservation, QuotaObservationTouch,
|
||||
QuotaState, QuotaWriteOutcome,
|
||||
};
|
||||
use gateway_core::routing::{AccountGroupId, ProviderKind};
|
||||
use gateway_store::{
|
||||
@@ -51,6 +51,22 @@ struct RecoveredAccountRow {
|
||||
access_token_expires_at: Option<chrono::DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[derive(sqlx::FromRow)]
|
||||
struct CoreRefreshRow {
|
||||
name: String,
|
||||
email: Option<String>,
|
||||
plan_type: Option<String>,
|
||||
provider_credentials_json: serde_json::Value,
|
||||
credential_revision: i64,
|
||||
credential_state: String,
|
||||
last_error_reason: Option<String>,
|
||||
last_error_message: Option<String>,
|
||||
quota_access_state: String,
|
||||
quota_evidence: Option<String>,
|
||||
quota_access_observed_at: Option<chrono::DateTime<Utc>>,
|
||||
quota_reset_at: Option<chrono::DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn postgres_provider_account_adapter_implements_core_port() {
|
||||
fn assert_port<T: ProviderAccountStore>() {}
|
||||
@@ -156,6 +172,115 @@ async fn core_quota_batch_reads_only_observed_accounts_in_one_contract_call() {
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_observation_touch_preserves_quota_state_and_advances_account_update_time() {
|
||||
let Some(database) = TestDatabase::create("provider_account_quota_touch").await else {
|
||||
return;
|
||||
};
|
||||
let repository = PgProviderAccountRepository::new(database.pool.clone());
|
||||
let account_id = ProviderAccountId::new("acct_quota_touch").expect("account id");
|
||||
repository
|
||||
.insert_provider_account(account(account_id.as_str(), "user-quota-touch"))
|
||||
.await
|
||||
.expect("insert quota fixture");
|
||||
let revision = CredentialRevision::new(1).expect("revision");
|
||||
let observed_at = SystemTime::now()
|
||||
.checked_sub(Duration::from_secs(60))
|
||||
.expect("old observation time");
|
||||
let touched_at = observed_at
|
||||
.checked_add(Duration::from_secs(90))
|
||||
.expect("new observation time");
|
||||
let reset_at = observed_at
|
||||
.checked_add(Duration::from_secs(3_600))
|
||||
.expect("quota reset time");
|
||||
let quota_json = json!({
|
||||
"rate_limit": {
|
||||
"allowed": false,
|
||||
"limit_reached": true,
|
||||
"primary_window": {"used_percent": 100, "reset_at": 1_900_000_000}
|
||||
}
|
||||
});
|
||||
assert_eq!(
|
||||
repository
|
||||
.compare_and_swap_quota(QuotaObservation {
|
||||
account_id: account_id.clone(),
|
||||
expected_revision: revision,
|
||||
quota: OpaqueProviderData::new(
|
||||
quota_json.as_object().expect("quota object").clone(),
|
||||
),
|
||||
observed_at,
|
||||
state: QuotaState::exhausted(
|
||||
QuotaEvidence::UsageLimitReached,
|
||||
observed_at,
|
||||
Some(reset_at),
|
||||
),
|
||||
})
|
||||
.await
|
||||
.expect("persist exhausted quota"),
|
||||
QuotaWriteOutcome::Updated
|
||||
);
|
||||
let account_updated_at: chrono::DateTime<Utc> = sqlx::query_scalar(
|
||||
"select updated_at from provider_accounts where id = 'acct_quota_touch'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load account update time before quota touch");
|
||||
|
||||
assert_eq!(
|
||||
repository
|
||||
.touch_quota_observation(QuotaObservationTouch {
|
||||
account_id,
|
||||
expected_revision: revision,
|
||||
observed_at: touched_at,
|
||||
})
|
||||
.await
|
||||
.expect("touch quota refresh time"),
|
||||
QuotaWriteOutcome::Updated
|
||||
);
|
||||
|
||||
let stored = sqlx::query_as::<
|
||||
_,
|
||||
(
|
||||
serde_json::Value,
|
||||
chrono::DateTime<Utc>,
|
||||
String,
|
||||
Option<String>,
|
||||
chrono::DateTime<Utc>,
|
||||
Option<chrono::DateTime<Utc>>,
|
||||
chrono::DateTime<Utc>,
|
||||
),
|
||||
>(
|
||||
"select provider_quota_json, quota_observed_at, quota_access_state,
|
||||
quota_evidence, quota_access_observed_at, quota_reset_at, updated_at
|
||||
from provider_accounts where id = 'acct_quota_touch'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load touched quota row");
|
||||
assert_eq!(stored.0, quota_json);
|
||||
assert_eq!(
|
||||
stored.1.timestamp_micros(),
|
||||
chrono::DateTime::<Utc>::from(touched_at).timestamp_micros()
|
||||
);
|
||||
assert_eq!(stored.2, "exhausted");
|
||||
assert_eq!(stored.3.as_deref(), Some("usage_limit_reached"));
|
||||
assert_eq!(
|
||||
stored.4.timestamp_micros(),
|
||||
chrono::DateTime::<Utc>::from(observed_at).timestamp_micros()
|
||||
);
|
||||
assert_eq!(
|
||||
stored.5.map(|value| value.timestamp_micros()),
|
||||
Some(chrono::DateTime::<Utc>::from(reset_at).timestamp_micros())
|
||||
);
|
||||
assert_eq!(
|
||||
stored.6.timestamp_micros(),
|
||||
chrono::DateTime::<Utc>::from(touched_at).timestamp_micros()
|
||||
);
|
||||
assert!(stored.6 > account_updated_at);
|
||||
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn delayed_provider_observations_cannot_overwrite_newer_state_or_quota() {
|
||||
let Some(database) = TestDatabase::create("provider_account_observation_fence").await else {
|
||||
@@ -280,6 +405,27 @@ async fn credential_error_message_is_persisted_and_cleared_on_recovery() {
|
||||
Some("upstream returned 402 payment required")
|
||||
);
|
||||
|
||||
repository
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: Some("Invalid refresh token.".to_owned()),
|
||||
account_id: account_id.clone(),
|
||||
expected_revision: revision,
|
||||
credential_state: CredentialState::Ready,
|
||||
observed_at: SystemTime::now(),
|
||||
error_reason: Some(AccountErrorReason::AccessTokenExpired),
|
||||
})
|
||||
.await
|
||||
.expect("persist recoverable refresh failure");
|
||||
let recoverable: (Option<String>, Option<String>) = sqlx::query_as(
|
||||
"select last_error_reason, last_error_message from provider_accounts where id = $1",
|
||||
)
|
||||
.bind(account_id.as_str())
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("read recoverable refresh failure");
|
||||
assert_eq!(recoverable.0.as_deref(), Some("access_token_expired"));
|
||||
assert_eq!(recoverable.1.as_deref(), Some("Invalid refresh token."));
|
||||
|
||||
repository
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
@@ -1208,8 +1354,37 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
|
||||
})
|
||||
.await
|
||||
.expect("seed provider account");
|
||||
sqlx::query(
|
||||
"update provider_accounts
|
||||
set last_error_reason = 'access_token_expired',
|
||||
last_error_message = 'stale refresh failure'
|
||||
where id = 'acct_core_refresh'",
|
||||
)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("seed refresh failure details");
|
||||
|
||||
let account_id = ProviderAccountId::new("acct_core_refresh").expect("account ID");
|
||||
let quota_observed_at = SystemTime::now();
|
||||
let quota_reset_at = quota_observed_at
|
||||
.checked_add(Duration::from_secs(3_600))
|
||||
.expect("quota reset time");
|
||||
assert_eq!(
|
||||
repository
|
||||
.apply_quota_access(QuotaAccessChange {
|
||||
account_id: account_id.clone(),
|
||||
expected_revision: CredentialRevision::new(1).expect("credential revision"),
|
||||
state: QuotaState::exhausted(
|
||||
QuotaEvidence::UsageLimitReached,
|
||||
quota_observed_at,
|
||||
Some(quota_reset_at),
|
||||
),
|
||||
})
|
||||
.await
|
||||
.expect("seed quota exhaustion"),
|
||||
QuotaWriteOutcome::Updated
|
||||
);
|
||||
let state_observed_at = SystemTime::now();
|
||||
let refreshed = CredentialCasUpdate::new(
|
||||
account_id.clone(),
|
||||
CredentialRevision::new(1).expect("credential revision"),
|
||||
@@ -1224,7 +1399,8 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
|
||||
Some(SystemTime::now() + Duration::from_secs(3_600)),
|
||||
Some(SystemTime::now() + Duration::from_secs(1_800)),
|
||||
)
|
||||
.expect("valid refresh update");
|
||||
.expect("valid refresh update")
|
||||
.with_account_state(CredentialState::Ready, state_observed_at, None, None);
|
||||
assert_eq!(
|
||||
repository
|
||||
.compare_and_swap_credential(refreshed)
|
||||
@@ -1235,24 +1411,37 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
|
||||
)
|
||||
);
|
||||
|
||||
let row: (
|
||||
String,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
serde_json::Value,
|
||||
i64,
|
||||
) = sqlx::query_as(
|
||||
"select name, email, plan_type, provider_credentials_json, credential_revision
|
||||
let row = sqlx::query_as::<_, CoreRefreshRow>(
|
||||
"select name, email, plan_type, provider_credentials_json, credential_revision,
|
||||
credential_state, last_error_reason, last_error_message,
|
||||
quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at
|
||||
from provider_accounts where id = 'acct_core_refresh'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load refreshed account");
|
||||
assert_eq!(row.0, "after refresh");
|
||||
assert_eq!(row.1.as_deref(), Some("after@example.invalid"));
|
||||
assert_eq!(row.2.as_deref(), Some("premium"));
|
||||
assert_eq!(row.3["access_token"], "after-secret");
|
||||
assert_eq!(row.4, 2);
|
||||
assert_eq!(row.name, "after refresh");
|
||||
assert_eq!(row.email.as_deref(), Some("after@example.invalid"));
|
||||
assert_eq!(row.plan_type.as_deref(), Some("premium"));
|
||||
assert_eq!(
|
||||
row.provider_credentials_json["access_token"],
|
||||
"after-secret"
|
||||
);
|
||||
assert_eq!(row.credential_revision, 2);
|
||||
assert_eq!(row.credential_state, "ready");
|
||||
assert_eq!(row.last_error_reason, None);
|
||||
assert_eq!(row.last_error_message, None);
|
||||
assert_eq!(row.quota_access_state, "exhausted");
|
||||
assert_eq!(row.quota_evidence.as_deref(), Some("usage_limit_reached"));
|
||||
assert_eq!(
|
||||
row.quota_access_observed_at
|
||||
.map(|value| value.timestamp_micros()),
|
||||
Some(chrono::DateTime::<Utc>::from(quota_observed_at).timestamp_micros())
|
||||
);
|
||||
assert_eq!(
|
||||
row.quota_reset_at.map(|value| value.timestamp_micros()),
|
||||
Some(chrono::DateTime::<Utc>::from(quota_reset_at).timestamp_micros())
|
||||
);
|
||||
|
||||
assert!(
|
||||
repository
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use gateway_core::engine::credential::OpaqueProviderData;
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderArtifactProfile, ProviderArtifactProfileCachePort, ProviderStoreErrorKind,
|
||||
};
|
||||
use gateway_core::routing::ProviderKind;
|
||||
use gateway_store::redis::RedisProviderArtifactProfileRepository;
|
||||
use redis::aio::ConnectionManager;
|
||||
use serde_json::{Map, Value};
|
||||
use uuid::Uuid;
|
||||
|
||||
#[test]
|
||||
fn artifact_profile_adapter_implements_provider_port() {
|
||||
fn assert_port<T: ProviderArtifactProfileCachePort>() {}
|
||||
assert_port::<RedisProviderArtifactProfileRepository>();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn artifact_profile_uses_one_expiring_key_and_rejects_rollback_or_conflict() {
|
||||
let Some((repository, mut connection, namespace)) = repository().await else {
|
||||
return;
|
||||
};
|
||||
let provider = ProviderKind::new("openai").expect("provider kind");
|
||||
let current = profile(provider.clone(), 6_570, "0.148.0-alpha.9");
|
||||
assert!(
|
||||
repository
|
||||
.replace_if_newer(current.clone(), Duration::from_secs(60))
|
||||
.await
|
||||
.expect("store current profile")
|
||||
);
|
||||
assert_eq!(
|
||||
repository.read(&provider).await.expect("read profile"),
|
||||
Some(current.clone())
|
||||
);
|
||||
|
||||
assert!(
|
||||
!repository
|
||||
.replace_if_newer(
|
||||
profile(provider.clone(), 6_415, "0.147.0-alpha.6.6"),
|
||||
Duration::from_secs(60),
|
||||
)
|
||||
.await
|
||||
.expect("reject older profile")
|
||||
);
|
||||
let conflict = repository
|
||||
.replace_if_newer(
|
||||
profile(provider.clone(), 6_570, "0.148.0-alpha.10"),
|
||||
Duration::from_secs(60),
|
||||
)
|
||||
.await
|
||||
.expect_err("same build with different Core must conflict");
|
||||
assert_eq!(conflict.kind(), ProviderStoreErrorKind::Conflict);
|
||||
assert_eq!(
|
||||
repository
|
||||
.read(&provider)
|
||||
.await
|
||||
.expect("read retained profile"),
|
||||
Some(current)
|
||||
);
|
||||
|
||||
let keys = redis::cmd("KEYS")
|
||||
.arg(format!("{namespace}:*"))
|
||||
.query_async::<Vec<String>>(&mut connection)
|
||||
.await
|
||||
.expect("list isolated keys");
|
||||
assert_eq!(keys.len(), 1);
|
||||
assert_eq!(
|
||||
keys[0],
|
||||
format!("{namespace}:provider:openai:artifact-profile:v1")
|
||||
);
|
||||
let ttl = redis::cmd("PTTL")
|
||||
.arg(&keys[0])
|
||||
.query_async::<i64>(&mut connection)
|
||||
.await
|
||||
.expect("read profile TTL");
|
||||
assert!((1..=60_000).contains(&ttl));
|
||||
}
|
||||
|
||||
fn profile(
|
||||
provider: ProviderKind,
|
||||
artifact_sequence: u64,
|
||||
codex_version: &str,
|
||||
) -> ProviderArtifactProfile {
|
||||
let mut fields = Map::new();
|
||||
fields.insert("schema_version".to_owned(), Value::from(1));
|
||||
fields.insert(
|
||||
"codex_version".to_owned(),
|
||||
Value::String(codex_version.to_owned()),
|
||||
);
|
||||
ProviderArtifactProfile::new(
|
||||
provider,
|
||||
artifact_sequence,
|
||||
SystemTime::UNIX_EPOCH + Duration::from_secs(1_800_000_000),
|
||||
OpaqueProviderData::new(fields),
|
||||
)
|
||||
}
|
||||
|
||||
async fn repository() -> Option<(
|
||||
RedisProviderArtifactProfileRepository,
|
||||
ConnectionManager,
|
||||
String,
|
||||
)> {
|
||||
let redis_url = crate::support::test_env("CPR_TEST_REDIS_URL")?;
|
||||
let client = redis::Client::open(redis_url).expect("valid CPR_TEST_REDIS_URL");
|
||||
let connection = client
|
||||
.get_connection_manager()
|
||||
.await
|
||||
.expect("connect test Redis");
|
||||
let namespace = format!("gateway-store-test-{}", Uuid::new_v4());
|
||||
let repository = RedisProviderArtifactProfileRepository::new(connection.clone(), &namespace)
|
||||
.expect("valid test namespace");
|
||||
Some((repository, connection, namespace))
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
mod artifact_profile;
|
||||
mod client_admission;
|
||||
mod coordination_buffer;
|
||||
mod credential_cooldown;
|
||||
|
||||
@@ -21,12 +21,18 @@ async fn store_bundle_worker_plan_and_leader_lease_are_single_use_and_fenced() {
|
||||
return;
|
||||
};
|
||||
let database = TestDatabase::create(&database_url).await;
|
||||
let config = store_config(&database.url, &redis_url);
|
||||
let runtime_data = tempfile::tempdir().expect("Store runtime data");
|
||||
let mut config = store_config(&database.url, &redis_url);
|
||||
config
|
||||
.resolve_and_validate(runtime_data.path())
|
||||
.expect("resolved Store config");
|
||||
let mut first = initialize(config.clone())
|
||||
.await
|
||||
.expect("first Store bundle");
|
||||
let second = initialize(config).await.expect("second Store bundle");
|
||||
|
||||
assert!(runtime_data.path().join("backup-staging").is_dir());
|
||||
|
||||
let kinds = first
|
||||
.take_worker_contributions()
|
||||
.into_iter()
|
||||
|
||||
@@ -17,6 +17,7 @@ cookie = "0.18.1"
|
||||
crypto_box = { version = "0.9.1", features = ["seal"] }
|
||||
ed25519-dalek = { version = "3.0.0", features = ["alloc", "pkcs8"] }
|
||||
futures = "0.3.32"
|
||||
flate2 = "1.1.9"
|
||||
gateway-admin = { path = "../../gateway-admin" }
|
||||
gateway-core = { path = "../../gateway-core" }
|
||||
gateway-protocol = { path = "../../gateway-protocol" }
|
||||
@@ -49,6 +50,7 @@ tokio-tungstenite = { git = "https://github.com/openai-oss-forks/tokio-tungsteni
|
||||
tracing = "0.1.44"
|
||||
tungstenite = { git = "https://github.com/openai-oss-forks/tungstenite-rs", rev = "4fffad30fe373adbdcffab9545e9e9bf4f2fc19f", features = ["deflate", "proxy"] }
|
||||
url = "2.5.8"
|
||||
urlencoding = "2.1.3"
|
||||
uuid = { version = "1.23.3", features = ["v4", "v7"] }
|
||||
# 官方 Codex app-server 对 /v1/responses 请求体默认做 zstd 压缩(level 3),
|
||||
# 并携带 Content-Encoding: zstd;保持同版本同级别以对齐上游指纹。
|
||||
|
||||
@@ -19,14 +19,15 @@ use gateway_admin::model::provider_credentials::{
|
||||
PreparedCredentialCreate, PreparedCredentialImport, PreparedCredentialRotation,
|
||||
PreparedCredentialRotationFacts, ProviderDocument, ProviderExport,
|
||||
ProviderExportCredentialInput, ProviderModel, ProviderModels, ProviderQuota,
|
||||
ProviderQuotaRequest, ProviderQuotaWindow, QuotaLocalUsageAttribution,
|
||||
ProviderQuotaRequest, ProviderQuotaWindow, ProviderQuotaWindowRole, QuotaLocalUsageAttribution,
|
||||
};
|
||||
use gateway_admin::model::{MutationActor, MutationContext, Revision};
|
||||
use gateway_admin::ports::provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind};
|
||||
use gateway_core::accounting::Money;
|
||||
use gateway_core::engine::credential::{
|
||||
CredentialRevision, LoadedCredential, NewProviderAccount, OpaqueProviderData,
|
||||
PlaintextCredential, ProviderAccount, ProviderAccountId, ProviderAccountStore,
|
||||
CredentialCasUpdateParts, CredentialRevision, LoadedCredential, NewProviderAccount,
|
||||
OpaqueProviderData, PlaintextCredential, ProviderAccount, ProviderAccountId,
|
||||
ProviderAccountStore,
|
||||
};
|
||||
use gateway_core::error::StoreErrorKind;
|
||||
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
|
||||
@@ -49,12 +50,14 @@ use crate::credential::{
|
||||
CompleteCodexOAuthAuthorization, CompletedCodexOAuthCredential, ExportManagedCodexCredential,
|
||||
StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
|
||||
};
|
||||
use crate::credential::{CodexCredentialCodec, CodexOAuthSecret, oauth_owner_ref};
|
||||
use crate::credential::{
|
||||
CodexCredentialCodec, CodexOAuthSecret, oauth_owner_ref, parse_access_token_expiration,
|
||||
};
|
||||
use crate::transport::CodexWebSocketPool;
|
||||
use crate::transport::openai_billing_breakdown;
|
||||
use crate::transport::profile::{
|
||||
CodexDesktopReleaseSnapshot, CodexDesktopReleaseStatus, CodexWireProfile, CodexWireProfileState,
|
||||
};
|
||||
use crate::transport::{OpenAiBillingUsage, openai_billing_breakdown};
|
||||
|
||||
const PROVIDER_NAME: &str = "openai";
|
||||
const PENDING_DOCUMENT_SCHEMA_VERSION: u64 = 3;
|
||||
@@ -215,10 +218,12 @@ impl ProviderAdmin for OpenAiAdminProvider {
|
||||
};
|
||||
let Some(breakdown) = openai_billing_breakdown(
|
||||
&input.upstream_model_id,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
input.cached_tokens.unwrap_or_default(),
|
||||
input.cache_write_tokens.unwrap_or_default(),
|
||||
OpenAiBillingUsage::new(
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
input.cached_tokens.unwrap_or_default(),
|
||||
input.cache_write_tokens.unwrap_or_default(),
|
||||
),
|
||||
input.service_tier.as_deref(),
|
||||
) else {
|
||||
return Ok(None);
|
||||
@@ -256,7 +261,7 @@ impl ProviderAdmin for OpenAiAdminProvider {
|
||||
))
|
||||
.await
|
||||
.inspect_err(|error| {
|
||||
log_import_failure("prepare_document", credential_admin_error_code(*error));
|
||||
log_import_failure("prepare_document", credential_admin_error_code(error));
|
||||
})
|
||||
.map_err(map_credential_admin_error)?;
|
||||
let observed_at = Utc::now();
|
||||
@@ -378,8 +383,10 @@ impl ProviderAdmin for OpenAiAdminProvider {
|
||||
.oauth()
|
||||
.and_then(|oauth| oauth.id_token.clone());
|
||||
}
|
||||
let access_token_expires_at =
|
||||
parse_access_token_expiration(secret.access_token.expose_secret());
|
||||
let prepared = CodexCredentialAdmin
|
||||
.prepare_refreshed_oauth_rotation(current, secret, None, None)
|
||||
.prepare_refreshed_oauth_rotation(current, secret, access_token_expires_at, None)
|
||||
.map_err(map_credential_admin_error)?;
|
||||
prepared_rotation(prepared, command.account.provider_kind)
|
||||
}
|
||||
@@ -544,15 +551,16 @@ fn prepared_rotation(
|
||||
provider_kind: ProviderKind,
|
||||
) -> Result<PreparedCredentialRotation, ProviderAdminError> {
|
||||
let (profile, credential, replacement_identity, guard) = prepared.into_parts();
|
||||
let (
|
||||
let CredentialCasUpdateParts {
|
||||
account_id,
|
||||
expected_revision,
|
||||
credential_profile,
|
||||
profile: credential_profile,
|
||||
credential,
|
||||
has_refresh_token,
|
||||
access_token_expires_at,
|
||||
next_refresh_at,
|
||||
) = credential.into_parts();
|
||||
account_state: _account_state,
|
||||
} = credential.into_parts();
|
||||
if profile != credential_profile || profile.account_id != account_id {
|
||||
return Err(provider_admin_error(ProviderAdminErrorKind::Internal));
|
||||
}
|
||||
@@ -726,39 +734,24 @@ fn project_quota_snapshot(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota
|
||||
.windows()
|
||||
.iter()
|
||||
.filter(|window| should_project_quota_window(window))
|
||||
.map(|window| {
|
||||
let mut data = Map::new();
|
||||
data.insert(
|
||||
"source".to_owned(),
|
||||
Value::String(window.source().to_owned()),
|
||||
);
|
||||
data.insert(
|
||||
"role".to_owned(),
|
||||
Value::String(quota_role(window.role()).to_owned()),
|
||||
);
|
||||
ProviderQuotaWindow {
|
||||
key: window.key().to_owned(),
|
||||
group: quota_group(window.kind()).to_owned(),
|
||||
label: codex_quota_window_label(
|
||||
window.limit_name(),
|
||||
window.kind(),
|
||||
window.role(),
|
||||
window.source(),
|
||||
window.window_seconds(),
|
||||
),
|
||||
source: Some(window.source().to_owned()),
|
||||
local_usage_attribution: if window.is_account_wide() {
|
||||
QuotaLocalUsageAttribution::AccountWide
|
||||
} else {
|
||||
QuotaLocalUsageAttribution::Unavailable
|
||||
},
|
||||
window_seconds: window.window_seconds(),
|
||||
used_percent: window.used_percent(),
|
||||
reset_at: window.reset_at(),
|
||||
limit_reached: window.limit_reached(),
|
||||
local_usage: None,
|
||||
provider_data: Some(ProviderDocument::new(OpaqueProviderData::new(data))),
|
||||
}
|
||||
.map(|window| ProviderQuotaWindow {
|
||||
key: window.key().to_owned(),
|
||||
group: quota_group(window.kind()).to_owned(),
|
||||
label: codex_quota_window_label(window.kind(), window.role(), window.window_seconds()),
|
||||
limit_id: Some(window.source().to_owned()),
|
||||
limit_name: window.limit_name().map(str::to_owned),
|
||||
role: Some(quota_role(window.role())),
|
||||
local_usage_attribution: if window.is_account_wide() {
|
||||
QuotaLocalUsageAttribution::AccountWide
|
||||
} else {
|
||||
QuotaLocalUsageAttribution::Unavailable
|
||||
},
|
||||
window_seconds: window.window_seconds(),
|
||||
used_percent: window.used_percent(),
|
||||
reset_at: window.reset_at(),
|
||||
limit_reached: window.limit_reached(),
|
||||
local_usage: None,
|
||||
provider_data: None,
|
||||
})
|
||||
.collect();
|
||||
// 快照级 limit_reached 只看滚动后的窗口触顶:顶层标记是观测事实,不能
|
||||
@@ -798,23 +791,21 @@ const fn quota_group(kind: CodexQuotaWindowKind) -> &'static str {
|
||||
}
|
||||
}
|
||||
|
||||
const fn quota_role(role: CodexQuotaWindowRole) -> &'static str {
|
||||
const fn quota_role(role: CodexQuotaWindowRole) -> ProviderQuotaWindowRole {
|
||||
match role {
|
||||
CodexQuotaWindowRole::Primary => "primary",
|
||||
CodexQuotaWindowRole::Secondary => "secondary",
|
||||
CodexQuotaWindowRole::Monthly => "monthly",
|
||||
CodexQuotaWindowRole::Primary => ProviderQuotaWindowRole::Primary,
|
||||
CodexQuotaWindowRole::Secondary => ProviderQuotaWindowRole::Secondary,
|
||||
CodexQuotaWindowRole::Monthly => ProviderQuotaWindowRole::Monthly,
|
||||
}
|
||||
}
|
||||
|
||||
/// 按窗口时长显示汉化额度名;非核心桶补上上游桶名称作为前缀。
|
||||
/// 按窗口时长显示汉化额度名;额度桶名称通过独立字段投影。
|
||||
fn codex_quota_window_label(
|
||||
limit_name: Option<&str>,
|
||||
kind: CodexQuotaWindowKind,
|
||||
role: CodexQuotaWindowRole,
|
||||
source: &str,
|
||||
window_seconds: Option<u64>,
|
||||
) -> String {
|
||||
let base = match kind {
|
||||
match kind {
|
||||
CodexQuotaWindowKind::Monthly => "月额度".to_owned(),
|
||||
CodexQuotaWindowKind::Weekly => "周额度".to_owned(),
|
||||
CodexQuotaWindowKind::ShortTerm => {
|
||||
@@ -825,29 +816,7 @@ fn codex_quota_window_label(
|
||||
}
|
||||
}
|
||||
CodexQuotaWindowKind::Other => custom_quota_window_label(window_seconds, role),
|
||||
};
|
||||
if let Some(prefix) = codex_quota_window_prefix(limit_name, source) {
|
||||
return format!("{prefix} · {base}");
|
||||
}
|
||||
base
|
||||
}
|
||||
|
||||
fn codex_quota_window_prefix(limit_name: Option<&str>, source: &str) -> Option<String> {
|
||||
let source_identifier = quota_identifier(source);
|
||||
if matches!(source_identifier.as_str(), "core" | "codex") {
|
||||
return None;
|
||||
}
|
||||
let label = limit_name.unwrap_or(source);
|
||||
match quota_identifier(label).as_str() {
|
||||
"premium" => Some("Premium".to_owned()),
|
||||
"code_review" => Some("代码审查".to_owned()),
|
||||
"spend_control" => Some("消费控制".to_owned()),
|
||||
_ => Some(label.replace('_', " ")),
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_identifier(value: &str) -> String {
|
||||
value.trim().to_ascii_lowercase().replace(['-', ' '], "_")
|
||||
}
|
||||
|
||||
fn custom_quota_window_label(window_seconds: Option<u64>, role: CodexQuotaWindowRole) -> String {
|
||||
@@ -1011,6 +980,7 @@ struct PendingDocument {
|
||||
state: String,
|
||||
nonce: String,
|
||||
code_verifier: String,
|
||||
installation_id: String,
|
||||
reauthorization_account_id: Option<String>,
|
||||
mutation: PendingMutationDocument,
|
||||
}
|
||||
@@ -1071,6 +1041,10 @@ fn encode_pending(pending: &CodexPendingAuthorization) -> Map<String, Value> {
|
||||
"code_verifier".to_owned(),
|
||||
Value::String(pending.code_verifier().expose_secret().to_owned()),
|
||||
);
|
||||
document.insert(
|
||||
"installation_id".to_owned(),
|
||||
Value::String(pending.installation_id().to_owned()),
|
||||
);
|
||||
document.insert(
|
||||
"reauthorization_account_id".to_owned(),
|
||||
pending
|
||||
@@ -1161,6 +1135,7 @@ fn decode_pending(
|
||||
state: SecretString::from(document.state),
|
||||
nonce: SecretString::from(document.nonce),
|
||||
code_verifier: SecretString::from(document.code_verifier),
|
||||
installation_id: document.installation_id,
|
||||
reauthorization_account_id: document.reauthorization_account_id,
|
||||
mutation: decode_mutation(document.mutation)?,
|
||||
})
|
||||
@@ -1267,31 +1242,37 @@ fn map_store_error(error: gateway_core::error::StoreError) -> ProviderAdminError
|
||||
|
||||
fn map_credential_admin_error(error: CodexCredentialAdminError) -> ProviderAdminError {
|
||||
use CodexCredentialAdminError as Error;
|
||||
provider_admin_error(match error {
|
||||
let upstream_message = error.upstream_message().map(ToOwned::to_owned);
|
||||
let kind = match error {
|
||||
Error::InvalidInput
|
||||
| Error::InvalidCredential
|
||||
| Error::MissingRefreshToken
|
||||
| Error::RefreshRejected
|
||||
| Error::AccountBanned => ProviderAdminErrorKind::Invalid,
|
||||
| Error::RefreshRejected { .. }
|
||||
| Error::AccountBanned { .. } => ProviderAdminErrorKind::Invalid,
|
||||
Error::NotFound => ProviderAdminErrorKind::NotFound,
|
||||
Error::RefreshLeaseUnavailable | Error::RefreshAmbiguous => {
|
||||
Error::RefreshLeaseUnavailable | Error::RefreshAmbiguous { .. } => {
|
||||
ProviderAdminErrorKind::Conflict
|
||||
}
|
||||
Error::RefreshUnavailable => ProviderAdminErrorKind::Unavailable,
|
||||
})
|
||||
};
|
||||
let error = provider_admin_error(kind);
|
||||
match upstream_message {
|
||||
Some(message) => error.with_message(message),
|
||||
None => error,
|
||||
}
|
||||
}
|
||||
|
||||
const fn credential_admin_error_code(error: CodexCredentialAdminError) -> &'static str {
|
||||
const fn credential_admin_error_code(error: &CodexCredentialAdminError) -> &'static str {
|
||||
match error {
|
||||
CodexCredentialAdminError::InvalidInput => "invalid_input",
|
||||
CodexCredentialAdminError::InvalidCredential => "invalid_credential",
|
||||
CodexCredentialAdminError::NotFound => "not_found",
|
||||
CodexCredentialAdminError::MissingRefreshToken => "missing_refresh_token",
|
||||
CodexCredentialAdminError::RefreshLeaseUnavailable => "refresh_lease_unavailable",
|
||||
CodexCredentialAdminError::RefreshRejected => "refresh_rejected",
|
||||
CodexCredentialAdminError::AccountBanned => "account_banned",
|
||||
CodexCredentialAdminError::RefreshRejected { .. } => "refresh_rejected",
|
||||
CodexCredentialAdminError::AccountBanned { .. } => "account_banned",
|
||||
CodexCredentialAdminError::RefreshUnavailable => "refresh_unavailable",
|
||||
CodexCredentialAdminError::RefreshAmbiguous => "refresh_ambiguous",
|
||||
CodexCredentialAdminError::RefreshAmbiguous { .. } => "refresh_ambiguous",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1360,9 +1341,7 @@ fn map_quota_error(error: CodexCredentialQuotaError) -> ProviderAdminError {
|
||||
fn map_catalog_error(error: CodexCredentialCatalogError) -> ProviderAdminError {
|
||||
use CodexCredentialCatalogError as Error;
|
||||
provider_admin_error(match error {
|
||||
Error::InvalidCredentialData | Error::ConflictingModelFacts | Error::InvalidEtag => {
|
||||
ProviderAdminErrorKind::Invalid
|
||||
}
|
||||
Error::InvalidCredentialData | Error::InvalidEtag => ProviderAdminErrorKind::Invalid,
|
||||
Error::NoEligibleCredential => ProviderAdminErrorKind::NotFound,
|
||||
Error::ConcurrentUpdate => ProviderAdminErrorKind::Conflict,
|
||||
Error::Upstream { .. } | Error::Cache => ProviderAdminErrorKind::Unavailable,
|
||||
|
||||
@@ -36,15 +36,17 @@ pub struct OpenAiConfig {
|
||||
}
|
||||
|
||||
impl OpenAiConfig {
|
||||
/// 校验 Provider-owned 字段,并定位旧版会话锚点持续使用的身份密钥。
|
||||
pub fn resolve_and_validate(&mut self, source_dir: &Path) -> Result<(), OpenAiConfigError> {
|
||||
/// 校验 Provider-owned 字段,并从统一运行数据目录定位会话身份密钥。
|
||||
pub fn resolve_and_validate(
|
||||
&mut self,
|
||||
runtime_data_dir: &Path,
|
||||
) -> Result<(), OpenAiConfigError> {
|
||||
self.api.validate()?;
|
||||
self.ws_pool.validate()?;
|
||||
self.quota.validate()?;
|
||||
self.auth.validate()?;
|
||||
self.wire_profile.validate()?;
|
||||
let project_root = source_dir.parent().unwrap_or(source_dir);
|
||||
self.identity_secret_path = project_root.join(".runtime/data/identity_hmac_secret");
|
||||
self.identity_secret_path = runtime_data_dir.join("identity_hmac_secret");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -101,7 +103,7 @@ impl Default for OpenAiConfig {
|
||||
quota: CodexQuotaSettings::default(),
|
||||
auth: CodexAuthSettings::default(),
|
||||
wire_profile: CodexWireProfileConfig::default(),
|
||||
identity_secret_path: PathBuf::from(".runtime/data/identity_hmac_secret"),
|
||||
identity_secret_path: PathBuf::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -269,7 +271,7 @@ impl CodexAuthSettings {
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct CodexWireProfileConfig {
|
||||
pub originator: String,
|
||||
/// npm `@openai/codex` 启动版本;运行时发布任务会更新此画像字段。
|
||||
/// 官方 Desktop ZIP 内嵌 Core 的启动基线;运行时按完整制品元组更新。
|
||||
pub codex_version: String,
|
||||
pub desktop_version: String,
|
||||
pub desktop_build: String,
|
||||
@@ -284,9 +286,9 @@ impl Default for CodexWireProfileConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
originator: "Codex Desktop".to_owned(),
|
||||
codex_version: "0.146.0".to_owned(),
|
||||
desktop_version: "26.727.51351".to_owned(),
|
||||
desktop_build: "6119".to_owned(),
|
||||
codex_version: "0.147.0-alpha.6.6".to_owned(),
|
||||
desktop_version: "26.803.81509".to_owned(),
|
||||
desktop_build: "6415".to_owned(),
|
||||
os_type: "Mac OS".to_owned(),
|
||||
os_version: "15.7.1".to_owned(),
|
||||
arch: "arm64".to_owned(),
|
||||
|
||||
@@ -55,6 +55,7 @@ pub struct ImportCodexOAuthCredential {
|
||||
pub(crate) struct UnresolvedCodexOAuthCredential {
|
||||
pub(crate) account_id: String,
|
||||
pub(crate) name: String,
|
||||
pub(crate) installation_id: String,
|
||||
pub(crate) secret: CodexOAuthSecret,
|
||||
pub(crate) metadata: CodexOAuthMetadata,
|
||||
pub(crate) access_token_expires_at: Option<DateTime<Utc>>,
|
||||
@@ -68,6 +69,7 @@ impl fmt::Debug for UnresolvedCodexOAuthCredential {
|
||||
.debug_struct("UnresolvedCodexOAuthCredential")
|
||||
.field("account_id", &self.account_id)
|
||||
.field("name", &self.name)
|
||||
.field("installation_id", &"<pseudonymous>")
|
||||
.field("secret", &"<redacted>")
|
||||
.field("metadata", &self.metadata)
|
||||
.field("access_token_expires_at", &self.access_token_expires_at)
|
||||
@@ -374,7 +376,7 @@ impl fmt::Debug for PreparedCodexCredentialRotationGuard {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Error)]
|
||||
pub enum CodexCredentialAdminError {
|
||||
#[error("Codex account input is invalid")]
|
||||
InvalidInput,
|
||||
@@ -387,13 +389,30 @@ pub enum CodexCredentialAdminError {
|
||||
#[error("Codex refresh lease is unavailable")]
|
||||
RefreshLeaseUnavailable,
|
||||
#[error("Codex refresh token was rejected")]
|
||||
RefreshRejected,
|
||||
RefreshRejected { message: Option<String> },
|
||||
#[error("Codex account is banned")]
|
||||
AccountBanned,
|
||||
AccountBanned { message: Option<String> },
|
||||
#[error("Codex refresh service is unavailable")]
|
||||
RefreshUnavailable,
|
||||
#[error("Codex refresh send state is ambiguous")]
|
||||
RefreshAmbiguous,
|
||||
RefreshAmbiguous { message: Option<String> },
|
||||
}
|
||||
|
||||
impl CodexCredentialAdminError {
|
||||
#[must_use]
|
||||
pub fn upstream_message(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::RefreshRejected { message }
|
||||
| Self::AccountBanned { message }
|
||||
| Self::RefreshAmbiguous { message } => message.as_deref(),
|
||||
Self::InvalidInput
|
||||
| Self::InvalidCredential
|
||||
| Self::NotFound
|
||||
| Self::MissingRefreshToken
|
||||
| Self::RefreshLeaseUnavailable
|
||||
| Self::RefreshUnavailable => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 无状态的 Codex Admin command preparer。
|
||||
@@ -473,8 +492,12 @@ impl CodexCredentialAdmin {
|
||||
chatgpt_account_id: upstream_account_id,
|
||||
} = input.metadata;
|
||||
let has_upstream_user_id = upstream_user_id.is_some();
|
||||
let credential = CodexCredentialCodec::encode_unresolved(&input.secret, Vec::new())
|
||||
.map_err(|_| CodexCredentialAdminError::InvalidCredential)?;
|
||||
let credential = CodexCredentialCodec::encode_unresolved(
|
||||
&input.secret,
|
||||
input.installation_id,
|
||||
Vec::new(),
|
||||
)
|
||||
.map_err(|_| CodexCredentialAdminError::InvalidCredential)?;
|
||||
let account = ProviderAccount::new(
|
||||
account_id,
|
||||
provider,
|
||||
@@ -794,8 +817,15 @@ impl CodexCredentialAdminService {
|
||||
.refresher
|
||||
.refresh(refresh_token.expose_secret())
|
||||
.await
|
||||
.inspect_err(|error| log_manual_refresh_failure(&account_id, error))
|
||||
.map_err(map_refresh_failure)?;
|
||||
let access_token_rotated = tokens.access_token.is_some();
|
||||
let access_token_expires_at = match tokens.access_token.as_deref() {
|
||||
Some(access_token) => parse_access_token_expiration(access_token),
|
||||
None => current
|
||||
.account
|
||||
.access_token_expires_at()
|
||||
.map(DateTime::<Utc>::from),
|
||||
};
|
||||
let secret = CodexOAuthSecret {
|
||||
access_token: tokens
|
||||
.access_token
|
||||
@@ -815,17 +845,6 @@ impl CodexCredentialAdminService {
|
||||
Some(account_id.as_str()),
|
||||
&secret,
|
||||
);
|
||||
let access_token_expires_at = if access_token_rotated {
|
||||
tokens
|
||||
.expires_in
|
||||
.and_then(|expires_in| SystemTime::now().checked_add(expires_in))
|
||||
.map(DateTime::<Utc>::from)
|
||||
} else {
|
||||
current
|
||||
.account
|
||||
.access_token_expires_at()
|
||||
.map(DateTime::<Utc>::from)
|
||||
};
|
||||
// 正常预刷新由 worker 读取当时的 runtime policy 动态判断;这里仅清除
|
||||
// 之前瞬态失败留下的 retry-not-before。
|
||||
let next_refresh_at = None;
|
||||
@@ -917,6 +936,7 @@ impl CodexCredentialAdminService {
|
||||
.or_else(|| metadata.email.clone())
|
||||
.filter(|name| !name.trim().is_empty())
|
||||
.unwrap_or_else(|| "Codex OAuth".to_owned()),
|
||||
installation_id: uuid::Uuid::new_v4().to_string(),
|
||||
secret,
|
||||
metadata,
|
||||
access_token_expires_at,
|
||||
@@ -991,6 +1011,7 @@ impl CodexCredentialAdminService {
|
||||
let access_token = tokens
|
||||
.access_token
|
||||
.ok_or(CodexCredentialAdminError::InvalidCredential)?;
|
||||
let access_token_expires_at = parse_access_token_expiration(&access_token);
|
||||
let secret = CodexOAuthSecret {
|
||||
access_token: SecretString::from(access_token),
|
||||
// RT 未轮换时仍保留导入提供的 RT,保证本次补全不会丢失后续刷新能力。
|
||||
@@ -1005,10 +1026,6 @@ impl CodexCredentialAdminService {
|
||||
Some(account_id.as_str()),
|
||||
&secret,
|
||||
);
|
||||
let access_token_expires_at = tokens
|
||||
.expires_in
|
||||
.and_then(|expires_in| SystemTime::now().checked_add(expires_in))
|
||||
.map(DateTime::<Utc>::from);
|
||||
Ok((secret, access_token_expires_at))
|
||||
}
|
||||
|
||||
@@ -1126,15 +1143,35 @@ fn china_rfc3339(value: DateTime<Utc>) -> String {
|
||||
.to_rfc3339()
|
||||
}
|
||||
|
||||
const fn map_refresh_failure(error: RefreshFailure) -> CodexCredentialAdminError {
|
||||
fn map_refresh_failure(error: RefreshFailure) -> CodexCredentialAdminError {
|
||||
match error {
|
||||
RefreshFailure::InvalidGrant => CodexCredentialAdminError::RefreshRejected,
|
||||
RefreshFailure::Banned => CodexCredentialAdminError::AccountBanned,
|
||||
RefreshFailure::RetryableTransport => CodexCredentialAdminError::RefreshUnavailable,
|
||||
RefreshFailure::Transport => CodexCredentialAdminError::RefreshAmbiguous,
|
||||
RefreshFailure::InvalidGrant { message, .. } => {
|
||||
CodexCredentialAdminError::RefreshRejected { message }
|
||||
}
|
||||
RefreshFailure::Banned { message, .. } => {
|
||||
CodexCredentialAdminError::AccountBanned { message }
|
||||
}
|
||||
RefreshFailure::RetryableTransport { .. } => CodexCredentialAdminError::RefreshUnavailable,
|
||||
RefreshFailure::Transport { message, .. } => {
|
||||
CodexCredentialAdminError::RefreshAmbiguous { message }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn log_manual_refresh_failure(account_id: &ProviderAccountId, error: &RefreshFailure) {
|
||||
let upstream = error.upstream();
|
||||
tracing::warn!(
|
||||
account_id = %account_id,
|
||||
failure_class = error.classification(),
|
||||
upstream_message = ?error.message(),
|
||||
upstream_status = ?upstream.map(super::token_client::RefreshUpstreamFailure::status),
|
||||
upstream_code = ?upstream.and_then(super::token_client::RefreshUpstreamFailure::code),
|
||||
upstream_type = ?upstream.and_then(super::token_client::RefreshUpstreamFailure::error_type),
|
||||
upstream_body = ?upstream.map(super::token_client::RefreshUpstreamFailure::body),
|
||||
"OpenAI OAuth manual refresh failed"
|
||||
);
|
||||
}
|
||||
|
||||
fn parse_import_document(
|
||||
payload: &Value,
|
||||
) -> Result<Vec<ParsedCodexImportAccount>, CodexCredentialAdminError> {
|
||||
|
||||
@@ -22,6 +22,8 @@ use url::Url;
|
||||
|
||||
use super::repository::{CodexCredentialRepository, CredentialRepositoryError};
|
||||
use super::security::{CodexRuntimeAuthentication, CodexRuntimeCredential};
|
||||
use crate::transport::headers::build_codex_profile_headers;
|
||||
use crate::transport::profile::CodexWireProfileState;
|
||||
use crate::transport::websocket::CodexWebSocketExchangeError;
|
||||
use crate::transport::{CodexClientError, CodexWebSocketPool};
|
||||
|
||||
@@ -167,13 +169,21 @@ pub trait CodexAgentIdentityTaskRegistrar: Send + Sync + 'static {
|
||||
pub struct OfficialCodexAgentIdentityTaskRegistrar {
|
||||
client: reqwest::Client,
|
||||
base_url: Url,
|
||||
profile: CodexWireProfileState,
|
||||
}
|
||||
|
||||
impl OfficialCodexAgentIdentityTaskRegistrar {
|
||||
pub fn new(client: reqwest::Client) -> Result<Self, CodexAgentIdentityError> {
|
||||
pub fn new(
|
||||
client: reqwest::Client,
|
||||
profile: CodexWireProfileState,
|
||||
) -> Result<Self, CodexAgentIdentityError> {
|
||||
let base_url = Url::parse(OFFICIAL_AGENT_IDENTITY_AUTH_BASE_URL)
|
||||
.map_err(|_| CodexAgentIdentityError::InvalidConfiguration)?;
|
||||
Ok(Self { client, base_url })
|
||||
Ok(Self {
|
||||
client,
|
||||
base_url,
|
||||
profile,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -197,9 +207,13 @@ impl CodexAgentIdentityTaskRegistrar for OfficialCodexAgentIdentityTaskRegistrar
|
||||
.path_segments_mut()
|
||||
.map_err(|_| CodexAgentIdentityError::InvalidConfiguration)?
|
||||
.extend(["v1", "agent", credential.runtime_id(), "task", "register"]);
|
||||
let profile = self.profile.snapshot();
|
||||
let headers = build_codex_profile_headers(&profile)
|
||||
.map_err(|_| CodexAgentIdentityError::InvalidConfiguration)?;
|
||||
let response = self
|
||||
.client
|
||||
.post(endpoint)
|
||||
.headers(headers)
|
||||
.header(reqwest::header::ACCEPT, "application/json")
|
||||
.json(&credential.registration(Utc::now()))
|
||||
.send()
|
||||
|
||||
@@ -138,8 +138,6 @@ pub enum CodexCredentialCatalogError {
|
||||
InvalidCredentialData,
|
||||
#[error("Codex model catalog upstream query failed: {detail}")]
|
||||
Upstream { detail: String },
|
||||
#[error("Codex model catalog contains conflicting account facts")]
|
||||
ConflictingModelFacts,
|
||||
#[error("Codex model catalog cache is unavailable")]
|
||||
Cache,
|
||||
#[error("Codex model catalog ETag is invalid")]
|
||||
@@ -359,11 +357,6 @@ impl CodexCredentialCatalogService {
|
||||
union_order.push(entry.key().clone());
|
||||
entry.insert(model.clone());
|
||||
}
|
||||
// 展示与元数据允许跨套餐漂移,union 取首见值;只有路由
|
||||
// 事实冲突才让整轮刷新失败。
|
||||
Entry::Occupied(entry) if conflicting_routing_facts(entry.get(), model) => {
|
||||
return Err(CodexCredentialCatalogError::ConflictingModelFacts);
|
||||
}
|
||||
Entry::Occupied(_) => {}
|
||||
}
|
||||
}
|
||||
@@ -661,13 +654,6 @@ fn model_ids(models: &[CodexCatalogModel]) -> Vec<String> {
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 路由正确性只依赖 Responses API 支持证据(决定 Generate operation 是否
|
||||
/// 编译进 `ModelCapabilities`);display_name/limits/metadata 等展示事实
|
||||
/// 跨套餐漂移不构成冲突。
|
||||
fn conflicting_routing_facts(existing: &CodexCatalogModel, candidate: &CodexCatalogModel) -> bool {
|
||||
existing.capabilities().responses_api() != candidate.capabilities().responses_api()
|
||||
}
|
||||
|
||||
fn catalog_candidates_by_scope(
|
||||
accounts: Vec<ProviderAccount>,
|
||||
) -> Result<BTreeMap<CodexCatalogScope, Vec<ProviderAccount>>, CodexCredentialCatalogError> {
|
||||
|
||||
@@ -19,6 +19,7 @@ pub(crate) use affinity::{
|
||||
derive_codex_cyber_policy_session_key, derive_codex_session_affinity_key,
|
||||
};
|
||||
pub(crate) use oauth::oauth_owner_ref;
|
||||
pub(crate) use types::parse_access_token_expiration;
|
||||
|
||||
pub use agent_identity::{
|
||||
CodexAgentIdentityError, CodexAgentIdentitySecret, CodexAgentIdentityTaskRegistrar,
|
||||
|
||||
@@ -1,10 +1,6 @@
|
||||
//! Codex Authorization Code + PKCE/OIDC 管理流。
|
||||
|
||||
use std::{
|
||||
fmt,
|
||||
sync::Arc,
|
||||
time::{Duration, SystemTime},
|
||||
};
|
||||
use std::{fmt, sync::Arc, time::Duration};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
@@ -35,10 +31,17 @@ use super::token_client::{
|
||||
AuthorizationCodeExchangeError, AuthorizationCodeExchanger, AuthorizationCodeGrant,
|
||||
OFFICIAL_CODEX_OAUTH_CLIENT_ID, OFFICIAL_CODEX_REDIRECT_URI,
|
||||
};
|
||||
use super::types::parse_chatgpt_jwt_claims;
|
||||
use super::types::{parse_access_token_expiration, parse_chatgpt_jwt_claims};
|
||||
use crate::transport::profile::CodexWireProfileState;
|
||||
|
||||
const AUTHORIZATION_ENDPOINT: &str = "https://auth.openai.com/oauth/authorize";
|
||||
const AUTHORIZATION_SCOPE: &str = "openid profile email offline_access";
|
||||
const DESKTOP_AUTH_ENDPOINT: &str = "https://chatgpt.com/codex/desktop-auth";
|
||||
const AUTHORIZATION_SCOPE: &str =
|
||||
"openid profile email offline_access api.connectors.read api.connectors.invoke";
|
||||
const AUTHORIZATION_ORIGINATOR: &str = "Codex Desktop";
|
||||
const STREAMLINED_LOGIN_QUERY: &str = "codex_streamlined_login";
|
||||
const NO_UNIVERSAL_LINKS_QUERY: &str = "no_universal_links";
|
||||
const SURFACE_STABLE_ID_DOMAIN: &[u8] = b"openai/codex-desktop/surface-stable-id/v1\0";
|
||||
const AUTHORIZATION_TTL: TimeDelta = TimeDelta::minutes(10);
|
||||
const AUTHORIZATION_CLAIM_TTL: Duration = Duration::from_secs(90);
|
||||
const MAX_CALLBACK_BYTES: usize = 16 * 1024;
|
||||
@@ -137,6 +140,7 @@ pub struct CodexPendingAuthorization {
|
||||
state: SecretString,
|
||||
nonce: SecretString,
|
||||
code_verifier: SecretString,
|
||||
installation_id: String,
|
||||
reauthorization: Option<ProviderAccountId>,
|
||||
mutation: PendingAuthorizationMutation,
|
||||
}
|
||||
@@ -150,6 +154,7 @@ pub struct StoredCodexPendingAuthorization {
|
||||
pub state: SecretString,
|
||||
pub nonce: SecretString,
|
||||
pub code_verifier: SecretString,
|
||||
pub installation_id: String,
|
||||
pub reauthorization_account_id: Option<String>,
|
||||
pub mutation: PendingAuthorizationMutation,
|
||||
}
|
||||
@@ -172,6 +177,7 @@ impl CodexPendingAuthorization {
|
||||
state: input.state,
|
||||
nonce: input.nonce,
|
||||
code_verifier: input.code_verifier,
|
||||
installation_id: input.installation_id,
|
||||
reauthorization,
|
||||
mutation: input.mutation,
|
||||
};
|
||||
@@ -183,6 +189,7 @@ impl CodexPendingAuthorization {
|
||||
|| !valid_secret(pending.state.expose_secret())
|
||||
|| !valid_secret(pending.nonce.expose_secret())
|
||||
|| !valid_secret(pending.code_verifier.expose_secret())
|
||||
|| !valid_installation_id(&pending.installation_id)
|
||||
|| !pending_mutation_matches(&pending)
|
||||
{
|
||||
return Err(CodexOAuthPendingStoreError::InvalidValue);
|
||||
@@ -230,6 +237,11 @@ impl CodexPendingAuthorization {
|
||||
&self.code_verifier
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn installation_id(&self) -> &str {
|
||||
&self.installation_id
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn reauthorization(&self) -> Option<&ProviderAccountId> {
|
||||
self.reauthorization.as_ref()
|
||||
@@ -253,6 +265,7 @@ impl fmt::Debug for CodexPendingAuthorization {
|
||||
.field("state", &"<redacted>")
|
||||
.field("nonce", &"<redacted>")
|
||||
.field("code_verifier", &"<redacted>")
|
||||
.field("installation_id", &"<pseudonymous>")
|
||||
.field("reauthorization", &self.reauthorization)
|
||||
.field("mutation", &self.mutation)
|
||||
.finish()
|
||||
@@ -348,6 +361,7 @@ pub struct CodexOAuthAdminService {
|
||||
store: Arc<dyn ProviderAccountStore>,
|
||||
credentials: CodexCredentialAdmin,
|
||||
oauth_client_id: String,
|
||||
profile: CodexWireProfileState,
|
||||
}
|
||||
|
||||
struct CodexOAuthAuthorizationCommitGuard {
|
||||
@@ -397,6 +411,7 @@ impl CodexOAuthAdminService {
|
||||
exchanger: Arc<dyn AuthorizationCodeExchanger>,
|
||||
store: Arc<dyn ProviderAccountStore>,
|
||||
credentials: CodexCredentialAdmin,
|
||||
profile: CodexWireProfileState,
|
||||
) -> Self {
|
||||
Self {
|
||||
pending,
|
||||
@@ -404,6 +419,7 @@ impl CodexOAuthAdminService {
|
||||
store,
|
||||
credentials,
|
||||
oauth_client_id: OFFICIAL_CODEX_OAUTH_CLIENT_ID.to_owned(),
|
||||
profile,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -420,8 +436,10 @@ impl CodexOAuthAdmin for CodexOAuthAdminService {
|
||||
&self,
|
||||
command: StartCodexOAuthAuthorization,
|
||||
) -> Result<CodexOAuthAuthorizationStarted, CodexOAuthAdminError> {
|
||||
let (name, reauthorization) = match command.mutation.target() {
|
||||
AuthorizationMutationTarget::Create { name } => (name.clone(), None),
|
||||
let (name, reauthorization, installation_id) = match command.mutation.target() {
|
||||
AuthorizationMutationTarget::Create { name } => {
|
||||
(name.clone(), None, Uuid::new_v4().to_string())
|
||||
}
|
||||
AuthorizationMutationTarget::Reauthorize { account_id } => {
|
||||
let current = self
|
||||
.store
|
||||
@@ -433,10 +451,15 @@ impl CodexOAuthAdmin for CodexOAuthAdminService {
|
||||
{
|
||||
return Err(CodexOAuthAdminError::NotFound);
|
||||
}
|
||||
(current.account.name().to_owned(), Some(account_id.clone()))
|
||||
let installation_id = current_installation_id(¤t)?;
|
||||
(
|
||||
current.account.name().to_owned(),
|
||||
Some(account_id.clone()),
|
||||
installation_id,
|
||||
)
|
||||
}
|
||||
};
|
||||
self.start_pending(command.mutation, name, reauthorization)
|
||||
self.start_pending(command.mutation, name, reauthorization, installation_id)
|
||||
.await
|
||||
}
|
||||
|
||||
@@ -488,6 +511,9 @@ impl CodexOAuthAdminService {
|
||||
.load_current_credential(target)
|
||||
.await
|
||||
.map_err(map_store_error)?;
|
||||
if current_installation_id(¤t)? != pending.installation_id() {
|
||||
return Err(CodexOAuthAdminError::Conflict);
|
||||
}
|
||||
Some(current)
|
||||
} else {
|
||||
None
|
||||
@@ -497,13 +523,12 @@ impl CodexOAuthAdminService {
|
||||
.map(current_oauth_refresh_token)
|
||||
.transpose()?
|
||||
.flatten();
|
||||
let (mut secret, id_token, expires_in) = self
|
||||
let (mut secret, id_token) = self
|
||||
.exchange_pending(&pending, callback_url, fallback_refresh_token)
|
||||
.await?;
|
||||
let mutation = pending.mutation.clone();
|
||||
let access_token_expires_at = expires_in
|
||||
.and_then(|expires_in| SystemTime::now().checked_add(expires_in))
|
||||
.map(DateTime::<Utc>::from);
|
||||
let access_token_expires_at =
|
||||
parse_access_token_expiration(secret.access_token.expose_secret());
|
||||
let metadata = parse_chatgpt_jwt_claims(id_token.expose_secret())
|
||||
.map_err(|_| CodexOAuthAdminError::TokenRejected)?;
|
||||
secret.id_token = Some(id_token);
|
||||
@@ -525,6 +550,7 @@ impl CodexOAuthAdminService {
|
||||
.prepare_unresolved_oauth(UnresolvedCodexOAuthCredential {
|
||||
account_id,
|
||||
name: pending.name,
|
||||
installation_id: pending.installation_id,
|
||||
secret,
|
||||
metadata,
|
||||
access_token_expires_at,
|
||||
@@ -544,6 +570,7 @@ impl CodexOAuthAdminService {
|
||||
mutation: PendingAuthorizationMutation,
|
||||
name: String,
|
||||
reauthorization: Option<ProviderAccountId>,
|
||||
installation_id: String,
|
||||
) -> Result<CodexOAuthAuthorizationStarted, CodexOAuthAdminError> {
|
||||
if !valid_text(&name) {
|
||||
return Err(CodexOAuthAdminError::InvalidInput);
|
||||
@@ -562,17 +589,21 @@ impl CodexOAuthAdminService {
|
||||
state: SecretString::from(random_secret()?),
|
||||
nonce: SecretString::from(random_secret()?),
|
||||
code_verifier: SecretString::from(random_secret()?),
|
||||
installation_id,
|
||||
reauthorization_account_id: reauthorization.as_ref().map(ToString::to_string),
|
||||
mutation,
|
||||
})
|
||||
.map_err(map_pending_error)?;
|
||||
let profile = self.profile.snapshot();
|
||||
let authorization_url =
|
||||
authorization_url(&pending, &self.oauth_client_id, &profile.desktop_version)?;
|
||||
self.pending
|
||||
.create(&pending)
|
||||
.await
|
||||
.map_err(map_pending_error)?;
|
||||
Ok(CodexOAuthAuthorizationStarted {
|
||||
flow_id: pending.flow_id.clone(),
|
||||
authorization_url: authorization_url(&pending, &self.oauth_client_id)?,
|
||||
authorization_url,
|
||||
expires_at,
|
||||
})
|
||||
}
|
||||
@@ -635,14 +666,7 @@ impl CodexOAuthAdminService {
|
||||
pending: &CodexPendingAuthorization,
|
||||
callback_url: &str,
|
||||
fallback_refresh_token: Option<SecretString>,
|
||||
) -> Result<
|
||||
(
|
||||
super::types::CodexOAuthSecret,
|
||||
SecretString,
|
||||
Option<Duration>,
|
||||
),
|
||||
CodexOAuthAdminError,
|
||||
> {
|
||||
) -> Result<(super::types::CodexOAuthSecret, SecretString), CodexOAuthAdminError> {
|
||||
let (code, callback_state) = callback_parts(callback_url)?;
|
||||
if !constant_time_equal(
|
||||
pending.state.expose_secret().as_bytes(),
|
||||
@@ -673,7 +697,7 @@ impl CodexOAuthAdminService {
|
||||
.as_ref()
|
||||
.map(ExposeSecret::expose_secret),
|
||||
);
|
||||
Ok((secret, tokens.id_token, tokens.expires_in))
|
||||
Ok((secret, tokens.id_token))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -688,27 +712,70 @@ fn current_oauth_refresh_token(
|
||||
.and_then(|oauth| oauth.refresh_token.clone()))
|
||||
}
|
||||
|
||||
fn current_installation_id(current: &LoadedCredential) -> Result<String, CodexOAuthAdminError> {
|
||||
let runtime = CodexCredentialCodec::decode(¤t.credential)
|
||||
.map_err(|_| CodexOAuthAdminError::Credential)?;
|
||||
if runtime.authentication.oauth().is_none() {
|
||||
return Err(CodexOAuthAdminError::Credential);
|
||||
}
|
||||
Ok(runtime.installation_id)
|
||||
}
|
||||
|
||||
fn derive_surface_stable_id(installation_id: &str) -> Result<String, CodexOAuthAdminError> {
|
||||
let installation_id =
|
||||
Uuid::parse_str(installation_id).map_err(|_| CodexOAuthAdminError::Credential)?;
|
||||
let mut digest = Sha256::new();
|
||||
digest.update(SURFACE_STABLE_ID_DOMAIN);
|
||||
digest.update(installation_id.as_bytes());
|
||||
let digest = digest.finalize();
|
||||
let mut bytes = [0_u8; 16];
|
||||
bytes.copy_from_slice(&digest[..16]);
|
||||
bytes[6] = (bytes[6] & 0x0f) | 0x40;
|
||||
bytes[8] = (bytes[8] & 0x3f) | 0x80;
|
||||
Ok(Uuid::from_bytes(bytes).hyphenated().to_string())
|
||||
}
|
||||
|
||||
fn authorization_url(
|
||||
pending: &CodexPendingAuthorization,
|
||||
oauth_client_id: &str,
|
||||
desktop_version: &str,
|
||||
) -> Result<String, CodexOAuthAdminError> {
|
||||
let mut url =
|
||||
if !valid_text(desktop_version) {
|
||||
return Err(CodexOAuthAdminError::InvalidInput);
|
||||
}
|
||||
let mut inner =
|
||||
Url::parse(AUTHORIZATION_ENDPOINT).map_err(|_| CodexOAuthAdminError::InvalidInput)?;
|
||||
let challenge = URL_SAFE_NO_PAD.encode(Sha256::digest(
|
||||
pending.code_verifier.expose_secret().as_bytes(),
|
||||
));
|
||||
url.query_pairs_mut()
|
||||
.append_pair("response_type", "code")
|
||||
.append_pair("client_id", oauth_client_id)
|
||||
.append_pair("redirect_uri", OFFICIAL_CODEX_REDIRECT_URI)
|
||||
.append_pair("scope", AUTHORIZATION_SCOPE)
|
||||
.append_pair("state", pending.state.expose_secret())
|
||||
.append_pair("nonce", pending.nonce.expose_secret())
|
||||
.append_pair("code_challenge", &challenge)
|
||||
.append_pair("code_challenge_method", "S256")
|
||||
.append_pair("id_token_add_organizations", "true")
|
||||
.append_pair("codex_cli_simplified_flow", "true");
|
||||
Ok(url.into())
|
||||
let stable_id = derive_surface_stable_id(pending.installation_id())?;
|
||||
{
|
||||
let mut query = inner.query_pairs_mut();
|
||||
query
|
||||
.append_pair("response_type", "code")
|
||||
.append_pair("client_id", oauth_client_id)
|
||||
.append_pair("redirect_uri", OFFICIAL_CODEX_REDIRECT_URI)
|
||||
.append_pair("scope", AUTHORIZATION_SCOPE)
|
||||
.append_pair("code_challenge", challenge.as_str())
|
||||
.append_pair("code_challenge_method", "S256")
|
||||
.append_pair("id_token_add_organizations", "true")
|
||||
.append_pair("codex_cli_simplified_flow", "true")
|
||||
.append_pair("state", pending.state.expose_secret())
|
||||
.append_pair("originator", AUTHORIZATION_ORIGINATOR)
|
||||
.append_pair("codex_app_version", desktop_version)
|
||||
.append_pair("source_surface_stable_id", &stable_id)
|
||||
.append_pair("codex_origin_stable_id", &stable_id)
|
||||
.append_pair(STREAMLINED_LOGIN_QUERY, "true");
|
||||
}
|
||||
|
||||
let mut outer =
|
||||
Url::parse(DESKTOP_AUTH_ENDPOINT).map_err(|_| CodexOAuthAdminError::InvalidInput)?;
|
||||
outer
|
||||
.query_pairs_mut()
|
||||
.append_pair("authorize_url", inner.as_str())
|
||||
.append_pair(STREAMLINED_LOGIN_QUERY, "true")
|
||||
.append_pair(NO_UNIVERSAL_LINKS_QUERY, "1");
|
||||
Ok(outer.into())
|
||||
}
|
||||
|
||||
fn callback_parts(value: &str) -> Result<(SecretString, SecretString), CodexOAuthAdminError> {
|
||||
@@ -760,6 +827,12 @@ fn valid_secret(value: &str) -> bool {
|
||||
valid_text(value) && value.len() >= 16
|
||||
}
|
||||
|
||||
fn valid_installation_id(value: &str) -> bool {
|
||||
Uuid::parse_str(value)
|
||||
.ok()
|
||||
.is_some_and(|uuid| uuid.get_version_num() == 4)
|
||||
}
|
||||
|
||||
pub(crate) fn oauth_owner_ref(owner: &AuthorizationOwner) -> String {
|
||||
let mut digest = Sha256::new();
|
||||
match owner {
|
||||
|
||||
@@ -0,0 +1,173 @@
|
||||
//! OpenAI quota 文档的规范化 map 与上游 `/usage` 输入适配。
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use gateway_protocol::openai::events::{RateLimitDetails, RateLimitKeySource};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
pub(super) const DEFAULT_CODEX_LIMIT_ID: &str = "codex";
|
||||
pub(super) const RATE_LIMITS_BY_LIMIT_ID: &str = "rate_limits_by_limit_id";
|
||||
|
||||
/// 官方新协议的多桶视图:每个 `limit_id` 对应一个独立快照。
|
||||
#[derive(Debug, Default)]
|
||||
pub(super) struct RateLimitSnapshotsByLimitId {
|
||||
snapshots: BTreeMap<String, Map<String, Value>>,
|
||||
}
|
||||
|
||||
impl RateLimitSnapshotsByLimitId {
|
||||
/// 接收上游 `/usage` 形状或已经规范化的 map;输出始终只保留 map。
|
||||
pub(super) fn take_from_document(quota: &mut Map<String, Value>) -> Self {
|
||||
let mut result = Self::default();
|
||||
|
||||
if let Some(by_limit_id) = quota
|
||||
.remove(RATE_LIMITS_BY_LIMIT_ID)
|
||||
.and_then(|value| value.as_object().cloned())
|
||||
{
|
||||
for (raw_limit_id, value) in by_limit_id {
|
||||
let Some(limit_id) = normalize_limit_id(&raw_limit_id) else {
|
||||
continue;
|
||||
};
|
||||
let mut snapshot = snapshot_object(value);
|
||||
snapshot.insert("limit_id".to_owned(), Value::String(limit_id.clone()));
|
||||
result.snapshots.insert(limit_id, snapshot);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(rate_limit) = quota.remove("rate_limit") {
|
||||
result.insert_usage_snapshot(DEFAULT_CODEX_LIMIT_ID, None, rate_limit);
|
||||
}
|
||||
if let Some(rate_limit) = quota.remove("code_review_rate_limit") {
|
||||
result.insert_usage_snapshot("code_review", Some("code_review"), rate_limit);
|
||||
}
|
||||
|
||||
let additional = quota
|
||||
.remove("additional_rate_limits")
|
||||
.and_then(|value| value.as_array().cloned())
|
||||
.unwrap_or_default();
|
||||
for (index, value) in additional.into_iter().enumerate() {
|
||||
let Value::Object(mut item) = value else {
|
||||
result.insert_usage_snapshot(&format!("additional_{index}"), None, value);
|
||||
continue;
|
||||
};
|
||||
let limit_id = item_limit_id(&item).unwrap_or_else(|| format!("additional_{index}"));
|
||||
let limit_name = item
|
||||
.get("limit_name")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned);
|
||||
let rate_limit = item.remove("rate_limit").unwrap_or(Value::Null);
|
||||
result.insert_usage_snapshot(&limit_id, limit_name.as_deref(), rate_limit);
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
fn insert_usage_snapshot(
|
||||
&mut self,
|
||||
limit_id: &str,
|
||||
limit_name: Option<&str>,
|
||||
rate_limit: Value,
|
||||
) {
|
||||
let Some(limit_id) = normalize_limit_id(limit_id) else {
|
||||
return;
|
||||
};
|
||||
let mut snapshot = snapshot_object(rate_limit);
|
||||
snapshot.insert("limit_id".to_owned(), Value::String(limit_id.clone()));
|
||||
if let Some(limit_name) = limit_name {
|
||||
snapshot.insert(
|
||||
"limit_name".to_owned(),
|
||||
Value::String(limit_name.to_owned()),
|
||||
);
|
||||
}
|
||||
// map 协议中一个 ID 只能有一个快照;规范化 map 与顶层账号桶优先。
|
||||
self.snapshots.entry(limit_id).or_insert(snapshot);
|
||||
}
|
||||
|
||||
pub(super) fn upsert(
|
||||
&mut self,
|
||||
limit_id: &str,
|
||||
details: &RateLimitDetails,
|
||||
rate_limit: Map<String, Value>,
|
||||
) {
|
||||
let mut snapshot = self.snapshots.remove(limit_id).unwrap_or_default();
|
||||
for field in [
|
||||
"allowed",
|
||||
"limit_reached",
|
||||
"primary_window",
|
||||
"secondary_window",
|
||||
] {
|
||||
snapshot.remove(field);
|
||||
}
|
||||
snapshot.extend(rate_limit);
|
||||
snapshot.insert("limit_id".to_owned(), Value::String(limit_id.to_owned()));
|
||||
if let Some(name) = details.limit_name.as_ref() {
|
||||
snapshot.insert("limit_name".to_owned(), Value::String(name.clone()));
|
||||
}
|
||||
self.snapshots.insert(limit_id.to_owned(), snapshot);
|
||||
}
|
||||
|
||||
/// 把 wire 观察解析到稳定 `limit_id`。
|
||||
///
|
||||
/// WebSocket 的 `additional_rate_limits` 可能只返回可读名称;这时只接受
|
||||
/// `/usage` 或 HTTP headers 已建立的唯一名称映射。无法唯一解析就不落库,
|
||||
/// 避免凭名称猜 ID 或污染默认 `codex` 桶。
|
||||
pub(super) fn resolve_limit_id(&self, details: &RateLimitDetails) -> Option<String> {
|
||||
if details.key_source == RateLimitKeySource::LimitId {
|
||||
return normalize_limit_id(&details.limit_id);
|
||||
}
|
||||
|
||||
let limit_name = details.limit_name.as_deref()?.trim();
|
||||
let mut matches = self
|
||||
.snapshots
|
||||
.iter()
|
||||
.filter(|(_, snapshot)| {
|
||||
snapshot
|
||||
.get("limit_name")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|stored| stored.trim().eq_ignore_ascii_case(limit_name))
|
||||
})
|
||||
.map(|(limit_id, _)| limit_id.clone());
|
||||
let resolved = matches.next()?;
|
||||
matches.next().is_none().then_some(resolved)
|
||||
}
|
||||
|
||||
pub(super) fn write_to_document(self, quota: &mut Map<String, Value>) {
|
||||
let by_limit_id = self
|
||||
.snapshots
|
||||
.into_iter()
|
||||
.map(|(limit_id, snapshot)| (limit_id, Value::Object(snapshot)))
|
||||
.collect::<Map<_, _>>();
|
||||
quota.insert(
|
||||
RATE_LIMITS_BY_LIMIT_ID.to_owned(),
|
||||
Value::Object(by_limit_id),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn canonicalize_rate_limit_document(
|
||||
mut quota: Map<String, Value>,
|
||||
) -> Map<String, Value> {
|
||||
let snapshots = RateLimitSnapshotsByLimitId::take_from_document(&mut quota);
|
||||
snapshots.write_to_document(&mut quota);
|
||||
quota
|
||||
}
|
||||
|
||||
fn snapshot_object(value: Value) -> Map<String, Value> {
|
||||
value.as_object().cloned().unwrap_or_else(|| {
|
||||
let mut snapshot = Map::new();
|
||||
snapshot.insert("rate_limit".to_owned(), value);
|
||||
snapshot
|
||||
})
|
||||
}
|
||||
|
||||
fn item_limit_id(item: &Map<String, Value>) -> Option<String> {
|
||||
item.get("limit_id")
|
||||
.or_else(|| item.get("metered_feature"))
|
||||
.or_else(|| item.get("limit_name"))
|
||||
.and_then(Value::as_str)
|
||||
.and_then(normalize_limit_id)
|
||||
}
|
||||
|
||||
fn normalize_limit_id(value: &str) -> Option<String> {
|
||||
let value = value.trim();
|
||||
(!value.is_empty() && value.len() <= 128 && !value.chars().any(char::is_control))
|
||||
.then(|| value.to_ascii_lowercase().replace('-', "_"))
|
||||
}
|
||||
@@ -1,8 +1,10 @@
|
||||
//! Codex quota 服务编排:权威额度事实、展示投影、主动/被动同步与 429 冷却。
|
||||
//!
|
||||
//! - [`document`]:上游 `/usage` 输入到多桶 map 的单向规范化。
|
||||
//! - [`snapshot`]:快照/窗口解析、聚合、O(1) 滚动与调度信号。
|
||||
//! - [`evidence`]:额度接口错误到凭据/额度事实的分类。
|
||||
|
||||
mod document;
|
||||
pub(crate) mod evidence;
|
||||
pub(crate) mod snapshot;
|
||||
|
||||
@@ -19,8 +21,8 @@ use chrono::Utc;
|
||||
use gateway_core::engine::credential::{
|
||||
AccountErrorReason, AccountQuotaSignals, CredentialRevision, CredentialState,
|
||||
OpaqueProviderData, ProviderAccount, ProviderAccountId, ProviderAccountStore,
|
||||
QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaObservation, QuotaState,
|
||||
QuotaWriteOutcome,
|
||||
QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaObservation, QuotaObservationTouch,
|
||||
QuotaState, QuotaWriteOutcome,
|
||||
};
|
||||
use gateway_core::provider_ports::{ProviderCooldown, ProviderCooldownPort};
|
||||
use gateway_protocol::openai::events::{
|
||||
@@ -38,6 +40,10 @@ use crate::transport::{CodexBackendClient, CodexClientError, CodexRequestContext
|
||||
|
||||
use super::agent_identity::{CodexAgentIdentityTaskService, PreparedCodexRuntimeCredential};
|
||||
use super::repository::{CodexCredentialRepository, CredentialRepositoryError};
|
||||
use document::{
|
||||
DEFAULT_CODEX_LIMIT_ID, RATE_LIMITS_BY_LIMIT_ID, RateLimitSnapshotsByLimitId,
|
||||
canonicalize_rate_limit_document,
|
||||
};
|
||||
use evidence::{QuotaEndpointFailure, classify_quota_endpoint_failure};
|
||||
use snapshot::{
|
||||
parse_account_quota_snapshot, quota_projection_ttl, quota_snapshot_from_observation,
|
||||
@@ -646,7 +652,25 @@ impl CodexCredentialQuotaService {
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if !quota_refresh_can_update(account, &snapshot, previous.as_ref()) {
|
||||
if account.quota().is_exhausted()
|
||||
&& !quota_refresh_allows_recovery(account, &snapshot, previous.as_ref())
|
||||
{
|
||||
let outcome = self
|
||||
.store
|
||||
.touch_quota_observation(QuotaObservationTouch {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
observed_at,
|
||||
})
|
||||
.await?;
|
||||
if outcome == QuotaWriteOutcome::Conflict {
|
||||
summary.stale += 1;
|
||||
return Ok(());
|
||||
}
|
||||
if let Some(previous) = previous {
|
||||
self.scheduling
|
||||
.observe(&previous.with_observed_at(observed_at));
|
||||
}
|
||||
summary.exhausted += 1;
|
||||
return Ok(());
|
||||
}
|
||||
@@ -688,10 +712,25 @@ impl CodexCredentialQuotaService {
|
||||
let Some(rate_limits) = parse_rate_limit_headers(headers) else {
|
||||
return Ok(false);
|
||||
};
|
||||
let has_quota_facts = rate_limits
|
||||
.limits
|
||||
.values()
|
||||
.any(|details| passive_rate_limit_snapshot(details).is_some());
|
||||
self.synchronize_passive_rate_limits(account, std::slice::from_ref(&rate_limits))
|
||||
.await
|
||||
}
|
||||
|
||||
/// 把一次推理响应中采集的结构化限流观察合并后单次落库。
|
||||
pub async fn synchronize_passive_rate_limits(
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
rate_limits: &[ParsedRateLimits],
|
||||
) -> Result<bool, CodexCredentialQuotaError> {
|
||||
if rate_limits.is_empty() {
|
||||
return Ok(false);
|
||||
}
|
||||
let has_quota_facts = rate_limits.iter().any(|observation| {
|
||||
observation
|
||||
.limits
|
||||
.values()
|
||||
.any(|details| passive_rate_limit_snapshot(details).is_some())
|
||||
});
|
||||
let existing = self
|
||||
.store
|
||||
.get_quotas(std::slice::from_ref(account.id()))
|
||||
@@ -716,14 +755,14 @@ impl CodexCredentialQuotaService {
|
||||
.compare_and_swap_quota(QuotaObservation {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
quota: OpaqueProviderData::new(merge_passive_quota(existing, &rate_limits)),
|
||||
quota: OpaqueProviderData::new(merge_passive_quota(existing, rate_limits)),
|
||||
observed_at: SystemTime::now(),
|
||||
state,
|
||||
})
|
||||
.await?;
|
||||
return Ok(outcome != QuotaWriteOutcome::Conflict);
|
||||
}
|
||||
let quota = merge_passive_quota(existing, &rate_limits);
|
||||
let quota = merge_passive_quota(existing, rate_limits);
|
||||
let observed_at = SystemTime::now();
|
||||
let snapshot = parse_account_quota_snapshot(
|
||||
account.id().clone(),
|
||||
@@ -930,8 +969,29 @@ impl CodexCredentialQuotaService {
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if !quota_refresh_can_update(&account, &snapshot, previous.as_ref()) {
|
||||
return Ok(previous.unwrap_or_else(|| snapshot.with_quota_state(account.quota())));
|
||||
if !quota_refresh_allows_recovery(&account, &snapshot, previous.as_ref()) {
|
||||
if self
|
||||
.store
|
||||
.touch_quota_observation(QuotaObservationTouch {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
observed_at,
|
||||
})
|
||||
.await?
|
||||
== QuotaWriteOutcome::Conflict
|
||||
{
|
||||
return Err(CodexCredentialQuotaError::RevisionConflict);
|
||||
}
|
||||
return Ok(match previous {
|
||||
Some(previous) => {
|
||||
let previous = previous.with_observed_at(observed_at);
|
||||
self.scheduling.observe(&previous);
|
||||
previous
|
||||
}
|
||||
None => snapshot
|
||||
.with_quota_state(account.quota())
|
||||
.with_observed_at(observed_at),
|
||||
});
|
||||
}
|
||||
let state = if account.quota().is_exhausted() {
|
||||
QuotaState::allowed(observed_at)
|
||||
@@ -1027,7 +1087,7 @@ impl CodexCredentialQuotaService {
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_refresh_can_update(
|
||||
fn quota_refresh_allows_recovery(
|
||||
account: &ProviderAccount,
|
||||
refreshed: &CodexAccountQuotaSnapshot,
|
||||
previous: Option<&CodexAccountQuotaSnapshot>,
|
||||
@@ -1138,18 +1198,48 @@ fn access_token_is_current(account: &ProviderAccount, now: SystemTime) -> bool {
|
||||
|
||||
fn merge_passive_quota(
|
||||
mut quota: Map<String, Value>,
|
||||
rate_limits: &ParsedRateLimits,
|
||||
observations: &[ParsedRateLimits],
|
||||
) -> Map<String, Value> {
|
||||
let active_limit = rate_limits
|
||||
.active_limit
|
||||
.as_deref()
|
||||
.or_else(|| rate_limits.limits.contains_key("codex").then_some("codex"));
|
||||
if let Some(active_limit) = active_limit {
|
||||
quota.insert(
|
||||
"active_limit".to_owned(),
|
||||
Value::String(active_limit.to_owned()),
|
||||
);
|
||||
let mut snapshots = RateLimitSnapshotsByLimitId::take_from_document(&mut quota);
|
||||
for rate_limits in observations {
|
||||
let default_is_named_alias = default_limit_is_named_alias(rate_limits);
|
||||
let mut resolved_limit_ids = BTreeMap::new();
|
||||
for (wire_limit_id, details) in &rate_limits.limits {
|
||||
// HTTP 与 WebSocket 都可能把活动具名桶镜像为默认 `codex` 窗口。
|
||||
// 同一 wire 观察里存在相同具名事实时丢弃镜像,不触碰 core 桶。
|
||||
if wire_limit_id == DEFAULT_CODEX_LIMIT_ID && default_is_named_alias {
|
||||
continue;
|
||||
}
|
||||
let Some(limit_id) = snapshots.resolve_limit_id(details) else {
|
||||
continue;
|
||||
};
|
||||
let Some(rate_limit) = passive_rate_limit_snapshot(details) else {
|
||||
continue;
|
||||
};
|
||||
snapshots.upsert(&limit_id, details, rate_limit);
|
||||
resolved_limit_ids.insert(wire_limit_id.as_str(), limit_id);
|
||||
}
|
||||
|
||||
let active_limit = rate_limits
|
||||
.active_limit
|
||||
.as_deref()
|
||||
.and_then(|wire_limit_id| {
|
||||
resolved_limit_ids.get(wire_limit_id).cloned().or_else(|| {
|
||||
(!rate_limits.limits.contains_key(wire_limit_id))
|
||||
.then(|| wire_limit_id.to_owned())
|
||||
})
|
||||
})
|
||||
.or_else(|| resolved_limit_ids.get(DEFAULT_CODEX_LIMIT_ID).cloned());
|
||||
if let Some(active_limit) = active_limit {
|
||||
quota.insert("active_limit".to_owned(), Value::String(active_limit));
|
||||
}
|
||||
merge_passive_metadata(&mut quota, rate_limits);
|
||||
}
|
||||
snapshots.write_to_document(&mut quota);
|
||||
quota
|
||||
}
|
||||
|
||||
fn merge_passive_metadata(quota: &mut Map<String, Value>, rate_limits: &ParsedRateLimits) {
|
||||
if let Some(plan_type) = rate_limits.plan_type.as_ref() {
|
||||
quota.insert("plan_type".to_owned(), Value::String(plan_type.clone()));
|
||||
}
|
||||
@@ -1177,48 +1267,19 @@ fn merge_passive_quota(
|
||||
}
|
||||
quota.insert("credits".to_owned(), Value::Object(value));
|
||||
}
|
||||
}
|
||||
|
||||
let mut additional = quota
|
||||
.remove("additional_rate_limits")
|
||||
.and_then(|value| value.as_array().cloned())
|
||||
.unwrap_or_default();
|
||||
for (limit_id, details) in &rate_limits.limits {
|
||||
let Some(rate_limit) = passive_rate_limit_snapshot(details) else {
|
||||
continue;
|
||||
};
|
||||
if Some(limit_id.as_str()) == active_limit {
|
||||
quota.insert("rate_limit".to_owned(), Value::Object(rate_limit));
|
||||
continue;
|
||||
}
|
||||
let index = additional.iter().position(|item| {
|
||||
item.get("metered_feature")
|
||||
.or_else(|| item.get("limit_name"))
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|existing| existing == limit_id)
|
||||
});
|
||||
let mut item = index
|
||||
.and_then(|index| additional.get(index))
|
||||
.and_then(Value::as_object)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
item.insert(
|
||||
"metered_feature".to_owned(),
|
||||
Value::String(limit_id.clone()),
|
||||
);
|
||||
if let Some(name) = details.limit_name.as_ref() {
|
||||
item.insert("limit_name".to_owned(), Value::String(name.clone()));
|
||||
}
|
||||
item.insert("rate_limit".to_owned(), Value::Object(rate_limit));
|
||||
match index {
|
||||
Some(index) => additional[index] = Value::Object(item),
|
||||
None => additional.push(Value::Object(item)),
|
||||
}
|
||||
}
|
||||
quota.insert(
|
||||
"additional_rate_limits".to_owned(),
|
||||
Value::Array(additional),
|
||||
);
|
||||
quota
|
||||
fn default_limit_is_named_alias(rate_limits: &ParsedRateLimits) -> bool {
|
||||
let Some(default_primary) = rate_limits
|
||||
.limits
|
||||
.get(DEFAULT_CODEX_LIMIT_ID)
|
||||
.and_then(|details| details.primary)
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
rate_limits.limits.iter().any(|(limit_id, details)| {
|
||||
limit_id != DEFAULT_CODEX_LIMIT_ID && details.primary == Some(default_primary)
|
||||
})
|
||||
}
|
||||
|
||||
/// 丢弃 core 限额中上游给出的无事实 `secondary_window` 占位。
|
||||
@@ -1230,7 +1291,13 @@ fn merge_passive_quota(
|
||||
/// 无事实值。带 reset、时长、正用量、触顶、未知或非法字段的次级窗口均完全按
|
||||
/// 原有额度逻辑保留。
|
||||
fn normalize_quota_window_placeholders(mut quota: Map<String, Value>) -> Map<String, Value> {
|
||||
if let Some(rate_limit) = quota.get_mut("rate_limit").and_then(Value::as_object_mut) {
|
||||
quota = canonicalize_rate_limit_document(quota);
|
||||
if let Some(rate_limit) = quota
|
||||
.get_mut(RATE_LIMITS_BY_LIMIT_ID)
|
||||
.and_then(Value::as_object_mut)
|
||||
.and_then(|limits| limits.get_mut(DEFAULT_CODEX_LIMIT_ID))
|
||||
.and_then(Value::as_object_mut)
|
||||
{
|
||||
drop_secondary_window_placeholder(rate_limit);
|
||||
}
|
||||
quota
|
||||
|
||||
@@ -12,6 +12,9 @@ use gateway_core::engine::credential::{
|
||||
};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use super::document::{
|
||||
DEFAULT_CODEX_LIMIT_ID, RATE_LIMITS_BY_LIMIT_ID, canonicalize_rate_limit_document,
|
||||
};
|
||||
use super::{CodexCredentialQuotaError, QUOTA_SCHEDULING_TTL};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
@@ -151,6 +154,13 @@ impl CodexAccountQuotaSnapshot {
|
||||
self.quota = quota;
|
||||
self
|
||||
}
|
||||
|
||||
/// 只更新展示与调度使用的最后成功查询时间,保留既有额度事实。
|
||||
#[must_use]
|
||||
pub(crate) const fn with_observed_at(mut self, observed_at: SystemTime) -> Self {
|
||||
self.observed_at = observed_at;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl CodexQuotaFact {
|
||||
@@ -166,9 +176,18 @@ impl CodexQuotaFact {
|
||||
}
|
||||
|
||||
pub fn parse_codex_quota_usage(usage: &Value) -> Result<CodexQuotaFact, CodexCredentialQuotaError> {
|
||||
let object = usage
|
||||
.as_object()
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
|
||||
let object = canonicalize_rate_limit_document(
|
||||
usage
|
||||
.as_object()
|
||||
.cloned()
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?,
|
||||
);
|
||||
parse_codex_quota_object(&object)
|
||||
}
|
||||
|
||||
fn parse_codex_quota_object(
|
||||
object: &Map<String, Value>,
|
||||
) -> Result<CodexQuotaFact, CodexCredentialQuotaError> {
|
||||
let mut aggregate = QuotaAggregate::default();
|
||||
for limit in canonical_rate_limits(object)? {
|
||||
aggregate.observe_rate_limit(limit.rate_limit)?;
|
||||
@@ -189,6 +208,15 @@ pub fn parse_codex_quota_usage(usage: &Value) -> Result<CodexQuotaFact, CodexCre
|
||||
})
|
||||
}
|
||||
|
||||
fn canonical_quota_object(usage: &Value) -> Result<Map<String, Value>, CodexCredentialQuotaError> {
|
||||
Ok(canonicalize_rate_limit_document(
|
||||
usage
|
||||
.as_object()
|
||||
.cloned()
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?,
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) fn quota_snapshot_from_observation(
|
||||
observation: &QuotaObservation,
|
||||
) -> Option<CodexAccountQuotaSnapshot> {
|
||||
@@ -270,13 +298,11 @@ pub(crate) fn parse_account_quota_snapshot(
|
||||
observed_at: SystemTime,
|
||||
usage: &Value,
|
||||
) -> Result<CodexAccountQuotaSnapshot, CodexCredentialQuotaError> {
|
||||
let fact = parse_codex_quota_usage(usage)?;
|
||||
let object = usage
|
||||
.as_object()
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
|
||||
let quota = authoritative_quota_state(object, observed_at)?;
|
||||
let object = canonical_quota_object(usage)?;
|
||||
let fact = parse_codex_quota_object(&object)?;
|
||||
let quota = authoritative_quota_state(&object, observed_at)?;
|
||||
let mut windows = Vec::new();
|
||||
for limit in canonical_rate_limits(object)? {
|
||||
for limit in canonical_rate_limits(&object)? {
|
||||
parse_rate_limit_windows(
|
||||
&limit.key,
|
||||
&limit.source,
|
||||
@@ -309,12 +335,16 @@ fn quota_source_order(source: &str) -> u8 {
|
||||
}
|
||||
}
|
||||
|
||||
/// 顶层 account-wide Codex 桶是唯一可改变账号额度访问结论的 quota 响应字段。
|
||||
/// 默认 `codex` 桶是唯一可改变账号额度访问结论的 quota 响应字段。
|
||||
fn authoritative_quota_state(
|
||||
usage: &Map<String, Value>,
|
||||
observed_at: SystemTime,
|
||||
) -> Result<QuotaState, CodexCredentialQuotaError> {
|
||||
let Some(rate_limit) = usage.get("rate_limit") else {
|
||||
let Some(rate_limit) = usage
|
||||
.get(RATE_LIMITS_BY_LIMIT_ID)
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|limits| limits.get(DEFAULT_CODEX_LIMIT_ID))
|
||||
else {
|
||||
return Ok(QuotaState::observed_unknown(observed_at));
|
||||
};
|
||||
if rate_limit.is_null() {
|
||||
@@ -391,64 +421,32 @@ fn canonical_rate_limits(
|
||||
object: &Map<String, Value>,
|
||||
) -> Result<Vec<CanonicalRateLimit<'_>>, CodexCredentialQuotaError> {
|
||||
let mut limits = Vec::new();
|
||||
let has_primary_codex_rate_limit = object
|
||||
.get("rate_limit")
|
||||
.is_some_and(|value| !value.is_null());
|
||||
if let Some(rate_limit) = object.get("rate_limit") {
|
||||
// 官方 Codex 将顶层 `rate_limit` 固定识别为 `codex`,而不把
|
||||
// `active_limit` 投影成另一个展示桶。后者只是当前活动套餐的元数据。
|
||||
limits.push(CanonicalRateLimit {
|
||||
key: "core".to_owned(),
|
||||
source: "codex".to_owned(),
|
||||
account_wide: true,
|
||||
limit_name: None,
|
||||
rate_limit,
|
||||
});
|
||||
}
|
||||
if let Some(rate_limit) = object
|
||||
.get("code_review_rate_limit")
|
||||
.filter(|value| !value.is_null())
|
||||
{
|
||||
limits.push(CanonicalRateLimit {
|
||||
key: "code-review".to_owned(),
|
||||
source: "code_review".to_owned(),
|
||||
account_wide: false,
|
||||
limit_name: Some("code_review".to_owned()),
|
||||
rate_limit,
|
||||
});
|
||||
}
|
||||
let Some(additional) = object
|
||||
.get("additional_rate_limits")
|
||||
.filter(|value| !value.is_null())
|
||||
else {
|
||||
let Some(by_limit_id) = object.get(RATE_LIMITS_BY_LIMIT_ID) else {
|
||||
return Ok(limits);
|
||||
};
|
||||
for (index, value) in additional
|
||||
.as_array()
|
||||
for (limit_id, rate_limit) in by_limit_id
|
||||
.as_object()
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?
|
||||
.iter()
|
||||
.enumerate()
|
||||
{
|
||||
let item = value
|
||||
let item = rate_limit
|
||||
.as_object()
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
|
||||
let (source, limit_name) = canonical_rate_limit_source(item, index);
|
||||
// 部分上游响应会在 `additional_rate_limits` 中重复默认的 `codex`
|
||||
// 桶。它没有独立的 limit id,不能在账号面板再次展示。
|
||||
if has_primary_codex_rate_limit && is_codex_limit_id(&source) {
|
||||
continue;
|
||||
}
|
||||
if let Some(rate_limit) = item.get("rate_limit") {
|
||||
// 非 Codex additional 桶保留为独立快照;额外用索引生成窗口键,
|
||||
// 避免同名 additional 产生重复前端 key。
|
||||
limits.push(CanonicalRateLimit {
|
||||
key: format!("additional-{index}-{source}"),
|
||||
account_wide: is_codex_limit_id(&source),
|
||||
source,
|
||||
limit_name,
|
||||
rate_limit,
|
||||
});
|
||||
}
|
||||
let source = item
|
||||
.get("limit_id")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or(limit_id)
|
||||
.to_owned();
|
||||
let limit_name = item
|
||||
.get("limit_name")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned);
|
||||
limits.push(CanonicalRateLimit {
|
||||
key: source.clone(),
|
||||
account_wide: is_codex_limit_id(&source),
|
||||
source,
|
||||
limit_name,
|
||||
rate_limit,
|
||||
});
|
||||
}
|
||||
Ok(limits)
|
||||
}
|
||||
@@ -457,30 +455,6 @@ fn is_codex_limit_id(value: &str) -> bool {
|
||||
value.trim().eq_ignore_ascii_case("codex")
|
||||
}
|
||||
|
||||
fn canonical_rate_limit_source(
|
||||
item: &Map<String, Value>,
|
||||
index: usize,
|
||||
) -> (String, Option<String>) {
|
||||
let source = item
|
||||
.get("metered_feature")
|
||||
.or_else(|| item.get("limit_name"))
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| {
|
||||
!value.is_empty() && value.len() <= 128 && !value.chars().any(char::is_control)
|
||||
})
|
||||
.map(str::to_owned)
|
||||
.unwrap_or_else(|| format!("additional-{index}"));
|
||||
let limit_name = item
|
||||
.get("limit_name")
|
||||
.or_else(|| item.get("metered_feature"))
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| {
|
||||
!value.is_empty() && value.len() <= 128 && !value.chars().any(char::is_control)
|
||||
})
|
||||
.map(str::to_owned);
|
||||
(source, limit_name)
|
||||
}
|
||||
|
||||
fn parse_rate_limit_windows(
|
||||
key_source: &str,
|
||||
source: &str,
|
||||
|
||||
@@ -20,8 +20,8 @@ use thiserror::Error;
|
||||
|
||||
use super::recovery_log::{CodexOAuthRecoveryOperation, record_oauth_recovery};
|
||||
use super::repository::{CodexCredentialRepository, CredentialRepositoryError};
|
||||
use super::token_client::{RefreshFailure, TokenPair, TokenRefresher};
|
||||
use super::types::CodexOAuthSecret;
|
||||
use super::token_client::{RefreshFailure, RefreshUpstreamFailure, TokenPair, TokenRefresher};
|
||||
use super::types::{CodexOAuthSecret, parse_access_token_expiration};
|
||||
|
||||
const PROVIDER_NAME: &str = "openai";
|
||||
const MAX_REFRESH_BATCH: u32 = 1_000;
|
||||
@@ -32,7 +32,7 @@ const REFRESH_BACKOFF_MAX_ATTEMPTS: u32 = 5;
|
||||
/// 耗尽指数退避后的 OAuth 恢复周期。
|
||||
const REFRESH_RECOVERY_DELAY: Duration = Duration::from_secs(10 * 60);
|
||||
/// 过期 AT 仍允许 RT 恢复的最长窗口。
|
||||
const REFRESH_RECOVERY_WINDOW: Duration = Duration::from_secs(24 * 60 * 60);
|
||||
const REFRESH_RECOVERY_WINDOW: Duration = Duration::from_secs(2 * 60 * 60);
|
||||
|
||||
pub(crate) fn refresh_recovery_deadline(
|
||||
access_token_expires_at: Option<SystemTime>,
|
||||
@@ -91,12 +91,19 @@ fn log_refresh_deferred(
|
||||
access_token_expires_at: Option<SystemTime>,
|
||||
attempt: u32,
|
||||
reason: &'static str,
|
||||
upstream_message: Option<&str>,
|
||||
upstream: Option<&RefreshUpstreamFailure>,
|
||||
retry_at: SystemTime,
|
||||
) {
|
||||
tracing::warn!(
|
||||
account_id = %account_id,
|
||||
attempt,
|
||||
reason,
|
||||
upstream_message = ?upstream_message,
|
||||
upstream_status = ?upstream.map(RefreshUpstreamFailure::status),
|
||||
upstream_code = ?upstream.and_then(RefreshUpstreamFailure::code),
|
||||
upstream_type = ?upstream.and_then(RefreshUpstreamFailure::error_type),
|
||||
upstream_body = ?upstream.map(RefreshUpstreamFailure::body),
|
||||
retry_at = %DateTime::<Utc>::from(retry_at),
|
||||
access_token_expires_at = ?access_token_expires_at.map(DateTime::<Utc>::from),
|
||||
recovery_deadline = ?refresh_recovery_deadline(access_token_expires_at)
|
||||
@@ -105,6 +112,26 @@ fn log_refresh_deferred(
|
||||
);
|
||||
}
|
||||
|
||||
struct RefreshFailureContext<'a> {
|
||||
reason: &'static str,
|
||||
message: Option<String>,
|
||||
upstream: Option<&'a RefreshUpstreamFailure>,
|
||||
}
|
||||
|
||||
impl<'a> RefreshFailureContext<'a> {
|
||||
const fn new(
|
||||
reason: &'static str,
|
||||
message: Option<String>,
|
||||
upstream: Option<&'a RefreshUpstreamFailure>,
|
||||
) -> Self {
|
||||
Self {
|
||||
reason,
|
||||
message,
|
||||
upstream,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn refresh_due_at(
|
||||
account: &ProviderAccount,
|
||||
policy: ProviderRefreshPolicy,
|
||||
@@ -248,20 +275,10 @@ impl CodexCredentialRefreshService {
|
||||
policy: ProviderRefreshPolicy,
|
||||
) -> Result<CodexCredentialRefreshOutcome, CodexCredentialRefreshError> {
|
||||
let account_id = due.account.id().to_string();
|
||||
if refresh_recovery_window_exhausted(
|
||||
let recovery_window_exhausted = refresh_recovery_window_exhausted(
|
||||
due.account.access_token_expires_at(),
|
||||
SystemTime::now(),
|
||||
) {
|
||||
return self
|
||||
.persist_terminal(
|
||||
&due.account,
|
||||
CredentialState::Expired,
|
||||
AccountErrorReason::CredentialExpired,
|
||||
"refresh_recovery_window_exhausted",
|
||||
CodexCredentialRefreshOutcome::Invalidated { account_id },
|
||||
)
|
||||
.await;
|
||||
}
|
||||
);
|
||||
let capacity = self
|
||||
.leases
|
||||
.try_acquire(ProviderLeaseRequest::RefreshCapacity(
|
||||
@@ -291,31 +308,72 @@ impl CodexCredentialRefreshService {
|
||||
.refresh_token
|
||||
.as_ref()
|
||||
.ok_or(CodexCredentialRefreshError::InvalidRefreshResponse)?;
|
||||
match self.refresher.refresh(refresh_token.expose_secret()).await {
|
||||
let refresh_result = self.refresher.refresh(refresh_token.expose_secret()).await;
|
||||
if recovery_window_exhausted && let Err(failure) = &refresh_result {
|
||||
let message = failure.message().map(str::to_owned);
|
||||
let upstream = failure.upstream();
|
||||
return match failure {
|
||||
RefreshFailure::Banned { .. } => {
|
||||
self.persist_terminal(
|
||||
&due.account,
|
||||
CredentialState::Banned,
|
||||
AccountErrorReason::AccountBanned,
|
||||
RefreshFailureContext::new(
|
||||
"refresh_recovery_window_exhausted",
|
||||
message,
|
||||
upstream,
|
||||
),
|
||||
CodexCredentialRefreshOutcome::Banned { account_id },
|
||||
)
|
||||
.await
|
||||
}
|
||||
RefreshFailure::InvalidGrant { .. }
|
||||
| RefreshFailure::RetryableTransport { .. }
|
||||
| RefreshFailure::Transport { .. } => {
|
||||
self.persist_terminal(
|
||||
&due.account,
|
||||
CredentialState::Expired,
|
||||
AccountErrorReason::CredentialExpired,
|
||||
RefreshFailureContext::new(
|
||||
"refresh_recovery_window_exhausted",
|
||||
message,
|
||||
upstream,
|
||||
),
|
||||
CodexCredentialRefreshOutcome::Invalidated { account_id },
|
||||
)
|
||||
.await
|
||||
}
|
||||
};
|
||||
}
|
||||
match refresh_result {
|
||||
Ok(tokens) => self.persist_success(due, tokens).await,
|
||||
Err(RefreshFailure::InvalidGrant) => {
|
||||
Err(RefreshFailure::InvalidGrant { message, upstream }) => {
|
||||
self.persist_terminal(
|
||||
&due.account,
|
||||
CredentialState::Expired,
|
||||
AccountErrorReason::CredentialExpired,
|
||||
"refresh_invalid_grant",
|
||||
RefreshFailureContext::new(
|
||||
"refresh_invalid_grant",
|
||||
message,
|
||||
upstream.as_deref(),
|
||||
),
|
||||
CodexCredentialRefreshOutcome::Invalidated { account_id },
|
||||
)
|
||||
.await
|
||||
}
|
||||
Err(RefreshFailure::Banned) => {
|
||||
Err(RefreshFailure::Banned { message, upstream }) => {
|
||||
self.persist_terminal(
|
||||
&due.account,
|
||||
CredentialState::Banned,
|
||||
AccountErrorReason::AccountBanned,
|
||||
"account_banned",
|
||||
RefreshFailureContext::new("account_banned", message, upstream.as_deref()),
|
||||
CodexCredentialRefreshOutcome::Banned { account_id },
|
||||
)
|
||||
.await
|
||||
}
|
||||
Err(RefreshFailure::RetryableTransport) => {
|
||||
Err(RefreshFailure::RetryableTransport { message }) => {
|
||||
if self
|
||||
.defer_refresh(&due.account, "transport-not-sent")
|
||||
.defer_refresh(&due.account, "transport-not-sent", Some(&message), None)
|
||||
.await?
|
||||
{
|
||||
Ok(CodexCredentialRefreshOutcome::Transient { account_id })
|
||||
@@ -323,11 +381,16 @@ impl CodexCredentialRefreshService {
|
||||
Ok(CodexCredentialRefreshOutcome::Stale { account_id })
|
||||
}
|
||||
}
|
||||
Err(RefreshFailure::Transport) => {
|
||||
// 上游瞬态(429/5xx/超时/畸形响应等,发送状态未知)保留现有凭据、
|
||||
// 推进退避重试;仅 invalid_grant/banned 才终态失效账号。
|
||||
Err(RefreshFailure::Transport { message, upstream }) => {
|
||||
// 上游瞬态(401/429/5xx/超时/畸形响应等)保留现有凭据、
|
||||
// 记录最近一次失败并推进有界退避。
|
||||
if self
|
||||
.defer_refresh(&due.account, "transport-ambiguous")
|
||||
.defer_refresh(
|
||||
&due.account,
|
||||
"transport-ambiguous",
|
||||
message.as_deref(),
|
||||
upstream.as_deref(),
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Ok(CodexCredentialRefreshOutcome::Transient { account_id })
|
||||
@@ -403,9 +466,11 @@ impl CodexCredentialRefreshService {
|
||||
access_token,
|
||||
refresh_token: rotated_refresh_token,
|
||||
id_token: rotated_id_token,
|
||||
expires_in,
|
||||
} = tokens;
|
||||
let access_token_rotated = access_token.is_some();
|
||||
let access_token_expires_at = match access_token.as_deref() {
|
||||
Some(access_token) => parse_access_token_expiration(access_token).map(SystemTime::from),
|
||||
None => due.account.access_token_expires_at(),
|
||||
};
|
||||
let CodexOAuthSecret {
|
||||
access_token: current_access_token,
|
||||
refresh_token: current_refresh_token,
|
||||
@@ -432,12 +497,6 @@ impl CodexCredentialRefreshService {
|
||||
.as_ref()
|
||||
.map(ExposeSecret::expose_secret),
|
||||
);
|
||||
let observed_at = SystemTime::now();
|
||||
let access_token_expires_at = if access_token_rotated {
|
||||
expires_in.and_then(|value| observed_at.checked_add(value))
|
||||
} else {
|
||||
due.account.access_token_expires_at()
|
||||
};
|
||||
// 成功轮换清掉仅用于失败退避的 retry-not-before;正常预刷新窗口由
|
||||
// worker 结合当前 runtime policy 动态判断,不写入账号时间字段。
|
||||
let next_refresh_at = None;
|
||||
@@ -483,9 +542,14 @@ impl CodexCredentialRefreshService {
|
||||
account: &ProviderAccount,
|
||||
credential_state: CredentialState,
|
||||
error_reason: AccountErrorReason,
|
||||
reason: &'static str,
|
||||
failure: RefreshFailureContext<'_>,
|
||||
outcome: CodexCredentialRefreshOutcome,
|
||||
) -> Result<CodexCredentialRefreshOutcome, CodexCredentialRefreshError> {
|
||||
let RefreshFailureContext {
|
||||
reason,
|
||||
message,
|
||||
upstream,
|
||||
} = failure;
|
||||
match self.repository.load_runtime_credential(account).await {
|
||||
Ok(_) => {}
|
||||
Err(CredentialRepositoryError::RevisionConflict) => {
|
||||
@@ -495,6 +559,12 @@ impl CodexCredentialRefreshService {
|
||||
}
|
||||
Err(error) => return Err(error.into()),
|
||||
}
|
||||
let persisted_message = message
|
||||
.as_ref()
|
||||
.filter(|message| !message.trim().is_empty())
|
||||
.map(String::as_str)
|
||||
.or_else(|| account.last_error_message())
|
||||
.map(str::to_owned);
|
||||
match self
|
||||
.repository
|
||||
.apply_state_with_reason(
|
||||
@@ -502,7 +572,7 @@ impl CodexCredentialRefreshService {
|
||||
credential_state,
|
||||
SystemTime::now(),
|
||||
Some(error_reason),
|
||||
Some(reason.to_owned()),
|
||||
persisted_message,
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -511,6 +581,11 @@ impl CodexCredentialRefreshService {
|
||||
account_id = %account.id(),
|
||||
?credential_state,
|
||||
reason,
|
||||
upstream_message = ?message.as_deref(),
|
||||
upstream_status = ?upstream.map(RefreshUpstreamFailure::status),
|
||||
upstream_code = ?upstream.and_then(RefreshUpstreamFailure::code),
|
||||
upstream_type = ?upstream.and_then(RefreshUpstreamFailure::error_type),
|
||||
upstream_body = ?upstream.map(RefreshUpstreamFailure::body),
|
||||
access_token_expires_at = ?account.access_token_expires_at()
|
||||
.map(DateTime::<Utc>::from),
|
||||
recovery_deadline = ?refresh_recovery_deadline(account.access_token_expires_at())
|
||||
@@ -532,6 +607,8 @@ impl CodexCredentialRefreshService {
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
reason: &'static str,
|
||||
upstream_message: Option<&str>,
|
||||
upstream: Option<&RefreshUpstreamFailure>,
|
||||
) -> Result<bool, CodexCredentialRefreshError> {
|
||||
let attempt = self
|
||||
.credential_state
|
||||
@@ -545,13 +622,28 @@ impl CodexCredentialRefreshService {
|
||||
attempt,
|
||||
reason,
|
||||
)?;
|
||||
match self.repository.defer_refresh(account, retry_at).await {
|
||||
let message = upstream_message
|
||||
.filter(|message| !message.trim().is_empty())
|
||||
.map(str::to_owned);
|
||||
let error_reason = upstream.map(|_| {
|
||||
account
|
||||
.credential_state()
|
||||
.error_reason()
|
||||
.unwrap_or(AccountErrorReason::AccessTokenExpired)
|
||||
});
|
||||
match self
|
||||
.repository
|
||||
.defer_refresh(account, retry_at, error_reason, message)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {
|
||||
log_refresh_deferred(
|
||||
account.id(),
|
||||
account.access_token_expires_at(),
|
||||
attempt,
|
||||
reason,
|
||||
upstream_message,
|
||||
upstream,
|
||||
retry_at,
|
||||
);
|
||||
Ok(true)
|
||||
|
||||
@@ -120,15 +120,18 @@ impl CodexCredentialRepository {
|
||||
access_token_expires_at,
|
||||
next_refresh_at,
|
||||
)
|
||||
.map_err(|_| CredentialRepositoryError::InvalidCredentialData)?;
|
||||
.map_err(|_| CredentialRepositoryError::InvalidCredentialData)?
|
||||
.with_account_state(CredentialState::Ready, SystemTime::now(), None, None);
|
||||
cas_revision(self.store.compare_and_swap_credential(update).await?)
|
||||
}
|
||||
|
||||
/// token 签名边界不可用时,以相同 credential 推进持久刷新退避。
|
||||
/// 以相同 credential 原子推进刷新退避及本次上游错误事实。
|
||||
pub async fn defer_refresh(
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
next_refresh_at: SystemTime,
|
||||
error_reason: Option<AccountErrorReason>,
|
||||
message: Option<String>,
|
||||
) -> Result<CredentialRevision, CredentialRepositoryError> {
|
||||
let current = self
|
||||
.store
|
||||
@@ -140,7 +143,7 @@ impl CodexCredentialRepository {
|
||||
let data = CodexCredentialCodec::decode_complete(¤t.credential)?;
|
||||
let has_refresh_token = data.has_refresh_token();
|
||||
let credential = CodexCredentialCodec::encode_complete(data)?;
|
||||
let update = CredentialCasUpdate::new(
|
||||
let mut update = CredentialCasUpdate::new(
|
||||
account.id().clone(),
|
||||
account.revision(),
|
||||
unchanged_profile(account),
|
||||
@@ -150,6 +153,14 @@ impl CodexCredentialRepository {
|
||||
Some(next_refresh_at),
|
||||
)
|
||||
.map_err(|_| CredentialRepositoryError::InvalidCredentialData)?;
|
||||
if let Some(error_reason) = error_reason {
|
||||
update = update.with_account_state(
|
||||
account.credential_state(),
|
||||
SystemTime::now(),
|
||||
Some(error_reason),
|
||||
message,
|
||||
);
|
||||
}
|
||||
cas_revision(self.store.compare_and_swap_credential(update).await?)
|
||||
}
|
||||
|
||||
@@ -245,13 +256,11 @@ impl CodexCredentialRepository {
|
||||
if !account.enabled() {
|
||||
return Ok(());
|
||||
}
|
||||
let (error_reason, message) = if credential_state == CredentialState::Ready {
|
||||
(None, None)
|
||||
let message = message.filter(|value| !value.trim().is_empty());
|
||||
let error_reason = if credential_state == CredentialState::Ready {
|
||||
message.as_ref().and(error_reason)
|
||||
} else {
|
||||
(
|
||||
error_reason.or_else(|| credential_state.error_reason()),
|
||||
message.filter(|value| !value.trim().is_empty()),
|
||||
)
|
||||
error_reason.or_else(|| credential_state.error_reason())
|
||||
};
|
||||
self.store
|
||||
.apply_state_change(AccountStateChange {
|
||||
|
||||
@@ -108,6 +108,7 @@ impl CodexCredentialCodec {
|
||||
oauth_subject: account.oauth_subject.clone(),
|
||||
poid: account.poid.clone(),
|
||||
}),
|
||||
uuid::Uuid::new_v4().to_string(),
|
||||
cookies,
|
||||
)
|
||||
}
|
||||
@@ -115,20 +116,22 @@ impl CodexCredentialCodec {
|
||||
/// 为尚未解析资料的 OAuth 账号编码凭据。
|
||||
pub(crate) fn encode_unresolved(
|
||||
secret: &CodexOAuthSecret,
|
||||
installation_id: String,
|
||||
cookies: Vec<CodexCookie>,
|
||||
) -> Result<PlaintextCredential, CodexCredentialDataError> {
|
||||
Self::encode_oauth(secret, None, cookies)
|
||||
Self::encode_oauth(secret, None, installation_id, cookies)
|
||||
}
|
||||
|
||||
fn encode_oauth(
|
||||
secret: &CodexOAuthSecret,
|
||||
principal: Option<CodexCredentialPrincipal>,
|
||||
installation_id: String,
|
||||
cookies: Vec<CodexCookie>,
|
||||
) -> Result<PlaintextCredential, CodexCredentialDataError> {
|
||||
Self::encode_complete(CodexCredentialData::OAuth(CodexOAuthCredentialData {
|
||||
schema_version: CODEX_CREDENTIAL_SCHEMA_VERSION,
|
||||
principal,
|
||||
installation_id: uuid::Uuid::new_v4().to_string(),
|
||||
installation_id,
|
||||
access_token: secret.access_token.expose_secret().to_owned(),
|
||||
refresh_token: secret
|
||||
.refresh_token
|
||||
|
||||
@@ -720,7 +720,15 @@ impl CodexCredentialSelector {
|
||||
.map(chrono::DateTime::<chrono::Utc>::from),
|
||||
"OpenAI access token expired; retaining account for bounded OAuth refresh recovery"
|
||||
);
|
||||
return Ok(());
|
||||
return self
|
||||
.apply_credential_state(
|
||||
account,
|
||||
CredentialState::Ready,
|
||||
AccountErrorReason::AccessTokenExpired,
|
||||
now,
|
||||
message,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
self.apply_credential_state(
|
||||
account,
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
use async_trait::async_trait;
|
||||
use reqwest::{Client, StatusCode, redirect::Policy};
|
||||
use secrecy::{ExposeSecret, SecretString};
|
||||
use serde::Deserialize;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::fmt;
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -27,9 +27,6 @@ pub struct TokenPair {
|
||||
pub refresh_token: Option<String>,
|
||||
/// 官方刷新响应省略时由持久化调用方保留当前 ID token。
|
||||
pub id_token: Option<String>,
|
||||
/// OAuth endpoint 可以省略或返回零 expiry;这不影响已经成功的 token
|
||||
/// exchange,只意味着调用方不能据此安排下一次刷新。
|
||||
pub expires_in: Option<Duration>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for TokenPair {
|
||||
@@ -45,22 +42,131 @@ impl fmt::Debug for TokenPair {
|
||||
&self.refresh_token.as_ref().map(|_| "[REDACTED]"),
|
||||
)
|
||||
.field("id_token", &self.id_token.as_ref().map(|_| "[REDACTED]"))
|
||||
.field("expires_in", &self.expires_in)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// Codex token 刷新的稳定失败分类。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
|
||||
#[derive(Clone, PartialEq, Eq, thiserror::Error)]
|
||||
pub enum RefreshFailure {
|
||||
#[error("refresh token is invalid or expired")]
|
||||
InvalidGrant,
|
||||
InvalidGrant {
|
||||
message: Option<String>,
|
||||
upstream: Option<Box<RefreshUpstreamFailure>>,
|
||||
},
|
||||
#[error("account is banned")]
|
||||
Banned,
|
||||
Banned {
|
||||
message: Option<String>,
|
||||
upstream: Option<Box<RefreshUpstreamFailure>>,
|
||||
},
|
||||
#[error("refresh transport failed before server processing")]
|
||||
RetryableTransport,
|
||||
RetryableTransport { message: String },
|
||||
#[error("refresh transport failed after possible server processing")]
|
||||
Transport,
|
||||
Transport {
|
||||
message: Option<String>,
|
||||
upstream: Option<Box<RefreshUpstreamFailure>>,
|
||||
},
|
||||
}
|
||||
|
||||
impl fmt::Debug for RefreshFailure {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("RefreshFailure")
|
||||
.field("classification", &self.classification())
|
||||
.field("message", &self.message().map(|_| "<redacted>"))
|
||||
.field("upstream", &self.upstream())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl RefreshFailure {
|
||||
#[must_use]
|
||||
pub fn message(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::InvalidGrant { message, .. }
|
||||
| Self::Banned { message, .. }
|
||||
| Self::Transport { message, .. } => message.as_deref(),
|
||||
Self::RetryableTransport { message } => Some(message),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn upstream(&self) -> Option<&RefreshUpstreamFailure> {
|
||||
match self {
|
||||
Self::InvalidGrant { upstream, .. }
|
||||
| Self::Banned { upstream, .. }
|
||||
| Self::Transport { upstream, .. } => upstream.as_deref(),
|
||||
Self::RetryableTransport { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn classification(&self) -> &'static str {
|
||||
match self {
|
||||
Self::InvalidGrant { .. } => "invalid-grant",
|
||||
Self::Banned { .. } => "account-banned",
|
||||
Self::RetryableTransport { .. } => "transport-not-sent",
|
||||
Self::Transport { .. } => "transport-ambiguous",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 当前 OAuth 刷新请求收到的完整非成功响应。
|
||||
///
|
||||
/// 该值不持久化;`Debug` 不输出正文。调用方仅在受控刷新失败日志中显式记录正文。
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub struct RefreshUpstreamFailure {
|
||||
status: u16,
|
||||
code: Option<String>,
|
||||
error_type: Option<String>,
|
||||
body: String,
|
||||
}
|
||||
|
||||
impl RefreshUpstreamFailure {
|
||||
fn new(status: StatusCode, body: &[u8], error: Option<&RefreshErrorResponse>) -> Self {
|
||||
Self {
|
||||
status: status.as_u16(),
|
||||
code: error
|
||||
.and_then(RefreshErrorResponse::code)
|
||||
.map(str::to_owned),
|
||||
error_type: error
|
||||
.and_then(RefreshErrorResponse::error_type)
|
||||
.map(str::to_owned),
|
||||
body: String::from_utf8_lossy(body).into_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn status(&self) -> u16 {
|
||||
self.status
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn code(&self) -> Option<&str> {
|
||||
self.code.as_deref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn error_type(&self) -> Option<&str> {
|
||||
self.error_type.as_deref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn body(&self) -> &str {
|
||||
&self.body
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for RefreshUpstreamFailure {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("RefreshUpstreamFailure")
|
||||
.field("status", &self.status)
|
||||
.field("code", &self.code)
|
||||
.field("error_type", &self.error_type)
|
||||
.field("body", &"<redacted>")
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// Codex token 刷新端口。
|
||||
@@ -92,7 +198,6 @@ impl fmt::Debug for AuthorizationCodeGrant {
|
||||
pub struct AuthorizationTokenSet {
|
||||
pub secret: crate::credential::CodexOAuthSecret,
|
||||
pub id_token: SecretString,
|
||||
pub expires_in: Option<Duration>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for AuthorizationTokenSet {
|
||||
@@ -101,7 +206,6 @@ impl fmt::Debug for AuthorizationTokenSet {
|
||||
.debug_struct("AuthorizationTokenSet")
|
||||
.field("secret", &"[REDACTED]")
|
||||
.field("id_token", &"[REDACTED]")
|
||||
.field("expires_in", &self.expires_in)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -183,7 +287,60 @@ struct RefreshTokenResponse {
|
||||
access_token: Option<String>,
|
||||
refresh_token: Option<String>,
|
||||
id_token: Option<String>,
|
||||
expires_in: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct RefreshTokenRequest<'a> {
|
||||
client_id: &'a str,
|
||||
grant_type: &'static str,
|
||||
refresh_token: &'a str,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RefreshErrorResponse {
|
||||
error: Option<RefreshError>,
|
||||
code: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(untagged)]
|
||||
enum RefreshError {
|
||||
Details(RefreshErrorDetails),
|
||||
Code(String),
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RefreshErrorDetails {
|
||||
code: Option<String>,
|
||||
message: Option<String>,
|
||||
#[serde(rename = "type")]
|
||||
error_type: Option<String>,
|
||||
}
|
||||
|
||||
impl RefreshErrorResponse {
|
||||
fn code(&self) -> Option<&str> {
|
||||
match self.error.as_ref() {
|
||||
Some(RefreshError::Details(error)) => error.code.as_deref(),
|
||||
Some(RefreshError::Code(code)) => Some(code.as_str()),
|
||||
None => None,
|
||||
}
|
||||
.or(self.code.as_deref())
|
||||
}
|
||||
|
||||
fn message(&self) -> Option<String> {
|
||||
match self.error.as_ref() {
|
||||
Some(RefreshError::Details(error)) => error.message.as_deref(),
|
||||
Some(RefreshError::Code(_)) | None => None,
|
||||
}
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn error_type(&self) -> Option<&str> {
|
||||
match self.error.as_ref() {
|
||||
Some(RefreshError::Details(error)) => error.error_type.as_deref(),
|
||||
Some(RefreshError::Code(_)) | None => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
@@ -191,7 +348,6 @@ struct AuthorizationCodeResponse {
|
||||
access_token: String,
|
||||
refresh_token: String,
|
||||
id_token: String,
|
||||
expires_in: Option<u64>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -200,11 +356,11 @@ impl TokenRefresher for OpenAiTokenClient {
|
||||
let response = self
|
||||
.client
|
||||
.post(&self.config.token_endpoint)
|
||||
.form(&[
|
||||
("grant_type", "refresh_token"),
|
||||
("client_id", self.config.client_id.as_str()),
|
||||
("refresh_token", refresh_token),
|
||||
])
|
||||
.json(&RefreshTokenRequest {
|
||||
client_id: self.config.client_id.as_str(),
|
||||
grant_type: "refresh_token",
|
||||
refresh_token,
|
||||
})
|
||||
.send()
|
||||
.await
|
||||
.map_err(|error| refresh_transport_failure(&error))?;
|
||||
@@ -212,7 +368,10 @@ impl TokenRefresher for OpenAiTokenClient {
|
||||
if !status.is_success() {
|
||||
return Err(classify_refresh_failure(status, &body));
|
||||
}
|
||||
parse_token_pair(&body).map_err(|()| RefreshFailure::Transport)
|
||||
parse_token_pair(&body).map_err(|()| RefreshFailure::Transport {
|
||||
message: Some("OpenAI OAuth refresh returned an invalid success response".to_owned()),
|
||||
upstream: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -260,7 +419,6 @@ impl AuthorizationCodeExchanger for OpenAiTokenClient {
|
||||
id_token: None,
|
||||
},
|
||||
id_token,
|
||||
expires_in: optional_expiry(tokens.expires_in),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -273,13 +431,21 @@ async fn read_bounded_response(
|
||||
while let Some(chunk) = response
|
||||
.chunk()
|
||||
.await
|
||||
.map_err(|_| RefreshFailure::Transport)?
|
||||
.map_err(|error| RefreshFailure::Transport {
|
||||
message: Some(error.to_string()),
|
||||
upstream: None,
|
||||
})?
|
||||
{
|
||||
let next_len = body
|
||||
.len()
|
||||
.checked_add(chunk.len())
|
||||
.filter(|length| *length <= MAX_OAUTH_RESPONSE_BYTES)
|
||||
.ok_or(RefreshFailure::Transport)?;
|
||||
.ok_or_else(|| RefreshFailure::Transport {
|
||||
message: Some(format!(
|
||||
"OpenAI OAuth response exceeded {MAX_OAUTH_RESPONSE_BYTES} bytes"
|
||||
)),
|
||||
upstream: None,
|
||||
})?;
|
||||
body.reserve(next_len.saturating_sub(body.len()));
|
||||
body.extend_from_slice(&chunk);
|
||||
}
|
||||
@@ -301,43 +467,68 @@ fn parse_token_pair(body: &[u8]) -> Result<TokenPair, ()> {
|
||||
refresh_token: tokens.refresh_token,
|
||||
// 官方刷新响应允许省略 ID token;缺失时由调用方保留当前值。
|
||||
id_token: tokens.id_token,
|
||||
expires_in: optional_expiry(tokens.expires_in),
|
||||
})
|
||||
}
|
||||
|
||||
fn optional_expiry(value: Option<u64>) -> Option<Duration> {
|
||||
value
|
||||
.filter(|seconds| *seconds > 0)
|
||||
.map(Duration::from_secs)
|
||||
}
|
||||
|
||||
fn classify_refresh_failure(status: StatusCode, body: &[u8]) -> RefreshFailure {
|
||||
// 5xx 与 429 按状态码判为瞬态:其正文(CDN/网关页等)不是权威 OAuth 错误响应,
|
||||
// 不据此终态失效账号。
|
||||
if status.is_server_error() || status == StatusCode::TOO_MANY_REQUESTS {
|
||||
return RefreshFailure::Transport;
|
||||
// 官方刷新错误的消息与错误码分别位于 `error.message`、`error.code`;
|
||||
// `error` 字符串与顶层 `code` 仅用于兼容官方客户端自身的错误码提取契约。
|
||||
let error = serde_json::from_slice::<RefreshErrorResponse>(body).ok();
|
||||
let message = error.as_ref().and_then(RefreshErrorResponse::message);
|
||||
let upstream = || {
|
||||
Some(Box::new(RefreshUpstreamFailure::new(
|
||||
status,
|
||||
body,
|
||||
error.as_ref(),
|
||||
)))
|
||||
};
|
||||
let normalized_code = error
|
||||
.as_ref()
|
||||
.and_then(RefreshErrorResponse::code)
|
||||
.map(str::to_ascii_lowercase);
|
||||
// 生产策略故意与官方 Codex 的“任意 401 立即终态”不同:
|
||||
// 显式 401 先进入有界恢复退避,避免瞬时授权故障直接失效账号。
|
||||
if status == StatusCode::UNAUTHORIZED {
|
||||
return RefreshFailure::Transport {
|
||||
message,
|
||||
upstream: upstream(),
|
||||
};
|
||||
}
|
||||
// 4xx 正文才是权威 OAuth 错误响应,据其区分真过期/封禁与其余瞬态。
|
||||
let lower = String::from_utf8_lossy(body).to_ascii_lowercase();
|
||||
if lower.contains("account has been deactivated") {
|
||||
return RefreshFailure::Banned;
|
||||
// 非 401 响应仍与官方一致:三个明确的 RT 原因是永久失败。
|
||||
if matches!(
|
||||
normalized_code.as_deref(),
|
||||
Some("refresh_token_expired" | "refresh_token_reused" | "refresh_token_invalidated")
|
||||
) {
|
||||
return RefreshFailure::InvalidGrant {
|
||||
message,
|
||||
upstream: upstream(),
|
||||
};
|
||||
}
|
||||
if lower.contains("invalid_grant")
|
||||
|| lower.contains("invalid_token")
|
||||
|| lower.contains("access_denied")
|
||||
|| lower.contains("refresh_token_expired")
|
||||
|| lower.contains("refresh_token_reused")
|
||||
{
|
||||
return RefreshFailure::InvalidGrant;
|
||||
if message.as_deref().is_some_and(|message| {
|
||||
message
|
||||
.to_ascii_lowercase()
|
||||
.contains("account has been deactivated")
|
||||
}) {
|
||||
return RefreshFailure::Banned {
|
||||
message,
|
||||
upstream: upstream(),
|
||||
};
|
||||
}
|
||||
RefreshFailure::Transport {
|
||||
message,
|
||||
upstream: upstream(),
|
||||
}
|
||||
RefreshFailure::Transport
|
||||
}
|
||||
|
||||
fn refresh_transport_failure(error: &reqwest::Error) -> RefreshFailure {
|
||||
let message = error.to_string();
|
||||
if is_safe_to_retry_refresh_transport(error) {
|
||||
RefreshFailure::RetryableTransport
|
||||
RefreshFailure::RetryableTransport { message }
|
||||
} else {
|
||||
RefreshFailure::Transport
|
||||
RefreshFailure::Transport {
|
||||
message: Some(message),
|
||||
upstream: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -22,8 +22,7 @@ use crate::credential::{
|
||||
CodexOAuthAdmin, CodexOAuthAdminService, OfficialCodexAgentIdentityTaskRegistrar,
|
||||
};
|
||||
use crate::transport::profile::{
|
||||
CodexCliReleaseService, CodexDesktopReleaseService, OfficialCodexCliReleaseTransport,
|
||||
OfficialCodexDesktopReleaseTransport,
|
||||
CodexArtifactProfileCache, CodexDesktopReleaseService, OfficialCodexDesktopReleaseTransport,
|
||||
};
|
||||
use crate::transport::{CodexWebSocketPool, build_reqwest_client};
|
||||
|
||||
@@ -37,8 +36,8 @@ pub mod credential;
|
||||
pub mod transport;
|
||||
|
||||
pub use transport::{
|
||||
CodexCanonicalDecoder, CodexCanonicalError, CodexRequestEncodeError, encode_generate_request,
|
||||
openai_billing_breakdown,
|
||||
CodexCanonicalDecoder, CodexCanonicalError, CodexRequestEncodeError, OpenAiBillingUsage,
|
||||
encode_generate_request, openai_billing_breakdown,
|
||||
};
|
||||
|
||||
/// OpenAI 初始化后交给组装根的最小能力集。
|
||||
@@ -63,6 +62,23 @@ pub async fn initialize(
|
||||
let runtime_policy = ports.runtime_policy();
|
||||
let credential_state = ports.credential_state();
|
||||
let profile = config.wire_profile_state();
|
||||
let artifact_cache =
|
||||
CodexArtifactProfileCache::new(provider_kind.clone(), ports.artifact_profiles());
|
||||
let configured_build = profile.snapshot().desktop_build.parse::<u64>().ok();
|
||||
let verified_profile = match artifact_cache.load().await {
|
||||
Ok(cached) => cached.filter(|cached| {
|
||||
cached
|
||||
.desktop_build
|
||||
.parse::<u64>()
|
||||
.ok()
|
||||
.zip(configured_build)
|
||||
.is_some_and(|(cached, configured)| cached >= configured)
|
||||
}),
|
||||
Err(error) => {
|
||||
tracing::warn!(error = %error, "OpenAI artifact profile cache could not be loaded");
|
||||
None
|
||||
}
|
||||
};
|
||||
let session_identity = config
|
||||
.session_identity()
|
||||
.map_err(|_| OpenAiInitializeError::SessionIdentity)?;
|
||||
@@ -73,15 +89,10 @@ pub async fn initialize(
|
||||
OfficialCodexDesktopReleaseTransport::new()
|
||||
.map_err(|_| OpenAiInitializeError::DesktopRelease)?,
|
||||
),
|
||||
artifact_cache,
|
||||
verified_profile,
|
||||
));
|
||||
let desktop_release_status = desktop_release.status();
|
||||
let cli_release = Arc::new(CodexCliReleaseService::new(
|
||||
profile.clone(),
|
||||
Arc::new(
|
||||
OfficialCodexCliReleaseTransport::new()
|
||||
.map_err(|_| OpenAiInitializeError::CliRelease)?,
|
||||
),
|
||||
));
|
||||
let repository = CodexCredentialRepository::new(Arc::clone(&accounts));
|
||||
let websocket_pool = Arc::new(CodexWebSocketPool::with_config(
|
||||
config.websocket_pool_config(),
|
||||
@@ -89,7 +100,7 @@ pub async fn initialize(
|
||||
let agent_identity = Arc::new(CodexAgentIdentityTaskService::new(
|
||||
repository.clone(),
|
||||
Arc::new(
|
||||
OfficialCodexAgentIdentityTaskRegistrar::new(http.clone())
|
||||
OfficialCodexAgentIdentityTaskRegistrar::new(http.clone(), profile.clone())
|
||||
.map_err(|_| OpenAiInitializeError::AgentIdentity)?,
|
||||
),
|
||||
Arc::clone(&websocket_pool),
|
||||
@@ -165,6 +176,7 @@ pub async fn initialize(
|
||||
exchanger,
|
||||
Arc::clone(&accounts),
|
||||
CodexCredentialAdmin,
|
||||
profile.clone(),
|
||||
)
|
||||
.with_oauth_client_id(config.oauth_client_id()),
|
||||
);
|
||||
@@ -187,7 +199,6 @@ pub async fn initialize(
|
||||
catalog,
|
||||
config.quota_refresh_policy(),
|
||||
config.oauth_refresh_enabled(),
|
||||
cli_release,
|
||||
desktop_release,
|
||||
)
|
||||
.map_err(|_| OpenAiInitializeError::Worker)?;
|
||||
@@ -243,8 +254,6 @@ pub enum OpenAiInitializeError {
|
||||
Refresh,
|
||||
#[error("OpenAI Desktop release service could not initialize")]
|
||||
DesktopRelease,
|
||||
#[error("OpenAI CLI release service could not initialize")]
|
||||
CliRelease,
|
||||
#[error("OpenAI worker plan is invalid")]
|
||||
Worker,
|
||||
}
|
||||
|
||||
@@ -6,6 +6,7 @@ use std::sync::Arc;
|
||||
use std::time::{Duration, Instant, SystemTime};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use bytes::Bytes;
|
||||
use futures::{StreamExt, future::BoxFuture};
|
||||
use gateway_core::engine::continuation::{ContinuationBinding, NativeContinuationScope};
|
||||
use gateway_core::engine::credential::{AccountFeedbackStats, ProviderAccount};
|
||||
@@ -33,6 +34,9 @@ use gateway_core::task::{
|
||||
WorkerDefinitionError, WorkerId, WorkerKind, WorkerLeaseRequest, WorkerRegistration,
|
||||
WorkerRunnable, WorkerSchedule, WorkerTaskError,
|
||||
};
|
||||
use gateway_protocol::openai::events::{
|
||||
ParsedRateLimits, parse_rate_limit_headers, rate_limits_to_header_pairs,
|
||||
};
|
||||
use reqwest::Client;
|
||||
use secrecy::{ExposeSecret, SecretString};
|
||||
use serde::{Deserialize, Serialize};
|
||||
@@ -57,23 +61,22 @@ use crate::transport::diagnostics::{
|
||||
CodexFailureCategory, CodexUpstreamFailure, CodexUpstreamSendPhase,
|
||||
};
|
||||
use crate::transport::profile::{
|
||||
APPCAST_POLL_INTERVAL, CodexCliReleaseService, CodexDesktopReleaseService,
|
||||
CodexWireProfileState,
|
||||
APPCAST_POLL_INTERVAL, CodexDesktopReleaseService, CodexWireProfileState,
|
||||
};
|
||||
use crate::transport::protocol::responses::{
|
||||
CodexResponsesRequest, PreviousResponseScope, ResponseEventSignals,
|
||||
};
|
||||
use crate::transport::request::{
|
||||
CodexRequestEncodeError, encode_generate_request, scope_request_to_account,
|
||||
CodexRequestEncodeError, RequestAccountScope, encode_generate_request, scope_request_to_account,
|
||||
};
|
||||
use crate::transport::session::CodexSessionIdentity;
|
||||
use crate::transport::usage::normalize_service_tier;
|
||||
use crate::transport::websocket::{CodexWebSocketExchangeError, PreviousResponseUnavailableReason};
|
||||
use crate::transport::{
|
||||
CODEX_RESPONSES_PATH, CodexAccountSelectionTelemetry, CodexBackendClient,
|
||||
CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError,
|
||||
CodexRateLimitHeaderUpdates, CodexRequestContext, CodexResponseMetadata, CodexTransportMetrics,
|
||||
CodexUpstreamDiagnostics, CodexWebSocketPool, endpoint_url,
|
||||
CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError, CodexRateLimitUpdates,
|
||||
CodexRequestContext, CodexResponseMetadata, CodexTransportMetrics, CodexUpstreamDiagnostics,
|
||||
CodexWebSocketPool, endpoint_url,
|
||||
};
|
||||
|
||||
const PROVIDER_NAME: &str = "openai";
|
||||
@@ -248,7 +251,12 @@ impl Provider for CodexProvider {
|
||||
if let Some(identity) = &self.session_identity {
|
||||
identity.prepare_local_conversation(&mut upstream_request);
|
||||
}
|
||||
if upstream_request.turn_state.is_none() {
|
||||
if previous_session.as_ref().is_some_and(|state| {
|
||||
same_client_turn(
|
||||
state.client_turn_id.as_deref(),
|
||||
upstream_request.client_turn_id.as_deref(),
|
||||
)
|
||||
}) {
|
||||
upstream_request.turn_state = previous_session
|
||||
.as_ref()
|
||||
.and_then(|state| state.turn_state.clone());
|
||||
@@ -278,57 +286,48 @@ impl Provider for CodexProvider {
|
||||
let output_started_at = Instant::now();
|
||||
let provider_kind = ProviderKind::new(PROVIDER_NAME)
|
||||
.map_err(|_| provider_error(ProviderErrorKind::Protocol, UpstreamSendState::NotSent))?;
|
||||
let cross_account = context
|
||||
let state_owner_cross_account = context
|
||||
.account_state_owner()
|
||||
.is_some_and(|owner| !owner.matches(&provider_kind, lease.account_id()))
|
||||
|| previous_session
|
||||
.as_ref()
|
||||
.is_some_and(|state| state.account_id != lease.account_id().as_str());
|
||||
let replay_previous_response = matches!(
|
||||
context.continuation_attempt(),
|
||||
ContinuationAttempt::ReplayOwner | ContinuationAttempt::ReplayAny
|
||||
) || previous_session.as_ref().is_some_and(|state| {
|
||||
state.continuation_scope == OpenAiContinuationScope::ReplayRequired
|
||||
});
|
||||
if replay_previous_response && previous_session.is_some() {
|
||||
// Codex 客户端 input 是完整历史;恢复时只解除上游续链绑定,避免重复叠加代理副本。
|
||||
upstream_request.set_previous_response_id(None);
|
||||
upstream_request.previous_response_scope = None;
|
||||
upstream_request.turn_state = None;
|
||||
}
|
||||
scope_request_to_account(
|
||||
&mut upstream_request,
|
||||
lease.installation_id(),
|
||||
cross_account,
|
||||
);
|
||||
if context.continuation_attempt() == ContinuationAttempt::Native
|
||||
&& !replay_previous_response
|
||||
let has_explicit_state_owner =
|
||||
context.account_state_owner().is_some() || previous_session.is_some();
|
||||
let cross_account =
|
||||
state_owner_cross_account || (!has_explicit_state_owner && lease.account_switch());
|
||||
if context.continuation_attempt() != ContinuationAttempt::None
|
||||
&& let Some(continuation) = context.continuation()
|
||||
{
|
||||
match continuation {
|
||||
ContinuationBinding::Pinned(continuation) => {
|
||||
let previous_response_scope = match previous_session
|
||||
.as_ref()
|
||||
.map(|state| state.continuation_scope)
|
||||
{
|
||||
Some(OpenAiContinuationScope::Persisted) => {
|
||||
PreviousResponseScope::Persisted
|
||||
}
|
||||
Some(OpenAiContinuationScope::ConnectionLocal) => {
|
||||
PreviousResponseScope::ConnectionLocal
|
||||
}
|
||||
Some(OpenAiContinuationScope::ReplayRequired) => {
|
||||
return Err(provider_error(
|
||||
ProviderErrorKind::Protocol,
|
||||
UpstreamSendState::NotSent,
|
||||
));
|
||||
}
|
||||
None => match continuation.scope() {
|
||||
NativeContinuationScope::Persisted => PreviousResponseScope::Persisted,
|
||||
NativeContinuationScope::ConnectionLocal => {
|
||||
let previous_response_scope = match context.continuation_attempt() {
|
||||
ContinuationAttempt::Native => match previous_session
|
||||
.as_ref()
|
||||
.map(|state| state.continuation_scope)
|
||||
{
|
||||
Some(OpenAiContinuationScope::Persisted) => {
|
||||
PreviousResponseScope::Persisted
|
||||
}
|
||||
Some(OpenAiContinuationScope::ConnectionLocal) => {
|
||||
PreviousResponseScope::ConnectionLocal
|
||||
}
|
||||
Some(OpenAiContinuationScope::ReplayRequired) => {
|
||||
PreviousResponseScope::ExternalUnknown
|
||||
}
|
||||
None => match continuation.scope() {
|
||||
NativeContinuationScope::Persisted => {
|
||||
PreviousResponseScope::Persisted
|
||||
}
|
||||
NativeContinuationScope::ConnectionLocal => {
|
||||
PreviousResponseScope::ConnectionLocal
|
||||
}
|
||||
},
|
||||
},
|
||||
ContinuationAttempt::ReplayOwner | ContinuationAttempt::ReplayAny => {
|
||||
PreviousResponseScope::ExternalUnknown
|
||||
}
|
||||
ContinuationAttempt::None => PreviousResponseScope::ExternalUnknown,
|
||||
};
|
||||
upstream_request.set_previous_response_id(Some(
|
||||
continuation.upstream_response_id().as_str().to_owned(),
|
||||
@@ -343,6 +342,18 @@ impl Provider for CodexProvider {
|
||||
}
|
||||
}
|
||||
}
|
||||
let account_scope = if !cross_account {
|
||||
RequestAccountScope::SameAccount
|
||||
} else if upstream_request.previous_response_id().is_some() {
|
||||
RequestAccountScope::CrossAccountContinuationProbe
|
||||
} else {
|
||||
RequestAccountScope::CrossAccountFullRequest
|
||||
};
|
||||
scope_request_to_account(
|
||||
&mut upstream_request,
|
||||
lease.installation_id(),
|
||||
account_scope,
|
||||
);
|
||||
let metadata = ProviderCallMetadata::new(
|
||||
provider_kind,
|
||||
candidate.upstream_model().clone(),
|
||||
@@ -360,6 +371,7 @@ impl Provider for CodexProvider {
|
||||
account_id: lease.account_id().as_str().to_owned(),
|
||||
conversation_id: upstream_request.local_conversation_id.clone(),
|
||||
turn_state: upstream_request.turn_state.clone(),
|
||||
client_turn_id: upstream_request.client_turn_id.clone(),
|
||||
response_store,
|
||||
continuation_scope: None,
|
||||
});
|
||||
@@ -412,6 +424,8 @@ struct OpenAiSessionState {
|
||||
conversation_id: Option<String>,
|
||||
#[serde(default)]
|
||||
turn_state: Option<String>,
|
||||
#[serde(default)]
|
||||
client_turn_id: Option<String>,
|
||||
continuation_scope: OpenAiContinuationScope,
|
||||
}
|
||||
|
||||
@@ -427,10 +441,17 @@ struct OpenAiSessionCapture {
|
||||
account_id: String,
|
||||
conversation_id: Option<String>,
|
||||
turn_state: Option<String>,
|
||||
client_turn_id: Option<String>,
|
||||
response_store: bool,
|
||||
continuation_scope: Option<OpenAiContinuationScope>,
|
||||
}
|
||||
|
||||
fn same_client_turn(previous: Option<&str>, current: Option<&str>) -> bool {
|
||||
previous
|
||||
.zip(current)
|
||||
.is_some_and(|(previous, current)| !previous.is_empty() && previous == current)
|
||||
}
|
||||
|
||||
fn decode_openai_session_state(request: &GenerateRequest) -> Option<OpenAiSessionState> {
|
||||
request
|
||||
.provider_session_state(PROVIDER_NAME)
|
||||
@@ -472,6 +493,7 @@ fn attach_openai_session_update(
|
||||
account_id: capture.account_id,
|
||||
conversation_id: capture.conversation_id,
|
||||
turn_state: capture.turn_state,
|
||||
client_turn_id: capture.client_turn_id,
|
||||
continuation_scope,
|
||||
}) else {
|
||||
return;
|
||||
@@ -699,9 +721,6 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
if let Some(observation) = observation_state.observation() {
|
||||
yield ProviderEvent::observation(observation);
|
||||
}
|
||||
if allows_account_state_mutation {
|
||||
synchronize_passive_quota("a, &active_account, &response.rate_limit_headers).await;
|
||||
}
|
||||
if let Some(etag) = response.response_metadata.models_etag.as_deref()
|
||||
&& let Err(error) = catalog.observe_response_etag(etag)
|
||||
{
|
||||
@@ -729,7 +748,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
let failure_diagnostics = response.diagnostics.clone();
|
||||
let failure_set_cookie_headers = response.set_cookie_headers.clone();
|
||||
let failure_rate_limit_headers = response.rate_limit_headers.clone();
|
||||
let rate_limit_updates = response.rate_limit_header_updates;
|
||||
let mut passive_quota_observation =
|
||||
OpenAiPassiveQuotaObservation::new(response.rate_limit_headers);
|
||||
let rate_limit_updates = response.rate_limit_updates;
|
||||
let turn_state_updates = response.turn_state_update;
|
||||
// OpenAI 线路为透明代理:HTTP SSE 与 WebSocket 两条上游均启用 raw 透传,
|
||||
// 下游按字节转发上游原文,避免 serde 往返改写数值/精度(大整数→f64、logprobs 等)。
|
||||
@@ -737,10 +758,19 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
// (transport/protocol/websocket.rs),push_frames 抽出的 data 即上游原文。
|
||||
let mut decoder = CodexCanonicalDecoder::new(upstream_model.as_str())
|
||||
.with_requested_service_tier(request.service_tier())
|
||||
.with_request_tool_pricing(upstream_model.as_str(), request.tools())
|
||||
.with_raw_sse_passthrough();
|
||||
let mut pre_commit_events = PreCommitClientEvents::new();
|
||||
loop {
|
||||
let Some(stream_deadline) = remaining(context.deadline()) else {
|
||||
if allows_account_state_mutation {
|
||||
synchronize_passive_quota(
|
||||
"a,
|
||||
&active_account,
|
||||
passive_quota_observation.rate_limits(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Err(provider_error(ProviderErrorKind::Timeout, UpstreamSendState::Sent))?;
|
||||
return;
|
||||
};
|
||||
@@ -773,15 +803,22 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
Err(failure) => {
|
||||
let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await;
|
||||
if !updates.is_empty() {
|
||||
if allows_account_state_mutation {
|
||||
synchronize_passive_quota("a, &active_account, &updates).await;
|
||||
}
|
||||
if observation_state.merge_rate_limit_headers(&updates)
|
||||
passive_quota_observation.observe(&updates);
|
||||
let update_headers = rate_limit_update_headers(&updates);
|
||||
if observation_state.merge_rate_limit_headers(&update_headers)
|
||||
&& let Some(observation) = observation_state.observation()
|
||||
{
|
||||
yield ProviderEvent::observation(observation);
|
||||
}
|
||||
}
|
||||
if allows_account_state_mutation {
|
||||
synchronize_passive_quota(
|
||||
"a,
|
||||
&active_account,
|
||||
passive_quota_observation.rate_limits(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
apply_failure(&failure_context, &active_account, &failure)
|
||||
.await;
|
||||
Err(failure.error)?;
|
||||
@@ -793,18 +830,19 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
let rate_limits_changed = if updates.is_empty() {
|
||||
false
|
||||
} else {
|
||||
if allows_account_state_mutation {
|
||||
synchronize_passive_quota("a, &active_account, &updates).await;
|
||||
}
|
||||
observation_state.merge_rate_limit_headers(&updates)
|
||||
passive_quota_observation.observe(&updates);
|
||||
observation_state.merge_rate_limit_headers(&rate_limit_update_headers(&updates))
|
||||
};
|
||||
if let (Some(capture), Some(updates)) =
|
||||
(session_capture.as_mut(), turn_state_updates.as_ref())
|
||||
// take:仅在上游写入过新 turn_state 时更新 capture,避免逐 chunk 克隆。
|
||||
let turn_state_changed = if let Some(updates) = turn_state_updates.as_ref()
|
||||
&& let Some(turn_state) = updates.lock().await.take()
|
||||
{
|
||||
capture.turn_state = Some(turn_state);
|
||||
}
|
||||
if let Some(capture) = session_capture.as_mut() {
|
||||
capture.turn_state = Some(turn_state.clone());
|
||||
}
|
||||
observation_state.merge_client_header("x-codex-turn-state", &turn_state)
|
||||
} else {
|
||||
false
|
||||
};
|
||||
let first_event_changed =
|
||||
observation_state.observe_stream_chunk(&chunk, output_started_at);
|
||||
let chunk_len = chunk.len();
|
||||
@@ -843,6 +881,14 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
.any(|event| matches!(event, GatewayEvent::Completed(_)));
|
||||
let terminal_changed = completed
|
||||
&& observation_state.mark_completed(terminal_response_is_incomplete(&events));
|
||||
if allows_account_state_mutation && (completed || terminal_failure.is_some()) {
|
||||
synchronize_passive_quota(
|
||||
"a,
|
||||
&active_account,
|
||||
passive_quota_observation.rate_limits(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
if let Some((failure, _)) = terminal_failure.as_ref() {
|
||||
apply_failure(&failure_context, &active_account, failure)
|
||||
.await;
|
||||
@@ -861,6 +907,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
if (rate_limits_changed
|
||||
|| service_tier_changed
|
||||
|| timing_changed
|
||||
|| turn_state_changed
|
||||
|| terminal_changed)
|
||||
&& let Some(observation) = observation_state.observation()
|
||||
{
|
||||
@@ -920,21 +967,31 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
let rate_limits_changed = if updates.is_empty() {
|
||||
false
|
||||
} else {
|
||||
if allows_account_state_mutation {
|
||||
synchronize_passive_quota("a, &active_account, &updates).await;
|
||||
}
|
||||
observation_state.merge_rate_limit_headers(&updates)
|
||||
passive_quota_observation.observe(&updates);
|
||||
observation_state.merge_rate_limit_headers(&rate_limit_update_headers(&updates))
|
||||
};
|
||||
if allows_account_state_mutation {
|
||||
synchronize_passive_quota(
|
||||
"a,
|
||||
&active_account,
|
||||
passive_quota_observation.rate_limits(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
if let Some((failure, _)) = terminal_failure.as_ref() {
|
||||
apply_failure(&failure_context, &active_account, failure)
|
||||
.await;
|
||||
}
|
||||
if let (Some(capture), Some(updates)) =
|
||||
(session_capture.as_mut(), turn_state_updates.as_ref())
|
||||
let turn_state_changed = if let Some(updates) = turn_state_updates.as_ref()
|
||||
&& let Some(turn_state) = updates.lock().await.take()
|
||||
{
|
||||
capture.turn_state = Some(turn_state);
|
||||
}
|
||||
if let Some(capture) = session_capture.as_mut() {
|
||||
capture.turn_state = Some(turn_state.clone());
|
||||
}
|
||||
observation_state.merge_client_header("x-codex-turn-state", &turn_state)
|
||||
} else {
|
||||
false
|
||||
};
|
||||
attach_openai_session_update(&mut events, &mut session_capture);
|
||||
let completed = events
|
||||
.iter()
|
||||
@@ -951,7 +1008,11 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
.observe_cyber_policy_success(cyber_policy_scope.as_ref())
|
||||
.await;
|
||||
}
|
||||
if (service_tier_changed || timing_changed || rate_limits_changed || terminal_changed)
|
||||
if (service_tier_changed
|
||||
|| timing_changed
|
||||
|| rate_limits_changed
|
||||
|| turn_state_changed
|
||||
|| terminal_changed)
|
||||
&& let Some(observation) = observation_state.observation()
|
||||
{
|
||||
yield ProviderEvent::observation(observation);
|
||||
@@ -1001,6 +1062,26 @@ struct OpenAiResponseObservationState {
|
||||
terminal: Option<OpenAiResponseTerminal>,
|
||||
}
|
||||
|
||||
struct OpenAiPassiveQuotaObservation {
|
||||
rate_limits: Vec<ParsedRateLimits>,
|
||||
}
|
||||
|
||||
impl OpenAiPassiveQuotaObservation {
|
||||
fn new(headers: Vec<(String, String)>) -> Self {
|
||||
Self {
|
||||
rate_limits: parse_rate_limit_headers(&headers).into_iter().collect(),
|
||||
}
|
||||
}
|
||||
|
||||
fn observe(&mut self, updates: &[ParsedRateLimits]) {
|
||||
self.rate_limits.extend_from_slice(updates);
|
||||
}
|
||||
|
||||
fn rate_limits(&self) -> &[ParsedRateLimits] {
|
||||
&self.rate_limits
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
enum OpenAiResponseTerminal {
|
||||
Completed,
|
||||
@@ -1120,6 +1201,26 @@ impl OpenAiResponseObservationState {
|
||||
changed
|
||||
}
|
||||
|
||||
fn merge_client_header(&mut self, name: &str, value: &str) -> bool {
|
||||
let value = Bytes::copy_from_slice(value.as_bytes());
|
||||
if let Some((_, existing)) = self
|
||||
.response_metadata
|
||||
.client_headers
|
||||
.iter_mut()
|
||||
.find(|(existing_name, _)| existing_name.eq_ignore_ascii_case(name))
|
||||
{
|
||||
if *existing == value {
|
||||
return false;
|
||||
}
|
||||
*existing = value;
|
||||
return true;
|
||||
}
|
||||
self.response_metadata
|
||||
.client_headers
|
||||
.push((name.to_owned(), value));
|
||||
true
|
||||
}
|
||||
|
||||
fn mark_completed(&mut self, incomplete: bool) -> bool {
|
||||
let terminal = if incomplete {
|
||||
OpenAiResponseTerminal::Incomplete
|
||||
@@ -1432,15 +1533,20 @@ fn insert_optional_metadata_millis(
|
||||
}
|
||||
}
|
||||
|
||||
async fn take_rate_limit_updates(
|
||||
updates: Option<&CodexRateLimitHeaderUpdates>,
|
||||
) -> Vec<(String, String)> {
|
||||
async fn take_rate_limit_updates(updates: Option<&CodexRateLimitUpdates>) -> Vec<ParsedRateLimits> {
|
||||
let Some(updates) = updates else {
|
||||
return Vec::new();
|
||||
};
|
||||
std::mem::take(&mut *updates.lock().await)
|
||||
}
|
||||
|
||||
fn rate_limit_update_headers(updates: &[ParsedRateLimits]) -> Vec<(String, String)> {
|
||||
updates
|
||||
.iter()
|
||||
.flat_map(rate_limits_to_header_pairs)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn terminal_response_is_incomplete(events: &[ProviderEvent]) -> bool {
|
||||
events
|
||||
.iter()
|
||||
@@ -1466,12 +1572,15 @@ fn terminal_response_is_incomplete(events: &[ProviderEvent]) -> bool {
|
||||
async fn synchronize_passive_quota(
|
||||
quota: &CodexCredentialQuotaService,
|
||||
account: &ProviderAccount,
|
||||
headers: &[(String, String)],
|
||||
rate_limits: &[ParsedRateLimits],
|
||||
) {
|
||||
if headers.is_empty() {
|
||||
if rate_limits.is_empty() {
|
||||
return;
|
||||
}
|
||||
if let Err(error) = quota.synchronize_passive_headers(account, headers).await {
|
||||
if let Err(error) = quota
|
||||
.synchronize_passive_rate_limits(account, rate_limits)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
account_id = %account.id(),
|
||||
error = %error,
|
||||
@@ -1480,6 +1589,17 @@ async fn synchronize_passive_quota(
|
||||
}
|
||||
}
|
||||
|
||||
async fn synchronize_passive_quota_headers(
|
||||
quota: &CodexCredentialQuotaService,
|
||||
account: &ProviderAccount,
|
||||
headers: &[(String, String)],
|
||||
) {
|
||||
let Some(rate_limits) = parse_rate_limit_headers(headers) else {
|
||||
return;
|
||||
};
|
||||
synchronize_passive_quota(quota, account, std::slice::from_ref(&rate_limits)).await;
|
||||
}
|
||||
|
||||
const fn actual_transport_name(transport: CodexBackendTransport) -> &'static str {
|
||||
match transport {
|
||||
CodexBackendTransport::HttpSse => HTTP_SSE_TRANSPORT,
|
||||
@@ -1492,12 +1612,10 @@ fn nonnegative_millis(value: Option<i64>) -> Option<u64> {
|
||||
}
|
||||
|
||||
fn compile_model_capabilities(model: &CodexCatalogModel) -> ProviderModelCapabilities {
|
||||
let evidence = model.capabilities();
|
||||
let mut operations = BTreeSet::new();
|
||||
if evidence.responses_api() == CodexCatalogCapabilityEvidence::DeclaredNative {
|
||||
operations.insert(OperationKind::Generate);
|
||||
}
|
||||
let capabilities = ModelCapabilities::new(operations, None).with_upstream_feature_validation();
|
||||
// Catalog membership is enough to publish Generate. `supported_in_api` is advisory;
|
||||
// the upstream response is authoritative for normal and diagnostic requests.
|
||||
let capabilities = ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), None)
|
||||
.with_upstream_feature_validation();
|
||||
ProviderModelCapabilities::new(model.request_model().clone(), capabilities)
|
||||
.with_presentation(codex_model_presentation(model))
|
||||
}
|
||||
@@ -1875,7 +1993,7 @@ async fn apply_failure(
|
||||
if !context.allows_account_state_mutation {
|
||||
return;
|
||||
}
|
||||
synchronize_passive_quota(context.quota, account, &failure.rate_limit_headers).await;
|
||||
synchronize_passive_quota_headers(context.quota, account, &failure.rate_limit_headers).await;
|
||||
let needs_authoritative_quota_refresh = matches!(
|
||||
failure.account_failure,
|
||||
Some(CodexAccountFailure::QuotaExhausted | CodexAccountFailure::UsageLimitExhausted { .. })
|
||||
@@ -2154,7 +2272,7 @@ fn map_upstream_failure(
|
||||
error = error.with_status(status.as_u16());
|
||||
}
|
||||
if replay_boundary.permits_provider_proof()
|
||||
&& (failure.replay_is_safe() || continuation_failure.is_some() || cyber_policy_failure)
|
||||
&& (failure.replay_is_safe() || cyber_policy_failure)
|
||||
{
|
||||
error = error.with_replay_safe();
|
||||
}
|
||||
@@ -2339,7 +2457,6 @@ const WORKER_LEASE_TTL: Duration = Duration::from_secs(15 * 60);
|
||||
const WORKER_LEASE_RENEWAL: Duration = Duration::from_secs(5 * 60);
|
||||
const OAUTH_REFRESH_INTERVAL: Duration = Duration::from_secs(30);
|
||||
const DESKTOP_RELEASE_WORKER_OWNER: &str = "openai-desktop-release";
|
||||
const CLI_RELEASE_WORKER_OWNER: &str = "openai-cli-release";
|
||||
const MODEL_ETAG_WORKER_OWNER: &str = "openai-model-etag";
|
||||
|
||||
pub(crate) fn worker_contributions(
|
||||
@@ -2348,7 +2465,6 @@ pub(crate) fn worker_contributions(
|
||||
catalog: Arc<CodexCredentialCatalogService>,
|
||||
quota_refresh_policy: CodexQuotaRefreshPolicy,
|
||||
oauth_refresh_enabled: bool,
|
||||
cli_release: Arc<CodexCliReleaseService>,
|
||||
desktop_release: Arc<CodexDesktopReleaseService>,
|
||||
) -> Result<Vec<WorkerContribution>, WorkerDefinitionError> {
|
||||
let refresh_id = WorkerId::try_new(WorkerKind::OAuthRefresh, PROVIDER_NAME)?;
|
||||
@@ -2356,8 +2472,6 @@ pub(crate) fn worker_contributions(
|
||||
let etag_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, MODEL_ETAG_WORKER_OWNER)?;
|
||||
let desktop_release_id =
|
||||
WorkerId::try_new(WorkerKind::QuotaCatalogHealth, DESKTOP_RELEASE_WORKER_OWNER)?;
|
||||
let cli_release_id =
|
||||
WorkerId::try_new(WorkerKind::QuotaCatalogHealth, CLI_RELEASE_WORKER_OWNER)?;
|
||||
let mut contributions = Vec::new();
|
||||
if oauth_refresh_enabled {
|
||||
contributions.push(WorkerContribution::Registration(scheduled_registration(
|
||||
@@ -2385,13 +2499,6 @@ pub(crate) fn worker_contributions(
|
||||
task: Box::new(OpenAiCatalogEtagTask { catalog }),
|
||||
},
|
||||
)?),
|
||||
WorkerContribution::Registration(scheduled_registration(
|
||||
cli_release_id,
|
||||
APPCAST_POLL_INTERVAL,
|
||||
Box::new(OpenAiCliReleaseTask {
|
||||
service: cli_release,
|
||||
}),
|
||||
)?),
|
||||
WorkerContribution::Registration(scheduled_registration(
|
||||
desktop_release_id,
|
||||
APPCAST_POLL_INTERVAL,
|
||||
@@ -2512,27 +2619,6 @@ struct OpenAiDesktopReleaseTask {
|
||||
service: Arc<CodexDesktopReleaseService>,
|
||||
}
|
||||
|
||||
struct OpenAiCliReleaseTask {
|
||||
service: Arc<CodexCliReleaseService>,
|
||||
}
|
||||
|
||||
impl ScheduledTask for OpenAiCliReleaseTask {
|
||||
fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
let refresh = self.service.refresh();
|
||||
tokio::pin!(refresh);
|
||||
let result = tokio::select! {
|
||||
() = context.cancellation().cancelled() => return Ok(()),
|
||||
result = &mut refresh => result,
|
||||
};
|
||||
if let Err(error) = result {
|
||||
tracing::warn!(error = %error, "OpenAI CLI release check failed");
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ScheduledTask for OpenAiDesktopReleaseTask {
|
||||
fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
|
||||
@@ -21,7 +21,10 @@ use thiserror::Error;
|
||||
use super::protocol::responses::{
|
||||
ResponseEventSignals, ResponsesSseFailure, response_event_signals,
|
||||
};
|
||||
use super::usage::{normalize_service_tier, openai_billing_breakdown};
|
||||
use super::usage::{
|
||||
OpenAiBillingUsage, WebSearchPricing, normalize_service_tier, openai_billing_breakdown,
|
||||
web_search_pricing,
|
||||
};
|
||||
|
||||
const CONTENTS_PER_OUTPUT: u32 = 1_024;
|
||||
|
||||
@@ -43,6 +46,7 @@ pub struct CodexCanonicalDecoder {
|
||||
semantic_output_seen: bool,
|
||||
requested_service_tier: Option<String>,
|
||||
response_service_tier: Option<String>,
|
||||
web_search_pricing: Option<WebSearchPricing>,
|
||||
timing_signals: ResponseEventSignals,
|
||||
raw_sse_passthrough: bool,
|
||||
}
|
||||
@@ -139,6 +143,7 @@ impl CodexCanonicalDecoder {
|
||||
semantic_output_seen: false,
|
||||
requested_service_tier: None,
|
||||
response_service_tier: None,
|
||||
web_search_pricing: None,
|
||||
timing_signals: ResponseEventSignals::default(),
|
||||
raw_sse_passthrough: false,
|
||||
}
|
||||
@@ -160,6 +165,13 @@ impl CodexCanonicalDecoder {
|
||||
self
|
||||
}
|
||||
|
||||
/// 记录请求工具类型,用于区分标准与预览 Web Search 的按次价格。
|
||||
#[must_use]
|
||||
pub fn with_request_tool_pricing(mut self, model: &str, tools: Option<&[Value]>) -> Self {
|
||||
self.web_search_pricing = web_search_pricing(model, tools);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn push(&mut self, chunk: &[u8]) -> CodexCanonicalOutcome {
|
||||
self.timing_signals = ResponseEventSignals::default();
|
||||
if self.raw_sse_passthrough {
|
||||
@@ -849,15 +861,15 @@ impl CodexCanonicalDecoder {
|
||||
.requested_service_tier
|
||||
.as_deref()
|
||||
.or(self.response_service_tier.as_deref());
|
||||
let web_search_calls = web_search_call_count(response);
|
||||
if let Some(breakdown) = usage
|
||||
.filter(|usage| billable_usage_is_complete(response, *usage))
|
||||
.and_then(|usage| {
|
||||
let web_search_calls = web_search_calls?;
|
||||
openai_billing_breakdown(
|
||||
&model,
|
||||
usage.input_tokens,
|
||||
usage.output_tokens,
|
||||
usage.cached_tokens,
|
||||
usage.cache_write_tokens,
|
||||
OpenAiBillingUsage::from(usage)
|
||||
.with_web_search_calls(web_search_calls, self.web_search_pricing),
|
||||
service_tier,
|
||||
)
|
||||
})
|
||||
@@ -984,6 +996,28 @@ fn core_usage(usage: TokenUsage) -> Usage {
|
||||
normalized
|
||||
}
|
||||
|
||||
fn web_search_call_count(response: &Value) -> Option<u64> {
|
||||
let mut calls = 0_u64;
|
||||
for item in response
|
||||
.get("output")
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter(|item| item.get("type").and_then(Value::as_str) == Some("web_search_call"))
|
||||
{
|
||||
match item
|
||||
.get("action")
|
||||
.and_then(|action| action.get("type"))
|
||||
.and_then(Value::as_str)
|
||||
{
|
||||
Some("search") => calls = calls.checked_add(1)?,
|
||||
Some("open_page" | "find_in_page") => {}
|
||||
Some(_) | None => return None,
|
||||
}
|
||||
}
|
||||
Some(calls)
|
||||
}
|
||||
|
||||
fn billable_usage_is_complete(response: &Value, usage: TokenUsage) -> bool {
|
||||
let Some(raw) = response.get("usage").filter(|value| value.is_object()) else {
|
||||
return false;
|
||||
|
||||
@@ -30,7 +30,7 @@ use super::response_meta::CodexResponseMetadata;
|
||||
use super::tls::{CustomCaError, build_reqwest_client_with_custom_ca, custom_ca_env_cache_key};
|
||||
use super::websocket::{
|
||||
CodexWebSocketExchangeError, CodexWebSocketPool, CodexWebSocketPoolKey,
|
||||
CodexWebSocketRateLimitHeaderUpdates, CodexWebSocketRequest, CodexWebSocketTurnStateUpdate,
|
||||
CodexWebSocketRateLimitUpdates, CodexWebSocketRequest, CodexWebSocketTurnStateUpdate,
|
||||
PreparedWebSocket, WebSocketOriginBreaker, WebSocketPoolDecision,
|
||||
};
|
||||
|
||||
@@ -519,8 +519,8 @@ pub struct CodexTransportMetrics {
|
||||
pub http_version: Option<String>,
|
||||
}
|
||||
|
||||
/// 响应头之后在 live 流中采集的限流更新。
|
||||
pub type CodexRateLimitHeaderUpdates = CodexWebSocketRateLimitHeaderUpdates;
|
||||
/// 响应头之后在 live 流中采集的结构化限流更新。
|
||||
pub type CodexRateLimitUpdates = CodexWebSocketRateLimitUpdates;
|
||||
|
||||
/// 响应头之后在 live 流中采集的 turn state 更新。
|
||||
pub type CodexTurnStateUpdate = CodexWebSocketTurnStateUpdate;
|
||||
@@ -537,8 +537,8 @@ pub struct CodexBackendStreamingResponse {
|
||||
pub set_cookie_headers: Vec<String>,
|
||||
/// 上游透传的限流头。
|
||||
pub rate_limit_headers: Vec<(String, String)>,
|
||||
/// live stream 期间捕获的限流头更新。
|
||||
pub rate_limit_header_updates: Option<CodexRateLimitHeaderUpdates>,
|
||||
/// live stream 期间捕获的结构化限流更新。
|
||||
pub rate_limit_updates: Option<CodexRateLimitUpdates>,
|
||||
/// live stream 期间捕获的 turn-state 更新。
|
||||
pub turn_state_update: Option<CodexTurnStateUpdate>,
|
||||
/// WebSocket 连接池决策。
|
||||
|
||||
@@ -147,14 +147,14 @@ impl CodexBackendClient {
|
||||
});
|
||||
}
|
||||
|
||||
let rate_limit_header_updates = Arc::new(tokio::sync::Mutex::new(Vec::new()));
|
||||
let rate_limit_updates = Arc::new(tokio::sync::Mutex::new(Vec::new()));
|
||||
Ok(CodexBackendStreamingResponse {
|
||||
body: http_sse_stream(response, Arc::clone(&rate_limit_header_updates)),
|
||||
body: http_sse_stream(response, Arc::clone(&rate_limit_updates)),
|
||||
transport: CodexBackendTransport::HttpSse,
|
||||
turn_state,
|
||||
set_cookie_headers,
|
||||
rate_limit_headers,
|
||||
rate_limit_header_updates: Some(rate_limit_header_updates),
|
||||
rate_limit_updates: Some(rate_limit_updates),
|
||||
turn_state_update: None,
|
||||
websocket_pool_decision: None,
|
||||
diagnostics,
|
||||
@@ -245,11 +245,12 @@ impl CodexBackendClient {
|
||||
self.websocket_pool.as_deref().zip(pool_key)
|
||||
};
|
||||
let fast_path_budget = match requirement {
|
||||
TransportRequirement::PersistedContinuation
|
||||
| TransportRequirement::ExternalUnknown
|
||||
| TransportRequirement::NewChain => Some(WEBSOCKET_FAST_PATH_BUDGET),
|
||||
TransportRequirement::PersistedContinuation | TransportRequirement::NewChain => {
|
||||
Some(WEBSOCKET_FAST_PATH_BUDGET)
|
||||
}
|
||||
TransportRequirement::ExplicitWebSocketWarmup
|
||||
| TransportRequirement::ExactWebSocketContinuation => None,
|
||||
| TransportRequirement::ExactWebSocketContinuation
|
||||
| TransportRequirement::ExternalUnknown => None,
|
||||
TransportRequirement::HttpRequired => None,
|
||||
};
|
||||
let prepare_started_at = Instant::now();
|
||||
@@ -515,7 +516,7 @@ impl CodexBackendClient {
|
||||
turn_state: exchange.turn_state,
|
||||
set_cookie_headers: exchange.set_cookie_headers,
|
||||
rate_limit_headers: exchange.rate_limit_headers,
|
||||
rate_limit_header_updates: Some(exchange.rate_limit_header_updates),
|
||||
rate_limit_updates: Some(exchange.rate_limit_updates),
|
||||
turn_state_update: Some(exchange.turn_state_update),
|
||||
websocket_pool_decision: exchange.pool_decision,
|
||||
diagnostics: exchange.diagnostics,
|
||||
@@ -677,7 +678,13 @@ impl CodexBackendClient {
|
||||
// API 透传黑名单也不能覆盖画像,避免下游客户端暴露不一致指纹。
|
||||
if matches!(
|
||||
name.as_str(),
|
||||
"originator" | "user-agent" | "x-oai-attestation" | "x-oai-is" | "x-oai-is-update"
|
||||
"openai-beta"
|
||||
| "originator"
|
||||
| "user-agent"
|
||||
| "x-oai-attestation"
|
||||
| "x-oai-is"
|
||||
| "x-oai-is-update"
|
||||
| "x-openai-internal-codex-residency"
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
@@ -696,6 +703,10 @@ impl CodexBackendClient {
|
||||
) -> CodexClientResult<HeaderMap> {
|
||||
let mut headers = self.request_headers_for_http_response(request, context)?;
|
||||
headers.remove(X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER);
|
||||
headers.insert(
|
||||
HeaderName::from_static("openai-beta"),
|
||||
HeaderValue::from_static("responses_websockets=2026-02-06"),
|
||||
);
|
||||
Ok(headers)
|
||||
}
|
||||
|
||||
@@ -709,14 +720,6 @@ impl CodexBackendClient {
|
||||
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
|
||||
insert_optional_header(&mut headers, "cookie", context.cookie_header)?;
|
||||
headers.insert(ACCEPT, HeaderValue::from_static("text/event-stream"));
|
||||
headers.insert(
|
||||
HeaderName::from_static("openai-beta"),
|
||||
HeaderValue::from_static("responses_websockets=2026-02-06"),
|
||||
);
|
||||
headers.insert(
|
||||
HeaderName::from_static("x-openai-internal-codex-residency"),
|
||||
HeaderValue::from_static("us"),
|
||||
);
|
||||
headers.insert(
|
||||
HeaderName::from_static("x-client-request-id"),
|
||||
HeaderValue::from_str(context.request_id)?,
|
||||
@@ -788,10 +791,6 @@ impl CodexBackendClient {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(USER_AGENT, HeaderValue::from_str(&profile.user_agent())?);
|
||||
headers.insert(AUTHORIZATION, HeaderValue::from_str(context.authorization)?);
|
||||
headers.insert(
|
||||
HeaderName::from_static("originator"),
|
||||
HeaderValue::from_str(&profile.originator)?,
|
||||
);
|
||||
insert_optional_header(&mut headers, "chatgpt-account-id", context.account_id)?;
|
||||
insert_optional_header(&mut headers, "cookie", context.cookie_header)?;
|
||||
headers.insert(ACCEPT, HeaderValue::from_static("application/json"));
|
||||
@@ -880,7 +879,7 @@ fn websocket_connection_profile(headers: &HeaderMap) -> String {
|
||||
|
||||
fn http_sse_stream(
|
||||
response: ReqwestResponse,
|
||||
rate_limit_updates: CodexRateLimitHeaderUpdates,
|
||||
rate_limit_updates: CodexRateLimitUpdates,
|
||||
) -> CodexBackendSseStream {
|
||||
let stream: CodexBackendSseStream =
|
||||
Box::pin(response.bytes_stream().map_err(CodexClientError::Http));
|
||||
@@ -903,7 +902,7 @@ fn http_sse_stream(
|
||||
|
||||
fn observe_http_sse_rate_limits(
|
||||
stream: CodexBackendSseStream,
|
||||
updates: CodexRateLimitHeaderUpdates,
|
||||
updates: CodexRateLimitUpdates,
|
||||
) -> CodexBackendSseStream {
|
||||
Box::pin(futures::stream::unfold(
|
||||
(stream, SseEventDecoder::default(), updates),
|
||||
@@ -927,9 +926,9 @@ fn observe_http_sse_rate_limits(
|
||||
|
||||
async fn append_http_sse_rate_limit_updates(
|
||||
frames: Vec<SseFrame>,
|
||||
updates: &CodexRateLimitHeaderUpdates,
|
||||
updates: &CodexRateLimitUpdates,
|
||||
) {
|
||||
let mut headers = Vec::new();
|
||||
let mut observations = Vec::new();
|
||||
for frame in frames {
|
||||
for event in frame.events() {
|
||||
if event
|
||||
@@ -942,10 +941,10 @@ async fn append_http_sse_rate_limit_updates(
|
||||
let Some(rate_limits) = events::parse_rate_limits_event_raw(&event.data) else {
|
||||
continue;
|
||||
};
|
||||
headers.extend(events::rate_limits_to_header_pairs(&rate_limits));
|
||||
observations.push(rate_limits);
|
||||
}
|
||||
}
|
||||
if !headers.is_empty() {
|
||||
updates.lock().await.extend(headers);
|
||||
if !observations.is_empty() {
|
||||
updates.lock().await.extend(observations);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,7 +1,5 @@
|
||||
/// `/codex/responses`
|
||||
pub const CODEX_RESPONSES_PATH: &str = "/codex/responses";
|
||||
/// `/codex/usage`
|
||||
pub const CODEX_USAGE_PATH: &str = "/codex/usage";
|
||||
/// `/api/codex/usage`
|
||||
pub const CODEX_USAGE_API_PATH: &str = "/api/codex/usage";
|
||||
/// `/wham/usage`
|
||||
@@ -16,20 +14,14 @@ pub fn endpoint_url(base_url: &str, endpoint_path: &str) -> String {
|
||||
)
|
||||
}
|
||||
|
||||
/// 返回 usage 相关 endpoint 的完整 URL 列表。
|
||||
pub fn usage_endpoint_urls(base_url: &str) -> Vec<String> {
|
||||
usage_endpoint_paths(base_url)
|
||||
.into_iter()
|
||||
.map(|path| endpoint_url(base_url, path))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn usage_endpoint_paths(base_url: &str) -> Vec<&'static str> {
|
||||
if has_backend_api_base_path(base_url) {
|
||||
vec![WHAM_USAGE_PATH, CODEX_USAGE_PATH]
|
||||
/// 返回与 base path 对应的唯一 usage endpoint。
|
||||
pub fn usage_endpoint_url(base_url: &str) -> String {
|
||||
let path = if has_backend_api_base_path(base_url) {
|
||||
WHAM_USAGE_PATH
|
||||
} else {
|
||||
vec![CODEX_USAGE_API_PATH, CODEX_USAGE_PATH]
|
||||
}
|
||||
CODEX_USAGE_API_PATH
|
||||
};
|
||||
endpoint_url(base_url, path)
|
||||
}
|
||||
|
||||
fn has_backend_api_base_path(base_url: &str) -> bool {
|
||||
|
||||
@@ -3,20 +3,33 @@ use reqwest::header::{AUTHORIZATION, HeaderMap, HeaderName, HeaderValue, USER_AG
|
||||
|
||||
use super::client::CodexClientResult;
|
||||
|
||||
const CODEX_RESIDENCY_HEADER: &str = "x-openai-internal-codex-residency";
|
||||
|
||||
/// 构造 Codex Core 为模型请求设置的稳定身份请求头。
|
||||
pub fn build_codex_base_headers(
|
||||
profile: &CodexWireProfile,
|
||||
authorization: &str,
|
||||
account_id: Option<&str>,
|
||||
) -> CodexClientResult<HeaderMap> {
|
||||
let mut headers = HeaderMap::new();
|
||||
let mut headers = build_codex_profile_headers(profile)?;
|
||||
headers.insert(AUTHORIZATION, HeaderValue::from_str(authorization)?);
|
||||
insert_optional_header(&mut headers, "chatgpt-account-id", account_id)?;
|
||||
Ok(headers)
|
||||
}
|
||||
|
||||
/// 构造不依赖访问凭据的 Codex Desktop 身份请求头。
|
||||
pub fn build_codex_profile_headers(profile: &CodexWireProfile) -> CodexClientResult<HeaderMap> {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(
|
||||
HeaderName::from_static("originator"),
|
||||
HeaderValue::from_str(&profile.originator)?,
|
||||
);
|
||||
headers.insert(USER_AGENT, HeaderValue::from_str(&profile.user_agent())?);
|
||||
// 此部署固定采用官方支持的 US residency requirement。
|
||||
headers.insert(
|
||||
HeaderName::from_static(CODEX_RESIDENCY_HEADER),
|
||||
HeaderValue::from_static("us"),
|
||||
);
|
||||
Ok(headers)
|
||||
}
|
||||
|
||||
|
||||
@@ -27,18 +27,18 @@ pub use self::{
|
||||
client::{
|
||||
CodexAccountSelectionTelemetry, CodexBackendClient, CodexBackendSseStream,
|
||||
CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError, CodexClientResult,
|
||||
CodexRateLimitHeaderUpdates, CodexRequestContext, CodexTransportDecision,
|
||||
CodexTransportMetrics, CodexTurnStateUpdate, build_reqwest_client,
|
||||
CodexRateLimitUpdates, CodexRequestContext, CodexTransportDecision, CodexTransportMetrics,
|
||||
CodexTurnStateUpdate, build_reqwest_client,
|
||||
},
|
||||
diagnostics::{CodexUpstreamDiagnostics, CodexUpstreamSendPhase},
|
||||
endpoints::{
|
||||
CODEX_RESPONSES_PATH, CODEX_USAGE_API_PATH, CODEX_USAGE_PATH, WHAM_USAGE_PATH,
|
||||
endpoint_url, usage_endpoint_urls,
|
||||
CODEX_RESPONSES_PATH, CODEX_USAGE_API_PATH, WHAM_USAGE_PATH, endpoint_url,
|
||||
usage_endpoint_url,
|
||||
},
|
||||
headers::build_codex_base_headers,
|
||||
request::{CodexRequestEncodeError, encode_generate_request},
|
||||
response_meta::CodexResponseMetadata,
|
||||
usage::{MAX_CODEX_USAGE_BODY_BYTES, openai_billing_breakdown},
|
||||
usage::{MAX_CODEX_USAGE_BODY_BYTES, OpenAiBillingUsage, openai_billing_breakdown},
|
||||
websocket::{
|
||||
CodexWebSocketPool, CodexWebSocketPoolConfig, CodexWebSocketPoolKey, WebSocketPoolDecision,
|
||||
},
|
||||
|
||||
@@ -0,0 +1,600 @@
|
||||
//! 有界读取 Codex Desktop ZIP 内嵌的 arm64 Core 版本。
|
||||
|
||||
use std::cmp;
|
||||
|
||||
use flate2::{Decompress, FlushDecompress, Status};
|
||||
use futures::StreamExt as _;
|
||||
use reqwest::header::{CONTENT_LENGTH, CONTENT_RANGE, RANGE};
|
||||
use reqwest::{Client, StatusCode};
|
||||
use url::Url;
|
||||
|
||||
const OFFICIAL_ARTIFACT_HOST: &str = "persistent.oaistatic.com";
|
||||
const EOCD_SIGNATURE: &[u8; 4] = b"PK\x05\x06";
|
||||
const CENTRAL_SIGNATURE: &[u8; 4] = b"PK\x01\x02";
|
||||
const LOCAL_SIGNATURE: &[u8; 4] = b"PK\x03\x04";
|
||||
const CORE_PATH_SUFFIX: &[u8] = b"/Contents/Resources/codex";
|
||||
const CORE_VERSION_ANCHOR: &[u8] = b"codex-mcp-client/";
|
||||
const MAX_EOCD_SEARCH_BYTES: u64 = 65_535 + 22;
|
||||
const MAX_CENTRAL_DIRECTORY_BYTES: u64 = 16 * 1024 * 1024;
|
||||
const MAX_RANGE_BYTES: u64 = 16 * 1024 * 1024;
|
||||
const CORE_RANGE_CHUNK_BYTES: u64 = 4 * 1024 * 1024;
|
||||
const MAX_CORE_COMPRESSED_BYTES: u64 = 256 * 1024 * 1024;
|
||||
const MAX_CORE_UNCOMPRESSED_BYTES: u64 = 512 * 1024 * 1024;
|
||||
const MAX_VERSION_BYTES: usize = 64;
|
||||
const MACHO_HEADER_BYTES: usize = 8;
|
||||
const MACHO_64_LE_MAGIC: [u8; 4] = [0xcf, 0xfa, 0xed, 0xfe];
|
||||
const MACHO_ARM64_CPU_TYPE: u32 = 0x0100_000c;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum CodexDesktopArtifactError {
|
||||
#[error("Codex Desktop artifact URL is invalid")]
|
||||
InvalidUrl,
|
||||
#[error("Codex Desktop artifact URL is not an approved official ZIP")]
|
||||
UnapprovedUrl,
|
||||
#[error("Codex Desktop artifact size is invalid")]
|
||||
InvalidSize,
|
||||
#[error("Codex Desktop artifact range request failed")]
|
||||
Http(#[from] reqwest::Error),
|
||||
#[error("Codex Desktop artifact did not honor an exact range request")]
|
||||
RangeUnsupported,
|
||||
#[error("Codex Desktop artifact returned an invalid range")]
|
||||
InvalidRange,
|
||||
#[error("Codex Desktop artifact ZIP metadata is invalid or unsupported")]
|
||||
InvalidArchive,
|
||||
#[error("Codex Desktop artifact does not contain exactly one bundled Core")]
|
||||
MissingCore,
|
||||
#[error("Codex Desktop bundled Core format is unsupported")]
|
||||
UnsupportedCore,
|
||||
#[error("Codex Desktop bundled Core exceeds the scan limit")]
|
||||
CoreTooLarge,
|
||||
#[error("Codex Desktop bundled Core version marker is missing or invalid")]
|
||||
MissingCoreVersion,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct CentralDirectory {
|
||||
offset: u64,
|
||||
size: u64,
|
||||
entries: u16,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(super) struct CoreEntry {
|
||||
pub(super) name: Vec<u8>,
|
||||
flags: u16,
|
||||
compression: u16,
|
||||
pub(super) compressed_size: u64,
|
||||
pub(super) uncompressed_size: u64,
|
||||
pub(super) local_header_offset: u64,
|
||||
}
|
||||
|
||||
pub(super) async fn fetch_codex_core_version(
|
||||
client: &Client,
|
||||
artifact_url: &str,
|
||||
artifact_size: u64,
|
||||
) -> Result<String, CodexDesktopArtifactError> {
|
||||
let url = validate_artifact(artifact_url, artifact_size)?;
|
||||
let tail_start = artifact_size.saturating_sub(MAX_EOCD_SEARCH_BYTES);
|
||||
let tail = read_range(client, &url, artifact_size, tail_start, artifact_size - 1).await?;
|
||||
let directory = parse_eocd(&tail, tail_start, artifact_size)?;
|
||||
let directory_end = directory
|
||||
.offset
|
||||
.checked_add(directory.size)
|
||||
.and_then(|end| end.checked_sub(1))
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?;
|
||||
let central = if directory.offset >= tail_start && directory_end < artifact_size {
|
||||
let start = usize::try_from(directory.offset - tail_start)
|
||||
.map_err(|_| CodexDesktopArtifactError::InvalidArchive)?;
|
||||
let end = usize::try_from(directory_end - tail_start + 1)
|
||||
.map_err(|_| CodexDesktopArtifactError::InvalidArchive)?;
|
||||
tail.get(start..end)
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?
|
||||
.to_vec()
|
||||
} else {
|
||||
read_range(client, &url, artifact_size, directory.offset, directory_end).await?
|
||||
};
|
||||
let entry = find_core_entry(¢ral, directory.entries)?;
|
||||
let data_offset = core_data_offset(client, &url, artifact_size, &entry).await?;
|
||||
scan_core_version(client, &url, artifact_size, data_offset, &entry).await
|
||||
}
|
||||
|
||||
pub(super) fn validate_codex_artifact(
|
||||
artifact_url: &str,
|
||||
artifact_size: u64,
|
||||
) -> Result<(), CodexDesktopArtifactError> {
|
||||
validate_artifact(artifact_url, artifact_size).map(drop)
|
||||
}
|
||||
|
||||
fn validate_artifact(
|
||||
artifact_url: &str,
|
||||
artifact_size: u64,
|
||||
) -> Result<Url, CodexDesktopArtifactError> {
|
||||
if artifact_size == 0 || artifact_size > u64::from(u32::MAX) {
|
||||
return Err(CodexDesktopArtifactError::InvalidSize);
|
||||
}
|
||||
let url = Url::parse(artifact_url).map_err(|_| CodexDesktopArtifactError::InvalidUrl)?;
|
||||
if url.scheme() != "https"
|
||||
|| url.host_str() != Some(OFFICIAL_ARTIFACT_HOST)
|
||||
|| !url.path().to_ascii_lowercase().ends_with(".zip")
|
||||
|| url.username() != ""
|
||||
|| url.password().is_some()
|
||||
|| url.fragment().is_some()
|
||||
{
|
||||
return Err(CodexDesktopArtifactError::UnapprovedUrl);
|
||||
}
|
||||
Ok(url)
|
||||
}
|
||||
|
||||
async fn read_range(
|
||||
client: &Client,
|
||||
url: &Url,
|
||||
artifact_size: u64,
|
||||
start: u64,
|
||||
end: u64,
|
||||
) -> Result<Vec<u8>, CodexDesktopArtifactError> {
|
||||
let expected = end
|
||||
.checked_sub(start)
|
||||
.and_then(|length| length.checked_add(1))
|
||||
.filter(|length| *length <= MAX_RANGE_BYTES)
|
||||
.ok_or(CodexDesktopArtifactError::InvalidRange)?;
|
||||
let response = client
|
||||
.get(url.clone())
|
||||
.header(RANGE, format!("bytes={start}-{end}"))
|
||||
.send()
|
||||
.await?;
|
||||
if response.status() != StatusCode::PARTIAL_CONTENT {
|
||||
return Err(CodexDesktopArtifactError::RangeUnsupported);
|
||||
}
|
||||
let content_range = response
|
||||
.headers()
|
||||
.get(CONTENT_RANGE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.ok_or(CodexDesktopArtifactError::InvalidRange)?;
|
||||
if parse_content_range(content_range) != Some((start, end, artifact_size)) {
|
||||
return Err(CodexDesktopArtifactError::InvalidRange);
|
||||
}
|
||||
if response
|
||||
.headers()
|
||||
.get(CONTENT_LENGTH)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| value.parse::<u64>().ok())
|
||||
!= Some(expected)
|
||||
{
|
||||
return Err(CodexDesktopArtifactError::InvalidRange);
|
||||
}
|
||||
|
||||
let capacity =
|
||||
usize::try_from(expected).map_err(|_| CodexDesktopArtifactError::InvalidRange)?;
|
||||
let mut body = Vec::with_capacity(capacity);
|
||||
let mut stream = response.bytes_stream();
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk?;
|
||||
if body
|
||||
.len()
|
||||
.checked_add(chunk.len())
|
||||
.is_none_or(|length| length > capacity)
|
||||
{
|
||||
return Err(CodexDesktopArtifactError::InvalidRange);
|
||||
}
|
||||
body.extend_from_slice(&chunk);
|
||||
}
|
||||
if body.len() != capacity {
|
||||
return Err(CodexDesktopArtifactError::InvalidRange);
|
||||
}
|
||||
Ok(body)
|
||||
}
|
||||
|
||||
pub(super) fn parse_content_range(value: &str) -> Option<(u64, u64, u64)> {
|
||||
let value = value.strip_prefix("bytes ")?;
|
||||
let (range, size) = value.split_once('/')?;
|
||||
let (start, end) = range.split_once('-')?;
|
||||
Some((start.parse().ok()?, end.parse().ok()?, size.parse().ok()?))
|
||||
}
|
||||
|
||||
fn parse_eocd(
|
||||
tail: &[u8],
|
||||
tail_start: u64,
|
||||
artifact_size: u64,
|
||||
) -> Result<CentralDirectory, CodexDesktopArtifactError> {
|
||||
let position = tail
|
||||
.windows(EOCD_SIGNATURE.len())
|
||||
.rposition(|window| window == EOCD_SIGNATURE)
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?;
|
||||
let eocd = tail
|
||||
.get(position..)
|
||||
.filter(|value| value.len() >= 22)
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?;
|
||||
let comment_length = usize::from(le_u16(eocd, 20)?);
|
||||
if eocd.len() != 22 + comment_length || le_u16(eocd, 4)? != 0 || le_u16(eocd, 6)? != 0 {
|
||||
return Err(CodexDesktopArtifactError::InvalidArchive);
|
||||
}
|
||||
let entries_on_disk = le_u16(eocd, 8)?;
|
||||
let entries = le_u16(eocd, 10)?;
|
||||
let size = u64::from(le_u32(eocd, 12)?);
|
||||
let offset = u64::from(le_u32(eocd, 16)?);
|
||||
if entries == 0
|
||||
|| entries != entries_on_disk
|
||||
|| entries == u16::MAX
|
||||
|| size == u64::from(u32::MAX)
|
||||
|| offset == u64::from(u32::MAX)
|
||||
|| size == 0
|
||||
|| size > MAX_CENTRAL_DIRECTORY_BYTES
|
||||
{
|
||||
return Err(CodexDesktopArtifactError::InvalidArchive);
|
||||
}
|
||||
let eocd_offset = tail_start
|
||||
.checked_add(
|
||||
u64::try_from(position).map_err(|_| CodexDesktopArtifactError::InvalidArchive)?,
|
||||
)
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?;
|
||||
if offset
|
||||
.checked_add(size)
|
||||
.is_none_or(|end| end > eocd_offset || end > artifact_size)
|
||||
{
|
||||
return Err(CodexDesktopArtifactError::InvalidArchive);
|
||||
}
|
||||
Ok(CentralDirectory {
|
||||
offset,
|
||||
size,
|
||||
entries,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn find_core_entry(
|
||||
central: &[u8],
|
||||
expected_entries: u16,
|
||||
) -> Result<CoreEntry, CodexDesktopArtifactError> {
|
||||
let mut cursor = 0_usize;
|
||||
let mut core = None;
|
||||
for _ in 0..expected_entries {
|
||||
let header = central
|
||||
.get(cursor..cursor.saturating_add(46))
|
||||
.filter(|header| header.get(..4) == Some(CENTRAL_SIGNATURE.as_slice()))
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?;
|
||||
let name_length = usize::from(le_u16(header, 28)?);
|
||||
let extra_length = usize::from(le_u16(header, 30)?);
|
||||
let comment_length = usize::from(le_u16(header, 32)?);
|
||||
let entry_length = 46_usize
|
||||
.checked_add(name_length)
|
||||
.and_then(|length| length.checked_add(extra_length))
|
||||
.and_then(|length| length.checked_add(comment_length))
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?;
|
||||
let entry = central
|
||||
.get(cursor..cursor.saturating_add(entry_length))
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?;
|
||||
let name = entry
|
||||
.get(46..46 + name_length)
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?;
|
||||
if name.ends_with(CORE_PATH_SUFFIX) {
|
||||
if core.is_some() {
|
||||
return Err(CodexDesktopArtifactError::MissingCore);
|
||||
}
|
||||
let flags = le_u16(header, 8)?;
|
||||
let compression = le_u16(header, 10)?;
|
||||
let compressed_size = u64::from(le_u32(header, 20)?);
|
||||
let uncompressed_size = u64::from(le_u32(header, 24)?);
|
||||
let local_header_offset = u64::from(le_u32(header, 42)?);
|
||||
if flags & 1 != 0
|
||||
|| compression != 8
|
||||
|| compressed_size == 0
|
||||
|| compressed_size == u64::from(u32::MAX)
|
||||
|| compressed_size > MAX_CORE_COMPRESSED_BYTES
|
||||
|| uncompressed_size == 0
|
||||
|| uncompressed_size == u64::from(u32::MAX)
|
||||
|| uncompressed_size > MAX_CORE_UNCOMPRESSED_BYTES
|
||||
|| local_header_offset == u64::from(u32::MAX)
|
||||
{
|
||||
return Err(CodexDesktopArtifactError::UnsupportedCore);
|
||||
}
|
||||
core = Some(CoreEntry {
|
||||
name: name.to_vec(),
|
||||
flags,
|
||||
compression,
|
||||
compressed_size,
|
||||
uncompressed_size,
|
||||
local_header_offset,
|
||||
});
|
||||
}
|
||||
cursor = cursor
|
||||
.checked_add(entry_length)
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?;
|
||||
}
|
||||
core.ok_or(CodexDesktopArtifactError::MissingCore)
|
||||
}
|
||||
|
||||
async fn core_data_offset(
|
||||
client: &Client,
|
||||
url: &Url,
|
||||
artifact_size: u64,
|
||||
entry: &CoreEntry,
|
||||
) -> Result<u64, CodexDesktopArtifactError> {
|
||||
let fixed_end = entry
|
||||
.local_header_offset
|
||||
.checked_add(29)
|
||||
.filter(|end| *end < artifact_size)
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?;
|
||||
let fixed = read_range(
|
||||
client,
|
||||
url,
|
||||
artifact_size,
|
||||
entry.local_header_offset,
|
||||
fixed_end,
|
||||
)
|
||||
.await?;
|
||||
if fixed.get(..4) != Some(LOCAL_SIGNATURE.as_slice())
|
||||
|| le_u16(&fixed, 6)? != entry.flags
|
||||
|| le_u16(&fixed, 8)? != entry.compression
|
||||
{
|
||||
return Err(CodexDesktopArtifactError::InvalidArchive);
|
||||
}
|
||||
let name_length = u64::from(le_u16(&fixed, 26)?);
|
||||
let extra_length = u64::from(le_u16(&fixed, 28)?);
|
||||
let variable_length = name_length
|
||||
.checked_add(extra_length)
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?;
|
||||
let variable_start = entry
|
||||
.local_header_offset
|
||||
.checked_add(30)
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?;
|
||||
let data_offset = variable_start
|
||||
.checked_add(variable_length)
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?;
|
||||
let data_end = data_offset
|
||||
.checked_add(entry.compressed_size)
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?;
|
||||
if data_end > artifact_size || name_length == 0 {
|
||||
return Err(CodexDesktopArtifactError::InvalidArchive);
|
||||
}
|
||||
let variable_end = data_offset
|
||||
.checked_sub(1)
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?;
|
||||
let variable = read_range(client, url, artifact_size, variable_start, variable_end).await?;
|
||||
let local_name_length =
|
||||
usize::try_from(name_length).map_err(|_| CodexDesktopArtifactError::InvalidArchive)?;
|
||||
if variable.get(..local_name_length) != Some(entry.name.as_slice()) {
|
||||
return Err(CodexDesktopArtifactError::InvalidArchive);
|
||||
}
|
||||
Ok(data_offset)
|
||||
}
|
||||
|
||||
async fn scan_core_version(
|
||||
client: &Client,
|
||||
url: &Url,
|
||||
artifact_size: u64,
|
||||
data_offset: u64,
|
||||
entry: &CoreEntry,
|
||||
) -> Result<String, CodexDesktopArtifactError> {
|
||||
let mut decompressor = Decompress::new(false);
|
||||
let mut scanner = CoreVersionScanner::default();
|
||||
let mut macho_header = Vec::with_capacity(MACHO_HEADER_BYTES);
|
||||
let mut compressed_read = 0_u64;
|
||||
let mut uncompressed_read = 0_u64;
|
||||
while compressed_read < entry.compressed_size {
|
||||
let chunk_length = cmp::min(
|
||||
CORE_RANGE_CHUNK_BYTES,
|
||||
entry.compressed_size - compressed_read,
|
||||
);
|
||||
let start = data_offset
|
||||
.checked_add(compressed_read)
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?;
|
||||
let end = start
|
||||
.checked_add(chunk_length - 1)
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)?;
|
||||
let compressed = read_range(client, url, artifact_size, start, end).await?;
|
||||
compressed_read = compressed_read
|
||||
.checked_add(chunk_length)
|
||||
.ok_or(CodexDesktopArtifactError::CoreTooLarge)?;
|
||||
let final_chunk = compressed_read == entry.compressed_size;
|
||||
let mut input_offset = 0_usize;
|
||||
loop {
|
||||
if input_offset == compressed.len() && !final_chunk {
|
||||
break;
|
||||
}
|
||||
let mut output = [0_u8; 64 * 1024];
|
||||
let input_before = decompressor.total_in();
|
||||
let output_before = decompressor.total_out();
|
||||
let status = decompressor
|
||||
.decompress(
|
||||
&compressed[input_offset..],
|
||||
&mut output,
|
||||
if final_chunk {
|
||||
FlushDecompress::Finish
|
||||
} else {
|
||||
FlushDecompress::None
|
||||
},
|
||||
)
|
||||
.map_err(|_| CodexDesktopArtifactError::UnsupportedCore)?;
|
||||
let consumed = usize::try_from(decompressor.total_in() - input_before)
|
||||
.map_err(|_| CodexDesktopArtifactError::CoreTooLarge)?;
|
||||
let produced = usize::try_from(decompressor.total_out() - output_before)
|
||||
.map_err(|_| CodexDesktopArtifactError::CoreTooLarge)?;
|
||||
if consumed == 0 && produced == 0 {
|
||||
return Err(CodexDesktopArtifactError::UnsupportedCore);
|
||||
}
|
||||
input_offset = input_offset
|
||||
.checked_add(consumed)
|
||||
.ok_or(CodexDesktopArtifactError::CoreTooLarge)?;
|
||||
uncompressed_read = uncompressed_read
|
||||
.checked_add(
|
||||
u64::try_from(produced).map_err(|_| CodexDesktopArtifactError::CoreTooLarge)?,
|
||||
)
|
||||
.filter(|read| {
|
||||
*read <= entry.uncompressed_size && *read <= MAX_CORE_UNCOMPRESSED_BYTES
|
||||
})
|
||||
.ok_or(CodexDesktopArtifactError::CoreTooLarge)?;
|
||||
let output = &output[..produced];
|
||||
if macho_header.len() < MACHO_HEADER_BYTES {
|
||||
let needed = MACHO_HEADER_BYTES - macho_header.len();
|
||||
macho_header.extend_from_slice(&output[..cmp::min(needed, output.len())]);
|
||||
if macho_header.len() == MACHO_HEADER_BYTES {
|
||||
validate_macho_header(&macho_header)?;
|
||||
}
|
||||
}
|
||||
if let Some(version) = scanner.push(output)? {
|
||||
if macho_header.len() != MACHO_HEADER_BYTES {
|
||||
return Err(CodexDesktopArtifactError::UnsupportedCore);
|
||||
}
|
||||
return Ok(version);
|
||||
}
|
||||
if status == Status::StreamEnd {
|
||||
if !final_chunk
|
||||
|| input_offset != compressed.len()
|
||||
|| uncompressed_read != entry.uncompressed_size
|
||||
|| macho_header.len() != MACHO_HEADER_BYTES
|
||||
{
|
||||
return Err(CodexDesktopArtifactError::UnsupportedCore);
|
||||
}
|
||||
return scanner
|
||||
.finish()?
|
||||
.ok_or(CodexDesktopArtifactError::MissingCoreVersion);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(CodexDesktopArtifactError::UnsupportedCore)
|
||||
}
|
||||
|
||||
fn validate_macho_header(header: &[u8]) -> Result<(), CodexDesktopArtifactError> {
|
||||
if header.get(..4) != Some(MACHO_64_LE_MAGIC.as_slice())
|
||||
|| header
|
||||
.get(4..8)
|
||||
.and_then(|bytes| bytes.try_into().ok())
|
||||
.map(u32::from_le_bytes)
|
||||
!= Some(MACHO_ARM64_CPU_TYPE)
|
||||
{
|
||||
return Err(CodexDesktopArtifactError::UnsupportedCore);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub(super) struct CoreVersionScanner {
|
||||
carry: Vec<u8>,
|
||||
}
|
||||
|
||||
impl CoreVersionScanner {
|
||||
pub(super) fn push(
|
||||
&mut self,
|
||||
bytes: &[u8],
|
||||
) -> Result<Option<String>, CodexDesktopArtifactError> {
|
||||
let mut window = Vec::with_capacity(self.carry.len() + bytes.len());
|
||||
window.extend_from_slice(&self.carry);
|
||||
window.extend_from_slice(bytes);
|
||||
let mut search_start = 0_usize;
|
||||
while let Some(relative_position) = find_bytes(&window[search_start..], CORE_VERSION_ANCHOR)
|
||||
{
|
||||
let position = search_start + relative_position;
|
||||
if let Some(version) =
|
||||
parse_codex_semver(&window[position + CORE_VERSION_ANCHOR.len()..], false)?
|
||||
{
|
||||
return Ok(Some(version));
|
||||
}
|
||||
search_start = position + CORE_VERSION_ANCHOR.len();
|
||||
}
|
||||
let carry_length = cmp::min(window.len(), CORE_VERSION_ANCHOR.len() + MAX_VERSION_BYTES);
|
||||
self.carry.clear();
|
||||
self.carry
|
||||
.extend_from_slice(&window[window.len() - carry_length..]);
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn finish(&self) -> Result<Option<String>, CodexDesktopArtifactError> {
|
||||
let mut search_start = 0_usize;
|
||||
while let Some(relative_position) =
|
||||
find_bytes(&self.carry[search_start..], CORE_VERSION_ANCHOR)
|
||||
{
|
||||
let position = search_start + relative_position;
|
||||
if let Some(version) =
|
||||
parse_codex_semver(&self.carry[position + CORE_VERSION_ANCHOR.len()..], true)?
|
||||
{
|
||||
return Ok(Some(version));
|
||||
}
|
||||
search_start = position + CORE_VERSION_ANCHOR.len();
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_codex_semver(
|
||||
bytes: &[u8],
|
||||
allow_end: bool,
|
||||
) -> Result<Option<String>, CodexDesktopArtifactError> {
|
||||
let mut cursor = 0_usize;
|
||||
for component in 0..3 {
|
||||
let start = cursor;
|
||||
while bytes.get(cursor).is_some_and(u8::is_ascii_digit) {
|
||||
cursor += 1;
|
||||
}
|
||||
if cursor == start || (cursor - start > 1 && bytes[start] == b'0') {
|
||||
return Ok(None);
|
||||
}
|
||||
if component < 2 {
|
||||
if bytes.get(cursor) != Some(&b'.') {
|
||||
return Ok(None);
|
||||
}
|
||||
cursor += 1;
|
||||
}
|
||||
}
|
||||
if bytes.get(cursor) == Some(&b'-') {
|
||||
cursor += 1;
|
||||
let channel = [b"alpha".as_slice(), b"beta".as_slice(), b"rc".as_slice()]
|
||||
.into_iter()
|
||||
.find(|channel| bytes.get(cursor..cursor + channel.len()) == Some(*channel));
|
||||
let Some(channel) = channel else {
|
||||
return Ok(None);
|
||||
};
|
||||
cursor += channel.len();
|
||||
let mut numeric_parts = 0_usize;
|
||||
while bytes.get(cursor) == Some(&b'.') {
|
||||
cursor += 1;
|
||||
let start = cursor;
|
||||
while bytes.get(cursor).is_some_and(u8::is_ascii_digit) {
|
||||
cursor += 1;
|
||||
}
|
||||
if cursor == start || (cursor - start > 1 && bytes[start] == b'0') {
|
||||
return Ok(None);
|
||||
}
|
||||
numeric_parts += 1;
|
||||
}
|
||||
if numeric_parts == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
if cursor == bytes.len() && !allow_end {
|
||||
return Ok(None);
|
||||
}
|
||||
if bytes.get(cursor).is_some_and(u8::is_ascii_digit)
|
||||
|| bytes.get(cursor) == Some(&b'.')
|
||||
|| bytes
|
||||
.get(cursor)
|
||||
.is_some_and(|byte| matches!(*byte, b'+' | b'-'))
|
||||
|| cursor > MAX_VERSION_BYTES
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
let version = std::str::from_utf8(&bytes[..cursor])
|
||||
.map_err(|_| CodexDesktopArtifactError::MissingCoreVersion)?;
|
||||
semver::Version::parse(version).map_err(|_| CodexDesktopArtifactError::MissingCoreVersion)?;
|
||||
Ok(Some(version.to_owned()))
|
||||
}
|
||||
|
||||
fn find_bytes(haystack: &[u8], needle: &[u8]) -> Option<usize> {
|
||||
haystack
|
||||
.windows(needle.len())
|
||||
.position(|window| window == needle)
|
||||
}
|
||||
|
||||
fn le_u16(bytes: &[u8], offset: usize) -> Result<u16, CodexDesktopArtifactError> {
|
||||
bytes
|
||||
.get(offset..offset + 2)
|
||||
.and_then(|value| value.try_into().ok())
|
||||
.map(u16::from_le_bytes)
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)
|
||||
}
|
||||
|
||||
fn le_u32(bytes: &[u8], offset: usize) -> Result<u32, CodexDesktopArtifactError> {
|
||||
bytes
|
||||
.get(offset..offset + 4)
|
||||
.and_then(|value| value.try_into().ok())
|
||||
.map(u32::from_le_bytes)
|
||||
.ok_or(CodexDesktopArtifactError::InvalidArchive)
|
||||
}
|
||||
+273
-107
@@ -5,31 +5,42 @@ use std::time::Duration;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use futures::{StreamExt as _, future::BoxFuture};
|
||||
use gateway_core::engine::credential::OpaqueProviderData;
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderArtifactProfile, ProviderArtifactProfileCachePort, ProviderStoreError,
|
||||
};
|
||||
use gateway_core::routing::ProviderKind;
|
||||
use reqwest::Client;
|
||||
use reqwest::redirect::Policy;
|
||||
use roxmltree::{Document, Node};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{Map, Value};
|
||||
use url::Url;
|
||||
|
||||
use self::desktop_artifact::{CodexDesktopArtifactError, fetch_codex_core_version};
|
||||
|
||||
mod desktop_artifact;
|
||||
|
||||
/// Codex Desktop 官方 appcast 地址。
|
||||
pub const CODEX_DESKTOP_APPCAST_URL: &str =
|
||||
"https://persistent.oaistatic.com/codex-app-prod/appcast.xml";
|
||||
pub const CODEX_CLI_RELEASE_URL: &str = "https://registry.npmjs.org/@openai%2Fcodex/latest";
|
||||
/// 官方发布元数据检查周期。
|
||||
pub const APPCAST_POLL_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60);
|
||||
|
||||
const APPCAST_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
const MAX_APPCAST_BYTES: usize = 1024 * 1024;
|
||||
const MAX_CLI_RELEASE_BYTES: usize = 64 * 1024;
|
||||
const ARTIFACT_PROFILE_CACHE_TTL: Duration = Duration::from_secs(30 * 24 * 60 * 60);
|
||||
const ARTIFACT_PROFILE_SCHEMA_VERSION: u64 = 1;
|
||||
|
||||
/// Codex Desktop 上游请求身份。
|
||||
///
|
||||
/// 启动配置提供经源码审计的 Core、运行环境和 Desktop 启动版本。npm CLI metadata
|
||||
/// 与官方 appcast 检查成功后,分别同步 Core 版本和 Desktop 版本及构建号。
|
||||
/// 启动配置提供经源码审计的 Core、运行环境和 Desktop 启动版本。运行时只会使用
|
||||
/// 同一个官方 Desktop ZIP 中核验出的 Core、Desktop 版本及构建号原子替换版本字段。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CodexWireProfile {
|
||||
/// `originator` 请求头及 User-Agent 产品名。
|
||||
pub originator: String,
|
||||
/// npm `@openai/codex` 启动版本;用于 `/codex/models?client_version=` 与 UA。
|
||||
/// Desktop ZIP 内嵌 Core 版本;用于 `/codex/models?client_version=` 与 UA。
|
||||
pub codex_version: String,
|
||||
/// Desktop 应用版本,用于 app-server `clientInfo.version` 对应的 UA 后缀。
|
||||
pub desktop_version: String,
|
||||
@@ -86,124 +97,148 @@ impl CodexWireProfileState {
|
||||
.clone()
|
||||
}
|
||||
|
||||
/// 使用官方 appcast 已验证的 Desktop 版本和构建号更新运行时画像。
|
||||
pub fn update_desktop_release(&self, desktop_version: &str, desktop_build: &str) {
|
||||
/// 原子发布同一 Desktop ZIP 中核验出的完整版本元组。
|
||||
pub fn update_bundled_release(&self, release: &CodexBundledReleaseProfile) {
|
||||
let mut profile = self
|
||||
.profile
|
||||
.write()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if profile.desktop_version == desktop_version && profile.desktop_build == desktop_build {
|
||||
return;
|
||||
profile.codex_version.clone_from(&release.codex_version);
|
||||
profile.desktop_version.clone_from(&release.desktop_version);
|
||||
profile.desktop_build.clone_from(&release.desktop_build);
|
||||
profile.verified_at = release.verified_at;
|
||||
}
|
||||
}
|
||||
|
||||
/// 同一个官方 Desktop ZIP 中核验出的原子版本元组。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CodexBundledReleaseProfile {
|
||||
pub codex_version: String,
|
||||
pub desktop_version: String,
|
||||
pub desktop_build: String,
|
||||
pub verified_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
impl CodexBundledReleaseProfile {
|
||||
fn artifact_sequence(&self) -> Result<u64, CodexDesktopReleaseError> {
|
||||
parse_build_sequence(&self.desktop_build)
|
||||
}
|
||||
|
||||
fn validate(&self) -> Result<(), CodexDesktopReleaseError> {
|
||||
semver::Version::parse(&self.codex_version)
|
||||
.map_err(|_| CodexDesktopReleaseError::InvalidCoreVersion)?;
|
||||
if !numeric_dotted_version(&self.desktop_version) {
|
||||
return Err(CodexDesktopReleaseError::InvalidVersion);
|
||||
}
|
||||
profile.desktop_version = desktop_version.to_owned();
|
||||
profile.desktop_build = desktop_build.to_owned();
|
||||
}
|
||||
|
||||
/// 使用官方 npm metadata 发布的 Codex CLI 版本更新运行时请求画像。
|
||||
pub fn update_cli_release(&self, codex_version: &str) {
|
||||
let mut profile = self
|
||||
.profile
|
||||
.write()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if profile.codex_version != codex_version {
|
||||
profile.codex_version = codex_version.to_owned();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub trait CodexCliReleaseTransport: Send + Sync {
|
||||
fn fetch(&self) -> BoxFuture<'_, Result<String, CodexCliReleaseError>>;
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct OfficialCodexCliReleaseTransport {
|
||||
client: Client,
|
||||
endpoint: Url,
|
||||
}
|
||||
|
||||
impl OfficialCodexCliReleaseTransport {
|
||||
pub fn new() -> Result<Self, CodexCliReleaseError> {
|
||||
let endpoint =
|
||||
Url::parse(CODEX_CLI_RELEASE_URL).map_err(|_| CodexCliReleaseError::InvalidEndpoint)?;
|
||||
let client = Client::builder()
|
||||
.https_only(true)
|
||||
.no_proxy()
|
||||
.redirect(Policy::none())
|
||||
.connect_timeout(Duration::from_secs(10))
|
||||
.timeout(APPCAST_TIMEOUT)
|
||||
.build()
|
||||
.map_err(|_| CodexCliReleaseError::ClientInitialization)?;
|
||||
Ok(Self { client, endpoint })
|
||||
}
|
||||
}
|
||||
|
||||
impl CodexCliReleaseTransport for OfficialCodexCliReleaseTransport {
|
||||
fn fetch(&self) -> BoxFuture<'_, Result<String, CodexCliReleaseError>> {
|
||||
Box::pin(async move {
|
||||
let response = self.client.get(self.endpoint.clone()).send().await?;
|
||||
if !response.status().is_success() {
|
||||
return Err(CodexCliReleaseError::HttpStatus(response.status().as_u16()));
|
||||
}
|
||||
if response
|
||||
.content_length()
|
||||
.is_some_and(|size| size > MAX_CLI_RELEASE_BYTES as u64)
|
||||
{
|
||||
return Err(CodexCliReleaseError::ResponseTooLarge);
|
||||
}
|
||||
let bytes = response.bytes().await?;
|
||||
if bytes.len() > MAX_CLI_RELEASE_BYTES {
|
||||
return Err(CodexCliReleaseError::ResponseTooLarge);
|
||||
}
|
||||
let document: serde_json::Value = serde_json::from_slice(&bytes)
|
||||
.map_err(|_| CodexCliReleaseError::InvalidDocument)?;
|
||||
let version = document
|
||||
.get("version")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.ok_or(CodexCliReleaseError::InvalidDocument)?;
|
||||
semver::Version::parse(version).map_err(|_| CodexCliReleaseError::InvalidVersion)?;
|
||||
Ok(version.to_owned())
|
||||
})
|
||||
self.artifact_sequence()?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct CodexCliReleaseService {
|
||||
transport: Arc<dyn CodexCliReleaseTransport>,
|
||||
profile: CodexWireProfileState,
|
||||
pub struct CodexArtifactProfileCache {
|
||||
provider_kind: ProviderKind,
|
||||
store: Arc<dyn ProviderArtifactProfileCachePort>,
|
||||
}
|
||||
|
||||
impl CodexCliReleaseService {
|
||||
impl CodexArtifactProfileCache {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
profile: CodexWireProfileState,
|
||||
transport: Arc<dyn CodexCliReleaseTransport>,
|
||||
provider_kind: ProviderKind,
|
||||
store: Arc<dyn ProviderArtifactProfileCachePort>,
|
||||
) -> Self {
|
||||
Self { transport, profile }
|
||||
Self {
|
||||
provider_kind,
|
||||
store,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn refresh(&self) -> Result<String, CodexCliReleaseError> {
|
||||
let version = self.transport.fetch().await?;
|
||||
self.profile.update_cli_release(&version);
|
||||
Ok(version)
|
||||
pub async fn load(
|
||||
&self,
|
||||
) -> Result<Option<CodexBundledReleaseProfile>, CodexDesktopReleaseError> {
|
||||
self.store
|
||||
.read(&self.provider_kind)
|
||||
.await
|
||||
.map_err(CodexDesktopReleaseError::ArtifactCache)?
|
||||
.map(decode_artifact_profile)
|
||||
.transpose()
|
||||
}
|
||||
|
||||
async fn replace_if_newer(
|
||||
&self,
|
||||
profile: &CodexBundledReleaseProfile,
|
||||
) -> Result<bool, CodexDesktopReleaseError> {
|
||||
profile.validate()?;
|
||||
self.store
|
||||
.replace_if_newer(
|
||||
encode_artifact_profile(self.provider_kind.clone(), profile)?,
|
||||
ARTIFACT_PROFILE_CACHE_TTL,
|
||||
)
|
||||
.await
|
||||
.map_err(CodexDesktopReleaseError::ArtifactCache)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum CodexCliReleaseError {
|
||||
#[error("Codex CLI release client initialization failed")]
|
||||
ClientInitialization,
|
||||
#[error("Codex CLI release endpoint is invalid")]
|
||||
InvalidEndpoint,
|
||||
#[error("Codex CLI release request failed")]
|
||||
Request(#[from] reqwest::Error),
|
||||
#[error("Codex CLI release endpoint returned HTTP {0}")]
|
||||
HttpStatus(u16),
|
||||
#[error("Codex CLI release response exceeded the size limit")]
|
||||
ResponseTooLarge,
|
||||
#[error("Codex CLI release document is invalid")]
|
||||
InvalidDocument,
|
||||
#[error("Codex CLI release version is invalid")]
|
||||
InvalidVersion,
|
||||
#[derive(Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct CachedArtifactProfile {
|
||||
schema_version: u64,
|
||||
codex_version: String,
|
||||
desktop_version: String,
|
||||
desktop_build: String,
|
||||
}
|
||||
|
||||
fn encode_artifact_profile(
|
||||
provider_kind: ProviderKind,
|
||||
profile: &CodexBundledReleaseProfile,
|
||||
) -> Result<ProviderArtifactProfile, CodexDesktopReleaseError> {
|
||||
let artifact_sequence = profile.artifact_sequence()?;
|
||||
let mut fields = Map::new();
|
||||
fields.insert(
|
||||
"schema_version".to_owned(),
|
||||
Value::from(ARTIFACT_PROFILE_SCHEMA_VERSION),
|
||||
);
|
||||
fields.insert(
|
||||
"codex_version".to_owned(),
|
||||
Value::String(profile.codex_version.clone()),
|
||||
);
|
||||
fields.insert(
|
||||
"desktop_version".to_owned(),
|
||||
Value::String(profile.desktop_version.clone()),
|
||||
);
|
||||
fields.insert(
|
||||
"desktop_build".to_owned(),
|
||||
Value::String(profile.desktop_build.clone()),
|
||||
);
|
||||
Ok(ProviderArtifactProfile::new(
|
||||
provider_kind,
|
||||
artifact_sequence,
|
||||
profile.verified_at.into(),
|
||||
OpaqueProviderData::new(fields),
|
||||
))
|
||||
}
|
||||
|
||||
fn decode_artifact_profile(
|
||||
profile: ProviderArtifactProfile,
|
||||
) -> Result<CodexBundledReleaseProfile, CodexDesktopReleaseError> {
|
||||
let wire: CachedArtifactProfile = serde_json::from_value(Value::Object(
|
||||
profile.profile().expose_to_provider().clone(),
|
||||
))
|
||||
.map_err(|_| CodexDesktopReleaseError::InvalidArtifactCache)?;
|
||||
if wire.schema_version != ARTIFACT_PROFILE_SCHEMA_VERSION {
|
||||
return Err(CodexDesktopReleaseError::InvalidArtifactCache);
|
||||
}
|
||||
let decoded = CodexBundledReleaseProfile {
|
||||
codex_version: wire.codex_version,
|
||||
desktop_version: wire.desktop_version,
|
||||
desktop_build: wire.desktop_build,
|
||||
verified_at: DateTime::<Utc>::from(profile.verified_at()),
|
||||
};
|
||||
decoded.validate()?;
|
||||
if decoded.artifact_sequence()? != profile.artifact_sequence() {
|
||||
return Err(CodexDesktopReleaseError::InvalidArtifactCache);
|
||||
}
|
||||
Ok(decoded)
|
||||
}
|
||||
|
||||
/// 官方 appcast 中按顺序出现的首个完整 Desktop 制品。
|
||||
@@ -266,6 +301,11 @@ impl CodexDesktopReleaseStatus {
|
||||
/// Desktop appcast 获取边界。生产实现固定访问官方 HTTPS,测试实现只替换此边界。
|
||||
pub trait CodexDesktopReleaseTransport: Send + Sync {
|
||||
fn fetch(&self) -> BoxFuture<'_, Result<CodexDesktopRelease, CodexDesktopReleaseError>>;
|
||||
|
||||
fn fetch_bundled_core_version<'a>(
|
||||
&'a self,
|
||||
release: &'a CodexDesktopRelease,
|
||||
) -> BoxFuture<'a, Result<String, CodexDesktopReleaseError>>;
|
||||
}
|
||||
|
||||
/// 固定访问官方 Desktop appcast 的生产 transport。
|
||||
@@ -325,6 +365,24 @@ impl CodexDesktopReleaseTransport for OfficialCodexDesktopReleaseTransport {
|
||||
parse_desktop_release(xml)
|
||||
})
|
||||
}
|
||||
|
||||
fn fetch_bundled_core_version<'a>(
|
||||
&'a self,
|
||||
release: &'a CodexDesktopRelease,
|
||||
) -> BoxFuture<'a, Result<String, CodexDesktopReleaseError>> {
|
||||
Box::pin(async move {
|
||||
let artifact_url = release
|
||||
.download_url
|
||||
.as_deref()
|
||||
.ok_or(CodexDesktopReleaseError::MissingArtifactIdentity)?;
|
||||
let artifact_size = release
|
||||
.download_size
|
||||
.ok_or(CodexDesktopReleaseError::MissingArtifactIdentity)?;
|
||||
fetch_codex_core_version(&self.client, artifact_url, artifact_size)
|
||||
.await
|
||||
.map_err(Into::into)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// 拉取 appcast 并原子发布请求画像的 Provider 服务。
|
||||
@@ -333,6 +391,8 @@ pub struct CodexDesktopReleaseService {
|
||||
transport: Arc<dyn CodexDesktopReleaseTransport>,
|
||||
status: CodexDesktopReleaseStatus,
|
||||
profile: CodexWireProfileState,
|
||||
cache: CodexArtifactProfileCache,
|
||||
verified: Arc<RwLock<Option<CodexBundledReleaseProfile>>>,
|
||||
}
|
||||
|
||||
impl CodexDesktopReleaseService {
|
||||
@@ -341,11 +401,18 @@ impl CodexDesktopReleaseService {
|
||||
pub fn new(
|
||||
profile: CodexWireProfileState,
|
||||
transport: Arc<dyn CodexDesktopReleaseTransport>,
|
||||
cache: CodexArtifactProfileCache,
|
||||
verified: Option<CodexBundledReleaseProfile>,
|
||||
) -> Self {
|
||||
if let Some(verified) = verified.as_ref() {
|
||||
profile.update_bundled_release(verified);
|
||||
}
|
||||
Self {
|
||||
transport,
|
||||
status: CodexDesktopReleaseStatus::default(),
|
||||
profile,
|
||||
cache,
|
||||
verified: Arc::new(RwLock::new(verified)),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -357,11 +424,9 @@ impl CodexDesktopReleaseService {
|
||||
/// 执行一次有界检查;失败只更新观察状态,不修改上一份成功画像。
|
||||
pub async fn refresh(&self) -> Result<CodexDesktopRelease, CodexDesktopReleaseError> {
|
||||
let checked_at = Utc::now();
|
||||
let result = self.transport.fetch().await;
|
||||
let result = self.refresh_inner(checked_at).await;
|
||||
match result {
|
||||
Ok(release) => {
|
||||
self.profile
|
||||
.update_desktop_release(&release.version, &release.build);
|
||||
self.status.record_success(checked_at, release.clone());
|
||||
Ok(release)
|
||||
}
|
||||
@@ -371,6 +436,91 @@ impl CodexDesktopReleaseService {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn refresh_inner(
|
||||
&self,
|
||||
checked_at: DateTime<Utc>,
|
||||
) -> Result<CodexDesktopRelease, CodexDesktopReleaseError> {
|
||||
let release = self.transport.fetch().await?;
|
||||
validate_release_artifact(&release)?;
|
||||
let release_sequence = parse_build_sequence(&release.build)?;
|
||||
let verified = self
|
||||
.verified
|
||||
.read()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.clone();
|
||||
if let Some(verified) = verified {
|
||||
let verified_sequence = verified.artifact_sequence()?;
|
||||
if release_sequence < verified_sequence {
|
||||
return Err(CodexDesktopReleaseError::ReleaseRollback);
|
||||
}
|
||||
if release_sequence == verified_sequence {
|
||||
if release.version != verified.desktop_version
|
||||
|| release.build != verified.desktop_build
|
||||
{
|
||||
return Err(CodexDesktopReleaseError::InconsistentRelease);
|
||||
}
|
||||
if !self.cache.replace_if_newer(&verified).await? {
|
||||
return Err(CodexDesktopReleaseError::ReleaseRollback);
|
||||
}
|
||||
self.profile.update_bundled_release(&verified);
|
||||
return Ok(release);
|
||||
}
|
||||
}
|
||||
|
||||
let codex_version = self.transport.fetch_bundled_core_version(&release).await?;
|
||||
let bundled = CodexBundledReleaseProfile {
|
||||
codex_version,
|
||||
desktop_version: release.version.clone(),
|
||||
desktop_build: release.build.clone(),
|
||||
verified_at: checked_at,
|
||||
};
|
||||
bundled.validate()?;
|
||||
let published = if self.cache.replace_if_newer(&bundled).await? {
|
||||
bundled
|
||||
} else {
|
||||
self.cache
|
||||
.load()
|
||||
.await?
|
||||
.filter(|cached| {
|
||||
cached
|
||||
.artifact_sequence()
|
||||
.is_ok_and(|sequence| sequence >= release_sequence)
|
||||
})
|
||||
.ok_or(CodexDesktopReleaseError::ReleaseRollback)?
|
||||
};
|
||||
self.profile.update_bundled_release(&published);
|
||||
*self
|
||||
.verified
|
||||
.write()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some(published);
|
||||
Ok(release)
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_release_artifact(
|
||||
release: &CodexDesktopRelease,
|
||||
) -> Result<(), CodexDesktopReleaseError> {
|
||||
if !release.signature_present {
|
||||
return Err(CodexDesktopReleaseError::MissingArtifactSignature);
|
||||
}
|
||||
let artifact_url = release
|
||||
.download_url
|
||||
.as_deref()
|
||||
.ok_or(CodexDesktopReleaseError::MissingArtifactIdentity)?;
|
||||
let artifact_size = release
|
||||
.download_size
|
||||
.ok_or(CodexDesktopReleaseError::MissingArtifactIdentity)?;
|
||||
desktop_artifact::validate_codex_artifact(artifact_url, artifact_size)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_build_sequence(value: &str) -> Result<u64, CodexDesktopReleaseError> {
|
||||
value
|
||||
.parse::<u64>()
|
||||
.ok()
|
||||
.filter(|sequence| *sequence > 0)
|
||||
.ok_or(CodexDesktopReleaseError::InvalidBuild)
|
||||
}
|
||||
|
||||
/// 解析 appcast 中按顺序出现的首个完整发布项。
|
||||
@@ -489,4 +639,20 @@ pub enum CodexDesktopReleaseError {
|
||||
InvalidPublishedAt,
|
||||
#[error("Codex Desktop appcast download size is invalid")]
|
||||
InvalidDownloadSize,
|
||||
#[error("Codex Desktop appcast artifact identity is incomplete")]
|
||||
MissingArtifactIdentity,
|
||||
#[error("Codex Desktop appcast artifact signature is missing")]
|
||||
MissingArtifactSignature,
|
||||
#[error("Codex Desktop bundled Core version is invalid")]
|
||||
InvalidCoreVersion,
|
||||
#[error(transparent)]
|
||||
Artifact(#[from] CodexDesktopArtifactError),
|
||||
#[error("Codex Desktop artifact profile cache operation failed")]
|
||||
ArtifactCache(#[source] ProviderStoreError),
|
||||
#[error("Codex Desktop artifact profile cache is invalid")]
|
||||
InvalidArtifactCache,
|
||||
#[error("Codex Desktop appcast attempted to roll back the verified artifact")]
|
||||
ReleaseRollback,
|
||||
#[error("Codex Desktop appcast release identity conflicts with the verified artifact")]
|
||||
InconsistentRelease,
|
||||
}
|
||||
@@ -5,6 +5,8 @@ use serde_json::{Map, Value};
|
||||
use crate::transport::protocol::responses::{CodexResponsesRequest, transport_requirement};
|
||||
|
||||
const REDACTED_PAYLOAD_VALUE: &str = "<redacted>";
|
||||
const PREVIOUS_RESPONSE_NOT_FOUND_CODE: &str = "previous_response_not_found";
|
||||
const INVALID_PREVIOUS_RESPONSE_ID_MESSAGE: &str = "Invalid `previous_response_id`.";
|
||||
|
||||
/// WebSocket 握手审计快照。
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
|
||||
@@ -54,16 +56,49 @@ pub fn websocket_event_to_sse_frame(raw: &str) -> Option<String> {
|
||||
websocket_event_frame(&value, raw)
|
||||
}
|
||||
|
||||
/// 同上,复用已解析的事件 JSON;`data` 段逐字节内嵌 `raw` 原文。
|
||||
/// 同上,复用已解析的事件 JSON;`data` 段通常逐字节内嵌 `raw` 原文。
|
||||
///
|
||||
/// 仅剥离传输层内部帧:`codex.rate_limits` 与承载 turn_state 的
|
||||
/// `response.metadata` 不下发客户端;业务事件一律原样转发,不做 schema 校验。
|
||||
/// `response.metadata` 不下发客户端。唯一的业务帧适配是为上游省略错误码的
|
||||
/// `Invalid previous_response_id` 补齐官方客户端用于完整历史重放的稳定错误码;
|
||||
/// 原始 status、type 与 message 均保持不变。
|
||||
pub(crate) fn websocket_event_frame(value: &Value, raw: &str) -> Option<String> {
|
||||
let event = websocket_event_type(value)?;
|
||||
if is_internal_websocket_event(event) || event == "response.metadata" {
|
||||
return None;
|
||||
}
|
||||
Some(encode_sse_event(event, raw))
|
||||
let normalized = previous_response_not_found_event(value);
|
||||
Some(encode_sse_event(
|
||||
event,
|
||||
normalized.as_deref().unwrap_or(raw),
|
||||
))
|
||||
}
|
||||
|
||||
fn previous_response_not_found_event(value: &Value) -> Option<String> {
|
||||
if websocket_event_type(value) != Some("error")
|
||||
|| value
|
||||
.get("status")
|
||||
.or_else(|| value.get("status_code"))
|
||||
.and_then(Value::as_u64)
|
||||
!= Some(400)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let error = value.get("error")?.as_object()?;
|
||||
if error.get("code").is_some_and(|code| !code.is_null())
|
||||
|| error.get("type").and_then(Value::as_str) != Some("invalid_request_error")
|
||||
|| error.get("message").and_then(Value::as_str)
|
||||
!= Some(INVALID_PREVIOUS_RESPONSE_ID_MESSAGE)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut normalized = value.clone();
|
||||
normalized.get_mut("error")?.as_object_mut()?.insert(
|
||||
"code".to_owned(),
|
||||
Value::String(PREVIOUS_RESPONSE_NOT_FOUND_CODE.to_owned()),
|
||||
);
|
||||
serde_json::to_string(&normalized).ok()
|
||||
}
|
||||
|
||||
pub(crate) fn websocket_event_type(value: &Value) -> Option<&str> {
|
||||
|
||||
@@ -61,6 +61,23 @@ const INSTALLATION_ID_KEYS: &[&str] = &[
|
||||
"x-codex-installation-id",
|
||||
];
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum RequestAccountScope {
|
||||
SameAccount,
|
||||
CrossAccountFullRequest,
|
||||
CrossAccountContinuationProbe,
|
||||
}
|
||||
|
||||
impl RequestAccountScope {
|
||||
const fn is_cross_account(self) -> bool {
|
||||
!matches!(self, Self::SameAccount)
|
||||
}
|
||||
|
||||
const fn preserves_continuation(self) -> bool {
|
||||
matches!(self, Self::CrossAccountContinuationProbe)
|
||||
}
|
||||
}
|
||||
|
||||
/// Provider 专属编码错误;不保存 prompt、schema 或 option 值。
|
||||
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
|
||||
pub enum CodexRequestEncodeError {
|
||||
@@ -289,12 +306,18 @@ fn normalize_conversation_anchor_text(text: &str) -> String {
|
||||
///
|
||||
/// 真实 account ID 与 installation ID 由随后构造的 `CodexRequestContext` 注入
|
||||
/// 请求头;正文只替换客户端原本声明过的 installation 字段,绝不接受客户端
|
||||
/// 提供的 token、cookie 或账号身份。
|
||||
/// 提供的 token、cookie 或账号身份。Continuation probe 保留 canonical
|
||||
/// previous-response 和原始增量 input,不能把它们误当成完整请求清理。
|
||||
pub(crate) fn scope_request_to_account(
|
||||
request: &mut CodexResponsesRequest,
|
||||
installation_id: &str,
|
||||
cross_account: bool,
|
||||
account_scope: RequestAccountScope,
|
||||
) {
|
||||
let cross_account = account_scope.is_cross_account();
|
||||
let continuation = account_scope
|
||||
.preserves_continuation()
|
||||
.then(|| request.previous_response_id().map(ToOwned::to_owned))
|
||||
.flatten();
|
||||
let client_metadata_turn_state = metadata_string(request, "x-codex-turn-state");
|
||||
let preserve_turn_state =
|
||||
!cross_account && (request.turn_state.is_some() || client_metadata_turn_state.is_some());
|
||||
@@ -306,23 +329,38 @@ pub(crate) fn scope_request_to_account(
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let turn_metadata = request
|
||||
.turn_metadata
|
||||
.as_deref()
|
||||
.and_then(|metadata| scope_turn_metadata(metadata, installation_id, cross_account));
|
||||
let client_metadata_turn_metadata = metadata_string(request, "x-codex-turn-metadata")
|
||||
.and_then(|metadata| scope_turn_metadata(&metadata, installation_id, cross_account));
|
||||
let turn_metadata = (!account_scope.preserves_continuation())
|
||||
.then(|| {
|
||||
request
|
||||
.turn_metadata
|
||||
.as_deref()
|
||||
.and_then(|metadata| scope_turn_metadata(metadata, installation_id, cross_account))
|
||||
})
|
||||
.flatten();
|
||||
let client_metadata_turn_metadata = (!account_scope.preserves_continuation())
|
||||
.then(|| {
|
||||
metadata_string(request, "x-codex-turn-metadata")
|
||||
.and_then(|metadata| scope_turn_metadata(&metadata, installation_id, cross_account))
|
||||
})
|
||||
.flatten();
|
||||
|
||||
if cross_account {
|
||||
request.passthrough_headers.remove("x-codex-turn-state");
|
||||
request.passthrough_headers.remove("x-codex-turn-metadata");
|
||||
sanitize_cross_account_input(request);
|
||||
if !account_scope.preserves_continuation() {
|
||||
sanitize_cross_account_input(request);
|
||||
}
|
||||
for key in CROSS_ACCOUNT_IDENTITY_KEYS
|
||||
.iter()
|
||||
.chain(ACCOUNT_BOUND_STATE_KEYS)
|
||||
{
|
||||
request.body_mut().remove(*key);
|
||||
}
|
||||
if account_scope.preserves_continuation() {
|
||||
for key in TURN_METADATA_KEYS {
|
||||
request.body_mut().remove(*key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for key in INSTALLATION_ID_KEYS {
|
||||
@@ -356,6 +394,11 @@ pub(crate) fn scope_request_to_account(
|
||||
{
|
||||
metadata.remove(*key);
|
||||
}
|
||||
if account_scope.preserves_continuation() {
|
||||
for key in TURN_METADATA_KEYS {
|
||||
metadata.remove(*key);
|
||||
}
|
||||
}
|
||||
}
|
||||
metadata.insert(
|
||||
"x-codex-installation-id".to_owned(),
|
||||
@@ -393,6 +436,9 @@ pub(crate) fn scope_request_to_account(
|
||||
|
||||
request.turn_state = turn_state;
|
||||
request.turn_metadata = turn_metadata;
|
||||
if let Some(continuation) = continuation {
|
||||
request.set_previous_response_id(Some(continuation));
|
||||
}
|
||||
}
|
||||
|
||||
fn sanitize_cross_account_input(request: &mut CodexResponsesRequest) {
|
||||
|
||||
@@ -193,7 +193,6 @@ fn client_response_header_is_forwardable(name: &str, connection_options: &[Strin
|
||||
| "x-openai-organization"
|
||||
| "x-openai-project"
|
||||
| "x-codex-installation-id"
|
||||
| "x-codex-turn-state"
|
||||
| "x-codex-turn-metadata"
|
||||
)
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
//! Codex 本地会话锚点的稳定派生。
|
||||
//!
|
||||
//! 旧版 Redis 会话亲和与 WebSocket 复用以 `.runtime/data/identity_hmac_secret`
|
||||
//! 旧版 Redis 会话亲和与 WebSocket 复用以运行数据目录下的 `identity_hmac_secret`
|
||||
//! 生成的 `lc_` 值为键。该密钥是持久运行数据,不是账号凭据;保留它可避免
|
||||
//! 架构迁移后同一客户端会话全部换键。
|
||||
|
||||
|
||||
@@ -2,19 +2,95 @@ use gateway_core::accounting::{
|
||||
CalculatedCostAmounts, CalculatedCostBreakdown, CalculatedCostRates, CurrencyCode, Decimal,
|
||||
Money,
|
||||
};
|
||||
use gateway_protocol::openai::events::retry_after_seconds_from_body;
|
||||
use gateway_protocol::openai::events::{TokenUsage, retry_after_seconds_from_body};
|
||||
use reqwest::StatusCode;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{
|
||||
CodexBackendClient, CodexClientError, CodexClientResult, CodexRequestContext,
|
||||
CodexUpstreamDiagnostics,
|
||||
client::{read_capped_response_body, retry_after_seconds, truncate_for_error},
|
||||
endpoints::usage_endpoint_urls,
|
||||
endpoints::usage_endpoint_url,
|
||||
response_meta,
|
||||
};
|
||||
|
||||
const LONG_CONTEXT_THRESHOLD: u64 = 272_000;
|
||||
const WEB_SEARCH_CALL_TICKS: u128 = 100_000_000;
|
||||
const WEB_SEARCH_PREVIEW_NON_REASONING_CALL_TICKS: u128 = 250_000_000;
|
||||
|
||||
/// OpenAI 公开 Token 价格计算所需的单次用量事实。
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct OpenAiBillingUsage {
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cached_tokens: u64,
|
||||
cache_write_tokens: u64,
|
||||
image_input_tokens: u64,
|
||||
image_output_tokens: u64,
|
||||
web_search_calls: u64,
|
||||
web_search_pricing: Option<WebSearchPricing>,
|
||||
}
|
||||
|
||||
impl OpenAiBillingUsage {
|
||||
/// 构造不含托管工具费用的 Token 用量。
|
||||
#[must_use]
|
||||
pub const fn new(
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cached_tokens: u64,
|
||||
cache_write_tokens: u64,
|
||||
) -> Self {
|
||||
Self {
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cached_tokens,
|
||||
cache_write_tokens,
|
||||
image_input_tokens: 0,
|
||||
image_output_tokens: 0,
|
||||
web_search_calls: 0,
|
||||
web_search_pricing: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) const fn with_web_search_calls(
|
||||
mut self,
|
||||
calls: u64,
|
||||
pricing: Option<WebSearchPricing>,
|
||||
) -> Self {
|
||||
self.web_search_calls = calls;
|
||||
self.web_search_pricing = pricing;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl From<TokenUsage> for OpenAiBillingUsage {
|
||||
fn from(usage: TokenUsage) -> Self {
|
||||
Self {
|
||||
input_tokens: usage.input_tokens,
|
||||
output_tokens: usage.output_tokens,
|
||||
cached_tokens: usage.cached_tokens,
|
||||
cache_write_tokens: usage.cache_write_tokens,
|
||||
image_input_tokens: usage.image_input_tokens,
|
||||
image_output_tokens: usage.image_output_tokens,
|
||||
web_search_calls: 0,
|
||||
web_search_pricing: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum WebSearchPricing {
|
||||
Standard,
|
||||
PreviewNonReasoning,
|
||||
}
|
||||
|
||||
impl WebSearchPricing {
|
||||
const fn price_per_call_ticks(self) -> u128 {
|
||||
match self {
|
||||
Self::Standard => WEB_SEARCH_CALL_TICKS,
|
||||
Self::PreviewNonReasoning => WEB_SEARCH_PREVIEW_NON_REASONING_CALL_TICKS,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
struct TokenRates {
|
||||
@@ -43,9 +119,11 @@ impl TokenRates {
|
||||
#[derive(Clone, Copy)]
|
||||
struct ModelPricing {
|
||||
standard: TokenRates,
|
||||
priority: TokenRates,
|
||||
long: TokenRates,
|
||||
long_priority: TokenRates,
|
||||
flex: TokenRates,
|
||||
fast: TokenRates,
|
||||
long_standard: TokenRates,
|
||||
long_flex: TokenRates,
|
||||
long_fast: TokenRates,
|
||||
cache_write_percent: u32,
|
||||
}
|
||||
|
||||
@@ -53,29 +131,37 @@ impl ModelPricing {
|
||||
const fn new(input: u128, output: u128, cache_read: u128) -> Self {
|
||||
Self {
|
||||
standard: TokenRates::new(input, output, cache_read),
|
||||
priority: TokenRates::ZERO,
|
||||
long: TokenRates::ZERO,
|
||||
long_priority: TokenRates::ZERO,
|
||||
flex: TokenRates::ZERO,
|
||||
fast: TokenRates::ZERO,
|
||||
long_standard: TokenRates::ZERO,
|
||||
long_flex: TokenRates::ZERO,
|
||||
long_fast: TokenRates::ZERO,
|
||||
cache_write_percent: 0,
|
||||
}
|
||||
}
|
||||
|
||||
const fn with_priority(mut self, input: u128, output: u128, cache_read: u128) -> Self {
|
||||
self.priority = TokenRates::new(input, output, cache_read);
|
||||
const fn with_flex(mut self, input: u128, output: u128, cache_read: u128) -> Self {
|
||||
self.flex = TokenRates::new(input, output, cache_read);
|
||||
self
|
||||
}
|
||||
|
||||
const fn with_long_context(
|
||||
mut self,
|
||||
input: u128,
|
||||
output: u128,
|
||||
cache_read: u128,
|
||||
priority_input: u128,
|
||||
priority_output: u128,
|
||||
priority_cache_read: u128,
|
||||
) -> Self {
|
||||
self.long = TokenRates::new(input, output, cache_read);
|
||||
self.long_priority = TokenRates::new(priority_input, priority_output, priority_cache_read);
|
||||
const fn with_fast(mut self, input: u128, output: u128, cache_read: u128) -> Self {
|
||||
self.fast = TokenRates::new(input, output, cache_read);
|
||||
self
|
||||
}
|
||||
|
||||
const fn with_long(mut self, input: u128, output: u128, cache_read: u128) -> Self {
|
||||
self.long_standard = TokenRates::new(input, output, cache_read);
|
||||
self
|
||||
}
|
||||
|
||||
const fn with_long_flex(mut self, input: u128, output: u128, cache_read: u128) -> Self {
|
||||
self.long_flex = TokenRates::new(input, output, cache_read);
|
||||
self
|
||||
}
|
||||
|
||||
const fn with_long_fast(mut self, input: u128, output: u128, cache_read: u128) -> Self {
|
||||
self.long_fast = TokenRates::new(input, output, cache_read);
|
||||
self
|
||||
}
|
||||
|
||||
@@ -83,6 +169,29 @@ impl ModelPricing {
|
||||
self.cache_write_percent = percent;
|
||||
self
|
||||
}
|
||||
|
||||
fn rates(self, tier: PricingTier, long_context: bool) -> Option<TokenRates> {
|
||||
// Only models with a published long-context column switch at the threshold.
|
||||
// A dash in that column is not a second, unavailable price tier for models
|
||||
// whose entire supported context is covered by the short-context price.
|
||||
let uses_long_rates = long_context && self.long_standard.is_configured();
|
||||
let rates = match (tier, uses_long_rates) {
|
||||
(PricingTier::Standard, false) => self.standard,
|
||||
(PricingTier::Flex, false) => self.flex,
|
||||
(PricingTier::Fast, false) => self.fast,
|
||||
(PricingTier::Standard, true) => self.long_standard,
|
||||
(PricingTier::Flex, true) => self.long_flex,
|
||||
(PricingTier::Fast, true) => self.long_fast,
|
||||
};
|
||||
rates.is_configured().then_some(rates)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum PricingTier {
|
||||
Standard,
|
||||
Flex,
|
||||
Fast,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
@@ -91,86 +200,178 @@ struct PricingRule {
|
||||
pricing: ModelPricing,
|
||||
}
|
||||
|
||||
// Source: https://developers.openai.com/api/docs/pricing (verified 2026-08-14).
|
||||
const PRICING_RULES: &[PricingRule] = &[
|
||||
PricingRule {
|
||||
model: "gpt-5.6-sol",
|
||||
pricing: ModelPricing::new(50_000, 300_000, 5_000)
|
||||
.with_cache_write(125)
|
||||
.with_priority(100_000, 600_000, 10_000)
|
||||
.with_long_context(100_000, 450_000, 10_000, 0, 0, 0),
|
||||
.with_flex(25_000, 150_000, 2_500)
|
||||
.with_fast(100_000, 600_000, 10_000)
|
||||
.with_long(100_000, 450_000, 10_000)
|
||||
.with_long_flex(50_000, 225_000, 5_000)
|
||||
.with_long_fast(200_000, 900_000, 20_000),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.6-terra",
|
||||
pricing: ModelPricing::new(20_000, 120_000, 2_000)
|
||||
.with_cache_write(125)
|
||||
.with_priority(40_000, 240_000, 4_000)
|
||||
.with_long_context(40_000, 180_000, 4_000, 0, 0, 0),
|
||||
.with_flex(10_000, 60_000, 1_000)
|
||||
.with_fast(40_000, 240_000, 4_000)
|
||||
.with_long(40_000, 180_000, 4_000)
|
||||
.with_long_flex(20_000, 90_000, 2_000)
|
||||
.with_long_fast(80_000, 360_000, 8_000),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.6-luna",
|
||||
pricing: ModelPricing::new(2_000, 12_000, 200)
|
||||
.with_cache_write(125)
|
||||
.with_priority(4_000, 24_000, 400)
|
||||
.with_long_context(4_000, 18_000, 400, 0, 0, 0),
|
||||
.with_flex(1_000, 6_000, 100)
|
||||
.with_fast(4_000, 24_000, 400)
|
||||
.with_long(4_000, 18_000, 400)
|
||||
.with_long_flex(2_000, 9_000, 200)
|
||||
.with_long_fast(8_000, 36_000, 800),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.6",
|
||||
pricing: ModelPricing::new(50_000, 300_000, 5_000)
|
||||
.with_cache_write(125)
|
||||
.with_priority(100_000, 600_000, 10_000)
|
||||
.with_long_context(100_000, 450_000, 10_000, 0, 0, 0),
|
||||
.with_flex(25_000, 150_000, 2_500)
|
||||
.with_fast(100_000, 600_000, 10_000)
|
||||
.with_long(100_000, 450_000, 10_000)
|
||||
.with_long_flex(50_000, 225_000, 5_000)
|
||||
.with_long_fast(200_000, 900_000, 20_000),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.5-pro",
|
||||
pricing: ModelPricing::new(300_000, 1_800_000, 0)
|
||||
.with_priority(750_000, 4_500_000, 0)
|
||||
.with_long_context(600_000, 2_700_000, 0, 1_500_000, 6_750_000, 0),
|
||||
.with_flex(150_000, 900_000, 0)
|
||||
.with_long(600_000, 2_700_000, 0),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.5",
|
||||
pricing: ModelPricing::new(50_000, 300_000, 5_000)
|
||||
.with_priority(125_000, 750_000, 12_500)
|
||||
.with_long_context(100_000, 450_000, 10_000, 250_000, 1_125_000, 25_000),
|
||||
.with_flex(25_000, 150_000, 2_500)
|
||||
.with_fast(125_000, 750_000, 12_500)
|
||||
.with_long(100_000, 450_000, 10_000)
|
||||
.with_long_flex(50_000, 225_000, 5_000),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.4-mini",
|
||||
pricing: ModelPricing::new(7_500, 45_000, 750),
|
||||
pricing: ModelPricing::new(7_500, 45_000, 750)
|
||||
.with_flex(3_750, 22_500, 375)
|
||||
.with_fast(15_000, 90_000, 1_500),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.4-nano",
|
||||
pricing: ModelPricing::new(2_000, 12_500, 200),
|
||||
pricing: ModelPricing::new(2_000, 12_500, 200).with_flex(1_000, 6_250, 100),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.4-pro",
|
||||
pricing: ModelPricing::new(300_000, 1_800_000, 0)
|
||||
.with_priority(750_000, 4_500_000, 0)
|
||||
.with_long_context(600_000, 2_700_000, 0, 1_500_000, 6_750_000, 0),
|
||||
.with_flex(150_000, 900_000, 0)
|
||||
.with_long(600_000, 2_700_000, 0)
|
||||
.with_long_flex(300_000, 1_350_000, 0),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.4",
|
||||
pricing: ModelPricing::new(25_000, 150_000, 2_500)
|
||||
.with_priority(50_000, 300_000, 5_000)
|
||||
.with_long_context(50_000, 225_000, 5_000, 100_000, 450_000, 10_000),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.3-codex-spark",
|
||||
pricing: ModelPricing::new(12_500, 100_000, 1_250).with_priority(25_000, 200_000, 2_500),
|
||||
.with_flex(12_500, 75_000, 1_300)
|
||||
.with_fast(50_000, 300_000, 5_000)
|
||||
.with_long(50_000, 225_000, 5_000)
|
||||
.with_long_flex(25_000, 112_500, 2_500),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.3-codex",
|
||||
pricing: ModelPricing::new(17_500, 140_000, 1_750).with_priority(35_000, 280_000, 3_500),
|
||||
pricing: ModelPricing::new(17_500, 140_000, 1_750).with_fast(35_000, 280_000, 3_500),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.2-pro",
|
||||
pricing: ModelPricing::new(210_000, 1_680_000, 0),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.2",
|
||||
pricing: ModelPricing::new(17_500, 140_000, 1_750).with_priority(35_000, 280_000, 3_500),
|
||||
pricing: ModelPricing::new(17_500, 140_000, 1_750)
|
||||
.with_flex(8_750, 70_000, 875)
|
||||
.with_fast(35_000, 280_000, 3_500),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5.1",
|
||||
pricing: ModelPricing::new(12_500, 100_000, 1_250)
|
||||
.with_flex(6_250, 50_000, 625)
|
||||
.with_fast(25_000, 200_000, 2_500),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5-mini",
|
||||
pricing: ModelPricing::new(2_500, 20_000, 250)
|
||||
.with_flex(1_250, 10_000, 125)
|
||||
.with_fast(4_500, 36_000, 450),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5-nano",
|
||||
pricing: ModelPricing::new(500, 4_000, 50).with_flex(250, 2_000, 25),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5-pro",
|
||||
pricing: ModelPricing::new(150_000, 1_200_000, 0),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-5",
|
||||
pricing: ModelPricing::new(12_500, 100_000, 1_250)
|
||||
.with_flex(6_250, 50_000, 625)
|
||||
.with_fast(25_000, 200_000, 2_500),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-4.1-mini",
|
||||
pricing: ModelPricing::new(4_000, 16_000, 1_000).with_fast(7_000, 28_000, 1_750),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-4.1-nano",
|
||||
pricing: ModelPricing::new(1_000, 4_000, 250).with_fast(2_000, 8_000, 500),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-4.1",
|
||||
pricing: ModelPricing::new(20_000, 80_000, 5_000).with_fast(35_000, 140_000, 8_750),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-4o-2024-05-13",
|
||||
pricing: ModelPricing::new(50_000, 150_000, 0).with_fast(87_500, 262_500, 0),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-4o-mini",
|
||||
pricing: ModelPricing::new(1_500, 6_000, 750),
|
||||
pricing: ModelPricing::new(1_500, 6_000, 750).with_fast(2_500, 10_000, 1_250),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-4o",
|
||||
pricing: ModelPricing::new(25_000, 100_000, 12_500),
|
||||
pricing: ModelPricing::new(25_000, 100_000, 12_500).with_fast(42_500, 170_000, 21_250),
|
||||
},
|
||||
PricingRule {
|
||||
model: "o1-pro",
|
||||
pricing: ModelPricing::new(1_500_000, 6_000_000, 0),
|
||||
},
|
||||
PricingRule {
|
||||
model: "o1",
|
||||
pricing: ModelPricing::new(150_000, 600_000, 75_000),
|
||||
},
|
||||
PricingRule {
|
||||
model: "o3-pro",
|
||||
pricing: ModelPricing::new(200_000, 800_000, 0),
|
||||
},
|
||||
PricingRule {
|
||||
model: "o3-mini",
|
||||
pricing: ModelPricing::new(11_000, 44_000, 5_500),
|
||||
},
|
||||
PricingRule {
|
||||
model: "o3",
|
||||
pricing: ModelPricing::new(20_000, 80_000, 5_000)
|
||||
.with_flex(10_000, 40_000, 2_500)
|
||||
.with_fast(35_000, 140_000, 8_750),
|
||||
},
|
||||
PricingRule {
|
||||
model: "o4-mini",
|
||||
pricing: ModelPricing::new(11_000, 44_000, 2_750)
|
||||
.with_flex(5_500, 22_000, 1_380)
|
||||
.with_fast(20_000, 80_000, 5_000),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-4-turbo",
|
||||
@@ -180,93 +381,86 @@ const PRICING_RULES: &[PricingRule] = &[
|
||||
model: "gpt-4",
|
||||
pricing: ModelPricing::new(300_000, 600_000, 0),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-3.5-turbo-instruct",
|
||||
pricing: ModelPricing::new(15_000, 20_000, 0),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-3.5-turbo-1106",
|
||||
pricing: ModelPricing::new(10_000, 20_000, 0),
|
||||
},
|
||||
PricingRule {
|
||||
model: "gpt-3.5-turbo",
|
||||
pricing: ModelPricing::new(5_000, 15_000, 0),
|
||||
},
|
||||
];
|
||||
|
||||
const UNPRICED_MODELS: &[&str] = &["gpt-5.3-codex-spark"];
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
struct TokenAmounts {
|
||||
input_ticks: u128,
|
||||
output_ticks: u128,
|
||||
cache_read_ticks: u128,
|
||||
cache_write_ticks: u128,
|
||||
total_ticks: u128,
|
||||
}
|
||||
|
||||
/// 按 OpenAI Provider 当前受控价格规则计算费用明细。
|
||||
#[must_use]
|
||||
pub fn openai_billing_breakdown(
|
||||
model: &str,
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cached_tokens: u64,
|
||||
cache_write_tokens: u64,
|
||||
usage: OpenAiBillingUsage,
|
||||
service_tier: Option<&str>,
|
||||
) -> Option<CalculatedCostBreakdown> {
|
||||
if cached_tokens > input_tokens {
|
||||
if usage.cached_tokens > usage.input_tokens
|
||||
|| usage.image_input_tokens > 0
|
||||
|| usage.image_output_tokens > 0
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let web_search_ticks = web_search_amount_ticks(usage)?;
|
||||
let pricing = model_pricing(model)?;
|
||||
let long_context = input_tokens > LONG_CONTEXT_THRESHOLD;
|
||||
let long_context = usage.input_tokens > LONG_CONTEXT_THRESHOLD;
|
||||
let normalized_tier = normalize_service_tier(service_tier);
|
||||
let priority = matches!(normalized_tier.as_deref(), Some("priority" | "fast"));
|
||||
let mut rates = if long_context && pricing.long.is_configured() {
|
||||
pricing.long
|
||||
} else {
|
||||
pricing.standard
|
||||
};
|
||||
let multiplier_percent = if priority {
|
||||
let priority_rates = if long_context && pricing.long_priority.is_configured() {
|
||||
pricing.long_priority
|
||||
} else {
|
||||
pricing.priority
|
||||
};
|
||||
if priority_rates.is_configured() {
|
||||
rates = priority_rates;
|
||||
100
|
||||
} else {
|
||||
200
|
||||
}
|
||||
} else if normalized_tier.as_deref() == Some("flex") {
|
||||
50
|
||||
} else {
|
||||
100
|
||||
};
|
||||
|
||||
let billed_cache_read = if rates.cache_read_ticks > 0 {
|
||||
cached_tokens.min(input_tokens)
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let cache_write_rate = if pricing.cache_write_percent > 0 {
|
||||
apply_percent(rates.input_ticks, pricing.cache_write_percent)?
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let billed_cache_write = if cache_write_rate > 0 {
|
||||
cache_write_tokens.min(input_tokens.saturating_sub(billed_cache_read))
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let uncached_input = input_tokens
|
||||
.saturating_sub(billed_cache_read)
|
||||
.saturating_sub(billed_cache_write);
|
||||
let input_amount = u128::from(uncached_input).checked_mul(rates.input_ticks)?;
|
||||
let output_amount = u128::from(output_tokens).checked_mul(rates.output_ticks)?;
|
||||
let cache_read_amount = u128::from(billed_cache_read).checked_mul(rates.cache_read_ticks)?;
|
||||
let cache_write_amount = u128::from(billed_cache_write).checked_mul(cache_write_rate)?;
|
||||
let standard_amount = input_amount
|
||||
.checked_add(output_amount)?
|
||||
.checked_add(cache_read_amount)?
|
||||
.checked_add(cache_write_amount)?;
|
||||
let total_amount = apply_percent(standard_amount, multiplier_percent)?;
|
||||
let tier = pricing_tier(normalized_tier.as_deref())?;
|
||||
let standard_rates = pricing.rates(PricingTier::Standard, long_context)?;
|
||||
let selected_rates = pricing.rates(tier, long_context)?;
|
||||
let mut standard = token_amounts(
|
||||
standard_rates,
|
||||
pricing.cache_write_percent,
|
||||
usage.input_tokens,
|
||||
usage.output_tokens,
|
||||
usage.cached_tokens,
|
||||
usage.cache_write_tokens,
|
||||
)?;
|
||||
let mut selected = token_amounts(
|
||||
selected_rates,
|
||||
pricing.cache_write_percent,
|
||||
usage.input_tokens,
|
||||
usage.output_tokens,
|
||||
usage.cached_tokens,
|
||||
usage.cache_write_tokens,
|
||||
)?;
|
||||
standard.total_ticks = standard.total_ticks.checked_add(web_search_ticks)?;
|
||||
selected.total_ticks = selected.total_ticks.checked_add(web_search_ticks)?;
|
||||
let multiplier_percent =
|
||||
effective_multiplier_percent(selected.total_ticks, standard.total_ticks)?;
|
||||
let cache_write_rate = cache_write_rate(selected_rates, pricing.cache_write_percent)?;
|
||||
|
||||
Some(CalculatedCostBreakdown::new(
|
||||
CalculatedCostAmounts::new(
|
||||
usd_money(input_amount)?,
|
||||
usd_money(output_amount)?,
|
||||
usd_money(cache_read_amount)?,
|
||||
usd_money(cache_write_amount)?,
|
||||
usd_money(standard_amount)?,
|
||||
usd_money(total_amount)?,
|
||||
usd_money(selected.input_ticks)?,
|
||||
usd_money(selected.output_ticks)?,
|
||||
usd_money(selected.cache_read_ticks)?,
|
||||
usd_money(selected.cache_write_ticks)?,
|
||||
usd_money(standard.total_ticks)?,
|
||||
usd_money(selected.total_ticks)?,
|
||||
),
|
||||
CalculatedCostRates::new(
|
||||
usd_price_per_million(rates.input_ticks)?,
|
||||
usd_price_per_million(rates.output_ticks)?,
|
||||
usd_price_per_million(rates.cache_read_ticks)?,
|
||||
usd_price_per_million(selected_rates.input_ticks)?,
|
||||
usd_price_per_million(selected_rates.output_ticks)?,
|
||||
usd_price_per_million(selected_rates.cache_read_ticks)?,
|
||||
usd_price_per_million(cache_write_rate)?,
|
||||
),
|
||||
Some(normalized_tier.unwrap_or_else(|| "default".to_owned())),
|
||||
@@ -276,6 +470,12 @@ pub fn openai_billing_breakdown(
|
||||
|
||||
fn model_pricing(model: &str) -> Option<ModelPricing> {
|
||||
let normalized = normalize_model_name(model);
|
||||
if UNPRICED_MODELS
|
||||
.iter()
|
||||
.any(|rule| model_matches_rule(&normalized, rule))
|
||||
{
|
||||
return None;
|
||||
}
|
||||
PRICING_RULES
|
||||
.iter()
|
||||
.filter(|rule| model_matches_rule(&normalized, rule.model))
|
||||
@@ -283,6 +483,120 @@ fn model_pricing(model: &str) -> Option<ModelPricing> {
|
||||
.map(|rule| rule.pricing)
|
||||
}
|
||||
|
||||
pub(crate) fn web_search_pricing(model: &str, tools: Option<&[Value]>) -> Option<WebSearchPricing> {
|
||||
let mut standard = false;
|
||||
let mut preview = false;
|
||||
for tool_type in tools
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(|tool| tool.get("type").and_then(Value::as_str))
|
||||
{
|
||||
if tool_type == "web_search_preview" || tool_type.starts_with("web_search_preview_") {
|
||||
preview = true;
|
||||
} else if tool_type == "web_search" || tool_type.starts_with("web_search_") {
|
||||
standard = true;
|
||||
}
|
||||
}
|
||||
match (standard, preview) {
|
||||
// These two models bill non-preview search content as a fixed 8K input block.
|
||||
// The response does not expose enough detail to replace that block without
|
||||
// double-counting ordinary input, so leave the request unpriced.
|
||||
(true, false) if fixed_block_web_search_model(model) => None,
|
||||
(true, false) => Some(WebSearchPricing::Standard),
|
||||
(false, true) if reasoning_model(model) => Some(WebSearchPricing::Standard),
|
||||
(false, true) => Some(WebSearchPricing::PreviewNonReasoning),
|
||||
(false, false) | (true, true) => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn fixed_block_web_search_model(model: &str) -> bool {
|
||||
let normalized = normalize_model_name(model);
|
||||
["gpt-4o-mini", "gpt-4.1-mini"]
|
||||
.iter()
|
||||
.any(|rule| model_matches_rule(&normalized, rule))
|
||||
}
|
||||
|
||||
fn reasoning_model(model: &str) -> bool {
|
||||
let normalized = normalize_model_name(model);
|
||||
normalized.starts_with("gpt-5")
|
||||
|| normalized.starts_with("o1")
|
||||
|| normalized.starts_with("o3")
|
||||
|| normalized.starts_with("o4")
|
||||
}
|
||||
|
||||
fn pricing_tier(service_tier: Option<&str>) -> Option<PricingTier> {
|
||||
match service_tier {
|
||||
None | Some("auto" | "default" | "standard") => Some(PricingTier::Standard),
|
||||
Some("flex") => Some(PricingTier::Flex),
|
||||
Some("fast" | "priority") => Some(PricingTier::Fast),
|
||||
Some(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn web_search_amount_ticks(usage: OpenAiBillingUsage) -> Option<u128> {
|
||||
if usage.web_search_calls == 0 {
|
||||
return Some(0);
|
||||
}
|
||||
u128::from(usage.web_search_calls).checked_mul(usage.web_search_pricing?.price_per_call_ticks())
|
||||
}
|
||||
|
||||
fn token_amounts(
|
||||
rates: TokenRates,
|
||||
cache_write_percent: u32,
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cached_tokens: u64,
|
||||
cache_write_tokens: u64,
|
||||
) -> Option<TokenAmounts> {
|
||||
let billed_cache_read = if rates.cache_read_ticks > 0 {
|
||||
cached_tokens.min(input_tokens)
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let cache_write_rate = cache_write_rate(rates, cache_write_percent)?;
|
||||
let billed_cache_write = if cache_write_rate > 0 {
|
||||
cache_write_tokens.min(input_tokens.saturating_sub(billed_cache_read))
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let uncached_input = input_tokens
|
||||
.saturating_sub(billed_cache_read)
|
||||
.saturating_sub(billed_cache_write);
|
||||
let input_ticks = u128::from(uncached_input).checked_mul(rates.input_ticks)?;
|
||||
let output_ticks = u128::from(output_tokens).checked_mul(rates.output_ticks)?;
|
||||
let cache_read_ticks = u128::from(billed_cache_read).checked_mul(rates.cache_read_ticks)?;
|
||||
let cache_write_ticks = u128::from(billed_cache_write).checked_mul(cache_write_rate)?;
|
||||
let total_ticks = input_ticks
|
||||
.checked_add(output_ticks)?
|
||||
.checked_add(cache_read_ticks)?
|
||||
.checked_add(cache_write_ticks)?;
|
||||
Some(TokenAmounts {
|
||||
input_ticks,
|
||||
output_ticks,
|
||||
cache_read_ticks,
|
||||
cache_write_ticks,
|
||||
total_ticks,
|
||||
})
|
||||
}
|
||||
|
||||
fn cache_write_rate(rates: TokenRates, percent: u32) -> Option<u128> {
|
||||
if percent == 0 {
|
||||
return Some(0);
|
||||
}
|
||||
apply_percent(rates.input_ticks, percent)
|
||||
}
|
||||
|
||||
fn effective_multiplier_percent(total: u128, standard: u128) -> Option<u32> {
|
||||
if standard == 0 {
|
||||
return Some(100);
|
||||
}
|
||||
let rounded = total
|
||||
.checked_mul(100)?
|
||||
.checked_add(standard / 2)?
|
||||
.checked_div(standard)?;
|
||||
u32::try_from(rounded).ok()
|
||||
}
|
||||
|
||||
fn normalize_model_name(model: &str) -> String {
|
||||
model
|
||||
.trim()
|
||||
@@ -336,78 +650,68 @@ pub const MAX_CODEX_USAGE_BODY_BYTES: usize = 1024 * 1024;
|
||||
impl CodexBackendClient {
|
||||
/// 获取 Codex usage JSON。
|
||||
pub async fn fetch_usage(&self, context: CodexRequestContext<'_>) -> CodexClientResult<Value> {
|
||||
let mut last_invalid_body = None;
|
||||
let headers = self.usage_request_headers(context)?;
|
||||
let response = self
|
||||
.client
|
||||
.get(usage_endpoint_url(&self.base_url))
|
||||
.headers(headers)
|
||||
.send()
|
||||
.await?;
|
||||
let status = response.status();
|
||||
let diagnostics = response_meta::diagnostics(Some(status.as_u16()), response.headers());
|
||||
let retry_after_seconds = retry_after_seconds(response.headers(), None);
|
||||
let body = read_capped_response_body(response, MAX_CODEX_USAGE_BODY_BYTES).await?;
|
||||
if body.limit_exceeded() {
|
||||
return Err(CodexClientError::Upstream {
|
||||
status: if status.is_success() {
|
||||
StatusCode::BAD_GATEWAY
|
||||
} else {
|
||||
status
|
||||
},
|
||||
retry_after_seconds,
|
||||
body: "upstream usage response exceeded the body limit".to_owned(),
|
||||
client_response: None,
|
||||
diagnostics: Box::new(diagnostics),
|
||||
set_cookie_headers: Vec::new(),
|
||||
rate_limit_headers: Vec::new(),
|
||||
transport: super::client::CodexBackendTransport::HttpSse,
|
||||
transport_metrics: Box::default(),
|
||||
send_phase: super::diagnostics::CodexUpstreamSendPhase::AfterPayload,
|
||||
});
|
||||
}
|
||||
let body = body.into_string();
|
||||
|
||||
for endpoint in usage_endpoint_urls(&self.base_url) {
|
||||
let headers = self.usage_request_headers(context)?;
|
||||
let response = self.client.get(endpoint).headers(headers).send().await?;
|
||||
let status = response.status();
|
||||
let diagnostics = response_meta::diagnostics(Some(status.as_u16()), response.headers());
|
||||
let retry_after_seconds = retry_after_seconds(response.headers(), None);
|
||||
let body = read_capped_response_body(response, MAX_CODEX_USAGE_BODY_BYTES).await?;
|
||||
if body.limit_exceeded() {
|
||||
return Err(CodexClientError::Upstream {
|
||||
status: if status.is_success() {
|
||||
StatusCode::BAD_GATEWAY
|
||||
} else {
|
||||
status
|
||||
},
|
||||
retry_after_seconds,
|
||||
body: "upstream usage response exceeded the body limit".to_owned(),
|
||||
client_response: None,
|
||||
diagnostics: Box::new(diagnostics),
|
||||
set_cookie_headers: Vec::new(),
|
||||
rate_limit_headers: Vec::new(),
|
||||
transport: super::client::CodexBackendTransport::HttpSse,
|
||||
transport_metrics: Box::default(),
|
||||
send_phase: super::diagnostics::CodexUpstreamSendPhase::AfterPayload,
|
||||
});
|
||||
}
|
||||
let body = body.into_string();
|
||||
|
||||
if status == StatusCode::NOT_FOUND {
|
||||
last_invalid_body = Some(body);
|
||||
continue;
|
||||
}
|
||||
if !status.is_success() {
|
||||
return Err(CodexClientError::Upstream {
|
||||
status,
|
||||
retry_after_seconds: retry_after_seconds
|
||||
.or_else(|| retry_after_seconds_from_body(&body)),
|
||||
body,
|
||||
client_response: None,
|
||||
diagnostics: Box::new(diagnostics),
|
||||
set_cookie_headers: Vec::new(),
|
||||
rate_limit_headers: Vec::new(),
|
||||
transport: super::client::CodexBackendTransport::HttpSse,
|
||||
transport_metrics: Box::default(),
|
||||
send_phase: super::diagnostics::CodexUpstreamSendPhase::AfterPayload,
|
||||
});
|
||||
}
|
||||
|
||||
match serde_json::from_str::<Value>(&body) {
|
||||
Ok(parsed) if is_usage_response(&parsed) => return Ok(parsed),
|
||||
_ => last_invalid_body = Some(body),
|
||||
}
|
||||
if !status.is_success() {
|
||||
return Err(CodexClientError::Upstream {
|
||||
status,
|
||||
retry_after_seconds: retry_after_seconds
|
||||
.or_else(|| retry_after_seconds_from_body(&body)),
|
||||
body,
|
||||
client_response: None,
|
||||
diagnostics: Box::new(diagnostics),
|
||||
set_cookie_headers: Vec::new(),
|
||||
rate_limit_headers: Vec::new(),
|
||||
transport: super::client::CodexBackendTransport::HttpSse,
|
||||
transport_metrics: Box::default(),
|
||||
send_phase: super::diagnostics::CodexUpstreamSendPhase::AfterPayload,
|
||||
});
|
||||
}
|
||||
|
||||
Err(CodexClientError::Upstream {
|
||||
status: StatusCode::BAD_GATEWAY,
|
||||
retry_after_seconds: None,
|
||||
body: last_invalid_body.map_or_else(
|
||||
|| "usage endpoint is unavailable".to_string(),
|
||||
|body| format!("invalid usage response: {}", truncate_for_error(&body)),
|
||||
),
|
||||
client_response: None,
|
||||
diagnostics: Box::new(CodexUpstreamDiagnostics::with_status(
|
||||
StatusCode::BAD_GATEWAY.as_u16(),
|
||||
)),
|
||||
set_cookie_headers: Vec::new(),
|
||||
rate_limit_headers: Vec::new(),
|
||||
transport: super::client::CodexBackendTransport::HttpSse,
|
||||
transport_metrics: Box::default(),
|
||||
send_phase: super::diagnostics::CodexUpstreamSendPhase::AfterPayload,
|
||||
})
|
||||
match serde_json::from_str::<Value>(&body) {
|
||||
Ok(parsed) if is_usage_response(&parsed) => Ok(parsed),
|
||||
_ => Err(CodexClientError::Upstream {
|
||||
status: StatusCode::BAD_GATEWAY,
|
||||
retry_after_seconds: None,
|
||||
body: format!("invalid usage response: {}", truncate_for_error(&body)),
|
||||
client_response: None,
|
||||
diagnostics: Box::new(diagnostics),
|
||||
set_cookie_headers: Vec::new(),
|
||||
rate_limit_headers: Vec::new(),
|
||||
transport: super::client::CodexBackendTransport::HttpSse,
|
||||
transport_metrics: Box::default(),
|
||||
send_phase: super::diagnostics::CodexUpstreamSendPhase::AfterPayload,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -112,6 +112,7 @@ fn is_sensitive_opening_header(name: &str) -> bool {
|
||||
| "x-codex-turn-metadata"
|
||||
| "x-codex-turn-state"
|
||||
| "x-codex-parent-thread-id"
|
||||
| "x-oai-attestation"
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ use std::{pin::Pin, sync::Arc, time::Duration};
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures::Stream;
|
||||
use gateway_protocol::openai::events::ParsedRateLimits;
|
||||
use tokio::sync::Mutex;
|
||||
use uuid::Uuid;
|
||||
|
||||
@@ -36,8 +37,8 @@ pub struct CodexWebSocketStreamingExchange {
|
||||
pub set_cookie_headers: Vec<String>,
|
||||
/// 上游握手响应里的限流头。
|
||||
pub rate_limit_headers: Vec<(String, String)>,
|
||||
/// 上游内部 rate-limit 事件里的动态更新。
|
||||
pub rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates,
|
||||
/// 上游内部 `codex.rate_limits` 事件里的结构化动态更新。
|
||||
pub rate_limit_updates: CodexWebSocketRateLimitUpdates,
|
||||
/// 上游内部 metadata 事件里的动态 turn state。
|
||||
pub turn_state_update: CodexWebSocketTurnStateUpdate,
|
||||
/// WebSocket 连接池决策。
|
||||
@@ -53,8 +54,8 @@ pub struct CodexWebSocketStreamingExchange {
|
||||
/// Responses WebSocket live SSE 字节流。
|
||||
pub type CodexWebSocketSseStream =
|
||||
Pin<Box<dyn Stream<Item = Result<Bytes, CodexWebSocketExchangeError>> + Send + 'static>>;
|
||||
/// live 流中的限流头动态更新。
|
||||
pub type CodexWebSocketRateLimitHeaderUpdates = Arc<Mutex<Vec<(String, String)>>>;
|
||||
/// live 流中的结构化限流动态更新。
|
||||
pub type CodexWebSocketRateLimitUpdates = Arc<Mutex<Vec<ParsedRateLimits>>>;
|
||||
/// live 流中的 turn state 动态更新。
|
||||
pub type CodexWebSocketTurnStateUpdate = Arc<Mutex<Option<String>>>;
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ pub(in crate::transport::websocket) enum WebSocketTerminalKind {
|
||||
}
|
||||
|
||||
pub(super) enum ExchangeAction {
|
||||
RateLimits(Vec<(String, String)>),
|
||||
RateLimits(events::ParsedRateLimits),
|
||||
TurnState(String),
|
||||
Forward {
|
||||
frame: String,
|
||||
@@ -41,8 +41,8 @@ pub(super) fn reduce_websocket_event(
|
||||
};
|
||||
if let Some(parsed) = events::parse_rate_limits_event(&value) {
|
||||
let headers = events::rate_limits_to_header_pairs(&parsed);
|
||||
metadata.rate_limit_headers.extend(headers.iter().cloned());
|
||||
return Ok(ExchangeAction::RateLimits(headers));
|
||||
metadata.rate_limit_headers.extend(headers);
|
||||
return Ok(ExchangeAction::RateLimits(parsed));
|
||||
}
|
||||
|
||||
response_meta::merge_response_metadata(
|
||||
|
||||
@@ -16,9 +16,8 @@ use super::super::{
|
||||
use super::io::{next_websocket_message, receive_idle_timeout, reused_stream_receive_error};
|
||||
use super::reducer::{ExchangeAction, WebSocketTerminalKind, reduce_websocket_event};
|
||||
use super::{
|
||||
CodexWebSocketExchangeError, CodexWebSocketRateLimitHeaderUpdates,
|
||||
CodexWebSocketStreamingExchange, CodexWebSocketTurnStateUpdate, WEBSOCKET_STREAM_BUFFER,
|
||||
reusable_websocket_metadata,
|
||||
CodexWebSocketExchangeError, CodexWebSocketRateLimitUpdates, CodexWebSocketStreamingExchange,
|
||||
CodexWebSocketTurnStateUpdate, WEBSOCKET_STREAM_BUFFER, reusable_websocket_metadata,
|
||||
};
|
||||
|
||||
pub(in crate::transport::websocket) struct WebSocketStreamPoolReturn {
|
||||
@@ -63,8 +62,8 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
|
||||
) -> CodexWebSocketStreamingExchange {
|
||||
let websocket_connection_id = websocket.connection_id();
|
||||
let response_metadata = metadata.clone();
|
||||
let rate_limit_header_updates = Arc::new(Mutex::new(Vec::new()));
|
||||
let rate_limit_header_updates_for_task = Arc::clone(&rate_limit_header_updates);
|
||||
let rate_limit_updates = Arc::new(Mutex::new(Vec::new()));
|
||||
let rate_limit_updates_for_task = Arc::clone(&rate_limit_updates);
|
||||
let turn_state_update = Arc::new(Mutex::new(metadata.turn_state.clone()));
|
||||
let turn_state_update_for_task = Arc::clone(&turn_state_update);
|
||||
let (tx, rx) = mpsc::channel(WEBSOCKET_STREAM_BUFFER);
|
||||
@@ -83,7 +82,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
|
||||
reused_connection,
|
||||
initial_event_timeout,
|
||||
shutdown,
|
||||
rate_limit_header_updates: rate_limit_header_updates_for_task,
|
||||
rate_limit_updates: rate_limit_updates_for_task,
|
||||
turn_state_update: turn_state_update_for_task,
|
||||
tx,
|
||||
})
|
||||
@@ -105,7 +104,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response(
|
||||
turn_state: response_metadata.turn_state,
|
||||
set_cookie_headers: response_metadata.set_cookie_headers,
|
||||
rate_limit_headers: response_metadata.rate_limit_headers,
|
||||
rate_limit_header_updates,
|
||||
rate_limit_updates,
|
||||
turn_state_update,
|
||||
pool_decision: None,
|
||||
connection_local_continuation: false,
|
||||
@@ -121,7 +120,7 @@ struct WebSocketStreamForwardState {
|
||||
reused_connection: bool,
|
||||
initial_event_timeout: Option<Duration>,
|
||||
shutdown: CancellationToken,
|
||||
rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates,
|
||||
rate_limit_updates: CodexWebSocketRateLimitUpdates,
|
||||
turn_state_update: CodexWebSocketTurnStateUpdate,
|
||||
tx: mpsc::Sender<Result<Bytes, CodexWebSocketExchangeError>>,
|
||||
}
|
||||
@@ -134,7 +133,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
reused_connection,
|
||||
initial_event_timeout,
|
||||
shutdown,
|
||||
rate_limit_header_updates,
|
||||
rate_limit_updates,
|
||||
turn_state_update,
|
||||
tx,
|
||||
} = state;
|
||||
@@ -231,8 +230,8 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
|
||||
};
|
||||
let (frame, terminal) = match reduce_websocket_event(&raw, &mut metadata, &mut continuation)
|
||||
{
|
||||
Ok(ExchangeAction::RateLimits(headers)) => {
|
||||
rate_limit_header_updates.lock().await.extend(headers);
|
||||
Ok(ExchangeAction::RateLimits(rate_limits)) => {
|
||||
rate_limit_updates.lock().await.push(rate_limits);
|
||||
continue;
|
||||
}
|
||||
Ok(ExchangeAction::TurnState(turn_state)) => {
|
||||
|
||||
@@ -21,8 +21,8 @@ pub use self::{
|
||||
},
|
||||
error::{CodexWebSocketCloseError, CodexWebSocketExchangeError, CodexWebSocketUpstreamError},
|
||||
exchange::{
|
||||
CodexWebSocketRateLimitHeaderUpdates, CodexWebSocketSseStream,
|
||||
CodexWebSocketStreamingExchange, CodexWebSocketTurnStateUpdate,
|
||||
CodexWebSocketRateLimitUpdates, CodexWebSocketSseStream, CodexWebSocketStreamingExchange,
|
||||
CodexWebSocketTurnStateUpdate,
|
||||
},
|
||||
handshake::responses_websocket_endpoint,
|
||||
model::{
|
||||
|
||||
@@ -3,6 +3,7 @@ use std::num::NonZeroU32;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
use chrono::{DateTime, TimeZone as _, Utc};
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_admin::model::accounts::AccountRecord;
|
||||
@@ -13,7 +14,7 @@ use gateway_admin::model::provider_credentials::{
|
||||
AuthorizationMutationTarget, AuthorizationOwnerBinding, CompleteAuthorization,
|
||||
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
|
||||
PrepareCredentialRotation, ProviderDocument, ProviderExportCredentialInput,
|
||||
ProviderQuotaRequest, QuotaLocalUsageAttribution,
|
||||
ProviderQuotaRequest, ProviderQuotaWindowRole, QuotaLocalUsageAttribution,
|
||||
};
|
||||
use gateway_admin::model::{MutationActor, MutationContext, Revision};
|
||||
use gateway_admin::ports::provider::ProviderAdminErrorKind;
|
||||
@@ -25,15 +26,15 @@ use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
|
||||
use gateway_core::provider_ports::{
|
||||
NewOAuthPendingFlow, OAuthPendingClaimOutcome, OAuthPendingConsumeOutcome,
|
||||
OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingReleaseOutcome,
|
||||
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCooldown, ProviderCooldownPort,
|
||||
ProviderCooldownScope, ProviderCredentialState, ProviderCredentialStatePort,
|
||||
ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderScopedCooldown, ProviderStoreError,
|
||||
ProviderStorePorts,
|
||||
ProviderArtifactProfile, ProviderArtifactProfileCachePort, ProviderCatalogCacheKey,
|
||||
ProviderCatalogCachePort, ProviderCooldown, ProviderCooldownPort, ProviderCooldownScope,
|
||||
ProviderCredentialState, ProviderCredentialStatePort, ProviderRefreshPolicy,
|
||||
ProviderRuntimePolicyPort, ProviderScopedCooldown, ProviderStoreError, ProviderStorePorts,
|
||||
};
|
||||
use gateway_core::routing::{ProviderKind, UpstreamModelId};
|
||||
use gateway_core::task::{WorkerContribution, WorkerKind, WorkerRunnable};
|
||||
use provider_openai::config::{CodexWireProfileConfig, OpenAiConfig};
|
||||
use provider_openai::credential::ImportCodexOAuthCredential;
|
||||
use provider_openai::credential::{CodexCredentialCodec, ImportCodexOAuthCredential};
|
||||
use provider_openai::transport::profile::APPCAST_POLL_INTERVAL;
|
||||
use secrecy::SecretString;
|
||||
use serde_json::{Map, Value, json};
|
||||
@@ -54,7 +55,7 @@ async fn openai_bundle_exposes_one_core_provider_and_drains_worker_contributions
|
||||
assert_eq!(bundle.core_provider().name(), "openai");
|
||||
assert_eq!(bundle.admin_provider().provider_kind().as_str(), "openai");
|
||||
let contributions = bundle.take_worker_contributions();
|
||||
assert_eq!(contributions.len(), 5);
|
||||
assert_eq!(contributions.len(), 4);
|
||||
assert!(
|
||||
contributions
|
||||
.iter()
|
||||
@@ -81,22 +82,6 @@ async fn openai_bundle_exposes_one_core_provider_and_drains_worker_contributions
|
||||
panic!("Desktop release worker must be scheduled");
|
||||
};
|
||||
assert_eq!(schedule.interval(), APPCAST_POLL_INTERVAL);
|
||||
let cli_release_worker = contributions
|
||||
.iter()
|
||||
.find_map(|contribution| match contribution {
|
||||
WorkerContribution::Registration(registration)
|
||||
if registration.id.owner() == "openai-cli-release" =>
|
||||
{
|
||||
Some(registration)
|
||||
}
|
||||
WorkerContribution::Registration(_) | WorkerContribution::Disabled { .. } => None,
|
||||
})
|
||||
.expect("CLI release worker");
|
||||
assert_eq!(cli_release_worker.id.kind(), WorkerKind::QuotaCatalogHealth);
|
||||
let WorkerRunnable::Scheduled { schedule, .. } = &cli_release_worker.runnable else {
|
||||
panic!("CLI release worker must be scheduled");
|
||||
};
|
||||
assert_eq!(schedule.interval(), APPCAST_POLL_INTERVAL);
|
||||
assert!(contributions.iter().any(|contribution| {
|
||||
matches!(
|
||||
contribution,
|
||||
@@ -164,15 +149,15 @@ async fn openai_admin_provider_exposes_live_wire_profile_and_validated_billing()
|
||||
cache_write_tokens: Some(0),
|
||||
total: CurrencyCost {
|
||||
currency: "USD".to_owned(),
|
||||
amount: "5".parse().expect("fast amount"),
|
||||
amount: "4.25".parse().expect("fast amount"),
|
||||
},
|
||||
})
|
||||
.expect("fast billing")
|
||||
.expect("known fast pricing");
|
||||
assert_eq!(fast_billing.service_tier.as_deref(), Some("priority"));
|
||||
assert_eq!(fast_billing.multiplier_percent, 200);
|
||||
assert_eq!(fast_billing.multiplier_percent, 170);
|
||||
assert_eq!(fast_billing.standard_amount.amount.as_str(), "2.5");
|
||||
assert_eq!(fast_billing.total_amount.amount.as_str(), "5");
|
||||
assert_eq!(fast_billing.total_amount.amount.as_str(), "4.25");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -249,6 +234,14 @@ async fn openai_admin_provider_persists_the_full_pending_envelope_and_binds_owne
|
||||
.get("reauthorization_credential_revision")
|
||||
.is_none()
|
||||
);
|
||||
assert!(
|
||||
payload
|
||||
.expose_to_provider()
|
||||
.get("installation_id")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(|value| uuid::Uuid::parse_str(value).ok())
|
||||
.is_some_and(|value| value.get_version_num() == 4)
|
||||
);
|
||||
assert_eq!(
|
||||
mutation.get("schema_version").and_then(Value::as_u64),
|
||||
Some(3)
|
||||
@@ -284,7 +277,7 @@ async fn openai_admin_provider_persists_the_full_pending_envelope_and_binds_owne
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_reauthorization_pending_payload_contains_only_stable_account_identity() {
|
||||
async fn openai_reauthorization_pending_payload_reuses_the_account_installation_id() {
|
||||
let accounts = Arc::new(MemoryAccountStore::default());
|
||||
accounts
|
||||
.seed_oauth_credential(ImportCodexOAuthCredential {
|
||||
@@ -296,6 +289,14 @@ async fn openai_reauthorization_pending_payload_contains_only_stable_account_ide
|
||||
enabled: true,
|
||||
})
|
||||
.await;
|
||||
let account_id = ProviderAccountId::new("acct_pending_reauth").expect("account id");
|
||||
let existing = accounts
|
||||
.load_current_credential(&account_id)
|
||||
.await
|
||||
.expect("seeded credential");
|
||||
let expected_installation_id = CodexCredentialCodec::decode(&existing.credential)
|
||||
.expect("decode seeded credential")
|
||||
.installation_id;
|
||||
let pending = Arc::new(TestOAuthPending::default());
|
||||
let config = valid_config();
|
||||
let bundle = provider_openai::initialize(
|
||||
@@ -313,9 +314,7 @@ async fn openai_reauthorization_pending_payload_contains_only_stable_account_ide
|
||||
.admin_provider()
|
||||
.start_authorization(PendingAuthorizationMutation::new(
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
AuthorizationMutationTarget::Reauthorize {
|
||||
account_id: ProviderAccountId::new("acct_pending_reauth").expect("account id"),
|
||||
},
|
||||
AuthorizationMutationTarget::Reauthorize { account_id },
|
||||
AuthorizationOwnerBinding::from_context(&context),
|
||||
))
|
||||
.await
|
||||
@@ -342,6 +341,10 @@ async fn openai_reauthorization_pending_payload_contains_only_stable_account_ide
|
||||
.and_then(Value::as_str),
|
||||
Some("acct_pending_reauth")
|
||||
);
|
||||
assert_eq!(
|
||||
document.get("installation_id").and_then(Value::as_str),
|
||||
Some(expected_installation_id.as_str())
|
||||
);
|
||||
assert!(
|
||||
document
|
||||
.get("reauthorization_credential_revision")
|
||||
@@ -563,7 +566,7 @@ async fn openai_admin_provider_projects_official_codex_quota_and_independent_buc
|
||||
assert_eq!(monthly.len(), 1);
|
||||
assert!(monthly.iter().any(|window| {
|
||||
window.label == "月额度"
|
||||
&& window.source.as_deref() == Some("codex")
|
||||
&& window.limit_id.as_deref() == Some("codex")
|
||||
&& window.used_percent == Some(91.0)
|
||||
}));
|
||||
assert!(
|
||||
@@ -576,20 +579,22 @@ async fn openai_admin_provider_projects_official_codex_quota_and_independent_buc
|
||||
let secondary = quota
|
||||
.windows
|
||||
.iter()
|
||||
.find(|window| window.key == "core-secondary")
|
||||
.find(|window| window.key == "codex-secondary")
|
||||
.expect("core secondary quota");
|
||||
assert_eq!(secondary.label, "次级额度");
|
||||
assert_eq!(secondary.used_percent, Some(88.0));
|
||||
let review = quota
|
||||
.windows
|
||||
.iter()
|
||||
.find(|window| window.source.as_deref() == Some("code_review"))
|
||||
.find(|window| window.limit_id.as_deref() == Some("code_review"))
|
||||
.expect("code review quota");
|
||||
assert_eq!(review.label, "代码审查 · 周额度");
|
||||
assert_eq!(review.label, "周额度");
|
||||
assert_eq!(review.limit_name.as_deref(), Some("code_review"));
|
||||
assert_eq!(review.role, Some(ProviderQuotaWindowRole::Primary));
|
||||
let spark = quota
|
||||
.windows
|
||||
.iter()
|
||||
.find(|window| window.source.as_deref() == Some("codex_bengalfox"))
|
||||
.find(|window| window.limit_id.as_deref() == Some("codex_bengalfox"))
|
||||
.expect("Spark quota");
|
||||
assert_eq!(
|
||||
(
|
||||
@@ -760,6 +765,61 @@ async fn openai_admin_provider_rejects_unprepared_mutations_before_store_commit(
|
||||
assert_eq!(refresh_error.kind(), ProviderAdminErrorKind::NotFound);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_rotation_preserves_the_new_access_token_jwt_expiration() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
store
|
||||
.seed_oauth_credential(ImportCodexOAuthCredential {
|
||||
account_id: "acct_admin_rotation_expiration".to_owned(),
|
||||
name: "admin rotation expiration".to_owned(),
|
||||
secret: secret("admin-rotation-access"),
|
||||
verified_account: profile("chatgpt-admin-rotation-expiration"),
|
||||
next_refresh_at: None,
|
||||
enabled: true,
|
||||
})
|
||||
.await;
|
||||
let account = store
|
||||
.account("acct_admin_rotation_expiration")
|
||||
.expect("stored account");
|
||||
let record = account_record(&account);
|
||||
let config = valid_config();
|
||||
let bundle = provider_openai::initialize(
|
||||
config.config.clone(),
|
||||
provider_ports_with(store, Arc::new(TestOAuthPending::default())),
|
||||
)
|
||||
.await
|
||||
.expect("OpenAI bundle");
|
||||
|
||||
let expires_at = Utc
|
||||
.timestamp_opt(2_000_000_000, 0)
|
||||
.single()
|
||||
.expect("valid test timestamp");
|
||||
let payload = URL_SAFE_NO_PAD.encode(
|
||||
serde_json::to_vec(&serde_json::json!({"exp": expires_at.timestamp()}))
|
||||
.expect("test JWT payload"),
|
||||
);
|
||||
let mut material = Map::new();
|
||||
material.insert(
|
||||
"access_token".to_owned(),
|
||||
Value::String(format!("unverified-header.{payload}.unverified-signature")),
|
||||
);
|
||||
material.insert(
|
||||
"refresh_token".to_owned(),
|
||||
Value::String("admin-rotation-refresh".to_owned()),
|
||||
);
|
||||
|
||||
let prepared = bundle
|
||||
.admin_provider()
|
||||
.prepare_rotation(PrepareCredentialRotation {
|
||||
account: record,
|
||||
provider_material: ProviderDocument::new(OpaqueProviderData::new(material)),
|
||||
})
|
||||
.await
|
||||
.expect("JWT rotation should be prepared");
|
||||
|
||||
assert_eq!(prepared.facts().access_token_expires_at, Some(expires_at));
|
||||
}
|
||||
|
||||
fn provider_ports() -> ProviderStorePorts {
|
||||
provider_ports_with(
|
||||
Arc::new(MemoryAccountStore::default()),
|
||||
@@ -785,6 +845,7 @@ fn provider_ports_with_catalog(
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
Arc::new(MemorySessionExclusions::default()),
|
||||
catalog_cache,
|
||||
Arc::new(TestArtifactProfiles),
|
||||
Arc::new(TestCredentialState),
|
||||
Arc::new(TestCooldown),
|
||||
Arc::new(TestRuntimePolicy),
|
||||
@@ -852,6 +913,25 @@ fn valid_config() -> TestOpenAiConfig {
|
||||
}
|
||||
}
|
||||
|
||||
struct TestArtifactProfiles;
|
||||
|
||||
impl ProviderArtifactProfileCachePort for TestArtifactProfiles {
|
||||
fn replace_if_newer(
|
||||
&self,
|
||||
_profile: ProviderArtifactProfile,
|
||||
_ttl: Duration,
|
||||
) -> BoxFuture<'_, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(true) })
|
||||
}
|
||||
|
||||
fn read<'a>(
|
||||
&'a self,
|
||||
_provider_kind: &'a ProviderKind,
|
||||
) -> BoxFuture<'a, Result<Option<ProviderArtifactProfile>, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(None) })
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct TestCatalogCache {
|
||||
values: Mutex<BTreeMap<String, OpaqueProviderData>>,
|
||||
|
||||
@@ -18,6 +18,16 @@ fn openai_config_builds_the_audited_wire_profile() {
|
||||
assert_eq!(profile.desktop_build, "19012345678");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_config_derives_identity_secret_from_runtime_data_dir() {
|
||||
let mut config = valid_config();
|
||||
config
|
||||
.resolve_and_validate(Path::new("/srv/gateway/runtime-data"))
|
||||
.expect("valid OpenAI config");
|
||||
|
||||
assert!(format!("{config:?}").contains("/srv/gateway/runtime-data/identity_hmac_secret"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_config_rejects_noncanonical_versions_and_empty_fields() {
|
||||
let mut config = valid_config();
|
||||
@@ -116,7 +126,7 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() {
|
||||
);
|
||||
assert_eq!(
|
||||
config.wire_profile_state().snapshot().user_agent(),
|
||||
"Codex Desktop/0.146.0 (Mac OS 15.7.1; arm64) unknown (Codex Desktop; 26.727.51351)"
|
||||
"Codex Desktop/0.147.0-alpha.6.6 (Mac OS 15.7.1; arm64) unknown (Codex Desktop; 26.803.81509)"
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -13,6 +13,7 @@ use gateway_core::engine::credential::{
|
||||
ProviderAccountStore, QuotaState,
|
||||
};
|
||||
use gateway_core::routing::ProviderKind;
|
||||
use provider_openai::OpenAiConfig;
|
||||
use provider_openai::credential::{
|
||||
CODEX_AUTHENTICATION_KIND_AGENT_IDENTITY, CodexAgentIdentityAuthMode,
|
||||
CodexAgentIdentityCredentialData, CodexAgentIdentityError, CodexAgentIdentitySecret,
|
||||
@@ -20,7 +21,9 @@ use provider_openai::credential::{
|
||||
CodexRuntimeAuthentication, is_agent_identity_task_invalid_response,
|
||||
};
|
||||
use provider_openai::transport::CodexWebSocketPool;
|
||||
use provider_openai::transport::headers::build_codex_profile_headers;
|
||||
use reqwest::StatusCode;
|
||||
use reqwest::header::USER_AGENT;
|
||||
use secrecy::ExposeSecret as _;
|
||||
use serde_json::Value;
|
||||
|
||||
@@ -144,6 +147,20 @@ fn agent_identity_task_error_classifier_requires_unauthorized_task_error() {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_identity_task_registration_should_use_the_shared_desktop_profile_headers() {
|
||||
let profile = OpenAiConfig::default().wire_profile_state().snapshot();
|
||||
let headers = build_codex_profile_headers(&profile).expect("Desktop profile headers");
|
||||
|
||||
assert_eq!(headers.len(), 3);
|
||||
assert_eq!(headers["originator"], "Codex Desktop");
|
||||
assert_eq!(headers["x-openai-internal-codex-residency"], "us");
|
||||
assert_eq!(
|
||||
headers[USER_AGENT],
|
||||
"Codex Desktop/0.147.0-alpha.6.6 (Mac OS 15.7.1; arm64) unknown (Codex Desktop; 26.803.81509)"
|
||||
);
|
||||
}
|
||||
|
||||
struct SequenceTaskRegistrar {
|
||||
task_ids: Mutex<VecDeque<String>>,
|
||||
calls: AtomicUsize,
|
||||
|
||||
@@ -253,7 +253,7 @@ async fn cross_plan_presentation_drift_unions_on_the_first_seen_model() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cross_plan_responses_api_conflict_still_fails_the_union() {
|
||||
async fn cross_plan_supported_in_api_drift_unions_on_the_first_seen_model() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
seed_account_with_plan(&store, "acct_conflict_plus", "plus").await;
|
||||
seed_account_with_plan(&store, "acct_conflict_pro", "pro").await;
|
||||
@@ -272,15 +272,12 @@ async fn cross_plan_responses_api_conflict_still_fails_the_union() {
|
||||
.await;
|
||||
let service = service_with_catalog_cache(&store, server.uri(), catalog_cache());
|
||||
|
||||
let error = service
|
||||
let snapshot = service
|
||||
.synchronize()
|
||||
.await
|
||||
.expect_err("routing facts must still agree across plans");
|
||||
.expect("supported_in_api drift must not block catalog publication");
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
CodexCredentialCatalogError::ConflictingModelFacts
|
||||
));
|
||||
assert_eq!(snapshot.models().len(), 1);
|
||||
server.verify().await;
|
||||
}
|
||||
|
||||
|
||||
@@ -812,13 +812,15 @@ fn credential_expired_failure_keeps_expired_oauth_for_bounded_refresh_recovery()
|
||||
))
|
||||
.expect("record credential expiry");
|
||||
|
||||
let retained = store
|
||||
.account("acct_primary")
|
||||
.expect("account retained for refresh");
|
||||
assert_eq!(retained.credential_state(), CredentialState::Ready);
|
||||
assert_eq!(
|
||||
store
|
||||
.account("acct_primary")
|
||||
.expect("account retained for refresh")
|
||||
.credential_state(),
|
||||
CredentialState::Ready
|
||||
retained.last_error_reason(),
|
||||
Some(AccountErrorReason::AccessTokenExpired)
|
||||
);
|
||||
assert_eq!(retained.last_error_message(), Some("token_expired"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
@@ -7,22 +8,26 @@ use gateway_admin::model::provider_credentials::{
|
||||
AuthorizationMutationTarget, AuthorizationOwnerBinding, PendingAuthorizationMutation,
|
||||
};
|
||||
use gateway_admin::model::{MutationActor, MutationContext};
|
||||
use gateway_core::engine::credential::AccountStatus;
|
||||
use gateway_core::engine::credential::{AccountStatus, ProviderAccountId};
|
||||
use gateway_core::routing::ProviderKind;
|
||||
use provider_openai::OpenAiConfig;
|
||||
use provider_openai::credential::token_client::{
|
||||
AuthorizationCodeExchangeError, AuthorizationCodeExchanger, AuthorizationCodeGrant,
|
||||
AuthorizationTokenSet,
|
||||
};
|
||||
use provider_openai::credential::{
|
||||
CodexCredentialAdmin, CodexOAuthAdmin, CodexOAuthAdminError, CodexOAuthAdminService,
|
||||
CodexOAuthPendingClaimOutcome, CodexOAuthPendingStore, CodexOAuthPendingStoreError,
|
||||
CodexOAuthSecret, CompleteCodexOAuthAuthorization, CompletedCodexOAuthCredential,
|
||||
StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
|
||||
CodexCredentialAdmin, CodexCredentialCodec, CodexOAuthAdmin, CodexOAuthAdminError,
|
||||
CodexOAuthAdminService, CodexOAuthPendingClaimOutcome, CodexOAuthPendingStore,
|
||||
CodexOAuthPendingStoreError, CodexOAuthSecret, CompleteCodexOAuthAuthorization,
|
||||
CompletedCodexOAuthCredential, ImportCodexOAuthCredential, StartCodexOAuthAuthorization,
|
||||
StoredCodexPendingAuthorization,
|
||||
};
|
||||
use provider_openai::transport::profile::CodexWireProfileState;
|
||||
use secrecy::SecretString;
|
||||
use url::Url;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::support::MemoryAccountStore;
|
||||
use crate::support::{MemoryAccountStore, profile as account_profile, secret};
|
||||
|
||||
#[derive(Default)]
|
||||
struct PendingStore {
|
||||
@@ -48,6 +53,7 @@ impl CodexOAuthPendingStore for PendingStore {
|
||||
state: pending.state().clone(),
|
||||
nonce: pending.nonce().clone(),
|
||||
code_verifier: pending.code_verifier().clone(),
|
||||
installation_id: pending.installation_id().to_owned(),
|
||||
reauthorization_account_id: pending.reauthorization().map(ToString::to_string),
|
||||
mutation: pending.mutation().clone(),
|
||||
});
|
||||
@@ -106,7 +112,6 @@ impl AuthorizationCodeExchanger for Exchanger {
|
||||
id_token: None,
|
||||
},
|
||||
id_token: SecretString::from(self.id_token.clone()),
|
||||
expires_in: Some(Duration::from_secs(60 * 60)),
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -125,18 +130,37 @@ fn mutation() -> PendingAuthorizationMutation {
|
||||
)
|
||||
}
|
||||
|
||||
fn reauthorization_mutation(
|
||||
account_id: ProviderAccountId,
|
||||
request_id: &str,
|
||||
) -> PendingAuthorizationMutation {
|
||||
let context = MutationContext {
|
||||
actor: MutationActor::System,
|
||||
request_id: request_id.to_owned(),
|
||||
};
|
||||
PendingAuthorizationMutation::new(
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
AuthorizationMutationTarget::Reauthorize { account_id },
|
||||
AuthorizationOwnerBinding::from_context(&context),
|
||||
)
|
||||
}
|
||||
|
||||
fn id_token(payload: serde_json::Value) -> String {
|
||||
let payload = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&payload).expect("payload JSON"));
|
||||
// 官方逻辑只读取 payload;header/signature 不参与本地 metadata 解析。
|
||||
format!("unverified-header.{payload}.unverified-signature")
|
||||
}
|
||||
|
||||
async fn complete(id_token: String) -> Result<CompletedCodexOAuthCredential, CodexOAuthAdminError> {
|
||||
async fn complete(
|
||||
id_token: String,
|
||||
) -> Result<(CompletedCodexOAuthCredential, String, String), CodexOAuthAdminError> {
|
||||
let pending_store = Arc::new(PendingStore::default());
|
||||
let service = CodexOAuthAdminService::new(
|
||||
Arc::new(PendingStore::default()),
|
||||
pending_store.clone(),
|
||||
Arc::new(Exchanger { id_token }),
|
||||
Arc::new(MemoryAccountStore::default()),
|
||||
CodexCredentialAdmin,
|
||||
profile(),
|
||||
);
|
||||
let started = service
|
||||
.start_authorization(StartCodexOAuthAuthorization {
|
||||
@@ -144,11 +168,26 @@ async fn complete(id_token: String) -> Result<CompletedCodexOAuthCredential, Cod
|
||||
})
|
||||
.await
|
||||
.expect("start OAuth authorization");
|
||||
let state = Url::parse(&started.authorization_url)
|
||||
.expect("authorization URL")
|
||||
let outer = Url::parse(&started.authorization_url).expect("authorization URL");
|
||||
let inner_url = outer
|
||||
.query_pairs()
|
||||
.find_map(|(key, value)| (key == "state").then(|| value.into_owned()))
|
||||
.expect("state parameter");
|
||||
.find_map(|(key, value)| (key == "authorize_url").then(|| value.into_owned()))
|
||||
.expect("inner authorization URL");
|
||||
let inner = Url::parse(&inner_url).expect("authorization URL");
|
||||
let parameters = inner.query_pairs().into_owned().collect::<BTreeMap<_, _>>();
|
||||
let state = parameters.get("state").expect("state parameter");
|
||||
let surface_stable_id = parameters
|
||||
.get("source_surface_stable_id")
|
||||
.expect("surface stable ID")
|
||||
.to_owned();
|
||||
let installation_id = pending_store
|
||||
.pending
|
||||
.lock()
|
||||
.expect("pending lock")
|
||||
.as_ref()
|
||||
.expect("stored pending authorization")
|
||||
.installation_id
|
||||
.clone();
|
||||
let completed = service
|
||||
.complete_authorization(CompleteCodexOAuthAuthorization {
|
||||
owner_ref: "test-owner".to_owned(),
|
||||
@@ -158,12 +197,279 @@ async fn complete(id_token: String) -> Result<CompletedCodexOAuthCredential, Cod
|
||||
)),
|
||||
})
|
||||
.await?;
|
||||
Ok(completed.credential)
|
||||
Ok((completed.credential, installation_id, surface_stable_id))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn authorize_url_matches_the_official_desktop_parameter_contract() {
|
||||
let service = CodexOAuthAdminService::new(
|
||||
Arc::new(PendingStore::default()),
|
||||
Arc::new(Exchanger {
|
||||
id_token: "unused".to_owned(),
|
||||
}),
|
||||
Arc::new(MemoryAccountStore::default()),
|
||||
CodexCredentialAdmin,
|
||||
profile(),
|
||||
);
|
||||
let started = service
|
||||
.start_authorization(StartCodexOAuthAuthorization {
|
||||
mutation: mutation(),
|
||||
})
|
||||
.await
|
||||
.expect("start OAuth authorization");
|
||||
let url = Url::parse(&started.authorization_url).expect("authorization URL");
|
||||
let outer_parameters = url.query_pairs().into_owned().collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
outer_parameters
|
||||
.iter()
|
||||
.map(|(name, _)| name.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
vec![
|
||||
"authorize_url",
|
||||
"codex_streamlined_login",
|
||||
"no_universal_links"
|
||||
]
|
||||
);
|
||||
let outer_parameters = outer_parameters.into_iter().collect::<BTreeMap<_, _>>();
|
||||
|
||||
assert_eq!(url.scheme(), "https");
|
||||
assert_eq!(url.host_str(), Some("chatgpt.com"));
|
||||
assert_eq!(url.path(), "/codex/desktop-auth");
|
||||
assert_eq!(outer_parameters.len(), 3);
|
||||
assert_eq!(
|
||||
outer_parameters
|
||||
.get("codex_streamlined_login")
|
||||
.map(String::as_str),
|
||||
Some("true")
|
||||
);
|
||||
assert_eq!(
|
||||
outer_parameters
|
||||
.get("no_universal_links")
|
||||
.map(String::as_str),
|
||||
Some("1")
|
||||
);
|
||||
let inner = Url::parse(
|
||||
outer_parameters
|
||||
.get("authorize_url")
|
||||
.expect("authorize_url parameter"),
|
||||
)
|
||||
.expect("inner authorization URL");
|
||||
let ordered_parameters = inner.query_pairs().into_owned().collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
ordered_parameters
|
||||
.iter()
|
||||
.map(|(name, _)| name.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
vec![
|
||||
"response_type",
|
||||
"client_id",
|
||||
"redirect_uri",
|
||||
"scope",
|
||||
"code_challenge",
|
||||
"code_challenge_method",
|
||||
"id_token_add_organizations",
|
||||
"codex_cli_simplified_flow",
|
||||
"state",
|
||||
"originator",
|
||||
"codex_app_version",
|
||||
"source_surface_stable_id",
|
||||
"codex_origin_stable_id",
|
||||
"codex_streamlined_login",
|
||||
]
|
||||
);
|
||||
let parameters = ordered_parameters.into_iter().collect::<BTreeMap<_, _>>();
|
||||
|
||||
assert_eq!(inner.host_str(), Some("auth.openai.com"));
|
||||
assert_eq!(inner.path(), "/oauth/authorize");
|
||||
assert_eq!(parameters.len(), 14);
|
||||
assert!(
|
||||
inner
|
||||
.query()
|
||||
.is_some_and(|query| query.contains("scope=openid+profile+email+offline_access"))
|
||||
);
|
||||
assert!(
|
||||
!inner
|
||||
.query()
|
||||
.is_some_and(|query| query.contains("scope=openid%20profile"))
|
||||
);
|
||||
assert_eq!(
|
||||
parameters.get("response_type").map(String::as_str),
|
||||
Some("code")
|
||||
);
|
||||
assert_eq!(
|
||||
parameters.get("redirect_uri").map(String::as_str),
|
||||
Some("http://localhost:1455/auth/callback")
|
||||
);
|
||||
assert_eq!(
|
||||
parameters.get("scope").map(String::as_str),
|
||||
Some("openid profile email offline_access api.connectors.read api.connectors.invoke")
|
||||
);
|
||||
assert_eq!(
|
||||
parameters.get("originator").map(String::as_str),
|
||||
Some("Codex Desktop")
|
||||
);
|
||||
assert_eq!(
|
||||
parameters.get("codex_app_version").map(String::as_str),
|
||||
Some("26.803.81509")
|
||||
);
|
||||
let surface_stable_id = parameters
|
||||
.get("source_surface_stable_id")
|
||||
.expect("surface stable ID");
|
||||
let parsed_surface_id = Uuid::parse_str(surface_stable_id).expect("UUID surface stable ID");
|
||||
assert_eq!(parsed_surface_id.get_version_num(), 4);
|
||||
assert_eq!(
|
||||
parameters.get("codex_origin_stable_id"),
|
||||
Some(surface_stable_id)
|
||||
);
|
||||
assert_eq!(
|
||||
parameters
|
||||
.get("codex_streamlined_login")
|
||||
.map(String::as_str),
|
||||
Some("true")
|
||||
);
|
||||
assert_eq!(
|
||||
parameters.get("code_challenge_method").map(String::as_str),
|
||||
Some("S256")
|
||||
);
|
||||
assert_eq!(
|
||||
parameters
|
||||
.get("id_token_add_organizations")
|
||||
.map(String::as_str),
|
||||
Some("true")
|
||||
);
|
||||
assert_eq!(
|
||||
parameters
|
||||
.get("codex_cli_simplified_flow")
|
||||
.map(String::as_str),
|
||||
Some("true")
|
||||
);
|
||||
assert!(
|
||||
parameters
|
||||
.get("client_id")
|
||||
.is_some_and(|value| !value.is_empty())
|
||||
);
|
||||
assert!(
|
||||
parameters
|
||||
.get("state")
|
||||
.is_some_and(|value| !value.is_empty())
|
||||
);
|
||||
assert!(
|
||||
parameters
|
||||
.get("code_challenge")
|
||||
.is_some_and(|value| !value.is_empty())
|
||||
);
|
||||
assert!(!parameters.contains_key("nonce"));
|
||||
}
|
||||
|
||||
fn profile() -> CodexWireProfileState {
|
||||
OpenAiConfig::default().wire_profile_state()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn separate_new_authorizations_use_distinct_surface_stable_ids() {
|
||||
let mut ids = Vec::new();
|
||||
for _ in 0..2 {
|
||||
let service = CodexOAuthAdminService::new(
|
||||
Arc::new(PendingStore::default()),
|
||||
Arc::new(Exchanger {
|
||||
id_token: "unused".to_owned(),
|
||||
}),
|
||||
Arc::new(MemoryAccountStore::default()),
|
||||
CodexCredentialAdmin,
|
||||
profile(),
|
||||
);
|
||||
let started = service
|
||||
.start_authorization(StartCodexOAuthAuthorization {
|
||||
mutation: mutation(),
|
||||
})
|
||||
.await
|
||||
.expect("start OAuth authorization");
|
||||
let outer = Url::parse(&started.authorization_url).expect("outer authorization URL");
|
||||
let inner = outer
|
||||
.query_pairs()
|
||||
.find_map(|(key, value)| (key == "authorize_url").then(|| value.into_owned()))
|
||||
.and_then(|value| Url::parse(&value).ok())
|
||||
.expect("inner authorization URL");
|
||||
let surface_stable_id = inner
|
||||
.query_pairs()
|
||||
.find_map(|(key, value)| {
|
||||
(key == "source_surface_stable_id").then(|| value.into_owned())
|
||||
})
|
||||
.expect("surface stable ID");
|
||||
ids.push(surface_stable_id);
|
||||
}
|
||||
assert_ne!(ids[0], ids[1]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn repeated_reauthorization_derives_the_same_surface_id_for_one_account() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
store
|
||||
.seed_oauth_credential(ImportCodexOAuthCredential {
|
||||
account_id: "acct_surface_reauth".to_owned(),
|
||||
name: "surface reauthorization".to_owned(),
|
||||
secret: secret("surface-reauth-access"),
|
||||
verified_account: account_profile("chatgpt-surface-reauth"),
|
||||
next_refresh_at: None,
|
||||
enabled: true,
|
||||
})
|
||||
.await;
|
||||
let account_id = ProviderAccountId::new("acct_surface_reauth").expect("account id");
|
||||
let mut surface_ids = Vec::new();
|
||||
for request_id in ["surface-reauth-first", "surface-reauth-second"] {
|
||||
let service = CodexOAuthAdminService::new(
|
||||
Arc::new(PendingStore::default()),
|
||||
Arc::new(Exchanger {
|
||||
id_token: "unused".to_owned(),
|
||||
}),
|
||||
store.clone(),
|
||||
CodexCredentialAdmin,
|
||||
profile(),
|
||||
);
|
||||
let started = service
|
||||
.start_authorization(StartCodexOAuthAuthorization {
|
||||
mutation: reauthorization_mutation(account_id.clone(), request_id),
|
||||
})
|
||||
.await
|
||||
.expect("start OAuth reauthorization");
|
||||
let outer = Url::parse(&started.authorization_url).expect("outer authorization URL");
|
||||
let inner = outer
|
||||
.query_pairs()
|
||||
.find_map(|(key, value)| (key == "authorize_url").then(|| value.into_owned()))
|
||||
.and_then(|value| Url::parse(&value).ok())
|
||||
.expect("inner authorization URL");
|
||||
surface_ids.push(
|
||||
inner
|
||||
.query_pairs()
|
||||
.find_map(|(key, value)| {
|
||||
(key == "source_surface_stable_id").then(|| value.into_owned())
|
||||
})
|
||||
.expect("surface stable ID"),
|
||||
);
|
||||
}
|
||||
|
||||
assert_eq!(surface_ids[0], surface_ids[1]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn first_exchange_persists_the_installation_id_without_reusing_it_as_the_surface_id() {
|
||||
let (credential, pending_installation_id, surface_stable_id) = complete(id_token(
|
||||
serde_json::json!({ "email": "identity@example.com" }),
|
||||
))
|
||||
.await
|
||||
.expect("completed OAuth exchange");
|
||||
let CompletedCodexOAuthCredential::Create(account) = credential else {
|
||||
panic!("expected account creation");
|
||||
};
|
||||
let runtime = CodexCredentialCodec::decode(&account.credential).expect("stored credential");
|
||||
|
||||
assert_eq!(runtime.installation_id, pending_installation_id);
|
||||
assert_ne!(runtime.installation_id, surface_stable_id);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn first_exchange_uses_official_id_token_claim_mapping_without_signature_validation() {
|
||||
let credential = complete(id_token(serde_json::json!({
|
||||
let (credential, _, _) = complete(id_token(serde_json::json!({
|
||||
"email": "top@example.com",
|
||||
"https://api.openai.com/profile": { "email": "fallback@example.com" },
|
||||
"https://api.openai.com/auth": {
|
||||
@@ -198,7 +504,7 @@ async fn first_exchange_uses_official_id_token_claim_mapping_without_signature_v
|
||||
|
||||
#[tokio::test]
|
||||
async fn first_exchange_permits_missing_identity_claims_but_requires_a_parseable_id_token() {
|
||||
let credential = complete(id_token(serde_json::json!({
|
||||
let (credential, _, _) = complete(id_token(serde_json::json!({
|
||||
"https://api.openai.com/profile": { "email": "profile@example.com" }
|
||||
})))
|
||||
.await
|
||||
|
||||
@@ -11,6 +11,7 @@ use gateway_core::engine::credential::{
|
||||
AccountStateChange, AccountStatus, CredentialState, ProviderAccount, ProviderAccountStore as _,
|
||||
QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaState,
|
||||
};
|
||||
use gateway_protocol::openai::events::parse_rate_limits_event;
|
||||
use provider_openai::OFFICIAL_CODEX_BASE_URL;
|
||||
use provider_openai::credential::{
|
||||
CodexCredentialQuotaError, CodexCredentialQuotaService, ImportCodexOAuthCredential,
|
||||
@@ -163,6 +164,224 @@ async fn successful_inference_headers_are_authoritative_even_at_full_display_usa
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn passive_quota_updates_keep_core_and_model_specific_buckets_independent() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let account_id = "acct_independent_passive_buckets";
|
||||
create_account(&store, account_id).await;
|
||||
let account = store.account(account_id).expect("account");
|
||||
let service = quota_service(&store);
|
||||
let initial_headers = vec![
|
||||
("x-codex-active-limit".to_owned(), "codex".to_owned()),
|
||||
("x-codex-primary-used-percent".to_owned(), "81".to_owned()),
|
||||
(
|
||||
"x-codex-primary-window-minutes".to_owned(),
|
||||
"10080".to_owned(),
|
||||
),
|
||||
(
|
||||
"x-codex-bengalfox-primary-used-percent".to_owned(),
|
||||
"0".to_owned(),
|
||||
),
|
||||
(
|
||||
"x-codex-bengalfox-primary-window-minutes".to_owned(),
|
||||
"10080".to_owned(),
|
||||
),
|
||||
(
|
||||
"x-codex-bengalfox-limit-name".to_owned(),
|
||||
"GPT-5.3-Codex-Spark".to_owned(),
|
||||
),
|
||||
];
|
||||
service
|
||||
.synchronize_passive_headers(&account, &initial_headers)
|
||||
.await
|
||||
.expect("initial passive quota");
|
||||
|
||||
let spark_headers = vec![
|
||||
// 生产响应会把当前具名桶同时复制到旧版 `x-codex-*` 默认头;这个
|
||||
// 兼容别名不能覆盖已有的账号 core 桶。
|
||||
("x-codex-active-limit".to_owned(), "codex".to_owned()),
|
||||
("x-codex-primary-used-percent".to_owned(), "1".to_owned()),
|
||||
(
|
||||
"x-codex-primary-window-minutes".to_owned(),
|
||||
"10080".to_owned(),
|
||||
),
|
||||
(
|
||||
"x-codex-primary-reset-at".to_owned(),
|
||||
"1900000100".to_owned(),
|
||||
),
|
||||
(
|
||||
"x-codex-bengalfox-primary-used-percent".to_owned(),
|
||||
"1".to_owned(),
|
||||
),
|
||||
(
|
||||
"x-codex-bengalfox-primary-window-minutes".to_owned(),
|
||||
"10080".to_owned(),
|
||||
),
|
||||
(
|
||||
"x-codex-bengalfox-primary-reset-at".to_owned(),
|
||||
"1900000100".to_owned(),
|
||||
),
|
||||
(
|
||||
"x-codex-bengalfox-limit-name".to_owned(),
|
||||
"GPT-5.3-Codex-Spark".to_owned(),
|
||||
),
|
||||
];
|
||||
service
|
||||
.synchronize_passive_headers(&account, &spark_headers)
|
||||
.await
|
||||
.expect("Spark passive quota");
|
||||
|
||||
let snapshot = service
|
||||
.read_account(account.id())
|
||||
.await
|
||||
.expect("read Spark quota")
|
||||
.expect("Spark quota snapshot");
|
||||
let core = snapshot
|
||||
.windows()
|
||||
.iter()
|
||||
.find(|window| window.source() == "codex")
|
||||
.expect("core quota window");
|
||||
let spark = snapshot
|
||||
.windows()
|
||||
.iter()
|
||||
.find(|window| window.source() == "codex_bengalfox")
|
||||
.expect("Spark quota window");
|
||||
assert_eq!(core.used_percent(), Some(81.0));
|
||||
assert_eq!(spark.used_percent(), Some(1.0));
|
||||
|
||||
let core_headers = vec![
|
||||
("x-codex-active-limit".to_owned(), "codex".to_owned()),
|
||||
("x-codex-primary-used-percent".to_owned(), "82".to_owned()),
|
||||
(
|
||||
"x-codex-primary-window-minutes".to_owned(),
|
||||
"10080".to_owned(),
|
||||
),
|
||||
];
|
||||
service
|
||||
.synchronize_passive_headers(&account, &core_headers)
|
||||
.await
|
||||
.expect("core passive quota");
|
||||
|
||||
let snapshot = service
|
||||
.read_account(account.id())
|
||||
.await
|
||||
.expect("read core quota")
|
||||
.expect("core quota snapshot");
|
||||
let core = snapshot
|
||||
.windows()
|
||||
.iter()
|
||||
.find(|window| window.source() == "codex")
|
||||
.expect("core quota window");
|
||||
let spark = snapshot
|
||||
.windows()
|
||||
.iter()
|
||||
.find(|window| window.source() == "codex_bengalfox")
|
||||
.expect("Spark quota window");
|
||||
assert_eq!(core.used_percent(), Some(82.0));
|
||||
assert_eq!(spark.used_percent(), Some(1.0));
|
||||
|
||||
let persisted = store.quota_json(account_id).expect("persisted quota");
|
||||
let by_limit_id = persisted
|
||||
.get("rate_limits_by_limit_id")
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.expect("rate-limit map");
|
||||
assert_eq!(
|
||||
by_limit_id.keys().map(String::as_str).collect::<Vec<_>>(),
|
||||
["codex", "codex_bengalfox"]
|
||||
);
|
||||
assert!(persisted.get("rate_limit").is_none());
|
||||
assert!(persisted.get("additional_rate_limits").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_additional_rate_limit_resolves_by_name_without_touching_core() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let account_id = "acct_websocket_additional_limit";
|
||||
create_account(&store, account_id).await;
|
||||
let account = store.account(account_id).expect("account");
|
||||
let service = quota_service(&store);
|
||||
let initial_headers = vec![
|
||||
("x-codex-active-limit".to_owned(), "codex".to_owned()),
|
||||
("x-codex-primary-used-percent".to_owned(), "84".to_owned()),
|
||||
(
|
||||
"x-codex-primary-window-minutes".to_owned(),
|
||||
"10080".to_owned(),
|
||||
),
|
||||
(
|
||||
"x-codex-primary-reset-at".to_owned(),
|
||||
"1900000000".to_owned(),
|
||||
),
|
||||
(
|
||||
"x-codex-bengalfox-primary-used-percent".to_owned(),
|
||||
"0".to_owned(),
|
||||
),
|
||||
(
|
||||
"x-codex-bengalfox-primary-window-minutes".to_owned(),
|
||||
"10080".to_owned(),
|
||||
),
|
||||
(
|
||||
"x-codex-bengalfox-primary-reset-at".to_owned(),
|
||||
"1900000100".to_owned(),
|
||||
),
|
||||
(
|
||||
"x-codex-bengalfox-limit-name".to_owned(),
|
||||
"GPT-5.3-Codex-Spark".to_owned(),
|
||||
),
|
||||
];
|
||||
service
|
||||
.synchronize_passive_headers(&account, &initial_headers)
|
||||
.await
|
||||
.expect("initial quota metadata");
|
||||
|
||||
let event = json!({
|
||||
"type": "codex.rate_limits",
|
||||
"plan_type": "pro",
|
||||
"rate_limits": {
|
||||
"allowed": true,
|
||||
"limit_reached": false,
|
||||
"primary": {
|
||||
"used_percent": 2,
|
||||
"window_minutes": 10080,
|
||||
"reset_at": 1900000100
|
||||
},
|
||||
"secondary": null
|
||||
},
|
||||
"additional_rate_limits": {
|
||||
"GPT-5.3-Codex-Spark": {
|
||||
"allowed": true,
|
||||
"limit_reached": false,
|
||||
"primary": {
|
||||
"used_percent": 2,
|
||||
"window_minutes": 10080,
|
||||
"reset_at": 1900000100
|
||||
},
|
||||
"secondary": null
|
||||
}
|
||||
},
|
||||
"credits": null
|
||||
});
|
||||
let observation = parse_rate_limits_event(&event).expect("WebSocket rate-limit event");
|
||||
service
|
||||
.synchronize_passive_rate_limits(&account, std::slice::from_ref(&observation))
|
||||
.await
|
||||
.expect("WebSocket passive quota");
|
||||
|
||||
let persisted = store.quota_json(account_id).expect("persisted quota");
|
||||
let by_limit_id = persisted["rate_limits_by_limit_id"]
|
||||
.as_object()
|
||||
.expect("rate-limit map");
|
||||
assert_eq!(
|
||||
by_limit_id["codex"]["primary_window"]["used_percent"],
|
||||
json!(84.0)
|
||||
);
|
||||
assert_eq!(
|
||||
by_limit_id["codex_bengalfox"]["primary_window"]["used_percent"],
|
||||
json!(2.0)
|
||||
);
|
||||
assert_eq!(persisted["active_limit"], json!("codex_bengalfox"));
|
||||
assert_eq!(by_limit_id.len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_refresh_persists_explicit_provider_denial() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
@@ -209,6 +428,14 @@ async fn quota_refresh_persists_explicit_provider_denial() {
|
||||
current.status_projection(SystemTime::now(), None).status,
|
||||
AccountStatus::QuotaExhausted
|
||||
);
|
||||
let persisted = store.quota_json(account_id).expect("persisted quota");
|
||||
assert!(
|
||||
persisted
|
||||
.get("rate_limits_by_limit_id")
|
||||
.and_then(|limits| limits.get("codex"))
|
||||
.is_some()
|
||||
);
|
||||
assert!(persisted.get("rate_limit").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -75,7 +75,7 @@ async fn quota_refresh_updates_access_fact_without_recovering_credential_error()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn exhausted_quota_worker_keeps_old_document_when_new_window_usage_is_high() {
|
||||
async fn exhausted_quota_worker_only_touches_refresh_time_when_recovery_is_rejected() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let account_id = "acct_worker_quota_recovery_gate";
|
||||
create_account(&store, account_id).await;
|
||||
@@ -101,7 +101,7 @@ async fn exhausted_quota_worker_keeps_old_document_when_new_window_usage_is_high
|
||||
server.uri(),
|
||||
);
|
||||
|
||||
service
|
||||
let old_snapshot = service
|
||||
.refresh_account(account.id())
|
||||
.await
|
||||
.expect("seed exhausted quota");
|
||||
@@ -125,6 +125,14 @@ async fn exhausted_quota_worker_keeps_old_document_when_new_window_usage_is_high
|
||||
assert_eq!(summary.exhausted, 1);
|
||||
assert_eq!(summary.updated, 0);
|
||||
assert_eq!(store.quota_json(account_id), Some(old_quota));
|
||||
let refreshed = service
|
||||
.read_account(account.id())
|
||||
.await
|
||||
.expect("read refreshed quota")
|
||||
.expect("refreshed quota snapshot");
|
||||
assert!(refreshed.observed_at() > old_snapshot.observed_at());
|
||||
assert_eq!(refreshed.fact(), old_snapshot.fact());
|
||||
assert_eq!(refreshed.quota().access(), QuotaAccessState::Exhausted);
|
||||
assert_eq!(
|
||||
store
|
||||
.account(account_id)
|
||||
@@ -137,10 +145,18 @@ async fn exhausted_quota_worker_keeps_old_document_when_new_window_usage_is_high
|
||||
|
||||
#[tokio::test]
|
||||
async fn manual_quota_refresh_only_recovers_after_reset_advances_below_ten_percent() {
|
||||
for (suffix, reset_at, used_percent, recovered) in [
|
||||
("same_reset", 1_900_000_000_i64, 0_u64, false),
|
||||
("exactly_ten", 1_900_003_600_i64, 10_u64, false),
|
||||
("below_ten", 1_900_003_600_i64, 9_u64, true),
|
||||
for (suffix, allowed, limit_reached, reset_at, used_percent, recovered) in [
|
||||
(
|
||||
"still_exhausted",
|
||||
false,
|
||||
true,
|
||||
1_900_003_600_i64,
|
||||
0_u8,
|
||||
false,
|
||||
),
|
||||
("same_reset", true, false, 1_900_000_000_i64, 0_u8, false),
|
||||
("exactly_ten", true, false, 1_900_003_600_i64, 10_u8, false),
|
||||
("below_ten", true, false, 1_900_003_600_i64, 9_u8, true),
|
||||
] {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let account_id = format!("acct_manual_quota_recovery_{suffix}");
|
||||
@@ -166,7 +182,7 @@ async fn manual_quota_refresh_only_recovers_after_reset_advances_below_ten_perce
|
||||
.expect("client"),
|
||||
server.uri(),
|
||||
);
|
||||
service
|
||||
let old_snapshot = service
|
||||
.refresh_account(account.id())
|
||||
.await
|
||||
.expect("seed exhausted quota");
|
||||
@@ -177,8 +193,8 @@ async fn manual_quota_refresh_only_recovers_after_reset_advances_below_ten_perce
|
||||
.and(path("/api/codex/usage"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"rate_limit": {
|
||||
"allowed": true,
|
||||
"limit_reached": false,
|
||||
"allowed": allowed,
|
||||
"limit_reached": limit_reached,
|
||||
"primary_window": {"used_percent": used_percent, "reset_at": reset_at}
|
||||
}
|
||||
})))
|
||||
@@ -197,11 +213,12 @@ async fn manual_quota_refresh_only_recovers_after_reset_advances_below_ten_perce
|
||||
assert_ne!(store.quota_json(&account_id), Some(old_quota));
|
||||
assert_eq!(current.quota().access(), QuotaAccessState::Allowed);
|
||||
} else {
|
||||
assert_eq!(snapshot.fact().remaining_percent(), Some(0));
|
||||
assert_eq!(snapshot.fact(), old_snapshot.fact());
|
||||
assert_eq!(snapshot.quota().access(), QuotaAccessState::Exhausted);
|
||||
assert_eq!(store.quota_json(&account_id), Some(old_quota));
|
||||
assert_eq!(current.quota().access(), QuotaAccessState::Exhausted);
|
||||
}
|
||||
assert!(snapshot.observed_at() > old_snapshot.observed_at());
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -396,11 +396,11 @@ async fn persisted_quota_orders_core_window_before_additional_limit() {
|
||||
.map(|window| window.source())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(sources, ["codex", "gpt-5.3-codex-spark"]);
|
||||
assert_eq!(sources, ["codex", "gpt_5.3_codex_spark"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn code_review_limit_projects_as_independent_window_with_limit_name() {
|
||||
async fn code_review_limit_projects_as_one_snapshot_per_limit_id() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_code_review").await;
|
||||
let account = store.account("acct_code_review").expect("account");
|
||||
@@ -455,21 +455,16 @@ async fn code_review_limit_projects_as_independent_window_with_limit_name() {
|
||||
.expect("read quota")
|
||||
.expect("quota snapshot");
|
||||
|
||||
// spend_control 不再生成窗口(只作 exhaustion 信号);顶层和 additional
|
||||
// 的 code_review 桶分别保留,避免一个快照覆盖另一个。
|
||||
// spend_control 不生成窗口(只作 exhaustion 信号);官方 map 协议中同一个
|
||||
// limit_id 只能保留一个快照,顶层 code_review 事实优先于重复 additional。
|
||||
let review = snapshot
|
||||
.windows()
|
||||
.iter()
|
||||
.filter(|window| window.source() == "code_review")
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(review.len(), 2);
|
||||
assert!(review.iter().any(|window| {
|
||||
window.limit_name() == Some("code_review") && window.used_percent() == Some(80.0)
|
||||
}));
|
||||
assert!(review.iter().any(|window| {
|
||||
window.limit_name() == Some("code review") && window.used_percent() == Some(55.0)
|
||||
}));
|
||||
assert_ne!(review[0].key(), review[1].key());
|
||||
assert_eq!(snapshot.windows().len(), 3);
|
||||
assert_eq!(review.len(), 1);
|
||||
assert_eq!(review[0].limit_name(), Some("code_review"));
|
||||
assert_eq!(review[0].used_percent(), Some(80.0));
|
||||
assert_eq!(snapshot.windows().len(), 2);
|
||||
assert_eq!(snapshot.quota().access(), QuotaAccessState::Unknown);
|
||||
}
|
||||
|
||||
@@ -4,9 +4,13 @@ use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
use chrono::{DateTime, TimeZone as _, Utc};
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_core::engine::credential::{ProviderAccountId, ProviderAccountStore};
|
||||
use gateway_core::engine::credential::{
|
||||
AccountErrorReason, AccountStateChange, AccountStatus, CredentialState, ProviderAccountId,
|
||||
ProviderAccountStore, QuotaAccessChange, QuotaEvidence, QuotaState, QuotaWriteOutcome,
|
||||
};
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderCredentialState, ProviderCredentialStatePort, ProviderLeaseAcquisition,
|
||||
@@ -14,12 +18,16 @@ use gateway_core::provider_ports::{
|
||||
ProviderSchedulingState, ProviderStoreError,
|
||||
};
|
||||
use gateway_core::routing::ProviderKind;
|
||||
use provider_openai::credential::token_client::{RefreshFailure, TokenPair, TokenRefresher};
|
||||
use provider_openai::credential::token_client::{
|
||||
OpenAiTokenClient, RefreshFailure, TokenClientConfig, TokenPair, TokenRefresher,
|
||||
};
|
||||
use provider_openai::credential::{
|
||||
CodexCredentialCodec, CodexCredentialRefreshOutcome, CodexCredentialRefreshService,
|
||||
ImportCodexOAuthCredential,
|
||||
};
|
||||
use secrecy::{ExposeSecret as _, SecretString};
|
||||
use wiremock::matchers::{method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
use crate::support::{MemoryAccountStore, profile, secret};
|
||||
|
||||
@@ -36,7 +44,6 @@ impl SingleUseRefresher {
|
||||
access_token: Some("refreshed-access-token".to_owned()),
|
||||
refresh_token: None,
|
||||
id_token: None,
|
||||
expires_in: Some(Duration::from_secs(60 * 60)),
|
||||
})),
|
||||
})
|
||||
}
|
||||
@@ -52,7 +59,17 @@ impl SingleUseRefresher {
|
||||
access_token: Some("refreshed-access-token".to_owned()),
|
||||
refresh_token: None,
|
||||
id_token: Some(id_token.to_owned()),
|
||||
expires_in: Some(Duration::from_secs(60 * 60)),
|
||||
})),
|
||||
})
|
||||
}
|
||||
|
||||
fn with_access_token(access_token: String) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
calls: AtomicUsize::new(0),
|
||||
response: Mutex::new(Some(TokenPair {
|
||||
access_token: Some(access_token),
|
||||
refresh_token: None,
|
||||
id_token: None,
|
||||
})),
|
||||
})
|
||||
}
|
||||
@@ -64,7 +81,6 @@ impl SingleUseRefresher {
|
||||
access_token: None,
|
||||
refresh_token: None,
|
||||
id_token: None,
|
||||
expires_in: Some(Duration::from_secs(60 * 60)),
|
||||
})),
|
||||
})
|
||||
}
|
||||
@@ -78,7 +94,21 @@ impl TokenRefresher for SingleUseRefresher {
|
||||
.lock()
|
||||
.expect("refresh response lock")
|
||||
.take()
|
||||
.ok_or(RefreshFailure::Transport)
|
||||
.ok_or_else(|| RefreshFailure::Transport {
|
||||
message: Some("test refresh response is exhausted".to_owned()),
|
||||
upstream: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct FailingRefresher {
|
||||
failure: RefreshFailure,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl TokenRefresher for FailingRefresher {
|
||||
async fn refresh(&self, _: &str) -> Result<TokenPair, RefreshFailure> {
|
||||
Err(self.failure.clone())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -267,6 +297,263 @@ async fn scheduled_refresh_uses_the_current_margin_without_persisting_a_normal_s
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn scheduled_refresh_rotates_tokens_while_quota_is_exhausted() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let policy = MutableRuntimePolicy::new(Duration::from_secs(5 * 60));
|
||||
let service = refresh_service(&store, SingleUseRefresher::new(), policy);
|
||||
let account_id = "acct_exhausted_quota_refresh";
|
||||
seed_refreshable_account(
|
||||
&store,
|
||||
account_id,
|
||||
SystemTime::now()
|
||||
.checked_add(Duration::from_secs(120))
|
||||
.expect("test expiry"),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
let account = store.account(account_id).expect("seeded account");
|
||||
assert_eq!(
|
||||
store
|
||||
.apply_quota_access(QuotaAccessChange {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
state: QuotaState::exhausted(
|
||||
QuotaEvidence::UsageLimitReached,
|
||||
SystemTime::now(),
|
||||
None,
|
||||
),
|
||||
})
|
||||
.await
|
||||
.expect("persist exhausted quota"),
|
||||
QuotaWriteOutcome::Updated
|
||||
);
|
||||
|
||||
let outcomes = service.refresh_due().await.expect("refresh cycle");
|
||||
|
||||
assert!(matches!(
|
||||
outcomes.as_slice(),
|
||||
[CodexCredentialRefreshOutcome::Refreshed {
|
||||
account_id: refreshed_account_id,
|
||||
..
|
||||
}] if refreshed_account_id == account_id
|
||||
));
|
||||
let account = store.account(account_id).expect("refreshed account");
|
||||
assert!(account.quota().is_exhausted());
|
||||
let loaded = store
|
||||
.load_credential(account.id(), account.revision())
|
||||
.await
|
||||
.expect("refreshed credential");
|
||||
let runtime = CodexCredentialCodec::decode(&loaded.credential).expect("runtime credential");
|
||||
let oauth = runtime.authentication.oauth().expect("OAuth credential");
|
||||
assert_eq!(oauth.access_token.expose_secret(), "refreshed-access-token");
|
||||
let expected_refresh_token = format!("rt-access-{account_id}");
|
||||
assert_eq!(
|
||||
oauth
|
||||
.refresh_token
|
||||
.as_ref()
|
||||
.map(SecretString::expose_secret),
|
||||
Some(expected_refresh_token.as_str())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn scheduled_refresh_persists_the_original_upstream_error_message() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let policy = MutableRuntimePolicy::new(Duration::from_secs(5 * 60));
|
||||
let upstream_message = "Refresh token has already been used.";
|
||||
let service = CodexCredentialRefreshService::new(
|
||||
store.repository(),
|
||||
Arc::new(FailingRefresher {
|
||||
failure: RefreshFailure::InvalidGrant {
|
||||
message: Some(upstream_message.to_owned()),
|
||||
upstream: None,
|
||||
},
|
||||
}),
|
||||
Arc::new(RefreshLeases),
|
||||
Arc::new(RefreshCredentialState),
|
||||
policy,
|
||||
);
|
||||
let account_id = "acct_invalid_refresh_token";
|
||||
seed_refreshable_account(
|
||||
&store,
|
||||
account_id,
|
||||
SystemTime::now()
|
||||
.checked_add(Duration::from_secs(120))
|
||||
.expect("test expiry"),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
let outcomes = service.refresh_due().await.expect("refresh cycle");
|
||||
|
||||
assert!(matches!(
|
||||
outcomes.as_slice(),
|
||||
[CodexCredentialRefreshOutcome::Invalidated {
|
||||
account_id: invalidated_account_id,
|
||||
}] if invalidated_account_id == account_id
|
||||
));
|
||||
assert_eq!(
|
||||
store
|
||||
.account(account_id)
|
||||
.expect("invalidated account")
|
||||
.last_error_message(),
|
||||
Some(upstream_message)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn scheduled_refresh_persists_retryable_message_inside_the_two_hour_window() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let policy = MutableRuntimePolicy::new(Duration::from_secs(5 * 60));
|
||||
let upstream_message = "Invalid refresh token.";
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth/token"))
|
||||
.respond_with(ResponseTemplate::new(401).set_body_json(serde_json::json!({
|
||||
"error": {
|
||||
"message": upstream_message,
|
||||
"type": "invalid_request_error",
|
||||
"code": "invalid_refresh_token"
|
||||
}
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let service = CodexCredentialRefreshService::new(
|
||||
store.repository(),
|
||||
Arc::new(OpenAiTokenClient::new(
|
||||
reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.build()
|
||||
.expect("test HTTP client"),
|
||||
TokenClientConfig {
|
||||
client_id: "test-public-client".to_owned(),
|
||||
token_endpoint: format!("{}/oauth/token", server.uri()),
|
||||
},
|
||||
)),
|
||||
Arc::new(RefreshLeases),
|
||||
Arc::new(RefreshCredentialState),
|
||||
policy,
|
||||
);
|
||||
let account_id = "acct_retryable_unauthorized";
|
||||
let expires_at = SystemTime::now()
|
||||
.checked_sub(Duration::from_secs(30 * 60))
|
||||
.expect("expired access token");
|
||||
seed_refreshable_account(&store, account_id, expires_at, None).await;
|
||||
|
||||
let outcomes = service.refresh_due().await.expect("refresh cycle");
|
||||
|
||||
assert!(matches!(
|
||||
outcomes.as_slice(),
|
||||
[CodexCredentialRefreshOutcome::Transient {
|
||||
account_id: deferred_account_id,
|
||||
}] if deferred_account_id == account_id
|
||||
));
|
||||
let account = store.account(account_id).expect("deferred account");
|
||||
assert_eq!(account.credential_state(), CredentialState::Ready);
|
||||
assert_eq!(
|
||||
account.last_error_reason(),
|
||||
Some(AccountErrorReason::AccessTokenExpired)
|
||||
);
|
||||
assert_eq!(account.last_error_message(), Some(upstream_message));
|
||||
let retry_at = account.next_refresh_at().expect("next refresh attempt");
|
||||
assert!(retry_at > SystemTime::now());
|
||||
assert!(
|
||||
retry_at
|
||||
<= expires_at
|
||||
.checked_add(Duration::from_secs(2 * 60 * 60))
|
||||
.expect("recovery deadline")
|
||||
);
|
||||
let projection = account.status_projection(SystemTime::now(), None);
|
||||
assert_eq!(projection.status, AccountStatus::Error);
|
||||
assert_eq!(projection.error_message.as_deref(), Some(upstream_message));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn scheduled_refresh_uses_the_final_message_after_the_two_hour_window() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let policy = MutableRuntimePolicy::new(Duration::from_secs(5 * 60));
|
||||
let upstream_message = "Invalid refresh token.";
|
||||
let service = CodexCredentialRefreshService::new(
|
||||
store.repository(),
|
||||
Arc::new(FailingRefresher {
|
||||
failure: RefreshFailure::Transport {
|
||||
message: Some(upstream_message.to_owned()),
|
||||
upstream: None,
|
||||
},
|
||||
}),
|
||||
Arc::new(RefreshLeases),
|
||||
Arc::new(RefreshCredentialState),
|
||||
policy,
|
||||
);
|
||||
let account_id = "acct_exhausted_unauthorized";
|
||||
seed_refreshable_account(
|
||||
&store,
|
||||
account_id,
|
||||
SystemTime::now()
|
||||
.checked_sub(Duration::from_secs(2 * 60 * 60 + 60))
|
||||
.expect("expired recovery window"),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
let outcomes = service.refresh_due().await.expect("refresh cycle");
|
||||
|
||||
assert!(matches!(
|
||||
outcomes.as_slice(),
|
||||
[CodexCredentialRefreshOutcome::Invalidated {
|
||||
account_id: invalidated_account_id,
|
||||
}] if invalidated_account_id == account_id
|
||||
));
|
||||
let account = store.account(account_id).expect("invalidated account");
|
||||
assert_eq!(account.credential_state(), CredentialState::Expired);
|
||||
assert_eq!(account.last_error_message(), Some(upstream_message));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn scheduled_refresh_success_clears_a_deferred_failure_message() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let policy = MutableRuntimePolicy::new(Duration::from_secs(5 * 60));
|
||||
let account_id = "acct_recovered_after_backoff";
|
||||
seed_refreshable_account(
|
||||
&store,
|
||||
account_id,
|
||||
SystemTime::now()
|
||||
.checked_sub(Duration::from_secs(60))
|
||||
.expect("expired access token"),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
let account = store.account(account_id).expect("seeded account");
|
||||
store
|
||||
.apply_state_change(AccountStateChange {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
credential_state: CredentialState::Ready,
|
||||
observed_at: SystemTime::now(),
|
||||
error_reason: Some(AccountErrorReason::AccessTokenExpired),
|
||||
message: Some("Invalid refresh token.".to_owned()),
|
||||
})
|
||||
.await
|
||||
.expect("seed deferred failure");
|
||||
let service = refresh_service(&store, SingleUseRefresher::new(), policy);
|
||||
|
||||
let outcomes = service.refresh_due().await.expect("refresh cycle");
|
||||
|
||||
assert!(matches!(
|
||||
outcomes.as_slice(),
|
||||
[CodexCredentialRefreshOutcome::Refreshed {
|
||||
account_id: refreshed_account_id,
|
||||
..
|
||||
}] if refreshed_account_id == account_id
|
||||
));
|
||||
let account = store.account(account_id).expect("refreshed account");
|
||||
assert_eq!(account.credential_state(), CredentialState::Ready);
|
||||
assert_eq!(account.last_error_reason(), None);
|
||||
assert_eq!(account.last_error_message(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn scheduled_refresh_persists_a_returned_id_token() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
@@ -302,6 +589,48 @@ async fn scheduled_refresh_persists_a_returned_id_token() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn scheduled_refresh_uses_the_rotated_access_token_jwt_expiration() {
|
||||
let expires_at = Utc
|
||||
.timestamp_opt(2_000_000_000, 0)
|
||||
.single()
|
||||
.expect("valid test timestamp");
|
||||
let payload = URL_SAFE_NO_PAD.encode(
|
||||
serde_json::to_vec(&serde_json::json!({"exp": expires_at.timestamp()}))
|
||||
.expect("test JWT payload"),
|
||||
);
|
||||
let refresher = SingleUseRefresher::with_access_token(format!(
|
||||
"unverified-header.{payload}.unverified-signature"
|
||||
));
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let service = refresh_service(
|
||||
&store,
|
||||
refresher,
|
||||
MutableRuntimePolicy::new(Duration::from_secs(5 * 60)),
|
||||
);
|
||||
let account_id = "acct_rotated_access_expiration";
|
||||
seed_refreshable_account(
|
||||
&store,
|
||||
account_id,
|
||||
SystemTime::now()
|
||||
.checked_add(Duration::from_secs(120))
|
||||
.expect("test expiry"),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
service.refresh_due().await.expect("refresh cycle");
|
||||
|
||||
assert_eq!(
|
||||
store
|
||||
.account(account_id)
|
||||
.expect("refreshed account")
|
||||
.access_token_expires_at()
|
||||
.map(DateTime::<Utc>::from),
|
||||
Some(expires_at)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn scheduled_refresh_preserves_the_stored_token_set_when_rotation_fields_are_omitted() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use provider_openai::credential::token_client::{
|
||||
AuthorizationCodeExchangeError, AuthorizationCodeExchanger, AuthorizationCodeGrant,
|
||||
OpenAiTokenClient, RefreshFailure, TokenClientConfig, TokenRefresher,
|
||||
@@ -46,13 +44,16 @@ async fn oversized_chunked_oauth_response_should_fail_closed_and_redact_body() {
|
||||
.expect_err("oversized response must fail closed before body classification");
|
||||
let diagnostic = format!("{failure:?} {failure}");
|
||||
|
||||
assert_eq!(failure, RefreshFailure::Transport);
|
||||
assert!(matches!(
|
||||
failure,
|
||||
RefreshFailure::Transport { upstream: None, .. }
|
||||
));
|
||||
assert!(!diagnostic.contains(marker));
|
||||
assert!(!diagnostic.contains("refresh-secret-request-marker"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bounded_oauth_response_should_parse_lifetime_and_rotated_token() {
|
||||
async fn bounded_oauth_response_should_parse_the_official_rotated_token_fields() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth/token"))
|
||||
@@ -74,7 +75,6 @@ async fn bounded_oauth_response_should_parse_lifetime_and_rotated_token() {
|
||||
assert_eq!(tokens.access_token.as_deref(), Some("access-rotated"));
|
||||
assert_eq!(tokens.refresh_token.as_deref(), Some("refresh-rotated"));
|
||||
assert_eq!(tokens.id_token.as_deref(), Some("header.e30.signature"));
|
||||
assert_eq!(tokens.expires_in, Some(Duration::from_secs(3600)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -95,7 +95,10 @@ async fn refresh_response_should_reject_a_malformed_rotated_id_token() {
|
||||
.await
|
||||
.expect_err("malformed ID token must not replace the stored token set");
|
||||
|
||||
assert_eq!(failure, RefreshFailure::Transport);
|
||||
assert!(matches!(
|
||||
failure,
|
||||
RefreshFailure::Transport { upstream: None, .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -116,7 +119,6 @@ async fn refresh_response_should_allow_all_rotated_tokens_to_be_omitted() {
|
||||
assert!(tokens.access_token.is_none());
|
||||
assert!(tokens.refresh_token.is_none());
|
||||
assert!(tokens.id_token.is_none());
|
||||
assert!(tokens.expires_in.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -142,7 +144,7 @@ async fn refresh_response_keeps_the_upstream_rotated_token_verbatim() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_should_exchange_the_official_form_fields() {
|
||||
async fn refresh_should_exchange_the_official_json_fields() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/oauth/token"))
|
||||
@@ -160,11 +162,24 @@ async fn refresh_should_exchange_the_official_form_fields() {
|
||||
.await
|
||||
.expect("refresh succeeds");
|
||||
let requests = server.received_requests().await.expect("received request");
|
||||
let body = String::from_utf8(requests[0].body.clone()).expect("form body is UTF-8");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(&requests[0].body).expect("JSON request body");
|
||||
|
||||
assert!(body.contains("grant_type=refresh_token"));
|
||||
assert!(body.contains("client_id=test-public-client"));
|
||||
assert!(body.contains("refresh_token=refresh+secret"));
|
||||
assert_eq!(
|
||||
body,
|
||||
serde_json::json!({
|
||||
"client_id": "test-public-client",
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": "refresh secret"
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0]
|
||||
.headers
|
||||
.get("content-type")
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("application/json")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -273,72 +288,165 @@ async fn authorization_code_exchange_requires_id_token_and_json_response() {
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_token_reuse_should_be_classified_as_invalid_grant() {
|
||||
let upstream_message = "Your refresh token has already been used to generate a new access token. Please try signing in again.";
|
||||
let failure = refresh_failure(
|
||||
400,
|
||||
r#"{"error":"invalid_grant","error_description":"refresh_token_reused"}"#,
|
||||
r#"{"error":{"message":"Your refresh token has already been used to generate a new access token. Please try signing in again.","type":"invalid_request_error","param":null,"code":"refresh_token_reused"}}"#,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(failure, RefreshFailure::InvalidGrant);
|
||||
let RefreshFailure::InvalidGrant { message, upstream } = failure else {
|
||||
panic!("refresh-token reuse must be terminal");
|
||||
};
|
||||
assert_eq!(message.as_deref(), Some(upstream_message));
|
||||
let upstream = upstream.expect("complete upstream failure");
|
||||
assert_eq!(upstream.status(), 400);
|
||||
assert_eq!(upstream.code(), Some("refresh_token_reused"));
|
||||
assert_eq!(upstream.error_type(), Some("invalid_request_error"));
|
||||
assert!(upstream.body().contains(r#""code":"refresh_token_reused""#));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn generic_invalid_grant_should_remain_transient_like_official_codex() {
|
||||
let body = r#"{"error":"invalid_grant","error_description":"not the official refresh error message field"}"#;
|
||||
let failure = refresh_failure(400, body).await;
|
||||
let diagnostic = format!("{failure:?} {failure}");
|
||||
|
||||
let RefreshFailure::Transport { message, upstream } = failure else {
|
||||
panic!("unknown official refresh code must remain transient");
|
||||
};
|
||||
assert_eq!(message, None);
|
||||
let upstream = upstream.expect("complete upstream failure");
|
||||
assert_eq!(upstream.code(), Some("invalid_grant"));
|
||||
assert_eq!(upstream.body(), body);
|
||||
assert!(!diagnostic.contains("error_description"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unauthorized_should_preserve_upstream_details_and_remain_retryable() {
|
||||
let body = r#"{
|
||||
"error": {
|
||||
"message": "Invalid refresh token.",
|
||||
"type": "invalid_request_error",
|
||||
"param": null,
|
||||
"code": "invalid_refresh_token"
|
||||
}
|
||||
}"#;
|
||||
let failure = refresh_failure(401, body).await;
|
||||
|
||||
let RefreshFailure::Transport { message, upstream } = failure else {
|
||||
panic!("production policy gives every 401 a bounded recovery window");
|
||||
};
|
||||
assert_eq!(message.as_deref(), Some("Invalid refresh token."));
|
||||
let upstream = upstream.expect("complete upstream failure");
|
||||
assert_eq!(upstream.status(), 401);
|
||||
assert_eq!(upstream.code(), Some("invalid_refresh_token"));
|
||||
assert_eq!(upstream.error_type(), Some("invalid_request_error"));
|
||||
assert_eq!(upstream.body(), body);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unauthorized_should_back_off_even_with_a_recognized_refresh_code() {
|
||||
let failure = refresh_failure(
|
||||
401,
|
||||
r#"{"error":{"code":"refresh_token_expired","message":"Refresh token expired."}}"#,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_transport_failure(
|
||||
&failure,
|
||||
401,
|
||||
Some("Refresh token expired."),
|
||||
r#"{"error":{"code":"refresh_token_expired","message":"Refresh token expired."}}"#,
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn deactivated_account_should_be_classified_as_banned() {
|
||||
let failure = refresh_failure(
|
||||
403,
|
||||
r#"{"error":"access_denied","error_description":"account has been deactivated"}"#,
|
||||
r#"{"error":{"code":"access_denied","message":"account has been deactivated"}}"#,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(failure, RefreshFailure::Banned);
|
||||
let RefreshFailure::Banned { message, upstream } = failure else {
|
||||
panic!("deactivated account must be banned");
|
||||
};
|
||||
assert_eq!(message.as_deref(), Some("account has been deactivated"));
|
||||
assert_eq!(upstream.expect("complete upstream failure").status(), 403);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn generic_banned_text_should_not_impersonate_the_deactivation_contract() {
|
||||
let failure = refresh_failure(403, "account is banned").await;
|
||||
|
||||
assert_eq!(failure, RefreshFailure::Transport);
|
||||
assert_transport_failure(&failure, 403, None, "account is banned");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unregistered_disabled_account_text_should_remain_a_transport_failure() {
|
||||
let failure = refresh_failure(400, "account disabled").await;
|
||||
|
||||
assert_eq!(failure, RefreshFailure::Transport);
|
||||
assert_transport_failure(&failure, 400, None, "account disabled");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_text_should_not_disable_the_oauth_credential() {
|
||||
let failure = refresh_failure(400, "quota exceeded").await;
|
||||
|
||||
assert_eq!(failure, RefreshFailure::Transport);
|
||||
assert_transport_failure(&failure, 400, None, "quota exceeded");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn token_revoked_text_without_invalid_grant_should_remain_temporary() {
|
||||
let failure = refresh_failure(400, "token_revoked").await;
|
||||
|
||||
assert_eq!(failure, RefreshFailure::Transport);
|
||||
assert_transport_failure(&failure, 400, None, "token_revoked");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn server_error_or_rate_limit_must_stay_transient_even_with_oauth_error_body() {
|
||||
// 5xx/429(含 CDN/网关页恰好嵌入 invalid_grant 字样)按状态码判瞬态,
|
||||
// 不因正文子串把账号永久终态——正文只在 4xx 时才是权威 OAuth 错误。
|
||||
async fn recognized_refresh_code_should_be_permanent_on_non_unauthorized_status() {
|
||||
for status in [500, 502, 503, 429] {
|
||||
let failure = refresh_failure(
|
||||
status,
|
||||
r#"{"error":"invalid_grant","error_description":"refresh_token_expired"}"#,
|
||||
r#"{"error":{"code":"refresh_token_expired","message":"Refresh token expired."}}"#,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(
|
||||
failure,
|
||||
RefreshFailure::Transport,
|
||||
"status {status} must classify as transient"
|
||||
assert!(
|
||||
matches!(failure, RefreshFailure::InvalidGrant { .. }),
|
||||
"official refresh code must take precedence for status {status}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unknown_server_error_or_rate_limit_should_remain_transient() {
|
||||
let body = r#"{"error":{"code":"temporarily_unavailable","message":"Try again."}}"#;
|
||||
for status in [500, 502, 503, 429] {
|
||||
let failure = refresh_failure(status, body).await;
|
||||
assert_transport_failure(&failure, status, Some("Try again."), body);
|
||||
}
|
||||
}
|
||||
|
||||
fn assert_transport_failure(
|
||||
failure: &RefreshFailure,
|
||||
status: u16,
|
||||
message: Option<&str>,
|
||||
body: &str,
|
||||
) {
|
||||
let RefreshFailure::Transport {
|
||||
message: actual_message,
|
||||
upstream,
|
||||
} = failure
|
||||
else {
|
||||
panic!("status {status} must classify as transient");
|
||||
};
|
||||
assert_eq!(actual_message.as_deref(), message);
|
||||
let upstream = upstream.as_deref().expect("complete upstream failure");
|
||||
assert_eq!(upstream.status(), status);
|
||||
assert_eq!(upstream.body(), body);
|
||||
}
|
||||
|
||||
async fn refresh_failure(status: u16, body: &str) -> RefreshFailure {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
|
||||
@@ -6,6 +6,9 @@ use std::time::{Duration, SystemTime};
|
||||
use base64::{Engine as _, engine::general_purpose::STANDARD};
|
||||
use chrono::Utc;
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use gateway_core::engine::continuation::{
|
||||
ContinuationBinding, NativeContinuationPin, PreviousResponseId,
|
||||
};
|
||||
use gateway_core::engine::credential::{
|
||||
AccountFeedbackStats, CredentialState, OpaqueProviderData, ProviderAccountId,
|
||||
ProviderAccountStore as _, QuotaAccessChange, QuotaAccessState, QuotaEvidence,
|
||||
@@ -16,9 +19,12 @@ use gateway_core::engine::{
|
||||
AccountAttemptContext, AttemptContext, CancellationToken, ContinuationAttempt, ModelRequestId,
|
||||
ProviderAccountStateOwner, RequestAttemptContext, UpstreamSendState,
|
||||
};
|
||||
use gateway_core::error::ProviderErrorKind;
|
||||
use gateway_core::error::{ContinuationFailure, ProviderErrorKind};
|
||||
use gateway_core::event::GatewayEvent;
|
||||
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload, ProviderSessionState};
|
||||
use gateway_core::operation::{
|
||||
CapabilityRequirements, GenerateRequest, Operation, OperationKind, ProtocolPayload,
|
||||
ProviderSessionState,
|
||||
};
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::routing::{
|
||||
ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind,
|
||||
@@ -267,7 +273,8 @@ fn planned_request(provider_name: &str, operation: Operation) -> ProviderRequest
|
||||
vec![ProviderModel::new(
|
||||
provider,
|
||||
upstream_model,
|
||||
ModelCapabilities::new(BTreeSet::from([operation.kind()]), Some(32_000)),
|
||||
ModelCapabilities::new(BTreeSet::from([operation.kind()]), Some(32_000))
|
||||
.with_upstream_feature_validation(),
|
||||
)],
|
||||
Vec::new(),
|
||||
)
|
||||
@@ -288,6 +295,8 @@ fn contract_account_scope() -> Arc<FrozenAccountScope> {
|
||||
let provider = ProviderKind::new("openai").expect("provider");
|
||||
let accounts = [
|
||||
"acct_affinity",
|
||||
"acct_affinity_switch_a",
|
||||
"acct_affinity_switch_b",
|
||||
"acct_atomic_failure",
|
||||
"acct_bare_atomic_failure",
|
||||
"acct_bounded_replay_grace",
|
||||
@@ -319,6 +328,7 @@ fn contract_account_scope() -> Arc<FrozenAccountScope> {
|
||||
"acct_truncated_stream",
|
||||
"acct_usage_limit_request_path",
|
||||
"acct_websocket_close",
|
||||
"acct_websocket_turn_state",
|
||||
]
|
||||
.into_iter()
|
||||
.map(|id| {
|
||||
@@ -389,6 +399,33 @@ fn context_with_state_owner(request_id: &str, owner_account_id: &str) -> Attempt
|
||||
)
|
||||
}
|
||||
|
||||
fn replay_any_context(request_id: &str, owner_account_id: &str) -> AttemptContext {
|
||||
let owner_account = ProviderAccountId::new(owner_account_id).expect("owner account id");
|
||||
let provider = ProviderKind::new("openai").expect("provider");
|
||||
let owner = ProviderAccountStateOwner::new(provider.clone(), owner_account.clone());
|
||||
let continuation = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("client-previous-response"),
|
||||
PreviousResponseId::new("upstream-previous-response"),
|
||||
ClientApiKeyId::new("key_openai_contract").expect("client key id"),
|
||||
provider,
|
||||
owner_account,
|
||||
);
|
||||
AttemptContext::new(
|
||||
RequestAttemptContext::new(
|
||||
ModelRequestId::new(request_id).expect("request id"),
|
||||
ClientApiKeyId::new("key_openai_contract").expect("client key id"),
|
||||
),
|
||||
NonZeroU32::new(2).expect("attempt"),
|
||||
SystemTime::now() + Duration::from_secs(30),
|
||||
account_policy(),
|
||||
AccountAttemptContext::new(BTreeSet::new(), None, Some(owner))
|
||||
.with_account_scope(contract_account_scope()),
|
||||
Some(ContinuationBinding::Pinned(continuation)),
|
||||
CancellationToken::new(),
|
||||
)
|
||||
.with_continuation_attempt(ContinuationAttempt::ReplayAny)
|
||||
}
|
||||
|
||||
async fn capture_scoped_http_request(
|
||||
request_id: &str,
|
||||
selected_account_id: &str,
|
||||
@@ -433,6 +470,76 @@ async fn capture_scoped_http_request(
|
||||
requests.pop().expect("single scoped request")
|
||||
}
|
||||
|
||||
async fn capture_turn_state_request(
|
||||
request_id: &str,
|
||||
previous_turn_id: Option<&str>,
|
||||
current_turn_id: Option<&str>,
|
||||
client_turn_state: Option<&str>,
|
||||
) -> wiremock::Request {
|
||||
let account_id = "acct_session_affinity";
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, account_id).await;
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
.set_body_string(CAPTURE_COMPLETED_SSE),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let mut session_state = Map::from_iter([
|
||||
("account_id".to_owned(), json!(account_id)),
|
||||
("conversation_id".to_owned(), json!("conversation")),
|
||||
("turn_state".to_owned(), json!("previous-turn-state")),
|
||||
("continuation_scope".to_owned(), json!("persisted")),
|
||||
]);
|
||||
if let Some(previous_turn_id) = previous_turn_id {
|
||||
session_state.insert("client_turn_id".to_owned(), json!(previous_turn_id));
|
||||
}
|
||||
let mut protocol_context = Map::from_iter([("use_websocket".to_owned(), json!(false))]);
|
||||
if let Some(current_turn_id) = current_turn_id {
|
||||
protocol_context.insert("turn_id".to_owned(), json!(current_turn_id));
|
||||
}
|
||||
if let Some(client_turn_state) = client_turn_state {
|
||||
protocol_context.insert("turn_state".to_owned(), json!(client_turn_state));
|
||||
}
|
||||
let operation = Operation::Generate(
|
||||
GenerateRequest::from_protocol_payload(
|
||||
ProtocolPayload::json_object(
|
||||
"openai",
|
||||
Map::from_iter([
|
||||
("model".to_owned(), json!("gpt-5.4")),
|
||||
("input".to_owned(), json!("current input")),
|
||||
]),
|
||||
)
|
||||
.expect("OpenAI payload")
|
||||
.with_context(protocol_context),
|
||||
)
|
||||
.with_provider_session_state(
|
||||
ProviderSessionState::new("openai", session_state).expect("provider session state"),
|
||||
),
|
||||
);
|
||||
let mut stream = provider_with_base_url(&store, server.uri())
|
||||
.execute(
|
||||
planned_request("openai", operation),
|
||||
context(request_id, CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
.expect("prepare turn-state provider stream");
|
||||
while let Some(event) = stream.next().await {
|
||||
event.expect("turn-state provider response");
|
||||
}
|
||||
let mut requests = server
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("captured turn-state request");
|
||||
assert_eq!(requests.len(), 1);
|
||||
requests.pop().expect("single turn-state request")
|
||||
}
|
||||
|
||||
fn captured_header_values(request: &wiremock::Request, name: &str) -> Vec<Vec<u8>> {
|
||||
request
|
||||
.headers
|
||||
@@ -789,6 +896,94 @@ async fn websocket_close_details_are_only_exposed_through_the_client_error() {
|
||||
assert_eq!(detail.error_type(), Some("websocket_close_error"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_turn_state_metadata_is_exposed_through_response_observation() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_websocket_turn_state").await;
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("bind WebSocket listener");
|
||||
let base_url = format!(
|
||||
"http://{}",
|
||||
listener.local_addr().expect("listener address")
|
||||
);
|
||||
let server = tokio::spawn(async move {
|
||||
let (stream, _) = listener
|
||||
.accept()
|
||||
.await
|
||||
.expect("accept WebSocket connection");
|
||||
let mut websocket = accept_codex_test_websocket(stream).await;
|
||||
let _request = websocket
|
||||
.next()
|
||||
.await
|
||||
.expect("WebSocket request")
|
||||
.expect("valid WebSocket request");
|
||||
websocket
|
||||
.send(Message::Text(
|
||||
json!({
|
||||
"type": "response.metadata",
|
||||
"headers": {
|
||||
"x-codex-turn-state": ["turn-state-from-websocket"]
|
||||
}
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.expect("send response metadata");
|
||||
websocket
|
||||
.send(Message::Text(
|
||||
json!({
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_websocket_turn_state",
|
||||
"model": "gpt-5.4",
|
||||
"status": "completed",
|
||||
"output": [],
|
||||
"usage": {
|
||||
"input_tokens": 1,
|
||||
"output_tokens": 1,
|
||||
"total_tokens": 2
|
||||
}
|
||||
}
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.expect("send completed response");
|
||||
websocket.close(None).await.expect("close WebSocket");
|
||||
});
|
||||
|
||||
let mut stream = provider_with_base_url(&store, base_url)
|
||||
.execute(
|
||||
planned_request("openai", generate_operation()),
|
||||
context("req_websocket_turn_state", CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
.expect("prepare WebSocket provider stream");
|
||||
let mut observed_turn_state = false;
|
||||
while let Some(event) = stream.next().await {
|
||||
let event = event.expect("provider event");
|
||||
if let Some(observation) = event.response_observation() {
|
||||
observed_turn_state |= observation.client_headers().iter().any(|header| {
|
||||
header.name().eq_ignore_ascii_case("x-codex-turn-state")
|
||||
&& header.value().as_ref() == b"turn-state-from-websocket"
|
||||
});
|
||||
}
|
||||
if event
|
||||
.canonical_facts()
|
||||
.iter()
|
||||
.any(|event| matches!(event, GatewayEvent::Completed(_)))
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
server.await.expect("WebSocket server");
|
||||
|
||||
assert!(observed_turn_state);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disabled_account_is_excluded_from_normal_scheduling() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
@@ -971,9 +1166,9 @@ async fn cross_account_scope_sanitizes_only_known_turn_metadata_fields() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn continuation_replay_should_use_the_client_full_input_without_legacy_transcript() {
|
||||
async fn cross_account_continuation_probe_should_preserve_delta_input_and_previous_response() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_client_history").await;
|
||||
create_account(&store, "acct_scope_new").await;
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
@@ -985,33 +1180,61 @@ async fn continuation_replay_should_use_the_client_full_input_without_legacy_tra
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let client_input = json!([
|
||||
{"role": "user", "content": "complete client history"},
|
||||
{"role": "assistant", "content": "previous answer"},
|
||||
{"role": "user", "content": "current turn"}
|
||||
let delta_input = json!([
|
||||
{
|
||||
"type": "message",
|
||||
"id": "delta-message-id",
|
||||
"encrypted_content": "opaque-delta-ciphertext",
|
||||
"role": "user",
|
||||
"content": "current turn only"
|
||||
},
|
||||
{
|
||||
"type": "compaction",
|
||||
"id": "delta-compaction-id",
|
||||
"encrypted_content": "opaque-compaction-ciphertext"
|
||||
}
|
||||
]);
|
||||
let generation = GenerateRequest::from_protocol_payload(
|
||||
ProtocolPayload::json_object(
|
||||
"openai",
|
||||
Map::from_iter([
|
||||
("model".to_owned(), json!("gpt-5.4")),
|
||||
("input".to_owned(), client_input.clone()),
|
||||
("input".to_owned(), delta_input.clone()),
|
||||
(
|
||||
"previous_response_id".to_owned(),
|
||||
json!("client-previous-response"),
|
||||
),
|
||||
("authorization".to_owned(), json!("client-extension-secret")),
|
||||
(
|
||||
"turnMetadata".to_owned(),
|
||||
json!(r#"{"account_id":"acct_scope_old","safe":true}"#),
|
||||
),
|
||||
(
|
||||
"client_metadata".to_owned(),
|
||||
json!({
|
||||
"x-codex-turn-state": "old-account-turn-state",
|
||||
"x-codex-turn-metadata": "old-account-turn-metadata",
|
||||
"future": "keep"
|
||||
}),
|
||||
),
|
||||
]),
|
||||
)
|
||||
.expect("OpenAI payload")
|
||||
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(false))])),
|
||||
.with_context(Map::from_iter([
|
||||
("use_websocket".to_owned(), json!(false)),
|
||||
("turn_id".to_owned(), json!("turn-probe")),
|
||||
("turn_state".to_owned(), json!("old-account-turn-state")),
|
||||
])),
|
||||
)
|
||||
.with_provider_session_state(
|
||||
ProviderSessionState::new(
|
||||
"openai",
|
||||
Map::from_iter([
|
||||
("account_id".to_owned(), json!("acct_client_history")),
|
||||
("account_id".to_owned(), json!("acct_scope_old")),
|
||||
("conversation_id".to_owned(), json!("conversation")),
|
||||
("continuation_scope".to_owned(), json!("replay_required")),
|
||||
(
|
||||
"transcript".to_owned(),
|
||||
json!([{"client_input": {"role": "user", "content": "stale proxy copy"}}]),
|
||||
),
|
||||
("turn_state".to_owned(), json!("old-account-turn-state")),
|
||||
("client_turn_id".to_owned(), json!("turn-probe")),
|
||||
("continuation_scope".to_owned(), json!("persisted")),
|
||||
]),
|
||||
)
|
||||
.expect("provider session state"),
|
||||
@@ -1019,21 +1242,84 @@ async fn continuation_replay_should_use_the_client_full_input_without_legacy_tra
|
||||
let mut stream = provider_with_base_url(&store, server.uri())
|
||||
.execute(
|
||||
planned_request("openai", Operation::Generate(generation)),
|
||||
context("req_client_history", CancellationToken::new())
|
||||
.with_continuation_attempt(ContinuationAttempt::ReplayAny),
|
||||
replay_any_context("req_cross_account_probe", "acct_scope_old"),
|
||||
)
|
||||
.await
|
||||
.expect("prepare replay stream");
|
||||
.expect("prepare continuation probe stream");
|
||||
while let Some(event) = stream.next().await {
|
||||
event.expect("replay response");
|
||||
event.expect("continuation probe response");
|
||||
}
|
||||
let requests = server
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("captured replay request");
|
||||
.expect("captured continuation probe request");
|
||||
let body: Value = captured_request_body(&requests[0]);
|
||||
|
||||
assert_eq!(body.get("input"), Some(&client_input));
|
||||
assert_eq!(body.get("input"), Some(&delta_input));
|
||||
assert_eq!(
|
||||
body.get("previous_response_id"),
|
||||
Some(&json!("upstream-previous-response"))
|
||||
);
|
||||
assert!(body.get("authorization").is_none());
|
||||
assert!(body.get("turnMetadata").is_none());
|
||||
assert!(
|
||||
body.pointer("/client_metadata/x-codex-turn-state")
|
||||
.is_none()
|
||||
);
|
||||
assert!(
|
||||
body.pointer("/client_metadata/x-codex-turn-metadata")
|
||||
.is_none()
|
||||
);
|
||||
assert_eq!(
|
||||
body.pointer("/client_metadata/future"),
|
||||
Some(&json!("keep"))
|
||||
);
|
||||
assert!(captured_header_values(&requests[0], "x-codex-turn-state").is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn matching_turn_id_should_restore_previous_turn_state() {
|
||||
let request = capture_turn_state_request(
|
||||
"req_restore_same_turn_state",
|
||||
Some("turn-same"),
|
||||
Some("turn-same"),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
captured_header_values(&request, "x-codex-turn-state"),
|
||||
vec![b"previous-turn-state".to_vec()]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn matching_turn_id_should_prefer_the_latest_provider_state_over_a_stale_client_echo() {
|
||||
let request = capture_turn_state_request(
|
||||
"req_replace_stale_client_turn_state",
|
||||
Some("turn-same"),
|
||||
Some("turn-same"),
|
||||
Some("stale-client-turn-state"),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
captured_header_values(&request, "x-codex-turn-state"),
|
||||
vec![b"previous-turn-state".to_vec()]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn new_or_unidentified_turn_should_not_restore_previous_turn_state() {
|
||||
for (request_id, previous_turn_id, current_turn_id) in [
|
||||
("req_new_turn_state", Some("turn-old"), Some("turn-new")),
|
||||
("req_unidentified_turn_state", None, Some("turn-new")),
|
||||
("req_missing_current_turn_state", Some("turn-old"), None),
|
||||
] {
|
||||
let request =
|
||||
capture_turn_state_request(request_id, previous_turn_id, current_turn_id, None).await;
|
||||
assert!(captured_header_values(&request, "x-codex-turn-state").is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -1259,6 +1545,134 @@ async fn explicit_session_id_should_override_turn_specific_prompt_cache_keys_for
|
||||
assert_eq!(keys[0], keys[1]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn affinity_quota_switch_should_clear_old_turn_state_without_a_provider_state_owner() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let first_account_id = "acct_affinity_switch_a";
|
||||
let second_account_id = "acct_affinity_switch_b";
|
||||
create_account(&store, first_account_id).await;
|
||||
let affinity = Arc::new(MemorySessionAffinity::default());
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
.set_body_string(CAPTURE_COMPLETED_SSE),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let (provider, quota) = provider_and_quota_with_affinity_and_base_url_and_leases(
|
||||
&store,
|
||||
Arc::clone(&affinity),
|
||||
server.uri(),
|
||||
Arc::new(TestLeaseCoordinator::default()),
|
||||
);
|
||||
let session_id = "stable-affinity-switch-session";
|
||||
let first_payload = ProtocolPayload::json_object(
|
||||
"openai",
|
||||
Map::from_iter([
|
||||
("model".to_owned(), json!("gpt-5.4")),
|
||||
("input".to_owned(), json!("first request")),
|
||||
("session_id".to_owned(), json!(session_id)),
|
||||
]),
|
||||
)
|
||||
.expect("OpenAI payload")
|
||||
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(false))]));
|
||||
let first = provider
|
||||
.execute(
|
||||
planned_request(
|
||||
"openai",
|
||||
Operation::Generate(GenerateRequest::from_protocol_payload(first_payload)),
|
||||
),
|
||||
context("req_affinity_switch_first", CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
.expect("prepare first affinity request");
|
||||
drop(first);
|
||||
let affinity_keys = affinity.lookup_keys();
|
||||
assert_eq!(affinity_keys.len(), 1);
|
||||
affinity.seed_binding(
|
||||
&ProviderKind::new("openai").expect("provider"),
|
||||
&affinity_keys[0],
|
||||
ProviderAccountId::new(first_account_id).expect("first account id"),
|
||||
);
|
||||
assert_eq!(affinity.binding_count(), 1);
|
||||
|
||||
create_account(&store, second_account_id).await;
|
||||
let first_account = store.account(first_account_id).expect("first account");
|
||||
let observed_at = SystemTime::now();
|
||||
store
|
||||
.compare_and_swap_quota(QuotaObservation {
|
||||
account_id: first_account.id().clone(),
|
||||
expected_revision: first_account.revision(),
|
||||
quota: OpaqueProviderData::new(Map::new()),
|
||||
observed_at,
|
||||
state: QuotaState::exhausted(QuotaEvidence::ProviderDenied, observed_at, None),
|
||||
})
|
||||
.await
|
||||
.expect("seed exhausted affinity account");
|
||||
let accounts = [
|
||||
store
|
||||
.account(first_account_id)
|
||||
.expect("first account after quota"),
|
||||
store.account(second_account_id).expect("second account"),
|
||||
];
|
||||
quota.prepare_scheduling(&accounts).await;
|
||||
|
||||
let second_payload = ProtocolPayload::json_object(
|
||||
"openai",
|
||||
Map::from_iter([
|
||||
("model".to_owned(), json!("gpt-5.4")),
|
||||
("input".to_owned(), json!("second request")),
|
||||
("session_id".to_owned(), json!(session_id)),
|
||||
]),
|
||||
)
|
||||
.expect("OpenAI payload")
|
||||
.with_context(Map::from_iter([
|
||||
("use_websocket".to_owned(), json!(false)),
|
||||
("turn_id".to_owned(), json!("turn-affinity-switch")),
|
||||
(
|
||||
"turn_state".to_owned(),
|
||||
json!("state-created-by-exhausted-account"),
|
||||
),
|
||||
]));
|
||||
let mut second = provider
|
||||
.execute(
|
||||
planned_request(
|
||||
"openai",
|
||||
Operation::Generate(GenerateRequest::from_protocol_payload(second_payload)),
|
||||
),
|
||||
context("req_affinity_switch_second", CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
.expect("prepare switched affinity request");
|
||||
while let Some(event) = second.next().await {
|
||||
event.expect("switched affinity response");
|
||||
}
|
||||
|
||||
let requests = server
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("captured affinity requests");
|
||||
assert_eq!(requests.len(), 1);
|
||||
let second_request = &requests[0];
|
||||
assert_eq!(
|
||||
captured_request_body(second_request).get("input"),
|
||||
Some(&json!("second request"))
|
||||
);
|
||||
assert!(captured_header_values(second_request, "x-codex-turn-state").is_empty());
|
||||
let expected_account_id = format!("chatgpt-{second_account_id}");
|
||||
assert_eq!(
|
||||
second_request
|
||||
.headers
|
||||
.get("chatgpt-account-id")
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(expected_account_id.as_str())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn thread_spawn_children_should_have_independent_session_affinity_keys() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
@@ -1931,6 +2345,77 @@ async fn bare_response_failed_should_remain_an_atomic_replay_safe_failure() {
|
||||
let _ = server.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn previous_response_not_found_should_remain_client_visible_without_internal_replay_proof() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_client_history").await;
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
.set_body_string(concat!(
|
||||
"event: response.failed\n",
|
||||
"data: {\"type\":\"response.failed\",\"status_code\":400,\"response\":{\"id\":\"resp_history_missing\",\"status\":\"failed\",\"error\":{\"code\":\"previous_response_not_found\",\"message\":\"Previous response was not found. Retrying the full request.\"}}}\n\n"
|
||||
)),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let payload = ProtocolPayload::json_object(
|
||||
"openai",
|
||||
Map::from_iter([
|
||||
("model".to_owned(), json!("gpt-5.4")),
|
||||
(
|
||||
"input".to_owned(),
|
||||
json!([{"role": "user", "content": "delta"}]),
|
||||
),
|
||||
(
|
||||
"previous_response_id".to_owned(),
|
||||
json!("missing-previous-response"),
|
||||
),
|
||||
]),
|
||||
)
|
||||
.expect("OpenAI payload")
|
||||
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(false))]));
|
||||
let mut stream = provider_with_base_url(&store, server.uri())
|
||||
.execute(
|
||||
planned_request(
|
||||
"openai",
|
||||
Operation::Generate(GenerateRequest::from_protocol_payload(payload)),
|
||||
),
|
||||
context("req_history_missing", CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
.expect("prepare missing-history stream");
|
||||
let mut failure = loop {
|
||||
match stream.next().await {
|
||||
Some(Ok(_)) => {}
|
||||
Some(Err(error)) => break error,
|
||||
None => panic!("previous_response_not_found must produce a typed failure"),
|
||||
}
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
failure.continuation_failure(),
|
||||
Some(ContinuationFailure::HistoryUnavailable)
|
||||
);
|
||||
assert!(!failure.replay_is_safe());
|
||||
assert_eq!(
|
||||
failure
|
||||
.upstream_code()
|
||||
.map(gateway_core::error::OpaqueUpstreamValue::as_str),
|
||||
Some("previous_response_not_found")
|
||||
);
|
||||
let events = failure.take_atomic_client_events();
|
||||
assert_eq!(events.len(), 1);
|
||||
assert_eq!(
|
||||
events[0].wire_event().and_then(|wire| wire.event_type()),
|
||||
Some("response.failed")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn continuation_prefetch_over_64_kib_should_commit_wire_without_protocol_failure() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
@@ -2424,3 +2909,36 @@ async fn provider_compiles_catalog_presentation_for_codex_models() {
|
||||
assert!(presentation.verbosity());
|
||||
assert!(!presentation.hidden());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_routes_catalog_model_when_supported_in_api_is_false() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_non_api_model").await;
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/codex/models"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "application/json")
|
||||
.set_body_raw(
|
||||
br#"{"models":[{"slug":"gpt-5.6-sol-wm","display_name":"GPT-5.6-Sol-WM","supported_in_api":false,"visibility":"hide"}]}"#,
|
||||
"application/json",
|
||||
),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let provider = provider_with_base_url(&store, server.uri());
|
||||
|
||||
let capabilities = provider
|
||||
.query_model_capabilities()
|
||||
.await
|
||||
.expect("catalog model capabilities");
|
||||
let model = capabilities.first().expect("catalog model");
|
||||
|
||||
assert!(
|
||||
model
|
||||
.capabilities()
|
||||
.match_requirements(&CapabilityRequirements::new(OperationKind::Generate))
|
||||
.is_some()
|
||||
);
|
||||
}
|
||||
|
||||
@@ -10,10 +10,10 @@ use async_trait::async_trait;
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_core::engine::credential::{
|
||||
AccountConcurrencyLimit, AccountErrorReason, AccountRuntimeSignals, AccountStateChange,
|
||||
AccountWeight, CredentialCasOutcome, CredentialCasUpdate, CredentialRevision, CredentialState,
|
||||
LoadedCredential, NewProviderAccount, OpaqueProviderData, ProviderAccount, ProviderAccountId,
|
||||
ProviderAccountStore, ProviderAccountUpdate, QuotaAccessChange, QuotaObservation, QuotaState,
|
||||
QuotaWriteOutcome,
|
||||
AccountWeight, CredentialCasOutcome, CredentialCasUpdate, CredentialCasUpdateParts,
|
||||
CredentialRevision, CredentialState, LoadedCredential, NewProviderAccount, OpaqueProviderData,
|
||||
ProviderAccount, ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate,
|
||||
QuotaAccessChange, QuotaObservation, QuotaObservationTouch, QuotaState, QuotaWriteOutcome,
|
||||
};
|
||||
use gateway_core::error::{StoreError, StoreErrorKind};
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
@@ -207,7 +207,7 @@ impl ProviderAccountStore for MemoryAccountStore {
|
||||
&self,
|
||||
update: CredentialCasUpdate,
|
||||
) -> Result<CredentialCasOutcome, StoreError> {
|
||||
let (
|
||||
let CredentialCasUpdateParts {
|
||||
account_id,
|
||||
expected_revision,
|
||||
profile,
|
||||
@@ -215,7 +215,8 @@ impl ProviderAccountStore for MemoryAccountStore {
|
||||
has_refresh_token,
|
||||
access_token_expires_at,
|
||||
next_refresh_at,
|
||||
) = update.into_parts();
|
||||
account_state,
|
||||
} = update.into_parts();
|
||||
let mut accounts = self.accounts.lock().expect("account store lock");
|
||||
let stored = accounts
|
||||
.get_mut(&account_id)
|
||||
@@ -226,15 +227,29 @@ impl ProviderAccountStore for MemoryAccountStore {
|
||||
let next = expected_revision
|
||||
.next()
|
||||
.map_err(|_| store_error(StoreErrorKind::Conflict))?;
|
||||
let credential_state = account_state
|
||||
.as_ref()
|
||||
.map_or(stored.account.credential_state(), |state| {
|
||||
state.credential_state
|
||||
});
|
||||
let last_error_reason = account_state.as_ref().map_or_else(
|
||||
|| stored.account.last_error_reason(),
|
||||
|state| state.error_reason,
|
||||
);
|
||||
let last_error_message = account_state.as_ref().map_or_else(
|
||||
|| stored.account.last_error_message().map(str::to_owned),
|
||||
|state| state.message.clone(),
|
||||
);
|
||||
let state_observed_at = account_state.as_ref().map(|state| state.observed_at);
|
||||
stored.account = rebuild_account(
|
||||
&stored.account,
|
||||
AccountRebuild {
|
||||
revision: next,
|
||||
enabled: stored.account.enabled(),
|
||||
credential_state: stored.account.credential_state(),
|
||||
credential_state,
|
||||
quota: stored.account.quota(),
|
||||
last_error_reason: stored.account.last_error_reason(),
|
||||
last_error_message: stored.account.last_error_message().map(str::to_owned),
|
||||
last_error_reason,
|
||||
last_error_message,
|
||||
access_token_expires_at,
|
||||
has_refresh_token,
|
||||
next_refresh_at,
|
||||
@@ -243,6 +258,9 @@ impl ProviderAccountStore for MemoryAccountStore {
|
||||
);
|
||||
stored.credential = credential;
|
||||
stored.quota = None;
|
||||
if let Some(observed_at) = state_observed_at {
|
||||
stored.state_observed_at = Some(observed_at);
|
||||
}
|
||||
Ok(CredentialCasOutcome::Updated(next))
|
||||
}
|
||||
|
||||
@@ -285,6 +303,27 @@ impl ProviderAccountStore for MemoryAccountStore {
|
||||
Ok(QuotaWriteOutcome::Updated)
|
||||
}
|
||||
|
||||
async fn touch_quota_observation(
|
||||
&self,
|
||||
touch: QuotaObservationTouch,
|
||||
) -> Result<QuotaWriteOutcome, StoreError> {
|
||||
let mut accounts = self.accounts.lock().expect("account store lock");
|
||||
let stored = accounts
|
||||
.get_mut(&touch.account_id)
|
||||
.ok_or_else(|| store_error(StoreErrorKind::InvalidData))?;
|
||||
if stored.account.revision() != touch.expected_revision {
|
||||
return Ok(QuotaWriteOutcome::Conflict);
|
||||
}
|
||||
let Some(quota) = &mut stored.quota else {
|
||||
return Ok(QuotaWriteOutcome::Conflict);
|
||||
};
|
||||
if quota.observed_at > touch.observed_at {
|
||||
return Ok(QuotaWriteOutcome::Conflict);
|
||||
}
|
||||
quota.observed_at = touch.observed_at;
|
||||
Ok(QuotaWriteOutcome::Updated)
|
||||
}
|
||||
|
||||
async fn apply_quota_access(
|
||||
&self,
|
||||
change: QuotaAccessChange,
|
||||
@@ -485,8 +524,11 @@ pub(crate) fn agent_identity_service_with_pool(
|
||||
.no_proxy()
|
||||
.build()
|
||||
.expect("agent task client");
|
||||
let registrar =
|
||||
OfficialCodexAgentIdentityTaskRegistrar::new(client).expect("agent task registrar");
|
||||
let registrar = OfficialCodexAgentIdentityTaskRegistrar::new(
|
||||
client,
|
||||
provider_openai::OpenAiConfig::default().wire_profile_state(),
|
||||
)
|
||||
.expect("agent task registrar");
|
||||
Arc::new(CodexAgentIdentityTaskService::new(
|
||||
store.repository(),
|
||||
Arc::new(registrar),
|
||||
@@ -580,6 +622,18 @@ impl MemorySessionAffinity {
|
||||
pub(crate) fn binding_count(&self) -> usize {
|
||||
self.bindings.lock().expect("session affinity lock").len()
|
||||
}
|
||||
|
||||
pub(crate) fn seed_binding(
|
||||
&self,
|
||||
provider_kind: &ProviderKind,
|
||||
key: &str,
|
||||
account_id: ProviderAccountId,
|
||||
) {
|
||||
self.bindings.lock().expect("session affinity lock").insert(
|
||||
(provider_kind.as_str().to_owned(), key.to_owned()),
|
||||
account_id,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderSessionAffinityPort for MemorySessionAffinity {
|
||||
|
||||
@@ -4,6 +4,7 @@ use provider_openai::transport::canonical::{
|
||||
CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalFailure, CodexCanonicalOutcome,
|
||||
};
|
||||
use provider_openai::transport::protocol::websocket::websocket_event_to_sse_frame;
|
||||
use serde_json::json;
|
||||
|
||||
const METADATA_PREFIX_FIXTURE: &str = include_str!(
|
||||
"../../../../gateway-api/tests/openai/responses/fixtures/http_sse/metadata_only_prefix.sse"
|
||||
@@ -262,6 +263,69 @@ fn decoder_should_use_response_service_tier_when_request_omits_it() {
|
||||
)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_add_standard_web_search_call_cost() {
|
||||
let body = concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_search_cost\",\"model\":\"gpt-5.4\"}}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_search_cost\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[{\"type\":\"web_search_call\",\"id\":\"ws_1\",\"status\":\"completed\",\"action\":{\"type\":\"search\"}}],\"usage\":{\"input_tokens\":0,\"output_tokens\":0,\"total_tokens\":0}}}\n\n",
|
||||
);
|
||||
let tools = vec![json!({ "type": "web_search" })];
|
||||
let events = CodexCanonicalDecoder::new("fallback")
|
||||
.with_request_tool_pricing("gpt-5.4", Some(&tools))
|
||||
.push(body.as_bytes())
|
||||
.expect("canonical web search response");
|
||||
|
||||
assert!(canonical_facts(&events).into_iter().any(|event| matches!(
|
||||
event,
|
||||
GatewayEvent::CalculatedCost(cost)
|
||||
if cost.total().amount().scaled() == 100_000_000
|
||||
)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_use_non_reasoning_preview_web_search_price() {
|
||||
let body = concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_preview_search_cost\",\"model\":\"gpt-4o\"}}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_preview_search_cost\",\"model\":\"gpt-4o\",\"status\":\"completed\",\"output\":[{\"type\":\"web_search_call\",\"id\":\"ws_1\",\"status\":\"completed\",\"action\":{\"type\":\"search\"}}],\"usage\":{\"input_tokens\":0,\"output_tokens\":0,\"total_tokens\":0}}}\n\n",
|
||||
);
|
||||
let tools = vec![json!({ "type": "web_search_preview" })];
|
||||
let events = CodexCanonicalDecoder::new("fallback")
|
||||
.with_request_tool_pricing("gpt-4o", Some(&tools))
|
||||
.push(body.as_bytes())
|
||||
.expect("canonical preview web search response");
|
||||
|
||||
assert!(canonical_facts(&events).into_iter().any(|event| matches!(
|
||||
event,
|
||||
GatewayEvent::CalculatedCost(cost)
|
||||
if cost.total().amount().scaled() == 250_000_000
|
||||
)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_fail_closed_for_fixed_block_web_search_content() {
|
||||
let body = concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_fixed_search_cost\",\"model\":\"gpt-4o-mini\"}}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_fixed_search_cost\",\"model\":\"gpt-4o-mini\",\"status\":\"completed\",\"output\":[{\"type\":\"web_search_call\",\"id\":\"ws_1\",\"status\":\"completed\",\"action\":{\"type\":\"search\"}}],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n",
|
||||
);
|
||||
let tools = vec![json!({ "type": "web_search" })];
|
||||
let events = CodexCanonicalDecoder::new("fallback")
|
||||
.with_request_tool_pricing("gpt-4o-mini", Some(&tools))
|
||||
.push(body.as_bytes())
|
||||
.expect("canonical fixed-block search response");
|
||||
|
||||
assert!(
|
||||
canonical_facts(&events)
|
||||
.into_iter()
|
||||
.all(|event| !matches!(event, GatewayEvent::CalculatedCost(_)))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn websocket_decoder_should_prefer_requested_service_tier_over_response_tier_for_billing() {
|
||||
let created = websocket_event_to_sse_frame(
|
||||
@@ -423,6 +487,25 @@ fn decoder_should_preserve_image_tool_tokens_in_canonical_usage() {
|
||||
assert_eq!(image_usage, Some((Some(31), Some(9))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_not_understate_cost_when_image_tool_price_is_ambiguous() {
|
||||
let body = concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_image_cost\",\"model\":\"gpt-5.4\"}}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_image_cost\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"usage\":{\"input_tokens\":12,\"output_tokens\":5,\"total_tokens\":17},\"tool_usage\":{\"image_gen\":{\"input_tokens\":31,\"output_tokens\":9}}}}\n\n",
|
||||
);
|
||||
let events = CodexCanonicalDecoder::new("fallback")
|
||||
.push(body.as_bytes())
|
||||
.expect("canonical image cost response");
|
||||
|
||||
assert!(
|
||||
canonical_facts(&events)
|
||||
.into_iter()
|
||||
.all(|event| !matches!(event, GatewayEvent::CalculatedCost(_)))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_drop_rate_limit_control_and_preserve_metadata_lifecycle_events() {
|
||||
let body = concat!(
|
||||
|
||||
@@ -223,6 +223,7 @@ async fn fetch_should_send_official_catalog_headers_and_query() {
|
||||
"user-agent",
|
||||
"codex_cli_rs/0.144.0 (linux 6.8; x86_64) xterm (codex_cli_rs; 1.0.0)",
|
||||
))
|
||||
.and(header("x-openai-internal-codex-residency", "us"))
|
||||
.and(header("accept", "application/json"))
|
||||
.and(header("x-codex-installation-id", "installation-123"))
|
||||
.respond_with(
|
||||
@@ -254,7 +255,6 @@ async fn fetch_should_send_official_catalog_headers_and_query() {
|
||||
for forbidden in [
|
||||
"content-type",
|
||||
"openai-beta",
|
||||
"x-openai-internal-codex-residency",
|
||||
"x-client-request-id",
|
||||
"session_id",
|
||||
"session-id",
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use provider_openai::transport::{endpoint_url, usage_endpoint_urls};
|
||||
use provider_openai::transport::{endpoint_url, usage_endpoint_url};
|
||||
|
||||
#[test]
|
||||
fn endpoints_should_join_origin_and_backend_paths_without_double_slashes() {
|
||||
@@ -9,12 +9,9 @@ fn endpoints_should_join_origin_and_backend_paths_without_double_slashes() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_endpoints_should_preserve_the_official_backend_api_fallback_order() {
|
||||
fn usage_endpoint_should_use_the_official_wham_path_for_backend_api() {
|
||||
assert_eq!(
|
||||
usage_endpoint_urls("https://chatgpt.com/backend-api"),
|
||||
vec![
|
||||
"https://chatgpt.com/backend-api/wham/usage".to_owned(),
|
||||
"https://chatgpt.com/backend-api/codex/usage".to_owned(),
|
||||
]
|
||||
usage_endpoint_url("https://chatgpt.com/backend-api"),
|
||||
"https://chatgpt.com/backend-api/wham/usage"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -178,6 +178,7 @@ fn websocket_opening_audit_should_redact_sensitive_headers() {
|
||||
"x-codex-turn-metadata".to_owned(),
|
||||
"{\"secret\":true}".to_owned(),
|
||||
),
|
||||
("x-oai-attestation".to_owned(), "device-token".to_owned()),
|
||||
],
|
||||
);
|
||||
let audit = connection.opening_audit_snapshot();
|
||||
@@ -190,6 +191,7 @@ fn websocket_opening_audit_should_redact_sensitive_headers() {
|
||||
"user-agent",
|
||||
"x-client-request-id",
|
||||
"x-codex-turn-metadata",
|
||||
"x-oai-attestation",
|
||||
]
|
||||
);
|
||||
assert!(
|
||||
@@ -313,6 +315,7 @@ async fn backend_websocket_should_forward_context_headers_and_preserve_payload_f
|
||||
("content-type", "application/json"),
|
||||
("accept", "text/event-stream"),
|
||||
("x-codex-installation-id", "install-123"),
|
||||
("openai-beta", "responses_websockets=2026-02-06"),
|
||||
("x-openai-internal-codex-residency", "us"),
|
||||
("x-codex-turn-state", "turn-state"),
|
||||
("x-codex-turn-metadata", "{\"thread_source\":\"subagent\"}"),
|
||||
@@ -382,10 +385,7 @@ async fn backend_http_should_restore_opaque_multivalue_header_bytes_and_lease_id
|
||||
raw_header_values(&raw, "x-openai-future-mode"),
|
||||
vec![b"future-ascii".to_vec(), b"\x80\xff".to_vec()]
|
||||
);
|
||||
assert_eq!(
|
||||
raw_header_values(&raw, "openai-beta"),
|
||||
vec![b"future_responses=v2".to_vec(), b"future_tools=v3".to_vec()]
|
||||
);
|
||||
assert!(raw_header_values(&raw, "openai-beta").is_empty());
|
||||
assert_eq!(
|
||||
raw_header_values(&raw, "x-codex-turn-state"),
|
||||
vec![b"turn-ascii".to_vec(), b"turn-\x80".to_vec()]
|
||||
@@ -396,10 +396,7 @@ async fn backend_http_should_restore_opaque_multivalue_header_bytes_and_lease_id
|
||||
"content-type",
|
||||
b"application/vnd.openai.request+json".as_slice(),
|
||||
),
|
||||
(
|
||||
"x-openai-internal-codex-residency",
|
||||
b"future-region".as_slice(),
|
||||
),
|
||||
("x-openai-internal-codex-residency", b"us".as_slice()),
|
||||
("x-still-valid", b"after-invalid".as_slice()),
|
||||
("authorization", b"Bearer lease-token".as_slice()),
|
||||
("chatgpt-account-id", b"lease-account".as_slice()),
|
||||
@@ -416,7 +413,12 @@ async fn backend_http_should_restore_opaque_multivalue_header_bytes_and_lease_id
|
||||
raw_header_values(&raw, "originator"),
|
||||
vec![b"codex_cli_rs".to_vec()]
|
||||
);
|
||||
for dropped in ["x-oai-attestation", "x-oai-is", "x-oai-is-update"] {
|
||||
for dropped in [
|
||||
"openai-beta",
|
||||
"x-oai-attestation",
|
||||
"x-oai-is",
|
||||
"x-oai-is-update",
|
||||
] {
|
||||
assert!(
|
||||
raw_header_values(&raw, dropped).is_empty(),
|
||||
"unexpected {dropped}"
|
||||
@@ -509,7 +511,7 @@ async fn backend_websocket_should_drop_only_unrepresentable_opaque_header_values
|
||||
);
|
||||
assert_eq!(
|
||||
values("openai-beta"),
|
||||
vec![b"future_responses=v2".to_vec(), b"future_tools=v3".to_vec()]
|
||||
vec![b"responses_websockets=2026-02-06".to_vec()]
|
||||
);
|
||||
assert_eq!(
|
||||
values("accept"),
|
||||
@@ -524,6 +526,10 @@ async fn backend_websocket_should_drop_only_unrepresentable_opaque_header_values
|
||||
vec![test_wire_profile().snapshot().user_agent().into_bytes()]
|
||||
);
|
||||
assert_eq!(values("originator"), vec![b"codex_cli_rs".to_vec()]);
|
||||
assert_eq!(
|
||||
values("x-openai-internal-codex-residency"),
|
||||
vec![b"us".to_vec()]
|
||||
);
|
||||
for dropped in ["x-oai-attestation", "x-oai-is", "x-oai-is-update"] {
|
||||
assert!(values(dropped).is_empty(), "unexpected {dropped}");
|
||||
}
|
||||
@@ -631,7 +637,6 @@ async fn backend_http_should_send_codex_context_without_browser_headers() {
|
||||
"content-type",
|
||||
"cookie",
|
||||
"accept",
|
||||
"openai-beta",
|
||||
"x-openai-internal-codex-residency",
|
||||
"x-client-request-id",
|
||||
"x-codex-installation-id",
|
||||
@@ -654,6 +659,7 @@ async fn backend_http_should_send_codex_context_without_browser_headers() {
|
||||
"sec-fetch-site",
|
||||
"sec-fetch-mode",
|
||||
"sec-fetch-dest",
|
||||
"openai-beta",
|
||||
] {
|
||||
assert!(header_names.iter().all(|name| name != forbidden));
|
||||
}
|
||||
@@ -947,7 +953,12 @@ async fn websocket_should_keep_an_exact_chain_while_new_connections_adopt_the_la
|
||||
)
|
||||
.await
|
||||
.expect("first response");
|
||||
profile.update_cli_release("1.2.4");
|
||||
profile.update_bundled_release(&CodexBundledReleaseProfile {
|
||||
codex_version: "1.2.4".to_owned(),
|
||||
desktop_version: "1.2.4".to_owned(),
|
||||
desktop_build: "124".to_owned(),
|
||||
verified_at: Utc::now(),
|
||||
});
|
||||
|
||||
let mut continuation = request.clone();
|
||||
continuation.set_previous_response_id(Some("resp_profile_first".to_owned()));
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user