From 9d1a0f7ad7c2fab5ff2c0fdba97dccb2829335c7 Mon Sep 17 00:00:00 2001 From: Will Chen Date: Tue, 3 Jun 2025 23:46:35 -0700 Subject: [PATCH] pro: show remaining credits (#329) Fixes #265 --- src/app/TitleBar.tsx | 84 ++++++++++++++++++++++++-------- src/hooks/useStreamChat.ts | 6 ++- src/hooks/useUserBudgetInfo.ts | 34 +++++++++++++ src/ipc/handlers/pro_handlers.ts | 61 +++++++++++++++++++++++ src/ipc/ipc_client.ts | 6 +++ src/ipc/ipc_host.ts | 2 + src/ipc/ipc_types.ts | 9 ++++ src/preload.ts | 1 + 8 files changed, 182 insertions(+), 21 deletions(-) create mode 100644 src/hooks/useUserBudgetInfo.ts create mode 100644 src/ipc/handlers/pro_handlers.ts diff --git a/src/app/TitleBar.tsx b/src/app/TitleBar.tsx index 38fffef..ee0c348 100644 --- a/src/app/TitleBar.tsx +++ b/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} - {isDyadPro && ( - - )} + {isDyadPro && } {showWindowControls && } @@ -192,3 +181,58 @@ function WindowsControls() { ); } + +export function DyadProButton({ + isDyadProEnabled, +}: { + isDyadProEnabled: boolean; +}) { + const { navigate } = useRouter(); + const { userBudget } = useUserBudgetInfo(); + return ( + + ); +} + +export function AICreditStatus({ userBudget }: { userBudget: UserBudgetInfo }) { + const remaining = Math.round( + userBudget.totalCredits - userBudget.usedCredits, + ); + return ( + + +
{remaining} credits left
+ + +
+

+ You have used {Math.round(userBudget.usedCredits)} credits out of{" "} + {userBudget.totalCredits}. +

+

+ Your budget resets on{" "} + {userBudget.budgetResetDate.toLocaleDateString()} +

+

Note: there is a slight delay in updating the credit status.

+
+ + + ); +} diff --git a/src/hooks/useStreamChat.ts b/src/hooks/useStreamChat.ts index 24668d8..2c11e74 100644 --- a/src/hooks/useStreamChat.ts +++ b/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 { diff --git a/src/hooks/useUserBudgetInfo.ts b/src/hooks/useUserBudgetInfo.ts new file mode 100644 index 0000000..b91e25b --- /dev/null +++ b/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, + }; +} diff --git a/src/ipc/handlers/pro_handlers.ts b/src/ipc/handlers/pro_handlers.ts new file mode 100644 index 0000000..0294127 --- /dev/null +++ b/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 => { + 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; + } + }); +} diff --git a/src/ipc/ipc_client.ts b/src/ipc/ipc_client.ts index 89d0c77..44ca9ea 100644 --- a/src/ipc/ipc_client.ts +++ b/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 { return this.ipcRenderer.invoke("clear-session-data"); } + + // Method to get user budget information + public async getUserBudget(): Promise { + return this.ipcRenderer.invoke("get-user-budget"); + } } diff --git a/src/ipc/ipc_host.ts b/src/ipc/ipc_host.ts index efbffe7..bef5965 100644 --- a/src/ipc/ipc_host.ts +++ b/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(); } diff --git a/src/ipc/ipc_types.ts b/src/ipc/ipc_types.ts index 744cefc..da75703 100644 --- a/src/ipc/ipc_types.ts +++ b/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; diff --git a/src/preload.ts b/src/preload.ts index 74e1226..bb70827 100644 --- a/src/preload.ts +++ b/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