refactor(predict): read deposit-wallet state via usePredictAccountState

matthewwalsh0 · matthewwalsh0 · commit d9d0d2944a61 · 2026-05-14T23:38:13.000+01:00
diff --git a/app/components/Views/confirmations/hooks/pay/useTransactionPayPostQuote.test.ts b/app/components/Views/confirmations/hooks/pay/useTransactionPayPostQuote.test.ts
@@ -6,6 +6,7 @@ import { useTransactionMetadataRequest } from '../transactions/useTransactionMet
 import { useTransactionPayWithdraw } from './useTransactionPayWithdraw';
 import Engine from '../../../../../core/Engine';
 import { computeProxyAddress } from '../../../../UI/Predict/providers/polymarket/safe/utils';
+import { usePredictAccountState } from '../../../../UI/Predict/hooks/usePredictAccountState';
 
 jest.mock('../transactions/useTransactionMetadataRequest');
 jest.mock('./useTransactionPayWithdraw');
@@ -14,14 +15,12 @@ jest.mock('../../../../../core/Engine', () => ({
     TransactionPayController: {
       setTransactionConfig: jest.fn(),
     },
-    PredictController: {
-      getAccountState: jest.fn(),
-    },
   },
 }));
 jest.mock('../../../../UI/Predict/providers/polymarket/safe/utils', () => ({
   computeProxyAddress: jest.fn(),
 }));
+jest.mock('../../../../UI/Predict/hooks/usePredictAccountState');
 
 const TRANSACTION_ID_MOCK = 'transaction-123';
 const FROM_MOCK = '0x1234567890123456789012345678901234567890' as Hex;
@@ -35,22 +34,24 @@ describe('useTransactionPayPostQuote', () => {
   const setTransactionConfigMock = jest.mocked(
     Engine.context.TransactionPayController.setTransactionConfig,
   );
-  const getAccountStateMock = jest.mocked(
-    Engine.context.PredictController.getAccountState,
-  );
+  const usePredictAccountStateMock = jest.mocked(usePredictAccountState);
   const computeProxyAddressMock = jest.mocked(computeProxyAddress);
 
-  const flushPromises = () => new Promise(setImmediate);
+  function mockAccountState(walletType: 'safe' | 'deposit-wallet'): void {
+    usePredictAccountStateMock.mockReturnValue({
+      data: {
+        address: '0xProxyAddress',
+        isDeployed: true,
+        walletType,
+      },
+      isLoading: false,
+    } as never);
+  }
 
   beforeEach(() => {
     jest.clearAllMocks();
     computeProxyAddressMock.mockReturnValue(PROXY_ADDRESS_MOCK);
-    getAccountStateMock.mockResolvedValue({
-      address: '0xProxyAddress',
-      isDeployed: true,
-      walletType: 'safe',
-      balance: 100,
-    } as never);
+    mockAccountState('safe');
     useTransactionPayWithdrawMock.mockReturnValue({
       isWithdraw: false,
       canSelectWithdrawToken: false,
@@ -292,66 +293,68 @@ describe('useTransactionPayPostQuote', () => {
       });
     });
 
-    it('flags transaction and clears refundTo when walletType is deposit-wallet', async () => {
-      getAccountStateMock.mockResolvedValue({
-        address: '0xDepositWallet',
-        isDeployed: true,
-        walletType: 'deposit-wallet',
-        balance: 50,
-      } as never);
+    it('flags transaction and skips refundTo when walletType is deposit-wallet', () => {
+      mockAccountState('deposit-wallet');
 
       renderHook(() => useTransactionPayPostQuote());
-      await flushPromises();
 
-      expect(setTransactionConfigMock).toHaveBeenCalledTimes(2);
+      expect(setTransactionConfigMock).toHaveBeenCalledTimes(1);
 
-      const followUpCallback = setTransactionConfigMock.mock.calls[1][1];
-      const config = {
-        refundTo: PROXY_ADDRESS_MOCK,
-      } as {
+      const callback = setTransactionConfigMock.mock.calls[0][1];
+      const config = {} as {
+        isPostQuote?: boolean;
         isPolymarketDepositWallet?: boolean;
         refundTo?: Hex;
       };
-      followUpCallback(config);
+      callback(config);
 
+      expect(config.isPostQuote).toBe(true);
       expect(config.isPolymarketDepositWallet).toBe(true);
       expect(config.refundTo).toBeUndefined();
+      expect(computeProxyAddressMock).not.toHaveBeenCalled();
     });
 
-    it('does not set deposit-wallet flag when walletType is safe', async () => {
-      getAccountStateMock.mockResolvedValue({
-        address: '0xSafeProxy',
-        isDeployed: true,
-        walletType: 'safe',
-        balance: 100,
-      } as never);
+    it('does not set deposit-wallet flag when walletType is safe', () => {
+      mockAccountState('safe');
 
       renderHook(() => useTransactionPayPostQuote());
-      await flushPromises();
 
       expect(setTransactionConfigMock).toHaveBeenCalledTimes(1);
+
+      const callback = setTransactionConfigMock.mock.calls[0][1];
+      const config = {} as {
+        isPostQuote?: boolean;
+        isPolymarketDepositWallet?: boolean;
+        refundTo?: Hex;
+      };
+      callback(config);
+
+      expect(config.isPolymarketDepositWallet).toBeUndefined();
     });
 
-    it('does not resolve account state for non-predictWithdraw flows', async () => {
-      useTransactionMetadataRequestMock.mockReturnValue({
-        id: TRANSACTION_ID_MOCK,
-        txParams: { from: FROM_MOCK },
-        type: TransactionType.perpsWithdraw,
+    it('defers setTransactionConfig until predict account state resolves', () => {
+      usePredictAccountStateMock.mockReturnValue({
+        data: undefined,
+        isLoading: true,
       } as never);
 
       renderHook(() => useTransactionPayPostQuote());
-      await flushPromises();
 
-      expect(getAccountStateMock).not.toHaveBeenCalled();
-      expect(setTransactionConfigMock).toHaveBeenCalledTimes(1);
+      expect(setTransactionConfigMock).not.toHaveBeenCalled();
     });
 
-    it('swallows getAccountState errors without setting the deposit-wallet flag', async () => {
-      getAccountStateMock.mockRejectedValue(new Error('boom'));
+    it('does not resolve account state for non-predictWithdraw flows', () => {
+      useTransactionMetadataRequestMock.mockReturnValue({
+        id: TRANSACTION_ID_MOCK,
+        txParams: { from: FROM_MOCK },
+        type: TransactionType.perpsWithdraw,
+      } as never);
 
       renderHook(() => useTransactionPayPostQuote());
-      await flushPromises();
 
+      expect(usePredictAccountStateMock).toHaveBeenCalledWith({
+        enabled: false,
+      });
       expect(setTransactionConfigMock).toHaveBeenCalledTimes(1);
     });
   });
diff --git a/app/components/Views/confirmations/hooks/pay/useTransactionPayPostQuote.ts b/app/components/Views/confirmations/hooks/pay/useTransactionPayPostQuote.ts
@@ -6,6 +6,7 @@ import { useTransactionPayWithdraw } from './useTransactionPayWithdraw';
 import { useTransactionMetadataRequest } from '../transactions/useTransactionMetadataRequest';
 import { computeProxyAddress } from '../../../../UI/Predict/providers/polymarket/safe/utils';
 import { hasTransactionType } from '../../utils/transaction';
+import { usePredictAccountState } from '../../../../UI/Predict/hooks/usePredictAccountState';
 
 const log = createProjectLogger('transaction-pay-post-quote');
 
@@ -37,6 +38,13 @@ export function useTransactionPayPostQuote(): void {
     TransactionType.predictWithdraw,
   ]);
 
+  const { data: accountState } = usePredictAccountState({
+    enabled: isPredictWithdraw,
+  });
+
+  const isDepositWalletWithdraw =
+    isPredictWithdraw && accountState?.walletType === 'deposit-wallet';
+
   useEffect(() => {
     if (
       !canSelectWithdrawToken ||
@@ -46,6 +54,10 @@ export function useTransactionPayPostQuote(): void {
       return;
     }
 
+    if (isPredictWithdraw && !accountState) {
+      return;
+    }
+
     try {
       const { TransactionPayController } = Engine.context;
       const from = transactionMeta?.txParams?.from as Hex | undefined;
@@ -55,7 +67,7 @@ export function useTransactionPayPostQuote(): void {
       // on the user's address directly (HyperCore -> Relay for perps; vault
       // teller -> user for money account).
       const refundTo =
-        isPerpsWithdraw || isMoneyAccountWithdraw
+        isPerpsWithdraw || isMoneyAccountWithdraw || isDepositWalletWithdraw
           ? undefined
           : from
             ? computeProxyAddress(from)
@@ -71,6 +83,10 @@ export function useTransactionPayPostQuote(): void {
         if (isPerpsWithdraw) {
           config.isHyperliquidSource = true;
         }
+
+        if (isDepositWalletWithdraw) {
+          config.isPolymarketDepositWallet = true;
+        }
       });
 
       isSet.current = transactionId;
@@ -80,61 +96,22 @@ export function useTransactionPayPostQuote(): void {
         refundTo,
         isPerpsWithdraw,
         isMoneyAccountWithdraw,
+        isDepositWalletWithdraw,
       });
-
-      if (isPredictWithdraw) {
-        applyDepositWalletFlag(transactionId);
-      }
     } catch (error) {
       log('Error initializing post-quote transaction', {
         error,
         transactionId,
       });
     }
   }, [
+    accountState,
     canSelectWithdrawToken,
+    isDepositWalletWithdraw,
     isMoneyAccountWithdraw,
     isPerpsWithdraw,
     isPredictWithdraw,
     transactionId,
     transactionMeta?.txParams?.from,
   ]);
 }
-
-// Deposit-wallet Predict withdrawals need a follow-up flag set after
-// resolving the user's Polymarket account state. The strategy switch
-// also drops refundTo because the bridge mints its own deposit address.
-async function applyDepositWalletFlag(transactionId: string): Promise<void> {
-  try {
-    const isDepositWallet = await isPolymarketDepositWalletWithdraw();
-    if (!isDepositWallet) {
-      return;
-    }
-    Engine.context.TransactionPayController.setTransactionConfig(
-      transactionId,
-      (config) => {
-        config.isPolymarketDepositWallet = true;
-        config.refundTo = undefined;
-      },
-    );
-    log('Marked transaction as Polymarket deposit-wallet withdraw', {
-      transactionId,
-    });
-  } catch (error) {
-    log('Failed to apply Polymarket deposit-wallet flag', {
-      error,
-      transactionId,
-    });
-  }
-}
-
-async function isPolymarketDepositWalletWithdraw(): Promise<boolean> {
-  try {
-    const accountState =
-      await Engine.context.PredictController.getAccountState();
-    return accountState.walletType === 'deposit-wallet';
-  } catch (error) {
-    log('Failed to resolve Polymarket account state', { error });
-    return false;
-  }
-}
diff --git a/app/selectors/transactionController.test.ts b/app/selectors/transactionController.test.ts
@@ -344,57 +344,6 @@ describe('TransactionController Selectors', () => {
       ]);
     });
 
-    it('keeps nonce-less transactions that share an actionId', () => {
-      const activeEvmAddress = '0x0000000000000000000000000000000000000001';
-      const state = {
-        engine: {
-          backgroundState: {
-            AccountsController: {
-              internalAccounts: {
-                selectedAccount: 'account-1',
-                accounts: {
-                  'account-1': {
-                    id: 'account-1',
-                    address: activeEvmAddress,
-                    type: 'eip155:eoa',
-                  },
-                },
-              },
-            },
-            TransactionController: {
-              transactions: [
-                {
-                  id: 'tx-a',
-                  chainId: '0x1',
-                  time: 100,
-                  txParams: {
-                    from: activeEvmAddress,
-                    actionId: 'shared-action',
-                  },
-                },
-                {
-                  id: 'tx-b',
-                  chainId: '0x1',
-                  time: 200,
-                  txParams: {
-                    from: activeEvmAddress,
-                    actionId: 'shared-action',
-                  },
-                },
-              ],
-            },
-          },
-        },
-        pendingSmartTransactionsForGroup: [],
-      } as unknown as RootState;
-
-      const ids = selectLocalTransactions(state).map(
-        (t) => (t as { id: string }).id,
-      );
-
-      expect(ids).toContain('tx-a');
-      expect(ids).toContain('tx-b');
-    });
   });
 
   describe('selectTransactionMetadataById', () => {
