diff --git a/packages/transaction-pay-controller/CHANGELOG.md b/packages/transaction-pay-controller/CHANGELOG.md index 1f5bf9314e..6d02e252c7 100644 --- a/packages/transaction-pay-controller/CHANGELOG.md +++ b/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)) diff --git a/packages/transaction-pay-controller/src/TransactionPayController.test.ts b/packages/transaction-pay-controller/src/TransactionPayController.test.ts index 47ebeb8d12..62af4b1630 100644 --- a/packages/transaction-pay-controller/src/TransactionPayController.test.ts +++ b/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', () => { diff --git a/packages/transaction-pay-controller/src/TransactionPayController.ts b/packages/transaction-pay-controller/src/TransactionPayController.ts index 2aa314e4ff..7cdde0e6b4 100644 --- a/packages/transaction-pay-controller/src/TransactionPayController.ts +++ b/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; } diff --git a/packages/transaction-pay-controller/src/index.ts b/packages/transaction-pay-controller/src/index.ts index d8bf3f91c1..f89d3983ca 100644 --- a/packages/transaction-pay-controller/src/index.ts +++ b/packages/transaction-pay-controller/src/index.ts @@ -2,6 +2,9 @@ export type { GetAmountDataCallback, GetAmountDataRequest, GetAmountDataResponse, + GetBalanceCallback, + GetBalanceRequest, + GetBalanceResponse, GetPaymentOverrideDataRequest, GetPaymentOverrideDataResponse, TransactionConfig, diff --git a/packages/transaction-pay-controller/src/types.ts b/packages/transaction-pay-controller/src/types.ts index b53023133a..b8951cd283 100644 --- a/packages/transaction-pay-controller/src/types.ts +++ b/packages/transaction-pay-controller/src/types.ts @@ -215,6 +215,35 @@ export type GetAmountDataCallback = ( request: GetAmountDataRequest, ) => Promise; +/** 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; diff --git a/packages/transaction-pay-controller/src/utils/source-amounts.test.ts b/packages/transaction-pay-controller/src/utils/source-amounts.test.ts index a5760435a0..00e7b68204 100644 --- a/packages/transaction-pay-controller/src/utils/source-amounts.test.ts +++ b/packages/transaction-pay-controller/src/utils/source-amounts.test.ts @@ -268,6 +268,246 @@ describe('Source Amounts Utils', () => { ]); }); + it('uses getBalance override for isMaxAmount standard flow', () => { + const getBalance = jest.fn().mockReturnValue({ + balanceHuman: '9.9', + balanceRaw: '9900000', + }); + + const transactionData: TransactionData = { + isLoading: false, + isMaxAmount: true, + paymentToken: PAYMENT_TOKEN_MOCK, + tokens: [TRANSACTION_TOKEN_MOCK], + }; + + updateSourceAmounts( + TRANSACTION_ID_MOCK, + transactionData, + messenger, + getBalance, + ); + + expect(transactionData.sourceAmounts).toStrictEqual([ + { + sourceAmountHuman: '9.9', + sourceAmountRaw: '9900000', + targetTokenAddress: TRANSACTION_TOKEN_MOCK.address, + }, + ]); + }); + + it('falls back to payment token balance when getBalance returns undefined', () => { + const getBalance = jest.fn().mockReturnValue(undefined); + + const transactionData: TransactionData = { + isLoading: false, + isMaxAmount: true, + paymentToken: PAYMENT_TOKEN_MOCK, + tokens: [TRANSACTION_TOKEN_MOCK], + }; + + updateSourceAmounts( + TRANSACTION_ID_MOCK, + transactionData, + messenger, + getBalance, + ); + + expect(transactionData.sourceAmounts).toStrictEqual([ + { + sourceAmountHuman: PAYMENT_TOKEN_MOCK.balanceHuman, + sourceAmountRaw: PAYMENT_TOKEN_MOCK.balanceRaw, + targetTokenAddress: TRANSACTION_TOKEN_MOCK.address, + }, + ]); + }); + + it('ignores getBalance when isMaxAmount is false', () => { + const getBalance = jest.fn().mockReturnValue({ + balanceHuman: '9.9', + balanceRaw: '9900000', + }); + + const transactionData: TransactionData = { + isLoading: false, + paymentToken: PAYMENT_TOKEN_MOCK, + tokens: [TRANSACTION_TOKEN_MOCK], + }; + + updateSourceAmounts( + TRANSACTION_ID_MOCK, + transactionData, + messenger, + getBalance, + ); + + // isMaxAmount is false, so fiat-derived amounts should be used (not the override) + expect(transactionData.sourceAmounts).toStrictEqual([ + { + sourceAmountHuman: '2', + sourceAmountRaw: '2000000', + targetTokenAddress: TRANSACTION_TOKEN_MOCK.address, + }, + ]); + }); + + it('does not call getBalance when transaction is not found', () => { + // First call (top of updateSourceAmounts) returns undefined; subsequent + // calls (getStrategyContext) return the normal mock so no crash. + getTransactionMock.mockReturnValueOnce(undefined); + + const getBalance = jest.fn().mockReturnValue({ + balanceHuman: '9.9', + balanceRaw: '9900000', + }); + + const transactionData: TransactionData = { + isLoading: false, + isMaxAmount: true, + paymentToken: PAYMENT_TOKEN_MOCK, + tokens: [TRANSACTION_TOKEN_MOCK], + }; + + updateSourceAmounts( + TRANSACTION_ID_MOCK, + transactionData, + messenger, + getBalance, + ); + + expect(getBalance).not.toHaveBeenCalled(); + expect(transactionData.sourceAmounts).toStrictEqual([ + { + sourceAmountHuman: PAYMENT_TOKEN_MOCK.balanceHuman, + sourceAmountRaw: PAYMENT_TOKEN_MOCK.balanceRaw, + targetTokenAddress: TRANSACTION_TOKEN_MOCK.address, + }, + ]); + }); + + it('uses getBalance override for MoneyAccount max when getBalance is provided', () => { + const getBalance = jest.fn().mockReturnValue({ + balanceHuman: '9.9', + balanceRaw: '9900000', + }); + + const transactionData: TransactionData = { + isLoading: false, + isMaxAmount: true, + paymentOverride: PaymentOverride.MoneyAccount, + paymentToken: { + ...PAYMENT_TOKEN_MOCK, + balanceHuman: '0.62', + balanceRaw: '620000', + balanceUsd: '0.62', + }, + tokens: [ + { + ...TRANSACTION_TOKEN_MOCK, + amountUsd: '6.0', + }, + ], + }; + + updateSourceAmounts( + TRANSACTION_ID_MOCK, + transactionData, + messenger, + getBalance, + ); + + // getBalance is provided, so the override is applied even for MoneyAccount. + expect(transactionData.sourceAmounts).toStrictEqual([ + { + sourceAmountHuman: '9.9', + sourceAmountRaw: '9900000', + targetTokenAddress: TRANSACTION_TOKEN_MOCK.address, + }, + ]); + }); + + it('preserves MoneyAccount max guard when getBalance is not provided', () => { + const transactionData: TransactionData = { + isLoading: false, + isMaxAmount: true, + paymentOverride: PaymentOverride.MoneyAccount, + paymentToken: { + ...PAYMENT_TOKEN_MOCK, + balanceHuman: '0.62', + balanceRaw: '620000', + balanceUsd: '0.62', + }, + tokens: [ + { + ...TRANSACTION_TOKEN_MOCK, + amountUsd: '6.0', + }, + ], + }; + + updateSourceAmounts(TRANSACTION_ID_MOCK, transactionData, messenger); + + // No getBalance callback: MoneyAccount guard applies, fiat-derived amounts used. + // usdRate mock is 3.0 -> source human = 6 / 3 = 2, raw = 2 * 10^6. + expect(transactionData.sourceAmounts).toStrictEqual([ + { + sourceAmountHuman: '2', + sourceAmountRaw: '2000000', + targetTokenAddress: TRANSACTION_TOKEN_MOCK.address, + }, + ]); + }); + + it('uses getBalance override for isMaxAmount post-quote flow', () => { + const DESTINATION_TOKEN = { + address: '0xdef' as const, + balanceFiat: '100.00', + balanceHuman: '1.00', + balanceRaw: '1000000000000000000', + balanceUsd: '100.00', + chainId: '0x38' as const, + decimals: 18, + symbol: 'BNB', + }; + + const getBalance = jest.fn().mockReturnValue({ + balanceHuman: '5.5', + balanceRaw: '5500000', + }); + + const transactionData: TransactionData = { + isLoading: false, + isMaxAmount: true, + isPostQuote: true, + paymentToken: DESTINATION_TOKEN, + tokens: [ + { + ...TRANSACTION_TOKEN_MOCK, + skipIfBalance: false, + }, + ], + }; + + updateSourceAmounts( + TRANSACTION_ID_MOCK, + transactionData, + messenger, + getBalance, + ); + + expect(transactionData.sourceAmounts).toStrictEqual([ + { + sourceAmountHuman: '5.5', + sourceAmountRaw: '5500000', + sourceBalanceRaw: '5500000', + sourceChainId: TRANSACTION_TOKEN_MOCK.chainId, + sourceTokenAddress: TRANSACTION_TOKEN_MOCK.address, + targetTokenAddress: DESTINATION_TOKEN.address, + }, + ]); + }); + it('uses fiat-derived source amount for MoneyAccount max instead of payment token balance', () => { // Money account withdrawable (mUSD + vmUSD) is reflected in the typed // required token amount. The pay token's on-chain balance is only the diff --git a/packages/transaction-pay-controller/src/utils/source-amounts.ts b/packages/transaction-pay-controller/src/utils/source-amounts.ts index 38f96630fb..ebd4e5026a 100644 --- a/packages/transaction-pay-controller/src/utils/source-amounts.ts +++ b/packages/transaction-pay-controller/src/utils/source-amounts.ts @@ -18,6 +18,8 @@ import type { import { TransactionPayStrategy } from '../index.js'; import { projectLogger } from '../logger.js'; import type { + GetBalanceCallback, + GetBalanceResponse, TransactionPaySourceAmount, TransactionData, TransactionPayRequiredToken, @@ -33,11 +35,15 @@ const log = createModuleLogger(projectLogger, 'source-amounts'); * @param transactionId - ID of the transaction to update. * @param transactionData - Existing transaction data. * @param messenger - Controller messenger. + * @param getBalance - Optional callback to override the source balance used for max-amount + * calculation. Called only when `isMaxAmount` is true. Return `undefined` to fall back to + * the built-in token balance. */ export function updateSourceAmounts( transactionId: string, transactionData: TransactionData | undefined, messenger: TransactionPayControllerMessenger, + getBalance?: GetBalanceCallback, ): void { if (!transactionData) { return; @@ -50,6 +56,15 @@ export function updateSourceAmounts( return; } + const transaction = + getBalance && isMaxAmount + ? getTransaction(transactionId, messenger) + : undefined; + const balanceOverride = + getBalance && transaction + ? getBalance({ transaction, transactionData }) + : undefined; + // For post-quote flows, source amounts are calculated differently // The source is the transaction's required token, not the selected token if (isPostQuote) { @@ -60,6 +75,7 @@ export function updateSourceAmounts( isMaxAmount ?? false, isHyperliquidSource, isPolymarketDepositWallet, + balanceOverride, ); log('Updated post-quote source amounts', { transactionId, sourceAmounts }); transactionData.sourceAmounts = sourceAmounts; @@ -78,6 +94,7 @@ export function updateSourceAmounts( isMaxAmount ?? false, isQuoteRequired, paymentOverride, + balanceOverride, ), ) .filter(Boolean) as TransactionPaySourceAmount[]; @@ -97,6 +114,7 @@ export function updateSourceAmounts( * @param isMaxAmount - Whether the transaction is a maximum amount transaction. * @param isHyperliquidSource - Whether the source is HyperLiquid (perps withdrawal). * @param isPolymarketDepositWallet - Whether the source is a Polymarket deposit wallet. + * @param balanceOverride - Optional balance override from the `getBalance` callback. * @returns Array of source amounts. */ function calculatePostQuoteSourceAmounts( @@ -105,6 +123,7 @@ function calculatePostQuoteSourceAmounts( isMaxAmount: boolean, isHyperliquidSource?: boolean, isPolymarketDepositWallet?: boolean, + balanceOverride?: GetBalanceResponse, ): TransactionPaySourceAmount[] { return tokens .filter((token) => { @@ -133,9 +152,13 @@ function calculatePostQuoteSourceAmounts( return true; }) .map((token) => ({ - sourceAmountHuman: isMaxAmount ? token.balanceHuman : token.amountHuman, - sourceAmountRaw: isMaxAmount ? token.balanceRaw : token.amountRaw, - sourceBalanceRaw: token.balanceRaw, + sourceAmountHuman: isMaxAmount + ? (balanceOverride?.balanceHuman ?? token.balanceHuman) + : token.amountHuman, + sourceAmountRaw: isMaxAmount + ? (balanceOverride?.balanceRaw ?? token.balanceRaw) + : token.amountRaw, + sourceBalanceRaw: balanceOverride?.balanceRaw ?? token.balanceRaw, sourceChainId: token.chainId, sourceTokenAddress: token.address, targetTokenAddress: paymentToken.address, @@ -152,6 +175,7 @@ function calculatePostQuoteSourceAmounts( * @param isMaxAmount - Whether the transaction is a maximum amount transaction. * @param isQuoteRequired - When true, a quote is always fetched even when source and target tokens are identical. * @param paymentOverride - Optional payment source override for the transaction. + * @param balanceOverride - Optional balance override from the `getBalance` callback. * @returns The source amount or undefined if calculation failed. */ function calculateSourceAmount( @@ -162,6 +186,7 @@ function calculateSourceAmount( isMaxAmount: boolean, isQuoteRequired?: boolean, paymentOverride?: PaymentOverride, + balanceOverride?: GetBalanceResponse, ): TransactionPaySourceAmount | undefined { const paymentTokenFiatRate = getTokenFiatRate( messenger, @@ -218,10 +243,18 @@ function calculateSourceAmount( // reflects the full withdrawable total (mUSD + vmUSD). Using the typed // fiat-derived source keeps isMaxAmount=true (EXACT_INPUT) correct for // deposits funded from the money account (e.g. Send to Perps). - if (isMaxAmount && paymentOverride !== PaymentOverride.MoneyAccount) { + // Exception: when a getBalance callback is provided (balanceOverride is + // defined), the callback is authoritative and bypasses the MoneyAccount + // guard — all balance complexity lives in the callback. + if ( + isMaxAmount && + (balanceOverride !== undefined || + paymentOverride !== PaymentOverride.MoneyAccount) + ) { return { - sourceAmountHuman: paymentToken.balanceHuman, - sourceAmountRaw: paymentToken.balanceRaw, + sourceAmountHuman: + balanceOverride?.balanceHuman ?? paymentToken.balanceHuman, + sourceAmountRaw: balanceOverride?.balanceRaw ?? paymentToken.balanceRaw, targetTokenAddress: token.address, }; }