yingqing/core/backend/apps/billing/services/ledger.py

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="释放未用预留额度",
        )