diff --git a/examples/core-connection/src/context/GlobalContext.tsx b/examples/core-connection/src/context/GlobalContext.tsx
index 0452590a..1783c6ca 100644
--- a/examples/core-connection/src/context/GlobalContext.tsx
+++ b/examples/core-connection/src/context/GlobalContext.tsx
@@ -27,7 +27,8 @@ const mockRecipients = [
 ];
 
 export const GlobalProvider = ({ children }: { children: ReactNode }) => {
-  const { account, handleSendSignRequestToPushWallet } = usePushWalletContext();
+  const { universalAddress, handleSendSignRequestToPushWallet } =
+    usePushWalletContext();
 
   const [pushNetwork, setPushNetwork] = useState<PushNetwork | null>(null);
   const [mockTx, setMockTx] = useState<Transaction | null>(null);
@@ -54,7 +55,7 @@ export const GlobalProvider = ({ children }: { children: ReactNode }) => {
   return (
     <GlobalContext.Provider
       value={{
-        account,
+        account: universalAddress ? universalAddress.address : null,
         pushNetwork,
         mockTx,
         handleSendSignRequestToPushWallet,
diff --git a/examples/core-connection/src/helpers/walletHelper.ts b/examples/core-connection/src/helpers/walletHelper.ts
index b3cbf7e4..fd021c02 100644
--- a/examples/core-connection/src/helpers/walletHelper.ts
+++ b/examples/core-connection/src/helpers/walletHelper.ts
@@ -1,3 +1,5 @@
+import { UniversalAddress } from '@pushprotocol/pushchain-ui-kit';
+
 export function centerMaskString(str: string) {
   if (str && str.length > 15) {
     const start = str.substring(0, 6);
@@ -61,3 +63,11 @@ export const convertCaipToObject = (
     };
   }
 };
+
+export const convertToCaip = ({
+  address,
+  chain,
+  chainId,
+}: UniversalAddress) => {
+  return `${chain}:${chainId}:${address}`;
+};
diff --git a/examples/core-connection/src/modules/landingPage/components/LandingPageLeftComponent.tsx b/examples/core-connection/src/modules/landingPage/components/LandingPageLeftComponent.tsx
index 189f9d7a..5a4e6e04 100644
--- a/examples/core-connection/src/modules/landingPage/components/LandingPageLeftComponent.tsx
+++ b/examples/core-connection/src/modules/landingPage/components/LandingPageLeftComponent.tsx
@@ -1,13 +1,18 @@
 import { css } from 'styled-components';
 import { Box, Button, Front, Sale, Text } from 'shared-components';
-import { ConnectPushWalletButton } from '@pushprotocol/pushchain-ui-kit';
 import { useGlobalContext } from '../../../context/GlobalContext';
 import { LandingPageBanner } from './LandingPageBanner';
 import { SimulateTxText } from './SimulateTxText';
+import {
+  PushWalletButton,
+  usePushWalletContext,
+} from '@pushprotocol/pushchain-ui-kit';
 
 const LandingPageLeftComponent = () => {
   const { pushNetwork, mockTx } = useGlobalContext();
 
+  const { universalAddress } = usePushWalletContext();
+
   const featuresCard = [
     {
       id: 1,
@@ -106,7 +111,13 @@ const LandingPageLeftComponent = () => {
               justifyContent="center"
               width="-webkit-fill-available"
             >
-              <ConnectPushWalletButton />
+              <PushWalletButton
+                universalAddress={universalAddress}
+                title="Connect Push Wallet"
+                styling={{
+                  width: 'inherit',
+                }}
+              />
             </Box>
           )}
 
diff --git a/examples/core-connection/src/modules/simulatePage/components/MockSendTransaction.tsx b/examples/core-connection/src/modules/simulatePage/components/MockSendTransaction.tsx
index 77d9b822..62bc1b7c 100644
--- a/examples/core-connection/src/modules/simulatePage/components/MockSendTransaction.tsx
+++ b/examples/core-connection/src/modules/simulatePage/components/MockSendTransaction.tsx
@@ -4,10 +4,14 @@ import { css } from 'styled-components';
 import { TransactionSnippet } from '../../../common/components';
 import { useGlobalContext } from '../../../context/GlobalContext';
 import { Transaction } from '@pushprotocol/push-chain/src/lib/generated/tx';
-import { centerMaskString } from '../../../helpers';
+import { centerMaskString, convertToCaip } from '../../../helpers';
+import { usePushWalletContext } from '@pushprotocol/pushchain-ui-kit';
+
 const MockSendTransaction = () => {
-  const { pushNetwork, mockTx, account, handleSendSignRequestToPushWallet } =
-    useGlobalContext();
+  const { pushNetwork, mockTx } = useGlobalContext();
+
+  const { handleSendSignRequestToPushWallet, universalAddress } =
+    usePushWalletContext();
 
   const [isSendingTxn, setIsSendingTxn] = useState(false);
   const [txnHash, setTxnHash] = useState<string | null>(null);
@@ -15,10 +19,10 @@ const MockSendTransaction = () => {
 
   const handleSendTransaction = async (mockTx: Transaction) => {
     try {
-      if (pushNetwork && account) {
+      if (pushNetwork && universalAddress) {
         setIsSendingTxn(true);
         const txHash = await pushNetwork.tx.send(mockTx, {
-          account,
+          account: convertToCaip(universalAddress), // This will get changed to utils usage in pushchain devnet package
           signMessage: async (data: Uint8Array) => {
             return await handleSendSignRequestToPushWallet(data);
           },
diff --git a/examples/core-connection/src/modules/simulatePage/components/SimulateHeader.tsx b/examples/core-connection/src/modules/simulatePage/components/SimulateHeader.tsx
index a2a0a748..2ba3ca44 100644
--- a/examples/core-connection/src/modules/simulatePage/components/SimulateHeader.tsx
+++ b/examples/core-connection/src/modules/simulatePage/components/SimulateHeader.tsx
@@ -1,13 +1,15 @@
 import { Box } from 'shared-components';
-import { useGlobalContext } from '../../../context/GlobalContext';
-import { TogglePushWalletButton } from '@pushprotocol/pushchain-ui-kit';
+import {
+  PushWalletButton,
+  usePushWalletContext,
+} from '@pushprotocol/pushchain-ui-kit';
 
 const SimulateHeader = () => {
-  const { account } = useGlobalContext();
+  const { universalAddress } = usePushWalletContext();
 
   return (
     <Box display="flex" justifyContent="end" width="100%">
-      {account && <TogglePushWalletButton account={account} />}
+      <PushWalletButton universalAddress={universalAddress} />
     </Box>
   );
 };
diff --git a/examples/core-connection/yarn.lock b/examples/core-connection/yarn.lock
index dd372d31..2b906a5d 100644
--- a/examples/core-connection/yarn.lock
+++ b/examples/core-connection/yarn.lock
@@ -6822,7 +6822,7 @@ __metadata:
 
 "shared-components@file:../../packages/shared-components::locator=core-connection%40workspace%3A.":
   version: 0.1.0
-  resolution: "shared-components@file:../../packages/shared-components#../../packages/shared-components::hash=7ec6fd&locator=core-connection%40workspace%3A."
+  resolution: "shared-components@file:../../packages/shared-components#../../packages/shared-components::hash=9fe31a&locator=core-connection%40workspace%3A."
   dependencies:
     "@emotion/react": "npm:^11.13.0"
     "@radix-ui/react-dialog": "npm:^1.1.1"
@@ -6843,7 +6843,7 @@ __metadata:
   peerDependencies:
     react: ^18.3.1
     react-dom: ^18.3.1
-  checksum: 10/043be10cccbe754a00ed3a41688d8cdbca487e56abe72b78aded4e7617ed723c5e08314960c3b410356b57475d28ac0162119ffff0c7516f0e465d97bc5f3d37
+  checksum: 10/7d96c35a77cdbd873e64bb718d9fa8f44607528009e0586c7100f24b6a77d47ba5a89b2ee715086738fc39e6413a3880a32f1e32e53adedd86b586cce29fcbda
   languageName: node
   linkType: hard
 
diff --git a/packages/ui-kit/src/lib/common/spinner/Spinner.tsx b/packages/ui-kit/src/lib/common/spinner/Spinner.tsx
index d920871a..db3b79ff 100644
--- a/packages/ui-kit/src/lib/common/spinner/Spinner.tsx
+++ b/packages/ui-kit/src/lib/common/spinner/Spinner.tsx
@@ -18,18 +18,18 @@ const Container = styled.div`
   align-items: center;
   justify-content: center;
   animation: ${spin} 1s linear infinite;
-  width: 24px;
-  height: 24px;
-  color:#D548EC;
+  width: 20px;
+  height: 20px;
+  color: #fff;
 `;
 
 // Spinner functional component
 const Spinner = () => {
-    return (
-        <Container>
-            <Ellipse />
-        </Container>
-    );
+  return (
+    <Container>
+      <Ellipse />
+    </Container>
+  );
 };
 
 export default Spinner;
diff --git a/packages/ui-kit/src/lib/wallet/components/ConnectPushWalletButton.tsx b/packages/ui-kit/src/lib/wallet/components/ConnectPushWalletButton.tsx
index 6f13e5c7..e6eab5b9 100644
--- a/packages/ui-kit/src/lib/wallet/components/ConnectPushWalletButton.tsx
+++ b/packages/ui-kit/src/lib/wallet/components/ConnectPushWalletButton.tsx
@@ -1,14 +1,18 @@
-import { FC } from 'react';
+import React, { FC } from 'react';
 import { usePushWalletContext } from './PushWalletProvider';
-import { walletConnectionButtonStatusMapper } from '../wallet.constants';
-import styled from 'styled-components'
+import styled from 'styled-components';
 import { Spinner } from '../../common';
 
 export type ConnectPushWalletButtonProps = {
   showLogOutButton?: boolean;
+  title?: string;
+  styling?: React.CSSProperties;
 };
 
-const ConnectPushWalletButton: FC<ConnectPushWalletButtonProps> = () => {
+const ConnectPushWalletButton: FC<ConnectPushWalletButtonProps> = ({
+  title,
+  styling,
+}) => {
   const {
     connectionStatus,
     handleConnectToPushWallet,
@@ -34,9 +38,14 @@ const ConnectPushWalletButton: FC<ConnectPushWalletButtonProps> = () => {
       <ConnectButton
         onClick={handleConnectWalletButton}
         disabled={isConnectButtonDisbaled || isLoading}
+        customStyle={styling}
       >
-        {walletConnectionButtonStatusMapper[connectionStatus]}
-        {isLoading && (<SpinnerContainer><Spinner /></SpinnerContainer>)}
+        {connectionStatus === 'notConnected' ? title : connectionStatus}
+        {isLoading && (
+          <SpinnerContainer>
+            <Spinner />
+          </SpinnerContainer>
+        )}
       </ConnectButton>
     </>
   );
@@ -44,31 +53,39 @@ const ConnectPushWalletButton: FC<ConnectPushWalletButtonProps> = () => {
 
 export { ConnectPushWalletButton };
 
-const ConnectButton = styled.button`
-    align-items: center;
-    cursor: pointer;
-    display: flex;
-    justify-content: center;
-    white-space: nowrap;
-    flex-shrink: 0;
-    border: none;
-    background-color: #D548EC;
-    color: rgba(255,255,255,1);
-    border-radius: 12px;
-    gap: 4px;
-    height: 48px;
-    padding: 16px 24px;
-    min-width: 100px;
-    leading-trim: both;
-    text-edge: cap;
-    font-family:FK Grotesk Neu;
-    font-size: 16px;
-    font-style: normal;
-    font-weight: 500;
-    line-height: 16px;
-    width:inherit;
+const ConnectButton = styled.button<{ customStyle?: React.CSSProperties }>`
+  align-items: center;
+  cursor: pointer;
+  display: flex;
+  justify-content: center;
+  white-space: nowrap;
+  flex-shrink: 0;
+  border: none;
+  background-color: #d548ec;
+  color: rgba(255, 255, 255, 1);
+  border-radius: 12px;
+  gap: 4px;
+  height: 48px;
+  padding: 16px 24px;
+  min-width: 100px;
+  leading-trim: both;
+  text-edge: cap;
+  font-family: FK Grotesk Neu;
+  font-size: 16px;
+  font-style: normal;
+  font-weight: 500;
+  line-height: 16px;
+  width: inherit;
 
-`
+  ${(props) =>
+    props.customStyle &&
+    Object.entries(props.customStyle)
+      .map(
+        ([key, value]) =>
+          `${key.replace(/([A-Z])/g, '-$1').toLowerCase()}: ${value};`
+      )
+      .join('\n')}
+`;
 const SpinnerContainer = styled.div`
-  padding:5px;
-`
\ No newline at end of file
+  padding: 4px;
+`;
diff --git a/packages/ui-kit/src/lib/wallet/components/PushWalletButton.tsx b/packages/ui-kit/src/lib/wallet/components/PushWalletButton.tsx
new file mode 100644
index 00000000..f16e0c41
--- /dev/null
+++ b/packages/ui-kit/src/lib/wallet/components/PushWalletButton.tsx
@@ -0,0 +1,28 @@
+import React, { ReactNode } from 'react';
+import { ConnectPushWalletButton } from './ConnectPushWalletButton';
+import { TogglePushWalletButton } from './TogglePushWalletButton';
+import { UniversalAddress } from '../wallet.types';
+
+type PushWalletButtonProps = {
+  universalAddress: UniversalAddress | null;
+  component?: ReactNode;
+  title?: string;
+  styling?: React.CSSProperties;
+};
+
+const PushWalletButton: React.FC<PushWalletButtonProps> = ({
+  universalAddress,
+  component,
+  title = 'Login',
+  styling,
+}) => {
+  // If universal address is present render the button
+  if (universalAddress) {
+    return <TogglePushWalletButton universalAddress={universalAddress} />;
+  } else if (component) {
+    // If no UA and custom component, then render the component
+    return <>{component}</>;
+  } else return <ConnectPushWalletButton title={title} styling={styling} />; // If no UA and no custom component, then render the connect button
+};
+
+export { PushWalletButton };
diff --git a/packages/ui-kit/src/lib/wallet/components/PushWalletIFrame.tsx b/packages/ui-kit/src/lib/wallet/components/PushWalletIFrame.tsx
deleted file mode 100644
index 2ca32cc2..00000000
--- a/packages/ui-kit/src/lib/wallet/components/PushWalletIFrame.tsx
+++ /dev/null
@@ -1,174 +0,0 @@
-import { FC } from 'react';
-import { usePushWalletContext } from './PushWalletProvider';
-import config from '../../config';
-import styled from 'styled-components';
-import { CrossIcon, DashIcon, Spinner } from '../../common';
-
-const PushWalletIFrame: FC = () => {
-  const {
-    env,
-    account,
-    iframeRef,
-    isWalletMinimised,
-    isWalletVisible,
-    setMinimiseWallet,
-    handleUserLogOutEvent,
-    isIframeLoading,
-    setIframeLoading,
-  } = usePushWalletContext();
-
-  return (
-    <>
-      {isWalletVisible ? (
-        <FrameContainer isWalletMinimised={isWalletMinimised} account={account}>
-          {isIframeLoading && (
-            <FrameLoadingContainer>
-              <CloseButtonContainer
-                onClick={() => {
-                  handleUserLogOutEvent();
-                }}
-              >
-                <CrossIcon />
-              </CloseButtonContainer>
-              <LoadingTextContainer>
-                <LoadingText>Loading...</LoadingText>
-                <Spinner />
-              </LoadingTextContainer>
-            </FrameLoadingContainer>
-          )}
-
-          <FrameSubContainer
-            isWalletMinimised={isWalletMinimised}
-            isIframeLoading={isIframeLoading}
-          >
-            <AccountContainer account={account}>
-              {account ? (
-                <DashButtonContainer onClick={() => setMinimiseWallet(true)}>
-                  <DashIcon />
-                </DashButtonContainer>
-              ) : (
-                <CloseButtonContainer
-                  onClick={() => {
-                    handleUserLogOutEvent();
-                  }}
-                >
-                  <CrossIcon />
-                </CloseButtonContainer>
-              )}
-            </AccountContainer>
-            <iframe
-              src={`${config.WALLET_URL[env]}/auth?app=${window.location.origin}`}
-              allow="publickey-credentials-create; publickey-credentials-get"
-              ref={iframeRef}
-              style={{
-                border: 'none',
-                width: '-webkit-fill-available',
-                height: '100vh',
-                borderBottomRightRadius: account ? '10px' : '0px',
-                borderBottomLeftRadius: account ? '10px' : '0px',
-              }}
-              onLoad={() => setIframeLoading(false)}
-            />
-          </FrameSubContainer>
-        </FrameContainer>
-      ) : null}
-    </>
-  );
-};
-
-export { PushWalletIFrame };
-
-const FrameContainer = styled.div<{
-  account: string | null;
-  isWalletMinimised: boolean;
-}>`
-  position: fixed;
-  display: flex;
-  flex-direction: column;
-  background-color: #17181b;
-  border-radius: 10px;
-  z-index: 99;
-
-  width: ${({ account, isWalletMinimised }) =>
-    isWalletMinimised ? '0px' : account ? '450px' : '100%'};
-  height: ${({ account, isWalletMinimised }) =>
-    isWalletMinimised ? '0px' : account ? '710px' : '100%'};
-  right: ${({ account }) => (account ? '24px' : '0')};
-  top: ${({ account }) => (account ? '24px' : '0')};
-
-  @media (max-width: 425px) {
-    width: ${({ account, isWalletMinimised }) =>
-      isWalletMinimised ? '0px' : account ? '96%' : '100%'};
-    right: ${({ account }) => (account ? '2%' : '0')};
-    top: ${({ account }) => (account ? '8%' : '0')};
-  }
-`;
-
-const CloseButtonContainer = styled.div`
-  display: flex;
-  align-items: center;
-  justify-content: flex-end;
-  cursor: pointer;
-  padding: 0 16px;
-`;
-
-const DashButtonContainer = styled.div`
-  display: flex;
-  align-items: center;
-  justify-content: center;
-  cursor: pointer;
-  border-radius: 1000px;
-  height: 14px;
-  width: 14px;
-  background-color: #ffbb16;
-`;
-
-const LoadingTextContainer = styled.div`
-  display: flex;
-  align-items: center;
-  justify-content: center;
-  gap: 16px;
-  width: -webkit-fill-available;
-  height: -webkit-fill-available;
-`;
-
-const LoadingText = styled.p`
-  font-size: 18px;
-  font-weight: 500;
-  line-height: 27px;
-  color: rgba(255, 255, 255, 1);
-  font-family: FK Grotesk Neu;
-  margin: 0px;
-  width: auto;
-`;
-
-const FrameLoadingContainer = styled.div`
-  height: -webkit-fill-available;
-  width: -webkit-fill-available;
-  flex-direction: column;
-  display: flex;
-  padding: 8px;
-  background-color: #17181b;
-`;
-
-const FrameSubContainer = styled.div<{
-  isWalletMinimised: boolean;
-  isIframeLoading: boolean;
-}>`
-  display: ${({ isWalletMinimised, isIframeLoading }) =>
-    isWalletMinimised || isIframeLoading ? 'none' : 'flex'};
-  width: -webkit-fill-available;
-  height: -webkit-fill-available;
-  flex-direction: column;
-`;
-
-const AccountContainer = styled.div<{ account: string | null }>`
-  width: -webkit-fill-available;
-  display: flex;
-  align-items: center;
-  justify-content: flex-end;
-  padding: var(--spacing-xxs) var(--spacing-xxs);
-  border-top-right-radius: ${({ account }) => (account ? '10px' : '0px')};
-  border-top-left-radius: ${({ account }) => (account ? '10px' : '0px')};
-  background-color: ${({ account }) => (account ? '#e3e3e3' : '#17181B')};
-`;
diff --git a/packages/ui-kit/src/lib/wallet/components/PushWalletProvider.tsx b/packages/ui-kit/src/lib/wallet/components/PushWalletProvider.tsx
index c2442349..726c5f21 100644
--- a/packages/ui-kit/src/lib/wallet/components/PushWalletProvider.tsx
+++ b/packages/ui-kit/src/lib/wallet/components/PushWalletProvider.tsx
@@ -7,30 +7,23 @@ import React, {
   useRef,
 } from 'react';
 import {
-  APP_TO_WALLET_ACTION,
   ConnectionStatus,
-  WALLET_TO_APP_ACTION,
+  UniversalAddress,
   WalletEventRespoonse,
 } from '../wallet.types';
+import { CONSTANTS } from '../../constants';
 import config, { ENV } from '../../config';
+import { getWalletDataFromAccount } from '../wallet.utils';
 
 // Define the context shape
 export type PushWalletContextType = {
-  account: string | null;
+  universalAddress: UniversalAddress | null;
   connectionStatus: ConnectionStatus;
   env: ENV;
-  iframeRef: React.MutableRefObject<HTMLIFrameElement | null>;
-  isWalletVisible: boolean;
-  isWalletMinimised: boolean;
-  setWalletVisibility: (isWalletVisible: boolean) => void;
   handleConnectToPushWallet: () => void;
   handleNewConnectionRequest: () => void;
-  handleSendSignRequestToPushWallet: (data: Uint8Array) => Promise<Uint8Array>;
-  setMinimiseWallet: React.Dispatch<React.SetStateAction<boolean>>;
-  handleUserLogOutEvent: () => void;
   handleLogOut: () => void;
-  isIframeLoading: boolean;
-  setIframeLoading: React.Dispatch<React.SetStateAction<boolean>>;
+  handleSendSignRequestToPushWallet: (data: Uint8Array) => Promise<Uint8Array>;
 };
 
 export type WalletProviderProps = { children: ReactNode; env: ENV };
@@ -45,19 +38,12 @@ export const PushWalletProvider: React.FC<WalletProviderProps> = ({
   children,
   env,
 }) => {
-  const [account, setAccount] =
-    useState<PushWalletContextType['account']>(null);
-
-  const [isWalletVisible, setWalletVisibility] = useState(false);
-
-  const [isWalletMinimised, setMinimiseWallet] = useState(false);
-
-  const [isIframeLoading, setIframeLoading] = useState(true);
+  const [universalAddress, setUniversalAddress] =
+    useState<PushWalletContextType['universalAddress']>(null);
 
   const [connectionStatus, setConnectionStatus] =
     useState<ConnectionStatus>('notConnected');
 
-  const iframeRef = useRef<HTMLIFrameElement | null>(null);
   const tabRef = useRef<Window | null>(null);
 
   const signatureResolverRef = useRef<{
@@ -84,7 +70,6 @@ export const PushWalletProvider: React.FC<WalletProviderProps> = ({
 
   const handleConnectToPushWallet = () => {
     handleOpenNewWalletTab();
-    setWalletVisibility(true);
     setConnectionStatus('connecting');
   };
 
@@ -99,18 +84,23 @@ export const PushWalletProvider: React.FC<WalletProviderProps> = ({
   };
 
   const handleNewConnectionRequest = () => {
-    setMinimiseWallet(false);
     setConnectionStatus('authenticating');
     sendMessageToPushWallet({
-      type: APP_TO_WALLET_ACTION.NEW_CONNECTION_REQUEST,
+      type: CONSTANTS.APP_TO_WALLET_ACTION.NEW_CONNECTION_REQUEST,
     });
   };
 
   const handleIsLoggedInAction = (response: WalletEventRespoonse) => {
     if (response?.account) {
       setConnectionStatus('connected');
-      setMinimiseWallet(true);
-      setAccount(response.account);
+
+      const result = getWalletDataFromAccount(response.account);
+
+      setUniversalAddress({
+        chainId: result.chainId,
+        chain: result.chain,
+        address: result.address,
+      });
     } else {
       handleNewConnectionRequest();
     }
@@ -118,25 +108,28 @@ export const PushWalletProvider: React.FC<WalletProviderProps> = ({
 
   const handleAppConnectionSuccess = (response: WalletEventRespoonse) => {
     setConnectionStatus('connected');
-    setMinimiseWallet(true);
-    setAccount(response.account!);
+    if (response.account) {
+      const result = getWalletDataFromAccount(response.account);
+      setUniversalAddress({
+        chainId: result.chainId,
+        chain: result.chain,
+        address: result.address,
+      });
+    }
   };
 
   const handleAppConnectionRejection = () => {
     setConnectionStatus('retry');
-    setAccount(null);
+    setUniversalAddress(null);
   };
 
   const handleAppConnectionRetry = () => {
-    setMinimiseWallet(true);
+    // setMinimiseWallet(true);
   };
 
   const handleUserLogOutEvent = () => {
     setConnectionStatus('notConnected');
-    setAccount(null);
-    setMinimiseWallet(false);
-    setWalletVisibility(false);
-    setIframeLoading(true);
+    setUniversalAddress(null);
   };
 
   const handleSendSignRequestToPushWallet = (
@@ -159,11 +152,9 @@ export const PushWalletProvider: React.FC<WalletProviderProps> = ({
         },
       };
 
-      setMinimiseWallet(false);
-
       // Send the sign request to the wallet tab
       sendMessageToPushWallet({
-        type: APP_TO_WALLET_ACTION.SIGN_MESSAGE,
+        type: CONSTANTS.APP_TO_WALLET_ACTION.SIGN_MESSAGE,
         data,
       });
     });
@@ -171,7 +162,7 @@ export const PushWalletProvider: React.FC<WalletProviderProps> = ({
 
   const handleLogOut = async () => {
     sendMessageToPushWallet({
-      type: APP_TO_WALLET_ACTION.LOG_OUT,
+      type: CONSTANTS.APP_TO_WALLET_ACTION.LOG_OUT,
       data: 'Log Out Event',
     });
   };
@@ -179,27 +170,27 @@ export const PushWalletProvider: React.FC<WalletProviderProps> = ({
   useEffect(() => {
     const messageHandler = (event: MessageEvent) => {
       switch (event.data.type) {
-        case WALLET_TO_APP_ACTION.IS_LOGGED_IN:
+        case CONSTANTS.WALLET_TO_APP_ACTION.IS_LOGGED_IN:
           handleIsLoggedInAction(event.data.data);
           break;
-        case WALLET_TO_APP_ACTION.APP_CONNECTION_SUCCESS:
+        case CONSTANTS.WALLET_TO_APP_ACTION.APP_CONNECTION_SUCCESS:
           handleAppConnectionSuccess(event.data.data);
           break;
-        case WALLET_TO_APP_ACTION.APP_CONNECTION_REJECTED:
+        case CONSTANTS.WALLET_TO_APP_ACTION.APP_CONNECTION_REJECTED:
           handleAppConnectionRejection();
           break;
-        case WALLET_TO_APP_ACTION.APP_CONNECTION_RETRY:
+        case CONSTANTS.WALLET_TO_APP_ACTION.APP_CONNECTION_RETRY:
           handleAppConnectionRetry();
           break;
-        case WALLET_TO_APP_ACTION.SIGNATURE:
+        case CONSTANTS.WALLET_TO_APP_ACTION.SIGNATURE:
           if (signatureResolverRef.current) {
             signatureResolverRef?.current?.success?.(event.data.data);
           }
           break;
-        case WALLET_TO_APP_ACTION.IS_LOGGED_OUT:
+        case CONSTANTS.WALLET_TO_APP_ACTION.IS_LOGGED_OUT:
           handleUserLogOutEvent();
           break;
-        case WALLET_TO_APP_ACTION.ERROR:
+        case CONSTANTS.WALLET_TO_APP_ACTION.ERROR:
           signatureResolverRef?.current?.error?.(event.data.data);
           break;
         default:
@@ -215,20 +206,12 @@ export const PushWalletProvider: React.FC<WalletProviderProps> = ({
   return (
     <PushWalletContext.Provider
       value={{
-        account,
+        universalAddress,
         connectionStatus,
         env,
-        iframeRef,
-        isWalletVisible,
-        setWalletVisibility,
-        isWalletMinimised,
-        setMinimiseWallet,
-        isIframeLoading,
-        setIframeLoading,
         handleConnectToPushWallet,
         handleNewConnectionRequest,
         handleSendSignRequestToPushWallet,
-        handleUserLogOutEvent,
         handleLogOut,
       }}
     >
diff --git a/packages/ui-kit/src/lib/wallet/components/TogglePushWalletButton.tsx b/packages/ui-kit/src/lib/wallet/components/TogglePushWalletButton.tsx
index 2c2711bc..676d441e 100644
--- a/packages/ui-kit/src/lib/wallet/components/TogglePushWalletButton.tsx
+++ b/packages/ui-kit/src/lib/wallet/components/TogglePushWalletButton.tsx
@@ -1,18 +1,19 @@
 import React, { useState } from 'react';
-import { centerMaskString, getWalletDataFromAccount } from '../wallet.utils';
+import { centerMaskString } from '../wallet.utils';
 import { CHAIN_LOGO } from '../../constants';
 import { usePushWalletContext } from './PushWalletProvider';
 import { PushLogo, PushMonotone } from '../../common';
 import styled from 'styled-components';
+import { UniversalAddress } from '../wallet.types';
 
 type TogglePushWalletButtonProps = {
-  account: string;
+  universalAddress: UniversalAddress;
 };
 const TogglePushWalletButton: React.FC<TogglePushWalletButtonProps> = ({
-  account,
+  universalAddress,
 }) => {
-  const { setMinimiseWallet, handleLogOut } = usePushWalletContext();
-  const { chainId, address } = getWalletDataFromAccount(account);
+  const { handleLogOut } = usePushWalletContext();
+  const { chainId, address } = universalAddress;
 
   const [isDropdownOpen, setDropdownOpen] = useState(false);
 
@@ -102,14 +103,3 @@ const DropdownMenu = styled.div`
   background-color: #17181b;
   color: rgba(255, 255, 255, 1);
 `;
-
-const DropdownItem = styled.button`
-  padding: 8px 16px;
-  cursor: pointer;
-  background: none;
-  border: none;
-  width: 100%;
-  text-align: left;
-  color: rgba(255, 255, 255, 1);
-  height: 48px;
-`;
diff --git a/packages/ui-kit/src/lib/wallet/index.ts b/packages/ui-kit/src/lib/wallet/index.ts
index 99566def..e4f36e8d 100644
--- a/packages/ui-kit/src/lib/wallet/index.ts
+++ b/packages/ui-kit/src/lib/wallet/index.ts
@@ -1,6 +1,7 @@
 export * from './components/PushWalletProvider';
 export * from './components/ConnectPushWalletButton';
 export * from './components/TogglePushWalletButton';
+export * from './components/PushWalletButton';
 
 export * from './wallet.types';
 export * from './wallet.constants';
diff --git a/packages/ui-kit/src/lib/wallet/wallet.types.ts b/packages/ui-kit/src/lib/wallet/wallet.types.ts
index 0cc854b9..cf16919e 100644
--- a/packages/ui-kit/src/lib/wallet/wallet.types.ts
+++ b/packages/ui-kit/src/lib/wallet/wallet.types.ts
@@ -26,6 +26,12 @@ export type ConnectionStatus =
   | 'connected'
   | 'retry';
 
+export type UniversalAddress = {
+  chainId: string;
+  chain: string;
+  address: string;
+};
+
 export type WalletEventRespoonse = {
   signature?: Uint8Array;
   account?: string;
diff --git a/packages/ui-kit/src/lib/wallet/wallet.utils.ts b/packages/ui-kit/src/lib/wallet/wallet.utils.ts
index 7afde79d..ede31650 100644
--- a/packages/ui-kit/src/lib/wallet/wallet.utils.ts
+++ b/packages/ui-kit/src/lib/wallet/wallet.utils.ts
@@ -1,19 +1,10 @@
 export const getWalletDataFromAccount = (
   account: string
 ): {
-  chainId: string | null;
-  chain: string | null;
+  chainId: string;
+  chain: string;
   address: string;
 } => {
-  // // Check if the input is a valid non-empty string
-  // if (!account || typeof account !== 'string') {
-  //   return {
-  //     chain: null,
-  //     chainId: null,
-  //     address: null,
-  //   };
-  // }
-
   const addressComponent = account.split(':');
 
   // Handle cases where there are exactly three components (chain, chainId, address)
@@ -28,15 +19,15 @@ export const getWalletDataFromAccount = (
   else if (addressComponent.length === 2) {
     return {
       chain: addressComponent[0],
-      chainId: null,
+      chainId: '',
       address: addressComponent[1],
     };
   }
   // If the input doesn't match the expected format, return the address only
   else {
     return {
-      chain: null,
-      chainId: null,
+      chain: '',
+      chainId: '',
       address: account,
     };
   }