From e14395d4806950caf53fb4df5870057c9a69dab9 Mon Sep 17 00:00:00 2001 From: Matthew Walsh Date: Tue, 13 Oct 2026 09:52:45 +0100 Subject: [PATCH 1/2] refactor(transaction-pay-controller): move Polymarket withdraw helpers out of the Relay strategy --- oxlint-suppressions.json | 28 +++++++++---------- .../src/strategy/relay/relay-quotes.ts | 10 +++---- .../src/strategy/relay/relay-submit.test.ts | 4 +-- .../src/strategy/relay/relay-submit.ts | 8 +++--- .../src/strategy/relay/relay-validation.ts | 8 +++--- .../relay => utils}/polymarket/calldata.ts | 0 .../relay => utils}/polymarket/constants.ts | 0 .../polymarket/withdraw.test.ts | 15 ++++++---- .../relay => utils}/polymarket/withdraw.ts | 20 ++++++------- 9 files changed, 48 insertions(+), 45 deletions(-) rename packages/transaction-pay-controller/src/{strategy/relay => utils}/polymarket/calldata.ts (100%) rename packages/transaction-pay-controller/src/{strategy/relay => utils}/polymarket/constants.ts (100%) rename packages/transaction-pay-controller/src/{strategy/relay => utils}/polymarket/withdraw.test.ts (97%) rename packages/transaction-pay-controller/src/{strategy/relay => utils}/polymarket/withdraw.ts (97%) diff --git a/oxlint-suppressions.json b/oxlint-suppressions.json index 8d49c12be2b..dfa57acad51 100644 --- a/oxlint-suppressions.json +++ b/oxlint-suppressions.json @@ -6630,20 +6630,6 @@ "count": 1 } }, - "packages/transaction-pay-controller/src/strategy/relay/polymarket/withdraw.test.ts": { - "typescript/no-unsafe-argument": { - "count": 16 - }, - "typescript/no-unsafe-assignment": { - "count": 4 - }, - "typescript/no-unsafe-call": { - "count": 8 - }, - "typescript/no-unsafe-member-access": { - "count": 24 - } - }, "packages/transaction-pay-controller/src/strategy/relay/relay-max.test.ts": { "typescript/no-unsafe-argument": { "count": 30 @@ -6818,6 +6804,20 @@ "count": 1 } }, + "packages/transaction-pay-controller/src/utils/polymarket/withdraw.test.ts": { + "typescript/no-unsafe-argument": { + "count": 16 + }, + "typescript/no-unsafe-assignment": { + "count": 4 + }, + "typescript/no-unsafe-call": { + "count": 8 + }, + "typescript/no-unsafe-member-access": { + "count": 24 + } + }, "packages/transaction-pay-controller/src/utils/provider.test.ts": { "no-unused-vars": { "count": 1 diff --git a/packages/transaction-pay-controller/src/strategy/relay/relay-quotes.ts b/packages/transaction-pay-controller/src/strategy/relay/relay-quotes.ts index e5190f70679..4fc9f4c6781 100644 --- a/packages/transaction-pay-controller/src/strategy/relay/relay-quotes.ts +++ b/packages/transaction-pay-controller/src/strategy/relay/relay-quotes.ts @@ -48,6 +48,11 @@ import { resolveGasStationCost, } from '../../utils/gas-payment.js'; import { calculateGasCost } from '../../utils/gas.js'; +import { + applyPolymarketDepositWalletOverrides, + getPredictWithdrawSafeAddress, + isPredictWithdraw, +} from '../../utils/polymarket/withdraw.js'; import { estimateQuoteGasLimits } from '../../utils/quote-gas.js'; import type { QuoteGasTransaction } from '../../utils/quote-gas.js'; import { resolveNonAtomicRecipient } from '../../utils/second-leg.js'; @@ -60,11 +65,6 @@ import { import { getQuotePricing } from '../../utils/trade-type.js'; import { TOKEN_TRANSFER_FOUR_BYTE } from './constants.js'; import { applyHyperliquidActivationFee } from './hyperliquid-activation.js'; -import { - applyPolymarketDepositWalletOverrides, - getPredictWithdrawSafeAddress, - isPredictWithdraw, -} from './polymarket/withdraw.js'; import { fetchRelayQuote } from './relay-api.js'; import { getRelayMaxQuote, diff --git a/packages/transaction-pay-controller/src/strategy/relay/relay-submit.test.ts b/packages/transaction-pay-controller/src/strategy/relay/relay-submit.test.ts index b3fc6520d2d..f493e062c88 100644 --- a/packages/transaction-pay-controller/src/strategy/relay/relay-submit.test.ts +++ b/packages/transaction-pay-controller/src/strategy/relay/relay-submit.test.ts @@ -46,7 +46,7 @@ jest.mock('../../utils/token'); jest.mock('../../utils/transaction'); jest.mock('../../utils/feature-flags'); jest.mock('./hyperliquid-withdraw'); -jest.mock('./polymarket/withdraw'); +jest.mock('../../utils/polymarket/withdraw'); jest.mock('./relay-submit-execute'); const NETWORK_CLIENT_ID_MOCK = 'networkClientIdMock'; const TRANSACTION_HASH_MOCK = '0x1234'; @@ -1797,7 +1797,7 @@ describe('Relay Submit Utils', () => { submitPolymarketWithdraw: jest.Mock; sweepPolymarketDepositWallet: jest.Mock; } { - const mod = jest.requireMock('./polymarket/withdraw'); + const mod = jest.requireMock('../../utils/polymarket/withdraw'); return { submitPolymarketWithdraw: mod.submitPolymarketWithdraw as jest.Mock, sweepPolymarketDepositWallet: diff --git a/packages/transaction-pay-controller/src/strategy/relay/relay-submit.ts b/packages/transaction-pay-controller/src/strategy/relay/relay-submit.ts index 1c500087ce9..82ff92c1fcd 100644 --- a/packages/transaction-pay-controller/src/strategy/relay/relay-submit.ts +++ b/packages/transaction-pay-controller/src/strategy/relay/relay-submit.ts @@ -30,6 +30,10 @@ import { logGasPaymentOutcome, resolveGasPayment, } from '../../utils/gas-payment.js'; +import { + sweepPolymarketDepositWallet, + submitPolymarketWithdraw, +} from '../../utils/polymarket/withdraw.js'; import { getNetworkClientId } from '../../utils/provider.js'; import { submitSecondLeg } from '../../utils/second-leg.js'; import { @@ -50,10 +54,6 @@ import { RELAY_PENDING_STATUSES, } from './constants.js'; import { submitHyperliquidWithdraw } from './hyperliquid-withdraw.js'; -import { - sweepPolymarketDepositWallet, - submitPolymarketWithdraw, -} from './polymarket/withdraw.js'; import { getRelayStatus } from './relay-api.js'; import { submitViaRelayExecute } from './relay-submit-execute.js'; import type { diff --git a/packages/transaction-pay-controller/src/strategy/relay/relay-validation.ts b/packages/transaction-pay-controller/src/strategy/relay/relay-validation.ts index 8b1c9337aca..d3ed5365daf 100644 --- a/packages/transaction-pay-controller/src/strategy/relay/relay-validation.ts +++ b/packages/transaction-pay-controller/src/strategy/relay/relay-validation.ts @@ -17,16 +17,16 @@ import { getEIP7702UpgradeContractAddress, isRelayValidationEnabled, } from '../../utils/feature-flags.js'; +import { + buildPolymarketDepositWalletSimulation, + isPredictWithdraw, +} from '../../utils/polymarket/withdraw.js'; import { validateQuoteExecution, QuoteError, isQuoteError, } from '../../utils/validation.js'; import type { QuoteSimulation } from '../../utils/validation.js'; -import { - buildPolymarketDepositWalletSimulation, - isPredictWithdraw, -} from './polymarket/withdraw.js'; import { getRelayExecuteRequest } from './relay-submit-execute.js'; import { getRelaySubmitCalls } from './relay-submit.js'; import type { RelayExecuteRequest, RelayQuote } from './types.js'; diff --git a/packages/transaction-pay-controller/src/strategy/relay/polymarket/calldata.ts b/packages/transaction-pay-controller/src/utils/polymarket/calldata.ts similarity index 100% rename from packages/transaction-pay-controller/src/strategy/relay/polymarket/calldata.ts rename to packages/transaction-pay-controller/src/utils/polymarket/calldata.ts diff --git a/packages/transaction-pay-controller/src/strategy/relay/polymarket/constants.ts b/packages/transaction-pay-controller/src/utils/polymarket/constants.ts similarity index 100% rename from packages/transaction-pay-controller/src/strategy/relay/polymarket/constants.ts rename to packages/transaction-pay-controller/src/utils/polymarket/constants.ts diff --git a/packages/transaction-pay-controller/src/strategy/relay/polymarket/withdraw.test.ts b/packages/transaction-pay-controller/src/utils/polymarket/withdraw.test.ts similarity index 97% rename from packages/transaction-pay-controller/src/strategy/relay/polymarket/withdraw.test.ts rename to packages/transaction-pay-controller/src/utils/polymarket/withdraw.test.ts index 2a3cae9c350..c5c3ffcdd38 100644 --- a/packages/transaction-pay-controller/src/strategy/relay/polymarket/withdraw.test.ts +++ b/packages/transaction-pay-controller/src/utils/polymarket/withdraw.test.ts @@ -6,11 +6,14 @@ import type { Hex } from '@metamask/utils'; import { POLYGON_PUSD_ADDRESS, POLYGON_USDCE_ADDRESS, -} from '../../../constants.js'; -import { getMessengerMock } from '../../../tests/messenger-mock.js'; -import type { QuoteRequest, TransactionPayQuote } from '../../../types.js'; -import { getLiveTokenBalance } from '../../../utils/token.js'; -import type { RelayQuote, RelayQuoteRequest } from '../types.js'; +} from '../../constants.js'; +import type { + RelayQuote, + RelayQuoteRequest, +} from '../../strategy/relay/types.js'; +import { getMessengerMock } from '../../tests/messenger-mock.js'; +import type { QuoteRequest, TransactionPayQuote } from '../../types.js'; +import { getLiveTokenBalance } from '../token.js'; import { POLYMARKET_COLLATERAL_OFFRAMP_POLYGON, POLYMARKET_COLLATERAL_ONRAMP_POLYGON, @@ -24,7 +27,7 @@ import { sweepPolymarketDepositWallet, } from './withdraw.js'; -jest.mock('../../../utils/token'); +jest.mock('../token'); const EOA_MOCK = '0x1111111111111111111111111111111111111111' as Hex; const DEPOSIT_WALLET_MOCK = '0x2222222222222222222222222222222222222222' as Hex; diff --git a/packages/transaction-pay-controller/src/strategy/relay/polymarket/withdraw.ts b/packages/transaction-pay-controller/src/utils/polymarket/withdraw.ts similarity index 97% rename from packages/transaction-pay-controller/src/strategy/relay/polymarket/withdraw.ts rename to packages/transaction-pay-controller/src/utils/polymarket/withdraw.ts index 3fa526df88b..eb075484211 100644 --- a/packages/transaction-pay-controller/src/strategy/relay/polymarket/withdraw.ts +++ b/packages/transaction-pay-controller/src/utils/polymarket/withdraw.ts @@ -10,21 +10,21 @@ import { CHAIN_ID_POLYGON, POLYGON_PUSD_ADDRESS, POLYGON_USDCE_ADDRESS, -} from '../../../constants.js'; -import { projectLogger } from '../../../logger.js'; -import type { - QuoteRequest, - TransactionPayControllerMessenger, - TransactionPayQuote, -} from '../../../types.js'; -import { getLiveTokenBalance } from '../../../utils/token.js'; -import type { QuoteSimulation } from '../../../utils/validation.js'; +} from '../../constants.js'; +import { projectLogger } from '../../logger.js'; import type { RelayQuote, RelayQuoteRequest, RelayStatus, RelayTransactionStep, -} from '../types.js'; +} from '../../strategy/relay/types.js'; +import type { + QuoteRequest, + TransactionPayControllerMessenger, + TransactionPayQuote, +} from '../../types.js'; +import { getLiveTokenBalance } from '../token.js'; +import type { QuoteSimulation } from '../validation.js'; import { encodeApprove, encodeUnwrap, From a3a798b3ad3589321afc8b52f2c8ff7f09cc7541 Mon Sep 17 00:00:00 2001 From: Matthew Walsh Date: Tue, 13 Oct 2026 10:00:47 +0100 Subject: [PATCH 2/2] refactor(transaction-pay-controller): decouple Polymarket withdraw helpers from Relay types --- oxlint-suppressions.json | 6 +- .../src/strategy/relay/deposit-step.test.ts | 52 ++++ .../src/strategy/relay/deposit-step.ts | 29 ++ .../src/strategy/relay/relay-quotes.ts | 28 +- .../src/strategy/relay/relay-submit.test.ts | 56 ++-- .../src/strategy/relay/relay-submit.ts | 17 +- .../src/strategy/relay/relay-validation.ts | 12 +- .../src/utils/polymarket/withdraw.test.ts | 289 +++++++----------- .../src/utils/polymarket/withdraw.ts | 246 ++++++++------- 9 files changed, 391 insertions(+), 344 deletions(-) create mode 100644 packages/transaction-pay-controller/src/strategy/relay/deposit-step.test.ts create mode 100644 packages/transaction-pay-controller/src/strategy/relay/deposit-step.ts diff --git a/oxlint-suppressions.json b/oxlint-suppressions.json index dfa57acad51..353b704d1df 100644 --- a/oxlint-suppressions.json +++ b/oxlint-suppressions.json @@ -6688,7 +6688,7 @@ "count": 1 }, "typescript/no-unsafe-assignment": { - "count": 34 + "count": 39 }, "typescript/no-unsafe-call": { "count": 17 @@ -6806,10 +6806,10 @@ }, "packages/transaction-pay-controller/src/utils/polymarket/withdraw.test.ts": { "typescript/no-unsafe-argument": { - "count": 16 + "count": 1 }, "typescript/no-unsafe-assignment": { - "count": 4 + "count": 6 }, "typescript/no-unsafe-call": { "count": 8 diff --git a/packages/transaction-pay-controller/src/strategy/relay/deposit-step.test.ts b/packages/transaction-pay-controller/src/strategy/relay/deposit-step.test.ts new file mode 100644 index 00000000000..0bd3cd30ee3 --- /dev/null +++ b/packages/transaction-pay-controller/src/strategy/relay/deposit-step.test.ts @@ -0,0 +1,52 @@ +import { getRelayDepositData, hasRelayDepositStep } from './deposit-step.js'; +import type { RelayQuote } from './types.js'; + +const DEPOSIT_DATA_MOCK = '0xa9059cbb'; + +function buildQuote(steps: unknown[]): RelayQuote { + return { steps } as RelayQuote; +} + +describe('Relay deposit step', () => { + describe('hasRelayDepositStep', () => { + it('returns true when a step is the deposit step', () => { + expect(hasRelayDepositStep(buildQuote([{ id: 'deposit' }]).steps)).toBe( + true, + ); + }); + + it('returns false for swap-only routes', () => { + expect(hasRelayDepositStep(buildQuote([{ id: 'swap' }]).steps)).toBe( + false, + ); + }); + }); + + describe('getRelayDepositData', () => { + it('returns the deposit transaction calldata', () => { + expect( + getRelayDepositData( + buildQuote([ + { + id: 'deposit', + items: [{ data: { data: DEPOSIT_DATA_MOCK } }], + kind: 'transaction', + }, + ]), + ), + ).toBe(DEPOSIT_DATA_MOCK); + }); + + it('returns undefined when the deposit step is not a transaction', () => { + expect( + getRelayDepositData( + buildQuote([{ id: 'deposit', items: [], kind: 'signature' }]), + ), + ).toBeUndefined(); + }); + + it('returns undefined when there is no deposit step', () => { + expect(getRelayDepositData(buildQuote([]))).toBeUndefined(); + }); + }); +}); diff --git a/packages/transaction-pay-controller/src/strategy/relay/deposit-step.ts b/packages/transaction-pay-controller/src/strategy/relay/deposit-step.ts new file mode 100644 index 00000000000..37229bf904c --- /dev/null +++ b/packages/transaction-pay-controller/src/strategy/relay/deposit-step.ts @@ -0,0 +1,29 @@ +import type { Hex } from '@metamask/utils'; + +import type { RelayQuote, RelayTransactionStep } from './types.js'; + +/** + * Whether a Relay quote routes through a deposit step, i.e. the user funds the + * request by transferring to Relay rather than calling a swap contract. + * + * @param steps - Relay quote steps. + * @returns True when the quote has a deposit step. + */ +export function hasRelayDepositStep(steps: RelayQuote['steps']): boolean { + return steps.some((step) => step.id === 'deposit'); +} + +/** + * Get the calldata of the Relay deposit transaction step. + * + * @param quote - Relay quote. + * @returns The deposit calldata, or `undefined` when there is none. + */ +export function getRelayDepositData(quote: RelayQuote): Hex | undefined { + const depositStep = quote.steps.find( + (step): step is RelayTransactionStep => + step.id === 'deposit' && step.kind === 'transaction', + ); + + return depositStep?.items[0]?.data?.data; +} diff --git a/packages/transaction-pay-controller/src/strategy/relay/relay-quotes.ts b/packages/transaction-pay-controller/src/strategy/relay/relay-quotes.ts index 4fc9f4c6781..9b2a0577f97 100644 --- a/packages/transaction-pay-controller/src/strategy/relay/relay-quotes.ts +++ b/packages/transaction-pay-controller/src/strategy/relay/relay-quotes.ts @@ -49,9 +49,9 @@ import { } from '../../utils/gas-payment.js'; import { calculateGasCost } from '../../utils/gas.js'; import { - applyPolymarketDepositWalletOverrides, + getPolymarketDepositWalletOverrides, + getPredictWithdrawFeeTokenAccount, getPredictWithdrawSafeAddress, - isPredictWithdraw, } from '../../utils/polymarket/withdraw.js'; import { estimateQuoteGasLimits } from '../../utils/quote-gas.js'; import type { QuoteGasTransaction } from '../../utils/quote-gas.js'; @@ -64,6 +64,7 @@ import { } from '../../utils/token.js'; import { getQuotePricing } from '../../utils/trade-type.js'; import { TOKEN_TRANSFER_FOUR_BYTE } from './constants.js'; +import { hasRelayDepositStep } from './deposit-step.js'; import { applyHyperliquidActivationFee } from './hyperliquid-activation.js'; import { fetchRelayQuote } from './relay-api.js'; import { @@ -363,11 +364,17 @@ async function getSingleQuote( }; if (effectiveRequest.isPolymarketDepositWallet) { - await applyPolymarketDepositWalletOverrides( - body, - effectiveRequest, + const overrides = await getPolymarketDepositWalletOverrides( + from, messenger, ); + const { depositWallet } = overrides; + + body.originCurrency = overrides.sourceTokenAddress; + body.user = depositWallet; + body.refundTo = depositWallet; + body.useDepositAddress = true; + body.strict = true; } const isAtomic = effectiveRequest.atomic !== false; @@ -926,8 +933,6 @@ async function calculateSourceNetworkCost( const { chainId, data, maxFeePerGas, maxPriorityFeePerGas, to, value } = relayParams[0]; - const isPredictWithdrawFlow = isPredictWithdraw(request, transaction); - // `fromOverride = Safe proxy` is only valid for deposit-style Relay routes // where the deposit contract reads the user's source-token balance directly. // Same-chain destinations route through DEX swap aggregators that frequently @@ -938,8 +943,8 @@ async function calculateSourceNetworkCost( // EOA `from` so simulation succeeds. const fromOverride = getPredictWithdrawSafeAddress( request, - quote.steps, transaction, + hasRelayDepositStep(quote.steps), ); // For post-quote flows the original transaction will be prepended to the @@ -999,9 +1004,10 @@ async function calculateSourceNetworkCost( // return nothing and force users to hold POL. // (`useFromOverride` only governs the gas-estimation `from` address, where // swap-style routes need EOA because DEX routers reject contract callers.) - const proxyFeeTokenAccount = isPredictWithdrawFlow - ? request.refundTo - : undefined; + const proxyFeeTokenAccount = getPredictWithdrawFeeTokenAccount( + request, + transaction, + ); const gasStationCost = await resolveGasStationCost({ accountSupports7702, diff --git a/packages/transaction-pay-controller/src/strategy/relay/relay-submit.test.ts b/packages/transaction-pay-controller/src/strategy/relay/relay-submit.test.ts index f493e062c88..b8af9724d6b 100644 --- a/packages/transaction-pay-controller/src/strategy/relay/relay-submit.test.ts +++ b/packages/transaction-pay-controller/src/strategy/relay/relay-submit.test.ts @@ -34,7 +34,7 @@ import { import { FALLBACK_HASH, RELAY_STATUS_URL } from './constants.js'; import { submitViaRelayExecute } from './relay-submit-execute.js'; import { getRelaySubmitCalls, submitRelayQuotes } from './relay-submit.js'; -import type { RelayQuote } from './types.js'; +import type { RelayQuote, RelayTransactionStep } from './types.js'; jest.mock('../../utils/second-leg', () => ({ ...jest.requireActual( @@ -1814,6 +1814,7 @@ describe('Relay Submit Utils', () => { }); sweepPolymarketDepositWallet.mockResolvedValue(undefined); request.quotes[0].request.isPolymarketDepositWallet = true; + request.quotes[0].original.steps[0].id = 'deposit'; request.quotes[0].original.steps[0].kind = 'transaction'; }); @@ -1822,28 +1823,32 @@ describe('Relay Submit Utils', () => { await submitRelayQuotes(request); - expect(submitPolymarketWithdraw).toHaveBeenCalledWith( - request.quotes[0], - FROM_MOCK, + expect(submitPolymarketWithdraw).toHaveBeenCalledWith({ + depositData: ( + request.quotes[0].original.steps[0] as RelayTransactionStep + ).items[0].data.data, + from: FROM_MOCK, messenger, - ); + sourceAmountRaw: request.quotes[0].sourceAmount.raw, + }); expect(addTransactionMock).not.toHaveBeenCalled(); expect(addTransactionBatchMock).not.toHaveBeenCalled(); }); - it('runs the USDC.e sweep with the success status on success', async () => { + it('runs the USDC.e sweep without waiting for a refund on success', async () => { const { sweepPolymarketDepositWallet } = getPolymarketMocks(); await submitRelayQuotes(request); - expect(sweepPolymarketDepositWallet).toHaveBeenCalledWith( - FROM_MOCK, + expect(sweepPolymarketDepositWallet).toHaveBeenCalledWith({ + from: FROM_MOCK, + isRefund: false, messenger, - { relayStatus: 'success', preSubmitUsdceBalance: 0n }, - ); + preSubmitUsdceBalance: 0n, + }); }); - it('passes the refund status and pre-submit balance to the sweep on refund', async () => { + it('sweeps as a refund with the pre-submit balance on refund', async () => { const { submitPolymarketWithdraw, sweepPolymarketDepositWallet } = getPolymarketMocks(); submitPolymarketWithdraw.mockResolvedValue({ @@ -1858,14 +1863,15 @@ describe('Relay Submit Utils', () => { await expect(submitRelayQuotes(request)).rejects.toThrow( 'Relay: Request failed with status: refund', ); - expect(sweepPolymarketDepositWallet).toHaveBeenCalledWith( - FROM_MOCK, + expect(sweepPolymarketDepositWallet).toHaveBeenCalledWith({ + from: FROM_MOCK, + isRefund: true, messenger, - { relayStatus: 'refund', preSubmitUsdceBalance: 1000000n }, - ); + preSubmitUsdceBalance: 1000000n, + }); }); - it('passes the refunded status and pre-submit balance to the sweep on refunded', async () => { + it('sweeps as a refund with the pre-submit balance on refunded', async () => { const { submitPolymarketWithdraw, sweepPolymarketDepositWallet } = getPolymarketMocks(); submitPolymarketWithdraw.mockResolvedValue({ @@ -1880,11 +1886,12 @@ describe('Relay Submit Utils', () => { await expect(submitRelayQuotes(request)).rejects.toThrow( 'Relay: Request failed with status: refunded', ); - expect(sweepPolymarketDepositWallet).toHaveBeenCalledWith( - FROM_MOCK, + expect(sweepPolymarketDepositWallet).toHaveBeenCalledWith({ + from: FROM_MOCK, + isRefund: true, messenger, - { relayStatus: 'refunded', preSubmitUsdceBalance: 2500000n }, - ); + preSubmitUsdceBalance: 2500000n, + }); }); it('returns timeout (tolerated) when Relay polling times out', async () => { @@ -1898,11 +1905,12 @@ describe('Relay Submit Utils', () => { await expect(submitRelayQuotes(request)).rejects.toThrow( 'Relay: Request failed with status: timeout', ); - expect(sweepPolymarketDepositWallet).toHaveBeenCalledWith( - FROM_MOCK, + expect(sweepPolymarketDepositWallet).toHaveBeenCalledWith({ + from: FROM_MOCK, + isRefund: false, messenger, - { relayStatus: 'timeout', preSubmitUsdceBalance: 0n }, - ); + preSubmitUsdceBalance: 0n, + }); }); }); diff --git a/packages/transaction-pay-controller/src/strategy/relay/relay-submit.ts b/packages/transaction-pay-controller/src/strategy/relay/relay-submit.ts index 82ff92c1fcd..1a0e4b31d55 100644 --- a/packages/transaction-pay-controller/src/strategy/relay/relay-submit.ts +++ b/packages/transaction-pay-controller/src/strategy/relay/relay-submit.ts @@ -31,8 +31,8 @@ import { resolveGasPayment, } from '../../utils/gas-payment.js'; import { - sweepPolymarketDepositWallet, submitPolymarketWithdraw, + sweepPolymarketDepositWallet, } from '../../utils/polymarket/withdraw.js'; import { getNetworkClientId } from '../../utils/provider.js'; import { submitSecondLeg } from '../../utils/second-leg.js'; @@ -53,6 +53,7 @@ import { RELAY_FAILURE_STATUSES, RELAY_PENDING_STATUSES, } from './constants.js'; +import { getRelayDepositData } from './deposit-step.js'; import { submitHyperliquidWithdraw } from './hyperliquid-withdraw.js'; import { getRelayStatus } from './relay-api.js'; import { submitViaRelayExecute } from './relay-submit-execute.js'; @@ -174,7 +175,12 @@ async function executeSingleQuote( await submitHyperliquidWithdraw(quote, quote.request.from, messenger); } else if (isPolymarket) { const { sourceHash, preSubmitUsdceBalance } = - await submitPolymarketWithdraw(quote, quote.request.from, messenger); + await submitPolymarketWithdraw({ + depositData: getRelayDepositData(quote.original), + from: quote.request.from, + messenger, + sourceAmountRaw: quote.sourceAmount.raw, + }); polymarketPreSubmitUsdceBalance = preSubmitUsdceBalance; setRelaySourceHash(transaction, messenger, sourceHash); } else { @@ -196,8 +202,11 @@ async function executeSingleQuote( log('Relay request completed', completion); if (isPolymarket) { - await sweepPolymarketDepositWallet(quote.request.from, messenger, { - relayStatus: completion.status, + await sweepPolymarketDepositWallet({ + from: quote.request.from, + isRefund: + completion.status === 'refund' || completion.status === 'refunded', + messenger, preSubmitUsdceBalance: polymarketPreSubmitUsdceBalance, }); diff --git a/packages/transaction-pay-controller/src/strategy/relay/relay-validation.ts b/packages/transaction-pay-controller/src/strategy/relay/relay-validation.ts index d3ed5365daf..d753830dc03 100644 --- a/packages/transaction-pay-controller/src/strategy/relay/relay-validation.ts +++ b/packages/transaction-pay-controller/src/strategy/relay/relay-validation.ts @@ -27,6 +27,7 @@ import { isQuoteError, } from '../../utils/validation.js'; import type { QuoteSimulation } from '../../utils/validation.js'; +import { getRelayDepositData } from './deposit-step.js'; import { getRelayExecuteRequest } from './relay-submit-execute.js'; import { getRelaySubmitCalls } from './relay-submit.js'; import type { RelayExecuteRequest, RelayQuote } from './types.js'; @@ -110,11 +111,12 @@ async function buildValidationSimulation( quote: TransactionPayQuote, ): Promise { if (quote.request.isPolymarketDepositWallet) { - return await buildPolymarketDepositWalletSimulation( - quote, - quote.request.from, - request.messenger, - ); + return await buildPolymarketDepositWalletSimulation({ + depositData: getRelayDepositData(quote.original), + from: quote.request.from, + messenger: request.messenger, + sourceAmountRaw: quote.sourceAmount.raw, + }); } const { calls } = await getRelaySubmitCalls({ diff --git a/packages/transaction-pay-controller/src/utils/polymarket/withdraw.test.ts b/packages/transaction-pay-controller/src/utils/polymarket/withdraw.test.ts index c5c3ffcdd38..da294e317c0 100644 --- a/packages/transaction-pay-controller/src/utils/polymarket/withdraw.test.ts +++ b/packages/transaction-pay-controller/src/utils/polymarket/withdraw.test.ts @@ -7,65 +7,40 @@ import { POLYGON_PUSD_ADDRESS, POLYGON_USDCE_ADDRESS, } from '../../constants.js'; -import type { - RelayQuote, - RelayQuoteRequest, -} from '../../strategy/relay/types.js'; import { getMessengerMock } from '../../tests/messenger-mock.js'; -import type { QuoteRequest, TransactionPayQuote } from '../../types.js'; +import type { QuoteRequest } from '../../types.js'; import { getLiveTokenBalance } from '../token.js'; import { POLYMARKET_COLLATERAL_OFFRAMP_POLYGON, POLYMARKET_COLLATERAL_ONRAMP_POLYGON, } from './constants.js'; import { - applyPolymarketDepositWalletOverrides, buildPolymarketDepositWalletSimulation, + getPolymarketDepositWalletOverrides, + getPredictWithdrawFeeTokenAccount, getPredictWithdrawSafeAddress, isPredictWithdraw, submitPolymarketWithdraw, sweepPolymarketDepositWallet, } from './withdraw.js'; +import type { DepositWalletWithdrawRequest } from './withdraw.js'; jest.mock('../token'); -const EOA_MOCK = '0x1111111111111111111111111111111111111111' as Hex; +const DEPOSIT_ADDRESS_MOCK = + '0x1234567890123456789012345678901234567890' as Hex; const DEPOSIT_WALLET_MOCK = '0x2222222222222222222222222222222222222222' as Hex; -const SOURCE_HASH_MOCK: Hex = `0x${'aa'.repeat(32)}`; +const EOA_MOCK = '0x1111111111111111111111111111111111111111' as Hex; const SOURCE_AMOUNT_RAW_MOCK = '1000000'; +const SOURCE_HASH_MOCK: Hex = `0x${'aa'.repeat(32)}`; -// transfer(0x1234...7890, 0) encoded calldata +// transfer(DEPOSIT_ADDRESS_MOCK, 1000000000) const TRANSFER_CALLDATA_MOCK = '0xa9059cbb0000000000000000000000001234567890123456789012345678901234567890000000000000000000000000000000000000000000000000000000003b9aca00' as Hex; -function buildQuote( - overrides: Partial = {}, -): TransactionPayQuote { - return { - original: { - steps: [ - { - id: 'deposit', - kind: 'transaction', - items: [ - { - data: { - data: TRANSFER_CALLDATA_MOCK, - }, - }, - ], - }, - ], - ...overrides, - }, - sourceAmount: { - raw: SOURCE_AMOUNT_RAW_MOCK, - human: '1', - fiat: '1', - usd: '1', - }, - } as TransactionPayQuote; -} +const WITHDRAW_TRANSACTION_MOCK = { + type: TransactionType.predictWithdraw, +} as TransactionMeta; describe('Polymarket withdraw', () => { const { @@ -75,6 +50,18 @@ describe('Polymarket withdraw', () => { } = getMessengerMock(); const getLiveTokenBalanceMock = jest.mocked(getLiveTokenBalance); + function buildRequest( + overrides: Partial = {}, + ): DepositWalletWithdrawRequest { + return { + depositData: TRANSFER_CALLDATA_MOCK, + from: EOA_MOCK, + messenger, + sourceAmountRaw: SOURCE_AMOUNT_RAW_MOCK, + ...overrides, + }; + } + beforeEach(() => { jest.resetAllMocks(); polymarketGetDepositWalletAddressMock.mockResolvedValue( @@ -86,31 +73,26 @@ describe('Polymarket withdraw', () => { getLiveTokenBalanceMock.mockResolvedValue('0'); }); - describe('applyPolymarketDepositWalletOverrides', () => { - it('rewrites the quote body for the deposit-wallet path', async () => { - const body = {} as RelayQuoteRequest; - const request = { from: EOA_MOCK } as QuoteRequest; - - await applyPolymarketDepositWalletOverrides(body, request, messenger); + describe('getPolymarketDepositWalletOverrides', () => { + it('returns the deposit wallet and USDC.e as the source token', async () => { + const result = await getPolymarketDepositWalletOverrides( + EOA_MOCK, + messenger, + ); expect(polymarketGetDepositWalletAddressMock).toHaveBeenCalledWith({ eoa: EOA_MOCK, }); - expect(body).toStrictEqual({ - originCurrency: POLYGON_USDCE_ADDRESS, - user: DEPOSIT_WALLET_MOCK, - refundTo: DEPOSIT_WALLET_MOCK, - useDepositAddress: true, - strict: true, + expect(result).toStrictEqual({ + depositWallet: DEPOSIT_WALLET_MOCK, + sourceTokenAddress: POLYGON_USDCE_ADDRESS, }); }); }); describe('submitPolymarketWithdraw', () => { it('submits the approve + unwrap batch via the relayer callback', async () => { - const quote = buildQuote(); - - const result = await submitPolymarketWithdraw(quote, EOA_MOCK, messenger); + const result = await submitPolymarketWithdraw(buildRequest()); expect(result).toStrictEqual({ sourceHash: SOURCE_HASH_MOCK, @@ -130,11 +112,7 @@ describe('Polymarket withdraw', () => { it('captures the pre-submit USDC.e balance', async () => { getLiveTokenBalanceMock.mockResolvedValue('2500000'); - const result = await submitPolymarketWithdraw( - buildQuote(), - EOA_MOCK, - messenger, - ); + const result = await submitPolymarketWithdraw(buildRequest()); expect(result.preSubmitUsdceBalance).toBe(2500000n); }); @@ -142,50 +120,32 @@ describe('Polymarket withdraw', () => { it('defaults pre-submit balance to zero when the balance read fails', async () => { getLiveTokenBalanceMock.mockRejectedValue(new Error('rpc down')); - const result = await submitPolymarketWithdraw( - buildQuote(), - EOA_MOCK, - messenger, - ); + const result = await submitPolymarketWithdraw(buildRequest()); expect(result.preSubmitUsdceBalance).toBe(0n); }); - it('throws when the Relay quote has no deposit step', async () => { - const quote = buildQuote({ steps: [] }); - + it('throws without submitting when there is no deposit calldata', async () => { await expect( - submitPolymarketWithdraw(quote, EOA_MOCK, messenger), - ).rejects.toThrow('Relay quote has no deposit step'); - }); + submitPolymarketWithdraw(buildRequest({ depositData: undefined })), + ).rejects.toThrow('quote has no deposit step calldata'); - it('throws when the Relay quote deposit step is missing calldata', async () => { - const quote = buildQuote({ - steps: [ - { - id: 'deposit', - kind: 'transaction', - items: [{ data: {} }], - }, - ], - } as unknown as Partial); - - await expect( - submitPolymarketWithdraw(quote, EOA_MOCK, messenger), - ).rejects.toThrow('deposit step is missing calldata'); + expect(polymarketSubmitDepositWalletBatchMock).not.toHaveBeenCalled(); }); }); describe('sweepPolymarketDepositWallet', () => { const successOptions = { - relayStatus: 'success' as const, + from: EOA_MOCK, + isRefund: false, + messenger, preSubmitUsdceBalance: 0n, }; it('wraps any USDC.e balance back into pUSD on the deposit wallet', async () => { getLiveTokenBalanceMock.mockResolvedValue('5000000'); - await sweepPolymarketDepositWallet(EOA_MOCK, messenger, successOptions); + await sweepPolymarketDepositWallet(successOptions); expect(polymarketSubmitDepositWalletBatchMock).toHaveBeenCalledTimes(1); const call = polymarketSubmitDepositWalletBatchMock.mock.calls[0][0]; @@ -199,7 +159,7 @@ describe('Polymarket withdraw', () => { it('is a no-op when the USDC.e balance is zero', async () => { getLiveTokenBalanceMock.mockResolvedValue('0'); - await sweepPolymarketDepositWallet(EOA_MOCK, messenger, successOptions); + await sweepPolymarketDepositWallet(successOptions); expect(polymarketSubmitDepositWalletBatchMock).not.toHaveBeenCalled(); }); @@ -208,7 +168,7 @@ describe('Polymarket withdraw', () => { getLiveTokenBalanceMock.mockRejectedValue(new Error('rpc down')); expect( - await sweepPolymarketDepositWallet(EOA_MOCK, messenger, successOptions), + await sweepPolymarketDepositWallet(successOptions), ).toBeUndefined(); expect(polymarketSubmitDepositWalletBatchMock).not.toHaveBeenCalled(); }); @@ -220,11 +180,17 @@ describe('Polymarket withdraw', () => { ); expect( - await sweepPolymarketDepositWallet(EOA_MOCK, messenger, successOptions), + await sweepPolymarketDepositWallet(successOptions), ).toBeUndefined(); }); - describe('when relayStatus is refund', () => { + describe('when the withdraw was refunded', () => { + const refundOptions = { + ...successOptions, + isRefund: true, + preSubmitUsdceBalance: 1000000n, + }; + beforeEach(() => { jest.useFakeTimers(); }); @@ -239,10 +205,7 @@ describe('Polymarket withdraw', () => { .mockResolvedValueOnce('1000000') .mockResolvedValueOnce('4000000'); - const sweepPromise = sweepPolymarketDepositWallet(EOA_MOCK, messenger, { - relayStatus: 'refund', - preSubmitUsdceBalance: 1000000n, - }); + const sweepPromise = sweepPolymarketDepositWallet(refundOptions); await jest.advanceTimersByTimeAsync(5000); await sweepPromise; @@ -253,30 +216,10 @@ describe('Polymarket withdraw', () => { expect(call.calls[1].target).toBe(POLYMARKET_COLLATERAL_ONRAMP_POLYGON); }); - it('also retries when relayStatus is refunded', async () => { - getLiveTokenBalanceMock - .mockResolvedValueOnce('1000000') - .mockResolvedValueOnce('4000000'); - - const sweepPromise = sweepPolymarketDepositWallet(EOA_MOCK, messenger, { - relayStatus: 'refunded', - preSubmitUsdceBalance: 1000000n, - }); - - await jest.advanceTimersByTimeAsync(4000); - await sweepPromise; - - expect(getLiveTokenBalanceMock).toHaveBeenCalledTimes(2); - expect(polymarketSubmitDepositWalletBatchMock).toHaveBeenCalledTimes(1); - }); - it('gives up after five attempts and sweeps the residual stale balance', async () => { getLiveTokenBalanceMock.mockResolvedValue('1000000'); - const sweepPromise = sweepPolymarketDepositWallet(EOA_MOCK, messenger, { - relayStatus: 'refund', - preSubmitUsdceBalance: 1000000n, - }); + const sweepPromise = sweepPolymarketDepositWallet(refundOptions); await jest.advanceTimersByTimeAsync(4000); await sweepPromise; @@ -288,23 +231,19 @@ describe('Polymarket withdraw', () => { }); describe('isPredictWithdraw', () => { - const withdrawTransaction = { - type: TransactionType.predictWithdraw, - } as TransactionMeta; - it('returns true for a post-quote predict withdraw', () => { expect( isPredictWithdraw( { isPostQuote: true } as QuoteRequest, - withdrawTransaction, + WITHDRAW_TRANSACTION_MOCK, ), ).toBe(true); }); it('returns false when not post-quote', () => { - expect(isPredictWithdraw({} as QuoteRequest, withdrawTransaction)).toBe( - false, - ); + expect( + isPredictWithdraw({} as QuoteRequest, WITHDRAW_TRANSACTION_MOCK), + ).toBe(false); }); it('returns false when the transaction is not a predict withdraw', () => { @@ -320,22 +259,17 @@ describe('Polymarket withdraw', () => { }); describe('getPredictWithdrawSafeAddress', () => { - const withdrawTransaction = { - type: TransactionType.predictWithdraw, - } as TransactionMeta; - const depositSteps = [ - { id: 'deposit', kind: 'transaction', items: [] }, - ] as unknown as RelayQuote['steps']; - const swapSteps = [ - { id: 'swap', kind: 'transaction', items: [] }, - ] as unknown as RelayQuote['steps']; + const safeRequest = { + isPostQuote: true, + refundTo: DEPOSIT_WALLET_MOCK, + } as QuoteRequest; it('returns refundTo for a deposit-style predict withdraw', () => { expect( getPredictWithdrawSafeAddress( - { isPostQuote: true, refundTo: DEPOSIT_WALLET_MOCK } as QuoteRequest, - depositSteps, - withdrawTransaction, + safeRequest, + WITHDRAW_TRANSACTION_MOCK, + true, ), ).toBe(DEPOSIT_WALLET_MOCK); }); @@ -343,13 +277,9 @@ describe('Polymarket withdraw', () => { it('returns undefined for the deposit-wallet variant (handled by its own simulation)', () => { expect( getPredictWithdrawSafeAddress( - { - isPostQuote: true, - refundTo: DEPOSIT_WALLET_MOCK, - isPolymarketDepositWallet: true, - } as QuoteRequest, - depositSteps, - withdrawTransaction, + { ...safeRequest, isPolymarketDepositWallet: true }, + WITHDRAW_TRANSACTION_MOCK, + true, ), ).toBeUndefined(); }); @@ -357,9 +287,9 @@ describe('Polymarket withdraw', () => { it('returns undefined for swap-only routes (no deposit step)', () => { expect( getPredictWithdrawSafeAddress( - { isPostQuote: true, refundTo: DEPOSIT_WALLET_MOCK } as QuoteRequest, - swapSteps, - withdrawTransaction, + safeRequest, + WITHDRAW_TRANSACTION_MOCK, + false, ), ).toBeUndefined(); }); @@ -367,24 +297,38 @@ describe('Polymarket withdraw', () => { it('returns undefined when not a predict withdraw', () => { expect( getPredictWithdrawSafeAddress( - { isPostQuote: true, refundTo: DEPOSIT_WALLET_MOCK } as QuoteRequest, - depositSteps, + safeRequest, { type: TransactionType.simpleSend } as TransactionMeta, + true, ), ).toBeUndefined(); }); }); - describe('buildPolymarketDepositWalletSimulation', () => { - const RELAY_DEPOSIT_ADDRESS_MOCK = - '0x1234567890123456789012345678901234567890' as Hex; + describe('getPredictWithdrawFeeTokenAccount', () => { + it('returns refundTo for a predict withdraw', () => { + expect( + getPredictWithdrawFeeTokenAccount( + { isPostQuote: true, refundTo: DEPOSIT_WALLET_MOCK } as QuoteRequest, + WITHDRAW_TRANSACTION_MOCK, + ), + ).toBe(DEPOSIT_WALLET_MOCK); + }); + it('returns undefined when not a predict withdraw', () => { + expect( + getPredictWithdrawFeeTokenAccount( + { refundTo: DEPOSIT_WALLET_MOCK } as QuoteRequest, + WITHDRAW_TRANSACTION_MOCK, + ), + ).toBeUndefined(); + }); + }); + + describe('buildPolymarketDepositWalletSimulation', () => { it('simulates the real approve + unwrap batch from the deposit wallet', async () => { - const simulation = await buildPolymarketDepositWalletSimulation( - buildQuote(), - EOA_MOCK, - messenger, - ); + const simulation = + await buildPolymarketDepositWalletSimulation(buildRequest()); expect(polymarketGetDepositWalletAddressMock).toHaveBeenCalledWith({ eoa: EOA_MOCK, @@ -402,53 +346,40 @@ describe('Polymarket withdraw', () => { expect(unwrap.value).toBe('0x0'); }); - it('encodes the approve for the offramp and the amount from the quote', async () => { - const simulation = await buildPolymarketDepositWalletSimulation( - buildQuote(), - EOA_MOCK, - messenger, - ); + it('approves the offramp and unwraps the source amount to the deposit address', async () => { + const simulation = + await buildPolymarketDepositWalletSimulation(buildRequest()); const [approve, unwrap] = simulation.transactions; - // approve(offramp, 1000000) - const approveInterface = new Interface([ + const decodedApprove = new Interface([ 'function approve(address spender, uint256 amount)', - ]); - const decodedApprove = approveInterface.decodeFunctionData( - 'approve', - approve.data as Hex, - ); + ]).decodeFunctionData('approve', approve.data as Hex); + expect(decodedApprove[0].toLowerCase()).toBe( POLYMARKET_COLLATERAL_OFFRAMP_POLYGON.toLowerCase(), ); expect(decodedApprove[1].toString()).toBe(SOURCE_AMOUNT_RAW_MOCK); - // unwrap(USDC.e, relayDepositAddress, 1000000) - const unwrapInterface = new Interface([ + const decodedUnwrap = new Interface([ 'function unwrap(address asset, address recipient, uint256 amount)', - ]); - const decodedUnwrap = unwrapInterface.decodeFunctionData( - 'unwrap', - unwrap.data as Hex, - ); + ]).decodeFunctionData('unwrap', unwrap.data as Hex); + expect(decodedUnwrap[0].toLowerCase()).toBe( POLYGON_USDCE_ADDRESS.toLowerCase(), ); expect(decodedUnwrap[1].toLowerCase()).toBe( - RELAY_DEPOSIT_ADDRESS_MOCK.toLowerCase(), + DEPOSIT_ADDRESS_MOCK.toLowerCase(), ); expect(decodedUnwrap[2].toString()).toBe(SOURCE_AMOUNT_RAW_MOCK); }); - it('throws when the Relay quote has no deposit step', async () => { + it('throws when there is no deposit calldata', async () => { await expect( buildPolymarketDepositWalletSimulation( - buildQuote({ steps: [] }), - EOA_MOCK, - messenger, + buildRequest({ depositData: undefined }), ), - ).rejects.toThrow('Relay quote has no deposit step'); + ).rejects.toThrow('quote has no deposit step calldata'); }); }); }); diff --git a/packages/transaction-pay-controller/src/utils/polymarket/withdraw.ts b/packages/transaction-pay-controller/src/utils/polymarket/withdraw.ts index eb075484211..ddac9ec9982 100644 --- a/packages/transaction-pay-controller/src/utils/polymarket/withdraw.ts +++ b/packages/transaction-pay-controller/src/utils/polymarket/withdraw.ts @@ -12,16 +12,9 @@ import { POLYGON_USDCE_ADDRESS, } from '../../constants.js'; import { projectLogger } from '../../logger.js'; -import type { - RelayQuote, - RelayQuoteRequest, - RelayStatus, - RelayTransactionStep, -} from '../../strategy/relay/types.js'; import type { QuoteRequest, TransactionPayControllerMessenger, - TransactionPayQuote, } from '../../types.js'; import { getLiveTokenBalance } from '../token.js'; import type { QuoteSimulation } from '../validation.js'; @@ -41,7 +34,19 @@ import { const log = createModuleLogger(projectLogger, 'polymarket-withdraw'); -type RelayQuoteStep = RelayQuote['steps'][number]; +/** + * A deposit-wallet Predict withdraw, funded by unwrapping pUSD from the + * deposit wallet straight to the provider's deposit address. + */ +export type DepositWalletWithdrawRequest = { + /** Calldata of the provider deposit step, an ERC-20 transfer to the deposit address. */ + depositData: Hex | undefined; + from: Hex; + messenger: TransactionPayControllerMessenger; + sourceAmountRaw: string; +}; + +type DepositWalletCall = { target: Hex; data: Hex }; /** * Whether a quote request + transaction represent a Predict withdraw post-quote @@ -65,34 +70,26 @@ export function isPredictWithdraw( * Resolve the Polymarket Safe proxy address that must act as `msg.sender` when * estimating gas for the withdraw leg of a Safe-based Predict withdraw. * - * Only applies to the Safe variant on deposit-style Relay routes: the source - * token lives in the Safe (carried as `request.refundTo`), and swap-only routes - * keep the EOA `from` because DEX aggregators reject contract callers. + * Only applies to the Safe variant on deposit-style routes: the source token + * lives in the Safe (carried as `request.refundTo`), and swap-only routes keep + * the EOA `from` because DEX aggregators reject contract callers. * * @param request - Quote request. - * @param steps - Relay quote steps (used to detect a deposit-style route). * @param transaction - Original transaction metadata. + * @param hasDepositStep - Whether the quote routes through a deposit step. * @returns The Safe proxy address, or `undefined` when the default EOA `from` * should be used. */ export function getPredictWithdrawSafeAddress( request: QuoteRequest, - steps: RelayQuoteStep[], transaction: TransactionMeta, + hasDepositStep: boolean, ): Hex | undefined { - if (!isPredictWithdraw(request, transaction)) { - return undefined; - } - - // Deposit-wallet withdraws are validated via their own simulation builder; - // `request.refundTo` is not the deposit wallet for that variant. - if (request.isPolymarketDepositWallet) { - return undefined; - } - - const hasDepositStep = steps.some((step) => step.id === 'deposit'); - - if (!hasDepositStep) { + if ( + !hasDepositStep || + request.isPolymarketDepositWallet || + !isPredictWithdraw(request, transaction) + ) { return undefined; } @@ -100,99 +97,124 @@ export function getPredictWithdrawSafeAddress( } /** - * Build the simulation for a Polymarket deposit-wallet Predict withdraw, - * mirroring the approve + unwrap batch broadcast by {@link submitPolymarketWithdraw} - * so validation checks the deposit wallet holds enough pUSD. + * Resolve the account holding the source token for gas-fee-token lookups. + * + * Unlike {@link getPredictWithdrawSafeAddress} this applies to every Predict + * withdraw route: the source token lives in the Safe and the EOA is empty until + * the Safe call runs mid-batch. + * + * @param request - Quote request. + * @param transaction - Original transaction metadata. + * @returns The Safe proxy address, or `undefined` for other flows. + */ +export function getPredictWithdrawFeeTokenAccount( + request: QuoteRequest, + transaction: TransactionMeta, +): Hex | undefined { + return isPredictWithdraw(request, transaction) ? request.refundTo : undefined; +} + +/** + * Resolve the quote overrides for a deposit-wallet Predict withdraw. The deposit + * wallet is both the sender and the refund address, and the provider is funded + * with the USDC.e produced by unwrapping pUSD. * - * @param quote - Relay quote (source amount + deposit step). * @param from - The user EOA that owns the deposit wallet. * @param messenger - Controller messenger. - * @returns Simulation for the approve + unwrap batch from the deposit wallet. + * @returns The deposit wallet and the source token to quote. */ -export async function buildPolymarketDepositWalletSimulation( - quote: TransactionPayQuote, +export async function getPolymarketDepositWalletOverrides( from: Hex, messenger: TransactionPayControllerMessenger, -): Promise { - const depositWalletAddress = await getDepositWalletAddress(messenger, from); - const relayDepositAddress = extractRelayDepositAddress(quote.original); - const amount = BigInt(quote.sourceAmount.raw); - +): Promise<{ depositWallet: Hex; sourceTokenAddress: Hex }> { return { - transactions: buildDepositWalletUnwrapCalls( - relayDepositAddress, - amount, - ).map((call) => ({ - from: depositWalletAddress, - to: call.target, - data: call.data, - value: '0x0', - })), + depositWallet: await getDepositWalletAddress(messenger, from), + sourceTokenAddress: POLYGON_USDCE_ADDRESS, }; } -export async function applyPolymarketDepositWalletOverrides( - body: Omit, - request: QuoteRequest, - messenger: TransactionPayControllerMessenger, -): Promise { - const depositWalletAddress = await getDepositWalletAddress( - messenger, +/** + * Build the simulation for a deposit-wallet Predict withdraw, mirroring the + * batch broadcast by {@link submitPolymarketWithdraw}. + * + * @param request - Withdraw request. + * @returns Simulation for the approve + unwrap batch from the deposit wallet. + */ +export async function buildPolymarketDepositWalletSimulation( + request: DepositWalletWithdrawRequest, +): Promise { + const depositWallet = await getDepositWalletAddress( + request.messenger, request.from, ); - body.originCurrency = POLYGON_USDCE_ADDRESS; - body.user = depositWalletAddress; - body.refundTo = depositWalletAddress; - body.useDepositAddress = true; - body.strict = true; + return { + transactions: buildUnwrapCalls(request).map((call) => ({ + data: call.data, + from: depositWallet, + to: call.target, + value: '0x0', + })), + }; } +/** + * Unwrap the deposit wallet's pUSD straight to the provider deposit address. + * + * @param request - Withdraw request. + * @returns The batch hash, and the USDC.e balance before submission so a + * later sweep can detect a refund. + */ export async function submitPolymarketWithdraw( - quote: TransactionPayQuote, - from: Hex, - messenger: TransactionPayControllerMessenger, + request: DepositWalletWithdrawRequest, ): Promise<{ sourceHash: Hex; preSubmitUsdceBalance: bigint }> { - const depositWalletAddress = await getDepositWalletAddress(messenger, from); - const relayDepositAddress = extractRelayDepositAddress(quote.original); - const amount = BigInt(quote.sourceAmount.raw); + const { from, messenger } = request; + const depositWallet = await getDepositWalletAddress(messenger, from); + const calls = buildUnwrapCalls(request); const preSubmitUsdceBalance = await readUsdceBalanceOrZero( messenger, - depositWalletAddress, + depositWallet, ); - log('Submitting unwrap batch to Relay deposit address', { - depositWalletAddress, - relayDepositAddress, - amount: amount.toString(), + log('Submitting unwrap batch to deposit address', { + depositWallet, preSubmitUsdceBalance: preSubmitUsdceBalance.toString(), + sourceAmountRaw: request.sourceAmountRaw, }); const result = await submitDepositWalletBatch(messenger, { + calls: calls.map((call) => ({ ...call, value: '0' })), + depositWallet, eoa: from, - depositWallet: depositWalletAddress, - calls: buildDepositWalletUnwrapCalls(relayDepositAddress, amount).map( - (call) => ({ ...call, value: '0' }), - ), }); return { ...result, preSubmitUsdceBalance }; } -export async function sweepPolymarketDepositWallet( - from: Hex, - messenger: TransactionPayControllerMessenger, - options: { - relayStatus: RelayStatus | 'timeout'; - preSubmitUsdceBalance: bigint; - }, -): Promise { - const isRefund = - options.relayStatus === 'refund' || options.relayStatus === 'refunded'; - const waitForBalanceAbove = isRefund - ? options.preSubmitUsdceBalance - : undefined; +/** + * Wrap any USDC.e left on the deposit wallet back into pUSD, such as a refund + * from a failed withdraw. + * + * @param options - Sweep options. + * @param options.from - The user EOA that owns the deposit wallet. + * @param options.isRefund - Whether the provider refunded the withdraw, in which + * case the sweep waits for the refunded balance to land. + * @param options.messenger - Controller messenger. + * @param options.preSubmitUsdceBalance - USDC.e balance before the withdraw. + */ +export async function sweepPolymarketDepositWallet({ + from, + isRefund, + messenger, + preSubmitUsdceBalance, +}: { + from: Hex; + isRefund: boolean; + messenger: TransactionPayControllerMessenger; + preSubmitUsdceBalance: bigint; +}): Promise { + const waitForBalanceAbove = isRefund ? preSubmitUsdceBalance : undefined; const depositWalletAddress = await getDepositWalletAddress(messenger, from); const usdceBalance = await readDepositWalletUsdceBalance( @@ -349,41 +371,29 @@ async function submitDepositWalletBatch( return result; } -function extractRelayDepositAddress(relayQuote: RelayQuote): Hex { - const depositStep = relayQuote.steps.find((step) => step.id === 'deposit'); - - if (depositStep?.kind !== 'transaction') { - throw new Error( - 'Polymarket deposit wallet withdraw: Relay quote has no deposit step', - ); - } - - const transactionStep = depositStep as RelayTransactionStep; - const depositCallData = transactionStep.items[0]?.data?.data; - - if (!depositCallData) { - throw new Error( - 'Polymarket deposit wallet withdraw: Relay quote deposit step is missing calldata', - ); - } - - return extractErc20TransferRecipient(depositCallData); -} - /** * Build the approve + unwrap batch that moves a deposit wallet's pUSD to the - * Relay deposit address as USDC.e. Shared by the simulation and the real submit - * so both model the exact same calls. + * provider deposit address as USDC.e. Shared by the simulation and the real + * submit so both model the exact same calls. * - * @param relayDepositAddress - The Relay deposit address that receives the - * unwrapped USDC.e. - * @param amount - The pUSD amount to unwrap. + * @param request - Withdraw request. + * @param request.depositData - Calldata of the provider deposit step. + * @param request.sourceAmountRaw - Amount of pUSD to unwrap. * @returns The two `{ target, data }` calls. */ -function buildDepositWalletUnwrapCalls( - relayDepositAddress: Hex, - amount: bigint, -): { target: Hex; data: Hex }[] { +function buildUnwrapCalls({ + depositData, + sourceAmountRaw, +}: DepositWalletWithdrawRequest): DepositWalletCall[] { + if (!depositData) { + throw new Error( + 'Polymarket deposit wallet withdraw: quote has no deposit step calldata', + ); + } + + const depositAddress = extractErc20TransferRecipient(depositData); + const amount = BigInt(sourceAmountRaw); + return [ { target: POLYGON_PUSD_ADDRESS, @@ -393,7 +403,7 @@ function buildDepositWalletUnwrapCalls( target: POLYMARKET_COLLATERAL_OFFRAMP_POLYGON, data: encodeUnwrap({ asset: POLYGON_USDCE_ADDRESS, - recipient: relayDepositAddress, + recipient: depositAddress, amount, }), },