feat(transaction-pay-controller): add optional getBalance callback for max-amount source-amount calculation

packages/transaction-pay-controller/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add optional `getBalance` callback to `TransactionPayControllerOptions` to override the source balance used for max-amount source-amount calculation ([#9802](https://github.com/MetaMask/core/pull/9802))
+
 ### Changed
 
 - Bump `@metamask/transaction-controller` from `^69.5.0` to `^69.5.1` ([#9798](https://github.com/MetaMask/core/pull/9798))

packages/transaction-pay-controller/src/TransactionPayController.test.ts

@@ -924,6 +924,7 @@ describe('TransactionPayController', () => {
           sourceAmounts: [{ sourceAmountHuman: '1.23' }],
         }),
         messenger,
+        undefined,
       );
 
       expect(updateQuotesMock).toHaveBeenCalledWith({
@@ -936,6 +937,32 @@ describe('TransactionPayController', () => {
         updateTransactionData: expect.any(Function),
       });
     });
+
+    it('forwards getBalance callback to updateSourceAmounts', () => {
+      const getBalance = jest
+        .fn()
+        .mockReturnValue({ balanceHuman: '9.9', balanceRaw: '9900000' });
+      const controller = createController({ getBalance });
+
+      controller.updatePaymentToken({
+        transactionId: TRANSACTION_ID_MOCK,
+        tokenAddress: TOKEN_ADDRESS_MOCK,
+        chainId: CHAIN_ID_MOCK,
+      });
+
+      const { updateTransactionData } = updatePaymentTokenMock.mock.calls[0][1];
+
+      updateTransactionData(TRANSACTION_ID_MOCK, (data) => {
+        data.isMaxAmount = true;
+      });
+
+      expect(updateSourceAmountsMock).toHaveBeenCalledWith(
+        TRANSACTION_ID_MOCK,
+        expect.any(Object),
+        messenger,
+        getBalance,
+      );
+    });
   });
 
   describe('transaction data removal', () => {

packages/transaction-pay-controller/src/TransactionPayController.ts

@@ -14,6 +14,7 @@ import {
 import { QuoteRefresher } from './helpers/QuoteRefresher.js';
 import type {
   GetAmountDataCallback,
+  GetBalanceCallback,
   GetDelegationTransactionCallback,
   GetPaymentOverrideDataCallback,
   PolymarketCallbacks,
@@ -68,6 +69,8 @@ export class TransactionPayController extends BaseController<
 > {
   readonly #getAmountData?: GetAmountDataCallback;
 
+  readonly #getBalance?: GetBalanceCallback;
+
   readonly #getDelegationTransaction: GetDelegationTransactionCallback;
 
   readonly #fiatOptions?: TransactionPayFiatOptions;
@@ -87,6 +90,7 @@ export class TransactionPayController extends BaseController<
   constructor({
     fiatOptions,
     getAmountData,
+    getBalance,
     getDelegationTransaction,
     getPaymentOverrideData,
     getStrategy,
@@ -103,6 +107,7 @@ export class TransactionPayController extends BaseController<
     });
 
     this.#getAmountData = getAmountData;
+    this.#getBalance = getBalance;
     this.#getDelegationTransaction = getDelegationTransaction;
     this.#fiatOptions = fiatOptions;
     this.#getPaymentOverrideData = getPaymentOverrideData;
@@ -369,7 +374,12 @@ export class TransactionPayController extends BaseController<
         isPostQuoteUpdated ||
         isAccountOverrideUpdated
       ) {
-        updateSourceAmounts(transactionId, current as never, this.messenger);
+        updateSourceAmounts(
+          transactionId,
+          current as never,
+          this.messenger,
+          this.#getBalance,
+        );
 
         shouldUpdateQuotes = true;
       }

packages/transaction-pay-controller/src/index.ts

@@ -2,6 +2,9 @@ export type {
   GetAmountDataCallback,
   GetAmountDataRequest,
   GetAmountDataResponse,
+  GetBalanceCallback,
+  GetBalanceRequest,
+  GetBalanceResponse,
   GetPaymentOverrideDataRequest,
   GetPaymentOverrideDataResponse,
   TransactionConfig,

packages/transaction-pay-controller/src/types.ts

@@ -215,6 +215,35 @@ export type GetAmountDataCallback = (
   request: GetAmountDataRequest,
 ) => Promise<GetAmountDataResponse>;
 
+/** Request passed to {@link GetBalanceCallback}. */
+export type GetBalanceRequest = {
+  /** Metadata of the transaction whose source balance is being resolved. */
+  transaction: TransactionMeta;
+  /** Pay-controller state for the transaction. */
+  transactionData: TransactionData;
+};
+
+/** Balance override returned by {@link GetBalanceCallback}. */
+export type GetBalanceResponse = {
+  /** Balance in human-readable format factoring token decimals. */
+  balanceHuman: string;
+  /** Balance in atomic format without factoring token decimals. */
+  balanceRaw: string;
+};
+
+/**
+ * Optional client-supplied callback that overrides the built-in
+ * pay-token / required-token balance lookup used for `isMaxAmount`
+ * source-amount calculation. Enables alternate balance sources
+ * (perps, predict, money-account, post-quote, etc.) without adding
+ * conditional branches inside the controller. MUST be synchronous:
+ * it runs inside the controller state-update block.
+ * Return `undefined` to fall back to the built-in token balance.
+ */
+export type GetBalanceCallback = (
+  request: GetBalanceRequest,
+) => GetBalanceResponse | undefined;
+
 /** Callback to update fiat payment state. */
 export type TransactionFiatPaymentCallback = (
   fiatPayment: TransactionFiatPayment,
@@ -254,6 +283,9 @@ export type TransactionPayControllerOptions = {
   /** Optional callback to re-encode nested transaction calldata for a given amount. */
   getAmountData?: GetAmountDataCallback;
 
+  /** Optional callback to override the source balance used for max-amount calculation. */
+  getBalance?: GetBalanceCallback;
+
   /** Callback to convert a transaction into a redeem delegation. */
   getDelegationTransaction: GetDelegationTransactionCallback;