i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 17:30: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<_>>();