Revert "Revert "feat(accounting): add balance update in create_accounting_record""

This reverts commit 92e1c92ed82031668002dc4b70bb546cf7ca30d2.

sageapi/api/accounting.py

@@ -31,7 +31,16 @@ async def create_accounting_record(
     request_id: str = '',
     transno: str = '',
 ) -> str:
-    """Create a new accounting record with idempotency via request_id."""
+    """Create accounting record: write detail + update balance.
+
+    The core job of accounting is:
+    1. Write account detail (accounting_records) based on journal entry
+    2. Write accounting log (status='accounted' in the record)
+    3. Update account balance (customer_balance)
+
+    amount > 0 means charge (balance decreases)
+    amount < 0 means credit/refund (balance increases)
+    """
     result: dict[str, Any] = {'success': False, 'record_id': None}
 
     try:
@@ -61,40 +70,94 @@ async def create_accounting_record(
                     result['duplicate'] = True
                     return json.dumps(result, ensure_ascii=False, default=str)
 
-        sql = """
-            INSERT INTO accounting_records
-                (id, customer_id, llmid, model_name, pricing_id,
-                 input_tokens, output_tokens, total_tokens, quantity,
-                 amount, currency, request_id, transno, status,
-                 created_at, updated_at)
-            VALUES
-                (${id}$, ${customer_id}$, ${llmid}$, ${model_name}$, ${pricing_id}$,
-                 ${input_tokens}$, ${output_tokens}$, ${total_tokens}$, ${quantity}$,
-                 ${amount}$, ${currency}$, ${request_id}$, ${transno}$, 'accounted',
-                 ${created_at}$, ${updated_at}$)
-        """
-        params = {
-            'id': record_id,
-            'customer_id': customer_id,
-            'llmid': llmid,
-            'model_name': model_name,
-            'pricing_id': pricing_id,
-            'input_tokens': input_tokens,
-            'output_tokens': output_tokens,
-            'total_tokens': total_tokens,
-            'quantity': quantity,
-            'amount': amount,
-            'currency': currency,
-            'request_id': request_id,
-            'transno': transno,
-            'created_at': now,
-            'updated_at': now,
-        }
-
         async with DBPools().sqlorContext(dbname) as sor:
+            # === Step 1 & 2: Write detail + log (accounting record) ===
+            sql = """
+                INSERT INTO accounting_records
+                    (id, customer_id, llmid, model_name, pricing_id,
+                     input_tokens, output_tokens, total_tokens, quantity,
+                     amount, currency, request_id, transno, status,
+                     created_at, updated_at)
+                VALUES
+                    (${id}$, ${customer_id}$, ${llmid}$, ${model_name}$, ${pricing_id}$,
+                     ${input_tokens}$, ${output_tokens}$, ${total_tokens}$, ${quantity}$,
+                     ${amount}$, ${currency}$, ${request_id}$, ${transno}$, 'accounted',
+                     ${created_at}$, ${updated_at}$)
+            """
+            params = {
+                'id': record_id,
+                'customer_id': customer_id,
+                'llmid': llmid,
+                'model_name': model_name,
+                'pricing_id': pricing_id,
+                'input_tokens': input_tokens,
+                'output_tokens': output_tokens,
+                'total_tokens': total_tokens,
+                'quantity': quantity,
+                'amount': amount,
+                'currency': currency,
+                'request_id': request_id,
+                'transno': transno,
+                'created_at': now,
+                'updated_at': now,
+            }
             await sor.sqlExe(sql, params)
+
+            # === Step 3: Update account balance ===
+            # First read current balance + credit_limit with lock
+            balance_rows = await sor.sqlExe(
+                "SELECT balance, credit_limit FROM customer_balance "
+                "WHERE id = ${customer_id}$",
+                {'customer_id': customer_id},
+            )
+
+            if isinstance(balance_rows, list) and balance_rows:
+                cur = balance_rows[0]
+                cur_balance = float(cur.get('balance', 0))
+                credit_limit = cur.get('credit_limit')
+            else:
+                # No balance record yet, initialize
+                cur_balance = 0.0
+                credit_limit = None
+
+            # amount > 0 = charge (deduct), amount < 0 = credit (add)
+            new_balance = cur_balance - amount
+
+            # Overdraft check: if balance goes negative, check credit limit
+            if new_balance < -0.0000001:
+                if credit_limit is not None and float(credit_limit) > 0:
+                    if abs(new_balance) > float(credit_limit):
+                        result['error'] = (
+                            f'Insufficient balance: balance={cur_balance}, '
+                            f'credit_limit={credit_limit}, charge={amount}'
+                        )
+                        return json.dumps(result, ensure_ascii=False, default=str)
+                else:
+                    result['error'] = (
+                        f'Insufficient balance: balance={cur_balance}, charge={amount}'
+                    )
+                    return json.dumps(result, ensure_ascii=False, default=str)
+
+            # Upsert balance record
+            balance_sql = """
+                INSERT INTO customer_balance
+                    (id, balance, currency, last_consumption, cached_at)
+                VALUES
+                    (${customer_id}$, ${new_balance}$, ${currency}$, NOW(), NOW())
+                ON DUPLICATE KEY UPDATE
+                    balance = ${new_balance}$,
+                    last_consumption = NOW(),
+                    cached_at = NOW()
+            """
+            await sor.sqlExe(balance_sql, {
+                'customer_id': customer_id,
+                'new_balance': new_balance,
+                'currency': currency,
+            })
+
             result['success'] = True
             result['record_id'] = record_id
+            result['new_balance'] = new_balance
 
     except Exception as e:
         error(f'create_accounting_record error: {e}')