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

feat(stats,rules): api_key 用量统计（total_tokens 字段 + 回填）与 body rules 新增 append/insert/regex_replace/name_style 操作

Browse Source 
- 新增迁移 20260422120000_add_api_key_usage_stats.sql，为 api_keys 表添加 total_tokens 列
- 新增回填 20260422120000_backfill_api_key_usage_stats.sql，按历史 usage 重建 api_key 维度汇总
- 在 UsageWriteRepository trait 中添加 rebuild_api_key_usage_stats，补齐 SQL/内存实现及上层调用链
- 内存实现中新增 apply_usage_stats_delta，支持增量更新 api_key 统计快照
- dev.sh：改进临时日志目录管理，并在网关异常退出时输出错误提示
- frontend EndpointFormDialog：将 append 操作从 insert 分支拆分，提供独立 path/value 输入 UI
- rules.rs：扩展 body rules 支持，新增 append/insert/regex_replace/name_style 操作及 WildcardSlice 路径段
This commit is contained in:
fawney19
2026-04-22 20:16:57 +08:00
parent 62153d7d36
commit cf6228f525
15 changed files with 1857 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
crates/aether-data/backfills/20260422120000_backfill_api_key_usage_stats.sql Normal file
View File

@@ -0,0 +1,38 @@
UPDATE public.api_keys
SET
total_requests = 0,
total_tokens = 0,
total_cost_usd = 0,
last_used_at = NULL;
WITH aggregated AS (
SELECT
usage.api_key_id,
COUNT(*)::INTEGER AS total_requests,
COALESCE(
SUM(
GREATEST(
COALESCE(
usage.total_tokens,
COALESCE(usage.input_tokens, 0) + COALESCE(usage.output_tokens, 0)
),
0
)::BIGINT
),
0
)::BIGINT AS total_tokens,
COALESCE(SUM(COALESCE(usage.total_cost_usd, 0)), 0)::NUMERIC(20,8) AS total_cost_usd,
MAX(usage.created_at) AS last_used_at
FROM public.usage
WHERE usage.api_key_id IS NOT NULL
AND BTRIM(usage.api_key_id) <> ''
GROUP BY usage.api_key_id
)
UPDATE public.api_keys
SET
total_requests = aggregated.total_requests,
total_tokens = aggregated.total_tokens,
total_cost_usd = aggregated.total_cost_usd,
last_used_at = aggregated.last_used_at
FROM aggregated
WHERE public.api_keys.id = aggregated.api_key_id;

1
crates/aether-data/bootstrap/20260413020000_baseline_v2.sql
View File

@@ -166,6 +166,7 @@ CREATE TABLE IF NOT EXISTS public.api_keys (
key_encrypted text,
name character varying(100),
total_requests integer DEFAULT 0,
total_tokens bigint DEFAULT '0'::bigint NOT NULL,
total_cost_usd numeric(20,8) DEFAULT '0'::double precision,
is_standalone boolean DEFAULT false NOT NULL,
allowed_providers json,

8
crates/aether-data/migrations/20260422120000_add_api_key_usage_stats.sql Normal file
View File

@@ -0,0 +1,8 @@
ALTER TABLE public.api_keys
ADD COLUMN IF NOT EXISTS total_tokens bigint DEFAULT '0'::bigint NOT NULL;
UPDATE public.provider_api_keys
SET
fingerprint = NULL,
updated_at = NOW()
WHERE fingerprint IS NOT NULL;

75
crates/aether-data/src/backfill.rs
View File

@@ -277,7 +277,7 @@ mod tests {
.into_iter()
.map(|item| item.version)
.collect::<Vec<_>>();
assert_eq!(versions, vec![20260422110000]);
assert_eq!(versions, vec![20260422110000, 20260422120000]);
}
#[test]
@@ -289,7 +289,7 @@ mod tests {
.into_iter()
.map(|item| item.version)
.collect::<Vec<_>>();
assert!(versions.is_empty());
assert_eq!(versions, vec![20260422120000]);
}
#[derive(Debug)]
@@ -506,11 +506,37 @@ mod tests {
.await
.expect("user fixture should insert");
query(
r#"
INSERT INTO public.api_keys (
id,
user_id,
key_hash,
name,
total_requests,
total_tokens,
total_cost_usd
) VALUES (
'api-key-backfill-1',
'user-backfill-1',
'hash-backfill-1',
'Alice CLI',
77,
7777,
77.77
)
"#,
)
.execute(&pool)
.await
.expect("api key fixture should insert");
query(
r#"
INSERT INTO public.usage (
id,
user_id,
api_key_id,
request_id,
provider_name,
model,
@@ -536,6 +562,7 @@ mod tests {
) VALUES (
'usage-backfill-1',
'user-backfill-1',
'api-key-backfill-1',
'req-backfill-1',
'openai',
'gpt-4o-mini',
@@ -568,8 +595,9 @@ mod tests {
let pending_before = pending_backfills(&pool)
.await
.expect("pending backfills should load");
assert_eq!(pending_before.len(), 1);
assert_eq!(pending_before.len(), 2);
assert_eq!(pending_before[0].version, 20260422110000);
assert_eq!(pending_before[1].version, 20260422120000);
run_backfills(&pool)
.await
@@ -585,7 +613,46 @@ mod tests {
.fetch_all(&pool)
.await
.expect("applied backfill versions should load");
assert_eq!(applied_versions, vec![20260422110000]);
assert_eq!(applied_versions, vec![20260422110000, 20260422120000]);
let api_key_total_requests: i64 = query_scalar(
"SELECT COALESCE(total_requests, 0)::BIGINT FROM public.api_keys WHERE id = 'api-key-backfill-1'",
)
.fetch_one(&pool)
.await
.expect("api key total requests should load");
assert_eq!(api_key_total_requests, 1);
let api_key_total_tokens: i64 = query_scalar(
"SELECT COALESCE(total_tokens, 0)::BIGINT FROM public.api_keys WHERE id = 'api-key-backfill-1'",
)
.fetch_one(&pool)
.await
.expect("api key total tokens should load");
assert_eq!(api_key_total_tokens, 150);
let api_key_total_cost: f64 = query_scalar(
"SELECT COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) FROM public.api_keys WHERE id = 'api-key-backfill-1'",
)
.fetch_one(&pool)
.await
.expect("api key total cost should load");
assert_eq!(api_key_total_cost, 1.25);
let api_key_last_used_at_unix_secs: Option<i64> = query_scalar(
"SELECT CAST(EXTRACT(EPOCH FROM last_used_at) AS BIGINT) FROM public.api_keys WHERE id = 'api-key-backfill-1'",
)
.fetch_one(&pool)
.await
.expect("api key last used should load");
assert_eq!(
api_key_last_used_at_unix_secs,
Some(
chrono::DateTime::parse_from_rfc3339("2024-05-06T07:08:09Z")
.expect("expected last used timestamp should parse")
.timestamp(),
)
);
let expected_finalized_unix_secs =
chrono::DateTime::parse_from_rfc3339("2024-05-06T07:18:09Z")

31
crates/aether-data/src/migrate.rs
View File

@@ -8,7 +8,7 @@ use tracing::{error, info, warn};
static MIGRATOR: Migrator = sqlx::migrate!("./migrations");
static BASELINE_V2_SQL: &str = include_str!("../bootstrap/20260413020000_baseline_v2.sql");
const BASELINE_V2_CUTOFF_VERSION: i64 = 20260422110000;
const BASELINE_V2_CUTOFF_VERSION: i64 = 20260422120000;
const MIGRATIONS_TABLE_EXISTS_SQL: &str =
"SELECT to_regclass('public._sqlx_migrations') IS NOT NULL";
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
@@ -586,6 +586,28 @@ mod tests {
.await
}
async fn column_exists(
pool: &PgPool,
table_name: &str,
column_name: &str,
) -> Result<bool, sqlx::Error> {
query_scalar::<_, bool>(
r#"
SELECT EXISTS (
SELECT 1
FROM information_schema.columns
WHERE table_schema = 'public'
AND table_name = $1
AND column_name = $2
)
"#,
)
.bind(table_name)
.bind(column_name)
.fetch_one(pool)
.await
}
#[test]
fn baseline_migration_restores_search_path_for_sqlx_bookkeeping() {
let baseline = MIGRATOR
@@ -639,6 +661,7 @@ mod tests {
20260418000000,
20260421000000,
20260422110000,
20260422120000,
]
);
}
@@ -671,6 +694,8 @@ mod tests {
assert!(BASELINE_V2_SQL.contains("idx_schema_backfills_applied_at"));
assert!(BASELINE_V2_SQL.contains("ALTER TABLE public.stats_hourly"));
assert!(BASELINE_V2_SQL.contains("response_time_sum_ms double precision"));
assert!(BASELINE_V2_SQL.contains("CREATE TABLE IF NOT EXISTS public.api_keys"));
assert!(BASELINE_V2_SQL.contains("total_tokens bigint DEFAULT '0'::bigint NOT NULL"));
assert!(
BASELINE_V2_SQL.contains("CREATE TABLE IF NOT EXISTS public.stats_user_daily_provider")
);
@@ -764,6 +789,7 @@ mod tests {
20260418000000,
20260421000000,
20260422110000,
20260422120000,
]
);
}
@@ -813,6 +839,9 @@ mod tests {
assert!(table_exists(&pool, "usage")
.await
.expect("usage lookup should succeed"));
assert!(column_exists(&pool, "api_keys", "total_tokens")
.await
.expect("api_keys.total_tokens lookup should succeed"));
let applied_count: i64 =
query_scalar("SELECT COUNT(*)::BIGINT FROM public._sqlx_migrations")

65
crates/aether-data/src/repository/auth/memory.rs
View File

@@ -9,6 +9,7 @@ use super::types::{
StandaloneApiKeyExportListQuery, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
UpdateStandaloneApiKeyBasicRecord, UpdateUserApiKeyBasicRecord,
};
use crate::repository::usage::{ApiKeyUsageContribution, ApiKeyUsageDelta};
use crate::DataLayerError;
#[derive(Debug, Default)]
@@ -111,6 +112,70 @@ impl InMemoryAuthApiKeySnapshotRepository {
.copied()
.unwrap_or(0)
}
pub(crate) fn apply_usage_stats_delta(
&self,
api_key_id: &str,
delta: &ApiKeyUsageDelta,
_recomputed_last_used_at_unix_secs: Option<u64>,
) {
let mut index = self
.index
.write()
.expect("auth api key snapshot repository lock");
let Some(record) = index.export_by_api_key_id.get_mut(api_key_id) else {
return;
};
record.total_requests = apply_i64_delta_to_u64(record.total_requests, delta.total_requests);
record.total_tokens = apply_i64_delta_to_u64(record.total_tokens, delta.total_tokens);
record.total_cost_usd = apply_f64_delta(record.total_cost_usd, delta.total_cost_usd);
}
pub(crate) fn rebuild_usage_stats(
&self,
contributions: &BTreeMap<String, ApiKeyUsageContribution>,
) {
let mut index = self
.index
.write()
.expect("auth api key snapshot repository lock");
for record in index.export_by_api_key_id.values_mut() {
record.total_requests = 0;
record.total_tokens = 0;
record.total_cost_usd = 0.0;
}
for (api_key_id, contribution) in contributions {
let Some(record) = index.export_by_api_key_id.get_mut(api_key_id) else {
continue;
};
record.total_requests = clamp_i64_to_u64(contribution.total_requests);
record.total_tokens = clamp_i64_to_u64(contribution.total_tokens);
record.total_cost_usd = contribution.total_cost_usd.max(0.0);
}
}
}
fn clamp_i64_to_u64(value: i64) -> u64 {
u64::try_from(value).unwrap_or_default()
}
fn apply_i64_delta_to_u64(current: u64, delta: i64) -> u64 {
clamp_i64_to_u64(
i64::try_from(current)
.unwrap_or(i64::MAX)
.saturating_add(delta),
)
}
fn apply_f64_delta(current: f64, delta: f64) -> f64 {
let next = current + delta;
if next.is_finite() {
next.max(0.0)
} else {
current.max(0.0)
}
}
#[async_trait]

338
crates/aether-data/src/repository/usage/memory.rs
View File

@@ -25,12 +25,14 @@ use chrono::Utc;
use serde_json::Value;
use super::{
provider_api_key_usage_contribution, strip_deprecated_usage_display_fields,
usage_can_recover_terminal_failure, ProviderApiKeyUsageContribution, ProviderApiKeyUsageDelta,
StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
StoredRequestUsageAudit, StoredUsageDailySummary, UpsertUsageRecord, UsageAuditListQuery,
UsageDailyHeatmapQuery, UsageReadRepository, UsageWriteRepository,
api_key_usage_contribution, provider_api_key_usage_contribution,
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
ApiKeyUsageContribution, ApiKeyUsageDelta, ProviderApiKeyUsageContribution,
ProviderApiKeyUsageDelta, StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary,
StoredProviderUsageWindow, StoredRequestUsageAudit, StoredUsageDailySummary, UpsertUsageRecord,
UsageAuditListQuery, UsageDailyHeatmapQuery, UsageReadRepository, UsageWriteRepository,
};
use crate::repository::auth::InMemoryAuthApiKeySnapshotRepository;
use crate::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use crate::DataLayerError;
@@ -39,6 +41,7 @@ pub struct InMemoryUsageReadRepository {
by_request_id: RwLock<BTreeMap<String, StoredRequestUsageAudit>>,
detached_bodies: RwLock<BTreeMap<String, Value>>,
provider_usage_windows: RwLock<Vec<StoredProviderUsageWindow>>,
auth_api_keys: Option<Arc<InMemoryAuthApiKeySnapshotRepository>>,
provider_catalog: Option<Arc<InMemoryProviderCatalogReadRepository>>,
}
@@ -56,6 +59,7 @@ impl InMemoryUsageReadRepository {
by_request_id: RwLock::new(by_request_id),
detached_bodies: RwLock::new(BTreeMap::new()),
provider_usage_windows: RwLock::new(Vec::new()),
auth_api_keys: None,
provider_catalog: None,
}
}
@@ -107,6 +111,7 @@ impl InMemoryUsageReadRepository {
by_request_id: RwLock::new(by_request_id),
detached_bodies: RwLock::new(detached_bodies),
provider_usage_windows: RwLock::new(Vec::new()),
auth_api_keys: None,
provider_catalog: None,
}
}
@@ -119,10 +124,19 @@ impl InMemoryUsageReadRepository {
by_request_id: self.by_request_id,
detached_bodies: self.detached_bodies,
provider_usage_windows: RwLock::new(items.into_iter().collect()),
auth_api_keys: self.auth_api_keys,
provider_catalog: self.provider_catalog,
}
}
pub fn with_auth_api_key_repository(
mut self,
repository: Arc<InMemoryAuthApiKeySnapshotRepository>,
) -> Self {
self.auth_api_keys = Some(repository);
self
}
pub fn with_provider_catalog_repository(
mut self,
repository: Arc<InMemoryProviderCatalogReadRepository>,
@@ -172,6 +186,31 @@ fn accumulate_provider_api_key_usage_contribution(
};
}
fn accumulate_api_key_usage_contribution(
aggregates: &mut BTreeMap<String, ApiKeyUsageContribution>,
contribution: ApiKeyUsageContribution,
) {
let entry = aggregates
.entry(contribution.api_key_id.clone())
.or_insert_with(|| ApiKeyUsageContribution {
api_key_id: contribution.api_key_id.clone(),
..ApiKeyUsageContribution::default()
});
entry.total_requests = entry
.total_requests
.saturating_add(contribution.total_requests);
entry.total_tokens = entry.total_tokens.saturating_add(contribution.total_tokens);
entry.total_cost_usd += contribution.total_cost_usd;
entry.last_used_at_unix_secs = match (
entry.last_used_at_unix_secs,
contribution.last_used_at_unix_secs,
) {
(Some(existing), Some(candidate)) => Some(existing.max(candidate)),
(None, Some(candidate)) => Some(candidate),
(existing, None) => existing,
};
}
fn usage_matches_list_query(item: &StoredRequestUsageAudit, query: &UsageAuditListQuery) -> bool {
// The field is historically named `created_at_unix_ms`, but usage audit rows
// across gateway handlers, SQL repositories and tests are stored as epoch seconds.
@@ -2407,6 +2446,44 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
};
by_request_id.insert(stored.request_id.clone(), stored.clone());
if let Some(auth_api_keys) = self.auth_api_keys.as_ref() {
let before_contribution = existing.as_ref().and_then(api_key_usage_contribution);
let after_contribution = api_key_usage_contribution(&stored);
match (before_contribution.as_ref(), after_contribution.as_ref()) {
(Some(before), Some(after)) if before.api_key_id == after.api_key_id => {
let delta = ApiKeyUsageDelta::between(before, after);
auth_api_keys.apply_usage_stats_delta(before.api_key_id.as_str(), &delta, None);
}
_ => {
if let Some(before) = before_contribution.as_ref() {
let delta = ApiKeyUsageDelta::removal(before);
let recomputed_last_used_at_unix_secs = by_request_id
.values()
.filter_map(|item| {
item.api_key_id
.as_deref()
.filter(|api_key_id| *api_key_id == before.api_key_id.as_str())
.map(|_| item.created_at_unix_ms)
})
.max();
auth_api_keys.apply_usage_stats_delta(
before.api_key_id.as_str(),
&delta,
recomputed_last_used_at_unix_secs,
);
}
if let Some(after) = after_contribution.as_ref() {
let delta = ApiKeyUsageDelta::addition(after);
auth_api_keys.apply_usage_stats_delta(
after.api_key_id.as_str(),
&delta,
None,
);
}
}
}
}
if let Some(provider_catalog) = self.provider_catalog.as_ref() {
let before_contribution = existing
.as_ref()
@@ -2450,6 +2527,23 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
Ok(stored)
}
async fn rebuild_api_key_usage_stats(&self) -> Result<u64, DataLayerError> {
let Some(auth_api_keys) = self.auth_api_keys.as_ref() else {
return Ok(0);
};
let by_request_id = self.by_request_id.read().expect("usage repository lock");
let mut aggregates = BTreeMap::new();
for usage in by_request_id.values() {
let Some(contribution) = api_key_usage_contribution(usage) else {
continue;
};
accumulate_api_key_usage_contribution(&mut aggregates, contribution);
}
auth_api_keys.rebuild_usage_stats(&aggregates);
Ok(aggregates.len() as u64)
}
async fn rebuild_provider_api_key_usage_stats(&self) -> Result<u64, DataLayerError> {
let Some(provider_catalog) = self.provider_catalog.as_ref() else {
return Ok(0);
@@ -2473,6 +2567,10 @@ mod tests {
use std::sync::Arc;
use super::InMemoryUsageReadRepository;
use crate::repository::auth::{
AuthApiKeyReadRepository, InMemoryAuthApiKeySnapshotRepository,
StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
};
use crate::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository, StoredProviderCatalogKey, StoredProviderCatalogProvider,
@@ -4028,6 +4126,67 @@ mod tests {
))
}
fn sample_auth_api_key_repository(
api_key_ids: &[&str],
) -> Arc<InMemoryAuthApiKeySnapshotRepository> {
let snapshots = api_key_ids.iter().map(|api_key_id| {
(
Some(format!("hash-{api_key_id}")),
StoredAuthApiKeySnapshot::new(
"user-1".to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
None,
None,
None,
(*api_key_id).to_string(),
Some(format!("Key {api_key_id}")),
true,
false,
false,
Some(120),
Some(8),
None,
None,
None,
None,
)
.expect("snapshot should build"),
)
});
let export_records = api_key_ids.iter().map(|api_key_id| {
StoredAuthApiKeyExportRecord::new(
"user-1".to_string(),
(*api_key_id).to_string(),
format!("hash-{api_key_id}"),
Some(format!("enc-{api_key_id}")),
Some(format!("Key {api_key_id}")),
None,
None,
None,
Some(120),
Some(8),
None,
true,
None,
false,
0,
0,
0.0,
false,
)
.expect("export record should build")
});
Arc::new(
InMemoryAuthApiKeySnapshotRepository::seed(snapshots)
.with_export_records(export_records),
)
}
#[tokio::test]
async fn upsert_syncs_linked_provider_key_stats_without_double_counting_request_count() {
let provider_catalog = sample_provider_catalog_repository(&["provider-key-1"]);
@@ -4078,6 +4237,50 @@ mod tests {
assert_eq!(key.last_used_at_unix_secs, Some(1_711_100_000));
}
#[tokio::test]
async fn upsert_syncs_linked_api_key_stats_without_double_counting_request_count() {
let auth_api_keys = sample_auth_api_key_repository(&["api-key-1"]);
let repository = InMemoryUsageReadRepository::default()
.with_auth_api_key_repository(Arc::clone(&auth_api_keys));
repository
.upsert(UpsertUsageRecord {
api_key_id: Some("api-key-1".to_string()),
total_tokens: Some(100),
total_cost_usd: Some(0.5),
created_at_unix_ms: Some(1_711_100_000),
updated_at_unix_secs: 1_711_100_000,
..sample_upsert_usage_record("req-api-key-1")
})
.await
.expect("pending upsert should succeed");
repository
.upsert(UpsertUsageRecord {
api_key_id: Some("api-key-1".to_string()),
status: "completed".to_string(),
billing_status: "settled".to_string(),
total_tokens: Some(180),
total_cost_usd: Some(0.75),
created_at_unix_ms: Some(1_711_100_000),
updated_at_unix_secs: 1_711_100_010,
finalized_at_unix_secs: Some(1_711_100_011),
..sample_upsert_usage_record("req-api-key-1")
})
.await
.expect("completed upsert should succeed");
let key = auth_api_keys
.list_export_api_keys_by_ids(&["api-key-1".to_string()])
.await
.expect("key list should succeed")
.into_iter()
.next()
.expect("api key should exist");
assert_eq!(key.total_requests, 1);
assert_eq!(key.total_tokens, 180);
assert_eq!(key.total_cost_usd, 0.75);
}
#[tokio::test]
async fn upsert_moves_linked_provider_key_stats_when_key_assignment_changes() {
let provider_catalog =
@@ -4146,6 +4349,63 @@ mod tests {
assert_eq!(key_b.last_used_at_unix_secs, Some(1_711_200_000));
}
#[tokio::test]
async fn upsert_moves_linked_api_key_stats_when_key_assignment_changes() {
let auth_api_keys = sample_auth_api_key_repository(&["api-key-a", "api-key-b"]);
let repository = InMemoryUsageReadRepository::default()
.with_auth_api_key_repository(Arc::clone(&auth_api_keys));
repository
.upsert(UpsertUsageRecord {
api_key_id: Some("api-key-a".to_string()),
status: "completed".to_string(),
billing_status: "settled".to_string(),
total_tokens: Some(120),
total_cost_usd: Some(0.4),
created_at_unix_ms: Some(1_711_200_000),
updated_at_unix_secs: 1_711_200_000,
finalized_at_unix_secs: Some(1_711_200_001),
..sample_upsert_usage_record("req-api-move-1")
})
.await
.expect("first upsert should succeed");
repository
.upsert(UpsertUsageRecord {
api_key_id: Some("api-key-b".to_string()),
status: "completed".to_string(),
billing_status: "settled".to_string(),
total_tokens: Some(140),
total_cost_usd: Some(0.6),
created_at_unix_ms: Some(1_711_200_000),
updated_at_unix_secs: 1_711_200_010,
finalized_at_unix_secs: Some(1_711_200_011),
..sample_upsert_usage_record("req-api-move-1")
})
.await
.expect("moved upsert should succeed");
let keys = auth_api_keys
.list_export_api_keys_by_ids(&["api-key-a".to_string(), "api-key-b".to_string()])
.await
.expect("key list should succeed");
let key_a = keys
.iter()
.find(|key| key.api_key_id == "api-key-a")
.expect("key a should exist");
let key_b = keys
.iter()
.find(|key| key.api_key_id == "api-key-b")
.expect("key b should exist");
assert_eq!(key_a.total_requests, 0);
assert_eq!(key_a.total_tokens, 0);
assert_eq!(key_a.total_cost_usd, 0.0);
assert_eq!(key_b.total_requests, 1);
assert_eq!(key_b.total_tokens, 140);
assert_eq!(key_b.total_cost_usd, 0.6);
}
#[tokio::test]
async fn rebuild_provider_key_usage_stats_resets_linked_catalog_to_current_usage() {
let provider_catalog = sample_provider_catalog_repository(&["provider-key-1"]);
@@ -4195,4 +4455,72 @@ mod tests {
assert_eq!(key.total_response_time_ms, Some(840));
assert_eq!(key.last_used_at_unix_secs, Some(1_711_300_250));
}
#[tokio::test]
async fn rebuild_api_key_usage_stats_resets_linked_auth_export_records_to_current_usage() {
let auth_api_keys = sample_auth_api_key_repository(&["api-key-1"]);
let mut stale_key = auth_api_keys
.list_export_api_keys_by_ids(&["api-key-1".to_string()])
.await
.expect("key list should succeed")
.into_iter()
.next()
.expect("api key should exist");
stale_key.total_requests = 99;
stale_key.total_tokens = 9_999;
stale_key.total_cost_usd = 42.0;
let auth_api_keys = Arc::new(
InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-api-key-1".to_string()),
StoredAuthApiKeySnapshot::new(
"user-1".to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
None,
None,
None,
"api-key-1".to_string(),
Some("Key api-key-1".to_string()),
true,
false,
false,
Some(120),
Some(8),
None,
None,
None,
None,
)
.expect("snapshot should build"),
)])
.with_export_records(vec![stale_key]),
);
let repository = InMemoryUsageReadRepository::seed(vec![
sample_usage("req-1", 1_711_300_000),
sample_usage("req-2", 1_711_300_250),
])
.with_auth_api_key_repository(Arc::clone(&auth_api_keys));
let rebuilt = repository
.rebuild_api_key_usage_stats()
.await
.expect("rebuild should succeed");
assert_eq!(rebuilt, 1);
let key = auth_api_keys
.list_export_api_keys_by_ids(&["api-key-1".to_string()])
.await
.expect("key list should succeed")
.into_iter()
.next()
.expect("api key should exist");
assert_eq!(key.total_requests, 2);
assert_eq!(key.total_tokens, 300);
assert_eq!(key.total_cost_usd, 0.24);
}
}

142
crates/aether-data/src/repository/usage/mod.rs
View File

@@ -26,6 +26,67 @@ pub(crate) use aether_data_contracts::repository::usage::{
pub use memory::InMemoryUsageReadRepository;
pub use sql::SqlxUsageReadRepository;
#[derive(Debug, Clone, PartialEq, Default)]
pub(crate) struct ApiKeyUsageContribution {
pub api_key_id: String,
pub total_requests: i64,
pub total_tokens: i64,
pub total_cost_usd: f64,
pub last_used_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Default)]
pub(crate) struct ApiKeyUsageDelta {
pub total_requests: i64,
pub total_tokens: i64,
pub total_cost_usd: f64,
pub candidate_last_used_at_unix_secs: Option<u64>,
pub removed_last_used_at_unix_secs: Option<u64>,
}
impl ApiKeyUsageDelta {
pub(crate) fn between(
before: &ApiKeyUsageContribution,
after: &ApiKeyUsageContribution,
) -> Self {
Self {
total_requests: after.total_requests - before.total_requests,
total_tokens: after.total_tokens - before.total_tokens,
total_cost_usd: after.total_cost_usd - before.total_cost_usd,
candidate_last_used_at_unix_secs: after.last_used_at_unix_secs,
removed_last_used_at_unix_secs: None,
}
}
pub(crate) fn addition(after: &ApiKeyUsageContribution) -> Self {
Self {
total_requests: after.total_requests,
total_tokens: after.total_tokens,
total_cost_usd: after.total_cost_usd,
candidate_last_used_at_unix_secs: after.last_used_at_unix_secs,
removed_last_used_at_unix_secs: None,
}
}
pub(crate) fn removal(before: &ApiKeyUsageContribution) -> Self {
Self {
total_requests: -before.total_requests,
total_tokens: -before.total_tokens,
total_cost_usd: -before.total_cost_usd,
candidate_last_used_at_unix_secs: None,
removed_last_used_at_unix_secs: before.last_used_at_unix_secs,
}
}
pub(crate) fn is_noop(&self) -> bool {
self.total_requests == 0
&& self.total_tokens == 0
&& self.total_cost_usd == 0.0
&& self.candidate_last_used_at_unix_secs.is_none()
&& self.removed_last_used_at_unix_secs.is_none()
}
}
#[derive(Debug, Clone, PartialEq, Default)]
pub(crate) struct ProviderApiKeyUsageContribution {
pub key_id: String,
@@ -200,13 +261,36 @@ pub(crate) fn provider_api_key_usage_contribution(
})
}
pub(crate) fn api_key_usage_contribution(
usage: &StoredRequestUsageAudit,
) -> Option<ApiKeyUsageContribution> {
let api_key_id = usage
.api_key_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())?
.to_string();
Some(ApiKeyUsageContribution {
api_key_id,
total_requests: 1,
total_tokens: i64::try_from(usage.total_tokens).unwrap_or(i64::MAX),
total_cost_usd: if usage.total_cost_usd.is_finite() {
usage.total_cost_usd.max(0.0)
} else {
0.0
},
last_used_at_unix_secs: Some(usage.created_at_unix_ms),
})
}
#[cfg(test)]
mod tests {
use super::{
incoming_usage_can_recover_terminal_failure, provider_api_key_usage_contribution,
provider_api_key_usage_is_error, provider_api_key_usage_is_success,
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
StoredRequestUsageAudit, UpsertUsageRecord,
api_key_usage_contribution, incoming_usage_can_recover_terminal_failure,
provider_api_key_usage_contribution, provider_api_key_usage_is_error,
provider_api_key_usage_is_success, strip_deprecated_usage_display_fields,
usage_can_recover_terminal_failure, StoredRequestUsageAudit, UpsertUsageRecord,
};
#[test]
@@ -426,4 +510,54 @@ mod tests {
assert_eq!(contribution.total_response_time_ms, 120);
assert_eq!(contribution.last_used_at_unix_secs, Some(123));
}
#[test]
fn api_key_usage_contribution_tracks_request_totals() {
let usage = StoredRequestUsageAudit::new(
"usage-1".to_string(),
"request-1".to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
None,
None,
"OpenAI".to_string(),
"gpt-5".to_string(),
None,
Some("provider-1".to_string()),
None,
None,
None,
None,
None,
None,
None,
None,
None,
false,
false,
12,
8,
20,
0.25,
0.25,
Some(200),
None,
None,
Some(120),
None,
"completed".to_string(),
"settled".to_string(),
123,
124,
Some(125),
)
.expect("usage should build");
let contribution = api_key_usage_contribution(&usage).expect("contribution should exist");
assert_eq!(contribution.api_key_id, "api-key-1");
assert_eq!(contribution.total_requests, 1);
assert_eq!(contribution.total_tokens, 20);
assert_eq!(contribution.total_cost_usd, 0.25);
assert_eq!(contribution.last_used_at_unix_secs, Some(123));
}
}

200
crates/aether-data/src/repository/usage/sql.rs
View File

@@ -33,11 +33,11 @@ use std::io::{Read, Write};
use uuid::Uuid;
use super::{
incoming_usage_can_recover_terminal_failure, provider_api_key_usage_contribution,
strip_deprecated_usage_display_fields, ProviderApiKeyUsageDelta,
StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary, StoredRequestUsageAudit,
StoredUsageDailySummary, UpsertUsageRecord, UsageAuditListQuery, UsageDailyHeatmapQuery,
UsageReadRepository, UsageWriteRepository,
api_key_usage_contribution, incoming_usage_can_recover_terminal_failure,
provider_api_key_usage_contribution, strip_deprecated_usage_display_fields, ApiKeyUsageDelta,
ProviderApiKeyUsageDelta, StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary,
StoredRequestUsageAudit, StoredUsageDailySummary, UpsertUsageRecord, UsageAuditListQuery,
UsageDailyHeatmapQuery, UsageReadRepository, UsageWriteRepository,
};
use crate::postgres::PostgresTransactionRunner;
use crate::{
@@ -1454,6 +1454,72 @@ WHERE id = ANY($1::TEXT[])
ORDER BY id ASC
"#;
const APPLY_API_KEY_USAGE_DELTA_SQL: &str = r#"
UPDATE api_keys
SET
total_requests = GREATEST(COALESCE(total_requests, 0) + $2, 0),
total_tokens = GREATEST(COALESCE(total_tokens, 0) + $3, 0),
total_cost_usd = CAST(
GREATEST(CAST(COALESCE(total_cost_usd, 0) AS DOUBLE PRECISION) + $4, 0) AS NUMERIC(20,8)
),
last_used_at = CASE
WHEN $5::double precision IS NOT NULL THEN CASE
WHEN last_used_at IS NULL THEN TO_TIMESTAMP($5::double precision)
ELSE GREATEST(last_used_at, TO_TIMESTAMP($5::double precision))
END
WHEN $6::double precision IS NOT NULL
AND last_used_at IS NOT NULL
AND EXTRACT(EPOCH FROM last_used_at)::BIGINT = $6::BIGINT
THEN (
SELECT MAX(created_at)
FROM "usage"
WHERE api_key_id = $1
)
ELSE last_used_at
END
WHERE id = $1
"#;
const RESET_API_KEY_USAGE_STATS_SQL: &str = r#"
UPDATE api_keys
SET
total_requests = 0,
total_tokens = 0,
total_cost_usd = 0,
last_used_at = NULL
"#;
const REBUILD_API_KEY_USAGE_STATS_SQL: &str = r#"
WITH aggregated AS (
SELECT
api_key_id,
COUNT(*)::INTEGER AS total_requests,
COALESCE(SUM(
GREATEST(
COALESCE(
total_tokens,
COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)
),
0
)::BIGINT
), 0)::BIGINT AS total_tokens,
COALESCE(SUM(COALESCE(total_cost_usd, 0)), 0)::NUMERIC(20,8) AS total_cost_usd,
MAX(created_at) AS last_used_at
FROM "usage"
WHERE api_key_id IS NOT NULL
AND BTRIM(api_key_id) <> ''
GROUP BY api_key_id
)
UPDATE api_keys
SET
total_requests = aggregated.total_requests,
total_tokens = aggregated.total_tokens,
total_cost_usd = aggregated.total_cost_usd,
last_used_at = aggregated.last_used_at
FROM aggregated
WHERE api_keys.id = aggregated.api_key_id
"#;
const APPLY_PROVIDER_API_KEY_USAGE_DELTA_SQL: &str = r#"
UPDATE provider_api_keys
SET
@@ -7660,11 +7726,48 @@ ORDER BY "usage".user_id ASC
stored.output_price_per_1m = settlement_pricing_snapshot.output_price_per_1m;
stored.request_metadata = request_metadata_value;
let before_contribution = previous_usage
let before_api_key_contribution =
previous_usage.as_ref().and_then(api_key_usage_contribution);
let after_api_key_contribution = api_key_usage_contribution(&stored);
match (
before_api_key_contribution.as_ref(),
after_api_key_contribution.as_ref(),
) {
(Some(before), Some(after)) if before.api_key_id == after.api_key_id => {
let delta = ApiKeyUsageDelta::between(before, after);
apply_api_key_usage_delta_in_tx(tx, before.api_key_id.as_str(), &delta)
.await?;
}
_ => {
if let Some(before) = before_api_key_contribution.as_ref() {
let delta = ApiKeyUsageDelta::removal(before);
apply_api_key_usage_delta_in_tx(
tx,
before.api_key_id.as_str(),
&delta,
)
.await?;
}
if let Some(after) = after_api_key_contribution.as_ref() {
let delta = ApiKeyUsageDelta::addition(after);
apply_api_key_usage_delta_in_tx(
tx,
after.api_key_id.as_str(),
&delta,
)
.await?;
}
}
}
let before_provider_contribution = previous_usage
.as_ref()
.and_then(provider_api_key_usage_contribution);
let after_contribution = provider_api_key_usage_contribution(&stored);
match (before_contribution.as_ref(), after_contribution.as_ref()) {
let after_provider_contribution = provider_api_key_usage_contribution(&stored);
match (
before_provider_contribution.as_ref(),
after_provider_contribution.as_ref(),
) {
(Some(before), Some(after)) if before.key_id == after.key_id => {
let delta = ProviderApiKeyUsageDelta::between(before, after);
apply_provider_api_key_usage_delta_in_tx(
@@ -7675,7 +7778,7 @@ ORDER BY "usage".user_id ASC
.await?;
}
_ => {
if let Some(before) = before_contribution.as_ref() {
if let Some(before) = before_provider_contribution.as_ref() {
let delta = ProviderApiKeyUsageDelta::removal(before);
apply_provider_api_key_usage_delta_in_tx(
tx,
@@ -7684,7 +7787,7 @@ ORDER BY "usage".user_id ASC
)
.await?;
}
if let Some(after) = after_contribution.as_ref() {
if let Some(after) = after_provider_contribution.as_ref() {
let delta = ProviderApiKeyUsageDelta::addition(after);
apply_provider_api_key_usage_delta_in_tx(
tx,
@@ -7701,6 +7804,25 @@ ORDER BY "usage".user_id ASC
.await
}
pub async fn rebuild_api_key_usage_stats(&self) -> Result<u64, DataLayerError> {
self.tx_runner
.run_read_write(|tx| {
Box::pin(async move {
sqlx::query(RESET_API_KEY_USAGE_STATS_SQL)
.execute(&mut **tx)
.await
.map_postgres_err()?;
let rows_affected = sqlx::query(REBUILD_API_KEY_USAGE_STATS_SQL)
.execute(&mut **tx)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected)
}) as BoxFuture<'_, Result<u64, DataLayerError>>
})
.await
}
pub async fn rebuild_provider_api_key_usage_stats(&self) -> Result<u64, DataLayerError> {
self.tx_runner
.run_read_write(|tx| {
@@ -7947,6 +8069,10 @@ impl UsageWriteRepository for SqlxUsageReadRepository {
Self::upsert(self, usage).await
}
async fn rebuild_api_key_usage_stats(&self) -> Result<u64, DataLayerError> {
Self::rebuild_api_key_usage_stats(self).await
}
async fn rebuild_provider_api_key_usage_stats(&self) -> Result<u64, DataLayerError> {
Self::rebuild_provider_api_key_usage_stats(self).await
}
@@ -7977,6 +8103,50 @@ async fn lock_usage_request_id_in_tx(
Ok(())
}
async fn apply_api_key_usage_delta_in_tx(
tx: &mut sqlx::Transaction<'_, Postgres>,
api_key_id: &str,
delta: &ApiKeyUsageDelta,
) -> Result<(), DataLayerError> {
if api_key_id.trim().is_empty() {
return Ok(());
}
if delta.is_noop() {
return Ok(());
}
let total_cost_usd_delta = if delta.total_cost_usd.is_finite() {
delta.total_cost_usd
} else {
0.0
};
sqlx::query(APPLY_API_KEY_USAGE_DELTA_SQL)
.bind(api_key_id)
.bind(i32::try_from(delta.total_requests).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"api_keys.total_requests delta exceeds i32: {}",
delta.total_requests
))
})?)
.bind(delta.total_tokens)
.bind(total_cost_usd_delta)
.bind(
delta
.candidate_last_used_at_unix_secs
.map(|value| value as f64),
)
.bind(
delta
.removed_last_used_at_unix_secs
.map(|value| value as f64),
)
.execute(&mut **tx)
.await
.map_postgres_err()?;
Ok(())
}
async fn apply_provider_api_key_usage_delta_in_tx(
tx: &mut sqlx::Transaction<'_, Postgres>,
key_id: &str,
@@ -9466,6 +9636,16 @@ mod tests {
.contains("lock_usage_request_id_in_tx(tx, &usage.request_id).await?;"));
}
#[test]
fn usage_sql_rebuild_matches_online_api_key_usage_semantics() {
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("COUNT(*)::INTEGER"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("COALESCE("));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("total_tokens,"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL
.contains("COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("AND BTRIM(api_key_id) <> ''"));
}
#[test]
fn usage_sql_rebuild_matches_online_provider_key_usage_semantics() {
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL