feat(billing): USD consume writes credit_ledger in one tx

Co-Authored-By: Claude <noreply@anthropic.com>

src/services/billing/mod.rs

@@ -2,7 +2,7 @@ pub mod ledger;
 
 use crate::{
     error::{AppError, AppResult},
-    services::shared::{enforce_key_request_limit, increment_key_usage},
+    services::shared::enforce_key_request_limit,
 };
 use chrono::Timelike;
 use serde::{Deserialize, Serialize};
@@ -56,25 +56,68 @@ impl UsageTracker {
         cost: f64,
     ) -> AppResult<()> {
         Self::ensure_schema(&self.db).await?;
+        crate::services::billing::ledger::ensure_ledger_schema(&self.db).await?;
         let prompt_tokens = prompt_tokens.max(0);
         let completion_tokens = completion_tokens.max(0);
         let total_tokens = prompt_tokens + completion_tokens;
+        let cost = cost.max(0.0);
+        let model = model.into();
 
+        let mut tx = self.db.begin().await?;
         sqlx::query(
             "INSERT INTO usage_logs
                 (downstream_key_id, model, prompt_tokens, completion_tokens, total_tokens, cost)
              VALUES (?, ?, ?, ?, ?, ?)",
         )
         .bind(key_id)
-        .bind(model.into())
+        .bind(&model)
         .bind(prompt_tokens)
         .bind(completion_tokens)
         .bind(total_tokens)
-        .bind(cost.max(0.0))
+        .bind(cost)
-        .execute(&*self.db)
+        .execute(&mut *tx)
         .await?;
 
-        increment_key_usage(&self.db, key_id, cost.max(0.0)).await?;
+        sqlx::query(
+            "UPDATE downstream_api_keys
+             SET used_requests = COALESCE(used_requests,0)+1,
+                 used_cost = COALESCE(used_cost,0)+?,
+                 last_used_at = datetime('now'), updated_at = datetime('now')
+             WHERE id = ?",
+        )
+        .bind(cost)
+        .bind(key_id)
+        .execute(&mut *tx)
+        .await?;
+
+        // Resolve owning user + new used_cost for the ledger snapshot.
+        let owner: Option<(Option<i64>, Option<f64>, Option<f64>)> = sqlx::query_as(
+            "SELECT user_id, max_cost, used_cost FROM downstream_api_keys WHERE id = ?",
+        )
+        .bind(key_id)
+        .fetch_optional(&mut *tx)
+        .await?;
+        if let Some((Some(user_id), max_cost, used_cost)) = owner {
+            let used = used_cost.unwrap_or(0.0);
+            let remaining = max_cost.map(|m| (m - used).max(0.0)).unwrap_or(used);
+            crate::services::billing::ledger::record_ledger_entry(
+                &mut tx,
+                crate::services::billing::ledger::LedgerEntry {
+                    user_id,
+                    track: "usd",
+                    entry_type: "consume",
+                    amount: -cost,
+                    balance_after: remaining,
+                    ref_type: Some("usage_log"),
+                    ref_id: None,
+                    model: Some(&model),
+                    operator: Some("system"),
+                    note: None,
+                },
+            )
+            .await?;
+        }
+        tx.commit().await?;
         Ok(())
     }
 
@@ -832,6 +875,7 @@ mod tests {
                 used_cost REAL DEFAULT 0,
                 max_requests INTEGER,
                 used_requests INTEGER DEFAULT 0,
+                user_id INTEGER,
                 last_used_at TEXT,
                 updated_at TEXT DEFAULT (datetime('now'))
             )",
@@ -851,6 +895,62 @@ mod tests {
         pool
     }
 
+    async fn billing_test_pool() -> SqlitePool {
+        let pool = SqlitePoolOptions::new()
+            .max_connections(1)
+            .connect("sqlite::memory:")
+            .await
+            .unwrap();
+        sqlx::query(
+            "CREATE TABLE users (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                username TEXT NOT NULL,
+                email TEXT NOT NULL,
+                password_hash TEXT NOT NULL,
+                role TEXT NOT NULL DEFAULT 'user',
+                status TEXT NOT NULL DEFAULT 'active',
+                created_at TEXT,
+                updated_at TEXT
+            )",
+        )
+        .execute(&pool)
+        .await
+        .unwrap();
+        sqlx::query(
+            "CREATE TABLE downstream_api_keys (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                name TEXT NOT NULL,
+                key TEXT NOT NULL,
+                enabled INTEGER DEFAULT 1,
+                max_cost REAL,
+                used_cost REAL DEFAULT 0,
+                max_requests INTEGER,
+                used_requests INTEGER DEFAULT 0,
+                user_id INTEGER,
+                last_used_at TEXT,
+                updated_at TEXT DEFAULT (datetime('now'))
+            )",
+        )
+        .execute(&pool)
+        .await
+        .unwrap();
+        sqlx::query(
+            "INSERT INTO users (id, username, email, password_hash, role, status)
+             VALUES (9, 'creator', 'creator@example.test', 'hash', 'user', 'active')",
+        )
+        .execute(&pool)
+        .await
+        .unwrap();
+        sqlx::query(
+            "INSERT INTO downstream_api_keys (id, name, key, max_cost, used_cost, user_id)
+             VALUES (7, 'default', 'sk-user-test', 10.0, 2.0, 9)",
+        )
+        .execute(&pool)
+        .await
+        .unwrap();
+        pool
+    }
+
     async fn create_proxy_logs(pool: &SqlitePool) {
         sqlx::query(
             "CREATE TABLE proxy_logs (
@@ -918,6 +1018,32 @@ mod tests {
         assert_eq!(summary.total_cost, 0.25);
     }
 
+    #[tokio::test]
+    async fn record_usage_writes_consume_ledger_entry() {
+        let pool = billing_test_pool().await;
+        let tracker = UsageTracker::new(std::sync::Arc::new(pool.clone()));
+        crate::services::billing::ledger::ensure_ledger_schema(&pool)
+            .await
+            .unwrap();
+        // key_id=7 属于 user 9，max_cost=10 used_cost=2（由 helper 建好）
+        tracker
+            .record_usage(7, "gpt-4o", 100, 50, 0.42)
+            .await
+            .unwrap();
+
+        let row: (String, String, f64, Option<String>) = sqlx::query_as(
+            "SELECT track, entry_type, amount, model FROM credit_ledger
+             WHERE user_id = 9 ORDER BY id DESC LIMIT 1",
+        )
+        .fetch_one(&pool)
+        .await
+        .unwrap();
+        assert_eq!(row.0, "usd");
+        assert_eq!(row.1, "consume");
+        assert_eq!(row.2, -0.42); // 消费为负
+        assert_eq!(row.3.as_deref(), Some("gpt-4o"));
+    }
+
     #[tokio::test]
     async fn rate_limit_checks_request_and_cost_limits() {
         let tracker = UsageTracker::new(Arc::new(pool().await));