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))
+
 ## [26.4.0]
 
 ### Changed
@@ -29,6 +33,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Read the `stableTokens` remote feature flag in `getStablecoins` instead of `stable-tokens` ([#9885](https://github.com/MetaMask/core/pull/9885))
 
+### Removed
+
+- **BREAKING:** Remove `resolveSourceAmount` constructor option from `TransactionPayController` and the associated `ResolveSourceAmountCallback`, `ResolveSourceAmountRequest`, and `ResolveSourceAmountResponse` types ([#9802](https://github.com/MetaMask/core/pull/9802))
+  - `resolveSourceAmount` is replaced by the more capable `getBalance` callback, which receives the full transaction and transaction data and returns `{ balanceHuman, balanceRaw }`. Migrate by replacing `resolveSourceAmount: ({ isMaxAmount, paymentOverride }) => ({ sourceAmountRaw })` with `getBalance: ({ transaction, transactionData }) => ({ balanceHuman, balanceRaw })`.
+
 ## [26.3.0]
 
 ### Added

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

@@ -938,9 +938,11 @@ describe('TransactionPayController', () => {
       });
     });
 
-    it('forwards the resolveSourceAmount option to updateSourceAmounts', () => {
-      const resolveSourceAmount = jest.fn();
-      const controller = createController({ resolveSourceAmount });
+    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,
@@ -951,16 +953,14 @@ describe('TransactionPayController', () => {
       const { updateTransactionData } = updatePaymentTokenMock.mock.calls[0][1];
 
       updateTransactionData(TRANSACTION_ID_MOCK, (data) => {
-        data.sourceAmounts = [
-          { sourceAmountHuman: '1.23' } as TransactionPaySourceAmount,
-        ];
+        data.isMaxAmount = true;
       });
 
       expect(updateSourceAmountsMock).toHaveBeenCalledWith(
         TRANSACTION_ID_MOCK,
         expect.any(Object),
         messenger,
-        resolveSourceAmount,
+        getBalance,
       );
     });
   });

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

@@ -14,10 +14,10 @@ import {
 import { QuoteRefresher } from './helpers/QuoteRefresher.js';
 import type {
   GetAmountDataCallback,
+  GetBalanceCallback,
   GetDelegationTransactionCallback,
   GetPaymentOverrideDataCallback,
   PolymarketCallbacks,
-  ResolveSourceAmountCallback,
   TransactionConfig,
   TransactionConfigCallback,
   TransactionData,
@@ -69,6 +69,8 @@ export class TransactionPayController extends BaseController<
 > {
   readonly #getAmountData?: GetAmountDataCallback;
 
+  readonly #getBalance?: GetBalanceCallback;
+
   readonly #getDelegationTransaction: GetDelegationTransactionCallback;
 
   readonly #fiatOptions?: TransactionPayFiatOptions;
@@ -85,18 +87,16 @@ export class TransactionPayController extends BaseController<
 
   readonly #polymarket?: PolymarketCallbacks;
 
-  readonly #resolveSourceAmount?: ResolveSourceAmountCallback;
-
   constructor({
     fiatOptions,
     getAmountData,
+    getBalance,
     getDelegationTransaction,
     getPaymentOverrideData,
     getStrategy,
     getStrategies,
     messenger,
     polymarket,
-    resolveSourceAmount,
     state,
   }: TransactionPayControllerOptions) {
     super({
@@ -107,13 +107,13 @@ export class TransactionPayController extends BaseController<
     });
 
     this.#getAmountData = getAmountData;
+    this.#getBalance = getBalance;
     this.#getDelegationTransaction = getDelegationTransaction;
     this.#fiatOptions = fiatOptions;
     this.#getPaymentOverrideData = getPaymentOverrideData;
     this.#getStrategy = getStrategy;
     this.#getStrategies = getStrategies;
     this.#polymarket = polymarket;
-    this.#resolveSourceAmount = resolveSourceAmount;
 
     this.messenger.registerMethodActionHandlers(
       this,
@@ -378,7 +378,7 @@ export class TransactionPayController extends BaseController<
           transactionId,
           current as never,
           this.messenger,
-          this.#resolveSourceAmount,
+          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,
@@ -19,9 +22,6 @@ export type {
   PolymarketCallbacks,
   QuoteErrorInfo,
   QuoteErrorReason,
-  ResolveSourceAmountCallback,
-  ResolveSourceAmountRequest,
-  ResolveSourceAmountResponse,
   TransactionPayControllerStateChangeEvent,
   TransactionPaymentToken,
   TransactionPayQuote,

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

@@ -215,36 +215,34 @@ export type GetAmountDataCallback = (
   request: GetAmountDataRequest,
 ) => Promise<GetAmountDataResponse>;
 
-/** Request passed to {@link ResolveSourceAmountCallback}. */
-export type ResolveSourceAmountRequest = {
-  /** Whether the user selected the maximum amount. */
-  isMaxAmount: boolean;
-
-  /** Optional payment source override for the transaction. */
-  paymentOverride?: PaymentOverride;
+/** 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;
 };
 
-/** Response returned by {@link ResolveSourceAmountCallback}. */
-export type ResolveSourceAmountResponse = {
-  /**
-   * Exact source token amount in atomic (raw) units. Used verbatim as the
-   * quote's source amount, bypassing the default fiat-derived calculation.
-   */
-  sourceAmountRaw: string;
+/** 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 callback that lets the client supply an exact atomic source amount
- * for a required token, bypassing the default fiat-derived source calculation.
- *
- * Returns `undefined` to fall back to the default calculation. Must be
- * synchronous: it is consumed during synchronous source-amount computation, so
- * the client should read from already-available (cached) state rather than
- * performing async lookups.
+ * 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 ResolveSourceAmountCallback = (
-  request: ResolveSourceAmountRequest,
-) => ResolveSourceAmountResponse | undefined;
+export type GetBalanceCallback = (
+  request: GetBalanceRequest,
+) => GetBalanceResponse | undefined;
 
 /** Callback to update fiat payment state. */
 export type TransactionFiatPaymentCallback = (
@@ -285,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;
 
@@ -309,12 +310,6 @@ export type TransactionPayControllerOptions = {
   /** Callbacks for the Polymarket relayer; required only for the Polymarket deposit-wallet flow. */
   polymarket?: PolymarketCallbacks;
 
-  /**
-   * Optional callback to supply an exact atomic source amount for a required
-   * token, bypassing the default fiat-derived source calculation.
-   */
-  resolveSourceAmount?: ResolveSourceAmountCallback;
-
   /** Initial state of the controller. */
   state?: Partial<TransactionPayControllerState>;
 };