i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-09 02:47:45 +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-admin/src/observability/usage.rs
+106 -1
View File
@@ -1295,6 +1295,8 @@ fn admin_usage_active_request_json(
let cache_creation_input_tokens = admin_usage_cache_creation_tokens(item);
let client_is_stream = admin_usage_client_is_stream(item);
let upstream_is_stream = admin_usage_upstream_is_stream(item);
let billing_multiplier = item.billing_multiplier();
let billing_cost = item.billing_cost().map(|cost| round_to(cost, 6));
let mut value = json!({
"id": item.id,
"status": item.status,
@@ -1334,6 +1336,11 @@ fn admin_usage_active_request_json(
"request_path_and_query": admin_usage_metadata_string(item, "request_path_and_query"),
"has_fallback": admin_usage_has_fallback(item),
});
value["billing_multiplier"] = json!(billing_multiplier);
value["billing_cost"] = json!(billing_cost);
value["routing_group_id"] = json!(item.routing_group_id());
value["routing_group_name"] = json!(item.routing_group_name());
value["rate_multiplier"] = json!(item.settlement_rate_multiplier());
value["end_to_end_time_ms"] = json!(admin_usage_metadata_u64(item, "end_to_end_time_ms"));
value["end_to_end_first_byte_time_ms"] = json!(admin_usage_metadata_u64(
item,
@@ -1395,6 +1402,8 @@ pub fn admin_usage_record_json(
.unwrap_or_else(|| "已删除用户".to_string());
let client_is_stream = admin_usage_client_is_stream(item);
let upstream_is_stream = admin_usage_upstream_is_stream(item);
let billing_multiplier = item.billing_multiplier();
let billing_cost = item.billing_cost().map(|cost| round_to(cost, 6));
let mut payload = json!({
"id": item.id,
@@ -1443,6 +1452,10 @@ pub fn admin_usage_record_json(
"provider_key_name": provider_key_name,
"model_version": Value::Null,
});
payload["billing_multiplier"] = json!(billing_multiplier);
payload["billing_cost"] = json!(billing_cost);
payload["routing_group_id"] = json!(item.routing_group_id());
payload["routing_group_name"] = json!(item.routing_group_name());
let object = payload
.as_object_mut()
.expect("admin usage record payload should be an object");
@@ -2723,7 +2736,7 @@ pub fn build_admin_usage_replay_plan_response(
mod tests {
use std::collections::BTreeMap;
use serde_json::json;
use serde_json::{json, Value};
use super::{
admin_usage_active_request_json, admin_usage_client_is_stream, admin_usage_has_body_value,
@@ -2848,6 +2861,98 @@ mod tests {
assert_eq!(record["client_is_stream"], false);
}
#[test]
fn admin_usage_payloads_preserve_routing_group_snapshot_and_precise_display_cost() {
for (metadata, multiplier, cost, group_name) in [
(None, 1.0, json!(0.0), Value::Null),
(
Some(json!({"routing_group_billing_multiplier": 0.0})),
0.0,
json!(0.0),
Value::Null,
),
(
Some(json!({
"routing_group_billing_multiplier": 2.5,
"routing_group_id": "group-1",
"routing_group_name": "请求时的分组",
"rate_multiplier": 0.5
})),
2.5,
json!(0.000004),
json!("请求时的分组"),
),
(
Some(json!({
"billing_multiplier_snapshot": {"version": 1, "factors": {"routing_group": 2.5, "user_group": 2.0}, "multiplier": 5.0},
"routing_group_billing_multiplier": 2.5,
"routing_group_id": "group-1",
"routing_group_name": "请求时的分组",
"rate_multiplier": 0.5
})),
5.0,
json!(0.000007),
json!("请求时的分组"),
),
] {
let item = StoredRequestUsageAudit {
total_cost_usd: 0.00000149,
actual_total_cost_usd: 0.0000002,
request_metadata: metadata,
..sample_usage("completed", Some(200), None)
};
let record = admin_usage_record_json(
&item,
&BTreeMap::new(),
&BTreeMap::new(),
false,
false,
None,
);
let active = admin_usage_active_request_json(&item, None, None, None);
let detail = build_admin_usage_detail_payload(
&item,
&BTreeMap::new(),
&BTreeMap::new(),
false,
false,
None,
false,
None,
&BTreeMap::new(),
);
for payload in [&record, &active, &detail] {
assert_eq!(payload["billing_multiplier"], multiplier);
assert_eq!(payload["billing_cost"], cost);
assert_eq!(payload["routing_group_name"], group_name);
assert_eq!(
payload["routing_group_id"],
if group_name.is_null() {
Value::Null
} else {
json!("group-1")
}
);
assert_eq!(
payload["rate_multiplier"],
if group_name.is_null() {
Value::Null
} else {
json!(0.5)
}
);
assert_eq!(payload["cost"], 0.000001);
assert_eq!(payload["actual_cost"], 0.0);
}
}
let item = StoredRequestUsageAudit {
total_cost_usd: f64::MAX,
request_metadata: Some(json!({"routing_group_billing_multiplier": 2.0})),
..sample_usage("completed", Some(200), None)
};
assert!(admin_usage_active_request_json(&item, None, None, None)["billing_cost"].is_null());
}
#[test]
fn admin_usage_payloads_expose_response_model_separately_from_mapping() {
let item = StoredRequestUsageAudit {
crates/aether-ai/formats/src/codex_profile.rs
+4 -1
View File
@@ -109,7 +109,10 @@ mod tests {
assert_eq!(profile.originator, "codex_cli_rs");
assert!(profile.user_agent.starts_with("codex_cli_rs/0.200.1 ("));
assert!(profile.user_agent.ends_with(") unknown"));
assert!(profile.user_agent.contains(std::env::consts::ARCH));
let architecture = super::OS_INFO
.architecture()
.unwrap_or(std::env::consts::ARCH);
assert!(profile.user_agent.contains(&format!("; {architecture})")));
}
#[test]
crates/aether-data/adapters/postgres/migrations/20261007000000_separate_customer_billing_amount.sql
+154
View File
@@ -0,0 +1,154 @@
-- Customer charges use the immutable request-time factor snapshot. Provider
-- procurement cost remains in actual_total_cost_usd for legacy reporting.
CREATE OR REPLACE FUNCTION public.usage_customer_billable_amount(
metadata jsonb, base_cost numeric, legacy_cost numeric
) RETURNS numeric LANGUAGE plpgsql IMMUTABLE PARALLEL SAFE AS $$
DECLARE factor jsonb; multiplier numeric; amount numeric;
factor_name text; factor_value jsonb; factor_number double precision;
expected_multiplier double precision := 1.0; factor_count integer := 0;
has_zero boolean := false;
BEGIN
IF metadata ? 'billing_multiplier_snapshot' THEN
IF jsonb_typeof(metadata->'billing_multiplier_snapshot') <> 'object'
OR metadata #> '{billing_multiplier_snapshot,version}' IS DISTINCT FROM '1'::jsonb
OR jsonb_typeof(metadata #> '{billing_multiplier_snapshot,factors}') IS DISTINCT FROM 'object'
THEN RETURN NULL; END IF;
factor := metadata #> '{billing_multiplier_snapshot,multiplier}';
FOR factor_name, factor_value IN
SELECT key, value FROM jsonb_each(metadata #> '{billing_multiplier_snapshot,factors}') ORDER BY key COLLATE "C"
LOOP
factor_count := factor_count + 1;
IF factor_count > 16 OR factor_name = '' OR length(factor_name) > 64
OR factor_name !~ '^[A-Za-z0-9_]+$'
OR jsonb_typeof(factor_value) IS DISTINCT FROM 'number'
THEN RETURN NULL; END IF;
factor_number := factor_value::text::double precision;
IF factor_number < 0 OR factor_number > 1.7976931348623157e308::double precision
THEN RETURN NULL; END IF;
has_zero := has_zero OR factor_number = 0;
END LOOP;
-- Rust short-circuits zero before multiplying any of the other factors.
IF has_zero THEN expected_multiplier := 0;
ELSE
FOR factor_name, factor_value IN
SELECT key, value FROM jsonb_each(metadata #> '{billing_multiplier_snapshot,factors}') ORDER BY key COLLATE "C"
LOOP
factor_number := factor_value::text::double precision;
BEGIN
expected_multiplier := expected_multiplier * factor_number;
EXCEPTION WHEN numeric_value_out_of_range THEN
-- PostgreSQL raises on float underflow; Rust rounds that product to 0.
IF expected_multiplier::numeric * factor_number::numeric > 1.7976931348623157e308::numeric
THEN RETURN NULL; END IF;
expected_multiplier := 0;
END;
END LOOP;
END IF;
ELSIF metadata ? 'routing_group_billing_multiplier' THEN
factor := metadata->'routing_group_billing_multiplier';
expected_multiplier := NULL;
ELSE
RETURN CASE WHEN legacy_cost NOT IN ('NaN'::numeric,'Infinity'::numeric,'-Infinity'::numeric)
THEN round(legacy_cost,8) END;
END IF;
IF jsonb_typeof(factor) IS DISTINCT FROM 'number' THEN RETURN NULL; END IF;
multiplier := factor::text::numeric;
factor_number := factor::text::double precision;
IF factor_number < 0
OR factor_number > 1.7976931348623157e308::double precision
OR (expected_multiplier IS NOT NULL AND factor_number <> expected_multiplier)
OR multiplier < 0 OR multiplier > 1.7976931348623157e308::numeric
OR base_cost IS NULL OR base_cost < 0
OR base_cost IN ('NaN'::numeric,'Infinity'::numeric,'-Infinity'::numeric)
THEN RETURN NULL; END IF;
amount := base_cost * multiplier;
IF amount > 1.7976931348623157e308::numeric THEN RETURN NULL; END IF;
RETURN round(amount,8);
EXCEPTION WHEN numeric_value_out_of_range OR invalid_text_representation THEN
-- Corrupt captured pricing must not abort an entire analytics query.
RETURN NULL;
END $$;
CREATE OR REPLACE VIEW public.usage_analytics_facts_v1 AS
SELECT u.request_id, COALESCE(u.id, u.request_id) AS id, u.created_at,
CASE WHEN identity.owner_id IS NOT NULL AND identity.is_standalone=false THEN identity.owner_id END AS actor_user_id,
identity.owner_id AS credential_owner_id,
CASE WHEN identity.owner_id IS NULL THEN 'unknown' WHEN identity.is_standalone THEN 'standalone'
WHEN NOT identity.is_standalone THEN 'employee' ELSE 'unknown' END AS attribution_kind,
CASE WHEN identity.owner_id IS NULL THEN 'unknown' WHEN identity.is_standalone THEN 'standalone_key'
WHEN NOT identity.is_standalone THEN 'user_account' ELSE 'unknown' END AS attribution_source,
COALESCE(a.record_kind, 'request') AS record_kind, a.parent_request_id,
u.api_key_id, u.model, u.target_model, u.provider_id, u.provider_name,
u.api_format, u.endpoint_kind, u.request_type, u.is_stream, u.has_format_conversion,
u.status, u.status_code, u.error_category, u.failure_origin, u.failure_stage, u.failure_reason,
u.failure_schema_version, u.response_time_ms, u.first_byte_time_ms,
COALESCE(s.billing_status, u.billing_status) AS settlement_status,
COALESCE(metadata.value->'usage_available', 'true'::jsonb) <> 'false'::jsonb AS usage_available,
COALESCE(metadata.value->'usage_pricing_available', 'true'::jsonb) <> 'false'::jsonb
AND (s.billing_total_cost_usd IS NOT NULL OR COALESCE(s.billing_status, u.billing_status) = 'settled') AS pricing_available,
CASE WHEN COALESCE(metadata.value->'usage_available', 'true'::jsonb) <> 'false'::jsonb
THEN b.input_tokens END AS input_tokens,
CASE WHEN COALESCE(metadata.value->'usage_available', 'true'::jsonb) <> 'false'::jsonb
THEN b.output_tokens END AS output_tokens,
CASE WHEN COALESCE(metadata.value->'usage_available', 'true'::jsonb) <> 'false'::jsonb
THEN b.total_tokens END AS total_tokens,
CASE WHEN COALESCE(metadata.value->'usage_available', 'true'::jsonb) <> 'false'::jsonb
THEN b.cache_read_input_tokens END AS cache_read_input_tokens,
CASE WHEN COALESCE(metadata.value->'usage_available', 'true'::jsonb) <> 'false'::jsonb
THEN b.cache_creation_input_tokens END AS cache_creation_input_tokens,
CASE WHEN COALESCE(metadata.value->'usage_pricing_available', 'true'::jsonb) <> 'false'::jsonb
AND (s.billing_total_cost_usd IS NOT NULL OR COALESCE(s.billing_status, u.billing_status) = 'settled')
THEN round(COALESCE(s.billing_total_cost_usd::numeric, u.total_cost_usd::numeric), 8) END AS rated_amount,
CASE WHEN COALESCE(metadata.value->'usage_pricing_available', 'true'::jsonb) <> 'false'::jsonb
AND (s.billing_actual_total_cost_usd IS NOT NULL OR COALESCE(s.billing_status, u.billing_status) = 'settled')
THEN public.usage_customer_billable_amount(metadata.value,
COALESCE(s.billing_total_cost_usd::numeric, u.total_cost_usd::numeric),
COALESCE(s.billing_actual_total_cost_usd::numeric, u.actual_total_cost_usd::numeric)) END AS billable_amount,
s.quota_covered_amount_usd AS quota_covered_amount,
s.wallet_consumed_amount_usd AS wallet_consumed_amount,
s.wallet_debit_amount_usd AS wallet_debit_amount,
s.wallet_recharge_debit_usd AS wallet_recharge_debit_amount,
s.wallet_gift_debit_usd AS wallet_gift_debit_amount,
s.wallet_overdraft_usd AS wallet_overdraft_amount,
s.allocation_status, s.finalized_at AS settled_at,
CASE WHEN s.billing_total_cost_usd IS NOT NULL THEN 'settlement_snapshot' ELSE 'legacy_float' END AS amount_source,
b.upstream_is_stream,
CASE WHEN metadata.value #>> '{analytics_measurement,source}' IN ('reported','estimated','mixed')
THEN metadata.value #>> '{analytics_measurement,source}' ELSE 'unknown' END AS token_source,
CASE WHEN COALESCE(metadata.value->'usage_available','true'::jsonb) <> 'false'::jsonb
AND COALESCE(metadata.value->'usage_pricing_available','true'::jsonb) <> 'false'::jsonb
AND s.input_price_per_1m IS NOT NULL AND s.billing_cache_read_cost_usd IS NOT NULL
THEN round(s.input_price_per_1m::numeric * b.cache_read_input_tokens::numeric / 1000000,8) END AS cache_estimated_full_cost_amount,
CASE WHEN COALESCE(metadata.value->'usage_available','true'::jsonb) <> 'false'::jsonb
AND COALESCE(metadata.value->'usage_pricing_available','true'::jsonb) <> 'false'::jsonb
AND s.input_price_per_1m IS NOT NULL AND s.billing_cache_read_cost_usd IS NOT NULL
THEN round(s.billing_cache_read_cost_usd::numeric,8) END AS cache_read_cost_amount,
CASE WHEN COALESCE(metadata.value->'usage_available','true'::jsonb) <> 'false'::jsonb
AND COALESCE(metadata.value->'usage_pricing_available','true'::jsonb) <> 'false'::jsonb
AND s.input_price_per_1m IS NOT NULL AND s.billing_cache_creation_cost_usd IS NOT NULL
THEN round(s.billing_cache_creation_cost_usd::numeric,8) END AS cache_creation_cost_amount
FROM public.usage u
-- OFFSET 0 keeps this projection from being flattened: large metadata is
-- detoasted and parsed once per request, rather than once per metric expression.
CROSS JOIN LATERAL (SELECT u.request_metadata::jsonb AS value OFFSET 0) metadata
LEFT JOIN public.usage_settlement_snapshots s USING (request_id)
LEFT JOIN public.usage_attribution_snapshots a USING (request_id)
JOIN public.usage_billing_facts b USING (request_id)
LEFT JOIN public.api_keys k ON k.id=u.api_key_id
CROSS JOIN LATERAL (
SELECT CASE WHEN a.request_id IS NOT NULL THEN a.credential_owner_id
WHEN EXISTS (SELECT 1 FROM public.users WHERE id=u.user_id AND NOT is_deleted) THEN u.user_id END AS owner_id,
COALESCE(k.is_standalone,
CASE WHEN jsonb_typeof(metadata.value #> '{analytics_attribution,is_standalone}')='boolean'
THEN (metadata.value #>> '{analytics_attribution,is_standalone}')::boolean END,
CASE WHEN jsonb_typeof(metadata.value->'api_key_is_standalone')='boolean'
THEN (metadata.value->>'api_key_is_standalone')::boolean END,
CASE WHEN a.attribution_source='user_account' THEN false
WHEN a.attribution_source='standalone_key' THEN true END,
CASE WHEN u.api_key_id IS NULL THEN false END) AS is_standalone
) identity;
-- Do not backfill existing rows or scan historical usage during the upgrade.
-- Historical daily totals retain their legacy charge through the read fallback;
-- normal daily aggregation writes billing_cost for newly aggregated days.
ALTER TABLE public.stats_daily ADD COLUMN IF NOT EXISTS billing_cost numeric(20,8);
crates/aether-data/adapters/postgres/src/auth.rs
+153 -4
View File
@@ -561,7 +561,18 @@ SET
rate_limit = CASE WHEN $7 THEN $8 ELSE rate_limit END,
concurrent_limit = CASE WHEN $9 THEN $10 ELSE concurrent_limit END,
ip_rules = CASE WHEN $11 THEN $12::jsonb ELSE ip_rules END,
feature_settings = CASE WHEN $13 THEN $14::jsonb ELSE feature_settings END,
feature_settings = CASE WHEN $16 THEN
NULLIF(
(COALESCE(CASE WHEN $13 THEN $14::jsonb ELSE feature_settings END, '{}'::jsonb)
- 'routing_group_id' - 'routing_group_name')
|| CASE WHEN $17 THEN
CASE WHEN $18::text IS NULL THEN '{}'::jsonb
ELSE jsonb_build_object('routing_group_id', $18::text) END
WHEN jsonb_typeof(feature_settings->'routing_group_id') = 'string' THEN
jsonb_build_object('routing_group_id', feature_settings->'routing_group_id')
ELSE '{}'::jsonb END,
'{}'::jsonb)
ELSE CASE WHEN $13 THEN $14::jsonb ELSE feature_settings END END,
updated_at = NOW()
WHERE user_id = $1
AND id = $2
@@ -1357,6 +1368,18 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.bind(record.feature_settings.is_some())
.bind(feature_settings)
.bind(false)
.bind(record.routing_group_selection.is_some())
.bind(
record
.routing_group_selection
.as_ref()
.is_some_and(|patch| patch.group_id.is_some()),
)
.bind(
record
.routing_group_selection
.and_then(|patch| patch.group_id.flatten()),
)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
@@ -1418,6 +1441,18 @@ WHERE id = $2
.bind(record.feature_settings.is_some())
.bind(feature_settings)
.bind(true)
.bind(record.routing_group_selection.is_some())
.bind(
record
.routing_group_selection
.as_ref()
.is_some_and(|patch| patch.group_id.is_some()),
)
.bind(
record
.routing_group_selection
.and_then(|patch| patch.group_id.flatten()),
)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
@@ -2143,12 +2178,126 @@ mod tests {
.contains("key_encrypted = CASE WHEN $3 THEN $4 ELSE key_encrypted END"));
assert!(UPDATE_USER_API_KEY_BASIC_SQL
.contains("ip_rules = CASE WHEN $11 THEN $12::jsonb ELSE ip_rules END"));
assert!(UPDATE_USER_API_KEY_BASIC_SQL.contains(
"feature_settings = CASE WHEN $13 THEN $14::jsonb ELSE feature_settings END"
));
assert!(UPDATE_USER_API_KEY_BASIC_SQL
.contains("CASE WHEN $13 THEN $14::jsonb ELSE feature_settings END"));
assert!(UPDATE_USER_API_KEY_BASIC_SQL.contains("AND ($15 = FALSE OR is_locked = FALSE)"));
}
#[tokio::test]
#[ignore = "requires AETHER_TEST_DATABASE_URL; uses only a temporary table"]
async fn live_api_key_routing_patch_preserves_concurrent_feature_edits() {
use aether_data_contracts::repository::auth::{
AuthApiKeyWriteRepository, UpdateApiKeyRoutingGroupSelection,
UpdateUserApiKeyBasicRecord,
};
use serde_json::json;
let pool = sqlx::postgres::PgPoolOptions::new()
.max_connections(1)
.connect(&std::env::var("AETHER_TEST_DATABASE_URL").unwrap())
.await
.unwrap();
// The repository's complete production UPDATE runs against a session-local table.
sqlx::raw_sql(
r#"
CREATE TEMP TABLE api_keys (
id text PRIMARY KEY, user_id text, key_hash text, key_encrypted text, name text,
allowed_providers json, allowed_api_formats json, allowed_models json,
ip_rules jsonb, rate_limit integer, concurrent_limit integer,
force_capabilities json, feature_settings jsonb, is_active boolean DEFAULT true,
is_locked boolean DEFAULT false, is_standalone boolean DEFAULT false,
expires_at timestamptz, auto_delete_on_expiry boolean DEFAULT false,
total_requests bigint DEFAULT 0, total_tokens bigint DEFAULT 0,
total_cost_usd numeric DEFAULT 0, last_used_at timestamptz,
created_at timestamptz DEFAULT NOW(), updated_at timestamptz DEFAULT NOW()
);
INSERT INTO api_keys (id,user_id,key_hash,name,feature_settings)
VALUES ('key-1','user-1','hash-1','key','{"routing_group_id":"a","pii":false}');
"#,
)
.execute(&pool)
.await
.unwrap();
let repository = SqlxAuthApiKeySnapshotReadRepository::new(pool.clone());
let patch = |features, group_id| 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: features,
routing_group_selection: Some(UpdateApiKeyRoutingGroupSelection { group_id }),
};
// Prepared before the group change: stale or injected group fields must not win.
let stale_feature_edit =
patch(Some(Some(json!({"routing_group_id":"a","pii":true}))), None);
repository
.update_user_api_key_basic_if_unlocked(patch(None, Some(Some("b".into()))))
.await
.unwrap()
.unwrap();
let edited = repository
.update_user_api_key_basic_if_unlocked(stale_feature_edit)
.await
.unwrap()
.unwrap();
assert_eq!(
edited.feature_settings,
Some(json!({"routing_group_id":"b","pii":true}))
);
let changed = repository
.update_user_api_key_basic_if_unlocked(patch(None, Some(Some("c".into()))))
.await
.unwrap()
.unwrap();
assert_eq!(
changed.feature_settings,
Some(json!({"routing_group_id":"c","pii":true}))
);
let cleared_features = repository
.update_user_api_key_basic_if_unlocked(patch(Some(None), None))
.await
.unwrap()
.unwrap();
assert_eq!(
cleared_features.feature_settings,
Some(json!({"routing_group_id":"c"}))
);
let cleared_group = repository
.update_user_api_key_basic_if_unlocked(patch(None, Some(None)))
.await
.unwrap()
.unwrap();
assert_eq!(cleared_group.feature_settings, None);
let mut admin = patch(Some(Some(json!({"admin":true}))), None);
admin.routing_group_selection = None;
assert_eq!(
repository
.update_user_api_key_basic(admin)
.await
.unwrap()
.unwrap()
.feature_settings,
Some(json!({"admin":true}))
);
sqlx::query("UPDATE api_keys SET is_locked=true")
.execute(&pool)
.await
.unwrap();
assert!(repository
.update_user_api_key_basic_if_unlocked(patch(None, Some(Some("d".into()))))
.await
.unwrap()
.is_none());
pool.close().await;
}
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
crates/aether-data/adapters/postgres/src/settlement.rs
+79
View File
@@ -1509,6 +1509,85 @@ mod tests {
(pool, schema)
}
#[tokio::test]
#[ignore = "requires AETHER_TEST_DATABASE_URL and PostgreSQL migrations"]
async fn live_composite_billing_settlement_preserves_provider_cost_and_is_idempotent() {
use super::*;
let (pool, schema) = isolated_settlement_test_pool().await;
let result = AssertUnwindSafe(async {
for table in ["wallets", "usage", "usage_settlement_snapshots", "usage_counter_deltas"] {
sqlx::query(&format!("CREATE TABLE {table} (LIKE public.{table} INCLUDING ALL)"))
.execute(&pool).await.expect("settlement fixture table should be created");
}
let repository = SqlxSettlementRepository::new(pool.clone());
for (scenario, charge, quota_covered) in [
("wallet", 20.0, 0.0),
("quota_and_wallet", 20.0, 7.0),
("zero_charge", 0.0, 0.0),
] {
sqlx::query("INSERT INTO users (id, username, email_verified) VALUES ($1, $1, false)")
.bind(scenario).execute(&pool).await.expect("user should insert");
sqlx::query("INSERT INTO wallets (id, user_id, balance, gift_balance, total_consumed, limit_mode, created_at, updated_at) VALUES ($1, $1, 100, 0, 0, 'finite', NOW(), NOW())")
.bind(scenario).execute(&pool).await.expect("wallet should insert");
// A zero-charge request must leave an active quota untouched too.
if scenario != "wallet" {
let grant = serde_json::json!([{
"type": "daily_quota", "daily_quota_usd": 7.0,
"reset_timezone": "UTC", "allow_wallet_overage": true,
}]);
sqlx::query("INSERT INTO billing_plans (id, title, price_amount, duration_unit, duration_value, entitlements_json, created_at, updated_at) VALUES ($1, $1, 10, 'month', 1, $2, NOW(), NOW())")
.bind(scenario).bind(&grant).execute(&pool).await.expect("plan should insert");
sqlx::query("INSERT INTO user_plan_entitlements (id, user_id, plan_id, payment_order_id, starts_at, expires_at, entitlements_snapshot, created_at, updated_at) VALUES ($1, $1, $1, $1, NOW() - INTERVAL '1 hour', NOW() + INTERVAL '1 day', $2, NOW(), NOW())")
.bind(scenario).bind(&grant).execute(&pool).await.expect("entitlement should insert");
}
let multiplier = charge / 10.0;
let metadata = serde_json::json!({"billing_multiplier_snapshot": {
"version": 1, "factors": {"routing_group": multiplier}, "multiplier": multiplier,
}});
sqlx::query("INSERT INTO usage (id, request_id, user_id, provider_id, provider_name, model, status, billing_status, total_cost_usd, actual_total_cost_usd, request_metadata) VALUES ($1, $1, $1, 'provider', 'Provider', 'model', 'completed', 'pending', 10, 5, $2)")
.bind(scenario).bind(metadata).execute(&pool).await.expect("usage should insert");
let input = UsageSettlementInput {
request_id: scenario.to_string(), user_id: Some(scenario.to_string()),
api_key_id: None, api_key_is_standalone: false,
provider_id: Some("provider".to_string()),
status: "completed".to_string(), billing_status: "pending".to_string(),
total_cost_usd: 10.0, actual_total_cost_usd: 5.0,
billing_cost_usd: Some(charge), finalized_at_unix_secs: None,
};
let settled = repository.settle_usage(input.clone()).await.unwrap().unwrap();
assert_eq!(settled.billing_status, "settled", "{scenario}");
assert_eq!(settled.wallet_balance_before, Some(100.0));
assert_eq!(settled.wallet_balance_after, Some(100.0 - (charge - quota_covered)));
assert_eq!(repository.settle_usage(input).await.unwrap(), Some(settled), "replayed {scenario}");
let wallet: (f64, f64) = sqlx::query_as("SELECT (balance + gift_balance)::double precision, total_consumed::double precision FROM wallets WHERE id = $1")
.bind(scenario).fetch_one(&pool).await.unwrap();
assert_eq!(wallet, (100.0 - (charge - quota_covered), charge - quota_covered), "{scenario}");
let quota: (i64, f64) = sqlx::query_as("SELECT COUNT(*), COALESCE(SUM(amount_usd), 0)::double precision FROM entitlement_usage_ledgers WHERE request_id = $1")
.bind(scenario).fetch_one(&pool).await.unwrap();
assert_eq!(quota, (if quota_covered > 0.0 { 1 } else { 0 }, quota_covered), "{scenario}");
let allocation: (f64, f64, f64, String) = sqlx::query_as("SELECT quota_covered_amount_usd::double precision, wallet_consumed_amount_usd::double precision, wallet_debit_amount_usd::double precision, allocation_status FROM usage_settlement_snapshots WHERE request_id = $1")
.bind(scenario).fetch_one(&pool).await.unwrap();
assert_eq!(allocation, (quota_covered, charge - quota_covered, charge - quota_covered, "complete".to_string()), "{scenario}");
let costs: (f64, f64) = sqlx::query_as("SELECT total_cost_usd::double precision, actual_total_cost_usd::double precision FROM usage WHERE request_id = $1")
.bind(scenario).fetch_one(&pool).await.unwrap();
assert_eq!(costs, (10.0, 5.0), "base and upstream cost must remain unchanged");
let provider_cost: (i64, f64) = sqlx::query_as("SELECT COUNT(*), COALESCE(SUM(total_cost_usd_delta), 0)::double precision FROM usage_counter_deltas WHERE request_id = $1 AND kind = 'provider_monthly' AND target_id = 'provider'")
.bind(scenario).fetch_one(&pool).await.unwrap();
assert_eq!(provider_cost, (1, 5.0), "upstream cost must be recorded once even for a zero-charge request");
}
}).catch_unwind().await;
sqlx::query(&format!("DROP SCHEMA {schema} CASCADE"))
.execute(&pool)
.await
.expect("isolated settlement schema should be removed");
pool.close().await;
if let Err(panic) = result {
std::panic::resume_unwind(panic);
}
}
#[tokio::test]
#[ignore = "requires AETHER_TEST_DATABASE_URL and PostgreSQL migrations"]
async fn live_usage_policy_window_aggregates_preserve_exact_admission_and_idempotency() {
crates/aether-data/adapters/postgres/src/usage/analytics_tests.rs
+121
View File
@@ -682,6 +682,7 @@ async fn live_overview_settlement_allocations_preserve_unlimited_and_finite_wall
billing_status: "pending".into(),
total_cost_usd: cost,
actual_total_cost_usd: cost,
billing_cost_usd: None,
finalized_at_unix_secs: None,
};
assert_eq!(
@@ -1266,6 +1267,19 @@ async fn live_overview_dashboard_total_matches_canonical_settlement_and_legacy_t
serde_json::json!({}),
1002,
),
(
"billing-snapshot",
"openai:chat",
120,
serde_json::json!({
"billing_multiplier_snapshot": {
"version": 1,
"factors": {"routing_group": 2.0, "user_group": 0.75},
"multiplier": 1.5
}
}),
120,
),
(
"unavailable",
"openai:chat",
@@ -1340,7 +1354,114 @@ async fn live_overview_dashboard_total_matches_canonical_settlement_and_legacy_t
.await
.unwrap();
assert_eq!(total.total_tokens, expected_tokens, "{case}");
if case == "billing-snapshot" {
assert_eq!(total.billable_amount.as_deref(), Some("0.37500000"));
}
assert_dashboard_total_matches_canonical(&total, &canonical);
}
tx.rollback().await.unwrap();
}
#[tokio::test]
#[ignore = "requires migrated isolated AETHER_TEST_DATABASE_URL"]
async fn live_customer_billing_amount_matches_canonical_and_dashboard_facts() {
let pool = sqlx::PgPool::connect(&std::env::var("AETHER_TEST_DATABASE_URL").unwrap())
.await
.unwrap();
let mut tx = pool.begin().await.unwrap();
let start = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap();
let composite = serde_json::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
});
for (case, metadata, expected) in [
("legacy", serde_json::json!({}), Some("0.50000000")),
("composite", composite.clone(), Some("3.00000000")),
("settlement-base", composite, Some("6.00000000")),
(
"free",
serde_json::json!({"routing_group_billing_multiplier": 0}),
Some("0.00000000"),
),
(
"null",
serde_json::json!({"billing_multiplier_snapshot": null, "routing_group_billing_multiplier": 1}),
None,
),
(
"negative-factor",
serde_json::json!({"billing_multiplier_snapshot": {"version": 1, "factors": {"routing_group": -1}, "multiplier": 1}}),
None,
),
(
"mismatch",
serde_json::json!({"billing_multiplier_snapshot": {"version": 1, "factors": {"routing_group": 2}, "multiplier": 1}}),
None,
),
(
"negative-legacy",
serde_json::json!({"routing_group_billing_multiplier": -1}),
None,
),
(
"string-legacy",
serde_json::json!({"routing_group_billing_multiplier": "1"}),
None,
),
(
"zero-before-overflow",
serde_json::json!({"billing_multiplier_snapshot": {"version": 1, "factors": {"a": 1e308, "b": 1e308, "z": 0}, "multiplier": 0}}),
Some("0.00000000"),
),
(
"overflow",
serde_json::json!({"billing_multiplier_snapshot": {"version": 1, "factors": {"a": 1e308, "b": 1e308}, "multiplier": 1}}),
None,
),
(
"bad-key",
serde_json::json!({"billing_multiplier_snapshot": {"version": 1, "factors": {"routing-group": 2}, "multiplier": 2}}),
None,
),
] {
let request = uuid::Uuid::new_v4().to_string();
sqlx::query("INSERT INTO usage(id,request_id,model,provider_name,status,billing_status,total_cost_usd,actual_total_cost_usd,created_at,request_metadata) VALUES($1,$1,$1,'billing-test','completed','settled',2,0.5,$2,$3)")
.bind(&request).bind(start).bind(metadata).execute(&mut *tx).await.unwrap();
if case == "settlement-base" {
sqlx::query("INSERT INTO usage_settlement_snapshots(request_id,billing_status,billing_total_cost_usd,billing_actual_total_cost_usd) VALUES($1,'settled',4,0.25)")
.bind(&request).execute(&mut *tx).await.unwrap();
}
let amount: Option<String> = sqlx::query_scalar(
"SELECT billable_amount::text FROM usage_analytics_facts_v1 WHERE request_id=$1",
)
.bind(&request)
.fetch_one(&mut *tx)
.await
.unwrap();
assert_eq!(amount.as_deref(), expected, "{case}");
let query = UsageAnalyticsQuery {
from_unix_ms: start.timestamp_millis() as u64,
to_unix_ms: (start + chrono::Duration::hours(1)).timestamp_millis() as u64,
model: Some(request),
..Default::default()
};
let canonical = super::analytics::read_analytics_metrics(&mut tx, &query, false)
.await
.unwrap();
let inline = super::dashboard::read_dashboard_total_metrics(&mut tx, &query, false)
.await
.unwrap();
assert_dashboard_total_matches_canonical(&inline, &canonical);
if let Some(expected) = expected {
assert_eq!(
canonical.billable_amount.as_deref(),
Some(expected),
"{case}"
);
}
}
tx.rollback().await.unwrap();
}
crates/aether-data/adapters/postgres/src/usage/dashboard.rs
+8 -8
View File
@@ -17,10 +17,10 @@ SELECT u.created_at, u.api_key_id, u.model, u.provider_id, u.api_format, u.endpo
u.request_type, u.status, u.is_stream, u.has_format_conversion, u.failure_origin,
'request'::text AS record_kind,
COALESCE(s.billing_status, u.billing_status) AS settlement_status,
COALESCE(availability.usage_available, 'true'::jsonb) <> 'false'::jsonb AS usage_available,
COALESCE(availability.usage_pricing_available, 'true'::jsonb) <> 'false'::jsonb
COALESCE(metadata.value->'usage_available', 'true'::jsonb) <> 'false'::jsonb AS usage_available,
COALESCE(metadata.value->'usage_pricing_available', 'true'::jsonb) <> 'false'::jsonb
AND (s.billing_total_cost_usd IS NOT NULL OR COALESCE(s.billing_status, u.billing_status) = 'settled') AS pricing_available,
CASE WHEN COALESCE(availability.usage_available, 'true'::jsonb) <> 'false'::jsonb THEN
CASE WHEN COALESCE(metadata.value->'usage_available', 'true'::jsonb) <> 'false'::jsonb THEN
GREATEST(
COALESCE(
CASE
@@ -89,15 +89,15 @@ SELECT u.created_at, u.api_key_id, u.model, u.provider_id, u.api_format, u.endpo
),
0
)::bigint END AS total_tokens,
CASE WHEN COALESCE(availability.usage_pricing_available, 'true'::jsonb) <> 'false'::jsonb
CASE WHEN COALESCE(metadata.value->'usage_pricing_available', 'true'::jsonb) <> 'false'::jsonb
AND (s.billing_actual_total_cost_usd IS NOT NULL OR COALESCE(s.billing_status, u.billing_status) = 'settled')
THEN round(COALESCE(s.billing_actual_total_cost_usd::numeric, u.actual_total_cost_usd::numeric), 8) END AS billable_amount,
THEN public.usage_customer_billable_amount(metadata.value,
COALESCE(s.billing_total_cost_usd::numeric, u.total_cost_usd::numeric),
COALESCE(s.billing_actual_total_cost_usd::numeric, u.actual_total_cost_usd::numeric)) END AS billable_amount,
s.allocation_status
FROM public.usage u
LEFT JOIN public.usage_settlement_snapshots s USING (request_id)
CROSS JOIN LATERAL json_to_record(
CASE WHEN json_typeof(u.request_metadata)='object' THEN u.request_metadata ELSE '{}'::json END
) AS availability(usage_available jsonb, usage_pricing_available jsonb)
CROSS JOIN LATERAL (SELECT u.request_metadata::jsonb AS value OFFSET 0) metadata
WHERE NOT EXISTS (SELECT 1 FROM public.usage_attribution_snapshots a
WHERE a.request_id=u.request_id AND a.record_kind='session')
) AS usage_analytics_facts_v1"#;
crates/aether-data/adapters/postgres/src/usage/dashboard_history_tests.rs
+23
View File
@@ -134,6 +134,29 @@ async fn live_dashboard_restores_legacy_history_without_replaying_or_double_coun
assert_eq!(advanced.activity_days, restored.activity_days);
assert_eq!(advanced.active_days, restored.active_days);
// New daily rollups retain customer charges independently after detail
// expires; older NULL daily charges retain their original legacy cost.
sqlx::query("UPDATE stats_daily SET billing_cost=1.5 WHERE id='recent'")
.execute(&pool).await.unwrap();
sqlx::query("DELETE FROM usage WHERE request_id='overlap'")
.execute(&pool).await.unwrap();
let billed_history = repo.query_dashboard_summary(&query).await.unwrap();
assert_eq!(billed_history.total.billable_amount.as_deref(), Some("123456791.59691357"));
assert_eq!(billed_history.total.request_count, restored.total.request_count);
assert_eq!(billed_history.today, restored.today);
let provider_cost: String = sqlx::query_scalar("SELECT actual_total_cost::text FROM stats_daily WHERE id='recent'")
.fetch_one(&pool).await.unwrap();
assert_eq!(provider_cost, "0.30000003");
// The pre-activation live prefix applies the same composite snapshot
// to its finalized base amount, independently of procurement cost.
sqlx::query("UPDATE usage SET request_metadata=$1 WHERE request_id='before-shared'")
.bind(serde_json::json!({"billing_multiplier_snapshot": {"version": 1, "factors": {"routing_group": 2, "user_group": 0.75}, "multiplier": 1.5}}))
.execute(&pool).await.unwrap();
let billed_prefix = repo.query_dashboard_summary(&query).await.unwrap();
assert_eq!(billed_prefix.total.billable_amount.as_deref(), Some("123456791.65864196"));
assert_eq!(billed_prefix.today.billable_amount.as_deref(), Some("0.93518517"));
// A summary cutoff without legacy daily history must leave the normal
// future-only projection and its requested calendar unchanged.
sqlx::query("DELETE FROM stats_daily").execute(&pool).await.unwrap();
crates/aether-data/adapters/postgres/src/usage/mod.rs
+28 -12
View File
@@ -42,18 +42,21 @@ use crate::{
PostgresTransactionRunner,
};
use aether_data_contracts::repository::usage::{
api_key_usage_contribution, model_usage_contribution, provider_api_key_usage_contribution,
sanitize_usage_capture_controls_for_persistence, sanitize_usage_for_persistence,
sanitize_usage_request_metadata, usage_can_recover_terminal_failure,
usage_error_category_for_status_code, usage_lifecycle_update_allowed, ApiKeyUsageDelta,
ModelUsageDelta, PendingUsageCleanupSummary, ProviderApiKeyUsageContribution,
ProviderApiKeyUsageDelta, ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary,
StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredRequestUsageAudit,
StoredUsageDailySummary, UpsertUsageRecord, UsageAuditListQuery, UsageCounterFlushSummary,
UsageCounterHealthSnapshot, UsageCounterPendingHealthSnapshot, UsageDailyHeatmapQuery,
UsageReadRepository, UsageWriteRepository, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY,
PROVIDER_REASONING_EFFORT_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY,
REQUESTED_REASONING_EFFORT_METADATA_KEY,
api_key_usage_contribution, model_usage_contribution, preserve_usage_routing_group_snapshot,
provider_api_key_usage_contribution, sanitize_usage_capture_controls_for_persistence,
sanitize_usage_for_persistence, sanitize_usage_request_metadata,
usage_can_recover_terminal_failure, usage_error_category_for_status_code,
usage_lifecycle_update_allowed, ApiKeyUsageDelta, ModelUsageDelta, PendingUsageCleanupSummary,
ProviderApiKeyUsageContribution, ProviderApiKeyUsageDelta, ProviderApiKeyWindowUsageRequest,
StoredProviderApiKeyUsageSummary, StoredProviderApiKeyWindowUsageSummary,
StoredProviderUsageSummary, StoredRequestUsageAudit, StoredUsageDailySummary,
UpsertUsageRecord, UsageAuditListQuery, UsageCounterFlushSummary, UsageCounterHealthSnapshot,
UsageCounterPendingHealthSnapshot, UsageDailyHeatmapQuery, UsageReadRepository,
UsageWriteRepository, 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 aether_data_contracts::DataLayerError;
@@ -8762,6 +8765,15 @@ ORDER BY "usage".user_id ASC
);
request_metadata_json = json_bind_text(request_metadata_value.as_ref())?;
}
if capture_update_allowed {
request_metadata_value = preserve_usage_routing_group_snapshot(
request_metadata_value,
previous_usage
.as_ref()
.and_then(|stored| stored.request_metadata.as_ref()),
);
request_metadata_json = json_bind_text(request_metadata_value.as_ref())?;
}
let _row = sqlx::query(UPSERT_SQL)
.bind(Uuid::new_v4().to_string())
.bind(&usage.request_id)
@@ -12554,6 +12566,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) = previous_metadata.get(key) {
retained.insert(key.to_string(), value.clone());
crates/aether-data/adapters/postgres/src/usage/queries/dashboard_history.sql
+6 -2
View File
@@ -11,7 +11,7 @@ WITH daily AS (
CASE WHEN effective_input_tokens=0 AND input_tokens>0 THEN input_tokens
ELSE effective_input_tokens END + cache_creation_tokens + cache_read_tokens
ELSE total_input_context END AS cache_input_tokens,
actual_total_cost::numeric AS billable_amount
COALESCE(billing_cost,actual_total_cost::numeric) AS billable_amount
FROM stats_daily
), facts AS MATERIALIZED (
SELECT (day AT TIME ZONE 'UTC')::date AS day, request_count,
@@ -22,7 +22,11 @@ WITH daily AS (
SELECT (b.created_at AT TIME ZONE 'UTC')::date, 1::bigint,
b.input_tokens, b.output_tokens, b.total_tokens, b.cache_creation_input_tokens,
b.cache_read_input_tokens, b.total_input_context,
COALESCE(s.billing_actual_total_cost_usd::numeric,u.actual_total_cost_usd::numeric)
CASE WHEN COALESCE(u.request_metadata::jsonb->'usage_pricing_available','true'::jsonb)<>'false'::jsonb
AND (s.billing_actual_total_cost_usd IS NOT NULL OR COALESCE(s.billing_status,u.billing_status)='settled')
THEN public.usage_customer_billable_amount(u.request_metadata::jsonb,
COALESCE(s.billing_total_cost_usd::numeric,u.total_cost_usd::numeric),
COALESCE(s.billing_actual_total_cost_usd::numeric,u.actual_total_cost_usd::numeric)) END
FROM usage_billing_facts b
JOIN usage u USING (request_id)
LEFT JOIN usage_settlement_snapshots s USING (request_id)
crates/aether-data/adapters/postgres/src/usage/queries/list_recent_usage_audits_prefix.sql
+20 -2
View File
@@ -176,6 +176,10 @@ SELECT
NULL::bytea AS client_response_body_compressed,
CASE
WHEN NULLIF(BTRIM("usage".request_metadata->>'client_ip'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'routing_group_id'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'routing_group_name'), '') IS NOT NULL
OR json_typeof("usage".request_metadata->'routing_group_billing_multiplier') = 'number'
OR "usage".request_metadata->'billing_multiplier_snapshot' IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'user_agent'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'request_path'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL
@@ -192,7 +196,17 @@ SELECT
OR ("usage".request_metadata->>'usage_pricing_available') IN ('true', 'false')
OR json_typeof("usage".request_metadata->'live_session') = 'object'
OR json_typeof("usage".request_metadata->'realtime_session') = 'object'
THEN jsonb_strip_nulls(jsonb_build_object(
THEN (jsonb_strip_nulls(jsonb_build_object(
'routing_group_id',
NULLIF(BTRIM("usage".request_metadata->>'routing_group_id'), ''),
'routing_group_name',
NULLIF(BTRIM("usage".request_metadata->>'routing_group_name'), ''),
'routing_group_billing_multiplier',
CASE
WHEN json_typeof("usage".request_metadata->'routing_group_billing_multiplier') = 'number'
THEN "usage".request_metadata->'routing_group_billing_multiplier'
ELSE NULL
END,
'client_ip',
NULLIF(BTRIM("usage".request_metadata->>'client_ip'), ''),
'user_agent',
@@ -255,7 +269,11 @@ SELECT
THEN "usage".request_metadata->'realtime_session'
ELSE NULL
END
))::json
)) || CASE
WHEN "usage".request_metadata->'billing_multiplier_snapshot' IS NOT NULL
THEN jsonb_build_object('billing_multiplier_snapshot', "usage".request_metadata->'billing_multiplier_snapshot')
ELSE '{}'::jsonb
END)::json
ELSE NULL::json
END AS request_metadata,
NULL::varchar AS http_request_body_ref,
crates/aether-data/adapters/postgres/src/usage/queries/list_usage_audits_prefix.sql
+20 -2
View File
@@ -176,6 +176,10 @@ SELECT
NULL::bytea AS client_response_body_compressed,
CASE
WHEN NULLIF(BTRIM("usage".request_metadata->>'client_ip'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'routing_group_id'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'routing_group_name'), '') IS NOT NULL
OR json_typeof("usage".request_metadata->'routing_group_billing_multiplier') = 'number'
OR "usage".request_metadata->'billing_multiplier_snapshot' IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'user_agent'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'request_path'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL
@@ -192,7 +196,17 @@ SELECT
OR ("usage".request_metadata->>'usage_pricing_available') IN ('true', 'false')
OR json_typeof("usage".request_metadata->'live_session') = 'object'
OR json_typeof("usage".request_metadata->'realtime_session') = 'object'
THEN jsonb_strip_nulls(jsonb_build_object(
THEN (jsonb_strip_nulls(jsonb_build_object(
'routing_group_id',
NULLIF(BTRIM("usage".request_metadata->>'routing_group_id'), ''),
'routing_group_name',
NULLIF(BTRIM("usage".request_metadata->>'routing_group_name'), ''),
'routing_group_billing_multiplier',
CASE
WHEN json_typeof("usage".request_metadata->'routing_group_billing_multiplier') = 'number'
THEN "usage".request_metadata->'routing_group_billing_multiplier'
ELSE NULL
END,
'client_ip',
NULLIF(BTRIM("usage".request_metadata->>'client_ip'), ''),
'user_agent',
@@ -255,7 +269,11 @@ SELECT
THEN "usage".request_metadata->'realtime_session'
ELSE NULL
END
))::json
)) || CASE
WHEN "usage".request_metadata->'billing_multiplier_snapshot' IS NOT NULL
THEN jsonb_build_object('billing_multiplier_snapshot', "usage".request_metadata->'billing_multiplier_snapshot')
ELSE '{}'::jsonb
END)::json
ELSE NULL::json
END AS request_metadata,
NULL::varchar AS http_request_body_ref,
crates/aether-data/contracts/src/repository/auth.rs
+39
View File
@@ -577,6 +577,45 @@ pub struct UpdateUserApiKeyBasicRecord {
/// unchanged. Keeping this patch in the basic mutation record lets repositories apply the
/// complete user-key update in one atomic write.
pub feature_settings: Option<Option<serde_json::Value>>,
/// Self-service updates merge routing selection separately against the
/// current stored settings. `None` retains administrative replacement semantics.
pub routing_group_selection: Option<UpdateApiKeyRoutingGroupSelection>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpdateApiKeyRoutingGroupSelection {
/// `None` preserves the latest stored group; `Some(None)` follows the
/// default; `Some(Some(id))` selects the validated public group.
pub group_id: Option<Option<String>>,
}
impl UpdateApiKeyRoutingGroupSelection {
/// Repositories must call this while holding the same write lock as the
/// surrounding API key mutation, so unrelated edits cannot restore a stale
/// group choice or a stale feature-settings object.
pub fn merge_feature_settings(
&self,
current: Option<&serde_json::Value>,
replacement: Option<Option<serde_json::Value>>,
) -> Option<serde_json::Value> {
let group_id = match &self.group_id {
None => current
.and_then(|value| value.get("routing_group_id"))
.and_then(serde_json::Value::as_str)
.map(|id| serde_json::Value::String(id.to_string())),
Some(group_id) => group_id.clone().map(serde_json::Value::String),
};
let mut settings = replacement
.unwrap_or_else(|| current.cloned())
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
settings.remove("routing_group_id");
settings.remove("routing_group_name");
if let Some(group_id) = group_id {
settings.insert("routing_group_id".to_string(), group_id);
}
(!settings.is_empty()).then_some(serde_json::Value::Object(settings))
}
}
impl std::fmt::Debug for UpdateUserApiKeyBasicRecord {
crates/aether-data/contracts/src/repository/settlement/types.rs
+57 -3
View File
@@ -346,6 +346,10 @@ pub struct UsageSettlementInput {
pub billing_status: String,
pub total_cost_usd: f64,
pub actual_total_cost_usd: f64,
/// Customer charge after all captured billing factors, independent of upstream cost.
/// Missing values retain the legacy charge based on `actual_total_cost_usd`.
#[serde(default)]
pub billing_cost_usd: Option<f64>,
pub finalized_at_unix_secs: Option<u64>,
}
@@ -366,6 +370,14 @@ impl UsageSettlementInput {
"settlement cost must be finite".to_string(),
));
}
if self
.billing_cost_usd
.is_some_and(|value| !value.is_finite() || value < 0.0)
{
return Err(crate::DataLayerError::InvalidInput(
"settlement billing_cost_usd must be finite and non-negative".to_string(),
));
}
Ok(())
}
}
@@ -511,14 +523,17 @@ pub fn settlement_billing_status_for_usage_status(status: &str) -> &'static str
}
pub fn settlement_billable_cost_usd(input: &UsageSettlementInput) -> f64 {
input.actual_total_cost_usd.max(0.0)
input
.billing_cost_usd
.unwrap_or(input.actual_total_cost_usd)
.max(0.0)
}
#[cfg(test)]
mod tests {
use super::{
validate_wallet_settlement_values, ReconcileUsagePolicyCostInput,
ReserveUsagePolicyCostInput, ReserveUsagePolicyRequestInput,
settlement_billable_cost_usd, validate_wallet_settlement_values,
ReconcileUsagePolicyCostInput, ReserveUsagePolicyCostInput, ReserveUsagePolicyRequestInput,
UsagePolicyCostReservationState, UsagePolicyCostWindow, UsagePolicyRequestWindow,
UsageSettlementInput,
};
@@ -535,11 +550,50 @@ mod tests {
billing_status: "pending".to_string(),
total_cost_usd: 0.1,
actual_total_cost_usd: 0.1,
billing_cost_usd: None,
finalized_at_unix_secs: None,
};
assert!(input.validate().is_err());
}
#[test]
fn explicit_customer_charge_is_validated_independently_of_upstream_cost() {
let mut input: UsageSettlementInput = serde_json::from_value(serde_json::json!({
"request_id": "billing-charge",
"user_id": "user-1",
"api_key_id": null,
"provider_id": "provider-1",
"status": "completed",
"billing_status": "pending",
"total_cost_usd": 2.0,
"actual_total_cost_usd": 0.5,
"finalized_at_unix_secs": null,
}))
.expect("legacy settlement input should deserialize");
assert_eq!(input.billing_cost_usd, None);
assert_eq!(settlement_billable_cost_usd(&input), 0.5);
assert!(input.validate().is_ok());
for charge in [3.0, 0.0] {
input.billing_cost_usd = Some(charge);
assert!(input.validate().is_ok());
assert_eq!(settlement_billable_cost_usd(&input), charge);
assert_eq!(input.actual_total_cost_usd, 0.5);
assert_eq!(
serde_json::from_value::<UsageSettlementInput>(
serde_json::to_value(&input).unwrap()
)
.unwrap(),
input
);
}
for charge in [-0.01, f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
input.billing_cost_usd = Some(charge);
assert!(input.validate().is_err());
}
}
#[test]
fn wallet_settlement_values_reject_corruption_and_overflow() {
assert!(validate_wallet_settlement_values(-3.0, 0.0, 12.0, 1.0).is_ok());
crates/aether-data/contracts/src/repository/usage/billing_multiplier.rs
+169
View File
@@ -0,0 +1,169 @@
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::DataLayerError;
use super::ROUTING_GROUP_BILLING_MULTIPLIER_METADATA_KEY;
pub const BILLING_MULTIPLIER_SNAPSHOT_METADATA_KEY: &str = "billing_multiplier_snapshot";
/// Immutable customer pricing factors. Provider Key rates belong to upstream cost,
/// not this snapshot. Add future factors (for example `user_group`) at admission.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct BillingMultiplierSnapshot {
version: u32,
factors: BTreeMap<String, f64>,
multiplier: f64,
}
impl BillingMultiplierSnapshot {
pub fn from_factors(factors: BTreeMap<String, f64>) -> Result<Self, DataLayerError> {
if factors.len() > 16
|| factors.iter().any(|(name, value)| {
name.is_empty()
|| name.len() > 64
|| !name
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || byte == b'_')
|| !value.is_finite()
|| *value < 0.0
})
{
return Err(invalid_snapshot());
}
let multiplier = if factors.values().any(|value| *value == 0.0) {
0.0
} else {
factors.values().product::<f64>()
};
if !multiplier.is_finite() {
return Err(invalid_snapshot());
}
Ok(Self {
version: 1,
factors,
multiplier,
})
}
pub fn validate(&self) -> Result<(), DataLayerError> {
let expected = Self::from_factors(self.factors.clone())?;
if self.version != 1 || self.multiplier != expected.multiplier {
return Err(invalid_snapshot());
}
Ok(())
}
pub fn multiplier(&self) -> f64 {
self.multiplier
}
pub fn cost(&self, base_cost: f64) -> Result<f64, DataLayerError> {
self.validate()?;
let cost = base_cost * self.multiplier;
if !base_cost.is_finite() || base_cost < 0.0 || !cost.is_finite() {
return Err(DataLayerError::InvalidInput(
"customer billing cost must be finite and non-negative".to_string(),
));
}
// Match wallet storage and usage-policy cost units (eight decimals).
// Scaling a finite large amount must not introduce infinity by itself.
let scaled = cost * 100_000_000.0;
Ok(if scaled.is_finite() {
scaled.round() / 100_000_000.0
} else {
cost
})
}
}
fn invalid_snapshot() -> DataLayerError {
DataLayerError::InvalidInput("invalid billing multiplier snapshot".to_string())
}
/// None preserves legacy charging. A malformed captured snapshot is an error,
/// never an instruction to silently charge a different rate.
pub fn billing_multiplier_snapshot(
metadata: Option<&Value>,
) -> Result<Option<BillingMultiplierSnapshot>, DataLayerError> {
let Some(metadata) = metadata.and_then(Value::as_object) else {
return Ok(None);
};
if let Some(value) = metadata.get(BILLING_MULTIPLIER_SNAPSHOT_METADATA_KEY) {
let snapshot: BillingMultiplierSnapshot =
serde_json::from_value(value.clone()).map_err(|_| invalid_snapshot())?;
snapshot.validate()?;
return Ok(Some(snapshot));
}
if let Some(value) = metadata.get(ROUTING_GROUP_BILLING_MULTIPLIER_METADATA_KEY) {
let multiplier = value.as_f64().ok_or_else(invalid_snapshot)?;
return BillingMultiplierSnapshot::from_factors(BTreeMap::from([(
"routing_group".to_string(),
multiplier,
)]))
.map(Some);
}
Ok(None)
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn composes_customer_factors_without_provider_cost_and_freezes_them() {
let snapshot = BillingMultiplierSnapshot::from_factors(BTreeMap::from([
("routing_group".to_string(), 2.0),
("user_group".to_string(), 0.75),
]))
.unwrap();
assert_eq!(snapshot.multiplier(), 1.5);
assert_eq!(snapshot.cost(10.0).unwrap(), 15.0);
assert_eq!(snapshot.cost(0.123456789).unwrap(), 0.18518518);
let metadata = json!({"billing_multiplier_snapshot": snapshot, "routing_group_billing_multiplier": 99, "rate_multiplier": 0.1});
assert_eq!(
billing_multiplier_snapshot(Some(&metadata)).unwrap(),
Some(snapshot)
);
assert_eq!(billing_multiplier_snapshot(None).unwrap(), None);
}
#[test]
fn rejects_corrupt_overflowing_snapshots_and_accepts_zero_rates() {
for factors in [
BTreeMap::from([("routing_group".into(), -1.0)]),
BTreeMap::from([("routing_group".into(), f64::INFINITY)]),
BTreeMap::from([
("routing_group".into(), f64::MAX),
("user_group".into(), 2.0),
]),
] {
assert!(BillingMultiplierSnapshot::from_factors(factors).is_err());
}
let zero = BillingMultiplierSnapshot::from_factors(BTreeMap::from([
("routing_group".into(), 0.0),
("user_group".into(), 2.0),
]))
.unwrap();
assert_eq!(zero.cost(10.0).unwrap(), 0.0);
for invalid in [
Value::Null,
json!({"version": 2, "factors": {}, "multiplier": 1}),
json!({"version": 1, "factors": {"routing_group": 2}, "multiplier": 1}),
] {
assert!(billing_multiplier_snapshot(Some(
&json!({"billing_multiplier_snapshot": invalid})
))
.is_err());
}
let doubled = BillingMultiplierSnapshot::from_factors(BTreeMap::from([(
"routing_group".into(),
2.0,
)]))
.unwrap();
assert!(doubled.cost(f64::MAX).is_err());
}
}
crates/aether-data/contracts/src/repository/usage/metadata_policy.rs
+209 -4
View File
@@ -8,14 +8,17 @@ use serde_json::{Map, Value};
use crate::repository::candidates::sanitize_request_candidate_skip_reason;
use super::{
normalize_provider_response_model, LIVE_SESSION_METADATA_KEY,
billing_multiplier_snapshot, normalize_provider_response_model,
BILLING_MULTIPLIER_SNAPSHOT_METADATA_KEY, LIVE_SESSION_METADATA_KEY,
PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY, PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY,
PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, PROVIDER_REASONING_EFFORT_METADATA_KEY,
PROVIDER_RESPONSE_MODEL_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY,
REALTIME_SESSION_METADATA_KEY, REQUESTED_REASONING_EFFORT_METADATA_KEY,
ROUTING_CANDIDATE_SKIP_REASON_METADATA_KEY, ROUTING_FAILURE_DIAGNOSTIC_METADATA_KEY,
USAGE_AVAILABLE_METADATA_KEY, USAGE_PRICING_AVAILABLE_METADATA_KEY,
WEBSOCKET_MODE_METADATA_KEY, WEBSOCKET_TRANSPORT_METADATA_KEY,
ROUTING_GROUP_BILLING_MULTIPLIER_METADATA_KEY, ROUTING_GROUP_ID_METADATA_KEY,
ROUTING_GROUP_NAME_METADATA_KEY, USAGE_AVAILABLE_METADATA_KEY,
USAGE_PRICING_AVAILABLE_METADATA_KEY, WEBSOCKET_MODE_METADATA_KEY,
WEBSOCKET_TRANSPORT_METADATA_KEY,
};
const UPSTREAM_IS_STREAM_KEY: &str = "upstream_is_stream";
@@ -43,8 +46,68 @@ pub fn sanitize_usage_request_metadata_ref(value: Option<&Value>) -> Option<Valu
sanitize_usage_request_metadata_object(value?.as_object()?)
}
/// Keep the request's first captured billing snapshot and reservation owner across retries.
pub fn preserve_usage_routing_group_snapshot(
incoming: Option<Value>,
previous: Option<&Value>,
) -> Option<Value> {
let Some(previous) = previous.and_then(Value::as_object) else {
return incoming;
};
let mut snapshot = Map::from_iter(
[
BILLING_MULTIPLIER_SNAPSHOT_METADATA_KEY,
ROUTING_GROUP_BILLING_MULTIPLIER_METADATA_KEY,
ROUTING_GROUP_ID_METADATA_KEY,
ROUTING_GROUP_NAME_METADATA_KEY,
PLAN_USAGE_RESERVATION_TOKEN_KEY,
]
.into_iter()
.filter_map(|key| {
previous
.get(key)
.map(|value| (key.to_string(), value.clone()))
}),
);
if !snapshot.contains_key(BILLING_MULTIPLIER_SNAPSHOT_METADATA_KEY)
&& snapshot.contains_key(ROUTING_GROUP_BILLING_MULTIPLIER_METADATA_KEY)
{
let captured = billing_multiplier_snapshot(Some(&Value::Object(snapshot.clone())))
.ok()
.flatten()
.and_then(|snapshot| serde_json::to_value(snapshot).ok())
.unwrap_or(Value::Null);
snapshot.insert(
BILLING_MULTIPLIER_SNAPSHOT_METADATA_KEY.to_string(),
captured,
);
}
let Some(Value::Object(snapshot)) = sanitize_usage_request_metadata_object(&snapshot) else {
return incoming;
};
let mut metadata = incoming
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
metadata.extend(snapshot);
Some(Value::Object(metadata))
}
pub fn sanitize_usage_request_metadata_object(source: &Map<String, Value>) -> Option<Value> {
let mut target = Map::new();
if let Some(snapshot) = source.get(BILLING_MULTIPLIER_SNAPSHOT_METADATA_KEY) {
let metadata = serde_json::json!({BILLING_MULTIPLIER_SNAPSHOT_METADATA_KEY: snapshot});
let snapshot = match billing_multiplier_snapshot(Some(&metadata)) {
Ok(Some(snapshot)) => serde_json::to_value(snapshot)
.expect("validated billing multiplier snapshot must serialize"),
// Preserve an invalid marker so malformed financial input cannot silently
// fall back to legacy billing after metadata projection.
_ => Value::Null,
};
target.insert(
BILLING_MULTIPLIER_SNAPSHOT_METADATA_KEY.to_string(),
snapshot,
);
}
if let Some(source) = source
.get("analytics_measurement")
.and_then(|value| value.get("source"))
@@ -81,6 +144,8 @@ pub fn sanitize_usage_request_metadata_object(source: &Map<String, Value>) -> Op
}
insert_token(source, &mut target, "trace_id", 128);
insert_token(source, &mut target, ROUTING_GROUP_ID_METADATA_KEY, 128);
insert_bounded_text(source, &mut target, ROUTING_GROUP_NAME_METADATA_KEY, 256);
insert_ip_address(source, &mut target, "client_ip");
insert_client_family(source, &mut target);
for key in [
@@ -181,6 +246,7 @@ pub fn sanitize_usage_request_metadata_object(source: &Map<String, Value>) -> Op
for key in [
"rate_multiplier",
ROUTING_GROUP_BILLING_MULTIPLIER_METADATA_KEY,
"input_price_per_1m",
"output_price_per_1m",
"cache_creation_price_per_1m",
@@ -189,6 +255,20 @@ pub fn sanitize_usage_request_metadata_object(source: &Map<String, Value>) -> Op
] {
insert_nonnegative_number(source, &mut target, key);
}
// An invalid legacy routing factor must remain a financial tombstone. Dropping it
// would make a subsequent reader silently fall back to the historical provider charge.
if source
.get(ROUTING_GROUP_BILLING_MULTIPLIER_METADATA_KEY)
.is_some_and(|value| {
!value
.as_f64()
.is_some_and(|value| value.is_finite() && value >= 0.0)
})
{
target
.entry(BILLING_MULTIPLIER_SNAPSHOT_METADATA_KEY.to_string())
.or_insert(Value::Null);
}
let billing_snapshot = source
.get("billing_snapshot")
@@ -1156,6 +1236,27 @@ fn insert_token(
target.insert(key.to_string(), Value::String(value.to_string()));
}
fn insert_bounded_text(
source: &Map<String, Value>,
target: &mut Map<String, Value>,
key: &str,
max_len: usize,
) {
let Some(value) = source
.get(key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| {
!value.is_empty()
&& value.chars().count() <= max_len
&& !value.chars().any(char::is_control)
})
else {
return;
};
target.insert(key.to_string(), Value::String(value.to_string()));
}
fn insert_dimension_token(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
let Some(value) = source
.get(key)
@@ -1263,7 +1364,111 @@ fn safe_version_value(value: &Value) -> Option<String> {
mod tests {
use serde_json::json;
use super::{sanitize_usage_request_metadata, sanitize_usage_request_metadata_ref};
use super::{
billing_multiplier_snapshot, preserve_usage_routing_group_snapshot,
sanitize_usage_request_metadata, sanitize_usage_request_metadata_ref,
};
#[test]
fn billing_multiplier_snapshot_projection_preserves_invalid_marker_and_immutable_factors() {
for snapshot in [
serde_json::Value::Null,
json!({"version": 1, "factors": {"routing_group": 2.0}, "multiplier": 1.0}),
json!({"version": 99, "factors": {}, "multiplier": 1.0}),
] {
let projected = sanitize_usage_request_metadata(Some(json!({
"billing_multiplier_snapshot": snapshot,
"routing_group_billing_multiplier": 0.5,
})))
.unwrap();
assert_eq!(
projected.get("billing_multiplier_snapshot"),
Some(&serde_json::Value::Null)
);
assert!(billing_multiplier_snapshot(Some(&projected)).is_err());
let preserved = preserve_usage_routing_group_snapshot(
Some(json!({"billing_multiplier_snapshot": {"version": 1, "factors": {}, "multiplier": 1.0}})),
Some(&projected),
).unwrap();
assert!(billing_multiplier_snapshot(Some(&preserved)).is_err());
}
let legacy =
json!({"routing_group_billing_multiplier": 0.25, "routing_group_name": "历史分组"});
let preserved = preserve_usage_routing_group_snapshot(
Some(json!({"billing_multiplier_snapshot": {"version": 1, "factors": {"routing_group": 99.0}, "multiplier": 99.0}})),
Some(&legacy),
).unwrap();
assert_eq!(
billing_multiplier_snapshot(Some(&preserved))
.unwrap()
.unwrap()
.multiplier(),
0.25
);
assert_eq!(preserved["routing_group_name"], "历史分组");
}
#[test]
fn billing_multiplier_snapshot_projection_rejects_malformed_legacy_factors() {
for factor in [serde_json::Value::Null, json!(-1), json!("2"), json!({})] {
let projected = sanitize_usage_request_metadata(Some(json!({
"routing_group_billing_multiplier": factor,
})))
.expect("invalid financial input must retain a tombstone");
assert_eq!(
projected["billing_multiplier_snapshot"],
serde_json::Value::Null
);
assert!(billing_multiplier_snapshot(Some(&projected)).is_err());
assert_eq!(
sanitize_usage_request_metadata(Some(projected.clone())),
Some(projected)
);
}
let generic = sanitize_usage_request_metadata(Some(json!({
"routing_group_billing_multiplier": -1,
"billing_multiplier_snapshot": {"version": 1, "factors": {"routing_group": 2}, "multiplier": 2},
})))
.unwrap();
assert_eq!(
billing_multiplier_snapshot(Some(&generic))
.unwrap()
.unwrap()
.multiplier(),
2.0
);
}
#[test]
fn billing_multiplier_snapshot_preserves_the_original_reservation_owner() {
let token_a = "550e8400-e29b-41d4-a716-446655440001";
let token_b = "550e8400-e29b-41d4-a716-446655440002";
let incoming = json!({
"plan_usage_reservation_token": token_b,
"billing_multiplier_snapshot": {"version": 1, "factors": {"routing_group": 3}, "multiplier": 3},
});
let previous = json!({
"plan_usage_reservation_token": token_a,
"billing_multiplier_snapshot": {"version": 1, "factors": {"routing_group": 0.5}, "multiplier": 0.5},
});
let preserved =
preserve_usage_routing_group_snapshot(Some(incoming.clone()), Some(&previous)).unwrap();
assert_eq!(preserved["plan_usage_reservation_token"], token_a);
assert_eq!(
billing_multiplier_snapshot(Some(&preserved))
.unwrap()
.unwrap()
.multiplier(),
0.5
);
for empty in [json!({}), json!({"plan_usage_reservation_token": " "})] {
let preserved =
preserve_usage_routing_group_snapshot(Some(incoming.clone()), Some(&empty))
.unwrap();
assert_eq!(preserved["plan_usage_reservation_token"], token_b);
}
}
#[test]
fn account_attribution_preserves_key_flag_without_custom_identity_or_purpose() {
crates/aether-data/contracts/src/repository/usage/mod.rs
+6 -2
View File
@@ -2,6 +2,7 @@ mod analytics;
#[cfg(test)]
mod analytics_tests;
mod attribution;
mod billing_multiplier;
mod capture_memory;
mod compression;
mod dashboard_summary;
@@ -12,6 +13,7 @@ mod types;
pub use analytics::*;
pub use attribution::*;
pub use billing_multiplier::*;
#[doc(hidden)]
pub use capture_memory::{
mark_usage_capture_memory_omitted, usage_json_heap_estimate, UsageCaptureMemoryBudget,
@@ -60,6 +62,8 @@ pub use types::{
PROVIDER_RESPONSE_MODEL_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY,
REALTIME_SESSION_METADATA_KEY, REQUESTED_REASONING_EFFORT_METADATA_KEY,
ROUTING_CANDIDATE_SKIP_REASON_METADATA_KEY, ROUTING_FAILURE_DIAGNOSTIC_METADATA_KEY,
USAGE_AVAILABLE_METADATA_KEY, USAGE_PRICING_AVAILABLE_METADATA_KEY,
WEBSOCKET_MODE_METADATA_KEY, WEBSOCKET_TRANSPORT_METADATA_KEY,
ROUTING_GROUP_BILLING_MULTIPLIER_METADATA_KEY, ROUTING_GROUP_ID_METADATA_KEY,
ROUTING_GROUP_NAME_METADATA_KEY, USAGE_AVAILABLE_METADATA_KEY,
USAGE_PRICING_AVAILABLE_METADATA_KEY, WEBSOCKET_MODE_METADATA_KEY,
WEBSOCKET_TRANSPORT_METADATA_KEY,
};
crates/aether-data/contracts/src/repository/usage/types.rs
+79
View File
@@ -11,6 +11,10 @@ pub const PROVIDER_RESPONSE_MODEL_METADATA_KEY: &str = "provider_response_model"
pub const PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY: &str = "provider_cache_ttl_minutes";
pub const ROUTING_CANDIDATE_SKIP_REASON_METADATA_KEY: &str = "routing_candidate_skip_reason";
pub const ROUTING_FAILURE_DIAGNOSTIC_METADATA_KEY: &str = "routing_failure_diagnostic";
/// Immutable routing-group multiplier captured when the request is planned.
pub const ROUTING_GROUP_BILLING_MULTIPLIER_METADATA_KEY: &str = "routing_group_billing_multiplier";
pub const ROUTING_GROUP_ID_METADATA_KEY: &str = "routing_group_id";
pub const ROUTING_GROUP_NAME_METADATA_KEY: &str = "routing_group_name";
pub const WEBSOCKET_MODE_METADATA_KEY: &str = "websocket_mode";
pub const WEBSOCKET_TRANSPORT_METADATA_KEY: &str = "websocket_transport";
pub const PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY: &str = "plan_usage_reservation_deferred";
@@ -828,6 +832,47 @@ impl StoredRequestUsageAudit {
self.request_metadata_number("rate_multiplier")
}
/// Historical requests without a captured multiplier retain the original 1x rate.
pub fn routing_group_billing_multiplier(&self) -> f64 {
self.request_metadata_number(ROUTING_GROUP_BILLING_MULTIPLIER_METADATA_KEY)
.filter(|value| value.is_finite() && *value >= 0.0)
.unwrap_or(1.0)
}
/// Routing factor projection retained for callers inspecting this individual factor.
pub fn routing_group_billing_cost(&self) -> Option<f64> {
let cost = self.total_cost_usd * self.routing_group_billing_multiplier();
cost.is_finite().then_some(cost)
}
pub fn billing_multiplier(&self) -> f64 {
super::billing_multiplier_snapshot(self.request_metadata.as_ref())
.ok()
.flatten()
.map(|snapshot| snapshot.multiplier())
.unwrap_or(1.0)
}
/// Customer charge is independent of upstream Key cost. Legacy rows keep their
/// original charge; no current configuration is consulted for historical usage.
pub fn billing_cost(&self) -> Option<f64> {
match super::billing_multiplier_snapshot(self.request_metadata.as_ref()).ok()? {
Some(snapshot) => snapshot.cost(self.total_cost_usd).ok(),
None => self
.actual_total_cost_usd
.is_finite()
.then_some(self.actual_total_cost_usd.max(0.0)),
}
}
pub fn routing_group_id(&self) -> Option<&str> {
self.request_metadata_string(ROUTING_GROUP_ID_METADATA_KEY)
}
pub fn routing_group_name(&self) -> Option<&str> {
self.request_metadata_string(ROUTING_GROUP_NAME_METADATA_KEY)
}
pub fn settlement_is_free_tier(&self) -> Option<bool> {
self.request_metadata_bool("is_free_tier")
}
@@ -3407,6 +3452,40 @@ mod tests {
assert!(record.validate().is_err());
}
#[test]
fn routing_group_snapshot_defaults_legacy_multiplier_without_inventing_a_group() {
let mut usage = sample_usage();
usage.total_cost_usd = 4.0;
assert_eq!(usage.routing_group_billing_multiplier(), 1.0);
assert_eq!(usage.routing_group_billing_cost(), Some(4.0));
assert_eq!(usage.routing_group_id(), None);
assert_eq!(usage.routing_group_name(), None);
for (value, multiplier, cost) in [
(json!(0), 0.0, 0.0),
(json!(0.25), 0.25, 1.0),
(json!(2.5), 2.5, 10.0),
(json!(-2), 1.0, 4.0),
(json!("Infinity"), 1.0, 4.0),
(json!(f64::INFINITY), 1.0, 4.0),
(json!(f64::NAN), 1.0, 4.0),
] {
usage.request_metadata = Some(json!({
"routing_group_billing_multiplier": value,
"routing_group_id": "group-recorded",
"routing_group_name": "请求时的分组",
"rate_multiplier": 0.75
}));
assert_eq!(usage.routing_group_billing_multiplier(), multiplier);
assert_eq!(usage.routing_group_billing_cost(), Some(cost));
assert_eq!(usage.routing_group_id(), Some("group-recorded"));
assert_eq!(usage.routing_group_name(), Some("请求时的分组"));
assert_eq!(usage.settlement_rate_multiplier(), Some(0.75));
}
usage.request_metadata = Some(json!({"routing_group_billing_multiplier": 2.0}));
usage.total_cost_usd = f64::MAX;
assert_eq!(usage.routing_group_billing_cost(), None);
}
#[test]
fn settlement_accessors_prefer_typed_metadata() {
let mut usage = sample_usage();
crates/aether-data/contracts/src/repository/wallet/types.rs
+1
View File
@@ -3286,6 +3286,7 @@ mod tests {
billing_status: "pending".to_string(),
total_cost_usd: 0.1,
actual_total_cost_usd: 0.1,
billing_cost_usd: None,
finalized_at_unix_secs: None,
};
assert!(input.validate().is_err());
crates/aether-data/runtime/schema/bootstrap/postgres/001_types_and_tables.sql
+1
View File
@@ -901,6 +901,7 @@ CREATE TABLE IF NOT EXISTS public.stats_daily (
cache_read_tokens bigint DEFAULT '0'::bigint NOT NULL,
total_cost numeric(20,8) DEFAULT '0'::double precision NOT NULL,
actual_total_cost numeric(20,8) DEFAULT '0'::double precision NOT NULL,
billing_cost numeric(20,8),
input_cost numeric(20,8) DEFAULT '0'::double precision NOT NULL,
output_cost numeric(20,8) DEFAULT '0'::double precision NOT NULL,
cache_creation_cost numeric(20,8) DEFAULT '0'::double precision NOT NULL,
crates/aether-data/runtime/schema/bootstrap/postgres/190_overview_analytics.sql
+74 -1
View File
@@ -201,6 +201,77 @@ DROP TRIGGER IF EXISTS overview_usage_delete_attribution ON public.usage;
CREATE TRIGGER overview_usage_delete_attribution BEFORE DELETE ON public.usage
FOR EACH ROW EXECUTE FUNCTION public.overview_delete_attribution();
-- Customer charges use the immutable request-time factor snapshot. Provider
-- procurement cost remains in actual_total_cost_usd for legacy reporting.
CREATE OR REPLACE FUNCTION public.usage_customer_billable_amount(
metadata jsonb, base_cost numeric, legacy_cost numeric
) RETURNS numeric LANGUAGE plpgsql IMMUTABLE PARALLEL SAFE AS $$
DECLARE factor jsonb; multiplier numeric; amount numeric;
factor_name text; factor_value jsonb; factor_number double precision;
expected_multiplier double precision := 1.0; factor_count integer := 0;
has_zero boolean := false;
BEGIN
IF metadata ? 'billing_multiplier_snapshot' THEN
IF jsonb_typeof(metadata->'billing_multiplier_snapshot') <> 'object'
OR metadata #> '{billing_multiplier_snapshot,version}' IS DISTINCT FROM '1'::jsonb
OR jsonb_typeof(metadata #> '{billing_multiplier_snapshot,factors}') IS DISTINCT FROM 'object'
THEN RETURN NULL; END IF;
factor := metadata #> '{billing_multiplier_snapshot,multiplier}';
FOR factor_name, factor_value IN
SELECT key, value FROM jsonb_each(metadata #> '{billing_multiplier_snapshot,factors}') ORDER BY key COLLATE "C"
LOOP
factor_count := factor_count + 1;
IF factor_count > 16 OR factor_name = '' OR length(factor_name) > 64
OR factor_name !~ '^[A-Za-z0-9_]+$'
OR jsonb_typeof(factor_value) IS DISTINCT FROM 'number'
THEN RETURN NULL; END IF;
factor_number := factor_value::text::double precision;
IF factor_number < 0 OR factor_number > 1.7976931348623157e308::double precision
THEN RETURN NULL; END IF;
has_zero := has_zero OR factor_number = 0;
END LOOP;
-- Rust short-circuits zero before multiplying any of the other factors.
IF has_zero THEN expected_multiplier := 0;
ELSE
FOR factor_name, factor_value IN
SELECT key, value FROM jsonb_each(metadata #> '{billing_multiplier_snapshot,factors}') ORDER BY key COLLATE "C"
LOOP
factor_number := factor_value::text::double precision;
BEGIN
expected_multiplier := expected_multiplier * factor_number;
EXCEPTION WHEN numeric_value_out_of_range THEN
-- PostgreSQL raises on float underflow; Rust rounds that product to 0.
IF expected_multiplier::numeric * factor_number::numeric > 1.7976931348623157e308::numeric
THEN RETURN NULL; END IF;
expected_multiplier := 0;
END;
END LOOP;
END IF;
ELSIF metadata ? 'routing_group_billing_multiplier' THEN
factor := metadata->'routing_group_billing_multiplier';
expected_multiplier := NULL;
ELSE
RETURN CASE WHEN legacy_cost NOT IN ('NaN'::numeric,'Infinity'::numeric,'-Infinity'::numeric)
THEN round(legacy_cost,8) END;
END IF;
IF jsonb_typeof(factor) IS DISTINCT FROM 'number' THEN RETURN NULL; END IF;
multiplier := factor::text::numeric;
factor_number := factor::text::double precision;
IF factor_number < 0
OR factor_number > 1.7976931348623157e308::double precision
OR (expected_multiplier IS NOT NULL AND factor_number <> expected_multiplier)
OR multiplier < 0 OR multiplier > 1.7976931348623157e308::numeric
OR base_cost IS NULL OR base_cost < 0
OR base_cost IN ('NaN'::numeric,'Infinity'::numeric,'-Infinity'::numeric)
THEN RETURN NULL; END IF;
amount := base_cost * multiplier;
IF amount > 1.7976931348623157e308::numeric THEN RETURN NULL; END IF;
RETURN round(amount,8);
EXCEPTION WHEN numeric_value_out_of_range OR invalid_text_representation THEN
-- Corrupt captured pricing must not abort an entire analytics query.
RETURN NULL;
END $$;
CREATE OR REPLACE VIEW public.usage_analytics_facts_v1 AS
SELECT u.request_id, COALESCE(u.id, u.request_id) AS id, u.created_at,
CASE WHEN identity.owner_id IS NOT NULL AND identity.is_standalone=false THEN identity.owner_id END AS actor_user_id,
@@ -233,7 +304,9 @@ SELECT u.request_id, COALESCE(u.id, u.request_id) AS id, u.created_at,
THEN round(COALESCE(s.billing_total_cost_usd::numeric, u.total_cost_usd::numeric), 8) END AS rated_amount,
CASE WHEN COALESCE(metadata.value->'usage_pricing_available', 'true'::jsonb) <> 'false'::jsonb
AND (s.billing_actual_total_cost_usd IS NOT NULL OR COALESCE(s.billing_status, u.billing_status) = 'settled')
THEN round(COALESCE(s.billing_actual_total_cost_usd::numeric, u.actual_total_cost_usd::numeric), 8) END AS billable_amount,
THEN public.usage_customer_billable_amount(metadata.value,
COALESCE(s.billing_total_cost_usd::numeric, u.total_cost_usd::numeric),
COALESCE(s.billing_actual_total_cost_usd::numeric, u.actual_total_cost_usd::numeric)) END AS billable_amount,
s.quota_covered_amount_usd AS quota_covered_amount,
s.wallet_consumed_amount_usd AS wallet_consumed_amount,
s.wallet_debit_amount_usd AS wallet_debit_amount,
crates/aether-data/runtime/schema/generated/postgres/baseline/007_stats.sql
+1
View File
@@ -172,6 +172,7 @@ CREATE TABLE IF NOT EXISTS public.stats_daily (
cache_read_tokens bigint DEFAULT 0 NOT NULL,
total_cost double precision DEFAULT 0 NOT NULL,
actual_total_cost double precision DEFAULT 0 NOT NULL,
billing_cost numeric(20,8),
input_cost double precision DEFAULT 0 NOT NULL,
output_cost double precision DEFAULT 0 NOT NULL,
cache_creation_cost double precision DEFAULT 0 NOT NULL,
crates/aether-data/runtime/schema/logical/007_stats.toml
+6
View File
@@ -499,6 +499,12 @@ name = "actual_total_cost"
type = "float64"
default = 0
[[table.stats_daily.columns]]
name = "billing_cost"
type = "decimal_money"
nullable = true
driver.postgres.type = "numeric(20,8)"
[[table.stats_daily.columns]]
name = "input_cost"
type = "float64"
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();
crates/aether-routing-core/src/model.rs
+112 -3
View File
@@ -132,6 +132,86 @@ fn is_false(value: &bool) -> bool {
mod execution_policy_tests {
use super::*;
#[test]
fn group_visibility_is_opt_in_and_round_trips_without_losing_policy() {
let legacy: RoutingGroupConfig = serde_json::from_str("{}").unwrap();
assert!(!legacy.user_visible);
assert!(!RoutingGroupConfig::default().user_visible);
for user_visible in [false, true] {
let config: RoutingGroupConfig = serde_json::from_value(serde_json::json!({
"user_visible": user_visible,
"billing_multiplier": 0.5,
"disabled_providers": ["private-provider"],
"default_policy": { "scheduling_mode": "fixed_order" }
}))
.unwrap();
let encoded = serde_json::to_value(&config).unwrap();
assert_eq!(encoded["user_visible"], user_visible);
assert_eq!(config.billing_multiplier, 0.5);
assert_eq!(config.disabled_providers, ["private-provider"]);
assert_eq!(
serde_json::from_value::<RoutingGroupConfig>(encoded).unwrap(),
config
);
}
for invalid in [
serde_json::json!(null),
serde_json::json!("true"),
serde_json::json!(1),
] {
assert!(
serde_json::from_value::<RoutingGroupConfig>(serde_json::json!({
"user_visible": invalid
}))
.is_err()
);
}
}
#[test]
fn group_billing_multiplier_defaults_to_one_and_rejects_invalid_values() {
let legacy: RoutingGroupConfig = serde_json::from_str("{}").unwrap();
assert_eq!(legacy.billing_multiplier, 1.0);
assert_eq!(RoutingGroupConfig::default().billing_multiplier, 1.0);
for multiplier in [0.0, 0.25, 1.0, 2.5] {
let config: RoutingGroupConfig = serde_json::from_value(serde_json::json!({
"billing_multiplier": multiplier
}))
.unwrap();
crate::validate_routing_group_config(&config).unwrap();
assert_eq!(config.billing_multiplier, multiplier);
assert_eq!(
serde_json::from_value::<RoutingGroupConfig>(
serde_json::to_value(&config).unwrap()
)
.unwrap(),
config
);
}
for multiplier in [-1.0, f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
let config = RoutingGroupConfig {
billing_multiplier: multiplier,
..RoutingGroupConfig::default()
};
assert!(matches!(
crate::validate_routing_group_config(&config),
Err(crate::RoutingValidationError::InvalidBillingMultiplier)
));
}
for value in [
serde_json::json!(null),
serde_json::json!("2"),
serde_json::json!(false),
] {
assert!(
serde_json::from_value::<RoutingGroupConfig>(serde_json::json!({
"billing_multiplier": value
}))
.is_err()
);
}
}
#[test]
fn routing_failover_configuration_round_trips_and_validates() {
let config: RoutingGroupConfig = serde_json::from_value(serde_json::json!({
@@ -188,6 +268,11 @@ pub struct RoutingModelPolicy {
pub allowed_providers: Vec<String>,
#[serde(default)]
pub allowed_keys: Vec<String>,
/// Per-model provider enablement. A `false` value adds a provider to this
/// model's exclusions and `true` removes an inherited exclusion, including
/// one from the legacy group-wide `disabled_providers` baseline.
#[serde(default)]
pub provider_enabled_overrides: BTreeMap<String, bool>,
#[serde(default)]
pub provider_priority_overrides: BTreeMap<String, i32>,
#[serde(default)]
@@ -222,10 +307,17 @@ pub struct RoutingRule {
pub stop_processing: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RoutingGroupConfig {
/// Providers excluded from every model in this group, including providers
/// otherwise selected by model policies or routing rules.
/// Whether authenticated users can discover and explicitly select this
/// group. Private bindings and automatic defaults remain independent.
#[serde(default)]
pub user_visible: bool,
/// Group-wide billing multiplier, snapshotted when a request is routed.
#[serde(default = "default_billing_multiplier")]
pub billing_multiplier: f64,
/// Legacy provider exclusion baseline for the group. Explicit per-model
/// enablement overrides may change it; allowlists and rules cannot.
#[serde(default)]
pub disabled_providers: Vec<String>,
/// The default policy is global for the selected strategy group. Model
@@ -238,6 +330,23 @@ pub struct RoutingGroupConfig {
pub rules: Vec<RoutingRule>,
}
pub(crate) fn default_billing_multiplier() -> f64 {
1.0
}
impl Default for RoutingGroupConfig {
fn default() -> Self {
Self {
user_visible: false,
billing_multiplier: default_billing_multiplier(),
disabled_providers: Vec::new(),
default_policy: RoutingDefaultPolicy::default(),
model_policies: Vec::new(),
rules: Vec::new(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingGroupRecord {
pub id: String,
crates/aether-routing-core/src/policy.rs
+171 -2
View File
@@ -47,10 +47,15 @@ pub struct MatchedRoutingRule {
pub phase: RoutingRulePhase,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ResolvedRoutingPolicy {
#[serde(default = "crate::model::default_billing_multiplier")]
pub billing_multiplier: f64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub group_id: Option<String>,
/// Display name captured alongside the selected group by the gateway.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub group_name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub group_version: Option<i64>,
pub selection_source: String,
@@ -80,7 +85,9 @@ pub fn resolve_routing_policy(
.map_err(|error| RoutingPolicyError::InvalidConfig(error.to_string()))?;
let mut policy = ResolvedRoutingPolicy {
billing_multiplier: config.billing_multiplier,
group_id: input.group_id.map(str::to_string),
group_name: None,
group_version: input.group_version,
selection_source: input.selection_source.to_string(),
requested_model: input.requested_model.to_string(),
@@ -105,6 +112,11 @@ pub fn resolve_routing_policy(
{
apply_model_policy(&mut policy, model_policy);
}
for model_policy in
matching_provider_enablement_policies(config, input.requested_model, input.resolved_model)
{
apply_provider_enable_overrides(&mut policy, model_policy);
}
let condition_context = RoutingConditionContext {
model: input.requested_model,
@@ -284,6 +296,57 @@ fn matching_model_policies<'a>(
.collect()
}
fn matching_provider_enablement_policies<'a>(
config: &'a RoutingGroupConfig,
requested_model: &str,
resolved_model: &str,
) -> Vec<&'a RoutingModelPolicy> {
let mut matches = matching_model_policies(config, requested_model, resolved_model)
.into_iter()
.filter(|policy| !policy.provider_enabled_overrides.is_empty())
.collect::<Vec<_>>();
// Enablement is a layered exception map: broad defaults first, then
// prefixes, then exact model entries. Other model-policy fields retain
// their historical configured-order merge semantics.
matches.sort_by_key(|policy| model_pattern_specificity(&policy.model));
matches
}
fn apply_provider_enable_overrides(
policy: &mut ResolvedRoutingPolicy,
model_policy: &RoutingModelPolicy,
) {
for (provider_id, enabled) in &model_policy.provider_enabled_overrides {
if *enabled {
policy
.ranking_overlay
.disabled_providers
.retain(|disabled| disabled != provider_id);
} else if !policy
.ranking_overlay
.disabled_providers
.iter()
.any(|disabled| disabled == provider_id)
{
policy
.ranking_overlay
.disabled_providers
.push(provider_id.clone());
}
}
}
fn model_pattern_specificity(pattern: &str) -> (u8, usize) {
let pattern = pattern.trim();
if pattern == "*" {
(0, 0)
} else if let Some(prefix) = pattern.strip_suffix('*') {
(1, prefix.len())
} else {
(2, 0)
}
}
fn model_allowed(patterns: &[String], requested_model: &str) -> bool {
patterns.is_empty()
|| patterns
@@ -405,7 +468,7 @@ mod tests {
}
#[test]
fn group_disabled_providers_apply_to_every_model_and_cannot_be_reenabled() {
fn legacy_group_exclusions_cannot_be_bypassed_by_allowlists_or_rule_actions() {
let config: RoutingGroupConfig = serde_json::from_value(json!({
"disabled_providers": ["provider-disabled"],
"model_policies": [{
@@ -478,6 +541,106 @@ mod tests {
}
}
#[test]
fn model_provider_enablement_is_scoped_and_specific_overrides_win() {
let mut config: RoutingGroupConfig = serde_json::from_value(json!({
"disabled_providers": ["provider-root"],
"model_policies": [
{
"model": "*",
"provider_enabled_overrides": {
"provider-model": false,
"provider-specific": false
}
},
{
"model": "model-*",
"provider_enabled_overrides": {
"provider-model": true,
"provider-specific": true
}
},
{
"model": "model-exact",
"provider_enabled_overrides": {
"provider-specific": false,
"provider-exact": true,
"provider-root": true
}
}
]
}))
.unwrap();
// Persisted order need not put broad defaults first. Only the new
// enablement map follows specificity; priorities retain their old order.
config.model_policies.reverse();
config.model_policies[0]
.provider_priority_overrides
.insert("provider-model".into(), 1);
config.model_policies[2]
.provider_priority_overrides
.insert("provider-model".into(), 9);
let for_model = |model: &str| {
resolve_routing_policy(
&config,
RoutingPolicyInput {
group_id: Some("group-1"),
group_version: Some(1),
selection_source: "test",
requested_model: model,
resolved_model: model,
api_format: "openai:chat",
user_id: None,
api_key_id: None,
headers: &json!({}),
body: &json!({}),
phase: RoutingRulePhase::ClientRequest,
},
)
.unwrap()
};
let exact = for_model("model-exact");
assert!(exact.ranking_overlay.provider_allowed("provider-model"));
assert_eq!(
exact.ranking_overlay.provider_priority_overrides["provider-model"],
9
);
assert!(!exact.ranking_overlay.provider_allowed("provider-specific"));
assert!(exact.ranking_overlay.provider_allowed("provider-exact"));
assert!(exact.ranking_overlay.provider_allowed("provider-root"));
let wildcard_prefix = for_model("model-other");
assert!(wildcard_prefix
.ranking_overlay
.provider_allowed("provider-model"));
assert!(wildcard_prefix
.ranking_overlay
.provider_allowed("provider-specific"));
assert!(!wildcard_prefix
.ranking_overlay
.provider_allowed("provider-root"));
let unrelated = for_model("other-model");
assert!(!unrelated.ranking_overlay.provider_allowed("provider-model"));
assert!(!unrelated
.ranking_overlay
.provider_allowed("provider-specific"));
assert!(!unrelated.ranking_overlay.provider_allowed("provider-root"));
let encoded = serde_json::to_value(&config).unwrap();
assert_eq!(
encoded["model_policies"][2]["provider_enabled_overrides"]["provider-model"],
false
);
assert_eq!(
serde_json::from_value::<RoutingGroupConfig>(encoded).unwrap(),
config
);
}
#[test]
fn all_model_scheduling_and_rankings_apply_to_future_models() {
let config: RoutingGroupConfig = serde_json::from_value(json!({
@@ -599,6 +762,8 @@ mod tests {
#[test]
fn resolves_model_policy_and_matching_rule() {
let config = RoutingGroupConfig {
user_visible: false,
billing_multiplier: 1.0,
disabled_providers: vec![],
default_policy: RoutingDefaultPolicy::default(),
model_policies: vec![RoutingModelPolicy {
@@ -668,6 +833,8 @@ mod tests {
#[test]
fn default_policy_applies_to_models_without_an_override() {
let config = RoutingGroupConfig {
user_visible: false,
billing_multiplier: 2.5,
disabled_providers: vec![],
default_policy: RoutingDefaultPolicy {
priority_mode: RoutingSetPriorityMode::GlobalKey,
@@ -707,6 +874,7 @@ mod tests {
assert_eq!(special.scheduling_mode, RoutingSchedulingMode::LoadBalance);
assert!(special.keep_priority_on_conversion);
assert_eq!(special.sticky_key_attempts, 3);
assert_eq!(special.billing_multiplier, 2.5);
assert_eq!(
special.ranking_overlay.allowed_providers,
vec!["provider-special"]
@@ -741,6 +909,7 @@ mod tests {
assert_eq!(ordinary.scheduling_mode, RoutingSchedulingMode::LoadBalance);
assert!(ordinary.keep_priority_on_conversion);
assert_eq!(ordinary.sticky_key_attempts, 3);
assert_eq!(ordinary.billing_multiplier, 2.5);
assert!(ordinary.ranking_overlay.allowed_providers.is_empty());
assert!(ordinary.ranking_overlay.allowed_keys.is_empty());
assert!(ordinary
crates/aether-routing-core/src/ranking.rs
+2 -1
View File
@@ -13,7 +13,8 @@ pub enum CandidateKind {
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct RankingOverlay {
/// Group-wide exclusions take precedence over every provider allowlist.
/// Effective provider exclusions after model overrides. These take
/// precedence over every provider allowlist.
#[serde(default)]
pub disabled_providers: Vec<String>,
#[serde(default)]
crates/aether-routing-core/src/trace.rs
+5 -1
View File
@@ -55,11 +55,15 @@ pub struct RoutingRuntimeFacts {
pub priority_mode: Option<RoutingSetPriorityMode>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct RoutingDecisionTrace {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub billing_multiplier: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub group_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub group_name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub group_version: Option<i64>,
pub selection_source: String,
#[serde(default)]
crates/aether-routing-core/src/validation.rs
+5
View File
@@ -28,6 +28,8 @@ const ROUTING_POOL_PRESETS: &[&str] = &[
#[derive(Debug, Error, Clone, PartialEq, Eq)]
pub enum RoutingValidationError {
#[error("routing group billing multiplier must be a non-negative finite number")]
InvalidBillingMultiplier,
#[error("routing failover rules are invalid: {0}")]
InvalidFailoverRules(String),
#[error("routing rule id is empty")]
@@ -71,6 +73,9 @@ pub enum RoutingValidationError {
pub fn validate_routing_group_config(
config: &RoutingGroupConfig,
) -> Result<(), RoutingValidationError> {
if !config.billing_multiplier.is_finite() || config.billing_multiplier < 0.0 {
return Err(RoutingValidationError::InvalidBillingMultiplier);
}
crate::validate_routing_failover_rules(&config.default_policy.execution_policy.failover_rules)
.map_err(RoutingValidationError::InvalidFailoverRules)?;
let mut rule_ids = BTreeSet::new();
crates/aether-testing/integration/src/bin/usage_settlement_hotspot_baseline.rs
+1
View File
@@ -527,6 +527,7 @@ fn settlement_input(index: usize) -> UsageSettlementInput {
billing_status: "pending".to_string(),
total_cost_usd: 0.0,
actual_total_cost_usd: COST_PER_REQUEST_USD,
billing_cost_usd: None,
finalized_at_unix_secs: Some(now_unix_secs().saturating_add(index as u64)),
}
}
crates/aether-usage/runtime/src/request_metadata.rs
+64 -4
View File
@@ -8,9 +8,11 @@ use aether_data_contracts::repository::usage::{
sanitize_usage_request_metadata_object as project_usage_request_metadata_object,
sanitize_usage_request_metadata_ref as project_usage_request_metadata_ref,
usage_body_capture_is_authoritative, UsageBodyCaptureState,
PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY,
PROVIDER_REASONING_EFFORT_METADATA_KEY, PROVIDER_RESPONSE_MODEL_METADATA_KEY,
PROVIDER_SERVICE_TIER_METADATA_KEY, REQUESTED_REASONING_EFFORT_METADATA_KEY,
BILLING_MULTIPLIER_SNAPSHOT_METADATA_KEY, PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY,
PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, PROVIDER_REASONING_EFFORT_METADATA_KEY,
PROVIDER_RESPONSE_MODEL_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 serde_json::{Map, Value};
@@ -112,6 +114,10 @@ pub(crate) fn retain_first_byte_request_metadata(value: Option<Value>) -> Option
| "model_id"
| "global_model_id"
| "global_model_name"
| ROUTING_GROUP_BILLING_MULTIPLIER_METADATA_KEY
| BILLING_MULTIPLIER_SNAPSHOT_METADATA_KEY
| ROUTING_GROUP_ID_METADATA_KEY
| ROUTING_GROUP_NAME_METADATA_KEY
)
});
(!metadata.is_empty()).then_some(Value::Object(metadata))
@@ -542,6 +548,10 @@ mod tests {
"request_path": "/v1/chat/completions",
"upstream_is_stream": true,
"proxy": {"mode": "manual", "node_id": "proxy-1"},
"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-1",
"routing_group_name": "默认调度策略",
"billing_snapshot": {"dimensions": [1, 2, 3]},
"settlement_snapshot": {"status": "pending"},
"stage_timings_ms": {"planning": 12}
@@ -555,7 +565,11 @@ mod tests {
"client_ip": "203.0.113.8",
"request_path": "/v1/chat/completions",
"request_path_and_query": "/v1/chat/completions",
"upstream_is_stream": true
"upstream_is_stream": true,
"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-1",
"routing_group_name": "默认调度策略"
})
);
}
@@ -644,6 +658,52 @@ mod tests {
.is_none());
}
#[test]
fn routing_group_snapshot_survives_seed_and_sanitization() {
for multiplier in [0.0, 0.25, 1.0, 2.5] {
let context = json!({
"routing_group_billing_multiplier": multiplier,
"billing_multiplier_snapshot": {"version": 1, "factors": {"routing_group": multiplier, "user_group": 2.0}, "multiplier": multiplier * 2.0},
"routing_group_id": "group-1",
"routing_group_name": "请求时的分组",
"rate_multiplier": 0.75,
"routing_trace": {"untrusted": true}
});
let metadata = build_usage_request_metadata_seed(&sample_plan(), context.as_object())
.expect("group snapshot should survive projection");
assert_eq!(metadata["routing_group_billing_multiplier"], multiplier);
assert_eq!(
metadata["billing_multiplier_snapshot"],
context["billing_multiplier_snapshot"]
);
assert_eq!(metadata["routing_group_id"], "group-1");
assert_eq!(metadata["routing_group_name"], "请求时的分组");
assert_eq!(metadata["rate_multiplier"], 0.75);
assert!(metadata.get("routing_trace").is_none());
assert_eq!(
sanitize_usage_request_metadata(Some(metadata.clone())),
Some(metadata)
);
}
for multiplier in [json!(-1), json!("Infinity"), json!(f64::NAN)] {
let context = json!({"routing_group_billing_multiplier": multiplier});
let metadata = build_usage_request_metadata_seed(&sample_plan(), context.as_object())
.expect(
"invalid pricing must retain a marker instead of falling back to legacy billing",
);
assert_eq!(
metadata.get("billing_multiplier_snapshot"),
Some(&Value::Null)
);
assert!(
aether_data_contracts::repository::usage::billing_multiplier_snapshot(Some(
&metadata
))
.is_err()
);
}
}
#[test]
fn builds_seed_from_context_and_allowlisted_metadata_only() {
let metadata = build_usage_request_metadata_seed(
crates/aether-usage/runtime/src/runtime.rs
+2 -21
View File
@@ -26,9 +26,7 @@ use crate::request_metadata::{
request_body_derived_facts_action, retain_first_byte_request_metadata,
RequestBodyDerivedFactsAction,
};
use crate::settlement::{
reconcile_usage_policy_cost_for_event_with_result, settle_usage_with_reconciled_cost,
};
use crate::settlement::settle_usage_after_upsert;
use crate::shutdown::{UsageBackgroundTasks, UsageShutdownState};
use crate::worker::{
build_usage_queue_worker_with_record_gate, UsageWorkerControl, UsageWorkerObservation,
@@ -5156,21 +5154,6 @@ impl UsageRuntime {
where
T: UsageRuntimeAccess,
{
let reconciled = match reconcile_usage_policy_cost_for_event_with_result(data, event).await
{
Ok(reconciled) => reconciled,
Err(err) => {
warn!(
event_name = "usage_event_cost_reconciliation_failed",
log_type = "event",
usage_event_type = ?event.event_type,
request_id = %event.request_id,
error = %err,
"usage runtime failed to reconcile plan cost before direct usage upsert"
);
return false;
}
};
match build_upsert_usage_record_from_event(event) {
Ok(record) => match catch_usage_writer_panic(
"direct usage upsert",
@@ -5179,9 +5162,7 @@ impl UsageRuntime {
.await
{
Ok(Some(stored)) => {
if let Err(err) =
settle_usage_with_reconciled_cost(data, &stored, reconciled).await
{
if let Err(err) = settle_usage_after_upsert(data, &stored, event).await {
warn!(
event_name = "usage_terminal_settlement_failed",
log_type = "event",
crates/aether-usage/runtime/src/settlement.rs
+113 -14
View File
@@ -7,7 +7,7 @@ use aether_data_contracts::repository::settlement::{
};
use aether_data_contracts::repository::usage::PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY;
use aether_data_contracts::repository::usage::{
cancelled_request_fee_is_billable, StoredRequestUsageAudit,
billing_multiplier_snapshot, cancelled_request_fee_is_billable, StoredRequestUsageAudit,
};
use aether_data_contracts::{DataLayerError, DataLayerError::InvalidInput};
use async_trait::async_trait;
@@ -72,16 +72,25 @@ pub(crate) async fn reconcile_usage_policy_cost_for_event_with_result(
return Ok(None);
};
let actual_cost_units = if terminal_state == UsagePolicyCostReservationState::Finalized {
let actual_cost_usd = event.data.actual_total_cost_usd.ok_or_else(|| {
let snapshot = billing_multiplier_snapshot(event.data.request_metadata.as_ref())?;
let cost = if snapshot.is_some() {
event.data.total_cost_usd
} else {
event.data.actual_total_cost_usd
};
let actual_cost_usd = cost.ok_or_else(|| {
InvalidInput(
"completed usage event with a plan reservation token is missing actual cost"
.to_string(),
)
})?;
nonnegative_usd_to_usage_policy_cost_units(finite_cost(actual_cost_usd)?.max(0.0))
.ok_or_else(|| {
InvalidInput("usage policy settlement cost exceeds the supported range".to_string())
})?
let actual_cost_usd = match snapshot {
Some(snapshot) => snapshot.cost(actual_cost_usd)?,
None => finite_cost(actual_cost_usd)?.max(0.0),
};
nonnegative_usd_to_usage_policy_cost_units(actual_cost_usd).ok_or_else(|| {
InvalidInput("usage policy settlement cost exceeds the supported range".to_string())
})?
} else {
0
};
@@ -115,6 +124,34 @@ pub async fn settle_usage_if_needed(
settle_usage_with_reconciled_cost(writer, usage, None).await
}
pub(crate) async fn settle_usage_after_upsert(
writer: &dyn UsageSettlementWriter,
usage: &StoredRequestUsageAudit,
event: &UsageEvent,
) -> Result<(), DataLayerError> {
// Different admissions can share a client request id. Finalize that event's
// own server-issued reservation without borrowing the other admission's rate.
if event_usage_policy_reservation_token(event).is_some()
&& event_usage_policy_reservation_token(event) != usage_policy_reservation_token(usage)
&& !plan_usage_reservation_reconciliation_is_deferred(event.data.request_metadata.as_ref())
{
let billable = event.event_type == UsageEventType::Completed
|| (event.event_type == UsageEventType::Cancelled
&& cancelled_request_fee_is_billable(event.data.request_metadata.as_ref()));
if billable
&& billing_multiplier_snapshot(event.data.request_metadata.as_ref())?.is_none()
&& billing_multiplier_snapshot(usage.request_metadata.as_ref())?.is_some()
{
return Err(InvalidInput(
"colliding usage admission is missing its own billing multiplier snapshot"
.to_string(),
));
}
reconcile_usage_policy_cost_for_event(writer, event).await?;
}
settle_usage_if_needed(writer, usage).await
}
pub(crate) async fn settle_usage_with_reconciled_cost(
writer: &dyn UsageSettlementWriter,
usage: &StoredRequestUsageAudit,
@@ -127,6 +164,11 @@ pub(crate) async fn settle_usage_with_reconciled_cost(
return Ok(());
}
let billing_cost_usd = match billing_multiplier_snapshot(usage.request_metadata.as_ref())? {
Some(snapshot) => snapshot.cost(usage.total_cost_usd)?,
None => finite_cost(usage.actual_total_cost_usd)?.max(0.0),
};
let finalized_at_unix_secs = usage
.finalized_at_unix_secs
.or(Some(usage.updated_at_unix_secs));
@@ -148,14 +190,14 @@ pub(crate) async fn settle_usage_with_reconciled_cost(
{
(
UsagePolicyCostReservationState::Finalized,
nonnegative_usd_to_usage_policy_cost_units(
finite_cost(usage.actual_total_cost_usd)?.max(0.0),
)
.ok_or_else(|| {
InvalidInput(
"usage policy settlement cost exceeds the supported range".to_string(),
)
})?,
nonnegative_usd_to_usage_policy_cost_units(billing_cost_usd).ok_or_else(
|| {
InvalidInput(
"usage policy settlement cost exceeds the supported range"
.to_string(),
)
},
)?,
)
} else {
(UsagePolicyCostReservationState::Released, 0)
@@ -194,6 +236,7 @@ pub(crate) async fn settle_usage_with_reconciled_cost(
billing_status: usage.billing_status.clone(),
total_cost_usd: finite_cost(usage.total_cost_usd)?,
actual_total_cost_usd: finite_cost(usage.actual_total_cost_usd)?,
billing_cost_usd: Some(billing_cost_usd),
finalized_at_unix_secs,
};
let _ = writer.settle_usage(input).await?;
@@ -449,6 +492,62 @@ mod tests {
);
}
#[tokio::test]
async fn composite_billing_rate_charges_customer_without_changing_provider_cost() {
for (group_rate, user_rate, expected_cost) in [(2.0, 0.75, 1.875), (0.0, 3.0, 0.0)] {
let writer = TestSettlementWriter {
has_writer: true,
..Default::default()
};
let mut usage = sample_usage();
let snapshot =
aether_data_contracts::repository::usage::BillingMultiplierSnapshot::from_factors(
std::collections::BTreeMap::from([
("routing_group".to_string(), group_rate),
("user_group".to_string(), user_rate),
]),
)
.unwrap();
usage.request_metadata.as_mut().unwrap()["billing_multiplier_snapshot"] =
json!(snapshot);
settle_usage_if_needed(&writer, &usage).await.unwrap();
let inputs = writer.inputs.lock().unwrap();
assert_eq!(inputs[0].billing_cost_usd, Some(expected_cost));
assert_eq!(inputs[0].total_cost_usd, 1.25);
assert_eq!(inputs[0].actual_total_cost_usd, 0.75);
assert_eq!(
writer.reconciliations.lock().unwrap()[0].actual_cost_units,
(expected_cost * 100_000_000.0).round() as u64
);
}
}
#[tokio::test]
async fn corrupt_or_overflowing_billing_rate_never_changes_wallet_or_reservation() {
for (base, snapshot) in [
(1.25, serde_json::Value::Null),
(
1.25,
json!({"version":1,"factors":{"routing_group":2},"multiplier":1}),
),
(
f64::MAX,
json!({"version":1,"factors":{"routing_group":2},"multiplier":2}),
),
] {
let writer = TestSettlementWriter {
has_writer: true,
..Default::default()
};
let mut usage = sample_usage();
usage.total_cost_usd = base;
usage.request_metadata.as_mut().unwrap()["billing_multiplier_snapshot"] = snapshot;
assert!(settle_usage_if_needed(&writer, &usage).await.is_err());
assert!(writer.inputs.lock().unwrap().is_empty());
assert!(writer.reconciliations.lock().unwrap().is_empty());
}
}
#[tokio::test]
async fn releases_pending_cancelled_usage_without_wallet_settlement() {
let writer = TestSettlementWriter {
crates/aether-usage/runtime/src/settlement_reuse_tests.rs
+256 -16
View File
@@ -66,6 +66,10 @@ impl UsageSettlementWriter for ReuseStore {
&self,
input: ReconcileUsagePolicyCostInput,
) -> Result<Option<StoredUsagePolicyCostReservation>, DataLayerError> {
assert!(
self.upserts.load(Ordering::Relaxed) > 0,
"a durable usage row must exist before a reservation is finalized"
);
input.validate()?;
self.reconciliations.lock().unwrap().push(input.clone());
tokio::task::yield_now().await;
@@ -185,7 +189,7 @@ async fn write(store: &ReuseStore, event: UsageEvent, direct: bool) {
}
#[tokio::test]
async fn worker_and_direct_writes_reuse_confirmed_reservation_and_still_settle_wallet() {
async fn worker_and_direct_writes_persist_before_reconciling_and_settling_wallet() {
for direct in [false, true] {
let store = ReuseStore::default();
write(&store, event(), direct).await;
@@ -198,11 +202,12 @@ async fn worker_and_direct_writes_reuse_confirmed_reservation_and_still_settle_w
assert_eq!(settlements.len(), 1);
assert_eq!(settlements[0].request_id, "req-1");
assert_eq!(settlements[0].actual_total_cost_usd, 0.75);
assert_eq!(settlements[0].billing_cost_usd, Some(0.75));
}
}
#[tokio::test]
async fn missing_or_different_reconciliation_results_keep_stored_usage_reconciliation() {
async fn missing_or_different_reconciliation_results_do_not_repeat_stored_usage_reconciliation() {
let changes: [fn(&mut StoredUsagePolicyCostReservation); 9] = [
|row| row.request_id = "other-request".to_string(),
|row| row.subject_id = "other-user".to_string(),
@@ -223,7 +228,7 @@ async fn missing_or_different_reconciliation_results_keep_stored_usage_reconcili
..Default::default()
};
write(&store, event(), direct).await;
assert_eq!(store.reconciliations.lock().unwrap().len(), 2);
assert_eq!(store.reconciliations.lock().unwrap().len(), 1);
assert_eq!(store.settlements.lock().unwrap().len(), 1);
}
}
@@ -252,19 +257,56 @@ async fn changed_stored_usage_is_reconciled_using_its_own_identity_cost_and_term
};
write(&store, event(), direct).await;
let reconciliations = store.reconciliations.lock().unwrap();
assert_eq!(reconciliations.len(), 2);
assert_eq!(reconciliations[1].request_id, stored.request_id);
let stored_token = stored.request_metadata.as_ref().unwrap()
["plan_usage_reservation_token"]
.as_str()
.unwrap();
assert_eq!(
reconciliations[1].subject_id,
reconciliations.len(),
if stored_token == RESERVATION_TOKEN {
1
} else {
2
}
);
if stored_token != RESERVATION_TOKEN {
assert_eq!(reconciliations[0].reservation_token, RESERVATION_TOKEN);
assert_eq!(reconciliations[0].actual_cost_units, 75_000_000);
}
let reconciliation = reconciliations.last().unwrap();
assert_eq!(reconciliation.request_id, stored.request_id);
assert_eq!(
reconciliation.subject_id,
stored.user_id.as_ref().unwrap().as_str()
);
assert_eq!(
reconciliations[1].reservation_token,
reconciliation.reservation_token,
stored.request_metadata.as_ref().unwrap()["plan_usage_reservation_token"]
.as_str()
.unwrap()
);
assert_ne!(reconciliations[0], reconciliations[1]);
assert_eq!(
reconciliation.actual_cost_units,
if stored.status == "failed" {
0
} else {
(stored.actual_total_cost_usd * 100_000_000.0) as u64
}
);
assert_eq!(
reconciliation.terminal_state,
if stored.status == "failed" {
UsagePolicyCostReservationState::Released
} else {
UsagePolicyCostReservationState::Finalized
}
);
assert_eq!(
reconciliation.finalized_at_unix_secs,
stored
.finalized_at_unix_secs
.unwrap_or(stored.updated_at_unix_secs)
);
let settlements = store.settlements.lock().unwrap();
assert_eq!(settlements.len(), 1);
assert_eq!(
@@ -276,6 +318,185 @@ async fn changed_stored_usage_is_reconciled_using_its_own_identity_cost_and_term
}
}
#[tokio::test]
async fn worker_and_direct_settle_customer_multiplier_snapshot_without_changing_provider_cost() {
for direct in [false, true] {
for (group_multiplier, promotion_multiplier, expected) in [
(0.0, 1.0, 0.0),
(1.0, 1.0, 1.25),
(2.0, 0.25, 0.625),
(3.0, 1.0, 3.75),
] {
let store = ReuseStore::default();
let mut event = event();
event.data.request_metadata.as_mut().unwrap()["billing_multiplier_snapshot"] = json!({
"version": 1,
"factors": {"routing_group": group_multiplier, "promotion": promotion_multiplier},
"multiplier": group_multiplier * promotion_multiplier
});
write(&store, event, direct).await;
let reconciliations = store.reconciliations.lock().unwrap();
assert_eq!(reconciliations.len(), 1, "direct={direct}");
assert_eq!(
reconciliations[0].actual_cost_units,
(expected * 100_000_000.0) as u64
);
let settlements = store.settlements.lock().unwrap();
assert_eq!(settlements.len(), 1);
assert_eq!(settlements[0].billing_cost_usd, Some(expected));
assert_eq!(settlements[0].total_cost_usd, 1.25);
assert_eq!(settlements[0].actual_total_cost_usd, 0.75);
}
}
}
#[tokio::test]
async fn sparse_terminal_event_uses_persisted_multiplier_and_cost_for_both_ledger_and_wallet() {
for direct in [false, true] {
let mut stored = sample_usage();
stored.total_cost_usd = 2.0;
stored.actual_total_cost_usd = 0.25;
stored.request_metadata.as_mut().unwrap()["billing_multiplier_snapshot"] = json!({
"version": 1,
"factors": {"routing_group": 3.0, "promotion": 0.5},
"multiplier": 1.5
});
let expected = stored.billing_cost().unwrap();
assert_eq!(expected, 3.0);
let store = ReuseStore {
stored_override: Some(stored),
..Default::default()
};
// Sparse asynchronous completion has neither the captured multiplier
// nor authoritative charges. Persistence restores the original snapshot.
let mut terminal = event();
terminal.data.total_cost_usd = None;
terminal.data.actual_total_cost_usd = None;
write(&store, terminal, direct).await;
let reconciliations = store.reconciliations.lock().unwrap();
assert_eq!(reconciliations.len(), 1, "direct={direct}");
assert_eq!(reconciliations[0].actual_cost_units, 300_000_000);
let settlements = store.settlements.lock().unwrap();
assert_eq!(settlements.len(), 1);
assert_eq!(settlements[0].billing_cost_usd, Some(expected));
assert_eq!(settlements[0].total_cost_usd, 2.0);
assert_eq!(settlements[0].actual_total_cost_usd, 0.25);
}
}
#[tokio::test]
async fn colliding_request_id_reconciles_new_token_with_its_own_snapshot_after_upsert() {
for direct in [false, true] {
let mut stored = sample_usage();
stored.total_cost_usd = 2.0;
stored.request_metadata = Some(json!({
"plan_usage_reservation_token": "previous-token",
"billing_multiplier_snapshot": {
"version": 1,
"factors": {"routing_group": 0.5},
"multiplier": 0.5
}
}));
let store = ReuseStore {
stored_override: Some(stored),
..Default::default()
};
let mut terminal = event();
terminal.data.request_metadata.as_mut().unwrap()["billing_multiplier_snapshot"] = json!({
"version": 1,
"factors": {"routing_group": 3.0},
"multiplier": 3.0
});
write(&store, terminal, direct).await;
assert_eq!(store.upserts.load(Ordering::Relaxed), 1);
let reconciliations = store.reconciliations.lock().unwrap();
assert_eq!(reconciliations.len(), 2, "direct={direct}");
assert_eq!(reconciliations[0].reservation_token, RESERVATION_TOKEN);
assert_eq!(reconciliations[0].actual_cost_units, 375_000_000);
assert_eq!(reconciliations[1].reservation_token, "previous-token");
assert_eq!(reconciliations[1].actual_cost_units, 100_000_000);
let settlements = store.settlements.lock().unwrap();
assert_eq!(settlements.len(), 1);
assert_eq!(settlements[0].billing_cost_usd, Some(1.0));
assert_eq!(settlements[0].actual_total_cost_usd, 0.75);
}
}
#[tokio::test]
async fn sparse_colliding_token_cannot_borrow_another_requests_multiplier() {
for direct in [false, true] {
for (event_type, billable_cancel) in [
(UsageEventType::Completed, false),
(UsageEventType::Cancelled, true),
] {
let mut stored = sample_usage();
stored.request_metadata = Some(json!({
"plan_usage_reservation_token": "previous-token",
"billing_multiplier_snapshot": {
"version": 1,
"factors": {"routing_group": 0.0},
"multiplier": 0.0
}
}));
let store = ReuseStore {
stored_override: Some(stored),
..Default::default()
};
let mut terminal = event();
terminal.event_type = event_type;
terminal.data.request_metadata.as_mut().unwrap()["cancelled_request_fee"] =
json!(billable_cancel);
if direct {
write(&store, terminal, true).await;
} else {
assert!(write_event_record(&store, &terminal).await.is_err());
}
assert_eq!(store.upserts.load(Ordering::Relaxed), 1);
assert!(
store.reconciliations.lock().unwrap().is_empty(),
"direct={direct}"
);
assert!(store.settlements.lock().unwrap().is_empty());
}
}
}
#[tokio::test]
async fn failed_colliding_token_is_released_without_borrowing_snapshot_or_charge() {
for direct in [false, true] {
let mut stored = sample_usage();
stored.billing_status = "settled".to_string();
stored.request_metadata = Some(json!({
"plan_usage_reservation_token": "previous-token",
"billing_multiplier_snapshot": {
"version": 1,
"factors": {"routing_group": 2.0},
"multiplier": 2.0
}
}));
let store = ReuseStore {
stored_override: Some(stored),
..Default::default()
};
let mut terminal = event();
terminal.event_type = UsageEventType::Failed;
terminal.data.total_cost_usd = None;
terminal.data.actual_total_cost_usd = None;
write(&store, terminal, direct).await;
let reconciliations = store.reconciliations.lock().unwrap();
assert_eq!(reconciliations.len(), 2, "direct={direct}");
assert_eq!(reconciliations[0].reservation_token, RESERVATION_TOKEN);
assert_eq!(
reconciliations[0].terminal_state,
UsagePolicyCostReservationState::Released
);
assert_eq!(reconciliations[0].actual_cost_units, 0);
assert_eq!(reconciliations[1].reservation_token, "previous-token");
assert_eq!(reconciliations[1].actual_cost_units, 250_000_000);
assert!(store.settlements.lock().unwrap().is_empty());
}
}
#[tokio::test]
async fn cancellation_release_billable_cancellation_and_zero_cost_preserve_settlement_rules() {
for direct in [false, true] {
@@ -329,7 +550,7 @@ async fn cancellation_release_billable_cancellation_and_zero_cost_preserve_settl
}
#[tokio::test]
async fn reconciliation_failure_stops_both_writes_before_upsert_and_wallet_settlement() {
async fn reconciliation_failure_keeps_durable_usage_but_stops_wallet_settlement() {
for direct in [false, true] {
let store = ReuseStore {
response: ReconcileResponse::Error,
@@ -341,13 +562,13 @@ async fn reconciliation_failure_stops_both_writes_before_upsert_and_wallet_settl
assert!(write_event_record(&store, &event()).await.is_err());
}
assert_eq!(store.reconciliations.lock().unwrap().len(), 1);
assert_eq!(store.upserts.load(Ordering::Relaxed), 0);
assert_eq!(store.upserts.load(Ordering::Relaxed), 1);
assert!(store.settlements.lock().unwrap().is_empty());
}
}
#[tokio::test]
async fn retry_after_upsert_failure_reconciles_again_before_settling() {
async fn retry_after_upsert_failure_reconciles_only_the_successfully_persisted_usage() {
for direct in [false, true] {
let store = ReuseStore {
fail_next_upsert: AtomicBool::new(true),
@@ -358,9 +579,10 @@ async fn retry_after_upsert_failure_reconciles_again_before_settling() {
} else {
assert!(write_event_record(&store, &event()).await.is_err());
}
assert!(store.reconciliations.lock().unwrap().is_empty());
assert!(store.settlements.lock().unwrap().is_empty());
write(&store, event(), direct).await;
assert_eq!(store.reconciliations.lock().unwrap().len(), 2);
assert_eq!(store.reconciliations.lock().unwrap().len(), 1);
assert_eq!(store.upserts.load(Ordering::Relaxed), 2);
assert_eq!(store.settlements.lock().unwrap().len(), 1);
}
@@ -481,11 +703,17 @@ async fn concurrent_duplicate_delivery_debits_real_memory_wallet_only_once() {
let store = store.clone();
let runtime = runtime.clone();
tasks.spawn(async move {
let mut terminal = event();
terminal.data.request_metadata.as_mut().unwrap()["billing_multiplier_snapshot"] = json!({
"version": 1,
"factors": {"routing_group": 3.0, "promotion": 0.5},
"multiplier": 1.5
});
if index % 2 == 0 {
write_event_record(store.as_ref(), &event()).await.unwrap();
write_event_record(store.as_ref(), &terminal).await.unwrap();
} else {
runtime
.record_terminal_event_direct(store.as_ref(), event())
.record_terminal_event_direct(store.as_ref(), terminal)
.await;
}
});
@@ -495,13 +723,25 @@ async fn concurrent_duplicate_delivery_debits_real_memory_wallet_only_once() {
}
assert_eq!(store.reconciliations.lock().unwrap().len(), 32);
assert_eq!(store.settlements.lock().unwrap().len(), 32);
assert!(store
.reconciliations
.lock()
.unwrap()
.iter()
.all(|input| input.actual_cost_units == 187_500_000));
assert!(store
.settlements
.lock()
.unwrap()
.iter()
.all(|input| input.billing_cost_usd == Some(1.875) && input.actual_total_cost_usd == 0.75));
let wallet = wallets
.find(WalletLookupKey::UserId("user-1"))
.await
.unwrap()
.unwrap();
assert_eq!(wallet.balance + wallet.gift_balance, 11.25);
assert_eq!(wallet.total_consumed, 0.75);
assert_eq!(wallet.balance + wallet.gift_balance, 10.125);
assert_eq!(wallet.total_consumed, 1.875);
assert!(matches!(
store
.repository
crates/aether-usage/runtime/src/worker.rs
+66 -52
View File
@@ -16,9 +16,7 @@ use crate::queue::UsageDeadLetterOutcome;
use crate::runtime::{
UsageBillingEventEnricher, UsageRuntimeAccess, UsageWorkerRecordConcurrencyGate,
};
use crate::settlement::{
reconcile_usage_policy_cost_for_event_with_result, settle_usage_with_reconciled_cost,
};
use crate::settlement::settle_usage_after_upsert;
use crate::{
build_upsert_usage_record_from_event, UsageEvent, UsageEventType, UsageQueue,
UsageRuntimeConfig, UsageSettlementWriter,
@@ -796,10 +794,11 @@ pub async fn write_event_record<T>(data: &T, event: &UsageEvent) -> Result<(), D
where
T: UsageRecordWriter + UsageSettlementWriter + Send + Sync,
{
let reconciled = reconcile_usage_policy_cost_for_event_with_result(data, event).await?;
let record = build_upsert_usage_record_from_event(event)?;
if let Some(stored) = data.upsert_usage_record(record).await? {
settle_usage_with_reconciled_cost(data, &stored, reconciled).await?;
// Sparse terminal events may omit pricing factors. Only the stored request
// snapshot is authoritative before finalizing an immutable cost reservation.
settle_usage_after_upsert(data, &stored, event).await?;
}
// Manual proxy traffic is counted at the actual transport-attempt boundary. Usage events are
// replayable, so emitting that side effect here would count normal requests and reclaims twice.
@@ -1240,49 +1239,49 @@ mod tests {
.lock()
.expect("records lock")
.push(record.clone());
Ok(Some(
StoredRequestUsageAudit::new(
"usage-1".to_string(),
record.request_id,
record.user_id,
record.api_key_id,
record.username,
record.api_key_name,
record.provider_name,
record.model,
record.target_model,
record.provider_id,
record.provider_endpoint_id,
record.provider_api_key_id,
record.request_type,
record.api_format,
record.api_family,
record.endpoint_kind,
record.endpoint_api_format,
record.provider_api_family,
record.provider_endpoint_kind,
record.has_format_conversion.unwrap_or(false),
record.is_stream.unwrap_or(false),
record.input_tokens.unwrap_or_default() as i32,
record.output_tokens.unwrap_or_default() as i32,
record.total_tokens.unwrap_or_default() as i32,
record.total_cost_usd.unwrap_or_default(),
record.actual_total_cost_usd.unwrap_or_default(),
record.status_code.map(i32::from),
record.error_message,
record.error_category,
record.response_time_ms.map(|value| value as i32),
record.first_byte_time_ms.map(|value| value as i32),
record.status,
record.billing_status,
record
.created_at_unix_ms
.unwrap_or(record.updated_at_unix_secs) as i64,
record.updated_at_unix_secs as i64,
record.finalized_at_unix_secs.map(|value| value as i64),
)
.expect("stored usage should build"),
))
let mut stored = StoredRequestUsageAudit::new(
"usage-1".to_string(),
record.request_id,
record.user_id,
record.api_key_id,
record.username,
record.api_key_name,
record.provider_name,
record.model,
record.target_model,
record.provider_id,
record.provider_endpoint_id,
record.provider_api_key_id,
record.request_type,
record.api_format,
record.api_family,
record.endpoint_kind,
record.endpoint_api_format,
record.provider_api_family,
record.provider_endpoint_kind,
record.has_format_conversion.unwrap_or(false),
record.is_stream.unwrap_or(false),
record.input_tokens.unwrap_or_default() as i32,
record.output_tokens.unwrap_or_default() as i32,
record.total_tokens.unwrap_or_default() as i32,
record.total_cost_usd.unwrap_or_default(),
record.actual_total_cost_usd.unwrap_or_default(),
record.status_code.map(i32::from),
record.error_message,
record.error_category,
record.response_time_ms.map(|value| value as i32),
record.first_byte_time_ms.map(|value| value as i32),
record.status,
record.billing_status,
record
.created_at_unix_ms
.unwrap_or(record.updated_at_unix_secs) as i64,
record.updated_at_unix_secs as i64,
record.finalized_at_unix_secs.map(|value| value as i64),
)
.expect("stored usage should build");
stored.request_metadata = record.request_metadata;
Ok(Some(stored))
}
}
@@ -1510,7 +1509,7 @@ mod tests {
}
#[tokio::test]
async fn same_request_id_terminal_events_reconcile_each_reservation_token_before_upsert() {
async fn same_request_id_terminal_events_reconcile_each_persisted_reservation_token() {
let store = TestUsageStore::default();
let mut first = sample_event();
first.request_id = "shared-client-trace".to_string();
@@ -1619,7 +1618,12 @@ mod tests {
worker.queue.ensure_consumer_group().await.expect("group");
let mut event = sample_event();
event.data.request_metadata = Some(serde_json::json!({
"plan_usage_reservation_token": "pricing-retry-reservation"
"plan_usage_reservation_token": "pricing-retry-reservation",
"billing_multiplier_snapshot": {
"version": 1,
"factors": {"routing_group": 3.0, "promotion": 0.5},
"multiplier": 1.5
}
}));
worker.queue.enqueue(&event).await.expect("enqueue");
let batch = worker
@@ -1676,17 +1680,27 @@ mod tests {
assert_eq!(records[0].total_cost_usd, Some(0.456));
assert_eq!(records[0].actual_total_cost_usd, Some(0.123));
assert_eq!(records[0].total_tokens, Some(10));
assert_eq!(
records[0].request_metadata.as_ref().unwrap()["billing_multiplier_snapshot"]
["multiplier"],
1.5
);
}
{
let reconciliations = store.reconciliations.lock().expect("reconciliations lock");
assert_eq!(reconciliations.len(), 1);
assert_eq!(reconciliations[0].actual_cost_units, 12_300_000);
assert_eq!(reconciliations[0].actual_cost_units, 68_400_000);
assert_eq!(
reconciliations[0].reservation_token,
"pricing-retry-reservation"
);
}
assert_eq!(store.settlements.lock().expect("settlements lock").len(), 1);
{
let settlements = store.settlements.lock().expect("settlements lock");
assert_eq!(settlements.len(), 1);
assert_eq!(settlements[0].billing_cost_usd, Some(0.456 * 1.5));
assert_eq!(settlements[0].actual_total_cost_usd, 0.123);
}
assert_eq!(
store.enrich_calls.lock().expect("enrich calls lock").len(),
2