Merge pull request #3211 from ResearchHub/fix-locked-balance-allocation

Fix locked balance allocation

Author: koutst

src/purchase/migrations/0050_balance_purchase.py

@@ -0,0 +1,25 @@
+# Generated by Django 5.2.12 on 2026-03-19 22:02
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("purchase", "0049_grant_add_pending_declined_moderation_fields"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="balance",
+            name="purchase",
+            field=models.ForeignKey(
+                blank=True,
+                null=True,
+                on_delete=django.db.models.deletion.SET_NULL,
+                related_name="balance_records",
+                to="purchase.purchase",
+            ),
+        ),
+    ]

src/purchase/related_models/balance_model.py

@@ -17,6 +17,16 @@ class LockType(models.TextChoices):
     object_id = models.PositiveIntegerField(null=True)
     source = GenericForeignKey("content_type", "object_id")
 
+    # Optional link to the purchase that triggered this balance record.
+    # Used for contribution refunds.
+    purchase = models.ForeignKey(
+        "purchase.Purchase",
+        on_delete=models.SET_NULL,
+        null=True,
+        blank=True,
+        related_name="balance_records",
+    )
+
     # TODO: why is this a char field?
     amount = models.CharField(max_length=255)
     testnet_amount = models.CharField(max_length=255, default=0, null=True, blank=True)

src/purchase/services/fundraise_service.py

@@ -217,10 +217,11 @@ def create_rsc_contribution(
         with transaction.atomic():
             user = User.objects.select_for_update().get(id=user.id)
 
-            # Check if user has enough balance in their wallet
-            # For fundraise contributions, we allow using locked balance
-            user_balance = user.get_balance(include_locked=True)
-            if user_balance - (amount + fee) < 0:
+            # Allocate the total spend (amount + fee) across locked/unlocked
+            # pools. Fundraise contributions are allowed to consume locked funds.
+            try:
+                allocations = user.allocate_spend(amount + fee, allow_locked=True)
+            except ValueError:
                 return None, "Insufficient balance"
 
             # Create purchase object
@@ -237,55 +238,39 @@ def create_rsc_contribution(
             # Deduct fees
             deduct_bounty_fees(user, fee, rh_fee, dao_fee, fee_object)
 
-            # Get user's available locked balance
-            available_locked_balance = user.get_locked_balance()
-
-            # Determine how to split the contribution amount
-            locked_amount_used = min(available_locked_balance, amount)
-            regular_amount_used = amount - locked_amount_used
-
-            # Determine how to split the fees using remaining locked balance
-            remaining_locked_balance = available_locked_balance - locked_amount_used
-            locked_fee_used = min(remaining_locked_balance, fee)
-            regular_fee_used = fee - locked_fee_used
-
-            # Create balance records for the contribution amount
-            if locked_amount_used > 0:
-                Balance.objects.create(
-                    user=user,
-                    content_type=ContentType.objects.get_for_model(Purchase),
-                    object_id=purchase.id,
-                    amount=f"-{locked_amount_used.to_eng_string()}",
-                    is_locked=True,
-                    lock_type=Balance.LockType.REFERRAL_BONUS,
-                )
-
-            if regular_amount_used > 0:
-                Balance.objects.create(
-                    user=user,
-                    content_type=ContentType.objects.get_for_model(Purchase),
-                    object_id=purchase.id,
-                    amount=f"-{regular_amount_used.to_eng_string()}",
-                )
-
-            # Create balance records for the fees
-            if locked_fee_used > 0:
-                Balance.objects.create(
-                    user=user,
-                    content_type=ContentType.objects.get_for_model(BountyFee),
-                    object_id=fee_object.id,
-                    amount=f"-{locked_fee_used.to_eng_string()}",
-                    is_locked=True,
-                    lock_type=Balance.LockType.REFERRAL_BONUS,
-                )
-
-            if regular_fee_used > 0:
-                Balance.objects.create(
-                    user=user,
-                    content_type=ContentType.objects.get_for_model(BountyFee),
-                    object_id=fee_object.id,
-                    amount=f"-{regular_fee_used.to_eng_string()}",
-                )
+            # Create balance debit records, splitting each allocation across
+            # the contribution amount and fee.
+            remaining_amount = amount
+
+            for alloc in allocations:
+                alloc_amount = alloc["amount"]
+
+                # Apply as much of this allocation to the contribution first,
+                # then the remainder to fees.
+                amount_used = min(alloc_amount, remaining_amount)
+                fee_used = alloc_amount - amount_used
+
+                if amount_used > 0:
+                    Balance.objects.create(
+                        user=user,
+                        content_type=ContentType.objects.get_for_model(Purchase),
+                        object_id=purchase.id,
+                        amount=f"-{amount_used.to_eng_string()}",
+                        is_locked=alloc["is_locked"],
+                        purchase=purchase,
+                    )
+
+                if fee_used > 0:
+                    Balance.objects.create(
+                        user=user,
+                        content_type=ContentType.objects.get_for_model(BountyFee),
+                        object_id=fee_object.id,
+                        amount=f"-{fee_used.to_eng_string()}",
+                        is_locked=alloc["is_locked"],
+                        purchase=purchase,
+                    )
+
+                remaining_amount -= amount_used
 
             # Track in Amplitude
             rh_fee_str = rh_fee.to_eng_string()
@@ -389,47 +374,53 @@ def create_usd_contribution(
 
         return contribution, None
 
+    def _refund_contribution_debit(
+        self, fundraise, user, debit, purchase_ct, bounty_fee_ct
+    ):
+        """Refund a single debit entry. Returns True on success, False on failure."""
+        abs_amount = abs(Decimal(debit.amount))
+        if abs_amount == 0:
+            return True
+
+        if debit.content_type == purchase_ct:
+            return fundraise.escrow.refund(user, abs_amount, is_locked=debit.is_locked)
+
+        if debit.content_type == bounty_fee_ct:
+            return self._refund_fee(user, debit, abs_amount)
+
+        return True
+
+    def _refund_fee(self, user, debit, abs_amount):
+        """Refund a fee debit from the revenue account. Returns True on success."""
+        rh_revenue_account = User.objects.get_revenue_account()
+        fee_object = BountyFee.objects.get(id=debit.object_id)
+        distribution = create_bounty_refund_distribution(abs_amount)
+        distributor = Distributor(
+            distribution,
+            user,
+            fee_object,
+            time.time(),
+            giver=rh_revenue_account,
+            is_locked=debit.is_locked,
+        )
+        record = distributor.distribute()
+        return record.distributed_status != "FAILED"
+
     def refund_rsc_contributions(self, fundraise: Fundraise) -> bool:
         """
         Refund all RSC contributions from escrow back to contributors.
         Also refunds the fees that were deducted when creating contributions.
+        Preserves the locked/unlocked status of the original funds.
         Returns True if all refunds successful, False if any fail.
         """
-        # Get all purchases (RSC contributions) for this fundraise
-        contributions = fundraise.purchases.all()
-
-        # Refund each RSC contributor
-        for contribution in contributions:
-            user = contribution.user
-            amount = Decimal(contribution.amount)
-
-            # Only refund what's still in escrow
-            if amount > 0:
-                success = fundraise.escrow.refund(user, amount)
-                if not success:
-                    # If a refund fails, we should abort the whole transaction
-                    return False
-
-            # Also refund the fees that were deducted when this contribution
-            # was made. Calculate the fee using the same logic used during
-            # contribution creation.
-            fee, _, _, fee_object = calculate_bounty_fees(amount)
-
-            if fee > 0:
-                # Create a refund for the fee
-                rh_revenue_account = User.objects.get_revenue_account()
-                distribution = create_bounty_refund_distribution(fee)
-                distributor = Distributor(
-                    distribution,
-                    user,
-                    fee_object,  # The BountyFee object
-                    time.time(),
-                    giver=rh_revenue_account,
-                )
-                record = distributor.distribute()
-                if record.distributed_status == "FAILED":
-                    # If fee refund fails, we should abort the whole
-                    # transaction
+        purchase_ct = ContentType.objects.get_for_model(Purchase)
+        bounty_fee_ct = ContentType.objects.get_for_model(BountyFee)
+
+        for contribution in fundraise.purchases.all():
+            for debit in Balance.objects.filter(purchase=contribution):
+                if not self._refund_contribution_debit(
+                    fundraise, contribution.user, debit, purchase_ct, bounty_fee_ct
+                ):
                     return False
 
         return True