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

Merge branch 'pr-390' into aether-rust-pioneer

Browse Source
This commit is contained in:
fawney19
2026-05-07 01:34:38 +08:00
parent 47a755a585 79cb9b502a
commit 2a124318b9
19 changed files with 1901 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
apps/aether-gateway/src/data/state/runtime.rs
View File

@@ -30,8 +30,9 @@ use super::{
WalletLookupKey, WalletMutationOutcome,
};
use aether_data_contracts::repository::usage::{
PendingUsageCleanupSummary, StoredUsageDailySummary, UsageAuditListQuery, UsageCleanupSummary,
UsageCleanupWindow, UsageDailyHeatmapQuery,
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest,
StoredProviderApiKeyWindowUsageSummary, StoredUsageDailySummary, UsageAuditListQuery,
UsageCleanupSummary, UsageCleanupWindow, UsageDailyHeatmapQuery,
};
use aether_video_tasks_core::read_data_backed_video_task_response;
@@ -1366,6 +1367,20 @@ impl GatewayDataState {
}
}
pub(crate) async fn summarize_usage_by_provider_api_key_windows(
&self,
requests: &[ProviderApiKeyWindowUsageRequest],
) -> Result<Vec<StoredProviderApiKeyWindowUsageSummary>, DataLayerError> {
match &self.usage_reader {
Some(repository) => {
repository
.summarize_usage_by_provider_api_key_windows(requests)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_users_by_ids(
&self,
user_ids: &[String],

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

@@ -9,7 +9,11 @@ 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::{
ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyWindowUsageSummary,
};
use serde_json::json;
use std::collections::BTreeMap;
fn admin_pool_string_list(value: Option<&serde_json::Value>) -> Option<Vec<String>> {
let values = value
@@ -289,6 +293,124 @@ fn admin_pool_quota_window<'a>(
})
}
pub(super) type AdminPoolCodexWindowUsageByKey =
BTreeMap<(String, String), StoredProviderApiKeyWindowUsageSummary>;
fn admin_pool_provider_type_is_codex(provider_type: &str) -> bool {
provider_type.trim().eq_ignore_ascii_case("codex")
}
fn admin_pool_codex_window_usage_code(
window: &serde_json::Map<String, serde_json::Value>,
) -> Option<&'static str> {
let code = window
.get("code")
.and_then(serde_json::Value::as_str)
.map(str::trim)?;
if code.eq_ignore_ascii_case("5h") {
Some("5h")
} else if code.eq_ignore_ascii_case("weekly") {
Some("weekly")
} else {
None
}
}
fn admin_pool_codex_window_usage_bounds(
window: &serde_json::Map<String, serde_json::Value>,
) -> Option<(u64, u64)> {
let reset_at = admin_pool_json_to_u64(window.get("reset_at"))?;
let window_minutes = admin_pool_json_to_u64(window.get("window_minutes"))?;
let window_seconds = window_minutes.checked_mul(60)?;
let start = reset_at.checked_sub(window_seconds)?;
(start < reset_at).then_some((start, reset_at))
}
pub(super) fn build_admin_pool_codex_window_usage_requests(
provider_type: &str,
keys: &[StoredProviderCatalogKey],
) -> Vec<ProviderApiKeyWindowUsageRequest> {
if !admin_pool_provider_type_is_codex(provider_type) {
return Vec::new();
}
let mut requests = Vec::new();
for key in keys {
let status_snapshot = provider_key_status_snapshot_payload(key, provider_type);
let Some(windows) = status_snapshot
.get("quota")
.and_then(serde_json::Value::as_object)
.and_then(|quota| quota.get("windows"))
.and_then(serde_json::Value::as_array)
else {
continue;
};
for window in windows.iter().filter_map(serde_json::Value::as_object) {
let Some(window_code) = admin_pool_codex_window_usage_code(window) else {
continue;
};
let Some((start_unix_secs, end_unix_secs)) =
admin_pool_codex_window_usage_bounds(window)
else {
continue;
};
requests.push(ProviderApiKeyWindowUsageRequest {
provider_api_key_id: key.id.clone(),
window_code: window_code.to_string(),
start_unix_secs,
end_unix_secs,
});
}
}
requests
}
fn admin_pool_codex_window_usage_payload(
usage: &StoredProviderApiKeyWindowUsageSummary,
) -> serde_json::Value {
json!({
"request_count": usage.request_count,
"total_tokens": usage.total_tokens,
"total_cost_usd": format!("{:.8}", usage.total_cost_usd),
})
}
fn admin_pool_attach_codex_window_usage(
status_snapshot: &mut serde_json::Value,
key_id: &str,
usage_by_key: &AdminPoolCodexWindowUsageByKey,
) {
if usage_by_key.is_empty() {
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(window_code) = admin_pool_codex_window_usage_code(window) else {
continue;
};
let lookup_key = (key_id.to_string(), window_code.to_string());
if let Some(usage) = usage_by_key.get(&lookup_key) {
window.insert(
"usage".to_string(),
admin_pool_codex_window_usage_payload(usage),
);
}
}
}
fn admin_pool_quota_window_used_percent(
window: &serde_json::Map<String, serde_json::Value>,
) -> Option<f64> {
@@ -804,6 +926,7 @@ pub(super) fn build_admin_pool_key_payload(
key: &StoredProviderCatalogKey,
runtime: &AdminProviderPoolRuntimeState,
pool_config: Option<AdminProviderPoolConfig>,
codex_window_usage_by_key: &AdminPoolCodexWindowUsageByKey,
) -> serde_json::Value {
let cooldown_reason = runtime.cooldown_reason_by_key.get(&key.id).cloned();
let cooldown_ttl_seconds = cooldown_reason
@@ -821,7 +944,14 @@ pub(super) fn build_admin_pool_key_payload(
admin_pool_derive_oauth_expires_at(provider_type, key, auth_config.as_ref());
let oauth_plan_type =
admin_pool_derive_oauth_plan_type(key, provider_type, auth_config.as_ref());
let status_snapshot = provider_key_status_snapshot_payload(key, provider_type);
let mut status_snapshot = provider_key_status_snapshot_payload(key, provider_type);
if admin_pool_provider_type_is_codex(provider_type) {
admin_pool_attach_codex_window_usage(
&mut status_snapshot,
&key.id,
codex_window_usage_by_key,
);
}
let account_snapshot = status_snapshot
.get("account")
.and_then(serde_json::Value::as_object);

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

@@ -253,6 +253,26 @@ pub(super) async fn build_admin_pool_list_keys_response(
}
_ => AdminProviderPoolRuntimeState::default(),
};
let codex_window_usage_requests =
pool_payloads::build_admin_pool_codex_window_usage_requests(&provider.provider_type, &keys);
let codex_window_usage_by_key: pool_payloads::AdminPoolCodexWindowUsageByKey =
if codex_window_usage_requests.is_empty() {
pool_payloads::AdminPoolCodexWindowUsageByKey::new()
} else {
state
.app()
.summarize_usage_by_provider_api_key_windows(&codex_window_usage_requests)
.await?
.into_iter()
.map(|usage| {
(
(usage.provider_api_key_id.clone(), usage.window_code.clone()),
usage,
)
})
.collect()
};
let items = keys
.into_iter()
.map(|key| {
@@ -263,6 +283,7 @@ pub(super) async fn build_admin_pool_list_keys_response(
&key,
&runtime,
pool_config.clone(),
&codex_window_usage_by_key,
)
})
.collect::<Vec<_>>();

10
apps/aether-gateway/src/state/runtime/usage_queries.rs
View File

@@ -345,6 +345,16 @@ impl AppState {
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn summarize_usage_by_provider_api_key_windows(
&self,
requests: &[usage::ProviderApiKeyWindowUsageRequest],
) -> Result<Vec<usage::StoredProviderApiKeyWindowUsageSummary>, GatewayError> {
self.data
.summarize_usage_by_provider_api_key_windows(requests)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_users_by_ids(
&self,
user_ids: &[String],

277
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};
@@ -21,6 +23,54 @@ use crate::constants::{
use crate::control::resolve_public_request_context;
use crate::data::GatewayDataState;
fn sample_pool_usage_row(
request_id: &str,
provider_api_key_id: &str,
created_at_unix_secs: i64,
total_tokens: i32,
total_cost_usd: f64,
) -> StoredRequestUsageAudit {
StoredRequestUsageAudit::new(
format!("usage-{request_id}"),
request_id.to_string(),
Some("user-codex".to_string()),
Some("api-key-codex".to_string()),
Some("codex-user".to_string()),
Some("codex-api-key".to_string()),
"codex".to_string(),
"gpt-5-codex".to_string(),
None,
Some("provider-codex".to_string()),
Some("endpoint-codex".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(240),
Some(80),
"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")
}
fn trusted_admin_headers() -> HeaderMap {
let mut headers = HeaderMap::new();
headers.insert(GATEWAY_HEADER, HeaderValue::from_static("rust-phase3b"));
@@ -793,6 +843,233 @@ async fn gateway_sorts_admin_pool_keys_by_imported_and_last_used_time() {
assert_eq!(last_used_names, vec!["active", "old", "fresh"]);
}
#[tokio::test]
async fn gateway_pool_list_adds_codex_cycle_usage_to_quota_windows() {
const RESET_AT: u64 = 1_711_000_000;
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 usage_key = sample_key(
"key-codex-cycle",
"provider-codex",
"openai:responses",
"oauth-placeholder",
);
usage_key.name = "codex cycle usage".to_string();
usage_key.auth_type = "oauth".to_string();
usage_key.request_count = Some(4);
usage_key.total_tokens = 999;
usage_key.total_cost_usd = 9.99;
usage_key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"code": "ok",
"label": serde_json::Value::Null,
"reason": serde_json::Value::Null,
"freshness": "fresh",
"source": "response_headers",
"observed_at": RESET_AT,
"exhausted": false,
"usage_ratio": 0.0,
"updated_at": RESET_AT,
"reset_seconds": serde_json::Value::Null,
"plan_type": "plus",
"windows": [
{
"code": "weekly",
"label": "周",
"scope": "account",
"unit": "percent",
"used_ratio": 0.0,
"remaining_ratio": 1.0,
"reset_at": RESET_AT,
"reset_seconds": 604_800,
"window_minutes": 10_080
},
{
"code": "5h",
"label": "5H",
"scope": "account",
"unit": "percent",
"used_ratio": 0.0,
"remaining_ratio": 1.0,
"reset_at": RESET_AT,
"reset_seconds": 18_000,
"window_minutes": 300
}
]
}
}));
let mut zero_key = sample_key(
"key-codex-zero",
"provider-codex",
"openai:responses",
"oauth-placeholder",
);
zero_key.name = "codex zero usage".to_string();
zero_key.auth_type = "oauth".to_string();
zero_key.status_snapshot = usage_key.status_snapshot.clone();
let mut invalid_key = sample_key(
"key-codex-invalid",
"provider-codex",
"openai:responses",
"oauth-placeholder",
);
invalid_key.name = "codex invalid window".to_string();
invalid_key.auth_type = "oauth".to_string();
invalid_key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"code": "ok",
"windows": [
{
"code": "weekly",
"label": "周",
"reset_at": serde_json::Value::Null,
"window_minutes": 10_080
},
{
"code": "5h",
"label": "5H",
"reset_at": RESET_AT,
"window_minutes": serde_json::Value::Null
}
]
}
}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![usage_key, zero_key, invalid_key],
));
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
sample_pool_usage_row(
"req-5h-a",
"key-codex-cycle",
RESET_AT as i64 - 60,
100,
0.10,
),
sample_pool_usage_row(
"req-5h-b",
"key-codex-cycle",
RESET_AT as i64 - 17_999,
125,
0.20,
),
sample_pool_usage_row(
"req-weekly-only",
"key-codex-cycle",
RESET_AT as i64 - 18_001,
150,
0.30,
),
sample_pool_usage_row(
"req-before-weekly",
"key-codex-cycle",
RESET_AT as i64 - 604_801,
200,
0.40,
),
]));
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 keys = payload["keys"].as_array().expect("keys should be array");
fn key_by_id<'a>(keys: &'a [serde_json::Value], key_id: &str) -> &'a serde_json::Value {
keys.iter()
.find(|key| key["key_id"] == json!(key_id))
.expect("key payload should exist")
}
fn window_by_code<'a>(key_payload: &'a serde_json::Value, code: &str) -> &'a serde_json::Value {
key_payload["status_snapshot"]["quota"]["windows"]
.as_array()
.expect("quota windows should be array")
.iter()
.find(|window| window["code"] == json!(code))
.expect("quota window should exist")
}
let usage_key_payload = key_by_id(keys, "key-codex-cycle");
let five_hour_window = window_by_code(usage_key_payload, "5h");
let weekly_window = window_by_code(usage_key_payload, "weekly");
assert_eq!(five_hour_window["usage"]["request_count"], json!(2));
assert_eq!(five_hour_window["usage"]["total_tokens"], json!(225));
assert_eq!(
five_hour_window["usage"]["total_cost_usd"],
json!("0.30000000")
);
assert_eq!(weekly_window["usage"]["request_count"], json!(3));
assert_eq!(weekly_window["usage"]["total_tokens"], json!(375));
assert_eq!(
weekly_window["usage"]["total_cost_usd"],
json!("0.60000000")
);
assert_eq!(usage_key_payload["request_count"], json!(4));
assert_eq!(usage_key_payload["total_tokens"], json!(999));
assert_eq!(usage_key_payload["total_cost_usd"], json!("9.99000000"));
let zero_key_payload = key_by_id(keys, "key-codex-zero");
assert_eq!(
window_by_code(zero_key_payload, "5h")["usage"]["request_count"],
json!(0)
);
assert_eq!(
window_by_code(zero_key_payload, "weekly")["usage"]["total_tokens"],
json!(0)
);
let invalid_key_payload = key_by_id(keys, "key-codex-invalid");
assert!(window_by_code(invalid_key_payload, "5h")
.get("usage")
.is_none());
assert!(window_by_code(invalid_key_payload, "weekly")
.get("usage")
.is_none());
}
#[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(