projectAIpopular/gateway/proxy/billing.py

"""计费两阶段（D-P11）：预扣-结算-退款，作为 Ledger 的 mixin。

拆分为独立模块的工程原因：ledger.py 承载 DDL/CRUD/查询；金额敏感的
两阶段语义集中在此，便于评审与测试聚焦。

并发正确性：全局锁内 check-then-act（单写者模型）+ 余额原子防线
（语句条件 balance_milli >= ?，受影响行数为 0 即不足）。
"""
from __future__ import annotations

import threading
import time
from typing import Any, Dict, List, Optional

from gateway.proxy.ledgerutil import _today

# 预扣：余额原子防线；日上限判定在锁内 Python 侧（单写者模型下等价安全）
# 占位顺序：est, spent_total, today, sid, est
_SQL_HOLD = """UPDATE students SET
  balance_milli = balance_milli - ?,
  spent_today_milli = ?, spent_date = ?
 WHERE id = ? AND balance_milli >= ?"""

_SQL_HOLD_MARK = """INSERT INTO usage_ledger
  (request_id, ts, key_id, model, bucket, charged_milli, status)
  VALUES (?, ?, ?, ?, ?, ?, 'holding')"""

_SQL_INSUFFICIENT = """INSERT OR REPLACE INTO usage_ledger
  (request_id, ts, key_id, model, bucket, charged_milli, status)
  VALUES (?, ?, ?, ?, ?, 0, 'insufficient')"""

_SQL_SETTLE_REFUND = """UPDATE students SET
  balance_milli = balance_milli + ?,
  spent_today_milli = MAX(0, spent_today_milli - ?)
 WHERE id = ?"""

_SQL_SETTLE_USAGE = """UPDATE usage_ledger SET
  in_miss_tok = ?, in_hit_tok = ?, out_tok = ?, gateway_cached = ?,
  upstream_cost_milli = ?, charged_milli = ?, margin_milli = ? - ?,
  ttfb_ms = ?, total_ms = ?, status = ?
 WHERE request_id = ?"""


class BillingMixin:
    """try_hold / settle / void（由 Ledger 继承；依赖 _lock/_connect）。"""

    _lock: threading.Lock

    def _connect(self):  # 由 Ledger 提供
        raise NotImplementedError

    def try_hold(self, request_id: str, key_id: int, student_id: int,
                 model: str, bucket: str, est_milli: int, ts: float) -> bool:
        """预扣：锁内读现状 -> 日上限判定 -> 余额原子扣减 -> 写 holding 流水。

        受影响行数为 0 即不足（余额不够）-> False（402）；
        request_id 幂等：重复 hold 直接返回 True（不重复扣）。
        """
        today = _today(ts)
        with self._lock, self._connect() as conn:
            dup = conn.execute(
                "SELECT 1 FROM usage_ledger WHERE request_id = ?",
                (request_id,)).fetchone()
            if dup:
                return True
            stu = conn.execute(
                "SELECT balance_milli, daily_cap_milli, spent_today_milli, spent_date"
                " FROM students WHERE id = ?", (student_id,)).fetchone()
            if stu is None:
                return False
            spent_base = stu["spent_today_milli"] if stu["spent_date"] == today else 0
            if stu["daily_cap_milli"] > 0 and spent_base + est_milli > stu["daily_cap_milli"]:
                conn.execute(_SQL_INSUFFICIENT,
                             (request_id, int(ts), key_id, model, bucket))
                return False
            cur = conn.execute(_SQL_HOLD, (
                est_milli, spent_base + est_milli, today, student_id, est_milli))
            if cur.rowcount == 0:
                conn.execute(_SQL_INSUFFICIENT,
                             (request_id, int(ts), key_id, model, bucket))
                return False
            conn.execute(_SQL_HOLD_MARK,
                         (request_id, int(ts), key_id, model, bucket, est_milli))
            return True

    def settle(self, request_id: str, actual_milli: int, *,
               in_miss_tok: int = 0, in_hit_tok: int = 0, out_tok: int = 0,
               gateway_cached: int = 0, upstream_cost_milli: int = 0,
               ttfb_ms: Optional[int] = None, total_ms: Optional[int] = None,
               status: str = "ok") -> bool:
        """结算：按真实值更新流水并回补（预扣额 − 实际额）差额。"""
        with self._lock, self._connect() as conn:
            row = conn.execute(
                "SELECT key_id, charged_milli FROM usage_ledger WHERE request_id = ?",
                (request_id,)).fetchone()
            if row is None:
                return False
            key_id = row["key_id"]
            est = row["charged_milli"]
            student = conn.execute(
                "SELECT student_id FROM proxy_keys WHERE id = ?", (key_id,)).fetchone()
            if student is None:
                return False
            refund = est - actual_milli
            if refund != 0:
                conn.execute(_SQL_SETTLE_REFUND,
                             (refund, refund, student["student_id"]))
            conn.execute(_SQL_SETTLE_USAGE, (
                in_miss_tok, in_hit_tok, out_tok, gateway_cached,
                upstream_cost_milli, actual_milli,
                actual_milli, upstream_cost_milli, ttfb_ms, total_ms,
                status, request_id))
            return True

    def void(self, request_id: str, status: str = "aborted") -> bool:
        """全额退款（上游失败）；流水保留审计。"""
        with self._lock, self._connect() as conn:
            row = conn.execute(
                "SELECT key_id, charged_milli, status FROM usage_ledger"
                " WHERE request_id = ?", (request_id,)).fetchone()
            if row is None or row["status"] not in ("holding", "ok"):
                return False
            est = row["charged_milli"]
            student = conn.execute(
                "SELECT student_id FROM proxy_keys WHERE id = ?",
                (row["key_id"],)).fetchone()
            if student is not None and est > 0:
                conn.execute(_SQL_SETTLE_REFUND,
                             (est, est, student["student_id"]))
            conn.execute(
                "UPDATE usage_ledger SET charged_milli = 0, status = ?"
                " WHERE request_id = ?", (status, request_id))
            return True

    # ---------- 流水查询 ----------
    def get_usage(self, request_id: str) -> Optional[Dict[str, Any]]:
        with self._lock, self._connect() as conn:
            row = conn.execute(
                "SELECT * FROM usage_ledger WHERE request_id = ?",
                (request_id,)).fetchone()
        return dict(row) if row else None

    def list_usage(self, student_id: Optional[int] = None,
                   limit: int = 50, offset: int = 0) -> List[Dict[str, Any]]:
        """流水分页（可按学生过滤，经其名下 key）；按学生过滤走联表常量语句。"""
        if student_id is None:
            with self._lock, self._connect() as conn:
                rows = conn.execute(
                    "SELECT * FROM usage_ledger ORDER BY ts DESC LIMIT ? OFFSET ?",
                    (limit, offset)).fetchall()
            return [dict(r) for r in rows]
        with self._lock, self._connect() as conn:
            rows = conn.execute(
                "SELECT u.* FROM usage_ledger u JOIN proxy_keys k ON k.id = u.key_id"
                " WHERE k.student_id = ? ORDER BY u.ts DESC LIMIT ? OFFSET ?",
                (student_id, limit, offset)).fetchall()
        return [dict(r) for r in rows]