From b77acb01623f9cd227190d4ebc1ad561e10dc4e5 Mon Sep 17 00:00:00 2001
From: Chase Fleming <1666730+chasefleming@users.noreply.github.com>
Date: Mon, 16 Jun 2025 11:23:16 -0700
Subject: [PATCH 1/2] Add token balance to profile

---
 packages/kit/src/components/Connect.tsx       | 23 ++++++++++++++++++-
 .../kit/src/hooks/useCrossVmTokenBalance.ts   |  2 +-
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/packages/kit/src/components/Connect.tsx b/packages/kit/src/components/Connect.tsx
index ff23d0429..9a0e9d1c6 100644
--- a/packages/kit/src/components/Connect.tsx
+++ b/packages/kit/src/components/Connect.tsx
@@ -1,31 +1,52 @@
 import React, {useState} from "react"
 import {useCurrentFlowUser} from "../hooks"
+import {
+  useCrossVmTokenBalance,
+  UseCrossVmTokenBalanceData,
+} from "../hooks/useCrossVmTokenBalance"
 import {Button, ButtonProps} from "./internal/Button"
 import {Dialog} from "./internal/Dialog"
 import {CircleUserRoundIcon} from "../icons/CircleUserRoundIcon"
 import {CopyIcon} from "../icons/CopyIcon"
 import {LogOutIcon} from "../icons/LogOutIcon"
 
+type BalanceType = keyof UseCrossVmTokenBalanceData
+
 interface ConnectProps {
   variant?: ButtonProps["variant"]
   onConnect?: () => void
   onDisconnect?: () => void
+  balanceType?: BalanceType
 }
 
 export const Connect: React.FC<ConnectProps> = ({
   variant = "primary",
   onConnect,
   onDisconnect,
+  balanceType = "cadence",
 }) => {
   const {user, authenticate, unauthenticate} = useCurrentFlowUser()
   const [open, setOpen] = useState(false)
   const [copied, setCopied] = useState(false)
 
+  const {data: balanceData} = useCrossVmTokenBalance({
+    owner: user?.addr,
+    vaultIdentifier: user?.addr ? `${user.addr}.FlowToken.Vault` : undefined,
+    query: {
+      enabled: !!user?.addr,
+    },
+  })
+
   const displayAddress =
     user?.loggedIn && user.addr
       ? `${user.addr.slice(0, 6)}...${user.addr.slice(-4)}`
       : ""
 
+  const displayBalance =
+    balanceData && typeof balanceData !== "string"
+      ? `${balanceData[balanceType].formatted} FLOW`
+      : "0.00 FLOW"
+
   const handleButtonClick = async () => {
     if (user?.loggedIn) {
       setOpen(true)
@@ -73,7 +94,7 @@ export const Connect: React.FC<ConnectProps> = ({
                 {displayAddress}
               </div>
               <div className="text-center text-sm text-gray-500 -mt-1">
-                1,416.00 FLOW
+                {displayBalance}
               </div>
             </div>
             <div className="flex gap-2 w-full">
diff --git a/packages/kit/src/hooks/useCrossVmTokenBalance.ts b/packages/kit/src/hooks/useCrossVmTokenBalance.ts
index 457f54f4e..11b22eefa 100644
--- a/packages/kit/src/hooks/useCrossVmTokenBalance.ts
+++ b/packages/kit/src/hooks/useCrossVmTokenBalance.ts
@@ -17,7 +17,7 @@ interface TokenBalance {
   precision: number
 }
 
-interface UseCrossVmTokenBalanceData {
+export interface UseCrossVmTokenBalanceData {
   cadence: TokenBalance
   evm: TokenBalance
   combined: TokenBalance

From eba757745554630ab1abd739b4d1ef6d5fb4b810 Mon Sep 17 00:00:00 2001
From: Chase Fleming <1666730+chasefleming@users.noreply.github.com>
Date: Mon, 16 Jun 2025 11:41:06 -0700
Subject: [PATCH 2/2] Fix vault identifier

---
 packages/kit/src/components/Connect.tsx | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/packages/kit/src/components/Connect.tsx b/packages/kit/src/components/Connect.tsx
index 9a0e9d1c6..dd6641305 100644
--- a/packages/kit/src/components/Connect.tsx
+++ b/packages/kit/src/components/Connect.tsx
@@ -4,6 +4,7 @@ import {
   useCrossVmTokenBalance,
   UseCrossVmTokenBalanceData,
 } from "../hooks/useCrossVmTokenBalance"
+import {useFlowChainId} from "../hooks/useFlowChainId"
 import {Button, ButtonProps} from "./internal/Button"
 import {Dialog} from "./internal/Dialog"
 import {CircleUserRoundIcon} from "../icons/CircleUserRoundIcon"
@@ -28,12 +29,13 @@ export const Connect: React.FC<ConnectProps> = ({
   const {user, authenticate, unauthenticate} = useCurrentFlowUser()
   const [open, setOpen] = useState(false)
   const [copied, setCopied] = useState(false)
+  const {data: chainId} = useFlowChainId()
 
   const {data: balanceData} = useCrossVmTokenBalance({
     owner: user?.addr,
-    vaultIdentifier: user?.addr ? `${user.addr}.FlowToken.Vault` : undefined,
+    vaultIdentifier: `A.${chainId === "testnet" ? "7e60df042a9c0868" : "1654653399040a61"}.FlowToken.Vault`,
     query: {
-      enabled: !!user?.addr,
+      enabled: !!user?.addr && !!chainId,
     },
   })