pro: show remaining credits (#329)

Fixes #265

src/app/TitleBar.tsx

@@ -13,6 +13,13 @@ import { useEffect, useState } from "react";
 import { DyadProSuccessDialog } from "@/components/DyadProSuccessDialog";
 import { useTheme } from "@/contexts/ThemeContext";
 import { IpcClient } from "@/ipc/ipc_client";
+import { useUserBudgetInfo } from "@/hooks/useUserBudgetInfo";
+import { UserBudgetInfo } from "@/ipc/ipc_types";
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipTrigger,
+} from "@/components/ui/tooltip";
 
 export const TitleBar = () => {
   const [selectedAppId] = useAtom(selectedAppIdAtom);
@@ -64,7 +71,7 @@ export const TitleBar = () => {
   };
 
   const isDyadPro = !!settings?.providerSettings?.auto?.apiKey?.value;
-  const isDyadProEnabled = settings?.enableDyadPro;
+  const isDyadProEnabled = Boolean(settings?.enableDyadPro);
 
   return (
     <>
@@ -82,25 +89,7 @@ export const TitleBar = () => {
         >
           {displayText}
         </Button>
-        {isDyadPro && (
+        {isDyadPro && <DyadProButton isDyadProEnabled={isDyadProEnabled} />}
-          <Button
-            data-testid="title-bar-dyad-pro-button"
-            onClick={() => {
-              navigate({
-                to: providerSettingsRoute.id,
-                params: { provider: "auto" },
-              });
-            }}
-            variant="outline"
-            className={cn(
-              "ml-4 no-app-region-drag h-7 bg-indigo-600 text-white dark:bg-indigo-600 dark:text-white",
-              !isDyadProEnabled && "bg-zinc-600 dark:bg-zinc-600",
-            )}
-            size="sm"
-          >
-            {isDyadProEnabled ? "Dyad Pro" : "Dyad Pro (disabled)"}
-          </Button>
-        )}
         {showWindowControls && <WindowsControls />}
       </div>
 
@@ -192,3 +181,58 @@ function WindowsControls() {
     </div>
   );
 }
+
+export function DyadProButton({
+  isDyadProEnabled,
+}: {
+  isDyadProEnabled: boolean;
+}) {
+  const { navigate } = useRouter();
+  const { userBudget } = useUserBudgetInfo();
+  return (
+    <Button
+      data-testid="title-bar-dyad-pro-button"
+      onClick={() => {
+        navigate({
+          to: providerSettingsRoute.id,
+          params: { provider: "auto" },
+        });
+      }}
+      variant="outline"
+      className={cn(
+        "ml-4 no-app-region-drag h-7 bg-indigo-600 text-white dark:bg-indigo-600 dark:text-white",
+        !isDyadProEnabled && "bg-zinc-600 dark:bg-zinc-600",
+      )}
+      size="sm"
+    >
+      {isDyadProEnabled ? "Dyad Pro" : "Dyad Pro (disabled)"}
+      {userBudget && <AICreditStatus userBudget={userBudget} />}
+    </Button>
+  );
+}
+
+export function AICreditStatus({ userBudget }: { userBudget: UserBudgetInfo }) {
+  const remaining = Math.round(
+    userBudget.totalCredits - userBudget.usedCredits,
+  );
+  return (
+    <Tooltip>
+      <TooltipTrigger>
+        <div className="text-xs mt-0.5">{remaining} credits left</div>
+      </TooltipTrigger>
+      <TooltipContent>
+        <div>
+          <p>
+            You have used {Math.round(userBudget.usedCredits)} credits out of{" "}
+            {userBudget.totalCredits}.
+          </p>
+          <p>
+            Your budget resets on{" "}
+            {userBudget.budgetResetDate.toLocaleDateString()}
+          </p>
+          <p>Note: there is a slight delay in updating the credit status.</p>
+        </div>
+      </TooltipContent>
+    </Tooltip>
+  );
+}

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,
+  };
+}

src/ipc/handlers/pro_handlers.ts

@@ -0,0 +1,61 @@
+import fetch from "node-fetch"; // Electron main process might need node-fetch
+import log from "electron-log";
+import { createLoggedHandler } from "./safe_handle";
+import { readSettings } from "../../main/settings"; // Assuming settings are read this way
+import { UserBudgetInfoSchema } from "../ipc_types";
+
+const logger = log.scope("pro_handlers");
+const handle = createLoggedHandler(logger);
+
+const CONVERSION_RATIO = (10 * 3) / 2;
+
+export function registerProHandlers() {
+  // This method should try to avoid throwing errors because this is auxiliary
+  // information and isn't critical to using the app
+  handle("get-user-budget", async (): Promise<UserBudgetInfo | null> => {
+    logger.info("Attempting to fetch user budget information.");
+
+    const settings = readSettings();
+
+    const apiKey = settings.providerSettings?.auto?.apiKey?.value;
+
+    if (!apiKey) {
+      logger.error("LLM Gateway API key (Dyad Pro) is not configured.");
+      return null;
+    }
+
+    const url = "https://llm-gateway.dyad.sh/user/info";
+    const headers = {
+      "Content-Type": "application/json",
+      Authorization: `Bearer ${apiKey}`,
+    };
+
+    try {
+      // Use native fetch if available, otherwise node-fetch will be used via import
+      const response = await fetch(url, {
+        method: "GET",
+        headers: headers,
+      });
+
+      if (!response.ok) {
+        const errorBody = await response.text();
+        logger.error(
+          `Failed to fetch user budget. Status: ${response.status}. Body: ${errorBody}`,
+        );
+        return null;
+      }
+
+      const data = await response.json();
+      const userInfoData = data["user_info"];
+      logger.info("Successfully fetched user budget information.");
+      return UserBudgetInfoSchema.parse({
+        usedCredits: userInfoData["spend"] * CONVERSION_RATIO,
+        totalCredits: userInfoData["max_budget"] * CONVERSION_RATIO,
+        budgetResetDate: new Date(userInfoData["budget_reset_at"]),
+      });
+    } catch (error: any) {
+      logger.error(`Error fetching user budget: ${error.message}`, error);
+      return null;
+    }
+  });
+}

src/ipc/ipc_client.ts

@@ -30,6 +30,7 @@ import type {
   ImportAppResult,
   ImportAppParams,
   RenameBranchParams,
+  UserBudgetInfo,
 } from "./ipc_types";
 import type { ProposalResult } from "@/lib/schemas";
 import { showError } from "@/lib/toast";
@@ -825,4 +826,9 @@ export class IpcClient {
   async clearSessionData(): Promise<void> {
     return this.ipcRenderer.invoke("clear-session-data");
   }
+
+  // Method to get user budget information
+  public async getUserBudget(): Promise<UserBudgetInfo | null> {
+    return this.ipcRenderer.invoke("get-user-budget");
+  }
 }

src/ipc/ipc_host.ts

@@ -18,6 +18,7 @@ import { registerLanguageModelHandlers } from "./handlers/language_model_handler
 import { registerReleaseNoteHandlers } from "./handlers/release_note_handlers";
 import { registerImportHandlers } from "./handlers/import_handlers";
 import { registerSessionHandlers } from "./handlers/session_handlers";
+import { registerProHandlers } from "./handlers/pro_handlers";
 
 export function registerIpcHandlers() {
   // Register all IPC handlers by category
@@ -41,4 +42,5 @@ export function registerIpcHandlers() {
   registerReleaseNoteHandlers();
   registerImportHandlers();
   registerSessionHandlers();
+  registerProHandlers();
 }

src/ipc/ipc_types.ts

@@ -1,3 +1,5 @@
+import { z } from "zod";
+
 export interface AppOutput {
   type: "stdout" | "stderr" | "info" | "client-error";
   message: string;
@@ -207,3 +209,10 @@ export interface RenameBranchParams {
   oldBranchName: string;
   newBranchName: string;
 }
+
+export const UserBudgetInfoSchema = z.object({
+  usedCredits: z.number(),
+  totalCredits: z.number(),
+  budgetResetDate: z.date(),
+});
+export type UserBudgetInfo = z.infer<typeof UserBudgetInfoSchema>;

src/preload.ts

@@ -75,6 +75,7 @@ const validInvokeChannels = [
   "check-app-name",
   "rename-branch",
   "clear-session-data",
+  "get-user-budget",
 ] as const;
 
 // Add valid receive channels