Fix: prefer last chainId over wallet chainId

src/hooks/__tests__/useChainId.test.ts

@@ -1,7 +1,5 @@
 import { useParams } from 'next/navigation'
 import useChainId from '@/hooks/useChainId'
-import { useAppDispatch } from '@/store'
-import { setLastChainId } from '@/store/sessionSlice'
 import { renderHook } from '@/tests/test-utils'
 import * as useWalletHook from '@/hooks/wallets/useWallet'
 import * as useChains from '@/hooks/useChains'
@@ -24,22 +22,6 @@ describe('useChainId hook', () => {
     })
   })
 
-  it('should read location.pathname if useRouter query.safe is empty', () => {
-    ;(useParams as any).mockImplementation(() => ({}))
-
-    Object.defineProperty(window, 'location', {
-      writable: true,
-      value: {
-        pathname: '/avax:0x0000000000000000000000000000000000000123/balances',
-        search: '',
-      },
-    })
-
-    const { result } = renderHook(() => useChainId())
-
-    expect(result.current).toEqual('43114')
-  })
-
   it('should read location.search if useRouter query.safe is empty', () => {
     ;(useParams as any).mockImplementation(() => ({}))
 
@@ -95,7 +77,7 @@ describe('useChainId hook', () => {
     expect(result.current).toBe('137')
   })
 
-  it('should return the wallet chain id if no chain in the URL and it is present in the chain configs', () => {
+  it('should return the wallet chain id if no chain in the URL and no last chain id', () => {
     ;(useParams as any).mockImplementation(() => ({}))
 
     jest.spyOn(useWalletHook, 'default').mockImplementation(
@@ -113,28 +95,10 @@ describe('useChainId hook', () => {
     expect(result.current).toBe('1337')
   })
 
-  it('should return the last used chain id if no chain in the URL and the connect wallet chain id is not present in the chain configs', () => {
-    ;(useParams as any).mockImplementation(() => ({}))
-
-    jest.spyOn(useWalletHook, 'default').mockImplementation(
-      () =>
-        ({
-          chainId: '1337',
-        } as ConnectedWallet),
-    )
-
-    jest.spyOn(useChains, 'default').mockImplementation(() => ({
-      configs: [],
-    }))
-
-    const { result } = renderHook(() => useChainId())
-    expect(result.current).toBe('11155111')
-  })
-
-  it('should return the last used chain id if no wallet is connected and there is no chain in the URL', () => {
+  it('should return the last chain id', () => {
     ;(useParams as any).mockImplementation(() => ({}))
 
-    renderHook(() => useAppDispatch()(setLastChainId('100')))
+    localStorage.setItem('SAFE_v2__session', `{"lastChainId": "100"}`)
 
     const { result } = renderHook(() => useChainId())
     expect(result.current).toBe('100')

src/hooks/useChainId.ts

@@ -5,7 +5,6 @@ import chains from '@/config/chains'
 import { useAppSelector } from '@/store'
 import { selectSession } from '@/store/sessionSlice'
 import { parsePrefixedAddress } from '@/utils/addresses'
-import { prefixedAddressRe } from '@/utils/url'
 import useWallet from './wallets/useWallet'
 import useChains from './useChains'
 
@@ -14,19 +13,7 @@ const defaultChainId = IS_PRODUCTION ? chains.eth : chains.sep
 // Use the location object directly because Next.js's router.query is available only on mount
 const getLocationQuery = (): ParsedUrlQuery => {
   if (typeof location === 'undefined') return {}
-
   const query = parse(location.search.slice(1))
-
-  if (!query.safe) {
-    const pathParam = location.pathname.split('/')[1]
-    const safeParam = prefixedAddressRe.test(pathParam) ? pathParam : ''
-
-    // Path param -> query param
-    if (prefixedAddressRe.test(pathParam)) {
-      query.safe = safeParam
-    }
-  }
-
   return query
 }
 
@@ -47,16 +34,20 @@ export const useUrlChainId = (): string | undefined => {
   return chains[shortName] || configs.find((item) => item.shortName === shortName)?.chainId
 }
 
-export const useChainId = (): string => {
-  const session = useAppSelector(selectSession)
-  const urlChainId = useUrlChainId()
+const useWalletChainId = (): string | undefined => {
   const wallet = useWallet()
   const { configs } = useChains()
-
   const walletChainId =
     wallet?.chainId && configs.some(({ chainId }) => chainId === wallet.chainId) ? wallet.chainId : undefined
+  return walletChainId
+}
+
+export const useChainId = (): string => {
+  const session = useAppSelector(selectSession)
+  const urlChainId = useUrlChainId()
+  const walletChainId = useWalletChainId()
 
-  return urlChainId || walletChainId || session.lastChainId || defaultChainId
+  return urlChainId || session.lastChainId || walletChainId || defaultChainId
 }
 
 export default useChainId