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

feat: add selectable routing groups and composite billing

Browse Source 
Support per-model provider enablement and compact model editing. Capture request-time billing factors, charge customer costs separately, and preserve historical statistics without backfills.
This commit is contained in:
elky
2026-10-07 14:49:57 +08:00
parent 310098a853
commit 911c7f8875
110 changed files with 6524 additions and 559 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/runtime/src/backend/stats/postgres_daily/mod.rs
+6
View File
@@ -159,6 +159,12 @@ async fn perform_stats_aggregation_for_day(
.execute(&mut *tx)
.await?;
sqlx::query(UPDATE_STATS_DAILY_BILLING_COST_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
.execute(&mut *tx)
.await?;
let model_rows =
upsert_stats_daily_model_rows(&mut tx, day_start_utc, day_end_utc, now_utc).await?;
let provider_rows =
crates/aether-data/runtime/src/backend/stats/postgres_daily/sql.rs
+12
View File
@@ -2404,6 +2404,18 @@ WHERE created_at >= $1
AND provider_name NOT IN ('unknown', 'pending')
"#;
// Keep customer consumption separate from the upstream procurement-cost rollup.
pub(super) const UPDATE_STATS_DAILY_BILLING_COST_SQL: &str = r#"
UPDATE stats_daily SET billing_cost=(
SELECT round(COALESCE(sum(billable_amount),0),8)
FROM usage_analytics_facts_v1
WHERE created_at >= $1 AND created_at < $2
AND status NOT IN ('pending','streaming')
AND provider_name NOT IN ('unknown','pending')
)
WHERE date=$1
"#;
#[cfg(test)]
mod tests {
use super::{
crates/aether-data/runtime/src/lifecycle/migrate/tests.rs
+2
View File
@@ -27,6 +27,7 @@ use crate::lifecycle::bootstrap::postgres::{
EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION, EMPTY_DATABASE_SNAPSHOT_SQL,
};
mod customer_billing_upgrade;
mod dashboard_user_anonymization;
mod legacy_overview_upgrade;
mod migration_deadlines;
@@ -1597,6 +1598,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260923000000,
20261001000000,
20261004000000,
20261007000000,
]
);
}
crates/aether-data/runtime/src/lifecycle/migrate/tests/customer_billing_upgrade.rs
+114
View File
@@ -0,0 +1,114 @@
use super::*;
const BILLING_VERSION: i64 = 20261007000000;
#[tokio::test]
async fn customer_billing_upgrade_preserves_history_and_aggregates_new_days() {
let Some(server) = ManagedPostgresServer::try_start().await.unwrap() else {
return;
};
let mut connection = PgConnection::connect(server.database_url()).await.unwrap();
connection.ensure_migrations_table().await.unwrap();
for migration in POSTGRES_MIGRATOR
.iter()
.filter(|migration| migration.version < BILLING_VERSION)
{
connection.apply(migration).await.unwrap();
}
let pool = PgPool::connect(server.database_url()).await.unwrap();
sqlx::raw_sql(
r#"
INSERT INTO stats_daily(id,date,total_requests,actual_total_cost,is_complete)
VALUES ('history','2026-07-17 00:00:00+00',1,0.5,true);
INSERT INTO usage(id,request_id,model,provider_name,status,billing_status,
total_cost_usd,actual_total_cost_usd,created_at,request_metadata)
VALUES ('history','history','m','p','completed','settled',2,0.5,
'2026-07-17 12:00:00+00','{"routing_group_billing_multiplier":2}');
"#,
)
.execute(&pool)
.await
.unwrap();
let history_before: serde_json::Value =
query_scalar("SELECT to_jsonb(d) FROM stats_daily d WHERE id='history'")
.fetch_one(&pool)
.await
.unwrap();
let usage_before: serde_json::Value =
query_scalar("SELECT to_jsonb(u) FROM usage u WHERE request_id='history'")
.fetch_one(&pool)
.await
.unwrap();
let migration = POSTGRES_MIGRATOR
.iter()
.find(|migration| migration.version == BILLING_VERSION)
.unwrap();
connection.apply(migration).await.unwrap();
// Even retained requests with captured factors must not rewrite old daily totals.
let history_after: serde_json::Value =
query_scalar("SELECT to_jsonb(d) - 'billing_cost' FROM stats_daily d WHERE id='history'")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(history_after, history_before);
let usage_after: serde_json::Value =
query_scalar("SELECT to_jsonb(u) FROM usage u WHERE request_id='history'")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(usage_after, usage_before);
let legacy_cost: (Option<String>, String) = sqlx::query_as(
"SELECT billing_cost::text, COALESCE(billing_cost,actual_total_cost::numeric)::text FROM stats_daily WHERE id='history'",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(legacy_cost, (None, "0.50000000".to_string()));
sqlx::raw_sql(
r#"
INSERT INTO usage(id,request_id,model,provider_name,status,billing_status,
total_cost_usd,actual_total_cost_usd,created_at,request_metadata)
VALUES ('new-billed','new-billed','m','p','completed','settled',4,1,
'2026-07-18 12:00:00+00',
'{"billing_multiplier_snapshot":{"version":1,"factors":{"routing_group":2,"user_group":0.75},"multiplier":1.5}}'),
('new-legacy','new-legacy','m','p','completed','settled',2,0.5,
'2026-07-18 13:00:00+00','{}');
"#,
)
.execute(&pool)
.await
.unwrap();
let new_day = historical_stats_day() + chrono::Duration::days(1);
let backend = postgres_backend(server.database_url());
let summary = backend
.aggregate_stats_daily(&crate::StatsDailyAggregationInput {
target_day_utc: new_day,
aggregated_at: new_day + chrono::Duration::days(1),
})
.await
.unwrap()
.unwrap();
assert_eq!(summary.day_start_utc, new_day);
assert_eq!(summary.total_requests, 2);
let new_costs: (String, String) = sqlx::query_as(
"SELECT billing_cost::text, actual_total_cost::text FROM stats_daily WHERE date=$1",
)
.bind(new_day)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(
new_costs,
("6.50000000".to_string(), "1.50000000".to_string())
);
assert!(query_scalar::<_, bool>(
"SELECT billing_cost IS NULL FROM stats_daily WHERE id='history'",
)
.fetch_one(&pool)
.await
.unwrap());
backend.pool().close().await;
pool.close().await;
}
crates/aether-data/runtime/src/lifecycle/migrate/tests/legacy_overview_upgrade.rs
+1
View File
@@ -115,6 +115,7 @@ WHERE version=20260919000000;
20260923000000,
20261001000000,
20261004000000,
20261007000000,
]
);
assert_eq!(
crates/aether-data/runtime/src/lifecycle/migrate/tests/overview_fact_metadata.rs
+1 -1
View File
@@ -168,7 +168,7 @@ VALUES('employee-key','owner',repeat('e',64),false),
let bootstrap =
include_str!("../../../../schema/bootstrap/postgres/190_overview_analytics.sql");
let view_start = bootstrap
.find("CREATE OR REPLACE VIEW public.usage_analytics_facts_v1 AS")
.find("CREATE OR REPLACE FUNCTION public.usage_customer_billable_amount(")
.unwrap();
sqlx::raw_sql(&bootstrap[view_start..])
.execute(&pool)
crates/aether-data/runtime/src/repository/auth/memory.rs
+143 -4
View File
@@ -1024,8 +1024,13 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
export.ip_rules = ip_rules;
}
}
if let Some(feature_settings) = record.feature_settings {
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
if let Some(selection) = record.routing_group_selection {
export.feature_settings = selection.merge_feature_settings(
export.feature_settings.as_ref(),
record.feature_settings,
);
} else if let Some(feature_settings) = record.feature_settings {
export.feature_settings = match feature_settings {
Some(serde_json::Value::Null) | None => None,
Some(value) => Some(value),
@@ -1089,8 +1094,13 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
export.ip_rules = ip_rules;
}
}
if let Some(feature_settings) = record.feature_settings {
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
if let Some(export) = index.export_by_api_key_id.get_mut(&record.api_key_id) {
if let Some(selection) = record.routing_group_selection {
export.feature_settings = selection.merge_feature_settings(
export.feature_settings.as_ref(),
record.feature_settings,
);
} else if let Some(feature_settings) = record.feature_settings {
export.feature_settings = match feature_settings {
Some(serde_json::Value::Null) | None => None,
Some(value) => Some(value),
@@ -2036,6 +2046,7 @@ mod tests {
concurrent_limit_present: false,
ip_rules: None,
feature_settings: Some(Some(serde_json::json!({"must_not_change": true}))),
routing_group_selection: None,
})
.await
.expect("locked basic update should resolve")
@@ -2103,6 +2114,7 @@ mod tests {
concurrent_limit_present: false,
ip_rules: None,
feature_settings: Some(Some(serde_json::json!({"admin": true}))),
routing_group_selection: None,
})
.await
.expect("administrator update should resolve")
@@ -2363,6 +2375,130 @@ mod tests {
);
}
#[tokio::test]
async fn user_key_feature_and_group_updates_merge_against_the_locked_current_record() {
use super::super::UpdateApiKeyRoutingGroupSelection;
fn patch() -> UpdateUserApiKeyBasicRecord {
UpdateUserApiKeyBasicRecord {
user_id: "user-1".into(),
api_key_id: "key-1".into(),
key_encrypted: None,
key_encrypted_present: false,
name: None,
name_present: false,
rate_limit: None,
rate_limit_present: false,
concurrent_limit: None,
concurrent_limit_present: false,
ip_rules: None,
feature_settings: None,
routing_group_selection: Some(UpdateApiKeyRoutingGroupSelection { group_id: None }),
}
}
// Both repository entry points must apply the same merge, with the
// self-service entry point additionally fencing locked keys.
for require_unlocked in [false, true] {
let repository = InMemoryAuthApiKeySnapshotRepository::seed([(
None,
sample_snapshot("key-1", "user-1"),
)]);
repository.set_user_api_key_feature_settings("user-1", "key-1", Some(serde_json::json!({
"routing_group_id": "group-a", "routing_group_name": "stale-name", "pii": {"enabled": false},
}))).await.unwrap().unwrap();
async fn apply(
repository: &InMemoryAuthApiKeySnapshotRepository,
record: UpdateUserApiKeyBasicRecord,
require_unlocked: bool,
) -> StoredAuthApiKeyExportRecord {
if require_unlocked {
repository
.update_user_api_key_basic_if_unlocked(record)
.await
.unwrap()
.unwrap()
} else {
repository
.update_user_api_key_basic(record)
.await
.unwrap()
.unwrap()
}
}
// The PII request is prepared while A is selected, but another
// request selects B before that prepared replacement is committed.
let mut prepared_pii = patch();
prepared_pii.feature_settings = Some(Some(serde_json::json!({
"pii": {"enabled": true}, "routing_group_id": "group-a", "routing_group_name": "injected-name",
})));
let mut select_b = patch();
select_b.routing_group_selection.as_mut().unwrap().group_id =
Some(Some("group-b".into()));
apply(&repository, select_b, require_unlocked).await;
let merged = apply(&repository, prepared_pii, require_unlocked).await;
assert_eq!(
merged.feature_settings,
Some(serde_json::json!({
"pii": {"enabled": true}, "routing_group_id": "group-b",
}))
);
// Conversely a group-only request prepared before a feature change
// must preserve the latest feature object when it reaches storage.
let mut prepared_group = patch();
prepared_group
.routing_group_selection
.as_mut()
.unwrap()
.group_id = Some(Some("group-c".into()));
let mut latest_pii = patch();
latest_pii.feature_settings = Some(Some(
serde_json::json!({ "pii": { "enabled": false, "mode": "strict" } }),
));
apply(&repository, latest_pii, require_unlocked).await;
let merged = apply(&repository, prepared_group, require_unlocked).await;
assert_eq!(
merged.feature_settings,
Some(serde_json::json!({
"pii": {"enabled": false, "mode": "strict"}, "routing_group_id": "group-c",
}))
);
let mut clear_features = patch();
clear_features.feature_settings = Some(None);
let cleared = apply(&repository, clear_features, require_unlocked).await;
assert_eq!(
cleared.feature_settings,
Some(serde_json::json!({"routing_group_id": "group-c"}))
);
let mut clear_group = patch();
clear_group
.routing_group_selection
.as_mut()
.unwrap()
.group_id = Some(None);
assert!(apply(&repository, clear_group, require_unlocked)
.await
.feature_settings
.is_none());
// Administrative callers can still replace the complete document.
let mut admin = patch();
admin.routing_group_selection = None;
admin.feature_settings = Some(Some(
serde_json::json!({ "routing_group_id": "admin-group", "admin": true }),
));
assert_eq!(
apply(&repository, admin, require_unlocked)
.await
.feature_settings,
Some(serde_json::json!({ "routing_group_id": "admin-group", "admin": true }))
);
}
}
#[tokio::test]
async fn update_user_api_key_basic_updates_concurrent_limit() {
let repository = InMemoryAuthApiKeySnapshotRepository::seed(vec![(
@@ -2384,6 +2520,7 @@ mod tests {
concurrent_limit_present: true,
ip_rules: None,
feature_settings: None,
routing_group_selection: None,
})
.await
.expect("update should succeed")
@@ -2419,6 +2556,7 @@ mod tests {
concurrent_limit_present: true,
ip_rules: None,
feature_settings: None,
routing_group_selection: None,
})
.await
.expect("nullable values should clear")
@@ -2441,6 +2579,7 @@ mod tests {
concurrent_limit_present: false,
ip_rules: None,
feature_settings: None,
routing_group_selection: None,
})
.await
.expect("zero rate limit should persist")
crates/aether-data/runtime/src/repository/auth/mod.rs
+2 -2
View File
@@ -6,8 +6,8 @@ pub use aether_data_contracts::repository::auth::{
AuthApiKeyLookupKey, AuthApiKeyReadRepository, AuthApiKeyWriteRepository, AuthRepository,
CompareAndSwapAuthApiKeyCiphertext, CreateStandaloneApiKeyRecord, CreateUserApiKeyRecord,
ResolvedAuthApiKeySnapshot, ResolvedAuthApiKeySnapshotReader, StandaloneApiKeyExportListQuery,
StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot, UpdateStandaloneApiKeyBasicRecord,
UpdateUserApiKeyBasicRecord,
StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot, UpdateApiKeyRoutingGroupSelection,
UpdateStandaloneApiKeyBasicRecord, UpdateUserApiKeyBasicRecord,
};
#[cfg(feature = "postgres")]
pub use aether_data_postgres::SqlxAuthApiKeySnapshotReadRepository;
crates/aether-data/runtime/src/repository/settlement/memory.rs
+120
View File
@@ -1086,6 +1086,116 @@ mod tests {
.expect("wallet should build")
}
fn group_billed_input(request_id: &str) -> UsageSettlementInput {
UsageSettlementInput {
request_id: request_id.to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()),
api_key_is_standalone: false,
provider_id: Some("provider-1".to_string()),
status: "completed".to_string(),
billing_status: "pending".to_string(),
total_cost_usd: 2.0,
actual_total_cost_usd: 0.5,
billing_cost_usd: Some(3.0),
finalized_at_unix_secs: Some(200),
}
}
#[tokio::test]
async fn group_customer_charge_debits_user_wallet_once_without_inflating_provider_cost() {
let repository =
InMemorySettlementRepository::seed(vec![sample_user_wallet("user-wallet", "user-1")]);
let input = group_billed_input("group-billed-user");
let first = repository
.settle_usage(input.clone())
.await
.unwrap()
.unwrap();
assert_eq!(first.wallet_id.as_deref(), Some("user-wallet"));
assert_eq!(first.wallet_balance_before, Some(12.0));
assert_eq!(first.wallet_balance_after, Some(9.0));
assert_eq!(first.provider_monthly_used_usd, Some(0.5));
assert_eq!(repository.settle_usage(input).await.unwrap(), Some(first));
repository.wallets.with_mut(|wallets| {
let wallet = &wallets["user-wallet"];
assert_eq!(wallet.balance, 7.0);
assert_eq!(wallet.gift_balance, 2.0);
assert_eq!(wallet.total_consumed, 3.0);
});
assert_eq!(
repository.provider_monthly_used.read().unwrap()["provider-1"],
0.5
);
}
#[tokio::test]
async fn zero_group_customer_charge_keeps_wallet_unchanged_and_records_provider_cost() {
let repository =
InMemorySettlementRepository::seed(vec![sample_user_wallet("user-wallet", "user-1")]);
let mut input = group_billed_input("group-billed-free");
input.billing_cost_usd = Some(0.0);
let first = repository
.settle_usage(input.clone())
.await
.unwrap()
.unwrap();
assert_eq!(first.billing_status, "settled");
assert_eq!(first.wallet_balance_before, Some(12.0));
assert_eq!(first.wallet_balance_after, Some(12.0));
assert_eq!(first.provider_monthly_used_usd, Some(0.5));
assert_eq!(repository.settle_usage(input).await.unwrap(), Some(first));
repository.wallets.with_mut(|wallets| {
let wallet = &wallets["user-wallet"];
assert_eq!(wallet.balance, 10.0);
assert_eq!(wallet.gift_balance, 2.0);
assert_eq!(wallet.total_consumed, 0.0);
});
assert_eq!(
repository.provider_monthly_used.read().unwrap()["provider-1"],
0.5
);
}
#[tokio::test]
async fn standalone_key_wallet_pays_group_customer_charge_without_debiting_owner() {
let repository = InMemorySettlementRepository::seed(vec![
sample_wallet(),
sample_user_wallet("owner-wallet", "user-1"),
]);
let mut input = group_billed_input("group-billed-standalone");
input.api_key_is_standalone = true;
let settlement = repository.settle_usage(input).await.unwrap().unwrap();
assert_eq!(settlement.wallet_id.as_deref(), Some("wallet-1"));
assert_eq!(settlement.wallet_balance_after, Some(9.0));
assert_eq!(settlement.provider_monthly_used_usd, Some(0.5));
repository.wallets.with_mut(|wallets| {
assert_eq!(wallets["wallet-1"].balance, 7.0);
assert_eq!(wallets["wallet-1"].total_consumed, 3.0);
assert_eq!(wallets["owner-wallet"].balance, 10.0);
assert_eq!(wallets["owner-wallet"].gift_balance, 2.0);
assert_eq!(wallets["owner-wallet"].total_consumed, 0.0);
});
}
#[tokio::test]
async fn invalid_customer_charge_rejects_settlement_before_mutating_financial_state() {
for charge in [-0.01, f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
let repository = InMemorySettlementRepository::seed(vec![sample_wallet()]);
let mut input = group_billed_input("group-billed-invalid");
input.billing_cost_usd = Some(charge);
assert!(repository.settle_usage(input).await.is_err());
repository.wallets.with_mut(|wallets| {
assert_eq!(wallets["wallet-1"].balance, 10.0);
assert_eq!(wallets["wallet-1"].gift_balance, 2.0);
assert_eq!(wallets["wallet-1"].total_consumed, 0.0);
});
assert!(repository.provider_monthly_used.read().unwrap().is_empty());
assert!(repository.settlements.read().unwrap().is_empty());
}
}
#[tokio::test]
async fn settles_usage_against_wallet_and_provider_quota() {
let repository = InMemorySettlementRepository::seed(vec![sample_wallet()]);
@@ -1100,6 +1210,7 @@ mod tests {
billing_status: "pending".to_string(),
total_cost_usd: 3.0,
actual_total_cost_usd: 6.0,
billing_cost_usd: None,
finalized_at_unix_secs: Some(200),
})
.await
@@ -1127,6 +1238,7 @@ mod tests {
billing_status: "pending".to_string(),
total_cost_usd: 3.0,
actual_total_cost_usd: 6.0,
billing_cost_usd: None,
finalized_at_unix_secs: Some(200),
})
.await
@@ -1152,6 +1264,7 @@ mod tests {
billing_status: "pending".to_string(),
total_cost_usd: 3.0,
actual_total_cost_usd: 6.0,
billing_cost_usd: None,
finalized_at_unix_secs: Some(200),
})
.await
@@ -1181,6 +1294,7 @@ mod tests {
billing_status: "pending".to_string(),
total_cost_usd: 3.0,
actual_total_cost_usd: 1.5,
billing_cost_usd: None,
finalized_at_unix_secs: Some(200),
})
.await
@@ -1207,6 +1321,7 @@ mod tests {
billing_status: "pending".to_string(),
total_cost_usd: 3.0,
actual_total_cost_usd: 15.0,
billing_cost_usd: None,
finalized_at_unix_secs: Some(200),
})
.await
@@ -1234,6 +1349,7 @@ mod tests {
billing_status: "pending".to_string(),
total_cost_usd: 3.0,
actual_total_cost_usd: 6.0,
billing_cost_usd: None,
finalized_at_unix_secs: Some(200),
};
@@ -1263,6 +1379,7 @@ mod tests {
billing_status: "pending".to_string(),
total_cost_usd: 3.0,
actual_total_cost_usd: 6.0,
billing_cost_usd: None,
finalized_at_unix_secs: Some(200),
};
let mut tasks = Vec::new();
@@ -1303,6 +1420,7 @@ mod tests {
billing_status: "pending".to_string(),
total_cost_usd: 1.0,
actual_total_cost_usd: 1.0,
billing_cost_usd: None,
finalized_at_unix_secs: Some(200),
})
.await;
@@ -1334,6 +1452,7 @@ mod tests {
billing_status: "pending".to_string(),
total_cost_usd: 2.0,
actual_total_cost_usd: 1.0,
billing_cost_usd: None,
finalized_at_unix_secs: Some(250),
})
.await
@@ -1351,6 +1470,7 @@ mod tests {
billing_status: "settled".to_string(),
total_cost_usd: 2.0,
actual_total_cost_usd: 1.0,
billing_cost_usd: None,
finalized_at_unix_secs: Some(250),
})
.await
crates/aether-data/runtime/src/repository/usage/memory.rs
+19 -7
View File
@@ -4,9 +4,9 @@ use std::sync::RwLock;
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
use aether_data_contracts::repository::usage::{
canonical_usage_body_ref_for, parse_usage_body_ref, sanitize_usage_request_metadata,
usage_body_ref, StoredUsageAuditAggregation, StoredUsageAuditSummary,
StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary,
canonical_usage_body_ref_for, parse_usage_body_ref, preserve_usage_routing_group_snapshot,
sanitize_usage_request_metadata, usage_body_ref, StoredUsageAuditAggregation,
StoredUsageAuditSummary, StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary,
StoredUsageCacheAffinityIntervalRow, StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary,
StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardProviderCount,
StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary,
@@ -22,9 +22,11 @@ use aether_data_contracts::repository::usage::{
UsageDashboardSummaryQuery, UsageErrorDistributionQuery, UsageLeaderboardGroupBy,
UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery,
UsagePerformancePercentilesQuery, UsageProviderPerformanceQuery, UsageSettledCostSummaryQuery,
UsageTimeSeriesGranularity, UsageTimeSeriesQuery, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY,
PROVIDER_REASONING_EFFORT_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY,
REQUESTED_REASONING_EFFORT_METADATA_KEY,
UsageTimeSeriesGranularity, UsageTimeSeriesQuery, BILLING_MULTIPLIER_SNAPSHOT_METADATA_KEY,
PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, PROVIDER_REASONING_EFFORT_METADATA_KEY,
PROVIDER_SERVICE_TIER_METADATA_KEY, REQUESTED_REASONING_EFFORT_METADATA_KEY,
ROUTING_GROUP_BILLING_MULTIPLIER_METADATA_KEY, ROUTING_GROUP_ID_METADATA_KEY,
ROUTING_GROUP_NAME_METADATA_KEY,
};
use async_trait::async_trait;
use chrono::Utc;
@@ -3035,6 +3037,10 @@ fn retain_previous_request_audit_metadata(
"request_path",
"request_query_string",
"request_path_and_query",
ROUTING_GROUP_BILLING_MULTIPLIER_METADATA_KEY,
BILLING_MULTIPLIER_SNAPSHOT_METADATA_KEY,
ROUTING_GROUP_ID_METADATA_KEY,
ROUTING_GROUP_NAME_METADATA_KEY,
] {
if let Some(value) = metadata.get(key) {
retained.insert(key.to_string(), value.clone());
@@ -3213,7 +3219,13 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
.and_then(|existing| existing.request_metadata.clone())
}
});
let request_metadata = sanitize_memory_request_metadata(request_metadata);
let request_metadata =
sanitize_memory_request_metadata(preserve_usage_routing_group_snapshot(
request_metadata,
existing
.as_ref()
.and_then(|stored| stored.request_metadata.as_ref()),
));
let (request_body, request_body_ref, request_body_state) = merge_usage_body_capture(
capture_usage.request_body.take(),
capture_usage.request_body_ref.take(),
crates/aether-data/runtime/src/repository/usage/memory/analytics.rs
+7 -5
View File
@@ -136,14 +136,16 @@ fn apply_allocations(
}
fn decimal_sum(
rows: &[&StoredRequestUsageAudit],
value: impl Fn(&StoredRequestUsageAudit) -> f64,
value: impl Fn(&StoredRequestUsageAudit) -> Option<f64>,
) -> Option<String> {
let amounts = rows
.iter()
.filter(|row| {
available(row, USAGE_PRICING_AVAILABLE_METADATA_KEY) && row.billing_status == "settled"
})
.map(|row| (value(row) * 100_000_000.0).round() as i128)
.filter_map(|row| value(row))
.filter(|amount| amount.is_finite())
.map(|amount| (amount * 100_000_000.0).round() as i128)
.collect::<Vec<_>>();
if amounts.is_empty() {
None
@@ -270,8 +272,8 @@ fn metrics(
metrics.first_byte_p90_ms = first_percentile(0.9);
metrics.first_byte_p99_ms = first_percentile(0.99);
metrics.usage_active_users = users.len() as u64;
metrics.rated_amount = decimal_sum(rows, |row| row.total_cost_usd);
metrics.billable_amount = decimal_sum(rows, |row| row.actual_total_cost_usd);
metrics.rated_amount = decimal_sum(rows, |row| Some(row.total_cost_usd));
metrics.billable_amount = decimal_sum(rows, |row| row.billing_cost());
metrics
}
@@ -281,7 +283,7 @@ fn dashboard_total_metrics(
) -> UsageAnalyticsMetrics {
let mut metrics = UsageAnalyticsMetrics {
request_count: rows.len() as u64,
billable_amount: decimal_sum(rows, |row| row.actual_total_cost_usd),
billable_amount: decimal_sum(rows, |row| row.billing_cost()),
..Default::default()
};
for row in rows {
crates/aether-data/runtime/src/repository/usage/memory/dashboard_summary.rs
+7 -3
View File
@@ -52,8 +52,10 @@ impl DashboardProjection {
!= Some(false)
};
let usage = available(USAGE_AVAILABLE_METADATA_KEY);
let priced =
available(USAGE_PRICING_AVAILABLE_METADATA_KEY) && row.billing_status == "settled";
let billing_cost = row.billing_cost();
let priced = available(USAGE_PRICING_AVAILABLE_METADATA_KEY)
&& row.billing_status == "settled"
&& billing_cost.is_some();
let stream = row
.request_metadata
.as_ref()
@@ -105,7 +107,9 @@ impl DashboardProjection {
actor: analytics::actor(row, keys).map(str::to_owned),
metrics,
billable_units: priced
.then(|| (row.actual_total_cost_usd * 100_000_000.0).round() as i128),
.then_some(billing_cost)
.flatten()
.map(|cost| (cost * 100_000_000.0).round() as i128),
},
);
}
crates/aether-data/runtime/src/repository/usage/memory/tests.rs
+109
View File
@@ -22,6 +22,63 @@ use aether_data_contracts::repository::usage::{
};
use serde_json::json;
#[tokio::test]
async fn customer_billing_statistics_use_frozen_factors_and_preserve_legacy_provider_cost() {
use aether_data_contracts::repository::usage::*;
let now = chrono::Utc::now();
let at = now - chrono::Duration::seconds(10);
let mut billed = sample_usage("customer-billed", at.timestamp());
billed.total_cost_usd = 2.0;
billed.actual_total_cost_usd = 0.5;
billed.request_metadata = Some(json!({
"billing_multiplier_snapshot": {
"version": 1,
"factors": {"routing_group": 2.0, "user_group": 0.75},
"multiplier": 1.5
},
"routing_group_billing_multiplier": 99.0,
"rate_multiplier": 0.25
}));
let mut legacy = sample_usage("customer-legacy", at.timestamp());
legacy.total_cost_usd = 2.0;
legacy.actual_total_cost_usd = 0.5;
let mut free = sample_usage("customer-free", at.timestamp());
free.total_cost_usd = 2.0;
free.actual_total_cost_usd = 0.5;
free.request_metadata = Some(json!({"routing_group_billing_multiplier": 0.0}));
let mut invalid = sample_usage("customer-invalid", at.timestamp());
invalid.total_cost_usd = 999.0;
invalid.actual_total_cost_usd = 999.0;
invalid.request_metadata = Some(json!({"billing_multiplier_snapshot": null}));
let repo = InMemoryUsageReadRepository::seed([billed, legacy, free, invalid])
.with_dashboard_stats_since(at - chrono::Duration::seconds(1));
let overview = repo
.query_usage_analytics(&UsageAnalyticsQuery {
from_unix_ms: (at - chrono::Duration::seconds(1)).timestamp_millis() as u64,
to_unix_ms: now.timestamp_millis() as u64,
timezone: "UTC".into(),
limit: 1,
..Default::default()
})
.await
.unwrap();
assert_eq!(
overview.summary.billable_amount.as_deref(),
Some("3.50000000")
);
let query = UsageDashboardAnalyticsQuery {
timezone: "UTC".into(),
};
let analytics = repo.query_dashboard_analytics(&query).await.unwrap();
assert_eq!(
analytics.total.summary.billable_amount.as_deref(),
Some("3.50000000")
);
let summary = repo.query_dashboard_summary(&query).await.unwrap();
assert_eq!(summary.total.billable_amount.as_deref(), Some("3.50000000"));
assert_eq!(summary.total.pricing_available_count, 3);
}
#[tokio::test]
async fn overview_model_performance_merges_provider_samples_without_pagination() {
use aether_data_contracts::repository::usage::*;
@@ -617,6 +674,58 @@ fn sample_upsert_usage_record(request_id: &str) -> UpsertUsageRecord {
}
}
#[tokio::test]
async fn upsert_preserves_routing_group_snapshot_across_terminal_metadata_replacement() {
for terminal_metadata in [
None,
Some(json!({"rate_multiplier": 0.5, "billing_snapshot": {"status": "complete"}})),
Some(json!({
"routing_group_billing_multiplier": 99.0,
"billing_multiplier_snapshot": {"version": 1, "factors": {"routing_group": 3.0}, "multiplier": 3.0},
"routing_group_id": "changed-group",
"routing_group_name": "changed-group-name",
"plan_usage_reservation_token": "550e8400-e29b-41d4-a716-446655440002",
"rate_multiplier": 0.5
})),
] {
let repository = InMemoryUsageReadRepository::default();
let mut pending = sample_upsert_usage_record("req-group-snapshot");
pending.request_metadata = Some(json!({
"routing_group_billing_multiplier": 0.25,
"billing_multiplier_snapshot": {"version": 1, "factors": {"routing_group": 0.25, "user_group": 2.0}, "multiplier": 0.5},
"routing_group_id": "group-original",
"routing_group_name": "请求时的分组",
"plan_usage_reservation_token": "550e8400-e29b-41d4-a716-446655440001"
}));
repository
.upsert(pending)
.await
.expect("pending usage should persist");
let mut terminal = sample_upsert_usage_record("req-group-snapshot");
terminal.status = "completed".to_string();
terminal.request_metadata = terminal_metadata;
terminal.updated_at_unix_secs += 1;
let stored = repository
.upsert(terminal)
.await
.expect("terminal usage should persist");
assert_eq!(stored.routing_group_billing_multiplier(), 0.25);
assert_eq!(stored.billing_multiplier(), 0.5);
assert_eq!(
stored.request_metadata.as_ref().unwrap()["billing_multiplier_snapshot"],
json!({
"version": 1, "factors": {"routing_group": 0.25, "user_group": 2.0}, "multiplier": 0.5
})
);
assert_eq!(stored.routing_group_id(), Some("group-original"));
assert_eq!(stored.routing_group_name(), Some("请求时的分组"));
assert_eq!(
stored.request_metadata.as_ref().unwrap()["plan_usage_reservation_token"],
"550e8400-e29b-41d4-a716-446655440001"
);
}
}
#[tokio::test]
async fn upsert_preserves_full_http_captures_across_lifecycle_updates() {
let repository = InMemoryUsageReadRepository::default();