from decimal import Decimal from django.db import transaction from django.db.models import Sum from django.utils import timezone from apps.billing.models import CreditAccount, CreditLedger, CreditReservation def _enforce_member_monthly_limit(*, team, user, amount: Decimal) -> None: """成员月度限额管控。团队页可设 monthly_credit_limit(0=不限),此前只存不管——纯装饰。 口径与成员列表的 month_charged 一致:自然月内该成员的 CHARGE 流水合计; 再加上其当前 ACTIVE 预留(在途任务),防止并发把限额冲穿。 调用方(reserve_credit)已持有 account 行锁,团队内 reserve 串行化,此检查无竞态。""" if user is None: return member = team.members.filter(user=user).first() if member is None: return limit = member.monthly_credit_limit or Decimal("0") if limit <= 0: return # 0 = 不限额 now = timezone.now() month_start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0) charged = ( CreditLedger.objects.filter(team=team, user=user, ledger_type=CreditLedger.Type.CHARGE, created_at__gte=month_start) .aggregate(s=Sum("amount"))["s"] or Decimal("0") ) reserved = ( CreditReservation.objects.filter(team=team, user=user, status=CreditReservation.Status.ACTIVE) .aggregate(s=Sum("amount"))["s"] or Decimal("0") ) if charged + reserved + amount > limit: raise ValueError(f"成员本月额度不足:限额 ¥{limit},本月已用 ¥{charged}(另有在途 ¥{reserved})") @transaction.atomic def reserve_credit(*, team, user, task, amount: Decimal) -> CreditReservation: account, _ = CreditAccount.objects.select_for_update().get_or_create(team=team) available = account.balance - account.reserved_balance if available < amount: raise ValueError("insufficient credit") _enforce_member_monthly_limit(team=team, user=user, amount=amount) account.reserved_balance += amount account.save(update_fields=["reserved_balance", "updated_at"]) reservation = CreditReservation.objects.create( team=team, user=user, project=task.project, task=task, amount=amount, ) CreditLedger.objects.create( team=team, user=user, project=task.project, task=task, ledger_type=CreditLedger.Type.RESERVE, amount=amount, balance_after=account.balance, reason="AI 任务预扣额度", ) return reservation @transaction.atomic def release_credit(*, reservation: CreditReservation, reason: str = "") -> None: account = CreditAccount.objects.select_for_update().get(team=reservation.team) # 锁内重读:与 charge 同款竞态——并发双 release(或 charge+release 交错)用陈旧快照都看到 ACTIVE, # reserved_balance 会被减两次直接记坏账。已终态(已扣/已释放)→ 幂等返回。 reservation = CreditReservation.objects.select_for_update().get(id=reservation.id) if reservation.status != CreditReservation.Status.ACTIVE: return account.reserved_balance -= reservation.amount account.save(update_fields=["reserved_balance", "updated_at"]) reservation.status = CreditReservation.Status.RELEASED reservation.save(update_fields=["status", "updated_at"]) CreditLedger.objects.create( team=reservation.team, user=reservation.user, project=reservation.project, task=reservation.task, ledger_type=CreditLedger.Type.RELEASE, amount=reservation.amount, balance_after=account.balance, reason=reason or "释放预留额度", ) @transaction.atomic def charge_reserved_credit(*, reservation: CreditReservation, actual_amount: Decimal) -> None: account = CreditAccount.objects.select_for_update().get(team=reservation.team) # 锁内重读 reservation:调用方传入的可能是并发竞态下的陈旧快照(双方都看到 ACTIVE), # 旧实现会对同一笔预留扣两次费(实测同秒两条 charge 流水)。已扣费 → 幂等直接返回。 reservation = CreditReservation.objects.select_for_update().get(id=reservation.id) if reservation.status == CreditReservation.Status.CHARGED: return if reservation.status != CreditReservation.Status.ACTIVE: raise ValueError("reservation is not active") if actual_amount > reservation.amount: raise ValueError("actual amount exceeds reserved amount") account.balance -= actual_amount account.reserved_balance -= reservation.amount account.save(update_fields=["balance", "reserved_balance", "updated_at"]) reservation.status = CreditReservation.Status.CHARGED reservation.save(update_fields=["status", "updated_at"]) CreditLedger.objects.create( team=reservation.team, user=reservation.user, project=reservation.project, task=reservation.task, ledger_type=CreditLedger.Type.CHARGE, amount=actual_amount, balance_after=account.balance, reason="AI 任务扣费", ) if reservation.amount > actual_amount: CreditLedger.objects.create( team=reservation.team, user=reservation.user, project=reservation.project, task=reservation.task, ledger_type=CreditLedger.Type.RELEASE, amount=reservation.amount - actual_amount, balance_after=account.balance, reason="释放未用预留额度", )