From f0af6c7458bfa5444678dba9d6e914438c434f10 Mon Sep 17 00:00:00 2001
From: Nikos Kontakis <wirednkod@gmail.com>
Date: Tue, 20 Aug 2024 12:28:04 +0200
Subject: [PATCH] minor fix

---
 src/contexts/AccountsContext.tsx |  7 ++++---
 src/header.tsx                   | 14 ++++++++++++--
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/src/contexts/AccountsContext.tsx b/src/contexts/AccountsContext.tsx
index 5cee3c2..fb40945 100644
--- a/src/contexts/AccountsContext.tsx
+++ b/src/contexts/AccountsContext.tsx
@@ -18,7 +18,7 @@ type AccountContextProps = {
 export interface IAccountContext {
   selectedAccount?: InjectedAccount
   accounts: InjectedAccount[]
-  selectAccount: (account: InjectedAccount) => void
+  selectAccount: (account: InjectedAccount | undefined) => void
 }
 
 const AccountContext = createContext<IAccountContext | undefined>(undefined)
@@ -27,8 +27,9 @@ const AccountContextProvider = ({ children }: AccountContextProps) => {
   const accounts = useRedotAccounts()
   const [selectedAccount, setSelected] = useState<InjectedAccount | undefined>()
 
-  const selectAccount = useCallback((account: InjectedAccount) => {
-    localStorage.setItem(LOCALSTORAGE_SELECTED_ACCOUNT_KEY, account?.address)
+  const selectAccount = useCallback((account: InjectedAccount | undefined) => {
+    account?.address &&
+      localStorage.setItem(LOCALSTORAGE_SELECTED_ACCOUNT_KEY, account?.address)
     setSelected(account)
   }, [])
 
diff --git a/src/header.tsx b/src/header.tsx
index f35967e..5d6650e 100644
--- a/src/header.tsx
+++ b/src/header.tsx
@@ -22,10 +22,18 @@ import {
   MenubarTrigger,
 } from '@/components/ui/menubar'
 import { useAccounts } from './contexts/AccountsContext'
+import { useEffect } from 'react'
 
 export const Header = () => {
   const { accounts, selectAccount, selectedAccount } = useAccounts()
   const [, disconnectAll] = useWalletDisconnector()
+
+  useEffect(() => {
+    if (!selectedAccount?.address && accounts.length > 0) {
+      selectAccount(accounts[0])
+    }
+  }, [accounts, selectAccount, selectedAccount?.address])
+
   return (
     <header className="sticky top-0 z-30 flex h-14 items-center gap-4 border-b bg-background px-4 sm:sticky sm:h-auto sm:border-0 sm:bg-transparent sm:px-6">
       <Sheet>
@@ -105,7 +113,6 @@ export const Header = () => {
                       onClick={() => selectAccount(account)}
                     >
                       <Polkicon
-                        copy
                         size={28}
                         address={account.address || ''}
                         className="mr-2"
@@ -118,7 +125,10 @@ export const Header = () => {
                 <DropdownMenuItem
                   className="cursor-pointer"
                   key={'logout'}
-                  onClick={() => disconnectAll()}
+                  onClick={() => {
+                    disconnectAll()
+                    selectAccount(undefined)
+                  }}
                 >
                   Logout
                 </DropdownMenuItem>