From 803a9801aedd1881989894f0200d066dfc9f8020 Mon Sep 17 00:00:00 2001 From: Matthew Walsh Date: Fri, 7 Aug 2026 14:37:56 +0100 Subject: [PATCH 1/4] feat: add optional getBalance callback to TransactionPayController --- .../transaction-pay-controller/CHANGELOG.md | 4 + .../src/TransactionPayController.test.ts | 25 ++++ .../src/TransactionPayController.ts | 7 +- .../transaction-pay-controller/src/index.ts | 3 + .../transaction-pay-controller/src/types.ts | 32 +++++ .../src/utils/source-amounts.test.ts | 114 ++++++++++++++++++ .../src/utils/source-amounts.ts | 22 +++- 7 files changed, 201 insertions(+), 6 deletions(-) diff --git a/packages/transaction-pay-controller/CHANGELOG.md b/packages/transaction-pay-controller/CHANGELOG.md index 1f5bf9314eb..cd7fcf1b400 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 + ### 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 47ebeb8d12f..7b536e4959f 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,30 @@ 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 2aa314e4ff9..a1dd2f80a0c 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,7 @@ 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 d8bf3f91c1e..f89d3983ca6 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 b53023133af..b8951cd2833 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 a5760435a0e..eb98a7e8c7a 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,120 @@ 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('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 38f96630fbb..92317c10a10 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, @@ -38,6 +40,7 @@ export function updateSourceAmounts( transactionId: string, transactionData: TransactionData | undefined, messenger: TransactionPayControllerMessenger, + getBalance?: GetBalanceCallback, ): void { if (!transactionData) { return; @@ -50,6 +53,11 @@ export function updateSourceAmounts( return; } + const transaction = getTransaction(transactionId, messenger); + const balanceOverride = getBalance && transaction + ? getBalance({ transaction: transaction as TransactionMeta, 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 +68,7 @@ export function updateSourceAmounts( isMaxAmount ?? false, isHyperliquidSource, isPolymarketDepositWallet, + balanceOverride, ); log('Updated post-quote source amounts', { transactionId, sourceAmounts }); transactionData.sourceAmounts = sourceAmounts; @@ -78,6 +87,7 @@ export function updateSourceAmounts( isMaxAmount ?? false, isQuoteRequired, paymentOverride, + balanceOverride, ), ) .filter(Boolean) as TransactionPaySourceAmount[]; @@ -105,6 +115,7 @@ function calculatePostQuoteSourceAmounts( isMaxAmount: boolean, isHyperliquidSource?: boolean, isPolymarketDepositWallet?: boolean, + balanceOverride?: GetBalanceResponse, ): TransactionPaySourceAmount[] { return tokens .filter((token) => { @@ -133,9 +144,9 @@ 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, @@ -162,6 +173,7 @@ function calculateSourceAmount( isMaxAmount: boolean, isQuoteRequired?: boolean, paymentOverride?: PaymentOverride, + balanceOverride?: GetBalanceResponse, ): TransactionPaySourceAmount | undefined { const paymentTokenFiatRate = getTokenFiatRate( messenger, @@ -220,8 +232,8 @@ function calculateSourceAmount( // deposits funded from the money account (e.g. Send to Perps). if (isMaxAmount && paymentOverride !== PaymentOverride.MoneyAccount) { return { - sourceAmountHuman: paymentToken.balanceHuman, - sourceAmountRaw: paymentToken.balanceRaw, + sourceAmountHuman: balanceOverride?.balanceHuman ?? paymentToken.balanceHuman, + sourceAmountRaw: balanceOverride?.balanceRaw ?? paymentToken.balanceRaw, targetTokenAddress: token.address, }; } From 561747afda7e52d4c314a2b6f819c295d328bf4c Mon Sep 17 00:00:00 2001 From: Matthew Walsh Date: Fri, 7 Aug 2026 14:46:27 +0100 Subject: [PATCH 2/4] =?UTF-8?q?fix:=20address=20review=20findings=20?= =?UTF-8?q?=E2=80=94=20prettier,=20getBalance=20guard,=20missing=20tests,?= =?UTF-8?q?=20changelog=20link?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../transaction-pay-controller/CHANGELOG.md | 2 +- .../src/TransactionPayController.test.ts | 4 +- .../src/TransactionPayController.ts | 7 +- .../src/utils/source-amounts.test.ts | 105 +++++++++++++++++- .../src/utils/source-amounts.ts | 23 +++- 5 files changed, 127 insertions(+), 14 deletions(-) diff --git a/packages/transaction-pay-controller/CHANGELOG.md b/packages/transaction-pay-controller/CHANGELOG.md index cd7fcf1b400..6d02e252c78 100644 --- a/packages/transaction-pay-controller/CHANGELOG.md +++ b/packages/transaction-pay-controller/CHANGELOG.md @@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Add optional `getBalance` callback to `TransactionPayControllerOptions` to override the source balance used for max-amount source-amount calculation +- 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 diff --git a/packages/transaction-pay-controller/src/TransactionPayController.test.ts b/packages/transaction-pay-controller/src/TransactionPayController.test.ts index 7b536e4959f..62af4b16302 100644 --- a/packages/transaction-pay-controller/src/TransactionPayController.test.ts +++ b/packages/transaction-pay-controller/src/TransactionPayController.test.ts @@ -939,7 +939,9 @@ describe('TransactionPayController', () => { }); it('forwards getBalance callback to updateSourceAmounts', () => { - const getBalance = jest.fn().mockReturnValue({ balanceHuman: '9.9', balanceRaw: '9900000' }); + const getBalance = jest + .fn() + .mockReturnValue({ balanceHuman: '9.9', balanceRaw: '9900000' }); const controller = createController({ getBalance }); controller.updatePaymentToken({ diff --git a/packages/transaction-pay-controller/src/TransactionPayController.ts b/packages/transaction-pay-controller/src/TransactionPayController.ts index a1dd2f80a0c..7cdde0e6b41 100644 --- a/packages/transaction-pay-controller/src/TransactionPayController.ts +++ b/packages/transaction-pay-controller/src/TransactionPayController.ts @@ -374,7 +374,12 @@ export class TransactionPayController extends BaseController< isPostQuoteUpdated || isAccountOverrideUpdated ) { - updateSourceAmounts(transactionId, current as never, this.messenger, this.#getBalance); + updateSourceAmounts( + transactionId, + current as never, + this.messenger, + this.#getBalance, + ); shouldUpdateQuotes = true; } 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 eb98a7e8c7a..eb0501f978c 100644 --- a/packages/transaction-pay-controller/src/utils/source-amounts.test.ts +++ b/packages/transaction-pay-controller/src/utils/source-amounts.test.ts @@ -268,7 +268,6 @@ describe('Source Amounts Utils', () => { ]); }); - it('uses getBalance override for isMaxAmount standard flow', () => { const getBalance = jest.fn().mockReturnValue({ balanceHuman: '9.9', @@ -282,7 +281,12 @@ describe('Source Amounts Utils', () => { tokens: [TRANSACTION_TOKEN_MOCK], }; - updateSourceAmounts(TRANSACTION_ID_MOCK, transactionData, messenger, getBalance); + updateSourceAmounts( + TRANSACTION_ID_MOCK, + transactionData, + messenger, + getBalance, + ); expect(transactionData.sourceAmounts).toStrictEqual([ { @@ -303,7 +307,12 @@ describe('Source Amounts Utils', () => { tokens: [TRANSACTION_TOKEN_MOCK], }; - updateSourceAmounts(TRANSACTION_ID_MOCK, transactionData, messenger, getBalance); + updateSourceAmounts( + TRANSACTION_ID_MOCK, + transactionData, + messenger, + getBalance, + ); expect(transactionData.sourceAmounts).toStrictEqual([ { @@ -326,7 +335,12 @@ describe('Source Amounts Utils', () => { tokens: [TRANSACTION_TOKEN_MOCK], }; - updateSourceAmounts(TRANSACTION_ID_MOCK, transactionData, messenger, getBalance); + updateSourceAmounts( + TRANSACTION_ID_MOCK, + transactionData, + messenger, + getBalance, + ); // isMaxAmount is false, so fiat-derived amounts should be used (not the override) expect(transactionData.sourceAmounts).toStrictEqual([ @@ -338,6 +352,82 @@ describe('Source Amounts Utils', () => { ]); }); + 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('ignores getBalance override when paymentOverride is MoneyAccount and isMaxAmount is true', () => { + 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, + ); + + // MoneyAccount branch uses fiat-derived amounts, not the getBalance override. + // 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, @@ -368,7 +458,12 @@ describe('Source Amounts Utils', () => { ], }; - updateSourceAmounts(TRANSACTION_ID_MOCK, transactionData, messenger, getBalance); + updateSourceAmounts( + TRANSACTION_ID_MOCK, + transactionData, + messenger, + getBalance, + ); expect(transactionData.sourceAmounts).toStrictEqual([ { diff --git a/packages/transaction-pay-controller/src/utils/source-amounts.ts b/packages/transaction-pay-controller/src/utils/source-amounts.ts index 92317c10a10..3f506379941 100644 --- a/packages/transaction-pay-controller/src/utils/source-amounts.ts +++ b/packages/transaction-pay-controller/src/utils/source-amounts.ts @@ -53,10 +53,16 @@ export function updateSourceAmounts( return; } - const transaction = getTransaction(transactionId, messenger); - const balanceOverride = getBalance && transaction - ? getBalance({ transaction: transaction as TransactionMeta, transactionData }) + const transaction = getBalance + ? getTransaction(transactionId, messenger) : undefined; + const balanceOverride = + getBalance && transaction && isMaxAmount + ? getBalance({ + transaction: transaction as TransactionMeta, + transactionData, + }) + : undefined; // For post-quote flows, source amounts are calculated differently // The source is the transaction's required token, not the selected token @@ -144,8 +150,12 @@ function calculatePostQuoteSourceAmounts( return true; }) .map((token) => ({ - sourceAmountHuman: isMaxAmount ? (balanceOverride?.balanceHuman ?? token.balanceHuman) : token.amountHuman, - sourceAmountRaw: isMaxAmount ? (balanceOverride?.balanceRaw ?? token.balanceRaw) : token.amountRaw, + 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, @@ -232,7 +242,8 @@ function calculateSourceAmount( // deposits funded from the money account (e.g. Send to Perps). if (isMaxAmount && paymentOverride !== PaymentOverride.MoneyAccount) { return { - sourceAmountHuman: balanceOverride?.balanceHuman ?? paymentToken.balanceHuman, + sourceAmountHuman: + balanceOverride?.balanceHuman ?? paymentToken.balanceHuman, sourceAmountRaw: balanceOverride?.balanceRaw ?? paymentToken.balanceRaw, targetTokenAddress: token.address, }; From 8b2dba5a28340f2c9906a155bff77bed65a1a4a7 Mon Sep 17 00:00:00 2001 From: Matthew Walsh Date: Fri, 7 Aug 2026 14:52:05 +0100 Subject: [PATCH 3/4] fix: add missing JSDoc params and remove unnecessary type assertion in source-amounts --- .../src/utils/source-amounts.ts | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/packages/transaction-pay-controller/src/utils/source-amounts.ts b/packages/transaction-pay-controller/src/utils/source-amounts.ts index 3f506379941..7b33050ed79 100644 --- a/packages/transaction-pay-controller/src/utils/source-amounts.ts +++ b/packages/transaction-pay-controller/src/utils/source-amounts.ts @@ -35,6 +35,9 @@ 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, @@ -53,15 +56,13 @@ export function updateSourceAmounts( return; } - const transaction = getBalance - ? getTransaction(transactionId, messenger) - : undefined; + const transaction = + getBalance && isMaxAmount + ? getTransaction(transactionId, messenger) + : undefined; const balanceOverride = - getBalance && transaction && isMaxAmount - ? getBalance({ - transaction: transaction as TransactionMeta, - transactionData, - }) + getBalance && transaction + ? getBalance({ transaction, transactionData }) : undefined; // For post-quote flows, source amounts are calculated differently @@ -113,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( @@ -173,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( From 6d62adf129392ce5f7131d4079d873b4422534e0 Mon Sep 17 00:00:00 2001 From: Matthew Walsh Date: Fri, 7 Aug 2026 16:20:51 +0100 Subject: [PATCH 4/4] feat: allow getBalance callback to override MoneyAccount max branch --- .../src/utils/source-amounts.test.ts | 35 +++++++++++++++++-- .../src/utils/source-amounts.ts | 9 ++++- 2 files changed, 41 insertions(+), 3 deletions(-) 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 eb0501f978c..00e7b682046 100644 --- a/packages/transaction-pay-controller/src/utils/source-amounts.test.ts +++ b/packages/transaction-pay-controller/src/utils/source-amounts.test.ts @@ -386,7 +386,7 @@ describe('Source Amounts Utils', () => { ]); }); - it('ignores getBalance override when paymentOverride is MoneyAccount and isMaxAmount is true', () => { + it('uses getBalance override for MoneyAccount max when getBalance is provided', () => { const getBalance = jest.fn().mockReturnValue({ balanceHuman: '9.9', balanceRaw: '9900000', @@ -417,7 +417,38 @@ describe('Source Amounts Utils', () => { getBalance, ); - // MoneyAccount branch uses fiat-derived amounts, not the getBalance override. + // 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([ { diff --git a/packages/transaction-pay-controller/src/utils/source-amounts.ts b/packages/transaction-pay-controller/src/utils/source-amounts.ts index 7b33050ed79..ebd4e5026ae 100644 --- a/packages/transaction-pay-controller/src/utils/source-amounts.ts +++ b/packages/transaction-pay-controller/src/utils/source-amounts.ts @@ -243,7 +243,14 @@ 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: balanceOverride?.balanceHuman ?? paymentToken.balanceHuman,