mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-04 16:37:46 +08:00
fix(codex): support dynamic quota windows
This commit is contained in:
+14
-4
@@ -106,7 +106,17 @@ fn reset_codex_cycle_usage_windows(status_snapshot: &mut Value, now_unix_secs: u
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
if !code.eq_ignore_ascii_case("5h") && !code.eq_ignore_ascii_case("weekly") {
|
||||
let scope = window
|
||||
.get("scope")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or("account");
|
||||
let has_zero_window = window.get("window_minutes").and_then(Value::as_u64) == Some(0);
|
||||
if code.is_empty()
|
||||
|| !scope.eq_ignore_ascii_case("account")
|
||||
|| code.to_ascii_lowercase().starts_with("spark_")
|
||||
|| has_zero_window
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -177,7 +187,7 @@ mod tests {
|
||||
}
|
||||
});
|
||||
|
||||
assert_eq!(reset_codex_cycle_usage_windows(&mut snapshot, 1_234), 2);
|
||||
assert_eq!(reset_codex_cycle_usage_windows(&mut snapshot, 1_234), 3);
|
||||
let windows = snapshot["quota"]["windows"].as_array().expect("windows");
|
||||
assert_eq!(windows[0]["usage_reset_at"], json!(1_234));
|
||||
assert_eq!(windows[0]["usage"]["request_count"], json!(0));
|
||||
@@ -187,7 +197,7 @@ mod tests {
|
||||
assert_eq!(windows[1]["usage"]["request_count"], json!(0));
|
||||
assert_eq!(windows[1]["usage"]["total_tokens"], json!(0));
|
||||
assert_eq!(windows[1]["usage"]["total_cost_usd"], json!("0.00000000"));
|
||||
assert!(windows[2].get("usage_reset_at").is_none());
|
||||
assert!(windows[2].get("usage").is_some());
|
||||
assert_eq!(windows[2]["usage_reset_at"], json!(1_234));
|
||||
assert_eq!(windows[2]["usage"]["request_count"], json!(0));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -255,12 +255,25 @@ fn admin_pool_quota_window_reset_seconds(
|
||||
|
||||
fn admin_pool_codex_quota_part_from_window(
|
||||
quota_snapshot: &serde_json::Map<String, serde_json::Value>,
|
||||
window_code: &str,
|
||||
label: &str,
|
||||
window: &serde_json::Map<String, serde_json::Value>,
|
||||
now_unix_secs: u64,
|
||||
show_reset_without_consumption: bool,
|
||||
) -> Option<String> {
|
||||
let window = admin_pool_quota_window(quota_snapshot, window_code)?;
|
||||
if admin_pool_json_to_u64(window.get("window_minutes")) == Some(0) {
|
||||
return None;
|
||||
}
|
||||
let label = window
|
||||
.get("label")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|label| !label.is_empty())
|
||||
.or_else(|| {
|
||||
window
|
||||
.get("code")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|code| !code.is_empty())
|
||||
})?;
|
||||
let used_percent = admin_pool_quota_window_used_percent(window)?;
|
||||
let reset_seconds =
|
||||
admin_pool_quota_window_reset_seconds(quota_snapshot, window, now_unix_secs);
|
||||
@@ -294,23 +307,35 @@ fn admin_pool_build_codex_account_quota_from_snapshot(
|
||||
.and_then(admin_provider_quota_pure::coerce_json_bool)
|
||||
.unwrap_or(false);
|
||||
|
||||
if let Some(part) = admin_pool_codex_quota_part_from_window(
|
||||
quota_snapshot,
|
||||
"weekly",
|
||||
"周",
|
||||
now_unix_secs,
|
||||
exhausted,
|
||||
) {
|
||||
parts.push(part);
|
||||
}
|
||||
if let Some(part) = admin_pool_codex_quota_part_from_window(
|
||||
quota_snapshot,
|
||||
"5h",
|
||||
"5H",
|
||||
now_unix_secs,
|
||||
exhausted,
|
||||
) {
|
||||
parts.push(part);
|
||||
if let Some(windows) = quota_snapshot
|
||||
.get("windows")
|
||||
.and_then(serde_json::Value::as_array)
|
||||
{
|
||||
for window in windows.iter().filter_map(serde_json::Value::as_object) {
|
||||
let code = window
|
||||
.get("code")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
let scope = window
|
||||
.get("scope")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or("account");
|
||||
if !scope.eq_ignore_ascii_case("account")
|
||||
|| code.to_ascii_lowercase().starts_with("spark_")
|
||||
{
|
||||
continue;
|
||||
}
|
||||
if let Some(part) = admin_pool_codex_quota_part_from_window(
|
||||
quota_snapshot,
|
||||
window,
|
||||
now_unix_secs,
|
||||
exhausted,
|
||||
) {
|
||||
parts.push(part);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !parts.is_empty() {
|
||||
@@ -345,6 +370,25 @@ fn admin_pool_current_unix_secs() -> u64 {
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn admin_pool_is_regular_codex_cycle_window(
|
||||
window: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> bool {
|
||||
let code = window
|
||||
.get("code")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
let scope = window
|
||||
.get("scope")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or("account");
|
||||
!code.is_empty()
|
||||
&& scope.eq_ignore_ascii_case("account")
|
||||
&& !code.to_ascii_lowercase().starts_with("spark_")
|
||||
&& admin_pool_json_to_u64(window.get("window_minutes")) != Some(0)
|
||||
}
|
||||
|
||||
fn admin_pool_prune_expired_codex_window_usage_at(
|
||||
status_snapshot: &mut serde_json::Value,
|
||||
now_unix_secs: u64,
|
||||
@@ -362,12 +406,7 @@ fn admin_pool_prune_expired_codex_window_usage_at(
|
||||
.iter_mut()
|
||||
.filter_map(serde_json::Value::as_object_mut)
|
||||
{
|
||||
let code = window
|
||||
.get("code")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
if !code.eq_ignore_ascii_case("5h") && !code.eq_ignore_ascii_case("weekly") {
|
||||
if !admin_pool_is_regular_codex_cycle_window(window) {
|
||||
continue;
|
||||
}
|
||||
let Some(reset_at) = admin_pool_json_to_u64(window.get("reset_at")) else {
|
||||
@@ -1405,7 +1444,9 @@ mod tests {
|
||||
"quota": {
|
||||
"windows": [
|
||||
{
|
||||
"code": "5h",
|
||||
"code": "monthly",
|
||||
"scope": "account",
|
||||
"window_minutes": 43_800,
|
||||
"reset_at": 1,
|
||||
"usage": {
|
||||
"request_count": 7,
|
||||
@@ -1451,6 +1492,31 @@ mod tests {
|
||||
assert_eq!(usage["total_cost_usd"], json!("0.60000000"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_monthly_quota_is_rendered_from_actual_snapshot_window() {
|
||||
let quota_snapshot = json!({
|
||||
"provider_type": "codex",
|
||||
"code": "ok",
|
||||
"exhausted": false,
|
||||
"windows": [
|
||||
{
|
||||
"code": "monthly",
|
||||
"label": "月",
|
||||
"scope": "account",
|
||||
"used_ratio": 0.14,
|
||||
"remaining_ratio": 0.86,
|
||||
"window_minutes": 43_800
|
||||
}
|
||||
]
|
||||
});
|
||||
let quota_snapshot = quota_snapshot.as_object().unwrap();
|
||||
|
||||
assert_eq!(
|
||||
admin_pool_build_account_quota("codex", Some(quota_snapshot)),
|
||||
Some("月剩余 86.0%".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn grok_model_quota_is_rendered_for_pool_rows() {
|
||||
let quota_snapshot = json!({
|
||||
|
||||
@@ -50,6 +50,8 @@ fn admin_pool_codex_default_window_minutes(code: &str) -> Option<u64> {
|
||||
Some(300)
|
||||
} else if code.eq_ignore_ascii_case("weekly") {
|
||||
Some(10_080)
|
||||
} else if code.eq_ignore_ascii_case("monthly") {
|
||||
Some(43_800)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
@@ -98,16 +100,27 @@ fn admin_pool_codex_cycle_usage_request(
|
||||
window: &serde_json::Map<String, serde_json::Value>,
|
||||
now_unix_secs: u64,
|
||||
) -> Option<ProviderApiKeyWindowUsageRequest> {
|
||||
let scope = window
|
||||
.get("scope")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or("account");
|
||||
if !scope.eq_ignore_ascii_case("account") {
|
||||
return None;
|
||||
}
|
||||
let window_code = window
|
||||
.get("code")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|code| code.eq_ignore_ascii_case("5h") || code.eq_ignore_ascii_case("weekly"))?
|
||||
.filter(|code| !code.is_empty() && !code.to_ascii_lowercase().starts_with("spark_"))?
|
||||
.to_ascii_lowercase();
|
||||
let reset_at = admin_pool_json_u64(window.get("reset_at"))?;
|
||||
let window_seconds = admin_pool_json_u64(window.get("window_minutes"))
|
||||
.or_else(|| admin_pool_codex_default_window_minutes(&window_code))?
|
||||
.checked_mul(60)?;
|
||||
let window_minutes = match admin_pool_json_u64(window.get("window_minutes")) {
|
||||
Some(0) => return None,
|
||||
Some(value) => value,
|
||||
None => admin_pool_codex_default_window_minutes(&window_code)?,
|
||||
};
|
||||
let window_seconds = window_minutes.checked_mul(60)?;
|
||||
if reset_at <= now_unix_secs {
|
||||
return None;
|
||||
}
|
||||
@@ -788,6 +801,57 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_cycle_usage_request_uses_actual_monthly_window_boundaries() {
|
||||
let key = sample_key("oauth");
|
||||
let reset_at = 5_000_000u64;
|
||||
let now = 3_000_000u64;
|
||||
let window = json!({
|
||||
"code": "monthly",
|
||||
"label": "月",
|
||||
"scope": "account",
|
||||
"reset_at": reset_at,
|
||||
"window_minutes": 43_800u64
|
||||
});
|
||||
|
||||
let request = admin_pool_codex_cycle_usage_request(
|
||||
&key,
|
||||
window.as_object().expect("window should be object"),
|
||||
now,
|
||||
)
|
||||
.expect("monthly usage request should build");
|
||||
|
||||
assert_eq!(request.window_code, "monthly");
|
||||
assert_eq!(request.start_unix_secs, reset_at - 43_800 * 60);
|
||||
assert_eq!(request.end_unix_secs, now);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_cycle_usage_request_ignores_zero_and_spark_windows() {
|
||||
let key = sample_key("oauth");
|
||||
for window in [
|
||||
json!({
|
||||
"code": "weekly",
|
||||
"scope": "account",
|
||||
"reset_at": 5_000_000u64,
|
||||
"window_minutes": 0
|
||||
}),
|
||||
json!({
|
||||
"code": "spark_weekly",
|
||||
"scope": "account",
|
||||
"reset_at": 5_000_000u64,
|
||||
"window_minutes": 10_080
|
||||
}),
|
||||
] {
|
||||
assert!(admin_pool_codex_cycle_usage_request(
|
||||
&key,
|
||||
window.as_object().expect("window should be object"),
|
||||
3_000_000,
|
||||
)
|
||||
.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn oauth_status_filter_prefers_catalog_key_expiry_over_auth_config_expiry() {
|
||||
let mut key = sample_key("oauth");
|
||||
|
||||
@@ -806,6 +806,50 @@ fn codex_default_window_minutes(code: &str) -> Option<u64> {
|
||||
}
|
||||
}
|
||||
|
||||
fn codex_quota_period_identity(window_minutes: u64) -> (String, String) {
|
||||
const MINUTES_PER_HOUR: u64 = 60;
|
||||
const MINUTES_PER_DAY: u64 = 24 * MINUTES_PER_HOUR;
|
||||
const MINUTES_PER_WEEK: u64 = 7 * MINUTES_PER_DAY;
|
||||
|
||||
if window_minutes == 5 * MINUTES_PER_HOUR {
|
||||
return ("5h".to_string(), "5H".to_string());
|
||||
}
|
||||
if window_minutes == MINUTES_PER_WEEK {
|
||||
return ("weekly".to_string(), "周".to_string());
|
||||
}
|
||||
if (28 * MINUTES_PER_DAY..=31 * MINUTES_PER_DAY).contains(&window_minutes) {
|
||||
return ("monthly".to_string(), "月".to_string());
|
||||
}
|
||||
|
||||
let label = if window_minutes % MINUTES_PER_WEEK == 0 {
|
||||
format!("{}周", window_minutes / MINUTES_PER_WEEK)
|
||||
} else if window_minutes % MINUTES_PER_DAY == 0 {
|
||||
format!("{}天", window_minutes / MINUTES_PER_DAY)
|
||||
} else if window_minutes % MINUTES_PER_HOUR == 0 {
|
||||
format!("{}H", window_minutes / MINUTES_PER_HOUR)
|
||||
} else {
|
||||
format!("{window_minutes}分钟")
|
||||
};
|
||||
(format!("window_{window_minutes}m"), label)
|
||||
}
|
||||
|
||||
fn codex_quota_window_identity(
|
||||
fallback_code: &str,
|
||||
fallback_label: &str,
|
||||
window_minutes: Option<u64>,
|
||||
) -> (String, String) {
|
||||
let Some(window_minutes) = window_minutes else {
|
||||
return (fallback_code.to_string(), fallback_label.to_string());
|
||||
};
|
||||
let is_spark = fallback_code.to_ascii_lowercase().starts_with("spark_");
|
||||
let (code, label) = codex_quota_period_identity(window_minutes);
|
||||
if is_spark {
|
||||
(format!("spark_{code}"), format!("Spark {label}"))
|
||||
} else {
|
||||
(code, label)
|
||||
}
|
||||
}
|
||||
|
||||
fn codex_quota_window_snapshot(
|
||||
metadata: &Map<String, Value>,
|
||||
prefix: &str,
|
||||
@@ -847,6 +891,10 @@ fn codex_quota_window_snapshot(
|
||||
.get(&window_minutes_key)
|
||||
.and_then(admin_provider_quota_pure::coerce_json_u64);
|
||||
|
||||
if explicit_window_minutes == Some(0) {
|
||||
return None;
|
||||
}
|
||||
|
||||
if used_percent.is_none()
|
||||
&& reset_at.is_none()
|
||||
&& reset_seconds.is_none()
|
||||
@@ -856,6 +904,7 @@ fn codex_quota_window_snapshot(
|
||||
}
|
||||
|
||||
let window_minutes = explicit_window_minutes.or_else(|| codex_default_window_minutes(code));
|
||||
let (code, label) = codex_quota_window_identity(code, label, window_minutes);
|
||||
let used_ratio = used_percent.map(|value| (value / 100.0).clamp(0.0, 1.0));
|
||||
let remaining_ratio = used_ratio.map(|value| (1.0 - value).max(0.0));
|
||||
|
||||
@@ -931,12 +980,16 @@ fn build_codex_quota_status_snapshot(
|
||||
let primary_windows = windows
|
||||
.iter()
|
||||
.filter(|window| {
|
||||
window
|
||||
let code = window
|
||||
.get("code")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|code| {
|
||||
code.eq_ignore_ascii_case("weekly") || code.eq_ignore_ascii_case("5h")
|
||||
})
|
||||
.unwrap_or_default();
|
||||
let scope = window
|
||||
.get("scope")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
scope.eq_ignore_ascii_case("account")
|
||||
&& !code.to_ascii_lowercase().starts_with("spark_")
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
@@ -3792,6 +3845,39 @@ mod tests {
|
||||
assert_eq!(five_h.get("window_minutes"), Some(&json!(300u64)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_provider_key_quota_status_snapshot_labels_actual_monthly_window() {
|
||||
let upstream_metadata = json!({
|
||||
"codex": {
|
||||
"updated_at": 1_784_287_450u64,
|
||||
"plan_type": "team",
|
||||
"primary_used_percent": 14.0,
|
||||
"primary_reset_at": 1_786_915_122u64,
|
||||
"primary_window_minutes": 43_800u64,
|
||||
"secondary_used_percent": 0.0,
|
||||
"secondary_reset_after_seconds": 0u64,
|
||||
"secondary_window_minutes": 0u64
|
||||
}
|
||||
});
|
||||
|
||||
let payload = sync_provider_key_quota_status_snapshot(
|
||||
None,
|
||||
"codex",
|
||||
Some(&upstream_metadata),
|
||||
"response_headers",
|
||||
)
|
||||
.expect("quota snapshot should sync");
|
||||
let windows = payload["quota"]["windows"]
|
||||
.as_array()
|
||||
.expect("quota windows should exist");
|
||||
|
||||
assert_eq!(windows.len(), 1);
|
||||
assert_eq!(windows[0]["code"], json!("monthly"));
|
||||
assert_eq!(windows[0]["label"], json!("月"));
|
||||
assert_eq!(windows[0]["window_minutes"], json!(43_800u64));
|
||||
assert_eq!(windows[0]["remaining_ratio"], json!(0.86));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_backfills_thin_ok_snapshot_from_upstream_metadata() {
|
||||
let mut key = sample_catalog_key();
|
||||
|
||||
Reference in New Issue
Block a user