refactor(transaction-pay-controller): decouple Polymarket withdraw helpers from the Relay strategy

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
@@ -6702,7 +6688,7 @@
       "count": 1
     },
     "typescript/no-unsafe-assignment": {
-      "count": 34
+      "count": 39
     },
     "typescript/no-unsafe-call": {
       "count": 17
@@ -6818,6 +6804,20 @@
       "count": 1
     }
   },
+  "packages/transaction-pay-controller/src/utils/polymarket/withdraw.test.ts": {
+    "typescript/no-unsafe-argument": {
+      "count": 1
+    },
+    "typescript/no-unsafe-assignment": {
+      "count": 6
+    },
+    "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

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();
+    });
+  });
+});

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;
+}

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 {
+  getPolymarketDepositWalletOverrides,
+  getPredictWithdrawFeeTokenAccount,
+  getPredictWithdrawSafeAddress,
+} 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';
@@ -59,12 +64,8 @@ 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 {
-  applyPolymarketDepositWalletOverrides,
-  getPredictWithdrawSafeAddress,
-  isPredictWithdraw,
-} from './polymarket/withdraw.js';
 import { fetchRelayQuote } from './relay-api.js';
 import {
   getRelayMaxQuote,
@@ -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,

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<typeof import('../../utils/second-leg.js')>(
@@ -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:
@@ -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,
+        });
       });
     });
 

packages/transaction-pay-controller/src/strategy/relay/relay-submit.ts

@@ -30,6 +30,10 @@ import {
   logGasPaymentOutcome,
   resolveGasPayment,
 } from '../../utils/gas-payment.js';
+import {
+  submitPolymarketWithdraw,
+  sweepPolymarketDepositWallet,
+} from '../../utils/polymarket/withdraw.js';
 import { getNetworkClientId } from '../../utils/provider.js';
 import { submitSecondLeg } from '../../utils/second-leg.js';
 import {
@@ -49,11 +53,8 @@ import {
   RELAY_FAILURE_STATUSES,
   RELAY_PENDING_STATUSES,
 } from './constants.js';
+import { getRelayDepositData } from './deposit-step.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 {
@@ -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,
     });