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(admin): restore pool key usage aggregates

Browse Source
This commit is contained in:
fawney19
2026-04-11 02:54:54 +08:00
parent cf02a10050
commit a54ac76688
9 changed files with 389 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -674,6 +674,26 @@ impl GatewayDataState {
}
}
pub(crate) async fn summarize_usage_by_provider_api_key_ids(
&self,
provider_api_key_ids: &[String],
) -> Result<
std::collections::BTreeMap<
String,
aether_data_contracts::repository::usage::StoredProviderApiKeyUsageSummary,
>,
DataLayerError,
> {
match &self.usage_reader {
Some(repository) => {
repository
.summarize_usage_by_provider_api_key_ids(provider_api_key_ids)
.await
}
None => Ok(std::collections::BTreeMap::new()),
}
}
pub(crate) async fn list_users_by_ids(
&self,
user_ids: &[String],

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

@@ -7,6 +7,7 @@ use super::{
ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL,
};
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::unix_secs_to_rfc3339;
use crate::GatewayError;
use axum::{
body::Body,
@@ -130,17 +131,38 @@ pub(super) async fn build_admin_pool_list_keys_response(
}
_ => AdminProviderPoolRuntimeState::default(),
};
let usage_summary_by_key_id = state
.app()
.summarize_usage_by_provider_api_key_ids(&key_ids)
.await?;
let items = keys
.into_iter()
.map(|key| {
pool_payloads::build_admin_pool_key_payload(
let mut payload = pool_payloads::build_admin_pool_key_payload(
state,
&provider.provider_type,
&key,
&runtime,
pool_config,
)
);
if let Some(summary) = usage_summary_by_key_id.get(&key.id) {
if let Some(object) = payload.as_object_mut() {
object.insert("request_count".to_string(), json!(summary.request_count));
object.insert("total_tokens".to_string(), json!(summary.total_tokens));
object.insert(
"total_cost_usd".to_string(),
json!(format!("{:.8}", summary.total_cost_usd)),
);
object.insert(
"last_used_at".to_string(),
json!(summary
.last_used_at_unix_secs
.and_then(unix_secs_to_rfc3339)),
);
}
}
payload
})
.collect::<Vec<_>>();

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

@@ -65,6 +65,19 @@ impl AppState {
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn summarize_usage_by_provider_api_key_ids(
&self,
provider_api_key_ids: &[String],
) -> Result<
std::collections::BTreeMap<String, usage::StoredProviderApiKeyUsageSummary>,
GatewayError,
> {
self.data
.summarize_usage_by_provider_api_key_ids(provider_api_key_ids)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_users_by_ids(
&self,
user_ids: &[String],

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

@@ -1,8 +1,11 @@
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::{
provider_catalog::InMemoryProviderCatalogReadRepository, 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};
@@ -467,6 +470,155 @@ async fn gateway_pool_list_includes_usage_totals_and_nullable_lru_score() {
assert!(keys[0]["lru_score"].is_null());
}
#[tokio::test]
async fn gateway_pool_list_uses_usage_aggregates_for_stats_and_last_used() {
let provider = sample_provider("provider-openai", "openai", 10).with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true
}
})),
);
let mut key = sample_key(
"key-openai-usage-live",
"provider-openai",
"openai:chat",
"sk-usage-live",
);
key.name = "live usage key".to_string();
key.request_count = Some(0);
key.total_tokens = 0;
key.total_cost_usd = 0.0;
key.last_used_at_unix_secs = None;
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key],
));
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
StoredRequestUsageAudit::new(
"usage-1".to_string(),
"request-1".to_string(),
None,
None,
None,
None,
"OpenAI".to_string(),
"gpt-5.4".to_string(),
None,
Some("provider-openai".to_string()),
None,
Some("key-openai-usage-live".to_string()),
None,
None,
None,
None,
None,
None,
None,
false,
false,
100,
50,
150,
0.5,
0.5,
Some(200),
None,
None,
None,
None,
"completed".to_string(),
"settled".to_string(),
1_711_000_000,
1_711_000_001,
Some(1_711_000_002),
)
.expect("usage row should build"),
StoredRequestUsageAudit::new(
"usage-2".to_string(),
"request-2".to_string(),
None,
None,
None,
None,
"OpenAI".to_string(),
"gpt-5.4".to_string(),
None,
Some("provider-openai".to_string()),
None,
Some("key-openai-usage-live".to_string()),
None,
None,
None,
None,
None,
None,
None,
false,
false,
200,
100,
300,
0.75,
0.75,
Some(200),
None,
None,
None,
None,
"completed".to_string(),
"settled".to_string(),
1_711_000_250,
1_711_000_251,
Some(1_711_000_252),
)
.expect("usage row should build"),
]));
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-openai/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");
assert_eq!(keys.len(), 1);
assert_eq!(keys[0]["request_count"], json!(2));
assert_eq!(keys[0]["total_tokens"], json!(450u64));
assert_eq!(keys[0]["total_cost_usd"], json!("1.25000000"));
assert_eq!(
keys[0]["last_used_at"],
json!(crate::handlers::shared::unix_secs_to_rfc3339(1_711_000_250))
);
}
#[tokio::test]
async fn gateway_handles_admin_pool_list_keys_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));