fix(wallets): cascade ongoing usage [ING-112]

## Context

Lago consumes wallet credits through two paths that diverge for
multi-wallet customers. Billing already cascades across wallets in
priority order, but ongoing usage stopped at the first wallet,
absorbing all usage there and going negative instead of spilling onto
the next. The negative ongoing balance is customer-visible and breaks
the two-wallet (paid + free) setup.

## Description

Add Wallets::Balance::AllocateOngoingUsageByWalletsService, which
mirrors the billing-side cascade: it nets current, draft, progressive
and pay-in-advance usage per fee, then distributes it across the
customer's active wallets in priority order. Each applicable wallet
fills to its balance and spills the overflow onto the next; the last
applicable wallet absorbs the remainder and may go negative, since
ongoing usage has no invoice to carry the excess. An over-billed fee
key reduces a per-currency budget the same way billing credits reduce
the invoice total, so wallets never absorb more than the customer's
net unbilled usage. A wallet with an active threshold-based recurring
rule keeps the previous behavior (absorbs everything, may go negative)
so the rule can fire and refill it.

RefreshWalletsService runs the allocator once and feeds each wallet its
precomputed amount. Because the cascade makes allocations
interdependent, every refresh now persists all wallets: the previous
target_wallet_ids narrowing would leave non-target wallets stale. The
per-fee single-wallet assignment and its FindApplicableOnFeesService
and BuildAllocationRulesService helpers are removed, and pre-existing
scenario specs documenting the old "whole fee to the first applicable
wallet" behavior are updated to the new billing-parity cascade.

Author: aquinofb

app/jobs/customers/refresh_wallet_job.rb

@@ -18,12 +18,13 @@ class RefreshWalletJob < ApplicationJob
     retry_on(*Integrations::Aggregator::BaseService.retryable_errors, wait: :polynomially_longer, attempts: 6)
 
     def perform(customer, wallet_ids: nil)
-      # When targeting specific wallets, the refresh is explicitly requested (e.g. balance increase)
-      # so we don't rely on the customer-wide awaiting_wallet_refresh flag.
+      # wallet_ids marks an explicitly requested refresh (e.g. balance increase) that must run
+      # even when the customer-wide awaiting_wallet_refresh flag is not set. The refresh itself
+      # always covers every wallet: the cascade makes allocations interdependent.
       return if wallet_ids.nil? && !customer.awaiting_wallet_refresh?
       return if customer.error_details.tax_error.exists?
 
-      Customers::RefreshWalletsService.call!(customer:, target_wallet_ids: wallet_ids)
+      Customers::RefreshWalletsService.call!(customer:)
     rescue BaseService::ValidationFailure => e
       tax_error = Array(e.messages[:tax_error])
 

app/services/customers/refresh_wallets_service.rb

@@ -2,68 +2,39 @@
 
 module Customers
   class RefreshWalletsService < BaseService
-    Result = BaseResult[:usage_amount_cents, :wallets, :allocation_rules]
+    Result = BaseResult[:wallets]
 
-    def initialize(customer:, include_generating_invoices: false, target_wallet_ids: nil)
+    def initialize(customer:, include_generating_invoices: false)
       @customer = customer
       @include_generating_invoices = include_generating_invoices
-      @target_wallet_ids = target_wallet_ids
 
       super
     end
 
     def call
-      usage_amount_cents = customer.active_subscriptions.map do |subscription|
-        invoice = ::Invoices::CustomerUsageService.call!(customer:, subscription:, usage_filters: UsageFilters::WITHOUT_PRESENTATION_FILTER).invoice
-
-        billed_progressive_invoice_subscriptions = ::Subscriptions::ProgressiveBilledAmount
-          .call(subscription:, include_generating_invoices:)
-          .invoice_subscriptions
-
-        {
-          billed_progressive_invoice_subscriptions:,
-          invoice:,
-          subscription:
-        }
+      wallet_allocations = Wallets::Balance::AllocateOngoingUsageByWalletsService.call!(
+        customer:,
+        wallets: all_wallets,
+        current_usage_fees:,
+        draft_invoices_fees:,
+        progressive_billing_fees:,
+        pay_in_advance_fees:
+      ).wallet_allocations
+
+      # The cascade makes every wallet's allocation depend on the others' balances, so all
+      # wallets must be persisted together; refreshing a subset would leave the rest stale.
+      all_wallets.each do |wallet|
+        Wallets::Balance::RefreshOngoingUsageService.call!(
+          wallet:,
+          ongoing_usage_amount_cents: wallet_allocations[wallet],
+          skip_single_wallet_update: true
+        )
       end
 
-      @allocation_rules = Wallets::BuildAllocationRulesService.call!(customer:).allocation_rules
-
-      # we need to get both: ALL fees and wallets_applicable_on_fees ({ fee_key => wallet_id, ... }) per each fee type
-      # to not rebuild wallet assigning per each fee when going per each separate wallet
-      usage_fees = usage_amount_cents.flat_map { |usage| usage[:invoice].fees }
-      wallets_applicable_on_usage_fees = assign_wallet_per_fee(usage_fees) # { usage_fee_key => wallet_id }
-
-      draft_invoice_fees = customer.invoices.draft.where.not(total_amount_cents: 0).includes(fees: :charge).flat_map(&:fees)
-      wallets_applicable_on_draft_fees = assign_wallet_per_fee(draft_invoice_fees) # { draft_fee_key => wallet_id }
-
-      progressive_billing_fees = usage_amount_cents.flat_map { |usage| usage[:billed_progressive_invoice_subscriptions].flat_map { it.invoice.fees } }
-      wallets_applicable_on_pb_fees = assign_wallet_per_fee(progressive_billing_fees) # { pb_fee_key => wallet_id }
-
-      pay_in_advance_fees = usage_amount_cents.flat_map { |usage| usage[:invoice].fees.select { |f| f.charge.pay_in_advance? } }
-      wallets_applicable_on_adv_fees = assign_wallet_per_fee(pay_in_advance_fees) # { adv_fee_key => wallet_id }
-
-      wallets_to_process = customer.wallets.active.includes(:recurring_transaction_rules)
-      wallets_to_process = wallets_to_process.where(id: target_wallet_ids) if target_wallet_ids
-      wallets_to_process.find_in_batches(batch_size: 100) do |wallets|
-        wallets.each do |wallet|
-          Wallets::Balance::RefreshOngoingUsageService.call!(
-            wallet:,
-            usage_amount_cents:,
-            skip_single_wallet_update: true,
-            current_usage_fees: applicable_fees(usage_fees, wallets_applicable_on_usage_fees, wallet),
-            draft_invoices_fees: applicable_fees(draft_invoice_fees, wallets_applicable_on_draft_fees, wallet),
-            progressive_billing_fees: applicable_fees(progressive_billing_fees, wallets_applicable_on_pb_fees, wallet),
-            pay_in_advance_fees: applicable_fees(pay_in_advance_fees, wallets_applicable_on_adv_fees, wallet)
-          )
-        end
-      end
-      wallets_to_process.update_all(last_ongoing_balance_sync_at: Time.current) # rubocop:disable Rails/SkipsModelValidations
+      Wallet.where(id: all_wallets.map(&:id)).update_all(last_ongoing_balance_sync_at: Time.current) # rubocop:disable Rails/SkipsModelValidations
 
       customer.update!(awaiting_wallet_refresh: false)
 
-      result.usage_amount_cents = usage_amount_cents
-      result.allocation_rules = allocation_rules
       result.wallets = customer.wallets.active.reload
       result
     rescue BaseService::FailedResult => e
@@ -72,32 +43,41 @@ def call
 
     private
 
-    attr_reader :customer, :include_generating_invoices, :allocation_rules, :target_wallet_ids
+    attr_reader :customer, :include_generating_invoices
 
-    def assign_wallet_per_fee(fees)
-      fee_wallet = {}
-      fee_targeting_wallets_enabled = customer.organization.events_targeting_wallets_enabled?
+    def all_wallets
+      @all_wallets ||= customer.wallets.active.includes(:recurring_transaction_rules, :wallet_targets).in_application_order.to_a
+    end
 
-      fees.each do |fee|
-        key = fee.item_key
+    def current_usage_fees
+      @current_usage_fees ||= subscription_usages.flat_map { |usage| usage[:invoice].fees }
+    end
 
-        if fee_targeting_wallets_enabled && fee.charge&.accepts_target_wallet && fee&.grouped_by&.dig("target_wallet_code").present?
-          targeted_wallet = customer.wallets.active.where(code: fee.grouped_by["target_wallet_code"]).ids.first
-          fee_wallet[key] = targeted_wallet
-          next if targeted_wallet
-        end
+    # pay-in-advance fees are a subset of the current usage fees, so filter rather than re-walk.
+    def pay_in_advance_fees
+      current_usage_fees.select { |fee| fee.charge.pay_in_advance? }
+    end
 
-        applicable_wallets = Wallets::FindApplicableOnFeesService
-          .call!(allocation_rules: allocation_rules, fee:, customer_id: customer.id, fee_targeting_wallets_enabled:)
-          .top_priority_wallet
-        fee_wallet[key] = applicable_wallets.presence
-      end
+    def draft_invoices_fees
+      customer.invoices.draft.where.not(total_amount_cents: 0).includes(fees: :charge).flat_map(&:fees)
+    end
 
-      fee_wallet
+    def progressive_billing_fees
+      subscription_usages.flat_map { |usage| usage[:billed_progressive_invoice_subscriptions].flat_map { it.invoice.fees.includes(:charge) } }
     end
 
-    def applicable_fees(fees, fee_map, wallet)
-      fees.select { |fee| fee_map[fee.item_key] == wallet.id && fee.amount_currency == wallet.balance_currency }
+    # One entry per active subscription: its current-usage invoice and the progressively
+    # billed invoice subscriptions used to net already-billed amounts out of ongoing usage.
+    def subscription_usages
+      @subscription_usages ||= customer.active_subscriptions.map do |subscription|
+        invoice = ::Invoices::CustomerUsageService.call!(customer:, subscription:, usage_filters: UsageFilters::WITHOUT_PRESENTATION_FILTER).invoice
+
+        billed_progressive_invoice_subscriptions = ::Subscriptions::ProgressiveBilledAmount
+          .call(subscription:, include_generating_invoices:)
+          .invoice_subscriptions
+
+        {billed_progressive_invoice_subscriptions:, invoice:, subscription:}
+      end
     end
   end
 end

app/services/wallets/balance/allocate_ongoing_usage_by_wallets_service.rb

@@ -0,0 +1,145 @@
+# frozen_string_literal: true
+
+module Wallets
+  module Balance
+    # Distributes ongoing (unbilled) usage across the customer's wallets in priority order,
+    # mirroring Credits::AllocatePrepaidCreditsByWalletsService. A wallet with an active
+    # threshold-based recurring rule is the exception: it absorbs everything and may go
+    # negative (no cascade) so the rule can fire and refill it.
+    class AllocateOngoingUsageByWalletsService < BaseService
+      Result = BaseResult[:wallet_allocations]
+
+      def initialize(customer:, wallets:, current_usage_fees:, draft_invoices_fees:, progressive_billing_fees:, pay_in_advance_fees:)
+        @customer = customer
+        @wallets = wallets
+        @current_usage_fees = current_usage_fees
+        @draft_invoices_fees = draft_invoices_fees
+        @progressive_billing_fees = progressive_billing_fees
+        @pay_in_advance_fees = pay_in_advance_fees
+
+        super
+      end
+
+      def call
+        result.wallet_allocations = calculate_wallet_allocations
+        result
+      end
+
+      private
+
+      attr_reader :customer, :wallets, :current_usage_fees, :draft_invoices_fees,
+        :progressive_billing_fees, :pay_in_advance_fees
+
+      def calculate_wallet_allocations
+        net_amounts = net_usage_by_fee_key
+        budgets = currency_budgets(net_amounts)
+        metas = wallets.map { |wallet| wallet_meta(wallet) }
+        allocations = wallets.index_with(0)
+
+        allocatable_pool(net_amounts).each do |fee_key, key_amount|
+          currency = fee_key.last
+          remaining = [key_amount, budgets[currency]].min
+          applicable = metas.select { |meta| applicable_fee?(fee_key:, wallet: meta[:wallet], targets: meta[:targets], types: meta[:types]) }
+
+          applicable.each_with_index do |meta, index|
+            break if remaining <= 0
+
+            if meta[:threshold] || index == applicable.length - 1
+              # A threshold wallet absorbs everything so its rule can refill it; the last
+              # applicable wallet absorbs the overflow. Both are allowed to go negative.
+              take = remaining
+            else
+              room = meta[:balance] - allocations[meta[:wallet]]
+              next if room <= 0
+              take = [remaining, room].min
+            end
+
+            allocations[meta[:wallet]] += take
+            remaining -= take
+            budgets[currency] -= take
+          end
+        end
+
+        allocations
+      end
+
+      def wallet_meta(wallet)
+        threshold = threshold_wallet?(wallet)
+        {
+          wallet:,
+          targets: wallet.wallet_targets.filter_map { |wt| ["charge", wt.billable_metric_id] if wt.billable_metric_id },
+          types: wallet.allowed_fee_types,
+          threshold:,
+          # Re-read the balance so a concurrent pay-in-advance DecreaseService can't make us cap
+          # against a stale value. Balance only, to keep the eager-loaded associations.
+          balance: threshold ? 0 : Wallet.where(id: wallet.id).pick(:balance_cents)
+        }
+      end
+
+      # Net the fee buckets into a signed amount per fee key. Keys whose net is <= 0
+      # (already fully billed) cannot receive allocations, but their negative nets still
+      # reduce the per-currency budget the same way billing credits reduce the invoice total.
+      def net_usage_by_fee_key
+        net_amounts = Hash.new(0)
+
+        add_to_pool(net_amounts, current_usage_fees) { |fee| fee.amount_cents + fee.taxes_amount_cents }
+        add_to_pool(net_amounts, draft_invoices_fees) { |fee| fee.amount_cents + fee.taxes_amount_cents - fee.precise_coupons_amount_cents }
+        add_to_pool(net_amounts, progressive_billing_fees) { |fee| -(fee.sub_total_excluding_taxes_amount_cents + fee.taxes_amount_cents) }
+        add_to_pool(net_amounts, pay_in_advance_fees) { |fee| -(fee.amount_cents + fee.taxes_amount_cents) }
+
+        net_amounts
+      end
+
+      def allocatable_pool(net_amounts)
+        net_amounts
+          .select { |_, amount| amount.positive? }
+          # Ties broken by fee key so the allocation order is stable across refreshes.
+          .sort_by { |fee_key, amount| [-amount, fee_key.map(&:to_s)] }
+          .to_h
+      end
+
+      # Mirrors billing's remaining_invoice_amount: an over-billed key's negative net offsets
+      # the other keys in its currency, the same way credits offset the invoice at billing.
+      def currency_budgets(net_amounts)
+        budgets = Hash.new(0)
+        net_amounts.each { |fee_key, amount| budgets[fee_key.last] += amount }
+        budgets.transform_values! { |amount| [amount, 0].max }
+      end
+
+      def add_to_pool(remaining, fees)
+        fees.each { |fee| remaining[fee_key(fee)] += yield(fee) }
+      end
+
+      def fee_key(fee)
+        target_wallet_code = if fee_targeting_wallets_enabled? && fee.charge&.accepts_target_wallet
+          fee.grouped_by&.dig("target_wallet_code")
+        end
+
+        [fee.fee_type, fee.charge&.billable_metric_id, target_wallet_code, fee.amount_currency]
+      end
+
+      def applicable_fee?(fee_key:, wallet:, targets:, types:)
+        fee_type, _billable_metric_id, target_wallet_code, currency = fee_key
+
+        return false unless wallet.balance_currency == currency
+        return wallet.code == target_wallet_code if target_wallet_code.present?
+
+        target_match = targets.include?(fee_key.first(2))
+        type_match = types.include?(fee_type)
+        unrestricted_wallet = targets.empty? && types.empty?
+
+        target_match || type_match || unrestricted_wallet
+      end
+
+      def threshold_wallet?(wallet)
+        wallet.recurring_transaction_rules.any? { |rule| rule.currently_active? && rule.threshold? }
+      end
+
+      def fee_targeting_wallets_enabled?
+        return @fee_targeting_wallets_enabled if defined?(@fee_targeting_wallets_enabled)
+
+        @fee_targeting_wallets_enabled = customer.organization.events_targeting_wallets_enabled?
+      end
+    end
+  end
+end

app/services/wallets/balance/refresh_ongoing_usage_service.rb

@@ -5,22 +5,15 @@ module Balance
     class RefreshOngoingUsageService < BaseService
       Result = BaseResult[:wallet]
 
-      def initialize(wallet:, usage_amount_cents:, current_usage_fees:, draft_invoices_fees:, progressive_billing_fees:, pay_in_advance_fees:, skip_single_wallet_update: false)
+      def initialize(wallet:, ongoing_usage_amount_cents:, skip_single_wallet_update: false)
         @wallet = wallet
-        @usage_amount_cents = usage_amount_cents
+        @ongoing_usage_amount_cents = ongoing_usage_amount_cents
         @skip_single_wallet_update = skip_single_wallet_update
-        @current_usage_fees = current_usage_fees
-        @draft_invoices_fees = draft_invoices_fees
-        @progressive_billing_fees = progressive_billing_fees
-        @pay_in_advance_fees = pay_in_advance_fees
 
         super
       end
 
       def call
-        @total_usage_amount_cents = calculate_total_usage_with_limitation
-        @total_billed_usage_amount_cents = calculate_total_billed_usage_amount_cents
-
         # Before this service is called, the wallet is already loaded in the memory. If while calculating current usage we received
         # a pay_in_advance_fee, wallet will be updated by Wallets::Balance::DecreaseService and current wallet version will throw an
         # `Attempted to update a stale object` error. To avoid this, we reload the wallet before updating it.
@@ -35,37 +28,7 @@ def call
 
       private
 
-      attr_reader :wallet, :total_usage_amount_cents, :total_billed_usage_amount_cents, :usage_amount_cents, :skip_single_wallet_update,
-        :current_usage_fees, :draft_invoices_fees, :progressive_billing_fees, :pay_in_advance_fees
-
-      delegate :customer, to: :wallet
-
-      def calculate_total_billed_usage_amount_cents
-        billed_progressive_invoices_amount_cents +
-          billed_pay_in_advance_amount_cents
-      end
-
-      def billed_progressive_invoices_amount_cents
-        progressive_billing_fees.sum do |fee|
-          fee.taxes_amount_cents + fee.sub_total_excluding_taxes_amount_cents
-        end
-      end
-
-      def draft_invoices_total_amount_cents
-        draft_invoices_fees.sum do |fee|
-          fee.amount_cents + fee.taxes_amount_cents - fee.precise_coupons_amount_cents
-        end
-      end
-
-      def billed_pay_in_advance_amount_cents
-        # Invoice that is returned from CustomerUsageService includes the taxes in total_usage
-        # so if the fees ae already paid, we should exclude fees AND their taxes
-        pay_in_advance_fees.sum { |fee| fee.amount_cents + fee.taxes_amount_cents }
-      end
-
-      def calculate_total_usage_with_limitation
-        current_usage_fees.sum { |fee| fee.amount_cents + fee.taxes_amount_cents }
-      end
+      attr_reader :wallet, :ongoing_usage_amount_cents, :skip_single_wallet_update
 
       def wallet_update_params
         params = {
@@ -89,9 +52,7 @@ def currency
       end
 
       def ongoing_usage_balance_cents
-        @ongoing_usage_balance_cents ||= total_usage_amount_cents +
-          draft_invoices_total_amount_cents -
-          total_billed_usage_amount_cents
+        ongoing_usage_amount_cents
       end
 
       def credits_ongoing_usage_balance