i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-04 00:17:45 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(codex): restore recovered quota status and stabilize lifecycle tests

Browse Source
This commit is contained in:
ZheFox
2026-09-24 10:33:45 +08:00
parent 57f53903f5
commit 31d5e2d172
14 changed files with 955 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/dispatch/pool_scheduler.rs
+21 -2
View File
@@ -4005,7 +4005,7 @@ mod tests {
async fn pool_key_cursor_reserve_minimum_quota_filters_pages_and_sticky_hits() {
for reserve_enabled in [false, true] {
for sticky in [false, true] {
for used_percent in [99.0, 98.0] {
for used_percent in [99.0, 98.0, 83.0] {
let provider_config = Some(json!({
"pool_advanced": {
"reserve_minimum_quota": reserve_enabled,
@@ -4016,8 +4016,27 @@ mod tests {
sample_codex_pool_provider("provider-pool", 0, provider_config.clone());
let endpoint = sample_codex_pool_endpoint("provider-pool", "endpoint-1");
let mut reserved = sample_codex_pool_key("provider-pool", "key-low");
reserved.status_snapshot = Some(json!({
"quota": {
"provider_type": "codex",
"updated_at": 100,
"allowed": false,
"exhausted": true,
"code": "exhausted",
"windows": [{
"code": "weekly",
"scope": "account",
"used_ratio": 1.0,
"reset_at": 4_102_444_800u64
}]
}
}));
reserved.upstream_metadata = Some(json!({
"codex": {"primary_used_percent": used_percent}
"codex": {
"updated_at": 200,
"primary_used_percent": used_percent,
"primary_reset_at": 4_102_444_800u64
}
}));
let ready = sample_codex_pool_key("provider-pool", "key-ready");
let rows = vec![
apps/aether-gateway/src/handlers/shared/catalog.rs
+135 -10
View File
@@ -21,6 +21,7 @@ use aether_crypto::{
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use aether_provider_pool::{
grok_pool_tier_from_quota_bucket, grok_supported_quota_windows_for_tier,
provider_pool_codex_metadata_has_account_quota,
};
use aether_scheduler_core::provider_key_circuit_payload_is_active_open_at;
use serde_json::{json, Map, Value};
@@ -1111,9 +1112,8 @@ fn build_codex_quota_status_snapshot(
source: &str,
) -> Option<Value> {
let metadata = provider_quota_metadata_bucket(upstream_metadata, "codex")?;
let observed_at_unix_secs = metadata
.get("updated_at")
.and_then(admin_provider_quota_pure::coerce_json_u64);
let observed_at_unix_secs = provider_quota_timestamp_unix_secs(metadata.get("observed_at"))
.or_else(|| provider_quota_timestamp_unix_secs(metadata.get("updated_at")));
let plan_type = metadata
.get("plan_type")
.and_then(Value::as_str)
@@ -2502,17 +2502,19 @@ fn codex_upstream_metadata_is_at_least_as_fresh(
let Some(metadata) = provider_quota_metadata_bucket(upstream_metadata, "codex") else {
return false;
};
let Some(metadata_updated_at) = metadata
.get("updated_at")
.and_then(admin_provider_quota_pure::coerce_json_u64)
// Identity, reset-credit, and model-only updates do not replace the
// account's quota observation, even when their timestamp is newer.
if !provider_pool_codex_metadata_has_account_quota(metadata) {
return false;
}
let Some(metadata_updated_at) = provider_quota_timestamp_unix_secs(metadata.get("observed_at"))
.or_else(|| provider_quota_timestamp_unix_secs(metadata.get("updated_at")))
else {
return false;
};
let snapshot_updated_at = quota_snapshot.and_then(|quota| {
quota
.get("updated_at")
.or_else(|| quota.get("observed_at"))
.and_then(admin_provider_quota_pure::coerce_json_u64)
provider_quota_timestamp_unix_secs(quota.get("observed_at"))
.or_else(|| provider_quota_timestamp_unix_secs(quota.get("updated_at")))
});
snapshot_updated_at.is_none_or(|updated_at| metadata_updated_at >= updated_at)
@@ -2611,6 +2613,19 @@ pub(crate) fn provider_key_status_snapshot_payload(
let mut snapshot = provider_key_status_snapshot_object(Some(&payload))
.or_else(|| default_provider_key_status_snapshot().as_object().cloned())
.unwrap_or_default();
// Legacy snapshots can retain an exhausted summary after a window reset or
// newer quota observation. Use the same decision as scheduling so the
// account list and its status filter do not keep displaying that stale block.
if provider_type.trim().eq_ignore_ascii_case("codex")
&& !aether_provider_pool::provider_pool_key_account_quota_exhausted(key, provider_type)
{
if let Some(quota) = snapshot.get_mut("quota").and_then(Value::as_object_mut) {
quota.insert("exhausted".to_string(), json!(false));
if quota.get("code").and_then(Value::as_str) == Some("exhausted") {
quota.insert("code".to_string(), json!("ok"));
}
}
}
snapshot.insert(
"oauth".to_string(),
build_provider_key_oauth_status_snapshot(key),
@@ -4213,6 +4228,116 @@ mod tests {
);
}
#[test]
fn provider_key_status_snapshot_payload_clears_stale_codex_exhaustion_summary() {
let mut key = sample_catalog_key();
key.status_snapshot = Some(json!({
"quota": {
"provider_type": "codex",
"code": "exhausted",
"exhausted": true,
"windows": [{
"code": "weekly",
"scope": "account",
"used_ratio": 0.83,
"remaining_ratio": 0.17,
"reset_at": 4_102_444_800u64
}]
}
}));
let payload = provider_key_status_snapshot_payload(&key, "codex");
assert_eq!(payload["quota"]["code"], "ok");
assert_eq!(payload["quota"]["exhausted"], false);
assert!(payload["quota"]["label"].is_null());
// An explicit current upstream refusal is not a stale percentage summary.
key.status_snapshot.as_mut().unwrap()["quota"]["allowed"] = json!(false);
let payload = provider_key_status_snapshot_payload(&key, "codex");
assert_eq!(payload["quota"]["code"], "exhausted");
assert_eq!(payload["quota"]["exhausted"], true);
// Missing capacity evidence must not clear an exhausted summary either.
key.status_snapshot = Some(json!({"quota": {
"provider_type": "codex", "code": "exhausted", "exhausted": true
}}));
let payload = provider_key_status_snapshot_payload(&key, "codex");
assert_eq!(payload["quota"]["code"], "exhausted");
assert_eq!(payload["quota"]["exhausted"], true);
}
#[test]
fn provider_key_status_snapshot_payload_refreshes_codex_timestamp_formats() {
for updated_at in [
json!(1_900_000_000u64),
json!(1_900_000_000_000u64),
json!("2030-03-17T17:46:40Z"),
] {
let mut key = sample_catalog_key();
key.upstream_metadata = Some(json!({
"codex": {
"updated_at": updated_at,
"primary_used_percent": 83.0,
"primary_reset_at": 4_102_444_800u64
}
}));
key.status_snapshot = Some(json!({
"quota": {
"provider_type": "codex",
"updated_at": 1_899_999_000u64,
"code": "exhausted",
"exhausted": true,
"allowed": false,
"windows": [{
"code": "weekly",
"scope": "account",
"used_ratio": 1.0,
"remaining_ratio": 0.0,
"reset_at": 4_102_444_800u64
}]
}
}));
let payload = provider_key_status_snapshot_payload(&key, "codex");
assert_eq!(payload["quota"]["code"], "ok");
assert_eq!(payload["quota"]["updated_at"], 1_900_000_000u64);
assert_eq!(payload["quota"]["windows"][0]["used_ratio"], 0.83);
assert!(payload["quota"]["allowed"].is_null());
}
}
#[test]
fn provider_key_status_snapshot_payload_preserves_codex_account_quota_on_unrelated_updates() {
for patch in [
json!({"plan_type": "pro"}),
json!({"spark_primary_used_percent": 83.0}),
json!({"credits_unlimited": false}),
json!({"windows": [{"code": "weekly", "reset_at": 4_102_444_800u64}]}),
] {
let mut metadata = patch;
metadata["updated_at"] = json!("2030-03-17T17:46:40Z");
let mut key = sample_catalog_key();
key.upstream_metadata = Some(json!({"codex": metadata}));
key.status_snapshot = Some(json!({"quota": {
"provider_type": "codex", "updated_at": 200,
"code": "exhausted", "exhausted": true, "allowed": false,
"windows": [{"code": "weekly", "scope": "account", "used_ratio": 1.0,
"remaining_ratio": 0.0, "reset_at": 4_102_444_800u64}]
}}));
let payload = provider_key_status_snapshot_payload(&key, "codex");
assert_eq!(payload["quota"]["code"], "exhausted", "{metadata}");
assert_eq!(payload["quota"]["exhausted"], true, "{metadata}");
assert_eq!(payload["quota"]["allowed"], false, "{metadata}");
assert_eq!(payload["quota"]["updated_at"], 200, "{metadata}");
assert_eq!(
payload["quota"]["windows"][0]["code"], "weekly",
"{metadata}"
);
assert_eq!(
payload["quota"]["windows"][0]["used_ratio"], 1.0,
"{metadata}"
);
}
}
#[test]
fn provider_key_status_snapshot_payload_restores_complete_codex_cache() {
let mut key = sample_catalog_key();
apps/aether-gateway/src/tests/control/admin/pool.rs
+131 -1
View File
@@ -2436,6 +2436,134 @@ async fn gateway_reserve_minimum_quota_marks_and_filters_codex_pool_keys() {
}
}
#[tokio::test]
async fn gateway_reserve_minimum_quota_recovers_after_config_update_with_cached_catalog() {
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
provider.config = Some(json!({
"pool_advanced": {"reserve_minimum_quota": true, "skip_exhausted_accounts": true}
}));
let mut key = sample_key(
"key-codex",
"provider-codex",
"openai:responses",
"oauth-placeholder",
);
key.auth_type = "oauth".to_string();
key.status_snapshot = Some(json!({
"quota": {
"provider_type": "codex", "updated_at": 100,
"code": "exhausted", "exhausted": true, "allowed": false,
"windows": [{
"code": "weekly", "scope": "account", "used_ratio": 1.0,
"remaining_ratio": 0.0, "reset_at": 4_102_444_800u64
}]
}
}));
key.upstream_metadata = Some(json!({
"codex": {"updated_at": 200, "primary_used_percent": 99.0,
"primary_reset_at": 4_102_444_800u64}
}));
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider.clone()],
Vec::new(),
vec![key],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(&repository))
.with_cached_provider_catalog_reader_for_tests(repository),
);
for reserve_enabled in [true, false, true] {
provider.config.as_mut().unwrap()["pool_advanced"]["reserve_minimum_quota"] =
json!(reserve_enabled);
state
.update_provider_catalog_provider(&provider)
.await
.expect("provider should update");
for status in ["all", "quota_exhausted", "available"] {
let response = local_admin_pool_response(
&state,
http::Method::GET,
&format!("/api/admin/pool/provider-codex/keys?status={status}"),
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = serde_json::from_slice(
&to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("json should parse");
let keys = payload["keys"].as_array().expect("keys should be array");
let visible = status == "all" || (status == "quota_exhausted") == reserve_enabled;
assert_eq!(keys.len(), usize::from(visible));
if visible {
assert_eq!(
keys[0]["scheduling_reason"] == "account_quota_exhausted",
reserve_enabled
);
assert_eq!(keys[0]["status_snapshot"]["quota"]["code"], "ok");
}
}
}
}
#[tokio::test]
async fn gateway_codex_pool_filter_clears_stale_exhausted_summary_with_remaining_quota() {
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
provider.config = Some(json!({"pool_advanced": {
"reserve_minimum_quota": false, "skip_exhausted_accounts": true
}}));
let mut key = sample_key(
"key-codex",
"provider-codex",
"openai:responses",
"oauth-placeholder",
);
key.auth_type = "oauth".to_string();
key.status_snapshot = Some(json!({"quota": {
"provider_type": "codex", "code": "exhausted", "exhausted": true,
"windows": [{"code": "weekly", "scope": "account", "used_ratio": 0.83,
"remaining_ratio": 0.17, "reset_at": 4_102_444_800u64}]
}}));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key],
)),
));
for status in ["all", "available", "quota_exhausted"] {
let response = local_admin_pool_response(
&state,
http::Method::GET,
&format!("/api/admin/pool/provider-codex/keys?status={status}"),
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = serde_json::from_slice(
&to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("json should parse");
let keys = payload["keys"].as_array().expect("keys should be array");
assert_eq!(keys.len(), usize::from(status != "quota_exhausted"));
if let Some(key) = keys.first() {
assert_eq!(key["scheduling_status"], "available");
assert_eq!(key["status_snapshot"]["quota"]["code"], "ok");
assert!(key["account_quota"].as_str().unwrap().contains("17.0%"));
}
}
}
#[tokio::test]
async fn gateway_lists_inherited_fixed_provider_api_formats_for_pool_keys() {
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(
@@ -2738,8 +2866,10 @@ async fn gateway_treats_stale_codex_exhausted_snapshot_as_available_when_windows
assert_eq!(keys[0]["scheduling_reason"], json!("available"));
assert_eq!(
keys[0]["account_quota"],
json!("周剩余 100.0% (7天0小时后重置) | 5H剩余 100.0% (5小时0分钟后重置)")
json!("周剩余 100.0% | 5H剩余 100.0%")
);
assert_eq!(keys[0]["status_snapshot"]["quota"]["code"], "ok");
assert_eq!(keys[0]["status_snapshot"]["quota"]["exhausted"], false);
}
#[tokio::test]