i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-10 11:19: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()?,
})
}