feat: add AirShelf core implementation

2026-06-05 10:21:40 +08:00
parent 2ba1058329
commit cfdcd84a30
252 changed files with 70828 additions and 0 deletions
@@ -0,0 +1,93 @@
+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",
+        )
+