From ac21ff8d6183284abb38ea4b2fc5693bde262ec2 Mon Sep 17 00:00:00 2001
From: kyranjamie <kyran.burraston@gmail.com>
Date: Thu, 1 Aug 2024 15:21:23 +0200
Subject: [PATCH] fix: account change in open frames

---
 src/app/common/hooks/use-key-actions.ts       |  7 +++++--
 src/app/features/container/container.tsx      |  6 +++++-
 src/app/routes/hooks/use-on-change-account.ts | 13 +++++++++++++
 src/shared/message-types.ts                   |  1 +
 src/shared/messages.ts                        |  4 +++-
 5 files changed, 27 insertions(+), 4 deletions(-)
 create mode 100644 src/app/routes/hooks/use-on-change-account.ts

diff --git a/src/app/common/hooks/use-key-actions.ts b/src/app/common/hooks/use-key-actions.ts
index c52aa2c1afa..5b67c98b926 100644
--- a/src/app/common/hooks/use-key-actions.ts
+++ b/src/app/common/hooks/use-key-actions.ts
@@ -5,6 +5,8 @@ import { generateSecretKey } from '@stacks/wallet-sdk';
 import { useBitcoinClient } from '@leather.io/query';
 
 import { logger } from '@shared/logger';
+import { InternalMethods } from '@shared/message-types';
+import { sendMessage } from '@shared/messages';
 import { clearChromeStorage } from '@shared/storage/redux-pesist';
 import { analytics } from '@shared/utils/analytics';
 
@@ -46,8 +48,9 @@ export function useKeyActions() {
         return dispatch(keyActions.unlockWalletAction(password));
       },
 
-      switchAccount(index: number) {
-        return dispatch(stxChainActions.switchAccount(index));
+      switchAccount(accountIndex: number) {
+        sendMessage({ method: InternalMethods.AccountChanged, payload: { accountIndex } });
+        return dispatch(stxChainActions.switchAccount(accountIndex));
       },
 
       async createNewAccount() {
diff --git a/src/app/features/container/container.tsx b/src/app/features/container/container.tsx
index b6a5c6afaaf..e6f1827dbb4 100644
--- a/src/app/features/container/container.tsx
+++ b/src/app/features/container/container.tsx
@@ -13,9 +13,11 @@ import { ContainerLayout } from '@app/components/layout';
 import { LoadingSpinner } from '@app/components/loading-spinner';
 import { SwitchAccountDialog } from '@app/features/dialogs/switch-account-dialog/switch-account-dialog';
 import { InAppMessages } from '@app/features/hiro-messages/in-app-messages';
+import { useOnChangeAccount } from '@app/routes/hooks/use-on-change-account';
 import { useOnSignOut } from '@app/routes/hooks/use-on-sign-out';
 import { useOnWalletLock } from '@app/routes/hooks/use-on-wallet-lock';
-import { useHasStateRehydrated } from '@app/store';
+import { useAppDispatch, useHasStateRehydrated } from '@app/store';
+import { stxChainSlice } from '@app/store/chains/stx-chain.slice';
 
 import { useRestoreFormState } from '../popup-send-form-restoration/use-restore-form-state';
 
@@ -23,6 +25,7 @@ export function Container() {
   const { pathname: locationPathname } = useLocation();
   const pathname = locationPathname as RouteUrls;
   const [isShowingSwitchAccount, setIsShowingSwitchAccount] = useState(false);
+  const dispatch = useAppDispatch();
 
   const hasStateRehydrated = useHasStateRehydrated();
 
@@ -31,6 +34,7 @@ export function Container() {
   useRestoreFormState();
   useInitalizeAnalytics();
   useHandleQueuedBackgroundAnalytics();
+  useOnChangeAccount(index => dispatch(stxChainSlice.actions.switchAccount(index)));
 
   useEffect(() => void analytics.page('view', `${pathname}`), [pathname]);
 
diff --git a/src/app/routes/hooks/use-on-change-account.ts b/src/app/routes/hooks/use-on-change-account.ts
new file mode 100644
index 00000000000..c1d3c58c1a0
--- /dev/null
+++ b/src/app/routes/hooks/use-on-change-account.ts
@@ -0,0 +1,13 @@
+import { InternalMethods } from '@shared/message-types';
+import type { BackgroundMessages } from '@shared/messages';
+
+import { useOnMount } from '@app/common/hooks/use-on-mount';
+
+export function useOnChangeAccount(handler: (accountIndex: number) => void) {
+  useOnMount(() => {
+    chrome.runtime.onMessage.addListener((message: BackgroundMessages, _sender, sendResponse) => {
+      if (message?.method === InternalMethods.AccountChanged) handler(message.payload.accountIndex);
+      sendResponse();
+    });
+  });
+}
diff --git a/src/shared/message-types.ts b/src/shared/message-types.ts
index c72f326e9ab..d59379e500b 100644
--- a/src/shared/message-types.ts
+++ b/src/shared/message-types.ts
@@ -28,6 +28,7 @@ export enum ExternalMethods {
 export enum InternalMethods {
   RequestDerivedStxAccounts = 'RequestDerivedStxAccounts',
   OriginatingTabClosed = 'OriginatingTabClosed',
+  AccountChanged = 'AccountChanged',
 }
 
 export type ExtensionMethods = ExternalMethods | InternalMethods;
diff --git a/src/shared/messages.ts b/src/shared/messages.ts
index 4ffb8c78cd2..90418eb7306 100644
--- a/src/shared/messages.ts
+++ b/src/shared/messages.ts
@@ -13,7 +13,9 @@ type OriginatingTabClosed = BackgroundMessage<
   { tabId: number }
 >;
 
-export type BackgroundMessages = OriginatingTabClosed;
+type AccountChanged = BackgroundMessage<InternalMethods.AccountChanged, { accountIndex: number }>;
+
+export type BackgroundMessages = OriginatingTabClosed | AccountChanged;
 
 export function sendMessage(message: BackgroundMessages) {
   return chrome.runtime.sendMessage(message);