From 5a7586eb61877d0e958030fe041d7ca22df0addc Mon Sep 17 00:00:00 2001
From: Emanuele Cesena <emanuele.cesena@gmail.com>
Date: Sun, 6 Oct 2024 19:09:16 -0700
Subject: [PATCH] flows: refresh balance after successful tx

---
 anchor/src/react/glam.tsx                | 8 +++++++-
 playground/src/app/flows/page.tsx        | 8 ++++++--
 playground/src/components/AssetInput.tsx | 6 ++++--
 3 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/anchor/src/react/glam.tsx b/anchor/src/react/glam.tsx
index 6400b81c..842aac52 100644
--- a/anchor/src/react/glam.tsx
+++ b/anchor/src/react/glam.tsx
@@ -31,8 +31,10 @@ interface GlamProviderContext {
   fund?: PublicKey;
   treasury?: FundCacheTreasury;
   fundsList: FundCache[];
+  //@ts-ignore
   allFunds: FundModel[];
   walletBalances: any;
+  walletBalancesQueryKey: any[];
   refresh?: () => void;
   setActiveFund: any;
   tokenList?: TokenListItem[];
@@ -175,10 +177,13 @@ export function GlamProvider({
     );
   }, [allFundsData]);
 
+  const walletBalancesQueryKey = ["balances", wallet?.publicKey];
   const { data: walletBalances } = useQuery({
-    queryKey: ["balances", wallet?.publicKey],
+    queryKey: walletBalancesQueryKey,
     enabled: !!wallet?.publicKey,
+    // refetchInterval: 2000, // 2s
     queryFn: async () => {
+      console.log("fetching walletBalances");
       const balanceLamports = await glamClient.provider.connection.getBalance(
         wallet?.publicKey || new PublicKey(0)
       );
@@ -228,6 +233,7 @@ export function GlamProvider({
     fundsList: useAtomValue(fundsListAtom),
     allFunds,
     walletBalances,
+    walletBalancesQueryKey,
     tokenList,
     setActiveFund,
   };
diff --git a/playground/src/app/flows/page.tsx b/playground/src/app/flows/page.tsx
index b10cc661..71c55102 100644
--- a/playground/src/app/flows/page.tsx
+++ b/playground/src/app/flows/page.tsx
@@ -4,7 +4,7 @@ import { useMemo } from "react";
 import { zodResolver } from "@hookform/resolvers/zod";
 import { useForm, SubmitHandler, FormProvider } from "react-hook-form";
 import { z } from "zod";
-import { useQuery } from "@tanstack/react-query";
+import { useQuery, useQueryClient } from "@tanstack/react-query";
 
 import { Button } from "@/components/ui/button";
 import {
@@ -372,7 +372,10 @@ function InvestorDisclaimers({
 }
 
 function InvestorWidget({ fundId }: { fundId: string }) {
-  const { glamClient, allFunds, walletBalances } = useGlam();
+  const { glamClient, allFunds, walletBalances, walletBalancesQueryKey } =
+    useGlam();
+  const queryClient = useQueryClient();
+
   const fund: any = fundId
     ? (allFunds || []).find((f: any) => f.idStr === fundId)
     : undefined;
@@ -508,6 +511,7 @@ function InvestorWidget({ fundId }: { fundId: string }) {
       txId = await glamClient.investor.subscribe(fund.id, asset, amount);
     }
 
+    queryClient.invalidateQueries({ queryKey: walletBalancesQueryKey });
     toast({
       title: `Successful ${direction}`,
       description: <ExplorerLink path={`tx/${txId}`} label={txId} />,
diff --git a/playground/src/components/AssetInput.tsx b/playground/src/components/AssetInput.tsx
index 57a7764c..a426128f 100644
--- a/playground/src/components/AssetInput.tsx
+++ b/playground/src/components/AssetInput.tsx
@@ -117,7 +117,7 @@ export const AssetInput: React.FC<AssetInputProps> = ({
       if (!isNaN(numericValue)) {
         setValue(name, numericValue);
       } else {
-        setValue(name, 0);
+        setValue(name, "");
       }
     }
   };
@@ -145,10 +145,12 @@ export const AssetInput: React.FC<AssetInputProps> = ({
             <div className="relative">
               <Input
                 {...field}
+                type="number"
+                step="any"
                 ref={inputRef}
                 value={getValues()[name]}
                 className="pr-20"
-                placeholder="0"
+                placeholder=""
                 onChange={(e) => handleInputChange(e.target.value)}
                 disabled={disableAmountInput}
               />