Fix: on close callback for WalletConnect (#2696)

* Refactor: on close callback for WalletConnect

* Use safeAddress + chainId instead of raw query

* Close modal on path change

src/components/tx-flow/index.tsx

@@ -1,9 +1,9 @@
-import { createContext, type ReactElement, type ReactNode, useState, useEffect, useCallback } from 'react'
-import TxModalDialog from '@/components/common/TxModalDialog'
+import { createContext, type ReactElement, type ReactNode, useState, useEffect, useCallback, useRef } from 'react'
 import { usePathname } from 'next/navigation'
-import useSafeInfo from '@/hooks/useSafeInfo'
-import { txDispatch, TxEvent } from '@/services/tx/txEvents'
+import TxModalDialog from '@/components/common/TxModalDialog'
 import { SuccessScreen } from './flows/SuccessScreen'
+import useSafeAddress from '@/hooks/useSafeAddress'
+import useChainId from '@/hooks/useChainId'
 
 const noop = () => {}
 
@@ -19,68 +19,74 @@ export const TxModalContext = createContext<TxModalContextType>({
   setFullWidth: noop,
 })
 
-const shouldClose = () => confirm('Closing this window will discard your current progress.')
+const confirmClose = () => {
+  return confirm('Closing this window will discard your current progress.')
+}
 
 export const TxModalProvider = ({ children }: { children: ReactNode }): ReactElement => {
   const [txFlow, setFlow] = useState<TxModalContextType['txFlow']>(undefined)
-  const [shouldWarn, setShouldWarn] = useState<boolean>(false)
-  const [, setOnClose] = useState<Parameters<TxModalContextType['setTxFlow']>[1]>(noop)
   const [fullWidth, setFullWidth] = useState<boolean>(false)
+  const shouldWarn = useRef<boolean>(true)
+  const onClose = useRef<() => void>(noop)
+  const safeId = useChainId() + useSafeAddress()
+  const prevSafeId = useRef<string>(safeId ?? '')
   const pathname = usePathname()
-  const [, setLastPath] = useState<string>(pathname)
-  const { safeAddress, safe } = useSafeInfo()
+  const prevPathname = useRef<string>(pathname)
 
   const handleModalClose = useCallback(() => {
-    setOnClose((prevOnClose) => {
-      prevOnClose?.()
-      return noop
-    })
+    onClose.current()
+    onClose.current = noop
     setFlow(undefined)
-    txDispatch(TxEvent.USER_QUIT, {})
-  }, [setFlow, setOnClose])
+  }, [])
 
   const handleShowWarning = useCallback(() => {
-    if (!shouldWarn) {
-      handleModalClose()
+    if (shouldWarn.current && !confirmClose()) {
       return
     }
-
-    if (!shouldClose()) return
-
     handleModalClose()
-  }, [shouldWarn, handleModalClose])
+  }, [handleModalClose])
 
   const setTxFlow = useCallback(
-    (newTxFlow: TxModalContextType['txFlow'], onClose?: () => void, newShouldWarn?: boolean) => {
-      // If flow is open and user opens a different one, show confirmation dialog if required
-      if (txFlow && newTxFlow && newTxFlow?.type !== SuccessScreen) {
-        if (shouldWarn && !shouldClose()) return
-        txDispatch(TxEvent.USER_QUIT, {})
-      }
-
-      setFlow(newTxFlow)
-      setOnClose(() => onClose ?? noop)
-      setShouldWarn(newShouldWarn ?? true)
+    (newTxFlow: TxModalContextType['txFlow'], newOnClose?: () => void, newShouldWarn?: boolean) => {
+      setFlow((prev) => {
+        if (prev === newTxFlow) return prev
+
+        // If a new flow is triggered, close the current one
+        if (prev && newTxFlow?.type !== SuccessScreen) {
+          if (shouldWarn.current && !confirmClose()) {
+            return prev
+          }
+          onClose.current()
+        }
+
+        onClose.current = newOnClose ?? noop
+        shouldWarn.current = newShouldWarn ?? true
+
+        return newTxFlow
+      })
     },
-    [txFlow, shouldWarn],
+    [],
   )
 
-  // Show the confirmation dialog if user navigates
+  // // Close the modal when the Safe changes
   useEffect(() => {
-    setLastPath((prev) => {
-      if (prev !== pathname && txFlow) {
-        handleShowWarning()
-      }
-      return pathname
-    })
-  }, [txFlow, handleShowWarning, pathname])
-
-  // Close the modal when the Safe changes
+    if (safeId === prevSafeId.current) return
+    prevSafeId.current = safeId
+
+    if (txFlow) {
+      handleShowWarning()
+    }
+  }, [txFlow, safeId])
+
+  // // Close the modal when the path changes
   useEffect(() => {
-    handleModalClose()
-    // Could have same address but different chain
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [safe.chainId, safeAddress])
+    if (pathname === prevPathname.current) return
+    prevPathname.current = pathname
+
+    if (txFlow) {
+      handleShowWarning()
+    }
+  }, [txFlow, pathname])
 
   return (
     <TxModalContext.Provider value={{ txFlow, setTxFlow, setFullWidth }}>

src/services/safe-wallet-provider/useSafeWalletProvider.tsx

@@ -66,18 +66,16 @@ export const _useTxFlowApi = (chainId: string, safeAddress: string): WalletSDK |
       const shouldSignOffChain =
         isOffchainEIP1271Supported(safe, currentChain) && !onChainSigning && settings.offChainSigning
 
-      if (shouldSignOffChain) {
-        setTxFlow(<SignMessageFlow logoUri={appInfo.iconUrl} name={appInfo.name} message={message} requestId={id} />)
-      } else {
-        setTxFlow(<SignMessageOnChainFlow props={{ appId: wcApp?.id, requestId: id, message, method }} />)
-      }
       const { title, options } = NotificationMessages.SIGNATURE_REQUEST(appInfo)
       showNotification(title, options)
 
       return new Promise((resolve, reject) => {
-        const unsubscribe = () => {
-          unsubscribeSignaturePrepared()
-          unsubscribeUserQuit()
+        let onClose = () => {
+          reject({
+            code: RpcErrorCode.USER_REJECTED,
+            message: 'User rejected signature',
+          })
+          unsubscribe()
         }
 
         const unsubscribeSignaturePrepared = safeMsgSubscribe(
@@ -90,13 +88,19 @@ export const _useTxFlowApi = (chainId: string, safeAddress: string): WalletSDK |
           },
         )
 
-        const unsubscribeUserQuit = txSubscribe(TxEvent.USER_QUIT, () => {
-          reject({
-            code: RpcErrorCode.USER_REJECTED,
-            message: 'User rejected signature request',
-          })
-          unsubscribe()
-        })
+        const unsubscribe = () => {
+          onClose = () => {}
+          unsubscribeSignaturePrepared()
+        }
+
+        if (shouldSignOffChain) {
+          setTxFlow(
+            <SignMessageFlow logoUri={appInfo.iconUrl} name={appInfo.name} message={message} requestId={id} />,
+            onClose,
+          )
+        } else {
+          setTxFlow(<SignMessageOnChainFlow props={{ appId: wcApp?.id, requestId: id, message, method }} />, onClose)
+        }
       })
     }
 
@@ -120,30 +124,16 @@ export const _useTxFlowApi = (chainId: string, safeAddress: string): WalletSDK |
           }
         })
 
-        setTxFlow(
-          <SafeAppsTxFlow
-            data={{
-              appId: undefined,
-              app: {
-                name: appInfo.name,
-                // Show WC details in transaction list
-                url: wcApp?.url ?? appInfo.url,
-                iconUrl: appInfo.iconUrl,
-              },
-              requestId: id,
-              txs: transactions,
-              params: params.params,
-            }}
-          />,
-        )
-
         const { title, options } = NotificationMessages.TRANSACTION_REQUEST(appInfo)
         showNotification(title, options)
 
         return new Promise((resolve, reject) => {
-          const unsubscribe = () => {
-            unsubscribeSignaturePrepared()
-            unsubscribeUserQuit()
+          let onClose = () => {
+            reject({
+              code: RpcErrorCode.USER_REJECTED,
+              message: 'User rejected transaction',
+            })
+            unsubscribe()
           }
 
           const unsubscribeSignaturePrepared = txSubscribe(
@@ -157,13 +147,28 @@ export const _useTxFlowApi = (chainId: string, safeAddress: string): WalletSDK |
             },
           )
 
-          const unsubscribeUserQuit = txSubscribe(TxEvent.USER_QUIT, () => {
-            reject({
-              code: RpcErrorCode.USER_REJECTED,
-              message: 'User rejected transaction',
-            })
-            unsubscribe()
-          })
+          const unsubscribe = () => {
+            onClose = () => {}
+            unsubscribeSignaturePrepared()
+          }
+
+          setTxFlow(
+            <SafeAppsTxFlow
+              data={{
+                appId: undefined,
+                app: {
+                  name: appInfo.name,
+                  // Show WC details in transaction list
+                  url: wcApp?.url ?? appInfo.url,
+                  iconUrl: appInfo.iconUrl,
+                },
+                requestId: id,
+                txs: transactions,
+                params: params.params,
+              }}
+            />,
+            onClose,
+          )
         })
       },
 

src/services/tx/txEvents.ts

@@ -2,7 +2,6 @@ import EventBus from '@/services/EventBus'
 import type { RequestId } from '@safe-global/safe-apps-sdk'
 
 export enum TxEvent {
-  USER_QUIT = 'USER_QUIT',
   SIGNED = 'SIGNED',
   SIGN_FAILED = 'SIGN_FAILED',
   PROPOSED = 'PROPOSED',
@@ -27,7 +26,6 @@ export enum TxEvent {
 type Id = { txId: string; groupKey?: string } | { txId?: string; groupKey: string }
 
 interface TxEvents {
-  [TxEvent.USER_QUIT]: {}
   [TxEvent.SIGNED]: { txId?: string }
   [TxEvent.SIGN_FAILED]: { txId?: string; error: Error }
   [TxEvent.PROPOSE_FAILED]: { error: Error }

src/services/walletconnect/WalletConnectContext.tsx

@@ -7,6 +7,7 @@ import useSafeWalletProvider from '@/services/safe-wallet-provider/useSafeWallet
 import WalletConnectWallet from './WalletConnectWallet'
 import { asError } from '../exceptions/utils'
 import { getPeerName, stripEip155Prefix } from './utils'
+import { IS_PRODUCTION } from '@/config/constants'
 
 const walletConnectSingleton = new WalletConnectWallet()
 
@@ -54,6 +55,10 @@ export const WalletConnectProvider = ({ children }: { children: ReactNode }) =>
     if (!walletConnect || !safeWalletProvider || !chainId) return
 
     return walletConnect.onRequest(async (event) => {
+      if (!IS_PRODUCTION) {
+        console.log('[WalletConnect] request', event)
+      }
+
       const { topic } = event
       const session = walletConnect.getActiveSessions().find((s) => s.topic === topic)
       const requestChainId = stripEip155Prefix(event.params.chainId)