pro: show remaining credits (#329)

Fixes #265

src/hooks/useStreamChat.ts

@@ -19,6 +19,7 @@ import { useProposal } from "./useProposal";
 import { useSearch } from "@tanstack/react-router";
 import { useRunApp } from "./useRunApp";
 import { useCountTokens } from "./useCountTokens";
+import { useUserBudgetInfo } from "./useUserBudgetInfo";
 
 export function getRandomNumberId() {
   return Math.floor(Math.random() * 1_000_000_000_000_000);
@@ -38,6 +39,7 @@ export function useStreamChat({
   const { refreshVersions } = useVersions(selectedAppId);
   const { refreshAppIframe } = useRunApp();
   const { countTokens } = useCountTokens();
+  const { refetchUserBudget } = useUserBudgetInfo();
 
   let chatId: number | undefined;
 
@@ -95,6 +97,8 @@ export function useStreamChat({
             }
             refreshProposal(chatId);
 
+            refetchUserBudget();
+
             // Keep the same as below
             setIsStreaming(false);
             refreshChats();
@@ -120,7 +124,7 @@ export function useStreamChat({
         setError(error instanceof Error ? error.message : String(error));
       }
     },
-    [setMessages, setIsStreaming, setIsPreviewOpen],
+    [setMessages, setIsStreaming, setIsPreviewOpen, refetchUserBudget],
   );
 
   return {

src/hooks/useUserBudgetInfo.ts

@@ -0,0 +1,34 @@
+import { useQuery } from "@tanstack/react-query";
+import { IpcClient } from "@/ipc/ipc_client";
+import type { UserBudgetInfo } from "@/ipc/ipc_types";
+
+const FIVE_MINUTES_IN_MS = 5 * 60 * 1000;
+
+export function useUserBudgetInfo() {
+  const queryKey = ["userBudgetInfo"];
+
+  const { data, isLoading, error, isFetching, refetch } = useQuery<
+    UserBudgetInfo | null,
+    Error,
+    UserBudgetInfo | null
+  >({
+    queryKey: queryKey,
+    queryFn: async () => {
+      const ipcClient = IpcClient.getInstance();
+      return ipcClient.getUserBudget();
+    },
+    // This data is not critical and can be stale for a bit
+    staleTime: FIVE_MINUTES_IN_MS,
+    // If an error occurs (e.g. API key not set), it returns null.
+    // We don't want react-query to retry automatically in such cases as it's not a transient network error.
+    retry: false,
+  });
+
+  return {
+    userBudget: data,
+    isLoadingUserBudget: isLoading,
+    userBudgetError: error,
+    isFetchingUserBudget: isFetching,
+    refetchUserBudget: refetch,
+  };
+}