from decimal import Decimal from django.db import transaction from apps.billing.models import CreditAccount, CreditLedger, CreditReservation @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") 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="reserve ai task credit", ) return reservation @transaction.atomic def release_credit(*, reservation: CreditReservation, reason: str = "") -> None: account = CreditAccount.objects.select_for_update().get(team=reservation.team) 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 "release reserved credit", ) @transaction.atomic def charge_reserved_credit(*, reservation: CreditReservation, actual_amount: Decimal) -> None: account = CreditAccount.objects.select_for_update().get(team=reservation.team) 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="charge ai task credit", ) 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="release unused reserved credit", )