feat(wallet): 钱包系统替代配额系统,新增支付与退款机制

- 新增钱包余额管理、充值、扣费、退款完整流程
- 新增支付网关抽象层(支持手动/支付宝/微信)
- 用量计费从配额系统迁移到钱包余额扣费
- 新增管理员钱包管理与支付订单管理页面
- 新增用户钱包中心页面
- 移除独立 Key 锁定机制,统一由钱包余额控制
- 新增相关 API 路由、序列化器与数据库迁移
- 新增钱包、支付、退款相关测试
This commit is contained in:
LewisPen
2026-03-08 00:05:48 +08:00
committed by fawney19
parent 9cdcce1b5f
commit 783f654953
108 changed files with 13152 additions and 3372 deletions

View File

@@ -428,7 +428,9 @@ class OAuthService:
or (email.split("@", 1)[0] if email else None)
or f"user_{uuid.uuid4().hex[:8]}"
)
default_quota = SystemConfigService.get_config(db, "default_user_quota_usd", default=10.0)
default_initial_gift = SystemConfigService.get_config(
db, "default_user_initial_gift_usd", default=None
)
# 生成唯一用户名 + 创建用户(简单重试)
user: User | None = None
@@ -445,9 +447,20 @@ class OAuthService:
role=UserRole.USER,
is_active=True,
last_login_at=now,
quota_usd=default_quota,
)
db.add(user)
db.flush()
from src.services.wallet import WalletService
WalletService.initialize_user_wallet(
db,
user=user,
initial_gift_usd=default_initial_gift,
unlimited=False,
description="OAuth 注册初始赠款",
)
db.commit()
db.refresh(user)
last_error = None

View File

@@ -31,7 +31,6 @@ if TYPE_CHECKING:
from src.models.database import ApiKey, User, UserRole
from src.services.auth.jwt_blacklist import JWTBlacklistService
from src.services.cache.user_cache import UserCacheService
from src.services.user.apikey import ApiKeyService
# API Key last_used_at 更新节流配置
# 同一个 API Key 在此时间间隔内只会更新一次 last_used_at
@@ -336,9 +335,9 @@ class AuthService:
username = f"{base_username}_ldap_{int(time.time())}{uuid.uuid4().hex[:4]}"
logger.info(f"LDAP 用户名冲突,使用新用户名: {ldap_user['username']} -> {username}")
# 读取系统配置的默认配额
default_quota = SystemConfigService.get_config(
db, "default_user_quota_usd", default=10.0
# 读取系统配置的默认初始赠款
default_initial_gift = SystemConfigService.get_config(
db, "default_user_initial_gift_usd", default=None
)
# 创建新用户
@@ -353,11 +352,22 @@ class AuthService:
role=UserRole.USER,
is_active=True,
last_login_at=datetime.now(timezone.utc),
quota_usd=default_quota,
)
try:
db.add(user)
db.flush()
from src.services.wallet import WalletService
WalletService.initialize_user_wallet(
db,
user=user,
initial_gift_usd=default_initial_gift,
unlimited=False,
description="LDAP 注册初始赠款",
)
db.commit()
db.refresh(user)
logger.info(f"LDAP 用户创建成功: {ldap_user['email']} (ID: {user.id})")
@@ -408,7 +418,7 @@ class AuthService:
logger.warning("API认证失败 - 密钥已禁用")
return None
if key_record.is_locked:
if key_record.is_locked and not key_record.is_standalone:
logger.warning("API认证失败 - 密钥已被管理员锁定")
raise ForbiddenException("该密钥已被管理员锁定,请联系管理员")
@@ -424,17 +434,6 @@ class AuthService:
logger.warning("API认证失败 - 密钥已过期")
return None
# 检查余额限制仅独立Key
is_balance_ok, remaining = ApiKeyService.check_balance(key_record)
if not is_balance_ok:
# 获取剩余余额用于日志
remaining_balance = ApiKeyService.get_remaining_balance(key_record)
logger.warning(
f"API认证失败 - 余额不足 "
f"(已用: ${key_record.balance_used_usd:.4f}, 剩余: ${remaining_balance:.4f})"
)
return None
# 获取用户
user = key_record.user
if not user.is_active:
@@ -467,23 +466,22 @@ class AuthService:
return user, key_record
@staticmethod
def check_user_quota(user: User, estimated_cost: float = 0) -> bool:
"""检查用户配额"""
if user.role == UserRole.ADMIN:
return True # 管理员无限制
def check_user_balance_access(user: User, estimated_cost: float = 0) -> bool:
"""按钱包余额/额度模式校验请求可用性。"""
from src.services.wallet import WalletService
# NULL 表示无限制
if user.quota_usd is None:
_ = estimated_cost
if user.role == UserRole.ADMIN:
return True
# 检查美元配额
if user.used_usd + estimated_cost > user.quota_usd:
logger.warning(
f"用户配额不足: {user.email} (已用: ${user.used_usd:.2f}, 配额: ${user.quota_usd:.2f})"
)
wallet = getattr(user, "wallet", None)
if wallet is None:
return False
return True
if wallet.status != "active":
return False
if WalletService.is_unlimited_wallet(wallet):
return True
return WalletService.get_spendable_balance_value(wallet) > 0
@staticmethod
def check_permission(user: User, required_role: UserRole = UserRole.USER) -> bool:

View File

@@ -31,6 +31,11 @@ def to_decimal(value: float | int | str | Decimal | None) -> Decimal:
return Decimal(str(value))
def to_money_decimal(value: float | int | str | Decimal | None) -> Decimal:
"""Convert values to Decimal and quantize to billing storage precision."""
return quantize_cost(to_decimal(value))
def quantize_decimal(value: Decimal, *, precision: int) -> Decimal:
"""Quantize a Decimal to the given number of decimal places (ROUND_HALF_UP)."""
quantizer = Decimal(10) ** -precision

View File

@@ -133,8 +133,6 @@ class UserCacheService:
"role": user.role.value if user.role else None,
"is_active": user.is_active,
"auth_source": user.auth_source.value if user.auth_source else None,
"quota_usd": float(user.quota_usd) if user.quota_usd is not None else None,
"used_usd": float(user.used_usd),
"created_at": user.created_at.isoformat() if user.created_at else None,
"last_login_at": user.last_login_at.isoformat() if user.last_login_at else None,
"model_capability_settings": user.model_capability_settings,
@@ -159,7 +157,6 @@ class UserCacheService:
email_verified=user_dict.get("email_verified", False),
username=user_dict["username"],
is_active=user_dict["is_active"],
used_usd=user_dict["used_usd"],
)
# 设置可选字段
@@ -169,9 +166,6 @@ class UserCacheService:
if user_dict.get("auth_source"):
user.auth_source = AuthSource(user_dict["auth_source"])
if user_dict.get("quota_usd") is not None:
user.quota_usd = user_dict["quota_usd"]
if user_dict.get("created_at"):
user.created_at = datetime.fromisoformat(user_dict["created_at"])

View File

@@ -0,0 +1,3 @@
from src.services.payment.service import PaymentService
__all__ = ["PaymentService"]

View File

@@ -0,0 +1,23 @@
from __future__ import annotations
from src.services.payment.gateway.alipay import AlipayGateway
from src.services.payment.gateway.base import PaymentGateway
from src.services.payment.gateway.manual import ManualGateway
from src.services.payment.gateway.wechat import WeChatGateway
_GATEWAYS: dict[str, PaymentGateway] = {
"alipay": AlipayGateway(),
"wechat": WeChatGateway(),
"manual": ManualGateway(),
}
def get_payment_gateway(payment_method: str) -> PaymentGateway:
key = (payment_method or "").strip().lower()
gateway = _GATEWAYS.get(key)
if gateway is None:
raise ValueError(f"unsupported payment_method: {payment_method}")
return gateway
__all__ = ["PaymentGateway", "get_payment_gateway"]

View File

@@ -0,0 +1,21 @@
from __future__ import annotations
from typing import Any
from src.services.payment.gateway.base import PaymentGateway
class AlipayGateway(PaymentGateway):
payment_method = "alipay"
display_name = "支付宝"
def create_checkout_payload(self, *, order: Any) -> dict[str, Any]:
gateway_order_id = getattr(order, "gateway_order_id", None) or f"ali_{order.order_no}"
return {
"gateway": self.payment_method,
"display_name": self.display_name,
"gateway_order_id": gateway_order_id,
"payment_url": f"/pay/mock/alipay/{order.order_no}",
"qr_code": f"mock://alipay/{order.order_no}",
"expires_at": getattr(order, "expires_at", None),
}

View File

@@ -0,0 +1,69 @@
from __future__ import annotations
import hashlib
import hmac
import json
from abc import ABC, abstractmethod
from typing import Any
class PaymentGateway(ABC):
"""支付网关抽象。
当前阶段只提供统一结构和占位返回,便于后续接入真实 SDK。
"""
payment_method: str
display_name: str
@abstractmethod
def create_checkout_payload(self, *, order: Any) -> dict[str, Any]:
"""为前端返回统一的支付指引结构。"""
@staticmethod
def build_callback_signature(
*,
payload: dict[str, Any] | None,
callback_secret: str | None,
) -> str | None:
if payload is None:
return None
if not callback_secret:
return None
canonical = json.dumps(
payload,
sort_keys=True,
ensure_ascii=False,
separators=(",", ":"),
default=str,
)
return hmac.new(
callback_secret.encode("utf-8"),
canonical.encode("utf-8"),
hashlib.sha256,
).hexdigest()
def verify_callback_payload(
self,
*,
payload: dict[str, Any] | None,
callback_signature: str | None = None,
callback_secret: str | None = None,
) -> bool:
"""校验回调。
默认使用 HMAC-SHA256 对 payload 进行签名校验。
真实接入时可由各支付渠道覆写该方法使用官方 SDK 验签。
"""
expected_signature = self.build_callback_signature(
payload=payload,
callback_secret=callback_secret,
)
if expected_signature is None:
return False
provided = (callback_signature or "").strip()
if not provided:
return False
if provided.lower().startswith("sha256="):
provided = provided.split("=", 1)[1]
return hmac.compare_digest(provided.lower(), expected_signature.lower())

View File

@@ -0,0 +1,22 @@
from __future__ import annotations
from typing import Any
from src.services.payment.gateway.base import PaymentGateway
class ManualGateway(PaymentGateway):
payment_method = "manual"
display_name = "人工打款"
def create_checkout_payload(self, *, order: Any) -> dict[str, Any]:
gateway_order_id = getattr(order, "gateway_order_id", None) or f"manual_{order.order_no}"
return {
"gateway": self.payment_method,
"display_name": self.display_name,
"gateway_order_id": gateway_order_id,
"payment_url": None,
"qr_code": None,
"instructions": "请线下确认到账后由管理员处理",
"expires_at": getattr(order, "expires_at", None),
}

View File

@@ -0,0 +1,21 @@
from __future__ import annotations
from typing import Any
from src.services.payment.gateway.base import PaymentGateway
class WeChatGateway(PaymentGateway):
payment_method = "wechat"
display_name = "微信支付"
def create_checkout_payload(self, *, order: Any) -> dict[str, Any]:
gateway_order_id = getattr(order, "gateway_order_id", None) or f"wx_{order.order_no}"
return {
"gateway": self.payment_method,
"display_name": self.display_name,
"gateway_order_id": gateway_order_id,
"payment_url": f"/pay/mock/wechat/{order.order_no}",
"qr_code": f"mock://wechat/{order.order_no}",
"expires_at": getattr(order, "expires_at", None),
}

View File

@@ -0,0 +1,478 @@
from __future__ import annotations
import hashlib
import json
from datetime import datetime, timedelta, timezone
from decimal import Decimal
from typing import Any
from uuid import uuid4
from sqlalchemy.orm import Session
from src.models.database import PaymentCallback, PaymentOrder, User, Wallet
from src.services.billing.precision import to_money_decimal
from src.services.payment.gateway import get_payment_gateway
from src.services.wallet import WalletService
class PaymentService:
"""支付订单与回调处理服务。
当前实现目标:
- 打通充值订单创建
- 打通支付回调幂等到账
- 真实网关签名/SDK 留给后续渠道适配层
"""
@staticmethod
def _build_order_no() -> str:
ts = datetime.now(timezone.utc).strftime("%Y%m%d%H%M%S%f")
return f"po_{ts}_{uuid4().hex[:12]}"
@staticmethod
def _build_payload_hash(payload: dict[str, Any] | None) -> str | None:
if payload is None:
return None
encoded = json.dumps(payload, sort_keys=True, ensure_ascii=False, default=str).encode(
"utf-8"
)
return hashlib.sha256(encoded).hexdigest()
@classmethod
def create_recharge_order(
cls,
db: Session,
*,
user: User,
amount_usd: Decimal | float | int | str,
payment_method: str,
pay_amount: Decimal | float | int | str | None = None,
pay_currency: str | None = None,
exchange_rate: Decimal | float | int | str | None = None,
expires_in_minutes: int = 30,
gateway_order_id: str | None = None,
gateway_response: dict[str, Any] | None = None,
) -> PaymentOrder:
amount = to_money_decimal(amount_usd)
if amount <= Decimal("0"):
raise ValueError("recharge amount must be positive")
if not payment_method:
raise ValueError("payment_method is required")
if payment_method == "admin_manual":
raise ValueError("admin_manual is reserved for admin recharge")
gateway = get_payment_gateway(payment_method)
wallet = WalletService.get_or_create_wallet(db, user=user)
if wallet is None:
raise ValueError("wallet not available")
if wallet.status != "active":
raise ValueError("wallet is not active")
now = datetime.now(timezone.utc)
order = PaymentOrder(
order_no=cls._build_order_no(),
wallet_id=wallet.id,
user_id=user.id,
amount_usd=amount,
pay_amount=to_money_decimal(pay_amount) if pay_amount is not None else None,
pay_currency=pay_currency,
exchange_rate=to_money_decimal(exchange_rate) if exchange_rate is not None else None,
refunded_amount_usd=Decimal("0"),
refundable_amount_usd=Decimal("0"),
payment_method=payment_method,
gateway_order_id=gateway_order_id,
gateway_response=gateway_response,
status="pending",
expires_at=now + timedelta(minutes=max(expires_in_minutes, 1)),
)
db.add(order)
db.flush()
checkout = gateway.create_checkout_payload(order=order)
order.gateway_order_id = order.gateway_order_id or checkout.get("gateway_order_id")
order.gateway_response = gateway_response if gateway_response is not None else checkout
return order
@classmethod
def refresh_order_status(cls, order: PaymentOrder | None) -> bool:
if order is None:
return False
if order.status != "pending":
return False
now = datetime.now(timezone.utc)
if order.expires_at is not None and order.expires_at < now:
order.status = "expired"
return True
return False
@staticmethod
def get_order(
db: Session,
*,
order_id: str | None = None,
order_no: str | None = None,
gateway_order_id: str | None = None,
) -> PaymentOrder | None:
if order_id:
return db.query(PaymentOrder).filter(PaymentOrder.id == order_id).first()
if order_no:
return db.query(PaymentOrder).filter(PaymentOrder.order_no == order_no).first()
if gateway_order_id:
return (
db.query(PaymentOrder)
.filter(PaymentOrder.gateway_order_id == gateway_order_id)
.first()
)
return None
@classmethod
def list_user_orders(
cls,
db: Session,
*,
user_id: str,
limit: int,
offset: int,
) -> tuple[list[PaymentOrder], int, bool]:
expired_count = cls.expire_overdue_pending_orders(db, user_id=user_id)
q = db.query(PaymentOrder).filter(PaymentOrder.user_id == user_id)
total = q.count()
items = q.order_by(PaymentOrder.created_at.desc()).offset(offset).limit(limit).all()
return items, total, expired_count > 0
@classmethod
def list_orders(
cls,
db: Session,
*,
status: str | None = None,
payment_method: str | None = None,
limit: int = 50,
offset: int = 0,
) -> tuple[list[PaymentOrder], int, bool]:
expired_count = 0
if status in {None, "pending", "expired"}:
expired_count = cls.expire_overdue_pending_orders(
db,
payment_method=payment_method,
)
q = db.query(PaymentOrder)
if status:
q = q.filter(PaymentOrder.status == status)
if payment_method:
q = q.filter(PaymentOrder.payment_method == payment_method)
total = q.count()
items = q.order_by(PaymentOrder.created_at.desc()).offset(offset).limit(limit).all()
return items, total, expired_count > 0
@staticmethod
def expire_overdue_pending_orders(
db: Session,
*,
user_id: str | None = None,
payment_method: str | None = None,
) -> int:
now = datetime.now(timezone.utc)
q = db.query(PaymentOrder).filter(
PaymentOrder.status == "pending",
PaymentOrder.expires_at.isnot(None),
PaymentOrder.expires_at < now,
)
if user_id:
q = q.filter(PaymentOrder.user_id == user_id)
if payment_method:
q = q.filter(PaymentOrder.payment_method == payment_method)
return int(q.update({PaymentOrder.status: "expired"}, synchronize_session=False) or 0)
@staticmethod
def list_callbacks(
db: Session,
*,
payment_method: str | None = None,
limit: int = 50,
offset: int = 0,
) -> tuple[list[PaymentCallback], int]:
q = db.query(PaymentCallback)
if payment_method:
q = q.filter(PaymentCallback.payment_method == payment_method)
total = q.count()
items = q.order_by(PaymentCallback.created_at.desc()).offset(offset).limit(limit).all()
return items, total
@staticmethod
def get_user_order(
db: Session,
*,
user_id: str,
order_id: str,
) -> PaymentOrder | None:
return (
db.query(PaymentOrder)
.filter(PaymentOrder.id == order_id, PaymentOrder.user_id == user_id)
.first()
)
@classmethod
def fail_order(
cls,
db: Session,
*,
order: PaymentOrder,
reason: str | None = None,
) -> PaymentOrder:
locked_order = (
db.query(PaymentOrder)
.filter(PaymentOrder.id == order.id)
.with_for_update()
.one_or_none()
)
if locked_order is None:
raise ValueError("payment order not found")
if locked_order.status == "credited":
raise ValueError("credited order cannot be failed")
locked_order.status = "failed"
payload = dict(locked_order.gateway_response or {})
if reason:
payload["failure_reason"] = reason
payload["failed_at"] = datetime.now(timezone.utc).isoformat()
locked_order.gateway_response = payload
return locked_order
@classmethod
def expire_order(
cls,
db: Session,
*,
order: PaymentOrder,
reason: str | None = None,
) -> tuple[PaymentOrder, bool]:
locked_order = (
db.query(PaymentOrder)
.filter(PaymentOrder.id == order.id)
.with_for_update()
.one_or_none()
)
if locked_order is None:
raise ValueError("payment order not found")
if locked_order.status == "credited":
raise ValueError("credited order cannot be expired")
if locked_order.status == "expired":
return locked_order, False
if locked_order.status != "pending":
raise ValueError(f"only pending order can be expired: {locked_order.status}")
locked_order.status = "expired"
payload = dict(locked_order.gateway_response or {})
if reason:
payload["expire_reason"] = reason
payload["expired_at"] = datetime.now(timezone.utc).isoformat()
locked_order.gateway_response = payload
return locked_order, True
@classmethod
def log_callback(
cls,
db: Session,
*,
payment_method: str,
callback_key: str,
order_no: str | None = None,
gateway_order_id: str | None = None,
payload: dict[str, Any] | None = None,
signature_valid: bool = False,
status: str = "received",
payment_order: PaymentOrder | None = None,
error_message: str | None = None,
) -> tuple[PaymentCallback, bool]:
existing = (
db.query(PaymentCallback).filter(PaymentCallback.callback_key == callback_key).first()
)
if existing is not None:
return existing, False
callback = PaymentCallback(
payment_order_id=payment_order.id if payment_order else None,
payment_method=payment_method,
callback_key=callback_key,
order_no=order_no,
gateway_order_id=gateway_order_id,
payload_hash=cls._build_payload_hash(payload),
signature_valid=signature_valid,
status=status,
payload=payload,
error_message=error_message,
)
db.add(callback)
db.flush()
return callback, True
@classmethod
def credit_order(
cls,
db: Session,
*,
order: PaymentOrder,
gateway_order_id: str | None = None,
gateway_response: dict[str, Any] | None = None,
pay_amount: Decimal | float | int | str | None = None,
pay_currency: str | None = None,
exchange_rate: Decimal | float | int | str | None = None,
) -> tuple[PaymentOrder, bool]:
locked_order = (
db.query(PaymentOrder)
.filter(PaymentOrder.id == order.id)
.with_for_update()
.one_or_none()
)
if locked_order is None:
raise ValueError("payment order not found")
if locked_order.status == "credited":
return locked_order, False
if locked_order.status in {"failed", "expired", "refunded"}:
raise ValueError(f"payment order is not creditable: {locked_order.status}")
now = datetime.now(timezone.utc)
if locked_order.expires_at is not None and locked_order.expires_at < now:
locked_order.status = "expired"
raise ValueError("payment order expired")
wallet = db.query(Wallet).filter(Wallet.id == locked_order.wallet_id).first()
if wallet is None:
raise ValueError("wallet not found")
if wallet.status != "active":
raise ValueError("wallet is not active")
if gateway_order_id:
locked_order.gateway_order_id = gateway_order_id
if gateway_response is not None:
locked_order.gateway_response = gateway_response
if pay_amount is not None:
locked_order.pay_amount = to_money_decimal(pay_amount)
if pay_currency is not None:
locked_order.pay_currency = pay_currency
if exchange_rate is not None:
locked_order.exchange_rate = to_money_decimal(exchange_rate)
locked_order.status = "paid"
locked_order.paid_at = locked_order.paid_at or now
locked_order.refundable_amount_usd = to_money_decimal(locked_order.amount_usd)
WalletService.create_wallet_transaction(
db,
wallet=wallet,
category="recharge",
reason_code="topup_gateway",
amount=locked_order.amount_usd,
balance_type="recharge",
link_type="payment_order",
link_id=locked_order.id,
description=f"充值到账({locked_order.payment_method})",
)
locked_order.status = "credited"
locked_order.credited_at = now
return locked_order, True
@classmethod
def handle_callback(
cls,
db: Session,
*,
payment_method: str,
callback_key: str,
payload: dict[str, Any] | None,
callback_signature: str | None,
callback_secret: str | None,
order_no: str | None = None,
gateway_order_id: str | None = None,
amount_usd: Decimal | float | int | str | None = None,
pay_amount: Decimal | float | int | str | None = None,
pay_currency: str | None = None,
exchange_rate: Decimal | float | int | str | None = None,
) -> dict[str, Any]:
gateway = get_payment_gateway(payment_method)
verified = gateway.verify_callback_payload(
payload=payload,
callback_signature=callback_signature,
callback_secret=callback_secret,
)
callback, created = cls.log_callback(
db,
payment_method=payment_method,
callback_key=callback_key,
order_no=order_no,
gateway_order_id=gateway_order_id,
payload=payload,
signature_valid=verified,
)
if not created and callback.status == "processed":
return {
"ok": True,
"duplicate": True,
"credited": False,
"order_id": callback.payment_order_id,
}
if not verified:
callback.status = "failed"
callback.error_message = "invalid callback signature"
callback.processed_at = datetime.now(timezone.utc)
return {"ok": False, "duplicate": not created, "error": callback.error_message}
order = cls.get_order(
db,
order_no=order_no or callback.order_no,
gateway_order_id=gateway_order_id or callback.gateway_order_id,
)
if order is None:
callback.status = "failed"
callback.error_message = "payment order not found"
callback.processed_at = datetime.now(timezone.utc)
return {"ok": False, "duplicate": not created, "error": callback.error_message}
callback.payment_order_id = order.id
callback.order_no = order.order_no
callback.gateway_order_id = gateway_order_id or order.gateway_order_id
if amount_usd is None:
callback.status = "failed"
callback.error_message = "callback amount is required"
callback.processed_at = datetime.now(timezone.utc)
return {"ok": False, "duplicate": not created, "error": callback.error_message}
expected = to_money_decimal(order.amount_usd)
actual = to_money_decimal(amount_usd)
if actual != expected:
callback.status = "failed"
callback.error_message = "callback amount mismatch"
callback.processed_at = datetime.now(timezone.utc)
return {"ok": False, "duplicate": not created, "error": callback.error_message}
try:
updated_order, credited = cls.credit_order(
db,
order=order,
gateway_order_id=gateway_order_id,
gateway_response=payload,
pay_amount=pay_amount,
pay_currency=pay_currency,
exchange_rate=exchange_rate,
)
except ValueError as exc:
callback.status = "failed"
callback.error_message = str(exc)
callback.processed_at = datetime.now(timezone.utc)
return {"ok": False, "duplicate": not created, "error": callback.error_message}
callback.status = "processed"
callback.error_message = None
callback.processed_at = datetime.now(timezone.utc)
return {
"ok": True,
"duplicate": not created,
"credited": credited,
"order_id": updated_order.id,
"order_no": updated_order.order_no,
"status": updated_order.status,
"wallet_id": updated_order.wallet_id,
}

View File

@@ -16,6 +16,7 @@ from typing import Any
from sqlalchemy.orm import Session
from src.core.enums import UserRole
from src.core.logger import logger
from src.database import create_session
@@ -80,7 +81,7 @@ class CacheWarmupService:
db = create_session()
# 获取一个管理员用户用于构造 context
admin_user = db.query(DBUser).filter(DBUser.role == "admin").first()
admin_user = db.query(DBUser).filter(DBUser.role == UserRole.ADMIN).first()
if not admin_user:
logger.info("缓存预热: 无管理员用户,跳过仪表盘统计预热")
return True
@@ -138,7 +139,7 @@ class CacheWarmupService:
db = create_session()
# 获取一个管理员用户
admin_user = db.query(DBUser).filter(DBUser.role == "admin").first()
admin_user = db.query(DBUser).filter(DBUser.role == UserRole.ADMIN).first()
if not admin_user:
logger.info("缓存预热: 无管理员用户,跳过每日统计预热")
return True

View File

@@ -94,6 +94,10 @@ class SystemConfigService:
"value": "AI Gateway",
"description": "站点副标题,显示在导航栏品牌名称下方",
},
"default_user_initial_gift_usd": {
"value": 10.0,
"description": "新用户默认初始赠款(美元)",
},
REQUEST_RECORD_LEVEL_KEY: {
"value": RequestRecordLevel.BASIC.value,
"description": "请求记录级别basic(基本信息), headers(含请求/响应头), full(完整请求/响应)",
@@ -143,38 +147,6 @@ class SystemConfigService:
"value": "01:05",
"description": "Provider 自动签到执行时间HH:MM 格式24小时制",
},
"enable_user_quota_reset": {
"value": False,
"description": "是否启用用户配额自动重置任务(按配置时间触发,按周期执行)",
},
"user_quota_reset_time": {
"value": "05:00",
"description": "用户配额自动重置执行时间HH:MM 格式24小时制",
},
"user_quota_reset_interval_days": {
"value": 1,
"description": "用户配额重置周期(天数)",
},
"enable_standalone_key_quota_reset": {
"value": False,
"description": "是否启用独立密钥额度自动重置任务(按配置时间触发,按周期执行)",
},
"standalone_key_quota_reset_time": {
"value": "05:00",
"description": "独立密钥额度自动重置执行时间HH:MM 格式24小时制",
},
"standalone_key_quota_reset_interval_days": {
"value": 1,
"description": "独立密钥额度重置周期(天数)",
},
"standalone_key_quota_reset_mode": {
"value": "all",
"description": "独立密钥额度重置模式all(全部独立密钥) 或 selected(指定密钥)",
},
"standalone_key_quota_reset_key_ids": {
"value": [],
"description": "独立密钥额度重置指定的密钥 ID 列表(仅 mode=selected 时生效)",
},
"provider_priority_mode": {
"value": "provider",
"description": "优先级策略provider(提供商优先模式) 或 global_key(全局Key优先模式)",

View File

@@ -25,7 +25,7 @@ from sqlalchemy import delete, text
from src.core.logger import logger
from src.database import create_session
from src.models.database import ApiKey, AuditLog, Provider, RequestCandidate, Usage
from src.models.database import AuditLog, Provider, RequestCandidate, Usage
from src.services.provider_ops.service import ProviderOpsService
from src.services.system.config import SystemConfigService
from src.services.system.scheduler import get_scheduler
@@ -39,10 +39,6 @@ class MaintenanceScheduler:
# 签到任务的 job_id
CHECKIN_JOB_ID = "provider_checkin"
# 用户配额重置任务的 job_id
USER_QUOTA_RESET_JOB_ID = "user_quota_reset"
# 独立密钥额度重置任务的 job_id
STANDALONE_KEY_QUOTA_RESET_JOB_ID = "standalone_key_quota_reset"
def __init__(self) -> None:
self.running = False
@@ -62,19 +58,6 @@ class MaintenanceScheduler:
finally:
db.close()
def _get_user_quota_reset_time(self) -> tuple[int, int]:
"""获取用户配额重置任务的执行时间
Returns:
(hour, minute) 元组
"""
db = create_session()
try:
time_str = SystemConfigService.get_config(db, "user_quota_reset_time", "05:00")
return self._parse_user_quota_reset_time_string(time_str)
finally:
db.close()
@staticmethod
def _parse_time_string(time_str: str) -> tuple[int, int]:
"""解析时间字符串为 (hour, minute) 元组
@@ -98,26 +81,6 @@ class MaintenanceScheduler:
except (ValueError, IndexError):
return (1, 5)
@staticmethod
def _parse_user_quota_reset_time_string(time_str: str) -> tuple[int, int]:
"""解析用户配额重置时间字符串为 (hour, minute) 元组
Returns:
(hour, minute) 元组,解析失败返回默认值 (5, 0)
"""
try:
if not time_str or ":" not in time_str:
return (5, 0)
parts = time_str.split(":")
hour = int(parts[0])
minute = int(parts[1])
# 验证范围
if 0 <= hour <= 23 and 0 <= minute <= 59:
return (hour, minute)
return (5, 0)
except (ValueError, IndexError):
return (5, 0)
def update_checkin_time(self, time_str: str) -> bool:
"""更新签到任务的执行时间
@@ -141,56 +104,6 @@ class MaintenanceScheduler:
return success
def update_user_quota_reset_time(self, time_str: str) -> bool:
"""更新用户配额重置任务的执行时间
Args:
time_str: HH:MM 格式的时间字符串
Returns:
是否成功更新
"""
hour, minute = self._parse_user_quota_reset_time_string(time_str)
scheduler = get_scheduler()
success = scheduler.reschedule_cron_job(
self.USER_QUOTA_RESET_JOB_ID,
hour=hour,
minute=minute,
)
if success:
logger.info(f"用户配额重置任务时间已更新为: {hour:02d}:{minute:02d}")
return success
def _get_standalone_key_quota_reset_time(self) -> tuple[int, int]:
"""获取独立密钥额度重置任务的执行时间"""
db = create_session()
try:
time_str = SystemConfigService.get_config(
db, "standalone_key_quota_reset_time", "05:00"
)
return self._parse_user_quota_reset_time_string(time_str)
finally:
db.close()
def update_standalone_key_quota_reset_time(self, time_str: str) -> bool:
"""更新独立密钥额度重置任务的执行时间"""
hour, minute = self._parse_user_quota_reset_time_string(time_str)
scheduler = get_scheduler()
success = scheduler.reschedule_cron_job(
self.STANDALONE_KEY_QUOTA_RESET_JOB_ID,
hour=hour,
minute=minute,
)
if success:
logger.info(f"独立密钥额度重置任务时间已更新为: {hour:02d}:{minute:02d}")
return success
def get_checkin_job_info(self) -> dict | None:
"""获取签到任务的信息
@@ -318,26 +231,6 @@ class MaintenanceScheduler:
name="Provider签到",
)
# 用户配额重置任务 - 根据配置时间执行(按周期配置决定是否执行)
quota_reset_hour, quota_reset_minute = self._get_user_quota_reset_time()
scheduler.add_cron_job(
self._scheduled_user_quota_reset,
hour=quota_reset_hour,
minute=quota_reset_minute,
job_id=self.USER_QUOTA_RESET_JOB_ID,
name="用户配额自动重置",
)
# 独立密钥额度重置任务 - 根据配置时间执行(按周期配置决定是否执行)
sk_reset_hour, sk_reset_minute = self._get_standalone_key_quota_reset_time()
scheduler.add_cron_job(
self._scheduled_standalone_key_quota_reset,
hour=sk_reset_hour,
minute=sk_reset_minute,
job_id=self.STANDALONE_KEY_QUOTA_RESET_JOB_ID,
name="独立密钥额度自动重置",
)
# 启动时执行一次初始化任务
asyncio.create_task(self._run_startup_tasks())
@@ -440,14 +333,6 @@ class MaintenanceScheduler:
"""Provider 签到任务(定时调用)"""
await self._perform_provider_checkin()
async def _scheduled_user_quota_reset(self) -> None:
"""用户配额重置任务(定时调用)"""
await self._perform_user_quota_reset()
async def _scheduled_standalone_key_quota_reset(self) -> None:
"""独立密钥额度重置任务(定时调用)"""
await self._perform_standalone_key_quota_reset()
# ========== 实际任务实现 ==========
async def _perform_stats_aggregation(self, backfill: bool = False) -> None:
@@ -858,227 +743,6 @@ class MaintenanceScheduler:
if db is not None:
db.close()
async def _perform_user_quota_reset(self) -> None:
"""执行用户配额自动重置任务
适用范围:
- 未删除is_deleted=false
- 仅对 quota_usd != NULL 的用户生效
"""
db = create_session()
try:
# 检查是否启用用户配额重置
if not SystemConfigService.get_config(db, "enable_user_quota_reset", False):
logger.info("用户配额自动重置已禁用,跳过任务")
return
# 重置周期(天数),不限制上限
interval_value = SystemConfigService.get_config(db, "user_quota_reset_interval_days", 1)
try:
interval_days = int(interval_value)
except Exception:
interval_days = 1
if interval_days < 1:
interval_days = 1
# 滚动计算根据上次执行日APP_TIMEZONE判断是否到期
last_reset_at = SystemConfigService.get_config(db, "user_quota_last_reset_at")
should_run = True
if last_reset_at:
last_dt: datetime | None = None
try:
if isinstance(last_reset_at, str):
last_dt = datetime.fromisoformat(last_reset_at)
except Exception:
last_dt = None
if last_dt is None:
logger.warning("user_quota_last_reset_at 格式无效,视为需要执行一次")
else:
if last_dt.tzinfo is None:
last_dt = last_dt.replace(tzinfo=timezone.utc)
from zoneinfo import ZoneInfo
from src.services.system.scheduler import APP_TIMEZONE
tz = ZoneInfo(APP_TIMEZONE)
now_local = datetime.now(tz)
last_local_date = last_dt.astimezone(tz).date()
days_since_reset = (now_local.date() - last_local_date).days
if days_since_reset < 0:
logger.warning("user_quota_last_reset_at 在未来,跳过本次用户配额自动重置")
should_run = False
elif days_since_reset < interval_days:
logger.info(
f"用户配额自动重置未到周期,跳过任务({days_since_reset}/{interval_days}天)"
)
should_run = False
if not should_run:
return
from src.models.database import User as DBUser
now_utc = datetime.now(timezone.utc)
reset_count = (
db.query(DBUser)
.filter(
DBUser.is_deleted.is_(False),
DBUser.quota_usd.isnot(None),
)
.update(
{
DBUser.used_usd: 0.0,
DBUser.updated_at: now_utc,
},
synchronize_session=False,
)
)
db.commit()
# 记录 last_reset_at成功执行后更新滚动计算用
SystemConfigService.set_config(
db,
"user_quota_last_reset_at",
now_utc.isoformat(),
"用户配额自动重置的上次执行时间UTC内部使用",
)
logger.info(
f"用户配额自动重置完成: interval_days={interval_days}, 重置用户数={reset_count}"
)
except Exception as e:
logger.exception(f"用户配额自动重置任务执行失败: {e}")
try:
db.rollback()
except Exception:
pass
finally:
db.close()
async def _perform_standalone_key_quota_reset(self) -> None:
"""执行独立密钥额度自动重置任务
适用范围:
- is_standalone=True 的密钥
- current_balance_usd != NULL有限额的密钥
- 支持 all全部和 selected指定密钥两种模式
"""
db = create_session()
try:
if not SystemConfigService.get_config(db, "enable_standalone_key_quota_reset", False):
logger.info("独立密钥额度自动重置已禁用,跳过任务")
return
# 重置周期
interval_value = SystemConfigService.get_config(
db, "standalone_key_quota_reset_interval_days", 1
)
try:
interval_days = int(interval_value)
except Exception:
interval_days = 1
if interval_days < 1:
interval_days = 1
# 滚动计算
last_reset_at = SystemConfigService.get_config(db, "standalone_key_quota_last_reset_at")
should_run = True
if last_reset_at:
last_dt: datetime | None = None
try:
if isinstance(last_reset_at, str):
last_dt = datetime.fromisoformat(last_reset_at)
except Exception:
last_dt = None
if last_dt is None:
logger.warning("standalone_key_quota_last_reset_at 格式无效,视为需要执行一次")
else:
if last_dt.tzinfo is None:
last_dt = last_dt.replace(tzinfo=timezone.utc)
from zoneinfo import ZoneInfo
from src.services.system.scheduler import APP_TIMEZONE
tz = ZoneInfo(APP_TIMEZONE)
now_local = datetime.now(tz)
last_local_date = last_dt.astimezone(tz).date()
days_since_reset = (now_local.date() - last_local_date).days
if days_since_reset < 0:
logger.warning("standalone_key_quota_last_reset_at 在未来,跳过本次重置")
should_run = False
elif days_since_reset < interval_days:
logger.info(
f"独立密钥额度自动重置未到周期,跳过任务"
f"{days_since_reset}/{interval_days}天)"
)
should_run = False
if not should_run:
return
# 确定重置范围
reset_mode = SystemConfigService.get_config(
db, "standalone_key_quota_reset_mode", "all"
)
now_utc = datetime.now(timezone.utc)
base_filter = [
ApiKey.is_standalone.is_(True),
ApiKey.current_balance_usd.isnot(None),
]
if reset_mode == "selected":
key_ids = SystemConfigService.get_config(
db, "standalone_key_quota_reset_key_ids", []
)
if not key_ids:
logger.info("独立密钥额度重置模式为 selected 但未选择任何密钥,跳过")
return
base_filter.append(ApiKey.id.in_(key_ids))
reset_count = (
db.query(ApiKey)
.filter(*base_filter)
.update(
{
ApiKey.balance_used_usd: 0.0,
ApiKey.updated_at: now_utc,
},
synchronize_session=False,
)
)
db.commit()
SystemConfigService.set_config(
db,
"standalone_key_quota_last_reset_at",
now_utc.isoformat(),
"独立密钥额度自动重置的上次执行时间UTC内部使用",
)
logger.info(
f"独立密钥额度自动重置完成: mode={reset_mode}, "
f"interval_days={interval_days}, 重置密钥数={reset_count}"
)
except Exception as e:
logger.exception(f"独立密钥额度自动重置任务执行失败: {e}")
try:
db.rollback()
except Exception:
pass
finally:
db.close()
async def _perform_candidate_cleanup(self) -> None:
"""清理过期的 request_candidates 记录"""

View File

@@ -6,6 +6,7 @@ from typing import Any
from sqlalchemy.orm import Session
from src.models.database import ApiKey, Usage, User
from src.services.billing.precision import to_money_decimal
from src.services.system.config import SystemConfigService
from src.services.usage._types import UsageCostInfo
from src.services.usage.error_classifier import classify_error
@@ -206,19 +207,19 @@ def build_usage_params(
"cache_read_input_tokens": cache_read_input_tokens,
"cache_creation_input_tokens_5m": cache_creation_input_tokens_5m,
"cache_creation_input_tokens_1h": cache_creation_input_tokens_1h,
"input_cost_usd": input_cost,
"output_cost_usd": output_cost,
"cache_cost_usd": cache_cost,
"cache_creation_cost_usd": cache_creation_cost,
"cache_read_cost_usd": cache_read_cost,
"request_cost_usd": request_cost,
"total_cost_usd": total_cost,
"actual_input_cost_usd": actual_input_cost,
"actual_output_cost_usd": actual_output_cost,
"actual_cache_creation_cost_usd": actual_cache_creation_cost,
"actual_cache_read_cost_usd": actual_cache_read_cost,
"actual_request_cost_usd": actual_request_cost,
"actual_total_cost_usd": actual_total_cost,
"input_cost_usd": to_money_decimal(input_cost),
"output_cost_usd": to_money_decimal(output_cost),
"cache_cost_usd": to_money_decimal(cache_cost),
"cache_creation_cost_usd": to_money_decimal(cache_creation_cost),
"cache_read_cost_usd": to_money_decimal(cache_read_cost),
"request_cost_usd": to_money_decimal(request_cost),
"total_cost_usd": to_money_decimal(total_cost),
"actual_input_cost_usd": to_money_decimal(actual_input_cost),
"actual_output_cost_usd": to_money_decimal(actual_output_cost),
"actual_cache_creation_cost_usd": to_money_decimal(actual_cache_creation_cost),
"actual_cache_read_cost_usd": to_money_decimal(actual_cache_read_cost),
"actual_request_cost_usd": to_money_decimal(actual_request_cost),
"actual_total_cost_usd": to_money_decimal(actual_total_cost),
"rate_multiplier": actual_rate_multiplier,
"input_price_per_1m": input_price,
"output_price_per_1m": output_price,

View File

@@ -7,10 +7,12 @@ from sqlalchemy.orm import Session
from src.core.logger import logger
from src.models.database import ApiKey, Usage, User
from src.services.billing.precision import to_money_decimal
from src.services.provider_keys.codex_quota_sync_dispatcher import (
dispatch_codex_quota_sync_from_response_headers,
)
from src.services.system.config import SystemConfigService
from src.services.wallet import WalletService
class UsageLifecycleMixin:
@@ -156,41 +158,32 @@ class UsageLifecycleMixin:
- 仅当 billing_status='pending' 时才会生效rowcount==1
- 不在本方法内 commit由调用方决定事务提交时机
"""
from sqlalchemy import update
now = datetime.now(timezone.utc)
cost = float(total_cost_usd)
request_cost = float(request_cost_usd) if request_cost_usd is not None else cost
cost = to_money_decimal(total_cost_usd)
request_cost = to_money_decimal(request_cost_usd) if request_cost_usd is not None else cost
result = db.execute(
update(Usage)
.where(
Usage.request_id == request_id,
Usage.billing_status == "pending",
)
.values(
billing_status="settled",
finalized_at=now,
total_cost_usd=cost,
request_cost_usd=request_cost,
status=status,
status_code=status_code,
error_message=error_message,
response_time_ms=response_time_ms,
)
)
if result.rowcount != 1:
usage = db.query(Usage).filter(Usage.request_id == request_id).with_for_update().first()
if not usage or usage.billing_status != "pending":
return False
usage.billing_status = "settled"
usage.finalized_at = now
usage.total_cost_usd = cost
usage.request_cost_usd = request_cost
usage.status = status
usage.status_code = status_code
usage.error_message = error_message
usage.response_time_ms = response_time_ms
if cost > 0:
WalletService.apply_usage_charge(db, usage=usage, amount_usd=cost)
# 写入审计快照(只在本次 finalize 生效时执行)
usage = db.query(Usage).filter(Usage.request_id == request_id).first()
if usage:
metadata = usage.request_metadata or {}
if billing_snapshot is not None:
metadata["billing_snapshot"] = billing_snapshot
if extra_metadata:
metadata.update(extra_metadata)
usage.request_metadata = cls._sanitize_request_metadata(metadata)
metadata = usage.request_metadata or {}
if billing_snapshot is not None:
metadata["billing_snapshot"] = billing_snapshot
if extra_metadata:
metadata.update(extra_metadata)
usage.request_metadata = cls._sanitize_request_metadata(metadata)
return True
@@ -210,27 +203,20 @@ class UsageLifecycleMixin:
- 仅当 billing_status='pending' 时才会生效rowcount==1
- 不在本方法内 commit由调用方决定事务提交时机
"""
from sqlalchemy import update
now = datetime.now(timezone.utc)
result = db.execute(
update(Usage)
.where(
Usage.request_id == request_id,
Usage.billing_status == "pending",
)
.values(
billing_status="void",
finalized_at=now,
total_cost_usd=0.0,
request_cost_usd=0.0,
status="cancelled",
status_code=status_code,
error_message=reason,
response_time_ms=None,
)
)
return result.rowcount == 1
usage = db.query(Usage).filter(Usage.request_id == request_id).with_for_update().first()
if not usage or usage.billing_status != "pending":
return False
usage.billing_status = "void"
usage.finalized_at = now
usage.total_cost_usd = to_money_decimal(0)
usage.request_cost_usd = to_money_decimal(0)
usage.status = "cancelled"
usage.status_code = status_code
usage.error_message = reason
usage.response_time_ms = None
return True
@classmethod
def finalize_submitted(
@@ -252,17 +238,13 @@ class UsageLifecycleMixin:
"""
异步任务提交成功时的幂等结算。
将 pending 使用记录标记为 settled费用暂时为 0
后续轮询完成后通过 update_settled_billing 更新实际费用
将 pending 使用记录保留为 pending仅补齐已知的 provider/响应信息
后续轮询完成后通过 update_settled_billing 一次性写入实际费用并扣钱包
约定:
- 仅当 billing_status='pending' 时才会生效rowcount==1
- 不在本方法内 commit由调用方决定事务提交时机
"""
from sqlalchemy import update
now = datetime.now(timezone.utc)
# 处理响应头和响应体
should_log_headers = SystemConfigService.should_log_headers(db)
should_log_body = SystemConfigService.should_log_body(db)
@@ -284,11 +266,7 @@ class UsageLifecycleMixin:
)
values: dict[str, Any] = {
"billing_status": "settled",
"finalized_at": now,
"total_cost_usd": 0.0,
"request_cost_usd": 0.0,
"status": "completed",
"status": "pending",
"status_code": status_code,
"response_time_ms": response_time_ms,
"provider_name": provider_name,
@@ -305,15 +283,12 @@ class UsageLifecycleMixin:
if processed_response_body is not None:
values["response_body"] = processed_response_body
result = db.execute(
update(Usage)
.where(
Usage.request_id == request_id,
Usage.billing_status == "pending",
)
.values(**values)
)
finalized = result.rowcount == 1
usage = db.query(Usage).filter(Usage.request_id == request_id).with_for_update().first()
if not usage or usage.billing_status != "pending":
return False
for key, value in values.items():
setattr(usage, key, value)
finalized = True
if finalized:
dispatch_codex_quota_sync_from_response_headers(
provider_api_key_id=provider_api_key_id,
@@ -338,54 +313,51 @@ class UsageLifecycleMixin:
extra_metadata: dict[str, Any] | None = None,
) -> bool:
"""
更新已结算记录的计费信息(用于异步任务轮询完成后)。
写入异步任务最终账单(轮询完成后调用)。
与 finalize_settled 不同
- finalize_settled: pending -> settled首次结算
- update_settled_billing: settled -> settled更新费用
语义
- 正常路径:pending -> settled / void(首次最终结算)
- 补写路径:已写入 0 成本但尚未扣钱包的记录,可补写一次最终值
- 已 void 的记录不可再结算
- 已扣钱包wallet_balance_after 已存在)的记录不可重复扣费
约定:
- 仅当 billing_status='settled' 时才会生效
- 不在本方法内 commit由调用方决定事务提交时机
"""
from sqlalchemy import update
now = datetime.now(timezone.utc)
cost = float(total_cost_usd)
request_cost = float(request_cost_usd) if request_cost_usd is not None else cost
cost = to_money_decimal(total_cost_usd)
request_cost = to_money_decimal(request_cost_usd) if request_cost_usd is not None else cost
values: dict[str, Any] = {
"total_cost_usd": cost,
"request_cost_usd": request_cost,
"status": status,
"status_code": status_code,
}
if error_message is not None:
values["error_message"] = error_message
if response_time_ms is not None:
values["response_time_ms"] = response_time_ms
result = db.execute(
update(Usage)
.where(
Usage.request_id == request_id,
Usage.billing_status == "settled",
)
.values(**values)
)
if result.rowcount != 1:
usage = db.query(Usage).filter(Usage.request_id == request_id).with_for_update().first()
if not usage or usage.billing_status == "void":
return False
if usage.billing_status == "settled" and usage.wallet_balance_after is not None:
return False
usage.total_cost_usd = cost
usage.request_cost_usd = request_cost
usage.status = status
usage.status_code = status_code
if error_message is not None:
usage.error_message = error_message
if response_time_ms is not None:
usage.response_time_ms = response_time_ms
usage.finalized_at = usage.finalized_at or now
if cost > 0:
usage.billing_status = "settled"
WalletService.apply_usage_charge(db, usage=usage, amount_usd=cost)
else:
usage.billing_status = "void" if status in {"failed", "cancelled"} else "settled"
# 写入审计快照
usage = db.query(Usage).filter(Usage.request_id == request_id).first()
if usage:
metadata = usage.request_metadata or {}
if billing_snapshot is not None:
metadata["billing_snapshot"] = billing_snapshot
if extra_metadata:
metadata.update(extra_metadata)
metadata["billing_updated_at"] = now.isoformat()
usage.request_metadata = cls._sanitize_request_metadata(metadata)
metadata = usage.request_metadata or {}
if billing_snapshot is not None:
metadata["billing_snapshot"] = billing_snapshot
if extra_metadata:
metadata.update(extra_metadata)
metadata["billing_updated_at"] = now.isoformat()
usage.request_metadata = cls._sanitize_request_metadata(metadata)
return True
@@ -409,26 +381,22 @@ class UsageLifecycleMixin:
- 仅当 billing_status='settled' 时才会生效
- 不在本方法内 commit由调用方决定事务提交时机
"""
from sqlalchemy import update
now = datetime.now(timezone.utc)
result = db.execute(
update(Usage)
.where(
Usage.request_id == request_id,
Usage.billing_status == "settled",
)
.values(
billing_status="void",
finalized_at=now,
total_cost_usd=0.0,
request_cost_usd=0.0,
status="cancelled",
status_code=status_code,
error_message=reason,
)
)
return result.rowcount == 1
usage = db.query(Usage).filter(Usage.request_id == request_id).with_for_update().first()
if not usage or usage.billing_status != "settled":
return False
if usage.wallet_balance_after is not None and to_money_decimal(usage.total_cost_usd) > 0:
# 已实际扣费的记录当前不做自动回滚,避免 silent inconsistency。
return False
usage.billing_status = "void"
usage.finalized_at = now
usage.total_cost_usd = to_money_decimal(0)
usage.request_cost_usd = to_money_decimal(0)
usage.status = "cancelled"
usage.status_code = status_code
usage.error_message = reason
return True
@classmethod
def update_usage_status(
@@ -549,11 +517,14 @@ class UsageLifecycleMixin:
db, provider_request_body, is_request=True
)
# 结算状态:当请求进入终态时,将 billing_status 标记为 settled
# 注意:取消是否应 VOID/部分结算由更高层策略决定;这里默认终态均视为已结算。
if status in ("completed", "failed", "cancelled"):
if getattr(usage, "billing_status", None) == "pending":
usage.billing_status = "settled"
# 仅在“明确不会收费”的终态下直接关闭账单。
# completed 的费用通常要由后续 record_usage / update_settled_billing 写入,
# 这里不能提前把 billing_status 置为 settled否则会阻断真正扣费。
if (
status in ("failed", "cancelled")
and getattr(usage, "billing_status", None) == "pending"
):
usage.billing_status = "void"
if getattr(usage, "finalized_at", None) is None:
usage.finalized_at = datetime.now(timezone.utc)

View File

@@ -7,7 +7,7 @@ from sqlalchemy import case, func
from sqlalchemy.orm import Session
from src.core.logger import logger
from src.models.database import ApiKey, Usage, User, UserRole
from src.models.database import ApiKey, Usage, User
class UsageQueryMixin:
@@ -125,66 +125,44 @@ class UsageQueryMixin:
return result
@staticmethod
def check_user_quota(
def check_request_balance(
db: Session,
user: User,
estimated_tokens: int = 0,
estimated_cost: float = 0,
api_key: ApiKey | None = None,
) -> tuple[bool, str]:
"""检查用户配额或独立Key余额
"""检查请求是否满足余额条件(支持独立 Key)。"""
from src.services.wallet import WalletService
Args:
db: 数据库会话
user: 用户对象
estimated_tokens: 预估token数
estimated_cost: 预估费用
api_key: API Key对象用于检查独立余额Key
wallet_access = WalletService.check_request_allowed(
db,
user=None if (api_key and api_key.is_standalone) else user,
api_key=api_key,
)
if wallet_access.allowed:
return True, "OK"
Returns:
(是否通过, 消息)
"""
if wallet_access.message == "钱包欠费,请先充值":
if api_key and api_key.is_standalone:
return False, "Key欠费请先调账或充值"
return False, "账户欠费,请先充值"
# 如果是独立余额Key检查Key的余额而不是用户配额
if wallet_access.message == "钱包不可用":
if api_key and api_key.is_standalone:
return False, "Key钱包不可用"
return False, "钱包不可用"
remaining = float(wallet_access.remaining) if wallet_access.remaining is not None else None
if api_key and api_key.is_standalone:
# 导入 ApiKeyService 以使用统一的余额计算方法
from src.services.user.apikey import ApiKeyService
if remaining is None:
return False, "Key余额不足"
return False, f"Key余额不足剩余: ${remaining:.2f}"
# NULL 表示无限制
if api_key.current_balance_usd is None:
return True, "OK"
# 使用统一的余额计算方法
remaining_balance = ApiKeyService.get_remaining_balance(api_key)
if remaining_balance is None:
return True, "OK"
# 检查余额是否充足
if remaining_balance < estimated_cost:
return (
False,
f"Key余额不足剩余: ${remaining_balance:.2f},需要: ${estimated_cost:.2f}",
)
return True, "OK"
# 普通Key检查用户配额
# 管理员无限制
if user.role == UserRole.ADMIN:
return True, "OK"
# NULL 表示无限制
if user.quota_usd is None:
return True, "OK"
# 有配额限制,检查是否超额
used_usd = float(user.used_usd or 0)
quota_usd = float(user.quota_usd)
if used_usd + estimated_cost > quota_usd:
remaining = quota_usd - used_usd
return False, f"配额不足(剩余: ${remaining:.2f}"
return True, "OK"
# admin 已在 WalletService.check_request_allowed 中放行,此处不再重复检查
if remaining is None:
return False, wallet_access.message or "余额不足"
return False, f"余额不足(剩余: ${remaining:.2f}"
@staticmethod
def get_usage_summary(
@@ -210,14 +188,14 @@ class UsageQueryMixin:
if end_date:
query = query.filter(Usage.created_at < end_date)
# 使用跨数据库兼容的日期函数
# 使用跨数据库可用的日期函数
from src.utils.database_helpers import date_trunc_portable
# 检测数据库方言
bind = db.bind
dialect = bind.dialect.name if bind is not None else "sqlite"
# 根据分组类型选择日期函数(兼容多种数据库)
# 根据分组类型选择日期函数(适配多种数据库)
if group_by == "day":
date_func = date_trunc_portable(dialect, "day", Usage.created_at)
elif group_by == "week":

View File

@@ -8,6 +8,7 @@ from sqlalchemy.orm import Session
from src.core.logger import logger
from src.models.database import ApiKey, Provider, ProxyNode, Usage, User, UserModelUsageCount
from src.services.billing.precision import to_money_decimal
from src.services.provider_keys.codex_quota_sync_dispatcher import (
dispatch_codex_quota_sync_from_response_headers,
)
@@ -20,6 +21,7 @@ from src.services.usage._recording_helpers import (
update_existing_usage,
)
from src.services.usage._types import UsageCostInfo, UsageRecordParams
from src.services.wallet import WalletService
def _extract_manual_proxy_node_id(metadata: dict[str, Any] | None) -> str | None:
@@ -69,12 +71,12 @@ def _increment_proxy_node_requests(
class UsageRecordingMixin(UsageBillingIntegrationMixin):
"""记录用量相关方法"""
# Metadata pruning configuration -- re-export from helpers for backward compatibility
# Metadata pruning configuration
_METADATA_PRUNE_KEYS: tuple[str, ...] = METADATA_PRUNE_KEYS
_METADATA_KEEP_KEYS: frozenset[str] = METADATA_KEEP_KEYS
# ------------------------------------------------------------------
# Backward-compatible thin wrappers
# Helper wrappers
# ------------------------------------------------------------------
@staticmethod
@@ -121,6 +123,56 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
"""元数据清理(委托到模块级函数)"""
return sanitize_request_metadata(metadata)
@staticmethod
def _is_terminal_status(status: str | None) -> bool:
return status in {"completed", "failed", "cancelled"}
@staticmethod
def _is_usage_finalized(usage: Usage) -> bool:
return (
getattr(usage, "billing_status", None) in {"settled", "void"}
and getattr(usage, "finalized_at", None) is not None
)
@classmethod
def _finalize_usage_billing(
cls,
db: Session,
*,
usage: Usage,
total_cost: float,
status: str | None,
finalized_at: datetime | None = None,
) -> tuple[bool, bool]:
"""完成 usage 的结算状态,并在需要时扣减钱包。
Returns:
(是否首次进入终态, 是否发生扣费)
"""
if not cls._is_terminal_status(status):
if getattr(usage, "billing_status", None) is None:
usage.billing_status = "pending"
return False, False
if (
getattr(usage, "billing_status", None) in {"settled", "void"}
and getattr(usage, "finalized_at", None) is not None
):
return False, False
now = finalized_at or datetime.now(timezone.utc)
charge_amount = to_money_decimal(total_cost)
usage.finalized_at = usage.finalized_at or now
if charge_amount > 0:
WalletService.apply_usage_charge(db, usage=usage, amount_usd=charge_amount)
usage.billing_status = "settled"
return True, True
usage.billing_status = "void" if status in {"failed", "cancelled"} else "settled"
return True, False
# ------------------------------------------------------------------
# Recording methods
# ------------------------------------------------------------------
@@ -221,7 +273,8 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
use_tiered_pricing=use_tiered_pricing,
target_model=target_model,
)
usage_params, _ = await cls._prepare_usage_record(params)
usage_params, total_cost = await cls._prepare_usage_record(params)
total_cost = to_money_decimal(total_cost)
# 创建 Usage 记录
usage = Usage(**usage_params)
@@ -243,17 +296,19 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
# 更新 Provider 月度使用量(原子操作)
if provider_id:
actual_total_cost = usage_params["actual_total_cost_usd"]
actual_total_cost = float(usage_params["actual_total_cost_usd"])
db.execute(
update(Provider)
.where(Provider.id == provider_id)
.values(monthly_used_usd=Provider.monthly_used_usd + actual_total_cost)
)
# 结算标记record_usage_async 写入的 Usage 通常为终态记录
if status not in ("pending", "streaming"):
usage.billing_status = "settled"
usage.finalized_at = datetime.now(timezone.utc)
cls._finalize_usage_billing(
db,
usage=usage,
total_cost=total_cost,
status=status,
)
dispatch_codex_quota_sync_from_response_headers(
provider_api_key_id=provider_api_key_id,
@@ -363,10 +418,20 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
target_model=target_model,
)
usage_params, total_cost = await cls._prepare_usage_record(params)
total_cost = to_money_decimal(total_cost)
# 检查是否已存在相同 request_id 的记录
existing_usage = db.query(Usage).filter(Usage.request_id == request_id).first()
existing_usage = (
db.query(Usage).filter(Usage.request_id == request_id).with_for_update().first()
)
if existing_usage:
if cls._is_usage_finalized(existing_usage):
logger.debug(
"request_id {} 已完成结算,跳过重复记账 (billing_status={})",
request_id,
getattr(existing_usage, "billing_status", None),
)
return existing_usage
logger.debug(
f"request_id {request_id} 已存在,更新现有记录 "
f"(status: {existing_usage.status} -> {status})"
@@ -389,75 +454,52 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
from src.models.database import ApiKey as ApiKeyModel
from src.models.database import GlobalModel
from src.models.database import User as UserModel
# 更新用户使用量(独立 Key 不计入创建者的使用记录)
if user and not (api_key and api_key.is_standalone):
db.execute(
update(UserModel)
.where(UserModel.id == user.id)
.values(
used_usd=UserModel.used_usd + total_cost,
total_usd=UserModel.total_usd + total_cost,
updated_at=sql_func.now(),
)
)
# 更新 API 密钥使用量
if api_key:
if api_key.is_standalone:
db.execute(
update(ApiKeyModel)
.where(ApiKeyModel.id == api_key.id)
.values(
total_requests=ApiKeyModel.total_requests + 1,
total_cost_usd=ApiKeyModel.total_cost_usd + total_cost,
balance_used_usd=ApiKeyModel.balance_used_usd + total_cost,
last_used_at=sql_func.now(),
updated_at=sql_func.now(),
)
)
else:
db.execute(
update(ApiKeyModel)
.where(ApiKeyModel.id == api_key.id)
.values(
total_requests=ApiKeyModel.total_requests + 1,
total_cost_usd=ApiKeyModel.total_cost_usd + total_cost,
last_used_at=sql_func.now(),
updated_at=sql_func.now(),
)
)
# 更新 GlobalModel 使用计数
db.execute(
update(GlobalModel)
.where(GlobalModel.name == model)
.values(usage_count=GlobalModel.usage_count + 1)
accounted, charge_applied = cls._finalize_usage_billing(
db,
usage=usage,
total_cost=total_cost,
status=status,
)
# 更新用户-模型调用次数计数器
cls._increment_user_model_usage(db, user, model)
if accounted:
# 更新 API 密钥使用量
if api_key:
values: dict[str, Any] = {
"total_requests": ApiKeyModel.total_requests + 1,
"last_used_at": sql_func.now(),
"updated_at": sql_func.now(),
}
if charge_applied:
values["total_cost_usd"] = ApiKeyModel.total_cost_usd + float(
to_money_decimal(total_cost)
)
db.execute(update(ApiKeyModel).where(ApiKeyModel.id == api_key.id).values(**values))
# 更新 Provider 月度使用量
if provider_id:
actual_total_cost = usage_params["actual_total_cost_usd"]
# 更新 GlobalModel 使用计数
db.execute(
update(Provider)
.where(Provider.id == provider_id)
.values(monthly_used_usd=Provider.monthly_used_usd + actual_total_cost)
update(GlobalModel)
.where(GlobalModel.name == model)
.values(usage_count=GlobalModel.usage_count + 1)
)
# 更新手动代理节点请求计数tunnel 节点由心跳上报,不在此处统计)
manual_node_id = _extract_manual_proxy_node_id(metadata)
if manual_node_id:
failed = {manual_node_id: 1} if status == "failed" else None
_increment_proxy_node_requests(db, {manual_node_id: 1}, failed)
# 更新用户-模型调用次数计数器
cls._increment_user_model_usage(db, user, model)
# 结算标记:终态请求写入 settled + finalized_at
if status not in ("pending", "streaming"):
usage.billing_status = "settled"
usage.finalized_at = datetime.now(timezone.utc)
# 更新 Provider 月度使用量Provider 端真实成本,无论钱包是否扣费)
if provider_id:
actual_total_cost = float(usage_params["actual_total_cost_usd"])
db.execute(
update(Provider)
.where(Provider.id == provider_id)
.values(monthly_used_usd=Provider.monthly_used_usd + actual_total_cost)
)
# 更新手动代理节点请求计数tunnel 节点由心跳上报,不在此处统计)
manual_node_id = _extract_manual_proxy_node_id(metadata)
if manual_node_id:
failed = {manual_node_id: 1} if status == "failed" else None
_increment_proxy_node_requests(db, {manual_node_id: 1}, failed)
dispatch_codex_quota_sync_from_response_headers(
provider_api_key_id=provider_api_key_id,
@@ -542,10 +584,14 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
cache_creation_cost = 0.0
cache_read_cost = 0.0
cache_cost = 0.0
request_cost = (
float(request_cost_usd) if request_cost_usd is not None else float(total_cost_usd)
request_cost_decimal = (
to_money_decimal(request_cost_usd)
if request_cost_usd is not None
else to_money_decimal(total_cost_usd)
)
total_cost = float(total_cost_usd)
total_cost_decimal = to_money_decimal(total_cost_usd)
request_cost = float(request_cost_decimal)
total_cost = float(total_cost_decimal)
usage_params = build_usage_params(
db=db,
@@ -598,10 +644,10 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
),
)
# Upsert并发幂等优先用 billing_status 作为结算闸门
from sqlalchemy import update
existing_usage = db.query(Usage).filter(Usage.request_id == request_id).first()
# Upsert并发幂等锁定 request_id 对应行,避免重复结算
existing_usage = (
db.query(Usage).filter(Usage.request_id == request_id).with_for_update().first()
)
if existing_usage:
# 避免重复记账:若已结算/作废,直接返回(防止并发重复加计数)
if getattr(existing_usage, "billing_status", None) in ("settled", "void"):
@@ -612,25 +658,6 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
)
return existing_usage
# 并发闸门:只有 billing_status='pending' 的那一次调用可以继续
now = datetime.now(timezone.utc)
claim = db.execute(
update(Usage)
.where(
Usage.request_id == request_id,
Usage.billing_status == "pending",
)
.values(billing_status="settled", finalized_at=now)
)
if claim.rowcount != 1:
# 已被其他 worker 抢先处理(或被 VOID
latest = db.query(Usage).filter(Usage.request_id == request_id).first()
return latest or existing_usage
# 同步 ORM 对象(避免后续代码读到旧值)
existing_usage.billing_status = "settled"
existing_usage.finalized_at = now
cls._update_existing_usage(existing_usage, usage_params, target_model)
usage = existing_usage
else:
@@ -649,69 +676,46 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
from src.models.database import ApiKey as ApiKeyModel
from src.models.database import GlobalModel
from src.models.database import User as UserModel
# 更新用户使用量(独立 Key 不计入创建者)
if user and not (api_key and api_key.is_standalone):
db.execute(
update(UserModel)
.where(UserModel.id == user.id)
.values(
used_usd=UserModel.used_usd + total_cost,
total_usd=UserModel.total_usd + total_cost,
updated_at=sql_func.now(),
)
)
# 更新 API 密钥使用量
if api_key:
if api_key.is_standalone:
db.execute(
update(ApiKeyModel)
.where(ApiKeyModel.id == api_key.id)
.values(
total_requests=ApiKeyModel.total_requests + 1,
total_cost_usd=ApiKeyModel.total_cost_usd + total_cost,
balance_used_usd=ApiKeyModel.balance_used_usd + total_cost,
last_used_at=sql_func.now(),
updated_at=sql_func.now(),
)
)
else:
db.execute(
update(ApiKeyModel)
.where(ApiKeyModel.id == api_key.id)
.values(
total_requests=ApiKeyModel.total_requests + 1,
total_cost_usd=ApiKeyModel.total_cost_usd + total_cost,
last_used_at=sql_func.now(),
updated_at=sql_func.now(),
)
)
# 更新 GlobalModel 使用计数
db.execute(
update(GlobalModel)
.where(GlobalModel.name == model)
.values(usage_count=GlobalModel.usage_count + 1)
accounted, charge_applied = cls._finalize_usage_billing(
db,
usage=usage,
total_cost=total_cost,
status=status,
)
# 更新用户-模型调用次数计数器
cls._increment_user_model_usage(db, user, model)
if accounted:
# 更新 API 密钥使用量
if api_key:
values: dict[str, Any] = {
"total_requests": ApiKeyModel.total_requests + 1,
"last_used_at": sql_func.now(),
"updated_at": sql_func.now(),
}
if charge_applied:
values["total_cost_usd"] = ApiKeyModel.total_cost_usd + float(
total_cost_decimal
)
db.execute(update(ApiKeyModel).where(ApiKeyModel.id == api_key.id).values(**values))
# 更新 Provider 月度使用量(使用 actual_total_cost
if provider_id:
actual_total_cost = usage_params["actual_total_cost_usd"]
# 更新 GlobalModel 使用计数
db.execute(
update(Provider)
.where(Provider.id == provider_id)
.values(monthly_used_usd=Provider.monthly_used_usd + actual_total_cost)
update(GlobalModel)
.where(GlobalModel.name == model)
.values(usage_count=GlobalModel.usage_count + 1)
)
# 结算标记record_usage_with_custom_cost 写入/更新的 Usage 通常为终态记录
if status not in ("pending", "streaming"):
usage.billing_status = "settled"
usage.finalized_at = datetime.now(timezone.utc)
# 更新用户-模型调用次数计数器
cls._increment_user_model_usage(db, user, model)
# 更新 Provider 月度使用量Provider 端真实成本,无论钱包是否扣费)
if provider_id:
actual_total_cost = float(usage_params["actual_total_cost_usd"])
db.execute(
update(Provider)
.where(Provider.id == provider_id)
.values(monthly_used_usd=Provider.monthly_used_usd + actual_total_cost)
)
dispatch_codex_quota_sync_from_response_headers(
provider_api_key_id=provider_api_key_id,
@@ -770,7 +774,6 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
from src.models.database import ApiKey as ApiKeyModel
from src.models.database import GlobalModel
from src.models.database import User as UserModel
# 分离需要更新和需要新建的记录
request_ids = [r.get("request_id") for r in records if r.get("request_id")]
@@ -782,15 +785,17 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
# 查询已存在的 Usage 记录(包括 pending/streaming 状态)
from sqlalchemy.orm import selectinload
existing_records = (
existing_query = (
db.query(Usage)
.options(
selectinload(Usage.user),
selectinload(Usage.api_key),
)
.filter(Usage.request_id.in_(request_ids))
.all()
)
if hasattr(existing_query, "with_for_update"):
existing_query = existing_query.with_for_update()
existing_records = existing_query.all()
existing_usages = {u.request_id: u for u in existing_records}
for record in records:
@@ -826,7 +831,6 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
)
usages: list[Usage] = []
user_costs: dict[str, float] = defaultdict(float) # user_id -> total_cost
apikey_stats: dict[str, dict[str, Any]] = defaultdict(
lambda: {"requests": 0, "cost": 0.0, "is_standalone": False}
)
@@ -858,6 +862,7 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
skipped_count = 0
updated_count = 0
inserted_count = 0
total_count = len(all_records)
# 辅助函数:构建 UsageRecordParams
@@ -949,7 +954,6 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
insert_results = prepared_results[len(update_params_list) :]
finalized_at = datetime.now(timezone.utc)
terminal_statuses = {"completed", "failed", "cancelled"}
# 1. 处理需要更新的记录
for i, (record, request_id, params) in enumerate(update_params_list):
@@ -965,42 +969,40 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
# 更新已存在的 Usage 记录
cls._update_existing_usage(existing_usage, usage_params, record.get("target_model"))
# 结算标记pending -> settled幂等闸门由 prefilter 控制)
if (
usage_params.get("status") in terminal_statuses
and getattr(existing_usage, "billing_status", None) == "pending"
):
existing_usage.billing_status = "settled"
if getattr(existing_usage, "finalized_at", None) is None:
existing_usage.finalized_at = finalized_at
accounted, charge_applied = cls._finalize_usage_billing(
db,
usage=existing_usage,
total_cost=total_cost,
status=usage_params.get("status"),
finalized_at=finalized_at,
)
usages.append(existing_usage)
updated_count += 1
# 聚合统计
model_name = record.get("model") or "unknown"
model_counts[model_name] += 1
if user:
user_model_counts[(str(user.id), model_name)] += 1
if accounted:
model_name = record.get("model") or "unknown"
model_counts[model_name] += 1
if user:
user_model_counts[(str(user.id), model_name)] += 1
provider_id = record.get("provider_id")
if provider_id:
actual_cost = usage_params.get("actual_total_cost_usd", 0)
provider_costs[provider_id] += actual_cost
provider_id = record.get("provider_id")
if charge_applied and provider_id:
actual_cost = usage_params.get("actual_total_cost_usd", 0)
provider_costs[provider_id] += actual_cost
if user and not (api_key and api_key.is_standalone):
user_costs[str(user.id)] += total_cost
if api_key:
key_id = str(api_key.id)
apikey_stats[key_id]["requests"] += 1
if charge_applied:
apikey_stats[key_id]["cost"] += total_cost
apikey_stats[key_id]["is_standalone"] = api_key.is_standalone
if api_key:
key_id = str(api_key.id)
apikey_stats[key_id]["requests"] += 1
apikey_stats[key_id]["cost"] += total_cost
apikey_stats[key_id]["is_standalone"] = api_key.is_standalone
manual_nid = _extract_manual_proxy_node_id(record.get("metadata"))
if manual_nid:
proxy_node_counts[manual_nid] += 1
if record.get("status") == "failed":
proxy_node_failed[manual_nid] += 1
manual_nid = _extract_manual_proxy_node_id(record.get("metadata"))
if manual_nid:
proxy_node_counts[manual_nid] += 1
if record.get("status") == "failed":
proxy_node_failed[manual_nid] += 1
provider_api_key_id = record.get("provider_api_key_id")
response_headers = record.get("response_headers")
@@ -1016,10 +1018,7 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
logger.warning("批量记录中更新失败: {}, request_id={}", e, request_id)
continue
# 2. 处理需要新建的记录(批量插入)
insert_mappings: list[dict[str, Any]] = []
insert_request_ids: list[str] = []
# 2. 处理需要新建的记录
for i, (record, request_id, params) in enumerate(insert_params_list):
try:
usage_params, total_cost, exc = insert_results[i]
@@ -1029,45 +1028,43 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
user = params.user
api_key = params.api_key
# 终态记录:补齐 settled/finalized_at非终态确保 billing_status=pending
status = usage_params.get("status")
if status in terminal_statuses:
if usage_params.get("billing_status") in (None, "pending"):
usage_params["billing_status"] = "settled"
usage_params.setdefault("finalized_at", finalized_at)
elif usage_params.get("billing_status") is None:
usage_params["billing_status"] = "pending"
insert_mappings.append(usage_params)
insert_request_ids.append(request_id)
usage = Usage(**usage_params)
db.add(usage)
accounted, charge_applied = cls._finalize_usage_billing(
db,
usage=usage,
total_cost=total_cost,
status=usage_params.get("status"),
finalized_at=finalized_at,
)
usages.append(usage)
inserted_count += 1
# 聚合统计
model_name = record.get("model") or "unknown"
model_counts[model_name] += 1
if user:
user_model_counts[(str(user.id), model_name)] += 1
if accounted:
model_name = record.get("model") or "unknown"
model_counts[model_name] += 1
if user:
user_model_counts[(str(user.id), model_name)] += 1
provider_id = record.get("provider_id")
if provider_id:
actual_cost = usage_params.get("actual_total_cost_usd", 0)
provider_costs[provider_id] += actual_cost
provider_id = record.get("provider_id")
if charge_applied and provider_id:
actual_cost = usage_params.get("actual_total_cost_usd", 0)
provider_costs[provider_id] += actual_cost
# 用户统计(独立 Key 不计入创建者)
if user and not (api_key and api_key.is_standalone):
user_costs[str(user.id)] += total_cost
# API Key 统计
if api_key:
key_id = str(api_key.id)
apikey_stats[key_id]["requests"] += 1
if charge_applied:
apikey_stats[key_id]["cost"] += total_cost
apikey_stats[key_id]["is_standalone"] = api_key.is_standalone
# API Key 统计
if api_key:
key_id = str(api_key.id)
apikey_stats[key_id]["requests"] += 1
apikey_stats[key_id]["cost"] += total_cost
apikey_stats[key_id]["is_standalone"] = api_key.is_standalone
manual_nid = _extract_manual_proxy_node_id(record.get("metadata"))
if manual_nid:
proxy_node_counts[manual_nid] += 1
if record.get("status") == "failed":
proxy_node_failed[manual_nid] += 1
manual_nid = _extract_manual_proxy_node_id(record.get("metadata"))
if manual_nid:
proxy_node_counts[manual_nid] += 1
if record.get("status") == "failed":
proxy_node_failed[manual_nid] += 1
provider_api_key_id = record.get("provider_api_key_id")
response_headers = record.get("response_headers")
@@ -1083,24 +1080,6 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
logger.warning("批量记录中跳过无效记录: {}, request_id={}", e, request_id)
continue
if insert_mappings:
try:
db.bulk_insert_mappings(Usage, insert_mappings)
# 仅用于保持返回值语义:将新建记录读回为 ORM 对象
inserted_records = (
db.query(Usage).filter(Usage.request_id.in_(insert_request_ids)).all()
)
inserted_map = {u.request_id: u for u in inserted_records}
for rid in insert_request_ids:
inserted_usage = inserted_map.get(rid)
if inserted_usage is not None:
usages.append(inserted_usage)
except Exception as e:
logger.error("批量插入 Usage 记录时出错: {}", e)
db.rollback()
raise
# 统计跳过的记录,失败率超过 10% 时提升日志级别
if skipped_count > 0:
skip_ratio = skipped_count / total_count if total_count > 0 else 0
@@ -1152,47 +1131,22 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
db.execute(
update(Provider)
.where(Provider.id == provider_id)
.values(monthly_used_usd=Provider.monthly_used_usd + cost)
)
# 批量更新用户使用量
for user_id, cost in user_costs.items():
if cost > 0:
db.execute(
update(UserModel)
.where(UserModel.id == user_id)
.values(
used_usd=UserModel.used_usd + cost,
total_usd=UserModel.total_usd + cost,
updated_at=sql_func.now(),
)
.values(monthly_used_usd=Provider.monthly_used_usd + float(cost))
)
# 批量更新 API Key 统计
for key_id, stats in apikey_stats.items():
if stats["is_standalone"]:
db.execute(
update(ApiKeyModel)
.where(ApiKeyModel.id == key_id)
.values(
total_requests=ApiKeyModel.total_requests + stats["requests"],
total_cost_usd=ApiKeyModel.total_cost_usd + stats["cost"],
balance_used_usd=ApiKeyModel.balance_used_usd + stats["cost"],
last_used_at=sql_func.now(),
updated_at=sql_func.now(),
)
)
else:
db.execute(
update(ApiKeyModel)
.where(ApiKeyModel.id == key_id)
.values(
total_requests=ApiKeyModel.total_requests + stats["requests"],
total_cost_usd=ApiKeyModel.total_cost_usd + stats["cost"],
last_used_at=sql_func.now(),
updated_at=sql_func.now(),
)
db.execute(
update(ApiKeyModel)
.where(ApiKeyModel.id == key_id)
.values(
total_requests=ApiKeyModel.total_requests + stats["requests"],
total_cost_usd=ApiKeyModel.total_cost_usd
+ float(to_money_decimal(stats["cost"])),
last_used_at=sql_func.now(),
updated_at=sql_func.now(),
)
)
# 批量更新手动代理节点请求计数
_increment_proxy_node_requests(db, proxy_node_counts, proxy_node_failed)
@@ -1208,7 +1162,6 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
# 单次提交所有更改
try:
db.commit()
inserted_count = len(insert_mappings)
total_written = updated_count + inserted_count
if updated_count > 0:
logger.debug("批量记录成功: 更新 {} 条, 新建 {}", updated_count, inserted_count)

View File

@@ -30,7 +30,6 @@ class ApiKeyService:
concurrent_limit: int = 5,
expire_days: int | None = None,
expires_at: datetime | None = None, # 直接传入过期时间,优先于 expire_days
initial_balance_usd: float | None = None,
is_standalone: bool = False,
auto_delete_on_expiry: bool = False,
) -> tuple[ApiKey, str]:
@@ -47,7 +46,6 @@ class ApiKeyService:
concurrent_limit: 并发限制
expire_days: 过期天数None = 永不过期
expires_at: 直接指定过期时间,优先于 expire_days
initial_balance_usd: 初始余额USD仅用于独立KeyNone = 无限制
is_standalone: 是否为独立余额Key仅管理员可创建
auto_delete_on_expiry: 过期后是否自动删除True=物理删除False=仅禁用)
"""
@@ -74,8 +72,6 @@ class ApiKeyService:
rate_limit=rate_limit,
concurrent_limit=concurrent_limit,
expires_at=final_expires_at,
balance_used_usd=0.0,
current_balance_usd=initial_balance_usd, # 直接使用初始余额None = 无限制
is_standalone=is_standalone,
auto_delete_on_expiry=auto_delete_on_expiry,
is_active=True,
@@ -87,7 +83,7 @@ class ApiKeyService:
logger.info(
f"创建API密钥: 用户ID {user_id}, 密钥名 {api_key.name}, "
f"独立Key={is_standalone}, 初始余额={initial_balance_usd}"
f"独立Key={is_standalone}"
)
return api_key, key # 返回密钥对象和明文密钥
@@ -143,7 +139,6 @@ class ApiKeyService:
"concurrent_limit",
"is_active",
"expires_at",
"balance_limit_usd",
"auto_delete_on_expiry",
]
@@ -188,48 +183,6 @@ class ApiKeyService:
logger.info(f"删除API密钥: ID {key_id}")
return True
@staticmethod
def get_remaining_balance(api_key: ApiKey) -> float | None:
"""计算剩余余额仅用于独立Key
Returns:
剩余余额None 表示无限制或非独立Key
"""
if not api_key.is_standalone:
return None
if api_key.current_balance_usd is None:
return None
# 剩余余额 = 当前余额 - 已使用余额
remaining = api_key.current_balance_usd - (api_key.balance_used_usd or 0)
return max(0, remaining) # 不能为负数
@staticmethod
def check_balance(api_key: ApiKey) -> tuple[bool, float | None]:
"""检查余额限制仅用于独立Key
Returns:
(is_allowed, remaining_balance): 是否允许请求剩余余额None表示无限制
"""
if not api_key.is_standalone:
# 非独立Key不检查余额
return True, None
# 使用新的预付费模式: current_balance_usd
if api_key.current_balance_usd is None:
# 无余额限制
return True, None
# 使用统一的余额计算方法
remaining = ApiKeyService.get_remaining_balance(api_key)
is_allowed = remaining > 0 if remaining is not None else True
if not is_allowed:
logger.warning(f"API密钥余额不足: Key ID {api_key.id}, " f"剩余余额 ${remaining:.4f}")
return is_allowed, remaining
@staticmethod
def check_rate_limit(db: Session, api_key: ApiKey, window_minutes: int = 1) -> tuple[bool, int]:
"""检查速率限制
@@ -263,57 +216,6 @@ class ApiKeyService:
return is_allowed, api_key.rate_limit - request_count
@staticmethod
def add_balance(db: Session, key_id: str, amount_usd: float) -> ApiKey | None:
"""为独立余额Key调整余额
Args:
db: 数据库会话
key_id: API Key ID
amount_usd: 要调整的余额金额USD正数为增加负数为扣除
Returns:
更新后的API Key对象如果Key不存在或不是独立Key则返回None
"""
api_key = db.query(ApiKey).filter(ApiKey.id == key_id).first()
if not api_key:
logger.warning(f"余额调整失败: Key ID {key_id} 不存在")
return None
if not api_key.is_standalone:
logger.warning(f"余额调整失败: Key ID {key_id} 不是独立余额Key")
return None
if amount_usd == 0:
logger.warning(f"余额调整失败: 调整金额不能为0当前值 ${amount_usd}")
return None
# 如果是扣除(负数),检查是否超过当前余额
if amount_usd < 0:
current = api_key.current_balance_usd or 0
if abs(amount_usd) > current:
logger.warning(
f"余额扣除失败: 扣除金额 ${abs(amount_usd):.4f} 超过当前余额 ${current:.4f}"
)
return None
# 调整当前余额
if api_key.current_balance_usd is None:
api_key.current_balance_usd = amount_usd if amount_usd > 0 else 0
else:
api_key.current_balance_usd = max(0, api_key.current_balance_usd + amount_usd)
api_key.updated_at = datetime.now(timezone.utc)
db.commit()
db.refresh(api_key)
action = "增加" if amount_usd > 0 else "扣除"
logger.info(
f"余额调整成功: Key ID {key_id}, {action} ${abs(amount_usd):.4f}, "
f"新余额 ${api_key.current_balance_usd:.4f}"
)
return api_key
@staticmethod
def cleanup_expired_keys(db: Session, auto_delete: bool = False) -> int:
"""清理过期的API密钥

View File

@@ -9,6 +9,7 @@ from sqlalchemy.orm import Session
from src.core.exceptions import NotFoundException
from src.core.logger import logger
from src.models.database import Provider, User, UserPreference
from src.services.wallet import WalletService
class PreferenceService:
@@ -98,6 +99,8 @@ class PreferenceService:
raise NotFoundException("User not found")
preferences = PreferenceService.get_or_create_preferences(db, user_id)
wallet = WalletService.get_wallet(db, user_id=user.id)
billing = WalletService.serialize_wallet_summary(wallet)
# 构建返回数据
user_data = {
@@ -125,12 +128,10 @@ class PreferenceService:
"announcements": preferences.announcement_notifications,
},
},
# 配额信息
"quota_usd": user.quota_usd,
"used_usd": user.used_usd,
"billing": billing,
"stats": {
"total_cost": user.used_usd,
"total_cost_all_time": user.total_usd,
"total_cost": billing["total_consumed"],
"total_cost_all_time": billing["total_consumed"],
"api_keys_count": len(user.api_keys),
},
}

View File

@@ -8,7 +8,7 @@ import asyncio
from datetime import datetime, timezone
from typing import Any
from sqlalchemy import and_, func
from sqlalchemy import and_, func, or_
from sqlalchemy.orm import Session
from src.core.logger import logger
@@ -30,13 +30,14 @@ class UserService:
username: str,
password: str,
role: UserRole = UserRole.USER,
quota_usd: float | None = 10.0,
initial_gift_usd: float | None = 10.0,
unlimited: bool = False,
email_verified: bool = False,
allowed_providers: list[str] | None = None,
allowed_api_formats: list[str] | None = None,
allowed_models: list[str] | None = None,
) -> User:
"""创建新用户quota_usd 为 None 表示无限制email 为 None 表示无邮箱"""
"""创建新用户"""
# 验证邮箱格式(仅当提供邮箱时)
if email is not None:
@@ -66,7 +67,6 @@ class UserService:
email_verified=email_verified if email else False,
username=username,
role=role,
quota_usd=quota_usd,
is_active=True,
allowed_providers=allowed_providers,
allowed_api_formats=allowed_api_formats,
@@ -75,6 +75,18 @@ class UserService:
user.set_password(password)
db.add(user)
db.flush()
from src.services.wallet import WalletService
WalletService.initialize_user_wallet(
db,
user=user,
initial_gift_usd=initial_gift_usd,
unlimited=unlimited,
description="用户初始赠款",
)
db.commit() # 立即提交事务,释放数据库锁
db.refresh(user)
@@ -91,7 +103,8 @@ class UserService:
password: str,
api_key_name: str = "默认密钥",
role: UserRole = UserRole.USER,
quota_usd: float | None = 10.0,
initial_gift_usd: float | None = 10.0,
unlimited: bool = False,
concurrent_limit: int = 5,
) -> tuple[User, ApiKey]:
"""
@@ -104,7 +117,8 @@ class UserService:
password: 密码
api_key_name: API密钥名称
role: 用户角色
quota_usd: USD配额None 表示无限制
initial_gift_usd: 初始赠款USD
unlimited: 是否无限制
concurrent_limit: 并发限制
Returns:
@@ -115,7 +129,13 @@ class UserService:
"""
# 创建用户
user = UserService.create_user(
db=db, email=email, username=username, password=password, role=role, quota_usd=quota_usd
db=db,
email=email,
username=username,
password=password,
role=role,
initial_gift_usd=initial_gift_usd,
unlimited=unlimited,
)
# 导入API密钥服务避免循环导入
@@ -173,7 +193,9 @@ class UserService:
if is_active is not None:
query = query.filter(User.is_active == is_active)
return query.offset(skip).limit(limit).all()
return (
query.order_by(User.created_at.desc(), User.id.desc()).offset(skip).limit(limit).all()
)
@staticmethod
@transactional()
@@ -187,7 +209,6 @@ class UserService:
updatable_fields = [
"email",
"username",
"quota_usd",
"is_active",
"role",
# 访问限制字段
@@ -198,7 +219,6 @@ class UserService:
# 允许设置为 None 的字段(表示无限制)
nullable_fields = [
"quota_usd",
"allowed_providers",
"allowed_api_formats",
"allowed_models",
@@ -237,16 +257,19 @@ class UserService:
"""删除用户(硬删除)
删除流程:
1. 手动删除关联的子记录(避免 SQLAlchemy ORM 与数据库 CASCADE 冲突)
2. 删除用户记录
3. 历史 Usage 记录保留user_id 会被数据库设为 NULL
4. 新用户注册时会有新的 UUID看不到旧用户的记录
1. 检查未完结账务,阻止删除
2. 手动删除 ORM cascade 冲突的子记录
3. 删除用户记录
4. 财务记录Wallet/PaymentOrder/RefundRequest/WalletTransaction
Usage 记录保留,外键 SET NULL由自动清理策略统一回收
"""
from src.models.database import (
AnnouncementRead,
ApiKey,
PaymentOrder,
RefundRequest,
UserPreference,
UserQuota,
Wallet,
)
user = db.query(User).filter(User.id == user_id).first()
@@ -256,16 +279,53 @@ class UserService:
# 记录删除信息用于日志
email = user.email
# 删除前阻断未完结账务,避免删除导致资金状态不一致。
wallet_ids = [
wallet_id
for (wallet_id,) in (
db.query(Wallet.id)
.outerjoin(ApiKey, Wallet.api_key_id == ApiKey.id)
.filter(or_(Wallet.user_id == user_id, ApiKey.user_id == user_id))
.all()
)
]
if wallet_ids:
pending_refund_count = (
db.query(RefundRequest)
.filter(
RefundRequest.wallet_id.in_(wallet_ids),
RefundRequest.status.in_(["pending_approval", "approved", "processing"]),
)
.count()
)
if pending_refund_count > 0:
raise ValueError("用户存在未完结退款,禁止删除")
pending_order_count = (
db.query(PaymentOrder)
.filter(
PaymentOrder.wallet_id.in_(wallet_ids),
PaymentOrder.status.in_(["pending", "paid"]),
)
.count()
)
if pending_order_count > 0:
raise ValueError("用户存在未完结充值订单,禁止删除")
# 手动删除子记录,避免 SQLAlchemy 的 ORM cascade 与数据库 CASCADE 冲突
# 这些表的数据库外键已经设置了 ON DELETE CASCADE但 SQLAlchemy 会先尝试 UPDATE 设置为 NULL
# 所以我们手动删除来避免这个问题
# UserPreference/AnnouncementRead 的数据库外键 ON DELETE CASCADE
# 但 SQLAlchemy 会先尝试 UPDATE SET NULL 导致冲突)
db.query(UserPreference).filter(UserPreference.user_id == user_id).delete(
synchronize_session=False
)
db.query(UserQuota).filter(UserQuota.user_id == user_id).delete(synchronize_session=False)
db.query(AnnouncementRead).filter(AnnouncementRead.user_id == user_id).delete(
synchronize_session=False
)
# 财务记录Wallet/WalletTransaction/PaymentOrder/RefundRequest/PaymentCallback
# 和 Usage 记录全部保留,数据库外键 SET NULL 自动断开关联,
# 由自动清理策略统一回收。
api_key_count = int(
db.query(func.count(ApiKey.id)).filter(ApiKey.user_id == user_id).scalar() or 0
)
@@ -326,29 +386,6 @@ class UserService:
logger.info(f"密码更改成功: 用户ID {user_id}")
return True, "密码更改成功"
@staticmethod
def update_user_quota(
db: Session,
user_id: str,
quota_usd: float | None = None,
) -> User | None:
"""更新用户配额"""
user = db.query(User).filter(User.id == user_id).first()
if not user:
return None
if quota_usd is not None:
user.quota_usd = quota_usd
db.commit()
db.refresh(user)
# 清除用户缓存
asyncio.create_task(UserCacheService.invalidate_user_cache(user.id, user.email))
logger.debug(f"更新用户配额: {user.email} (USD: {quota_usd})")
return user
@staticmethod
def get_user_usage_stats(
db: Session,

View File

@@ -0,0 +1,3 @@
from src.services.wallet.service import WalletAccessResult, WalletService
__all__ = ["WalletAccessResult", "WalletService"]

View File

@@ -0,0 +1,942 @@
from __future__ import annotations
from dataclasses import dataclass
from datetime import datetime, timezone
from decimal import Decimal
from typing import Literal
from uuid import uuid4
from sqlalchemy import func
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session
from src.core.enums import UserRole
from src.models.database import (
ApiKey,
PaymentOrder,
RefundRequest,
Usage,
User,
Wallet,
WalletTransaction,
)
from src.services.billing.precision import to_money_decimal
WalletCategory = Literal["recharge", "gift", "adjust", "refund"]
WalletBalanceBucket = Literal["recharge", "gift"]
REASON_TOPUP_ADMIN_MANUAL = "topup_admin_manual"
REASON_TOPUP_GATEWAY = "topup_gateway"
REASON_TOPUP_CARD_CODE = "topup_card_code"
REASON_GIFT_INITIAL = "gift_initial"
REASON_GIFT_CAMPAIGN = "gift_campaign"
REASON_GIFT_EXPIRE_RECLAIM = "gift_expire_reclaim"
REASON_ADJUST_ADMIN = "adjust_admin"
REASON_ADJUST_SYSTEM = "adjust_system"
REASON_REFUND_OUT = "refund_out"
REASON_REFUND_REVERT = "refund_revert"
@dataclass(slots=True)
class WalletAccessResult:
allowed: bool
remaining: Decimal | None
message: str
wallet: Wallet | None = None
class WalletService:
"""统一钱包服务。"""
@staticmethod
def get_limit_mode(wallet: Wallet | None) -> str:
if wallet is None:
return "finite"
limit_mode = getattr(wallet, "limit_mode", None)
if limit_mode in {"finite", "unlimited"}:
return str(limit_mode)
return "finite"
@classmethod
def is_unlimited_wallet(cls, wallet: Wallet | None) -> bool:
return cls.get_limit_mode(wallet) == "unlimited"
@classmethod
def get_recharge_balance_value(cls, wallet: Wallet | None) -> Decimal:
if wallet is None:
return Decimal("0")
return to_money_decimal(wallet.balance)
@classmethod
def get_gift_balance_value(cls, wallet: Wallet | None) -> Decimal:
if wallet is None:
return Decimal("0")
return to_money_decimal(getattr(wallet, "gift_balance", None))
@classmethod
def get_spendable_balance_value(cls, wallet: Wallet | None) -> Decimal:
return cls.get_recharge_balance_value(wallet) + cls.get_gift_balance_value(wallet)
@classmethod
def get_refundable_balance_value(cls, wallet: Wallet | None) -> Decimal:
# 赠款余额不可退款,仅充值余额可退。
return cls.get_recharge_balance_value(wallet)
@classmethod
def serialize_wallet_summary(cls, wallet: Wallet | None) -> dict[str, object]:
recharge_balance = cls.get_recharge_balance_value(wallet)
gift_balance = cls.get_gift_balance_value(wallet)
spendable_balance = recharge_balance + gift_balance
limit_mode = cls.get_limit_mode(wallet)
return {
"id": wallet.id if wallet else None,
"balance": float(spendable_balance),
"recharge_balance": float(recharge_balance),
"gift_balance": float(gift_balance),
"refundable_balance": float(recharge_balance),
"currency": wallet.currency if wallet else "USD",
"status": wallet.status if wallet else "active",
"limit_mode": limit_mode,
"unlimited": limit_mode == "unlimited",
"total_recharged": float(wallet.total_recharged or 0) if wallet else 0.0,
"total_consumed": float(wallet.total_consumed or 0) if wallet else 0.0,
"total_refunded": float(wallet.total_refunded or 0) if wallet else 0.0,
"total_adjusted": float(wallet.total_adjusted or 0) if wallet else 0.0,
"updated_at": wallet.updated_at if wallet else None,
}
@staticmethod
def _build_order_no(prefix: str) -> str:
ts = datetime.now(timezone.utc).strftime("%Y%m%d%H%M%S%f")
return f"{prefix}_{ts}_{uuid4().hex[:12]}"
@classmethod
def initialize_user_wallet(
cls,
db: Session,
*,
user: User,
initial_gift_usd: Decimal | float | int | str | None,
unlimited: bool = False,
description: str = "用户初始赠款",
) -> Wallet | None:
"""初始化用户钱包,并按需要写入初始赠款。"""
if not user.id:
return None
wallet = cls.get_wallet(db, user_id=user.id)
if wallet is None:
wallet = Wallet(
user_id=user.id,
balance=Decimal("0"),
gift_balance=Decimal("0"),
total_recharged=Decimal("0"),
total_consumed=Decimal("0"),
total_refunded=Decimal("0"),
total_adjusted=Decimal("0"),
limit_mode="unlimited" if unlimited else "finite",
currency="USD",
status="active",
)
db.add(wallet)
db.flush()
else:
wallet.limit_mode = "unlimited" if unlimited else "finite"
gift_amount = to_money_decimal(initial_gift_usd)
if not unlimited and gift_amount > Decimal("0"):
cls.create_wallet_transaction(
db,
wallet=wallet,
category="gift",
reason_code=REASON_GIFT_INITIAL,
amount=gift_amount,
balance_type="gift",
link_type="system_task",
link_id=user.id,
description=description,
)
return wallet
@classmethod
def initialize_api_key_wallet(
cls,
db: Session,
*,
api_key: ApiKey,
initial_balance_usd: Decimal | float | int | str | None,
unlimited: bool = False,
operator_id: str | None = None,
description: str = "初始调账",
) -> Wallet | None:
"""初始化独立 Key 钱包,并按需执行初始调账。
设计目标:
- 初始化语义与用户钱包保持一致(均由 WalletService 统一入口完成)
- 独立 Key 不支持充值,余额变动统一通过调账流水实现
"""
if not api_key.id:
return None
wallet = cls.get_wallet(db, api_key_id=api_key.id)
if wallet is None:
wallet = Wallet(
api_key_id=api_key.id,
balance=Decimal("0"),
gift_balance=Decimal("0"),
total_recharged=Decimal("0"),
total_consumed=Decimal("0"),
total_refunded=Decimal("0"),
total_adjusted=Decimal("0"),
limit_mode="unlimited" if unlimited else "finite",
currency="USD",
status="active",
)
db.add(wallet)
db.flush()
else:
wallet.limit_mode = "unlimited" if unlimited else "finite"
initial_amount = to_money_decimal(initial_balance_usd)
if not unlimited and initial_amount > Decimal("0"):
cls.create_wallet_transaction(
db,
wallet=wallet,
category="adjust",
reason_code=REASON_ADJUST_SYSTEM,
amount=initial_amount,
balance_type="recharge",
link_type="system_task",
link_id=api_key.id,
operator_id=operator_id,
description=description,
)
return wallet
@classmethod
def get_wallet(
cls,
db: Session,
*,
user_id: str | None = None,
api_key_id: str | None = None,
) -> Wallet | None:
if api_key_id:
wallet = db.query(Wallet).filter(Wallet.api_key_id == api_key_id).first()
if wallet is not None:
return wallet
if user_id:
return db.query(Wallet).filter(Wallet.user_id == user_id).first()
return None
@classmethod
def get_or_create_wallet(
cls,
db: Session,
*,
user: User | None = None,
api_key: ApiKey | None = None,
user_id: str | None = None,
api_key_id: str | None = None,
) -> Wallet | None:
if user is None and user_id:
user = db.query(User).filter(User.id == user_id).first()
if api_key is None and api_key_id:
api_key = db.query(ApiKey).filter(ApiKey.id == api_key_id).first()
owner_user_id = user.id if user else user_id
owner_api_key_id = api_key.id if api_key else api_key_id
# owner 解析规则:
# - 独立 Key: 归属 API Key 钱包
# - 普通 Key + 用户: 归属用户钱包(避免 user_id/api_key_id 同时写入)
# - 仅提供 API Key: 归属 API Key 钱包
api_key_is_standalone = bool(getattr(api_key, "is_standalone", False)) if api_key else False
if owner_user_id is not None and not api_key_is_standalone:
owner_api_key_id = None
elif owner_api_key_id is not None:
owner_user_id = None
wallet = cls.get_wallet(db, user_id=owner_user_id, api_key_id=owner_api_key_id)
if wallet:
return wallet
if owner_user_id is None and owner_api_key_id is None:
return None
bootstrap = Wallet(
user_id=owner_user_id,
api_key_id=owner_api_key_id,
balance=Decimal("0"),
gift_balance=Decimal("0"),
total_recharged=Decimal("0"),
total_consumed=Decimal("0"),
total_refunded=Decimal("0"),
total_adjusted=Decimal("0"),
limit_mode="finite",
currency="USD",
status="active",
created_at=datetime.now(timezone.utc),
updated_at=datetime.now(timezone.utc),
)
try:
with db.begin_nested():
db.add(bootstrap)
db.flush()
return bootstrap
except IntegrityError:
# 并发创建时可能触发唯一约束,回查已创建的钱包并复用。
wallet = cls.get_wallet(db, user_id=owner_user_id, api_key_id=owner_api_key_id)
if wallet is not None:
return wallet
raise
@classmethod
def check_request_allowed(
cls,
db: Session,
*,
user: User | None,
api_key: ApiKey | None = None,
) -> WalletAccessResult:
if user and user.role == UserRole.ADMIN:
return WalletAccessResult(True, None, "OK", None)
wallet = cls.get_or_create_wallet(db, user=user, api_key=api_key)
if wallet is None:
return WalletAccessResult(False, Decimal("0"), "钱包不存在", None)
remaining = cls.get_spendable_balance_value(wallet)
recharge_balance = cls.get_recharge_balance_value(wallet)
if wallet.status != "active":
return WalletAccessResult(False, remaining, "钱包不可用", wallet)
# 充值余额为负视为欠费,禁止继续消费(即使总可用余额仍为正)。
if recharge_balance < Decimal("0"):
return WalletAccessResult(False, recharge_balance, "钱包欠费,请先充值", wallet)
if cls.is_unlimited_wallet(wallet):
return WalletAccessResult(True, None, "OK", wallet)
if remaining <= Decimal("0"):
return WalletAccessResult(False, remaining, "钱包余额不足", wallet)
return WalletAccessResult(True, remaining, "OK", wallet)
@classmethod
def get_balance_snapshot(
cls,
db: Session,
*,
user: User | None,
api_key: ApiKey | None = None,
) -> Decimal | None:
wallet = cls.get_or_create_wallet(db, user=user, api_key=api_key)
if wallet is None:
return None
recharge_balance = cls.get_recharge_balance_value(wallet)
if recharge_balance < Decimal("0"):
return recharge_balance
if cls.is_unlimited_wallet(wallet):
return None
return cls.get_spendable_balance_value(wallet)
@classmethod
def _resolve_wallet_for_usage(cls, db: Session, usage: Usage) -> Wallet | None:
if usage.wallet_id:
wallet = db.query(Wallet).filter(Wallet.id == usage.wallet_id).first()
if wallet:
return wallet
api_key = None
if usage.api_key_id:
api_key = db.query(ApiKey).filter(ApiKey.id == usage.api_key_id).first()
if api_key and api_key.is_standalone:
return cls.get_or_create_wallet(db, api_key=api_key)
if usage.user_id:
user = db.query(User).filter(User.id == usage.user_id).first()
return cls.get_or_create_wallet(db, user=user, api_key=api_key)
return None
@classmethod
def apply_usage_charge(
cls,
db: Session,
*,
usage: Usage,
amount_usd: Decimal | float | int | str,
) -> tuple[Decimal | None, Decimal | None]:
amount = to_money_decimal(amount_usd)
if amount <= Decimal("0"):
return None, None
wallet = cls._resolve_wallet_for_usage(db, usage)
if wallet is None:
return None, None
locked_wallet = (
db.query(Wallet).filter(Wallet.id == wallet.id).with_for_update().one_or_none()
)
if locked_wallet is None:
return None, None
before_recharge = cls.get_recharge_balance_value(locked_wallet)
before_gift = cls.get_gift_balance_value(locked_wallet)
before_total = before_recharge + before_gift
if cls.is_unlimited_wallet(locked_wallet):
locked_wallet.total_consumed = to_money_decimal(locked_wallet.total_consumed) + amount
locked_wallet.updated_at = datetime.now(timezone.utc)
usage.wallet_id = locked_wallet.id
usage.wallet_balance_before = before_total
usage.wallet_balance_after = before_total
usage.wallet_recharge_balance_before = before_recharge
usage.wallet_recharge_balance_after = before_recharge
usage.wallet_gift_balance_before = before_gift
usage.wallet_gift_balance_after = before_gift
return before_total, before_total
# 赠款优先扣减:赠款不可退款,优先消耗可避免与充值余额混淆。
gift_deduction = min(max(before_gift, Decimal("0")), amount)
recharge_deduction = amount - gift_deduction
after_gift = before_gift - gift_deduction
after_recharge = before_recharge - recharge_deduction
after_total = after_recharge + after_gift
locked_wallet.balance = after_recharge
locked_wallet.gift_balance = after_gift
locked_wallet.total_consumed = to_money_decimal(locked_wallet.total_consumed) + amount
locked_wallet.updated_at = datetime.now(timezone.utc)
usage.wallet_id = locked_wallet.id
usage.wallet_balance_before = before_total
usage.wallet_balance_after = after_total
usage.wallet_recharge_balance_before = before_recharge
usage.wallet_recharge_balance_after = after_recharge
usage.wallet_gift_balance_before = before_gift
usage.wallet_gift_balance_after = after_gift
return before_total, after_total
@classmethod
def set_wallet_limit_mode(
cls,
db: Session,
*,
wallet: Wallet,
limit_mode: Literal["finite", "unlimited"],
) -> Wallet:
if limit_mode not in {"finite", "unlimited"}:
raise ValueError("limit_mode must be finite or unlimited")
locked_wallet = (
db.query(Wallet).filter(Wallet.id == wallet.id).with_for_update().one_or_none()
)
if locked_wallet is None:
raise ValueError("wallet not found")
locked_wallet.limit_mode = limit_mode
locked_wallet.updated_at = datetime.now(timezone.utc)
db.flush()
return locked_wallet
@classmethod
def create_wallet_transaction(
cls,
db: Session,
*,
wallet: Wallet,
category: WalletCategory,
reason_code: str,
amount: Decimal | float | int | str,
balance_type: WalletBalanceBucket | None = None,
link_type: str | None = None,
link_id: str | None = None,
operator_id: str | None = None,
description: str | None = None,
) -> WalletTransaction:
if category not in {"recharge", "gift", "adjust", "refund"}:
raise ValueError("category must be recharge/gift/adjust/refund")
if not reason_code:
raise ValueError("reason_code is required")
locked_wallet = (
db.query(Wallet).filter(Wallet.id == wallet.id).with_for_update().one_or_none()
)
if locked_wallet is None:
raise ValueError("wallet not found")
delta = to_money_decimal(amount)
bucket = balance_type
if bucket is None:
bucket = "gift" if category == "gift" else "recharge"
before_recharge = cls.get_recharge_balance_value(locked_wallet)
before_gift = cls.get_gift_balance_value(locked_wallet)
before_total = before_recharge + before_gift
after_recharge = before_recharge
after_gift = before_gift
if bucket == "recharge":
after_recharge = before_recharge + delta
else:
after_gift = before_gift + delta
after_total = after_recharge + after_gift
if category == "refund" and bucket != "recharge":
raise ValueError("refund transaction must use recharge balance")
if category == "refund" and delta < Decimal("0") and after_recharge < Decimal("0"):
raise ValueError("refund amount exceeds refundable recharge balance")
if bucket == "gift" and delta < Decimal("0") and after_gift < Decimal("0"):
raise ValueError("gift balance cannot be negative")
if bucket == "gift" and locked_wallet.api_key_id is not None:
raise ValueError("api key wallet does not support gift balance")
locked_wallet.balance = after_recharge
locked_wallet.gift_balance = after_gift
locked_wallet.updated_at = datetime.now(timezone.utc)
if category == "recharge":
locked_wallet.total_recharged = to_money_decimal(locked_wallet.total_recharged) + delta
elif category == "refund":
# refund_out 为负值累计退款增加refund_revert 为正值(累计退款回退)。
next_total_refunded = to_money_decimal(locked_wallet.total_refunded) - delta
locked_wallet.total_refunded = max(next_total_refunded, Decimal("0"))
elif category in {"gift", "adjust"}:
locked_wallet.total_adjusted = to_money_decimal(locked_wallet.total_adjusted) + delta
tx = WalletTransaction(
wallet_id=locked_wallet.id,
category=category,
reason_code=reason_code,
amount=delta,
balance_before=before_total,
balance_after=after_total,
recharge_balance_before=before_recharge,
recharge_balance_after=after_recharge,
gift_balance_before=before_gift,
gift_balance_after=after_gift,
link_type=link_type,
link_id=link_id,
operator_id=operator_id,
description=description,
)
db.add(tx)
db.flush()
return tx
@classmethod
def create_manual_recharge_order(
cls,
db: Session,
*,
wallet: Wallet,
amount_usd: Decimal | float | int | str,
payment_method: str = "admin_manual",
operator_id: str | None = None,
description: str | None = None,
reason_code: str | None = None,
link_type: str = "payment_order",
link_id: str | None = None,
) -> PaymentOrder:
amount = to_money_decimal(amount_usd)
if amount <= Decimal("0"):
raise ValueError("recharge amount must be positive")
if wallet.api_key_id is not None:
raise ValueError("api key wallet does not support recharge, use adjust instead")
now = datetime.now(timezone.utc)
order = PaymentOrder(
order_no=cls._build_order_no("po"),
wallet_id=wallet.id,
user_id=wallet.user_id,
amount_usd=amount,
refunded_amount_usd=Decimal("0"),
refundable_amount_usd=amount,
payment_method=payment_method,
status="credited",
paid_at=now,
credited_at=now,
gateway_response={
"source": "manual",
"operator_id": operator_id,
"description": description,
},
)
db.add(order)
db.flush()
tx_reason = reason_code
if tx_reason is None:
if payment_method in {"card_code", "gift_code", "card_recharge"}:
tx_reason = REASON_TOPUP_CARD_CODE
else:
tx_reason = REASON_TOPUP_ADMIN_MANUAL
cls.create_wallet_transaction(
db,
wallet=wallet,
category="recharge",
reason_code=tx_reason,
amount=amount,
balance_type="recharge",
link_type=link_type,
link_id=link_id or order.id,
operator_id=operator_id,
description=description or "管理员充值",
)
return order
@classmethod
def admin_adjust_balance(
cls,
db: Session,
*,
wallet: Wallet,
amount_usd: Decimal | float | int | str,
balance_type: Literal["recharge", "gift"] = "recharge",
operator_id: str | None = None,
description: str | None = None,
) -> WalletTransaction:
amount = to_money_decimal(amount_usd)
if amount == Decimal("0"):
raise ValueError("adjust amount must not be zero")
if balance_type not in {"recharge", "gift"}:
raise ValueError("balance_type must be recharge or gift")
if balance_type == "gift" and wallet.api_key_id is not None:
raise ValueError("api key wallet does not support gift balance")
# 正向调账:加给谁就加给谁,不做抵充。
if amount > Decimal("0"):
return cls.create_wallet_transaction(
db,
wallet=wallet,
category="adjust",
reason_code=REASON_ADJUST_ADMIN,
amount=amount,
balance_type=balance_type,
link_type="admin_action",
link_id=wallet.id,
operator_id=operator_id,
description=description or "管理员调账",
)
# 负向调账:先扣所选账户,再扣另一账户;若仍不足,继续计入充值余额(可为负)。
locked_wallet = (
db.query(Wallet).filter(Wallet.id == wallet.id).with_for_update().one_or_none()
)
if locked_wallet is None:
raise ValueError("wallet not found")
before_recharge = cls.get_recharge_balance_value(locked_wallet)
before_gift = cls.get_gift_balance_value(locked_wallet)
before_total = before_recharge + before_gift
after_recharge = before_recharge
after_gift = before_gift
remaining = -amount
def consume_positive_bucket(
balance: Decimal, to_consume: Decimal
) -> tuple[Decimal, Decimal]:
if to_consume <= Decimal("0"):
return balance, Decimal("0")
available = max(balance, Decimal("0"))
consumed = min(available, to_consume)
return balance - consumed, to_consume - consumed
if balance_type == "gift":
after_gift, remaining = consume_positive_bucket(after_gift, remaining)
after_recharge, remaining = consume_positive_bucket(after_recharge, remaining)
else:
after_recharge, remaining = consume_positive_bucket(after_recharge, remaining)
after_gift, remaining = consume_positive_bucket(after_gift, remaining)
if remaining > Decimal("0"):
after_recharge = after_recharge - remaining
if after_gift < Decimal("0"):
raise ValueError("gift balance cannot be negative")
after_total = after_recharge + after_gift
locked_wallet.balance = after_recharge
locked_wallet.gift_balance = after_gift
locked_wallet.updated_at = datetime.now(timezone.utc)
locked_wallet.total_adjusted = to_money_decimal(locked_wallet.total_adjusted) + amount
tx = WalletTransaction(
wallet_id=locked_wallet.id,
category="adjust",
reason_code=REASON_ADJUST_ADMIN,
amount=amount,
balance_before=before_total,
balance_after=after_total,
recharge_balance_before=before_recharge,
recharge_balance_after=after_recharge,
gift_balance_before=before_gift,
gift_balance_after=after_gift,
link_type="admin_action",
link_id=wallet.id,
operator_id=operator_id,
description=description or "管理员调账",
)
db.add(tx)
db.flush()
return tx
@classmethod
def _get_pending_refund_reserved_amount(
cls,
db: Session,
*,
wallet_id: str | None = None,
payment_order_id: str | None = None,
) -> Decimal:
query = db.query(func.coalesce(func.sum(RefundRequest.amount_usd), 0)).filter(
RefundRequest.status.in_(["pending_approval", "approved"])
)
if wallet_id is not None:
query = query.filter(RefundRequest.wallet_id == wallet_id)
if payment_order_id is not None:
query = query.filter(RefundRequest.payment_order_id == payment_order_id)
return to_money_decimal(query.scalar() or 0)
@classmethod
def create_refund_request(
cls,
db: Session,
*,
wallet: Wallet,
user_id: str | None,
amount_usd: Decimal | float | int | str,
refund_no: str,
source_type: str,
source_id: str | None,
refund_mode: str,
payment_order: PaymentOrder | None = None,
reason: str | None = None,
requested_by: str | None = None,
idempotency_key: str | None = None,
) -> RefundRequest:
amount = to_money_decimal(amount_usd)
if amount <= Decimal("0"):
raise ValueError("refund amount must be positive")
locked_wallet = (
db.query(Wallet).filter(Wallet.id == wallet.id).with_for_update().one_or_none()
)
if locked_wallet is None:
raise ValueError("wallet not found")
refundable_balance = cls.get_refundable_balance_value(locked_wallet)
reserved_wallet_amount = cls._get_pending_refund_reserved_amount(
db,
wallet_id=locked_wallet.id,
)
available_refundable_balance = refundable_balance - reserved_wallet_amount
if amount > available_refundable_balance:
raise ValueError("refund amount exceeds available refundable recharge balance")
locked_payment_order = None
if payment_order is not None:
locked_payment_order = (
db.query(PaymentOrder)
.filter(PaymentOrder.id == payment_order.id)
.with_for_update()
.one_or_none()
)
if locked_payment_order is None:
raise ValueError("payment order not found")
if locked_payment_order.wallet_id != locked_wallet.id:
raise ValueError("payment order does not belong to wallet")
if locked_payment_order.status != "credited":
raise ValueError("payment order is not refundable")
refundable_amount = to_money_decimal(locked_payment_order.refundable_amount_usd)
reserved_order_amount = cls._get_pending_refund_reserved_amount(
db,
payment_order_id=locked_payment_order.id,
)
available_refundable_amount = refundable_amount - reserved_order_amount
if amount > available_refundable_amount:
raise ValueError("refund amount exceeds available refundable amount")
refund = RefundRequest(
refund_no=refund_no,
wallet_id=locked_wallet.id,
user_id=user_id,
payment_order_id=locked_payment_order.id if locked_payment_order else None,
source_type=source_type,
source_id=source_id,
refund_mode=refund_mode,
amount_usd=amount,
status="pending_approval",
reason=reason,
requested_by=requested_by,
idempotency_key=idempotency_key,
)
db.add(refund)
db.flush()
return refund
@classmethod
def move_refund_to_processing(
cls,
db: Session,
*,
refund: RefundRequest,
operator_id: str | None = None,
) -> WalletTransaction:
locked_refund = (
db.query(RefundRequest)
.filter(RefundRequest.id == refund.id)
.with_for_update()
.one_or_none()
)
if locked_refund is None:
raise ValueError("refund not found")
if locked_refund.status not in {"approved", "pending_approval"}:
raise ValueError("refund status is not approvable")
locked_wallet = (
db.query(Wallet)
.filter(Wallet.id == locked_refund.wallet_id)
.with_for_update()
.one_or_none()
)
if locked_wallet is None:
raise ValueError("wallet not found")
payment_order = None
if locked_refund.payment_order_id:
payment_order = (
db.query(PaymentOrder)
.filter(PaymentOrder.id == locked_refund.payment_order_id)
.with_for_update()
.one_or_none()
)
if payment_order is None:
raise ValueError("payment order not found")
refund_amount = to_money_decimal(locked_refund.amount_usd)
refundable_amount = to_money_decimal(payment_order.refundable_amount_usd)
if refund_amount > refundable_amount:
raise ValueError("refund amount exceeds refundable amount")
tx = cls.create_wallet_transaction(
db,
wallet=locked_wallet,
category="refund",
reason_code=REASON_REFUND_OUT,
amount=-to_money_decimal(locked_refund.amount_usd),
balance_type="recharge",
link_type="refund_request",
link_id=locked_refund.id,
operator_id=operator_id,
description="退款占款",
)
if payment_order is not None:
delta = to_money_decimal(locked_refund.amount_usd)
payment_order.refunded_amount_usd = (
to_money_decimal(payment_order.refunded_amount_usd) + delta
)
payment_order.refundable_amount_usd = (
to_money_decimal(payment_order.refundable_amount_usd) - delta
)
locked_refund.status = "processing"
locked_refund.approved_by = operator_id
locked_refund.processed_by = operator_id
locked_refund.processed_at = datetime.now(timezone.utc)
locked_refund.updated_at = datetime.now(timezone.utc)
return tx
@classmethod
def fail_refund(
cls,
db: Session,
*,
refund: RefundRequest,
reason: str,
operator_id: str | None = None,
) -> WalletTransaction | None:
locked_refund = (
db.query(RefundRequest)
.filter(RefundRequest.id == refund.id)
.with_for_update()
.one_or_none()
)
if locked_refund is None:
raise ValueError("refund not found")
if locked_refund.status in {"pending_approval", "approved"}:
locked_refund.status = "failed"
locked_refund.failure_reason = reason
locked_refund.updated_at = datetime.now(timezone.utc)
return None
if locked_refund.status != "processing":
raise ValueError(f"cannot fail refund in status: {locked_refund.status}")
wallet = db.query(Wallet).filter(Wallet.id == locked_refund.wallet_id).first()
if wallet is None:
raise ValueError("wallet not found")
tx = cls.create_wallet_transaction(
db,
wallet=wallet,
category="refund",
reason_code=REASON_REFUND_REVERT,
amount=to_money_decimal(locked_refund.amount_usd),
balance_type="recharge",
link_type="refund_request",
link_id=locked_refund.id,
operator_id=operator_id,
description="退款失败回补",
)
if locked_refund.payment_order_id:
payment_order = (
db.query(PaymentOrder)
.filter(PaymentOrder.id == locked_refund.payment_order_id)
.with_for_update()
.one_or_none()
)
if payment_order is not None:
delta = to_money_decimal(locked_refund.amount_usd)
payment_order.refunded_amount_usd = (
to_money_decimal(payment_order.refunded_amount_usd) - delta
)
payment_order.refundable_amount_usd = (
to_money_decimal(payment_order.refundable_amount_usd) + delta
)
locked_refund.status = "failed"
locked_refund.failure_reason = reason
locked_refund.updated_at = datetime.now(timezone.utc)
return tx
@classmethod
def complete_refund(
cls,
db: Session,
*,
refund: RefundRequest,
gateway_refund_id: str | None = None,
payout_reference: str | None = None,
payout_proof: dict | None = None,
) -> RefundRequest:
locked_refund = (
db.query(RefundRequest)
.filter(RefundRequest.id == refund.id)
.with_for_update()
.one_or_none()
)
if locked_refund is None:
raise ValueError("refund not found")
if locked_refund.status != "processing":
raise ValueError("refund status must be processing before completion")
locked_refund.status = "succeeded"
locked_refund.gateway_refund_id = gateway_refund_id
locked_refund.payout_reference = payout_reference
locked_refund.payout_proof = payout_proof
locked_refund.completed_at = datetime.now(timezone.utc)
locked_refund.updated_at = datetime.now(timezone.utc)
return locked_refund