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

fix(usage): canonicalize cached token totals

Browse Source
This commit is contained in:
elky
2026-07-16 16:38:30 +08:00
parent e369e4aab1
commit 8fe4d24408
17 changed files with 1804 additions and 1124 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/runtime/src/lifecycle/backfill/tests.rs
+29 -6
View File
@@ -27,7 +27,8 @@ fn pending_backfills_from_applied_returns_all_versions_when_none_applied() {
20260422120000,
20260504120000,
20260505120000,
20260517012000
20260517012000,
20260716010000
]
);
}
@@ -47,7 +48,8 @@ fn pending_backfills_from_applied_skips_versions_already_applied() {
20260422120000,
20260504120000,
20260505120000,
20260517012000
20260517012000,
20260716010000
]
);
}
@@ -184,7 +186,7 @@ impl ManagedPostgresServer {
.arg("-c")
.arg("max_connections=8")
.arg("-c")
.arg("dynamic_shared_memory_type=none")
.arg("dynamic_shared_memory_type=mmap")
.arg("-c")
.arg("autovacuum=off")
.stdout(Stdio::from(stdout))
@@ -417,12 +419,13 @@ async fn run_backfills_rebuilds_stats_and_records_execution() {
let pending_before = pending_backfills(&pool)
.await
.expect("pending backfills should load");
assert_eq!(pending_before.len(), 5);
assert_eq!(pending_before.len(), 6);
assert_eq!(pending_before[0].version, 20260422110000);
assert_eq!(pending_before[1].version, 20260422120000);
assert_eq!(pending_before[2].version, 20260504120000);
assert_eq!(pending_before[3].version, 20260505120000);
assert_eq!(pending_before[4].version, 20260517012000);
assert_eq!(pending_before[5].version, 20260716010000);
run_backfills(&pool)
.await
@@ -445,7 +448,8 @@ async fn run_backfills_rebuilds_stats_and_records_execution() {
20260422120000,
20260504120000,
20260505120000,
20260517012000
20260517012000,
20260716010000
]
);
@@ -841,6 +845,25 @@ async fn run_backfills_rebuilds_stats_and_records_execution() {
.expect("user daily model-provider stats should load");
assert_eq!(user_daily_model_provider_requests, 1);
let canonical_rollup_totals: Vec<i64> = query_scalar(
r#"
SELECT total_tokens FROM public.stats_daily_model_provider
UNION ALL
SELECT total_tokens FROM public.stats_user_daily_model
UNION ALL
SELECT total_tokens FROM public.stats_user_daily_provider
UNION ALL
SELECT total_tokens FROM public.stats_user_daily_api_format
UNION ALL
SELECT total_tokens FROM public.stats_user_daily_model_provider
ORDER BY total_tokens
"#,
)
.fetch_all(&pool)
.await
.expect("canonical persisted rollup totals should load");
assert_eq!(canonical_rollup_totals, vec![150, 150, 150, 150, 150]);
let user_daily_ephemeral_5m_tokens: i64 = query_scalar(
"SELECT COALESCE(SUM(cache_creation_ephemeral_5m_tokens), 0)::BIGINT FROM public.stats_user_daily",
)
@@ -1100,5 +1123,5 @@ async fn run_backfills_rebuilds_stats_and_records_execution() {
.fetch_one(&pool)
.await
.expect("backfill count should load");
assert_eq!(applied_count, 5);
assert_eq!(applied_count, 6);
}
crates/aether-data/runtime/src/lifecycle/bootstrap/postgres.rs
+1 -1
View File
@@ -7,7 +7,7 @@ use tracing::info;
// Generated by build.rs from schema/bootstrap/postgres.
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260715130100;
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260716000000;
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT
crates/aether-data/runtime/src/lifecycle/migrate/tests.rs
+206
View File
@@ -325,6 +325,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260715000000,
20260715130000,
20260715130100,
20260716000000,
]
);
}
@@ -1577,6 +1578,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260715000000,
20260715130000,
20260715130100,
20260716000000,
]
);
}
@@ -2119,6 +2121,210 @@ async fn prepare_database_for_startup_bootstraps_clean_database() {
);
}
#[tokio::test]
async fn postgres_usage_billing_facts_total_tokens_counts_cached_input_once() {
const LEGACY_VIEW_MIGRATION_VERSION: i64 = 20260505130000;
const FIX_MIGRATION_VERSION: i64 = 20260716000000;
let Some(server) = ManagedPostgresServer::try_start()
.await
.expect("postgres billing-facts migration test should start or skip")
else {
return;
};
let pool = PgPool::connect(server.database_url())
.await
.expect("pool should connect");
let pending = prepare_database_for_startup(&pool)
.await
.expect("clean database bootstrap should succeed");
assert!(pending.is_empty());
let legacy_view_migration = POSTGRES_MIGRATOR
.iter()
.find(|migration| migration.version == LEGACY_VIEW_MIGRATION_VERSION)
.expect("legacy billing-facts view migration should be embedded");
assert!(POSTGRES_MIGRATOR
.iter()
.any(|migration| migration.version == FIX_MIGRATION_VERSION));
query(
r#"
INSERT INTO public."usage" (
id,
request_id,
api_key_id,
provider_name,
model,
api_format,
endpoint_api_format,
input_tokens,
output_tokens,
cache_creation_input_tokens,
cache_read_input_tokens,
total_tokens,
status,
billing_status
) VALUES
(
'billing-facts-cache-read',
'billing-facts-cache-read',
'billing-facts-api-key',
'openai',
'gpt-5',
'openai:chat',
'openai:chat',
166103,
94,
0,
164608,
999999,
'completed',
'settled'
),
(
'billing-facts-cache-create',
'billing-facts-cache-create',
'billing-facts-api-key',
'openai',
'gpt-5',
'openai:chat',
'openai:chat',
1000,
20,
100,
400,
999999,
'completed',
'settled'
),
(
'billing-facts-legacy',
'billing-facts-legacy',
'billing-facts-legacy-api-key',
'legacy',
'legacy-model',
NULL,
NULL,
50,
27,
0,
0,
77,
'completed',
'settled'
)
"#,
)
.execute(&pool)
.await
.expect("usage fixtures should be inserted");
query(
r#"
INSERT INTO public.usage_settlement_snapshots (
request_id,
billing_status,
billing_input_tokens,
billing_effective_input_tokens,
billing_output_tokens,
billing_cache_creation_tokens,
billing_cache_read_tokens,
billing_total_input_context
) VALUES
(
'billing-facts-cache-read',
'settled',
166103,
1495,
94,
0,
164608,
166103
),
(
'billing-facts-cache-create',
'settled',
1000,
500,
20,
100,
400,
900
)
"#,
)
.execute(&pool)
.await
.expect("settlement fixtures should be inserted");
sqlx::raw_sql(legacy_view_migration.sql.as_ref())
.execute(&pool)
.await
.expect("legacy billing-facts view should be restored for the upgrade fixture");
let duplicated_total: i64 = query_scalar(
"SELECT total_tokens FROM public.usage_billing_facts WHERE request_id = 'billing-facts-cache-read'",
)
.fetch_one(&pool)
.await
.expect("legacy billing-facts total should be readable");
assert_eq!(duplicated_total, 330805);
query("DELETE FROM public._sqlx_migrations WHERE version = $1")
.bind(FIX_MIGRATION_VERSION)
.execute(&pool)
.await
.expect("billing-facts fix migration stamp should be reset");
super::run_migrations(&pool)
.await
.expect("billing-facts fix migration should rebuild the view");
let cache_read_facts: (i64, i64, i64, i64, i64) = sqlx::query_as(
r#"
SELECT
effective_input_tokens,
cache_creation_input_tokens,
cache_read_input_tokens,
total_input_context,
total_tokens
FROM public.usage_billing_facts
WHERE request_id = 'billing-facts-cache-read'
"#,
)
.fetch_one(&pool)
.await
.expect("canonical cache-read billing facts should be readable");
assert_eq!(cache_read_facts, (1495, 0, 164608, 166103, 166197));
let cache_creation_facts: (i64, i64) = sqlx::query_as(
r#"
SELECT total_input_context, total_tokens
FROM public.usage_billing_facts
WHERE request_id = 'billing-facts-cache-create'
"#,
)
.fetch_one(&pool)
.await
.expect("canonical cache-creation billing facts should be readable");
assert_eq!(cache_creation_facts, (1000, 1020));
let legacy_total: i64 = query_scalar(
"SELECT total_tokens FROM public.usage_billing_facts WHERE request_id = 'billing-facts-legacy'",
)
.fetch_one(&pool)
.await
.expect("legacy billing-facts total should be readable");
assert_eq!(legacy_total, 77);
let repository = aether_data_postgres::SqlxUsageReadRepository::new(pool.clone());
let totals = repository
.summarize_total_tokens_by_api_key_ids(&["billing-facts-api-key".to_string()])
.await
.expect("API-key totals should use canonical billing facts");
assert_eq!(totals.get("billing-facts-api-key"), Some(&167217));
}
#[tokio::test]
async fn postgres_migrations_repair_invalid_concurrent_cleanup_index() {
const MIGRATION_VERSION: i64 = 20260715000000;