feat(wallets): Support multiple wallets (#4678)

## Context

We’re introducing the priority field to wallets to support custom
execution or display ordering. This PR builds on the base work from
#4391, which added the field to
the Wallet model and exposed it via the API.

The goal of this PR is to implement the logic that uses the priority
value when handling wallets — for example, when determining which
wallets to apply or process first.

## Description

This PR includes:
• Logic to order wallets by priority (and created_at as a secondary
key).
• Ensures active_wallets_in_application_order returns wallets in the
correct priority sequence.
• Updates any internal wallet processing to respect the defined order.

This builds directly on the foundation laid in PR
#4391 and assumes the priority
field and its validations are already in place.

---------

Co-authored-by: Miguel Pinto <darkymiguel@gmail.com>

Author: floganz

.junie/guidelines.md

@@ -6,52 +6,113 @@ class AppliedPrepaidCreditsService < BaseService
 
     def initialize(invoice:, wallets:, max_wallet_decrease_attempts: DEFAULT_MAX_WALLET_DECREASE_ATTEMPTS)
       @invoice = invoice
-      @wallet = wallets.first
+      @wallets = wallets
       @max_wallet_decrease_attempts = max_wallet_decrease_attempts
       raise ArgumentError, "max_wallet_decrease_attempts must be between 1 and #{DEFAULT_MAX_WALLET_DECREASE_ATTEMPTS} (inclusive)" if max_wallet_decrease_attempts < 1 || max_wallet_decrease_attempts > DEFAULT_MAX_WALLET_DECREASE_ATTEMPTS
 
       super(nil)
     end
 
-    activity_loggable(
-      action: "wallet_transaction.created",
-      record: -> { result.wallet_transaction }
-    )
-
     def call
-      if already_applied?
+      if wallets_already_applied?
         return result.service_failure!(code: "already_applied", message: "Prepaid credits already applied")
       end
 
-      amount_cents = compute_amount
-
-      ApplicationRecord.transaction do
-        wallet_transaction = create_wallet_transaction(amount_cents)
-        result.wallet_transaction = wallet_transaction
-        amount_cents = wallet_transaction.amount_cents
-
-        with_optimistic_lock_retry(wallet) do
-          Wallets::Balance::DecreaseService.call(wallet:, wallet_transaction:)
+      result.prepaid_credit_amount_cents ||= 0
+      result.wallet_transactions ||= []
+
+      ActiveRecord::Base.transaction do
+        ordered_remaining_amounts = calculate_amounts_for_fees_by_type_and_bm
+        wallets.each do |wallet|
+          wallet.reload
+          wallet_fee_transactions = []
+          wallet_targets_array = wallet.wallet_targets.map do |wt|
+            if wt&.billable_metric_id
+              ["charge", wt.billable_metric_id]
+            end
+          end
+          wallet_types_array = wallet.allowed_fee_types
+
+          ordered_remaining_amounts.each do |fee_key, remaining_amount|
+            next if remaining_amount <= 0
+
+            next unless applicable_fee?(fee_key:, targets: wallet_targets_array, types: wallet_types_array)
+
+            used_amount = wallet_fee_transactions.sum { |t| t[:amount_cents] }
+            remaining_wallet_balance = wallet.balance_cents - used_amount
+            next if remaining_wallet_balance <= 0
+
+            transaction_amount = [remaining_amount, remaining_wallet_balance].min
+            next if transaction_amount <= 0
+
+            ordered_remaining_amounts[fee_key] -= transaction_amount
+            wallet_fee_transactions << {
+              fee_key: fee_key,
+              amount_cents: transaction_amount
+            }
+          end
+
+          total_amount_cents = wallet_fee_transactions.sum { |t| t[:amount_cents] }
+          next if total_amount_cents <= 0
+
+          wallet_transaction = create_wallet_transaction(wallet, total_amount_cents)
+          amount_cents = wallet_transaction.amount_cents
+
+          with_optimistic_lock_retry(wallet) do
+            Wallets::Balance::DecreaseService.call(wallet:, wallet_transaction:, skip_refresh: true)
+          end
+
+          result.wallet_transactions << wallet_transaction
+          result.prepaid_credit_amount_cents += amount_cents
+          invoice.prepaid_credit_amount_cents += amount_cents
         end
-      end
 
-      result.prepaid_credit_amount_cents = amount_cents
-      invoice.prepaid_credit_amount_cents += amount_cents
-
-      SendWebhookJob.perform_after_commit("wallet_transaction.created", result.wallet_transaction)
+        Customers::RefreshWalletsService.call(customer:, include_generating_invoices: true)
+        invoice.save! if invoice.changed?
+      end
 
+      schedule_webhook_notifications(result.wallet_transactions)
       result
     rescue ActiveRecord::RecordInvalid => e
       result.record_validation_failure!(record: e.record)
     end
 
     private
 
-    attr_accessor :invoice, :wallet, :max_wallet_decrease_attempts
+    attr_accessor :invoice, :wallets, :max_wallet_decrease_attempts
 
-    delegate :balance_cents, to: :wallet
+    delegate :customer, to: :invoice
+
+    def schedule_webhook_notifications(wallet_transactions)
+      wallet_transactions.each do |wt|
+        Utils::ActivityLog.produce_after_commit(wt, "wallet_transaction.created")
+        SendWebhookJob.perform_after_commit("wallet_transaction.created", wt)
+      end
+    end
+
+    def calculate_amounts_for_fees_by_type_and_bm
+      remaining = Hash.new(0)
+
+      invoice.fees.includes(:charge).find_each do |fee|
+        cap = fee.sub_total_excluding_taxes_amount_cents +
+          fee.taxes_precise_amount_cents -
+          fee.precise_credit_notes_amount_cents
+
+        next if cap <= 0
+        key = [fee.fee_type, fee.charge&.billable_metric_id]
+        remaining[key] += cap
+      end
+
+      remaining.sort_by { |_, v| -v }.to_h
+    end
 
-    def create_wallet_transaction(amount_cents)
+    def wallets_already_applied?
+      return false unless invoice
+
+      WalletTransaction.exists?(invoice_id: invoice.id, wallet_id: wallets.map(&:id))
+    end
+
+    def create_wallet_transaction(wallet, amount_cents)
       wallet_credit = WalletCredit.from_amount_cents(wallet:, amount_cents:)
 
       result = WalletTransactions::CreateService.call!(
@@ -82,45 +143,12 @@ def with_optimistic_lock_retry(wallet, &block)
       end
     end
 
-    def already_applied?
-      invoice&.wallet_transactions&.exists?
-    end
-
-    def compute_amount
-      if wallet.limited_to_billable_metrics? && billable_metric_limited_fees
-        bm_limited_fees = billable_metric_limited_fees
-        remaining_fees = invoice.fees - bm_limited_fees
-        remaining_fees = remaining_fees.reject { |fee| fee.fee_type == "charge" }
-      else
-        bm_limited_fees = []
-        remaining_fees = invoice.fees
-      end
-
-      fee_type_limited_fees = if wallet.limited_fee_types?
-        remaining_fees.filter { |fee| wallet.allowed_fee_types.include?(fee.fee_type) }
-      elsif wallet.limited_to_billable_metrics? && billable_metric_limited_fees
-        []
-      else
-        remaining_fees
-      end
-
-      if wallet.limited_fee_types? || wallet.limited_to_billable_metrics?
-        [balance_cents, limited_fees_total(bm_limited_fees + fee_type_limited_fees)].min
-      else
-        [balance_cents, invoice.total_amount_cents].min
-      end
-    end
-
-    def billable_metric_limited_fees
-      @billable_metric_limited_fees ||= invoice.fees
-        .joins(charge: :billable_metric)
-        .where(billable_metric: {id: wallet.wallet_targets.pluck(:billable_metric_id)})
-    end
+    def applicable_fee?(fee_key:, targets:, types:)
+      target_match = targets.include?(fee_key)
+      type_match = types.include?(fee_key.first)
+      unrestricted_wallet = targets.empty? && types.empty?
 
-    def limited_fees_total(applicable_fees)
-      applicable_fees.sum do |f|
-        f.sub_total_excluding_taxes_precise_amount_cents + f.taxes_precise_amount_cents - f.precise_credit_notes_amount_cents
-      end
+      target_match || type_match || unrestricted_wallet
     end
   end
 end

app/services/credits/applied_prepaid_credits_service.rb

@@ -2,7 +2,7 @@
 
 module Customers
   class RefreshWalletsService < BaseService
-    Result = BaseResult[:usage_amount_cents, :wallets]
+    Result = BaseResult[:usage_amount_cents, :wallets, :allocation_rules]
 
     def initialize(customer:, include_generating_invoices: false)
       @customer = customer
@@ -27,11 +27,18 @@ def call
         }
       end
 
+      allocation_rules = Wallets::BuildAllocationRulesService.call!(customer:).allocation_rules
+
       customer.wallets.active.find_each do |wallet|
-        Wallets::Balance::RefreshOngoingUsageService.call!(wallet:, usage_amount_cents:)
+        Wallets::Balance::RefreshOngoingUsageService.call!(
+          wallet:,
+          usage_amount_cents:,
+          allocation_rules:
+        )
       end
 
       result.usage_amount_cents = usage_amount_cents
+      result.allocation_rules = allocation_rules
       result.wallets = customer.wallets.active.reload
       result
     rescue BaseService::FailedResult => e

app/services/customers/refresh_wallets_service.rb

@@ -292,7 +292,8 @@ def should_create_charge_fees?(subscription)
     end
 
     def wallets
-      @wallets ||= customer.wallets.active.with_positive_balance
+      @wallets ||= customer.wallets.active.includes(:wallet_targets)
+        .with_positive_balance.in_application_order
     end
 
     def should_create_credit_note_credit?

app/services/invoices/calculate_fees_service.rb

@@ -119,7 +119,8 @@ def should_deliver_email?
     end
 
     def wallets
-      @wallets ||= customer.wallets.active.with_positive_balance
+      @wallets ||= customer.wallets.active.includes(:wallet_targets)
+        .with_positive_balance.in_application_order
     end
 
     def should_create_applied_prepaid_credit?

app/services/invoices/create_pay_in_advance_charge_service.rb

@@ -146,13 +146,13 @@ def create_credit_note_credit
     end
 
     def create_applied_prepaid_credit
-      wallets = subscription.customer.wallets.active.with_positive_balance
+      wallets = subscription.customer.wallets.active.includes(:wallet_targets)
+        .with_positive_balance.in_application_order
 
       return if wallets.none?
       return unless invoice.total_amount_cents.positive?
 
       prepaid_credit_result = Credits::AppliedPrepaidCreditsService.call!(invoice:, wallets:)
-
       invoice.total_amount_cents -= prepaid_credit_result.prepaid_credit_amount_cents
     end
   end

app/services/invoices/progressive_billing_service.rb

@@ -85,7 +85,8 @@ def should_deliver_email?
       end
 
       def wallets
-        @wallets ||= customer.wallets.active.with_positive_balance
+        @wallets ||= customer.wallets.active.includes(:wallet_targets)
+          .with_positive_balance.in_application_order
       end
 
       def should_create_credit_note_credit?

app/services/invoices/provider_taxes/pull_taxes_and_apply_service.rb

@@ -79,12 +79,12 @@ def should_create_applied_prepaid_credit?
     end
 
     def wallets
-      @wallets ||= voided_invoice.customer.wallets.active.with_positive_balance
+      @wallets ||= voided_invoice.customer.wallets.active.includes(:wallet_targets)
+        .with_positive_balance.in_application_order
     end
 
     def create_applied_prepaid_credit
       prepaid_credit_result = Credits::AppliedPrepaidCreditsService.call!(invoice: regenerated_invoice, wallets:)
-
       refresh_amounts(credit_amount_cents: prepaid_credit_result.prepaid_credit_amount_cents)
     end
 

app/services/invoices/regenerate_from_voided_service.rb

@@ -3,9 +3,10 @@
 module Wallets
   module Balance
     class DecreaseService < BaseService
-      def initialize(wallet:, wallet_transaction:)
+      def initialize(wallet:, wallet_transaction:, skip_refresh: false)
         @wallet = wallet.reload
         @wallet_transaction = wallet_transaction
+        @skip_refresh = skip_refresh
 
         super
       end
@@ -23,7 +24,12 @@ def call
           last_consumed_credit_at: Time.current
         )
 
-        Customers::RefreshWalletsService.call(customer: wallet.customer, include_generating_invoices: true)
+        unless skip_refresh
+          Customers::RefreshWalletsService.call(
+            customer: wallet.customer,
+            include_generating_invoices: true
+          )
+        end
 
         after_commit { SendWebhookJob.perform_later("wallet.updated", wallet) }
 
@@ -33,7 +39,7 @@ def call
 
       private
 
-      attr_reader :wallet, :wallet_transaction
+      attr_reader :wallet, :wallet_transaction, :skip_refresh
     end
   end
 end

app/services/wallets/balance/decrease_service.rb

@@ -93,18 +93,33 @@ def credits_ongoing_balance
         ongoing_balance_cents.to_f.fdiv(currency.subunit_to_unit).fdiv(wallet.rate_amount)
       end
 
-      def calculate_total_usage_with_limitation(usage_amount_cents)
-        return usage_amount_cents.sum { |e| e[:total_usage_amount_cents] } unless wallet.limited_to_billable_metrics?
+      def assign_wallet_per_fee(fees)
+        fee_wallet = {}
+
+        allocation_rules = Wallets::BuildAllocationRulesService.call!(customer:).allocation_rules
+
+        fees.each do |fee|
+          key = fee.id || fee.object_id
+
+          applicable_wallets = Wallets::FindApplicableOnFeesService
+            .call!(allocation_rules:, fee:)
+            .top_priority_wallet
+
+          fee_wallet[key] = applicable_wallets.presence
+        end
+
+        fee_wallet
+      end
 
-        # current usage fees are not persisted so we can't use join
+      def calculate_total_usage_with_limitation(usage_amount_cents)
         all_fees = usage_amount_cents.flat_map { |usage| usage[:invoice].fees }
-        charge_ids = Charge.where(id: all_fees.map(&:charge_id)).where(billable_metric_id: wallet.wallet_targets.pluck(:billable_metric_id)).pluck(:id)
+        return 0 if all_fees.empty?
 
-        return usage_amount_cents.sum { |e| e[:total_usage_amount_cents] } if charge_ids.empty?
+        wallets_applicable_on_fees = assign_wallet_per_fee(all_fees) # { fee_key => wallet_id }
 
-        all_fees
-          .select { |f| charge_ids.include?(f.charge_id) }
-          .sum { |f| f.amount_cents + f.taxes_amount_cents }
+        all_fees.sum do |f|
+          (wallets_applicable_on_fees[(f.id || f.object_id)] == wallet.id) ? (f.amount_cents + f.taxes_amount_cents) : 0
+        end
       end
     end
   end