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

feat: allow wallet fallback after plan quota exhaustion

Browse Source
This commit is contained in:
elky
2026-10-09 21:21:53 +08:00
parent bad13237d2
commit 63520c8c27
39 changed files with 1474 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/adapters/postgres/migrations/20261009000000_user_preferences_allow_wallet_overage.sql
+3
View File
@@ -0,0 +1,3 @@
-- Wallet fallback is opt-in for both existing and newly created preferences.
ALTER TABLE public.user_preferences
ADD COLUMN IF NOT EXISTS allow_wallet_overage boolean NOT NULL DEFAULT false;
crates/aether-data/adapters/postgres/src/billing.rs
+9 -26
View File
@@ -1193,10 +1193,8 @@ WHERE id = $1
r#"
SELECT
user_plan_entitlements.id,
user_plan_entitlements.entitlements_snapshot,
billing_plans.entitlements_json AS plan_entitlements_json
user_plan_entitlements.entitlements_snapshot
FROM user_plan_entitlements
JOIN billing_plans ON billing_plans.id = user_plan_entitlements.plan_id
WHERE user_plan_entitlements.user_id = $1
AND user_plan_entitlements.status = 'active'
AND user_plan_entitlements.starts_at <= NOW()
@@ -1216,12 +1214,9 @@ ORDER BY user_plan_entitlements.expires_at ASC,
let entitlement_id: String = row.try_get("id").map_postgres_err()?;
let entitlements: serde_json::Value =
row.try_get("entitlements_snapshot").map_postgres_err()?;
let plan_entitlements: serde_json::Value =
row.try_get("plan_entitlements_json").map_postgres_err()?;
grants.extend(daily_quota_grants_from_entitlement(
&entitlement_id,
&entitlements,
daily_quota_wallet_overage_policy(&plan_entitlements),
now,
)?);
}
@@ -1229,9 +1224,15 @@ ORDER BY user_plan_entitlements.expires_at ASC,
let mut total_quota_usd = 0.0;
let mut used_usd = 0.0;
let mut remaining_usd = 0.0;
let mut allow_wallet_overage = true;
let allow_wallet_overage = sqlx::query_scalar::<_, bool>(
"SELECT allow_wallet_overage FROM user_preferences WHERE user_id = $1",
)
.bind(user_id)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?
.unwrap_or(false);
for grant in &grants {
allow_wallet_overage &= grant.allow_wallet_overage;
let used = sqlx::query_scalar::<_, Option<f64>>(
r#"
SELECT CAST(COALESCE(SUM(amount_usd), 0) AS DOUBLE PRECISION)
@@ -1336,7 +1337,6 @@ struct DailyQuotaGrant {
entitlement_id: String,
daily_quota_usd: f64,
usage_date: String,
allow_wallet_overage: bool,
}
fn daily_quota_usage_date(
@@ -1355,7 +1355,6 @@ fn daily_quota_usage_date(
fn daily_quota_grants_from_entitlement(
entitlement_id: &str,
entitlements: &serde_json::Value,
current_allow_wallet_overage: Option<bool>,
now: chrono::DateTime<chrono::Utc>,
) -> Result<Vec<DailyQuotaGrant>, DataLayerError> {
let mut grants = Vec::new();
@@ -1381,27 +1380,11 @@ fn daily_quota_grants_from_entitlement(
.and_then(serde_json::Value::as_str),
now,
)?,
allow_wallet_overage: current_allow_wallet_overage.unwrap_or_else(|| {
item.get("allow_wallet_overage")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
}),
});
}
Ok(grants)
}
fn daily_quota_wallet_overage_policy(entitlements: &serde_json::Value) -> Option<bool> {
entitlements.as_array()?.iter().find_map(|item| {
(item.get("type").and_then(serde_json::Value::as_str) == Some("daily_quota"))
.then(|| {
item.get("allow_wallet_overage")
.and_then(serde_json::Value::as_bool)
})
.flatten()
})
}
fn map_payment_gateway_config_row(
row: &sqlx::postgres::PgRow,
) -> Result<PaymentGatewayConfigRecord, DataLayerError> {
crates/aether-data/adapters/postgres/src/settlement.rs
+93 -55
View File
@@ -18,6 +18,7 @@ use crate::PostgresTransactionRunner;
const FIND_USAGE_FOR_SETTLEMENT_SQL: &str = r#"
SELECT
usage_record.request_id,
usage_record.request_metadata,
COALESCE(usage_settlement_snapshots.wallet_id, usage_record.wallet_id) AS wallet_id,
COALESCE(usage_settlement_snapshots.billing_status, usage_record.billing_status) AS billing_status,
COALESCE(
@@ -465,7 +466,6 @@ struct DailyQuotaGrant {
entitlement_id: String,
daily_quota_usd: f64,
usage_date: String,
allow_wallet_overage: bool,
}
fn daily_quota_usage_date(
@@ -484,7 +484,6 @@ fn daily_quota_usage_date(
fn daily_quota_grants_from_entitlement(
entitlement_id: &str,
entitlements: &serde_json::Value,
current_allow_wallet_overage: Option<bool>,
now: chrono::DateTime<chrono::Utc>,
) -> Result<Vec<DailyQuotaGrant>, DataLayerError> {
let mut grants = Vec::new();
@@ -511,27 +510,11 @@ fn daily_quota_grants_from_entitlement(
entitlement_id: entitlement_id.to_string(),
daily_quota_usd,
usage_date,
allow_wallet_overage: current_allow_wallet_overage.unwrap_or_else(|| {
item.get("allow_wallet_overage")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
}),
});
}
Ok(grants)
}
fn daily_quota_wallet_overage_policy(entitlements: &serde_json::Value) -> Option<bool> {
entitlements.as_array()?.iter().find_map(|item| {
(item.get("type").and_then(serde_json::Value::as_str) == Some("daily_quota"))
.then(|| {
item.get("allow_wallet_overage")
.and_then(serde_json::Value::as_bool)
})
.flatten()
})
}
async fn consume_daily_quota_postgres(
tx: &mut crate::PostgresTransaction,
user_id: &str,
@@ -549,16 +532,13 @@ async fn consume_daily_quota_postgres(
return Ok(DailyQuotaDebitResult::default());
}
let now = chrono::Utc::now();
// Serialize each entitlement's debits. Read the shared plan's current overage policy
// from this statement's snapshot without locking every subscriber's plan row.
// Serialize each entitlement's debits without locking the shared plan row.
let entitlement_rows = sqlx::query(
r#"
SELECT
user_plan_entitlements.id,
user_plan_entitlements.entitlements_snapshot,
billing_plans.entitlements_json AS plan_entitlements_json
user_plan_entitlements.entitlements_snapshot
FROM user_plan_entitlements
JOIN billing_plans ON billing_plans.id = user_plan_entitlements.plan_id
WHERE user_plan_entitlements.user_id = $1
AND user_plan_entitlements.status = 'active'
AND user_plan_entitlements.starts_at <= NOW()
@@ -578,12 +558,9 @@ FOR UPDATE OF user_plan_entitlements
let entitlement_id: String = row.try_get("id").map_postgres_err()?;
let entitlements: serde_json::Value =
row.try_get("entitlements_snapshot").map_postgres_err()?;
let plan_entitlements: serde_json::Value =
row.try_get("plan_entitlements_json").map_postgres_err()?;
grants.extend(daily_quota_grants_from_entitlement(
&entitlement_id,
&entitlements,
daily_quota_wallet_overage_policy(&plan_entitlements),
now,
)?);
}
@@ -593,9 +570,16 @@ FOR UPDATE OF user_plan_entitlements
let mut grants_with_remaining = Vec::new();
let mut total_remaining = 0.0;
let mut allow_wallet_overage = true;
// Wallet fallback is the user's explicit choice, never the plan's legacy default.
let allow_wallet_overage = sqlx::query_scalar::<_, bool>(
"SELECT allow_wallet_overage FROM user_preferences WHERE user_id = $1",
)
.bind(user_id)
.fetch_optional(&mut **tx)
.await
.map_postgres_err()?
.unwrap_or(false);
for grant in grants {
allow_wallet_overage &= grant.allow_wallet_overage;
let used = sqlx::query_scalar::<_, Option<f64>>(
r#"
SELECT CAST(COALESCE(SUM(amount_usd), 0) AS DOUBLE PRECISION)
@@ -1293,33 +1277,43 @@ LIMIT 1
}
let billable_cost_usd = settlement_billable_cost_usd(&input);
let wallet_debit_cost_usd = if !api_key_is_standalone {
if let Some(user_id) =
input.user_id.as_deref().filter(|value| !value.is_empty())
{
let quota = consume_daily_quota_postgres(
tx,
user_id,
&input.request_id,
billable_cost_usd,
wallet_available_usd,
wallet_can_overdraft,
)
.await?;
if quota.insufficient {
final_billing_status = "insufficient_quota".to_string();
settlement.billing_status = final_billing_status.clone();
0.0
let plan_wallet_fallback = usage_row
.try_get::<Option<serde_json::Value>, _>("request_metadata")
.map_postgres_err()?
.and_then(|metadata| {
metadata
.get("plan_wallet_fallback")
.and_then(serde_json::Value::as_bool)
})
.unwrap_or(false);
let wallet_debit_cost_usd =
if !api_key_is_standalone && !plan_wallet_fallback {
if let Some(user_id) =
input.user_id.as_deref().filter(|value| !value.is_empty())
{
let quota = consume_daily_quota_postgres(
tx,
user_id,
&input.request_id,
billable_cost_usd,
wallet_available_usd,
wallet_can_overdraft,
)
.await?;
if quota.insufficient {
final_billing_status = "insufficient_quota".to_string();
settlement.billing_status = final_billing_status.clone();
0.0
} else {
quota_covered = quota.debited_usd;
(billable_cost_usd - quota.debited_usd).max(0.0)
}
} else {
quota_covered = quota.debited_usd;
(billable_cost_usd - quota.debited_usd).max(0.0)
billable_cost_usd
}
} else {
billable_cost_usd
}
} else {
billable_cost_usd
};
};
if final_billing_status != "settled" {
sync_usage_settlement_snapshot(&mut **tx, &settlement).await?;
sqlx::query(FINALIZE_USAGE_BILLING_SQL)
@@ -1496,6 +1490,7 @@ mod tests {
"user_plan_entitlements",
"entitlement_usage_ledgers",
"users",
"user_preferences",
"usage_request_admissions",
"usage_cost_reservations",
] {
@@ -1524,17 +1519,20 @@ mod tests {
for (scenario, charge, quota_covered) in [
("wallet", 20.0, 0.0),
("quota_and_wallet", 20.0, 7.0),
("wallet_fallback", 20.0, 0.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 user_preferences (id, user_id, allow_wallet_overage) VALUES ($1, $1, true)")
.bind(scenario).execute(&pool).await.expect("wallet fallback preference 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,
"reset_timezone": "UTC", "allow_wallet_overage": false,
}]);
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");
@@ -1542,7 +1540,7 @@ mod tests {
.bind(scenario).bind(&grant).execute(&pool).await.expect("entitlement should insert");
}
let multiplier = charge / 10.0;
let metadata = serde_json::json!({"billing_multiplier_snapshot": {
let metadata = serde_json::json!({"plan_wallet_fallback": scenario == "wallet_fallback", "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)")
@@ -1701,6 +1699,8 @@ mod tests {
.await
.expect("shared plan should insert");
for user_id in ["user-a", "user-b"] {
sqlx::query("INSERT INTO users (id, username, email_verified) VALUES ($1, $1, false)")
.bind(user_id).execute(&pool).await.expect("user 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, 'shared-plan', $1, NOW() - INTERVAL '1 hour', NOW() + INTERVAL '1 day', $2, NOW(), NOW())",
)
@@ -1786,14 +1786,52 @@ mod tests {
held.rollback().await.expect("held quota debit should roll back");
let mut after_edit = pool.begin().await.expect("fresh transaction should start");
let unchanged_policy = super::consume_daily_quota_postgres(
&mut after_edit, "user-a", "request-plan-policy-after", 2.0, Some(5.0), true,
)
.await
.expect("plan changes must not opt the user into wallet charges");
assert!(unchanged_policy.insufficient);
after_edit.rollback().await.expect("default policy verification should roll back");
sqlx::query("INSERT INTO user_preferences (id, user_id, allow_wallet_overage) VALUES ('user-a', 'user-a', true)")
.execute(&pool).await.expect("user should enable wallet fallback");
let mut after_edit = pool.begin().await.expect("user preference transaction should start");
let updated_policy = super::consume_daily_quota_postgres(
&mut after_edit, "user-a", "request-policy-after", 2.0, Some(5.0), true,
)
.await
.expect("fresh quota read should use current plan configuration");
.expect("fresh quota read should use the user preference");
assert!(!updated_policy.insufficient);
assert_eq!(updated_policy.debited_usd, 1.0);
after_edit.rollback().await.expect("policy verification should roll back");
after_edit.commit().await.expect("enabled wallet fallback should commit");
let mut exhausted = pool.begin().await.expect("exhausted quota transaction should start");
let debit = super::consume_daily_quota_postgres(
&mut exhausted, "user-a", "request-exhausted", 2.0, Some(5.0), false,
).await.expect("exhausted quota should fall back to wallet");
assert!(!debit.insufficient);
assert_eq!(debit.debited_usd, 0.0);
let debit = super::consume_daily_quota_postgres(
&mut exhausted, "user-a", "request-exhausted-no-wallet", 2.0, Some(0.0), false,
).await.expect("empty wallet should reject exhausted quota fallback");
assert!(debit.insufficient);
exhausted.rollback().await.expect("exhausted quota verification should roll back");
use aether_data_contracts::repository::billing::BillingReadRepository;
let billing = crate::SqlxBillingReadRepository::new(pool.clone());
let quota = billing.find_user_daily_quota_availability("user-a").await.unwrap().unwrap();
assert!(quota.allow_wallet_overage);
assert_eq!(quota.remaining_usd, 0.0);
sqlx::query("UPDATE user_preferences SET allow_wallet_overage = false WHERE user_id = 'user-a'")
.execute(&pool).await.expect("user should disable wallet fallback");
let quota = billing.find_user_daily_quota_availability("user-a").await.unwrap().unwrap();
assert!(!quota.allow_wallet_overage);
let mut disabled = pool.begin().await.expect("disabled fallback transaction should start");
let debit = super::consume_daily_quota_postgres(
&mut disabled, "user-a", "request-disabled", 2.0, Some(5.0), true,
).await.expect("disabled preference should prevent wallet charges");
assert!(debit.insufficient);
assert_eq!(debit.debited_usd, 0.0);
disabled.rollback().await.expect("disabled fallback verification should roll back");
})
.catch_unwind()
.await;
crates/aether-data/adapters/postgres/src/usage/analytics_tests.rs
+2
View File
@@ -665,6 +665,8 @@ async fn live_overview_settlement_allocations_preserve_unlimited_and_finite_wall
_ => 0.0,
};
if quota > 0.0 {
sqlx::query("INSERT INTO user_preferences(id,user_id,allow_wallet_overage) VALUES($1,$1,true)")
.bind(&user).execute(&pool).await.unwrap();
let grant = serde_json::json!([{"type":"daily_quota","daily_quota_usd":quota,"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,'test',1,'month',1,$2,NOW(),NOW())").bind(&user).bind(&grant).execute(&pool).await.unwrap();
sqlx::query("INSERT INTO payment_orders(id,order_no,wallet_id,user_id,amount_usd,payment_method,created_at) VALUES($1,$1,$2,$1,1,'test',NOW())").bind(&user).bind(&wallet).execute(&pool).await.unwrap();
crates/aether-data/adapters/postgres/src/users.rs
+11 -3
View File
@@ -524,7 +524,8 @@ SELECT
up.timezone,
up.email_notifications,
up.usage_alerts,
up.announcement_notifications
up.announcement_notifications,
up.allow_wallet_overage
FROM user_preferences up
LEFT JOIN providers p
ON p.id = up.default_provider_id
@@ -546,6 +547,7 @@ WITH upserted AS (
email_notifications,
usage_alerts,
announcement_notifications,
allow_wallet_overage,
created_at,
updated_at
) VALUES (
@@ -560,6 +562,7 @@ WITH upserted AS (
$9,
$10,
$11,
$12,
NOW(),
NOW()
)
@@ -573,6 +576,7 @@ WITH upserted AS (
email_notifications = EXCLUDED.email_notifications,
usage_alerts = EXCLUDED.usage_alerts,
announcement_notifications = EXCLUDED.announcement_notifications,
allow_wallet_overage = EXCLUDED.allow_wallet_overage,
updated_at = NOW()
RETURNING
user_id,
@@ -584,7 +588,8 @@ WITH upserted AS (
timezone,
email_notifications,
usage_alerts,
announcement_notifications
announcement_notifications,
allow_wallet_overage
)
SELECT
upserted.user_id,
@@ -597,7 +602,8 @@ SELECT
upserted.timezone,
upserted.email_notifications,
upserted.usage_alerts,
upserted.announcement_notifications
upserted.announcement_notifications,
upserted.allow_wallet_overage
FROM upserted
LEFT JOIN providers p
ON p.id = upserted.default_provider_id
@@ -2931,6 +2937,7 @@ LIMIT 1
.bind(preferences.email_notifications)
.bind(preferences.usage_alerts)
.bind(preferences.announcement_notifications)
.bind(preferences.allow_wallet_overage)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
@@ -3220,6 +3227,7 @@ fn map_user_preference_row(
announcement_notifications: row
.try_get("announcement_notifications")
.map_postgres_err()?,
allow_wallet_overage: row.try_get("allow_wallet_overage").map_postgres_err()?,
})
}
crates/aether-data/contracts/src/repository/usage/metadata_policy.rs
+46 -2
View File
@@ -82,13 +82,36 @@ pub fn preserve_usage_routing_group_snapshot(
captured,
);
}
let Some(Value::Object(snapshot)) = sanitize_usage_request_metadata_object(&snapshot) else {
let snapshot = sanitize_usage_request_metadata_object(&snapshot)
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
if snapshot.is_empty()
&& previous
.get("plan_wallet_fallback")
.and_then(Value::as_bool)
!= Some(true)
{
return incoming;
};
}
let mut metadata = incoming
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
metadata.extend(snapshot);
if previous
.get("plan_wallet_fallback")
.and_then(Value::as_bool)
== Some(true)
{
metadata.insert("plan_wallet_fallback".to_string(), Value::Bool(true));
}
if metadata
.get("plan_wallet_fallback")
.and_then(Value::as_bool)
== Some(true)
{
metadata.remove(PLAN_USAGE_RESERVATION_TOKEN_KEY);
metadata.remove(PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY);
}
Some(Value::Object(metadata))
}
@@ -152,6 +175,7 @@ pub fn sanitize_usage_request_metadata_object(source: &Map<String, Value>) -> Op
"client_requested_stream",
UPSTREAM_IS_STREAM_KEY,
"api_key_is_standalone",
"plan_wallet_fallback",
WEBSOCKET_MODE_METADATA_KEY,
PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY,
"transport_error",
@@ -1369,6 +1393,26 @@ mod tests {
sanitize_usage_request_metadata, sanitize_usage_request_metadata_ref,
};
#[test]
fn wallet_fallback_survives_sparse_video_completion_and_removes_plan_reservation() {
let token = "550e8400-e29b-41d4-a716-446655440000";
let previous = json!({"plan_wallet_fallback": true, "plan_usage_reservation_token": token});
let preserved = preserve_usage_routing_group_snapshot(None, Some(&previous)).unwrap();
assert_eq!(preserved["plan_wallet_fallback"], true);
assert!(preserved.get("plan_usage_reservation_token").is_none());
let incoming = json!({"plan_wallet_fallback": true});
let previous =
json!({"plan_wallet_fallback": false, "plan_usage_reservation_token": token});
let preserved =
preserve_usage_routing_group_snapshot(Some(incoming), Some(&previous)).unwrap();
assert_eq!(preserved["plan_wallet_fallback"], true);
assert!(preserved.get("plan_usage_reservation_token").is_none());
assert!(
sanitize_usage_request_metadata(Some(json!({"plan_wallet_fallback": "true"})))
.is_none()
);
}
#[test]
fn billing_multiplier_snapshot_projection_preserves_invalid_marker_and_immutable_factors() {
for snapshot in [
crates/aether-data/contracts/src/repository/users.rs
+12
View File
@@ -776,6 +776,8 @@ pub struct StoredUserPreferenceRecord {
pub email_notifications: bool,
pub usage_alerts: bool,
pub announcement_notifications: bool,
#[serde(default)]
pub allow_wallet_overage: bool,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
@@ -925,6 +927,7 @@ impl StoredUserPreferenceRecord {
email_notifications: true,
usage_alerts: true,
announcement_notifications: true,
allow_wallet_overage: false,
}
}
}
@@ -1967,5 +1970,14 @@ mod tests {
assert!(record.email_notifications);
assert!(record.usage_alerts);
assert!(record.announcement_notifications);
assert!(!record.allow_wallet_overage);
let mut legacy_json = serde_json::to_value(&record).expect("preferences should serialize");
legacy_json
.as_object_mut()
.expect("preferences object")
.remove("allow_wallet_overage");
let restored: StoredUserPreferenceRecord =
serde_json::from_value(legacy_json).expect("legacy preferences should deserialize");
assert!(!restored.allow_wallet_overage);
}
}
crates/aether-data/runtime/schema/bootstrap/postgres/001_types_and_tables.sql
+1
View File
@@ -1349,6 +1349,7 @@ CREATE TABLE IF NOT EXISTS public.user_preferences (
email_notifications boolean DEFAULT true,
usage_alerts boolean DEFAULT true,
announcement_notifications boolean DEFAULT true,
allow_wallet_overage boolean DEFAULT false NOT NULL,
created_at timestamp with time zone DEFAULT now() NOT NULL,
updated_at timestamp with time zone DEFAULT now() NOT NULL
);
crates/aether-data/runtime/schema/generated/postgres/baseline/001_identity.sql
+1
View File
@@ -191,6 +191,7 @@ CREATE TABLE IF NOT EXISTS public.user_preferences (
email_notifications boolean DEFAULT true NOT NULL,
usage_alerts boolean DEFAULT true NOT NULL,
announcement_notifications boolean DEFAULT true NOT NULL,
allow_wallet_overage boolean DEFAULT false NOT NULL,
created_at bigint NOT NULL,
updated_at bigint NOT NULL
);
crates/aether-data/runtime/schema/logical/001_identity.toml
+5
View File
@@ -788,6 +788,11 @@ name = "announcement_notifications"
type = "bool"
default = true
[[table.user_preferences.columns]]
name = "allow_wallet_overage"
type = "bool"
default = false
[[table.user_preferences.columns]]
name = "created_at"
type = "unix_seconds"
crates/aether-data/runtime/src/lifecycle/migrate/tests.rs
+2
View File
@@ -37,6 +37,7 @@ mod overview_migration_safety;
mod policy_nulls;
mod provider_expenses;
mod scoped_provider_creation;
mod user_preferences;
/// A clean PostgreSQL database is bootstrapped from the schema snapshot first;
/// migrations after the privacy/security frontier are intentionally left
@@ -1599,6 +1600,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20261001000000,
20261004000000,
20261007000000,
20261009000000,
]
);
}
crates/aether-data/runtime/src/lifecycle/migrate/tests/legacy_overview_upgrade.rs
+1
View File
@@ -116,6 +116,7 @@ WHERE version=20260919000000;
20261001000000,
20261004000000,
20261007000000,
20261009000000,
]
);
assert_eq!(
crates/aether-data/runtime/src/lifecycle/migrate/tests/user_preferences.rs
+93
View File
@@ -0,0 +1,93 @@
use aether_data_contracts::repository::users::StoredUserPreferenceRecord;
use aether_data_postgres::SqlxUserReadRepository;
use sqlx::{query, query_scalar, PgPool};
use super::{ManagedPostgresServer, POSTGRES_MIGRATOR};
#[tokio::test]
async fn wallet_overage_preference_upgrade_defaults_off_and_round_trips() {
let Some(server) = ManagedPostgresServer::try_start()
.await
.expect("postgres should start")
else {
return;
};
let pool = PgPool::connect(server.database_url())
.await
.expect("database should connect");
sqlx::raw_sql(
r#"
CREATE TABLE providers (id text PRIMARY KEY, name text NOT NULL);
CREATE TABLE user_preferences (
id text PRIMARY KEY, user_id text UNIQUE NOT NULL,
avatar_url text, bio text, default_provider_id text,
theme text NOT NULL DEFAULT 'light', language text NOT NULL DEFAULT 'zh-CN',
timezone text NOT NULL DEFAULT 'Asia/Shanghai',
email_notifications boolean NOT NULL DEFAULT true,
usage_alerts boolean NOT NULL DEFAULT true,
announcement_notifications boolean NOT NULL DEFAULT true,
created_at timestamptz DEFAULT now(), updated_at timestamptz DEFAULT now()
);
INSERT INTO user_preferences (id, user_id, theme) VALUES ('legacy-pref', 'legacy-user', 'dark');
"#,
)
.execute(&pool)
.await
.expect("legacy preferences should seed");
let migration = POSTGRES_MIGRATOR
.iter()
.find(|migration| migration.version == 20261009000000)
.expect("wallet overage migration should be embedded");
for _ in 0..2 {
sqlx::raw_sql(&migration.sql)
.execute(&pool)
.await
.expect("migration should apply idempotently");
}
let repository = SqlxUserReadRepository::new(pool.clone());
let mut preferences = repository
.read_user_preferences("legacy-user")
.await
.expect("legacy preferences should read")
.expect("legacy preferences should exist");
assert!(!preferences.allow_wallet_overage);
assert_eq!(preferences.theme, "dark");
for enabled in [true, false] {
preferences.allow_wallet_overage = enabled;
assert_eq!(
repository
.write_user_preferences(&preferences)
.await
.expect("preferences should update"),
Some(preferences.clone())
);
assert_eq!(
repository
.read_user_preferences("legacy-user")
.await
.expect("preferences should read"),
Some(preferences.clone())
);
}
query("INSERT INTO user_preferences (id, user_id) VALUES ('new-pref', 'new-user')")
.execute(&pool)
.await
.expect("new preferences should insert");
let default: bool = query_scalar(
"SELECT allow_wallet_overage FROM user_preferences WHERE user_id = 'new-user'",
)
.fetch_one(&pool)
.await
.expect("new default should read");
assert!(!default);
let mut enabled = StoredUserPreferenceRecord::default_for_user("enabled-user");
enabled.allow_wallet_overage = true;
assert_eq!(
repository
.write_user_preferences(&enabled)
.await
.expect("new enabled preferences should insert"),
Some(enabled)
);
pool.close().await;
}
crates/aether-data/runtime/src/repository/billing/memory.rs
+2 -27
View File
@@ -92,14 +92,12 @@ fn billing_plan_from_input(
fn daily_quota_availability_from_entitlements(
entitlements: impl IntoIterator<Item = UserPlanEntitlementRecord>,
billing_plans: &BTreeMap<String, BillingPlanRecord>,
now: u64,
) -> UserDailyQuotaAvailabilityRecord {
let mut has_active_daily_quota = false;
let mut total_quota_usd = 0.0;
let used_usd = 0.0;
let mut remaining_usd = 0.0;
let mut allow_wallet_overage = true;
for entitlement in entitlements {
if entitlement.status != "active"
|| entitlement.starts_at_unix_secs > now
@@ -110,9 +108,6 @@ fn daily_quota_availability_from_entitlements(
let Some(items) = entitlement.entitlements_snapshot.as_array() else {
continue;
};
let current_allow_wallet_overage = billing_plans
.get(&entitlement.plan_id)
.and_then(|plan| daily_quota_wallet_overage_policy(&plan.entitlements_json));
for item in items {
if item.get("type").and_then(serde_json::Value::as_str) != Some("daily_quota") {
continue;
@@ -127,11 +122,6 @@ fn daily_quota_availability_from_entitlements(
has_active_daily_quota = true;
total_quota_usd += daily_quota_usd;
remaining_usd += daily_quota_usd;
allow_wallet_overage &= current_allow_wallet_overage.unwrap_or_else(|| {
item.get("allow_wallet_overage")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
});
}
}
UserDailyQuotaAvailabilityRecord {
@@ -139,21 +129,11 @@ fn daily_quota_availability_from_entitlements(
total_quota_usd,
used_usd,
remaining_usd,
allow_wallet_overage,
// The gateway combines this quota with the user's stored preference.
allow_wallet_overage: false,
}
}
fn daily_quota_wallet_overage_policy(entitlements: &serde_json::Value) -> Option<bool> {
entitlements.as_array()?.iter().find_map(|item| {
(item.get("type").and_then(serde_json::Value::as_str) == Some("daily_quota"))
.then(|| {
item.get("allow_wallet_overage")
.and_then(serde_json::Value::as_bool)
})
.flatten()
})
}
#[async_trait]
impl BillingReadRepository for InMemoryBillingReadRepository {
async fn find_model_context(
@@ -578,13 +558,8 @@ impl BillingReadRepository for InMemoryBillingReadRepository {
.filter(|item| item.user_id == user_id)
.cloned()
.collect::<Vec<_>>();
let billing_plans = self
.billing_plans_by_id
.read()
.expect("billing repository lock");
Ok(Some(daily_quota_availability_from_entitlements(
entitlements,
&billing_plans,
now,
)))
}
crates/aether-data/runtime/src/repository/users/memory.rs
+2
View File
@@ -3226,6 +3226,7 @@ mod tests {
email_notifications: true,
usage_alerts: true,
announcement_notifications: true,
allow_wallet_overage: true,
};
let repository = InMemoryUserReadRepository::seed_auth_users([admin.clone()])
.with_export_users([export_row])
@@ -4029,6 +4030,7 @@ mod tests {
email_notifications: false,
usage_alerts: true,
announcement_notifications: false,
allow_wallet_overage: true,
};
assert!(repository
crates/aether-usage/runtime/src/request_metadata.rs
+1
View File
@@ -103,6 +103,7 @@ pub(crate) fn retain_first_byte_request_metadata(value: Option<Value>) -> Option
| "client_requested_stream"
| "upstream_is_stream"
| "api_key_is_standalone"
| "plan_wallet_fallback"
| "plan_usage_reservation_token"
| "request_path"
| "request_query_string"
crates/aether-usage/runtime/src/settlement.rs
+19
View File
@@ -286,6 +286,15 @@ fn usage_api_key_is_standalone(usage: &StoredRequestUsageAudit) -> bool {
}
fn usage_policy_reservation_token(usage: &StoredRequestUsageAudit) -> Option<&str> {
if usage
.request_metadata
.as_ref()
.and_then(|metadata| metadata.get("plan_wallet_fallback"))
.and_then(serde_json::Value::as_bool)
== Some(true)
{
return None;
}
usage
.request_metadata
.as_ref()
@@ -296,6 +305,16 @@ fn usage_policy_reservation_token(usage: &StoredRequestUsageAudit) -> Option<&st
}
fn event_usage_policy_reservation_token(event: &UsageEvent) -> Option<&str> {
if event
.data
.request_metadata
.as_ref()
.and_then(|metadata| metadata.get("plan_wallet_fallback"))
.and_then(serde_json::Value::as_bool)
== Some(true)
{
return None;
}
event
.data
.request_metadata
crates/aether-usage/runtime/src/write.rs
+6
View File
@@ -2246,6 +2246,12 @@ fn build_runtime_request_metadata_seed_from_parts(
Value::Bool(api_key_is_standalone),
);
}
if let Some(wallet_fallback) = context_bool(context, "plan_wallet_fallback") {
metadata.insert(
"plan_wallet_fallback".to_string(),
Value::Bool(wallet_fallback),
);
}
if let Some(websocket_mode) = context_bool(context, WEBSOCKET_MODE_METADATA_KEY) {
metadata.insert(
WEBSOCKET_MODE_METADATA_KEY.to_string(),