update wallet balance from websocket

components/Balances.tsx

@@ -12,13 +12,13 @@ const Balances = () => {
     return quotient.toNumber() + remainder.toNumber() / divisor.toNumber()
   }
 
-  function calculateBalance(a) {
-    const mint = mints[a.account.mint.toBase58()]
-    return mint ? fixedPointToNumber(a.account.amount, mint.decimals) : 0
+  function calculateBalance(acc) {
+    const mint = mints[acc.mint.toBase58()]
+    return mint ? fixedPointToNumber(acc.amount, mint.decimals) : 0
   }
 
-  const displayedBalances = tokenAccounts
-    .map((a) => ({ id: a.publicKey.toBase58(), balance: calculateBalance(a) }))
+  const displayedBalances = Object.entries(tokenAccounts)
+    .map(([id, acc]) => ({ id, balance: calculateBalance(acc) }))
     .sort((a, b) => (a.id > b.id ? 1 : -1))
 
   return (

hooks/useWallet.tsx

@@ -127,9 +127,8 @@ export default function useWallet() {
   }, [wallet, setWalletStore])
 
   useInterval(async () => {
-    await actions.fetchWalletTokenAccounts()
     await actions.fetchWalletMints()
-  }, 20 * SECONDS)
+  }, 120 * SECONDS)
 
   return { connected, wallet }
 }

stores/useWalletStore.tsx

@@ -6,7 +6,7 @@ import { EndpointInfo, WalletAdapter } from '../@types/types'
 import {
   getOwnedTokenAccounts,
   getMint,
-  ProgramAccount,
+  subscribeToTokenAccount,
   TokenAccount,
   MintAccount,
 } from '../utils/tokens'
@@ -39,8 +39,9 @@ interface WalletStore extends State {
   }
   current: WalletAdapter | undefined
   providerUrl: string
-  tokenAccounts: ProgramAccount<TokenAccount>[]
+  tokenAccounts: { [pubkey: string]: TokenAccount }
   mints: { [pubkey: string]: MintAccount }
+  subscriptions: { [pubkey: string]: () => void }
   set: (x: any) => void
   actions: any
 }
@@ -55,28 +56,43 @@ const useWalletStore = create<WalletStore>((set, get) => ({
   },
   current: null,
   providerUrl: null,
-  tokenAccounts: [],
+  tokenAccounts: {},
   mints: {},
+  subscriptions: {},
   actions: {
     async fetchWalletTokenAccounts() {
       const connection = get().connection.current
       const connected = get().connected
+      const subscriptions = get().subscriptions
       const wallet = get().current
       const walletOwner = wallet?.publicKey
       const set = get().set
 
       if (connected && walletOwner) {
-        const ownedTokenAccounts = await getOwnedTokenAccounts(
+        const ownedAccountsResult = await getOwnedTokenAccounts(
           connection,
           walletOwner
         )
 
+        const newTokenAccounts: { [pubkey: string]: TokenAccount } = {}
+        ownedAccountsResult.forEach((r) => {
+          newTokenAccounts[r.publicKey.toBase58()] = r.account
+        })
+
+        // cancel all subscriptions
+        Object.values(subscriptions).forEach((s) => s())
+
         set((state) => {
-          state.tokenAccounts = ownedTokenAccounts
+          state.subscriptions = {}
+          state.tokenAccounts = newTokenAccounts
+        })
+
+        ownedAccountsResult.forEach((r) => {
+          this.subscribeToTokenAccount(r.publicKey)
         })
       } else {
         set((state) => {
-          state.tokenAccounts = []
+          state.tokenAccounts = {}
         })
       }
     },
@@ -87,14 +103,14 @@ const useWalletStore = create<WalletStore>((set, get) => ({
       const set = get().set
 
       if (connected) {
-        const fetchMints = tokenAccounts.map((a) =>
-          getMint(connection, a.account.mint)
+        const fetchMints = Object.values(tokenAccounts).map((a) =>
+          getMint(connection, a.mint)
         )
         const mintResults = await Promise.all(fetchMints)
 
         const newMints: { [pubkey: string]: MintAccount } = {}
         mintResults.forEach(
-          (m) => (newMints[m.publicKey.toBase58()] = m.account)
+          (r) => (newMints[r.publicKey.toBase58()] = r.account)
         )
 
         set((state) => {
@@ -106,6 +122,23 @@ const useWalletStore = create<WalletStore>((set, get) => ({
         })
       }
     },
+    async subscribeToTokenAccount(pubkey) {
+      const connection = get().connection.websocket
+      const connected = get().connected
+      const set = get().set
+
+      if (connected) {
+        const sub = subscribeToTokenAccount(connection, pubkey, (r) => {
+          set((s) => {
+            s.tokenAccounts[pubkey.toBase58()] = r.account
+          })
+        })
+
+        set((s) => {
+          s.subscriptions[pubkey.toBase58()] = sub
+        })
+      }
+    },
   },
   set: (fn) => set(produce(fn)),
 }))

utils/tokens.tsx

@@ -9,7 +9,7 @@ import {
 
 export type TokenAccount = AccountInfo
 export type MintAccount = MintInfo
-export type ProgramAccount<T> = {
+export type AccountResponse<T> = {
   publicKey: PublicKey
   account: T
 }
@@ -17,7 +17,7 @@ export type ProgramAccount<T> = {
 export async function getOwnedTokenAccounts(
   connection: Connection,
   publicKey: PublicKey
-): Promise<ProgramAccount<TokenAccount>[]> {
+): Promise<AccountResponse<TokenAccount>[]> {
   const results = await connection.getTokenAccountsByOwner(publicKey, {
     programId: TOKEN_PROGRAM_ID,
   })
@@ -32,7 +32,7 @@ export async function getOwnedTokenAccounts(
 export async function getMint(
   connection: Connection,
   publicKey: PublicKey
-): Promise<ProgramAccount<MintAccount>> {
+): Promise<AccountResponse<MintAccount>> {
   const result = await connection.getAccountInfo(publicKey)
   const data = Buffer.from(result.data)
   const account = parseMintAccountData(data)
@@ -42,6 +42,21 @@ export async function getMint(
   }
 }
 
+export function subscribeToTokenAccount(
+  connection: Connection,
+  publicKey: PublicKey,
+  cb: (r: AccountResponse<TokenAccount>) => void
+): () => void {
+  const id = connection.onAccountChange(publicKey, (info) => {
+    const data = Buffer.from(info.data)
+    const account = parseTokenAccountData(publicKey, data)
+    cb({ publicKey, account })
+  })
+  return () => {
+    connection.removeAccountChangeListener(id)
+  }
+}
+
 // copied from @solana/spl-token
 
 const TOKEN_PROGRAM_ID = new PublicKey(