Merge pull request #512 from Entropy-Xu/codex/fix-wallet-overdraft-settlement

[codex] 修复钱包余额不足后重复消费

crates/aether-data/src/repository/settlement/memory.rs

@@ -161,15 +161,9 @@ impl SettlementWriteRepository for InMemorySettlementRepository {
                             before_gift,
                             input.total_cost_usd,
                         );
-                        if debit_plan.covered_usd() + SETTLEMENT_EPSILON_USD < input.total_cost_usd
+                        (wallet.balance, wallet.gift_balance) =
-                        {
+                            debit_plan.after_balances(before_recharge, before_gift);
-                            final_billing_status = "insufficient_quota".to_string();
+                        wallet.total_consumed += input.total_cost_usd;
-                            settlement.billing_status = final_billing_status.clone();
-                        } else {
-                            wallet.balance = before_recharge - debit_plan.recharge_deduction;
-                            wallet.gift_balance = before_gift - debit_plan.gift_deduction;
-                            wallet.total_consumed += input.total_cost_usd;
-                        }
                     }
                 }
 
@@ -359,7 +353,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn finite_wallet_insufficient_balance_does_not_overdraw() {
+    async fn finite_wallet_insufficient_balance_overdraws_and_settles() {
         let repository = InMemorySettlementRepository::seed(vec![sample_wallet()]);
         let settlement = repository
             .settle_usage(UsageSettlementInput {
@@ -378,12 +372,12 @@ mod tests {
             .expect("settlement should succeed")
             .expect("settlement should exist");
 
-        assert_eq!(settlement.billing_status, "insufficient_quota");
+        assert_eq!(settlement.billing_status, "settled");
         assert_eq!(settlement.wallet_balance_before, Some(12.0));
-        assert_eq!(settlement.wallet_balance_after, Some(12.0));
+        assert_eq!(settlement.wallet_balance_after, Some(-3.0));
-        assert_eq!(settlement.wallet_recharge_balance_after, Some(10.0));
+        assert_eq!(settlement.wallet_recharge_balance_after, Some(-3.0));
-        assert_eq!(settlement.wallet_gift_balance_after, Some(2.0));
+        assert_eq!(settlement.wallet_gift_balance_after, Some(0.0));
-        assert_eq!(settlement.provider_monthly_used_usd, None);
+        assert_eq!(settlement.provider_monthly_used_usd, Some(7.5));
     }
 
     #[tokio::test]

crates/aether-data/src/repository/settlement/mod.rs

@@ -9,11 +9,15 @@ const SETTLEMENT_EPSILON_USD: f64 = 0.000_000_01;
 struct WalletDebitPlan {
     recharge_deduction: f64,
     gift_deduction: f64,
+    recharge_overdraft: f64,
 }
 
 impl WalletDebitPlan {
-    fn covered_usd(self) -> f64 {
+    fn after_balances(self, recharge_balance: f64, gift_balance: f64) -> (f64, f64) {
-        self.recharge_deduction + self.gift_deduction
+        (
+            recharge_balance - self.recharge_deduction - self.recharge_overdraft,
+            gift_balance - self.gift_deduction,
+        )
     }
 }
 
@@ -26,13 +30,15 @@ fn plan_finite_wallet_debit(
     gift_balance: f64,
     requested_usd: f64,
 ) -> WalletDebitPlan {
-    let recharge_deduction = recharge_balance.max(0.0).min(requested_usd.max(0.0));
+    let requested_usd = requested_usd.max(0.0);
-    let gift_deduction = gift_balance
+    let recharge_deduction = recharge_balance.max(0.0).min(requested_usd);
-        .max(0.0)
+    let after_recharge_remaining = (requested_usd - recharge_deduction).max(0.0);
-        .min((requested_usd - recharge_deduction).max(0.0));
+    let gift_deduction = gift_balance.max(0.0).min(after_recharge_remaining);
+    let recharge_overdraft = (after_recharge_remaining - gift_deduction).max(0.0);
     WalletDebitPlan {
         recharge_deduction,
         gift_deduction,
+        recharge_overdraft,
     }
 }
 

crates/aether-data/src/repository/settlement/mysql.rs

@@ -207,6 +207,7 @@ async fn consume_daily_quota_mysql(
     request_id: &str,
     total_cost_usd: f64,
     wallet_available_usd: Option<f64>,
+    wallet_can_overdraft: bool,
     now_unix_secs: i64,
 ) -> Result<DailyQuotaDebitResult, DataLayerError> {
     if total_cost_usd <= 0.0 {
@@ -280,6 +281,7 @@ WHERE user_entitlement_id = ?
         });
     }
     if allow_wallet_overage
+        && !wallet_can_overdraft
         && wallet_available_usd.is_some_and(|available| {
             total_remaining + available + SETTLEMENT_EPSILON_USD < total_cost_usd
         })
@@ -448,6 +450,7 @@ FOR UPDATE
                 None
             };
 
+            let wallet_can_overdraft = wallet_row.is_some();
             let wallet_available_usd = match wallet_row.as_ref() {
                 Some(row) => {
                     let limit_mode: String = row.try_get("limit_mode").map_sql_err()?;
@@ -484,6 +487,7 @@ FOR UPDATE
                         &input.request_id,
                         input.total_cost_usd,
                         wallet_available_usd,
+                        wallet_can_overdraft,
                         updated_at,
                     )
                     .await?;
@@ -544,14 +548,8 @@ FOR UPDATE
                             before_gift,
                             wallet_debit_cost_usd,
                         );
-                        if debit_plan.covered_usd() + SETTLEMENT_EPSILON_USD < wallet_debit_cost_usd
+                        (after_recharge, after_gift) =
-                        {
+                            debit_plan.after_balances(before_recharge, before_gift);
-                            final_billing_status = "insufficient_quota".to_string();
-                            settlement.billing_status = final_billing_status.clone();
-                        } else {
-                            after_recharge = before_recharge - debit_plan.recharge_deduction;
-                            after_gift = before_gift - debit_plan.gift_deduction;
-                        }
                     }
                     if final_billing_status == "settled" {
                         sqlx::query(

crates/aether-data/src/repository/settlement/postgres.rs

@@ -312,6 +312,7 @@ async fn consume_daily_quota_postgres(
     request_id: &str,
     total_cost_usd: f64,
     wallet_available_usd: Option<f64>,
+    wallet_can_overdraft: bool,
 ) -> Result<DailyQuotaDebitResult, DataLayerError> {
     if total_cost_usd <= 0.0 {
         return Ok(DailyQuotaDebitResult::default());
@@ -379,6 +380,7 @@ WHERE user_entitlement_id = $1
         });
     }
     if allow_wallet_overage
+        && !wallet_can_overdraft
         && wallet_available_usd.is_some_and(|available| {
             total_remaining + available + SETTLEMENT_EPSILON_USD < total_cost_usd
         })
@@ -560,6 +562,7 @@ LIMIT 1
                             None
                         };
 
+                        let wallet_can_overdraft = wallet_row.is_some();
                         let wallet_available_usd = match wallet_row.as_ref() {
                             Some(row) => {
                                 let limit_mode: String =
@@ -600,6 +603,7 @@ LIMIT 1
                                     &input.request_id,
                                     input.total_cost_usd,
                                     wallet_available_usd,
+                                    wallet_can_overdraft,
                                 )
                                 .await?;
                                 if quota.insufficient {
@@ -646,16 +650,8 @@ LIMIT 1
                                         before_gift,
                                         wallet_debit_cost_usd,
                                     );
-                                    if debit_plan.covered_usd() + SETTLEMENT_EPSILON_USD
+                                    (after_recharge, after_gift) =
-                                        < wallet_debit_cost_usd
+                                        debit_plan.after_balances(before_recharge, before_gift);
-                                    {
-                                        final_billing_status = "insufficient_quota".to_string();
-                                        settlement.billing_status = final_billing_status.clone();
-                                    } else {
-                                        after_recharge =
-                                            before_recharge - debit_plan.recharge_deduction;
-                                        after_gift = before_gift - debit_plan.gift_deduction;
-                                    }
                                 }
                                 if final_billing_status == "settled" {
                                     sqlx::query(

crates/aether-data/src/repository/settlement/sqlite.rs

@@ -221,6 +221,7 @@ async fn consume_daily_quota_sqlite(
     request_id: &str,
     total_cost_usd: f64,
     wallet_available_usd: Option<f64>,
+    wallet_can_overdraft: bool,
     now_unix_secs: i64,
 ) -> Result<DailyQuotaDebitResult, DataLayerError> {
     if total_cost_usd <= 0.0 {
@@ -293,6 +294,7 @@ WHERE user_entitlement_id = ?
         });
     }
     if allow_wallet_overage
+        && !wallet_can_overdraft
         && wallet_available_usd.is_some_and(|available| {
             total_remaining + available + SETTLEMENT_EPSILON_USD < total_cost_usd
         })
@@ -459,6 +461,7 @@ LIMIT 1
                 None
             };
 
+            let wallet_can_overdraft = wallet_row.is_some();
             let wallet_available_usd = match wallet_row.as_ref() {
                 Some(row) => {
                     let limit_mode: String = row.try_get("limit_mode").map_sql_err()?;
@@ -495,6 +498,7 @@ LIMIT 1
                         &input.request_id,
                         input.total_cost_usd,
                         wallet_available_usd,
+                        wallet_can_overdraft,
                         updated_at,
                     )
                     .await?;
@@ -555,14 +559,8 @@ LIMIT 1
                             before_gift,
                             wallet_debit_cost_usd,
                         );
-                        if debit_plan.covered_usd() + SETTLEMENT_EPSILON_USD < wallet_debit_cost_usd
+                        (after_recharge, after_gift) =
-                        {
+                            debit_plan.after_balances(before_recharge, before_gift);
-                            final_billing_status = "insufficient_quota".to_string();
-                            settlement.billing_status = final_billing_status.clone();
-                        } else {
-                            after_recharge = before_recharge - debit_plan.recharge_deduction;
-                            after_gift = before_gift - debit_plan.gift_deduction;
-                        }
                     }
                     if final_billing_status == "settled" {
                         sqlx::query(
@@ -813,6 +811,55 @@ mod tests {
         assert_eq!(wallet_total, 12.0);
     }
 
+    #[tokio::test]
+    async fn sqlite_repository_overdraws_finite_wallet_and_settles_usage() {
+        let pool = sqlx::sqlite::SqlitePoolOptions::new()
+            .max_connections(1)
+            .connect("sqlite::memory:")
+            .await
+            .expect("sqlite pool should connect");
+        run_sqlite_migrations(&pool)
+            .await
+            .expect("sqlite migrations should run");
+        seed_settlement_rows(&pool).await;
+
+        let repository = SqliteSettlementRepository::new(pool.clone());
+        let settlement = repository
+            .settle_usage(UsageSettlementInput {
+                request_id: "request-overdraw".to_string(),
+                user_id: Some("user-1".to_string()),
+                api_key_id: None,
+                api_key_is_standalone: false,
+                provider_id: Some("provider-1".to_string()),
+                status: "completed".to_string(),
+                billing_status: "pending".to_string(),
+                total_cost_usd: 15.0,
+                actual_total_cost_usd: 7.5,
+                finalized_at_unix_secs: Some(1_236),
+            })
+            .await
+            .expect("settlement should run")
+            .expect("usage should exist");
+
+        assert_eq!(settlement.billing_status, "settled");
+        assert_eq!(settlement.wallet_id.as_deref(), Some("wallet-1"));
+        assert_eq!(settlement.wallet_balance_before, Some(12.0));
+        assert_eq!(settlement.wallet_balance_after, Some(-3.0));
+        assert_eq!(settlement.wallet_recharge_balance_after, Some(-3.0));
+        assert_eq!(settlement.wallet_gift_balance_after, Some(0.0));
+        assert_eq!(settlement.provider_monthly_used_usd, Some(12.5));
+
+        let wallet = sqlx::query(
+            "SELECT balance, gift_balance, total_consumed FROM wallets WHERE id = 'wallet-1'",
+        )
+        .fetch_one(&pool)
+        .await
+        .expect("wallet should load");
+        assert_eq!(wallet.try_get::<f64, _>("balance").unwrap(), -3.0);
+        assert_eq!(wallet.try_get::<f64, _>("gift_balance").unwrap(), 0.0);
+        assert_eq!(wallet.try_get::<f64, _>("total_consumed").unwrap(), 15.0);
+    }
+
     #[tokio::test]
     async fn sqlite_repository_records_wallet_for_quota_covered_user_usage() {
         let pool = sqlx::sqlite::SqlitePoolOptions::new()
@@ -883,7 +930,8 @@ INSERT INTO "usage" (
 )
 VALUES
   ('request-1', 'user-1', 'provider-1', 'completed', 'pending', 3.0, 2.0),
-  ('request-2', 'user-1', 'provider-1', 'failed', 'pending', 3.0, 2.0);
+  ('request-2', 'user-1', 'provider-1', 'failed', 'pending', 3.0, 2.0),
+  ('request-overdraw', 'user-1', 'provider-1', 'completed', 'pending', 15.0, 7.5);
 "#,
         )
         .execute(pool)