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

Fix pool cycle usage display

Browse Source
This commit is contained in:
RWDai
2026-05-07 22:25:52 +08:00
parent 6f620d92be
commit 71742d5ad2
3 changed files with 403 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
apps/aether-gateway/src/handlers/admin/provider/pool_admin/payloads.rs
View File

@@ -9,7 +9,9 @@ use aether_admin::provider::quota as admin_provider_quota_pure;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
};
use aether_data_contracts::repository::usage::StoredProviderApiKeyWindowUsageSummary;
use serde_json::json;
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
fn admin_pool_string_list(value: Option<&serde_json::Value>) -> Option<Vec<String>> {
@@ -411,8 +413,10 @@ fn admin_pool_current_unix_secs() -> u64 {
.unwrap_or(0)
}
fn admin_pool_prune_expired_codex_window_usage(status_snapshot: &mut serde_json::Value) {
let now_unix_secs = admin_pool_current_unix_secs();
fn admin_pool_prune_expired_codex_window_usage_at(
status_snapshot: &mut serde_json::Value,
now_unix_secs: u64,
) {
let Some(windows) = status_snapshot
.get_mut("quota")
.and_then(serde_json::Value::as_object_mut)
@@ -451,6 +455,55 @@ fn admin_pool_prune_expired_codex_window_usage(status_snapshot: &mut serde_json:
}
}
fn admin_pool_prune_expired_codex_window_usage(status_snapshot: &mut serde_json::Value) {
admin_pool_prune_expired_codex_window_usage_at(status_snapshot, admin_pool_current_unix_secs());
}
fn admin_pool_apply_codex_window_usage_summaries(
status_snapshot: &mut serde_json::Value,
usage_by_code: Option<&BTreeMap<String, StoredProviderApiKeyWindowUsageSummary>>,
) {
let Some(usage_by_code) = usage_by_code else {
return;
};
let Some(windows) = status_snapshot
.get_mut("quota")
.and_then(serde_json::Value::as_object_mut)
.and_then(|quota| quota.get_mut("windows"))
.and_then(serde_json::Value::as_array_mut)
else {
return;
};
for window in windows
.iter_mut()
.filter_map(serde_json::Value::as_object_mut)
{
let Some(summary) = window
.get("code")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.map(str::to_ascii_lowercase)
.and_then(|code| usage_by_code.get(&code))
else {
continue;
};
let total_cost_usd = if summary.total_cost_usd.is_finite() {
summary.total_cost_usd.max(0.0)
} else {
0.0
};
window.insert(
"usage".to_string(),
json!({
"request_count": summary.request_count,
"total_tokens": summary.total_tokens,
"total_cost_usd": format!("{total_cost_usd:.8}"),
}),
);
}
}
fn admin_pool_quota_windows<'a>(
quota_snapshot: &'a serde_json::Map<String, serde_json::Value>,
) -> Vec<&'a serde_json::Map<String, serde_json::Value>> {
@@ -853,6 +906,8 @@ pub(super) fn build_admin_pool_key_payload(
key: &StoredProviderCatalogKey,
runtime: &AdminProviderPoolRuntimeState,
pool_config: Option<AdminProviderPoolConfig>,
codex_cycle_usage_by_code: Option<&BTreeMap<String, StoredProviderApiKeyWindowUsageSummary>>,
now_unix_secs: u64,
) -> serde_json::Value {
let cooldown_reason = runtime.cooldown_reason_by_key.get(&key.id).cloned();
let cooldown_ttl_seconds = cooldown_reason
@@ -872,7 +927,11 @@ pub(super) fn build_admin_pool_key_payload(
admin_pool_derive_oauth_plan_type(key, provider_type, auth_config.as_ref());
let mut status_snapshot = provider_key_status_snapshot_payload(key, provider_type);
if provider_type.trim().eq_ignore_ascii_case("codex") {
admin_pool_prune_expired_codex_window_usage(&mut status_snapshot);
admin_pool_apply_codex_window_usage_summaries(
&mut status_snapshot,
codex_cycle_usage_by_code,
);
admin_pool_prune_expired_codex_window_usage_at(&mut status_snapshot, now_unix_secs);
}
let account_snapshot = status_snapshot
.get("account")

139
apps/aether-gateway/src/handlers/admin/provider/pool_admin/read_routes/keys.rs
View File

@@ -11,6 +11,9 @@ use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::GatewayError;
use aether_admin::provider::pool as admin_provider_pool_pure;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use aether_data_contracts::repository::usage::{
ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyWindowUsageSummary,
};
use axum::{
body::Body,
http,
@@ -18,7 +21,131 @@ use axum::{
Json,
};
use serde_json::json;
use std::cmp::Ordering;
use std::{
cmp::Ordering,
collections::BTreeMap,
time::{SystemTime, UNIX_EPOCH},
};
type AdminPoolCodexCycleUsageByKey =
BTreeMap<String, BTreeMap<String, StoredProviderApiKeyWindowUsageSummary>>;
fn admin_pool_json_u64(value: Option<&serde_json::Value>) -> Option<u64> {
match value {
Some(serde_json::Value::Number(number)) => number.as_u64(),
Some(serde_json::Value::String(text)) => text.trim().parse::<u64>().ok(),
_ => None,
}
}
fn admin_pool_codex_default_window_minutes(code: &str) -> Option<u64> {
if code.eq_ignore_ascii_case("5h") {
Some(300)
} else if code.eq_ignore_ascii_case("weekly") {
Some(10_080)
} else {
None
}
}
fn admin_pool_current_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|duration| duration.as_secs())
.unwrap_or(0)
}
fn admin_pool_codex_cycle_usage_request(
key: &StoredProviderCatalogKey,
window: &serde_json::Map<String, serde_json::Value>,
now_unix_secs: u64,
) -> Option<ProviderApiKeyWindowUsageRequest> {
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"))?
.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)?;
if reset_at <= now_unix_secs {
return None;
}
let mut start_unix_secs = reset_at.checked_sub(window_seconds)?;
if let Some(usage_reset_at) = admin_pool_json_u64(window.get("usage_reset_at")) {
start_unix_secs = start_unix_secs.max(usage_reset_at);
}
if start_unix_secs >= reset_at || start_unix_secs >= now_unix_secs {
return None;
}
Some(ProviderApiKeyWindowUsageRequest {
provider_api_key_id: key.id.clone(),
window_code,
start_unix_secs,
end_unix_secs: now_unix_secs,
})
}
fn admin_pool_codex_cycle_usage_requests(
keys: &[StoredProviderCatalogKey],
now_unix_secs: u64,
) -> Vec<ProviderApiKeyWindowUsageRequest> {
keys.iter()
.flat_map(|key| {
key.status_snapshot
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
.and_then(serde_json::Value::as_object)
.and_then(|quota| quota.get("windows"))
.and_then(serde_json::Value::as_array)
.into_iter()
.flatten()
.filter_map(serde_json::Value::as_object)
.filter_map(|window| {
admin_pool_codex_cycle_usage_request(key, window, now_unix_secs)
})
.collect::<Vec<_>>()
})
.collect()
}
async fn read_admin_pool_codex_cycle_usage_by_key(
state: &AdminAppState<'_>,
provider_type: &str,
keys: &[StoredProviderCatalogKey],
now_unix_secs: u64,
) -> Result<AdminPoolCodexCycleUsageByKey, GatewayError> {
if !provider_type.trim().eq_ignore_ascii_case("codex") || keys.is_empty() {
return Ok(BTreeMap::new());
}
let requests = admin_pool_codex_cycle_usage_requests(keys, now_unix_secs);
if requests.is_empty() {
return Ok(BTreeMap::new());
}
let summaries = state
.app()
.summarize_usage_by_provider_api_key_windows(&requests)
.await?;
let mut usage_by_key = AdminPoolCodexCycleUsageByKey::new();
for summary in summaries {
let window_code = summary.window_code.trim().to_ascii_lowercase();
if window_code.is_empty() {
continue;
}
usage_by_key
.entry(summary.provider_api_key_id.clone())
.or_default()
.insert(window_code, summary);
}
Ok(usage_by_key)
}
fn admin_pool_compare_optional_unix_secs(
left: Option<u64>,
@@ -253,6 +380,14 @@ pub(super) async fn build_admin_pool_list_keys_response(
}
_ => AdminProviderPoolRuntimeState::default(),
};
let now_unix_secs = admin_pool_current_unix_secs();
let codex_cycle_usage_by_key = read_admin_pool_codex_cycle_usage_by_key(
state,
&provider.provider_type,
&keys,
now_unix_secs,
)
.await?;
let items = keys
.into_iter()
@@ -264,6 +399,8 @@ pub(super) async fn build_admin_pool_list_keys_response(
&key,
&runtime,
pool_config.clone(),
codex_cycle_usage_by_key.get(&key.id),
now_unix_secs,
)
})
.collect::<Vec<_>>();

203
apps/aether-gateway/src/tests/control/admin/pool.rs
View File

@@ -2,7 +2,9 @@ use std::sync::{Arc, Mutex};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::usage::InMemoryUsageReadRepository;
use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository;
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
use axum::body::{to_bytes, Body, Bytes};
use axum::routing::{any, get, post};
use axum::{extract::Request, Router};
@@ -70,6 +72,56 @@ async fn local_admin_pool_response(
.expect("pool route should resolve locally")
}
fn sample_provider_key_usage_row(
id: &str,
request_id: &str,
provider_id: &str,
provider_api_key_id: &str,
created_at_unix_secs: i64,
total_tokens: i32,
total_cost_usd: f64,
) -> StoredRequestUsageAudit {
StoredRequestUsageAudit::new(
id.to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("user key".to_string()),
"codex".to_string(),
"gpt-5".to_string(),
None,
Some(provider_id.to_string()),
Some("endpoint-1".to_string()),
Some(provider_api_key_id.to_string()),
Some("responses".to_string()),
Some("openai:responses".to_string()),
Some("openai".to_string()),
Some("responses".to_string()),
Some("openai:responses".to_string()),
Some("openai".to_string()),
Some("responses".to_string()),
false,
false,
total_tokens,
0,
total_tokens,
total_cost_usd,
total_cost_usd,
Some(200),
None,
None,
Some(120),
Some(40),
"completed".to_string(),
"settled".to_string(),
created_at_unix_secs,
created_at_unix_secs + 1,
Some(created_at_unix_secs + 2),
)
.expect("usage row should build")
}
#[tokio::test]
async fn gateway_handles_admin_pool_overview_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -1049,6 +1101,157 @@ async fn gateway_pool_list_reads_materialized_codex_cycle_usage_from_quota_windo
.is_none());
}
#[tokio::test]
async fn gateway_pool_list_overrides_stale_codex_cycle_usage_from_usage_facts() {
let now_unix_secs = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("system time should be after unix epoch")
.as_secs();
let reset_at = now_unix_secs + 3_600;
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true
}
})),
);
provider.provider_type = "codex".to_string();
let mut key = sample_key(
"key-codex-stale-cycle",
"provider-codex",
"openai:responses",
"oauth-placeholder",
);
key.name = "codex stale cycle".to_string();
key.auth_type = "oauth".to_string();
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"code": "ok",
"updated_at": reset_at,
"exhausted": false,
"windows": [
{
"code": "weekly",
"label": "周",
"reset_at": reset_at,
"window_minutes": 10_080,
"usage": {
"request_count": 1,
"total_tokens": 50,
"total_cost_usd": "0.05000000"
}
},
{
"code": "5h",
"label": "5H",
"reset_at": reset_at,
"window_minutes": 300,
"usage_reset_at": now_unix_secs.saturating_sub(500),
"usage": {
"request_count": 9,
"total_tokens": 700,
"total_cost_usd": "0.70000000"
}
}
]
}
}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key],
));
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
sample_provider_key_usage_row(
"usage-weekly-only",
"req-weekly-only",
"provider-codex",
"key-codex-stale-cycle",
now_unix_secs.saturating_sub(400_000) as i64,
1_000,
1.25,
),
sample_provider_key_usage_row(
"usage-five-hour-before-manual-reset",
"req-five-hour-before-manual-reset",
"provider-codex",
"key-codex-stale-cycle",
now_unix_secs.saturating_sub(1_000) as i64,
999,
9.99,
),
sample_provider_key_usage_row(
"usage-five-hour",
"req-five-hour",
"provider-codex",
"key-codex-stale-cycle",
now_unix_secs.saturating_sub(100) as i64,
200,
0.75,
),
]));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_and_usage_reader_for_tests(
provider_catalog_repository,
usage_repository,
),
);
let response = local_admin_pool_response(
&state,
http::Method::GET,
"/api/admin/pool/provider-codex/keys?page=1&page_size=50&status=all",
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 body should parse");
let key_payload = payload["keys"]
.as_array()
.expect("keys should be array")
.first()
.expect("key should exist");
let windows = key_payload["status_snapshot"]["quota"]["windows"]
.as_array()
.expect("quota windows should exist");
let weekly = windows
.iter()
.find(|window| window["code"] == json!("weekly"))
.expect("weekly window should exist");
let five_hour = windows
.iter()
.find(|window| window["code"] == json!("5h"))
.expect("5h window should exist");
assert_eq!(weekly["usage"]["request_count"], json!(3));
assert_eq!(weekly["usage"]["total_tokens"], json!(2_199));
assert_eq!(weekly["usage"]["total_cost_usd"], json!("11.99000000"));
assert_eq!(five_hour["usage"]["request_count"], json!(1));
assert_eq!(five_hour["usage"]["total_tokens"], json!(200));
assert_eq!(five_hour["usage"]["total_cost_usd"], json!("0.75000000"));
}
#[tokio::test]
async fn gateway_marks_account_blocked_pool_key_in_list_keys_response() {
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(