From a0149d8565fce6c9edf2f2aa8c3765c5a8c387d1 Mon Sep 17 00:00:00 2001
From: Jordan Ribbink <jribbink@telus.net>
Date: Thu, 8 Aug 2024 12:55:46 -0700
Subject: [PATCH] Fix last used

---
 components/views/ConnectExtension.tsx         |  8 ++++
 components/views/ScanConnect.tsx              |  8 +++-
 helpers/rpc.ts                                |  1 +
 ...{useLocalStorage.js => useLocalStorage.ts} |  0
 hooks/useWalletHistory.ts                     |  6 +--
 hooks/useWallets.ts                           |  2 +-
 hooks/useWcUri.ts                             | 46 +++++++++++++++----
 7 files changed, 57 insertions(+), 14 deletions(-)
 rename hooks/{useLocalStorage.js => useLocalStorage.ts} (100%)

diff --git a/components/views/ConnectExtension.tsx b/components/views/ConnectExtension.tsx
index cfed137b..fc90831b 100644
--- a/components/views/ConnectExtension.tsx
+++ b/components/views/ConnectExtension.tsx
@@ -5,6 +5,8 @@ import { FCL_SERVICE_METHODS } from '../../helpers/constants'
 import { useRpc } from '../../contexts/FclContext'
 import { FclRequest } from '../../helpers/rpc'
 import WalletIcon from '../Icons/WalletIcon'
+import { useWalletHistory } from '../../hooks/useWalletHistory'
+import { handleCancel } from '../../helpers/window'
 
 type ConnectExtensionProps = {
   wallet: Wallet
@@ -15,6 +17,7 @@ export default function ConnectExtension({ wallet }: ConnectExtensionProps) {
   const [isConnecting, setIsConnecting] = useState(false)
   const hasAttemptedConnection = useRef(true)
   const showSpinner = !rpc || isConnecting
+  const { setLastUsed } = useWalletHistory()
 
   function connect() {
     setIsConnecting(true)
@@ -22,6 +25,11 @@ export default function ConnectExtension({ wallet }: ConnectExtensionProps) {
       if (service.method === FCL_SERVICE_METHODS.EXT) {
         rpc
           .request(FclRequest.EXEC_SERVICE, { service })
+          .then(() => {
+            console.log('Connected to extension')
+            setLastUsed(wallet)
+            handleCancel()
+          })
           .catch(e => {
             console.error('Failed to connect to extension', e)
           })
diff --git a/components/views/ScanConnect.tsx b/components/views/ScanConnect.tsx
index d3aa7efb..fb33d487 100644
--- a/components/views/ScanConnect.tsx
+++ b/components/views/ScanConnect.tsx
@@ -4,6 +4,8 @@ import QRCode from '../QRCode'
 import CopyButton from '../CopyButton'
 import HybridButton from '../HybridButton'
 import { useWcUri } from '../../hooks/useWcUri'
+import { useWalletHistory } from '../../hooks/useWalletHistory'
+import { handleCancel } from '../../helpers/window'
 
 interface ScanConnectProps {
   wallet: Wallet
@@ -11,7 +13,11 @@ interface ScanConnectProps {
 }
 
 export default function ScanConnect({ wallet, onGetWallet }: ScanConnectProps) {
-  const { uri, error, isLoading } = useWcUri()
+  const { setLastUsed } = useWalletHistory()
+  const { uri, error, isLoading } = useWcUri(() => {
+    setLastUsed(wallet)
+    handleCancel()
+  })
 
   return (
     <Stack
diff --git a/helpers/rpc.ts b/helpers/rpc.ts
index 64dc26f3..fbc75265 100644
--- a/helpers/rpc.ts
+++ b/helpers/rpc.ts
@@ -3,6 +3,7 @@ import { Service } from '../types'
 
 export enum DiscoveryNotification {
   NOTIFY_QRCODE_ERROR = 'notifyQRCodeError',
+  NOTIFY_QRCODE_CONNECTED = 'notifyQRCodeConnected',
 }
 
 export enum FclRequest {
diff --git a/hooks/useLocalStorage.js b/hooks/useLocalStorage.ts
similarity index 100%
rename from hooks/useLocalStorage.js
rename to hooks/useLocalStorage.ts
diff --git a/hooks/useWalletHistory.ts b/hooks/useWalletHistory.ts
index f7dda3dc..98ff7e13 100644
--- a/hooks/useWalletHistory.ts
+++ b/hooks/useWalletHistory.ts
@@ -6,7 +6,7 @@ import { useLocalStorage } from './useLocalStorage'
 export function useWalletHistory() {
   const [, setLastUsedState] = useLocalStorage(
     LOCAL_STORAGE_KEYS.LAST_USED,
-    null
+    null,
   )
   const [lastUsedUid] = useLocalStorage(LOCAL_STORAGE_KEYS.LAST_USED, null)
 
@@ -14,12 +14,12 @@ export function useWalletHistory() {
     (wallet: Wallet) => {
       setLastUsedState(wallet.uid)
     },
-    [setLastUsedState]
+    [setLastUsedState],
   )
 
   const isLastUsed = useCallback(
     (wallet: Wallet) => wallet.uid === lastUsedUid,
-    [lastUsedUid]
+    [lastUsedUid],
   )
 
   return { isLastUsed, setLastUsed, lastUsedUid }
diff --git a/hooks/useWallets.ts b/hooks/useWallets.ts
index 367428e4..317a0a68 100644
--- a/hooks/useWallets.ts
+++ b/hooks/useWallets.ts
@@ -42,7 +42,7 @@ export function useWallets() {
   }
 
   const { data: wallets, error } = useSWR(genKey(requestUrl, body), url =>
-    fetcher<Wallet[]>(url, body)
+    fetcher<Wallet[]>(url, body),
   )
 
   return { wallets, error, isLoading: !wallets && !error }
diff --git a/hooks/useWcUri.ts b/hooks/useWcUri.ts
index 0d60b8c4..430785ce 100644
--- a/hooks/useWcUri.ts
+++ b/hooks/useWcUri.ts
@@ -1,8 +1,9 @@
 import useSWR from 'swr'
 import { useRpc } from '../contexts/FclContext'
 import { DiscoveryNotification, FclRequest } from '../helpers/rpc'
+import { useEffect } from 'react'
 
-export function useWcUri() {
+export function useWcUri(onConnected?: () => void) {
   const rpc = useRpc()
   const {
     data: uri,
@@ -14,17 +15,44 @@ export function useWcUri() {
       {},
     )
 
-    // Subscribe to QR error notifications (e.g. user declined, QR expired)
-    const onError = ({ error: _error, uri: _uri }) => {
-      if (_uri === uri) {
-        // QR code is no longer valid, reset the URI and unsubscribe
+    return uri
+  })
+
+  useEffect(() => {
+    if (!rpc || !uri) {
+      return
+    }
+
+    const connectHandler = ({ uri: _uri }: { uri: string }) => {
+      if (uri === _uri) {
+        onConnected?.()
+      }
+    }
+
+    const errorHandler = ({
+      uri: errorUri,
+      error,
+    }: {
+      uri: string
+      error: string
+    }) => {
+      if (uri === errorUri) {
+        console.error('URI connection error:', error)
         mutate()
-        rpc.unsubscribe(DiscoveryNotification.NOTIFY_QRCODE_ERROR, onError)
       }
     }
-    rpc.subscribe(DiscoveryNotification.NOTIFY_QRCODE_ERROR, onError)
 
-    return uri
-  })
+    rpc.subscribe(DiscoveryNotification.NOTIFY_QRCODE_CONNECTED, connectHandler)
+    rpc.subscribe(DiscoveryNotification.NOTIFY_QRCODE_ERROR, errorHandler)
+
+    return () => {
+      rpc.unsubscribe(
+        DiscoveryNotification.NOTIFY_QRCODE_CONNECTED,
+        connectHandler,
+      )
+      rpc.unsubscribe(DiscoveryNotification.NOTIFY_QRCODE_ERROR, errorHandler)
+    }
+  }, [uri, rpc])
+
   return { uri, error, isLoading: !uri && !error }
 }