Browse Source

Display credit balance for all accounts (#4992)

* Display credit balance for all accounts

The credit balance display was previously only shown for personal accounts. This change removes the check for `activeOrganization === null` and displays the credit balance and "Add Credits" button for all account types, including organization accounts. A divider is added above the balance section for visual separation once the backend change is deployed.

* changeset

* Improve refresh logic

Refactors the `AccountView` component to properly display and manage credits for both user and organization accounts. It introduces the `getOrganizationCredits` API call to fetch organization-specific credits and updates the UI accordingly. The refresh logic has also been improved to ensure data consistency and prevent unnecessary API calls.

Key changes:

- Implemented `getOrganizationCredits` to fetch credits for the active organization.
- Modified the credit display to show organization credits when an organization is active.
- Updated the refresh logic to use `useCallback` and `debounce` for better performance and to prevent race conditions.
- Added a periodic refresh to update account data every 30 seconds.
- Improved error handling and loading state management.
- Removed the interval ref and replaced it with a simpler useEffect for periodic refresh.
- Added last fetch time to the UI.

* clean up

* deepEqual

* org management

* prevent race condition
Bee 5 months ago
parent
commit
db6d288efa

+ 5 - 0
.changeset/bright-olives-leave.md

@@ -0,0 +1,5 @@
+---
+"claude-dev": patch
+---
+
+Display account balance for all org members

+ 0 - 1
webview-ui/.eslintrc.json

@@ -23,7 +23,6 @@
 		"@typescript-eslint/no-explicit-any": "off",
 		"@typescript-eslint/no-empty-object-type": "off",
 		"no-case-declarations": "off",
-		"react-hooks/exhaustive-deps": "off",
 		"prefer-const": "off",
 		"no-extra-semi": "off",
 		"eslint-rules/no-direct-vscode-api": "warn",

+ 158 - 116
webview-ui/src/components/account/AccountView.tsx

@@ -16,6 +16,9 @@ import {
 import { memo, useCallback, useEffect, useMemo, useRef, useState } from "react"
 import ClineLogoWhite from "../../assets/ClineLogoWhite"
 import CreditsHistoryTable from "./CreditsHistoryTable"
+import { GetOrganizationCreditsRequest } from "@shared/proto/account"
+import debounce from "debounce"
+import deepEqual from "fast-deep-equal"
 import VSCodeButtonLink from "../common/VSCodeButtonLink"
 
 // Custom hook for animated credit display with styled decimals
@@ -123,10 +126,10 @@ export const ClineAccountView = () => {
 	const [userOrganizations, setUserOrganizations] = useState<UserOrganization[]>([])
 	const [activeOrganization, setActiveOrganization] = useState<UserOrganization | null>(null)
 	const [isLoading, setIsLoading] = useState(true)
-	const [isSwitchingOrg, setIsSwitchingOrg] = useState(false)
+	const [isSwitchingProfile, setIsSwitchingProfile] = useState(false)
 	const [usageData, setUsageData] = useState<ClineAccountUsageTransaction[]>([])
 	const [paymentsData, setPaymentsData] = useState<PaymentTransaction[]>([])
-	const intervalRef = useRef<NodeJS.Timeout | null>(null)
+	const [lastFetchTime, setLastFetchTime] = useState<number>(Date.now())
 
 	const clineUris = useMemo(() => {
 		const base = new URL(clineUser?.appBaseUrl || CLINE_APP_URL)
@@ -141,113 +144,166 @@ export const ClineAccountView = () => {
 		}
 	}, [clineUser?.appBaseUrl, activeOrganization])
 
-	async function getUserCredits() {
-		setIsLoading(true)
-		try {
-			const response = await AccountServiceClient.getUserCredits(EmptyRequest.create())
-			setBalance(response.balance?.currentBalance ?? null)
-			setUsageData(convertProtoUsageTransactions(response.usageTransactions))
-			setPaymentsData(response.paymentTransactions)
-		} catch (error) {
-			console.error("Failed to fetch user credits data:", error)
-			setBalance(null)
-			setUsageData([])
-			setPaymentsData([])
-		} finally {
-			setIsLoading(false)
-		}
-	}
+	// Add a ref to track the intended organization during transitions
+	const pendingOrganizationRef = useRef<string | null>(null)
 
-	async function getUserOrganizations() {
-		setIsLoading(true)
+	const getUserOrganizations = useCallback(async () => {
 		try {
-			const response = await AccountServiceClient.getUserOrganizations(EmptyRequest.create())
-			setUserOrganizations(response.organizations || [])
-			setActiveOrganization(response.organizations.find((org: UserOrganization) => org.active) || null)
-		} catch (error) {
-			console.error("Failed to fetch user organizations:", error)
-			setUserOrganizations([])
-			setActiveOrganization(null)
-		} finally {
-			setIsLoading(false)
-		}
-	}
-
-	// Fetch all account data when component mounts using gRPC
-	useEffect(() => {
-		if (!user) return
-
-		const fetchUserData = async () => {
-			try {
-				Promise.all([getUserCredits(), getUserOrganizations()])
-			} catch (error) {
-				console.error("Failed to fetch user data:", error)
+			if (!clineUser?.uid) {
 				setBalance(null)
-				setUsageData([])
-				setPaymentsData([])
-			} finally {
+				setUserOrganizations([])
+				setActiveOrganization(null)
+				setIsSwitchingProfile(false)
 				setIsLoading(false)
+				return
+			}
+			const response = await AccountServiceClient.getUserOrganizations(EmptyRequest.create())
+			if (response.organizations && !deepEqual(userOrganizations, response.organizations)) {
+				setUserOrganizations(response.organizations)
+
+				// Only update activeOrganization if we're not in the middle of a switch
+				// or if the server response matches our pending change
+				const serverActiveOrg = response.organizations.find((org: UserOrganization) => org.active)
+				const serverActiveOrgId = serverActiveOrg?.organizationId || ""
+
+				if (!isSwitchingProfile || pendingOrganizationRef.current === serverActiveOrgId) {
+					if (serverActiveOrgId !== (activeOrganization?.organizationId || "")) {
+						setActiveOrganization(serverActiveOrg || null)
+					}
+					// Clear pending ref if the server state matches
+					if (pendingOrganizationRef.current === serverActiveOrgId) {
+						pendingOrganizationRef.current = null
+					}
+				}
 			}
+		} catch (error) {
+			console.error("Failed to fetch user organizations:", error)
 		}
+	}, [clineUser?.uid, userOrganizations, isSwitchingProfile, activeOrganization?.organizationId])
 
-		fetchUserData()
-	}, [user])
+	const fetchCreditBalance = useCallback(async () => {
+		try {
+			setIsLoading(true)
 
-	// Periodic refresh while component is mounted
-	useEffect(() => {
-		if (!user) return
+			// Use the pending organization if we're switching, otherwise use current active org
+			const targetOrgId =
+				pendingOrganizationRef.current !== null ? pendingOrganizationRef.current : activeOrganization?.organizationId
 
-		intervalRef.current = setInterval(() => {
-			getUserCredits().catch((err) => console.error("Auto-refresh failed:", err))
-		}, 10_000)
+			const response = targetOrgId
+				? await AccountServiceClient.getOrganizationCredits(
+						GetOrganizationCreditsRequest.fromPartial({
+							organizationId: targetOrgId,
+						}),
+					)
+				: await AccountServiceClient.getUserCredits(EmptyRequest.create())
 
-		return () => {
-			if (intervalRef.current) clearInterval(intervalRef.current)
-		}
-	}, [user])
+			// Update balance if changed
+			const newBalance = response.balance?.currentBalance
+			if (newBalance !== balance) {
+				setBalance(newBalance ?? null)
+			}
 
-	const handleManualRefresh = async () => {
-		await getUserCredits()
+			const clineUsage = convertProtoUsageTransactions(response.usageTransactions)
+			setUsageData(clineUsage || [])
 
-		if (intervalRef.current) {
-			clearInterval(intervalRef.current)
-			intervalRef.current = setInterval(() => {
-				getUserCredits().catch((err) => console.error("Auto-refresh failed:", err))
-			}, 10_000)
+			if (activeOrganization?.organizationId) {
+				setPaymentsData([]) // Organizations don't have payment transactions
+			} else {
+				// Check if response is UserCreditsData type
+				if (typeof response !== "object" || !("paymentTransactions" in response)) {
+					return
+				}
+				const paymentsData = response.paymentTransactions || []
+				// Check if paymentTransactions is part of the response
+				if (response.paymentTransactions?.length !== paymentsData?.length) {
+					setPaymentsData(paymentsData)
+				}
+			}
+		} finally {
+			setLastFetchTime(Date.now())
+			setIsLoading(false)
 		}
-	}
-
-	const handleLogin = () => {
-		handleSignIn()
-	}
+	}, [activeOrganization?.organizationId])
 
-	const handleLogout = () => {
-		handleSignOut()
-	}
+	const handleManualRefresh = useCallback(
+		debounce(() => !isLoading && fetchCreditBalance(), 500, { immediate: true }),
+		[fetchCreditBalance, isLoading],
+	)
 
 	const handleOrganizationChange = useCallback(
 		async (event: any) => {
 			const newOrgId = (event.target as VSCodeDropdownChangeEvent["target"]).value
+			const currentOrgId = activeOrganization?.organizationId || ""
+
+			if (currentOrgId !== newOrgId) {
+				setIsSwitchingProfile(true)
+				setBalance(null)
 
-			if (activeOrganization?.organizationId !== newOrgId) {
-				setIsSwitchingOrg(true) // Disable dropdown
+				// Set the pending organization immediately to prevent race conditions
+				pendingOrganizationRef.current = newOrgId
 
 				try {
+					// Update local state immediately for UI responsiveness
+					if (newOrgId === "") {
+						setActiveOrganization(null)
+					} else {
+						const org = userOrganizations.find((org: UserOrganization) => org.organizationId === newOrgId)
+						if (org) {
+							setActiveOrganization(org)
+						}
+					}
+
+					// Send the change to the server
 					await AccountServiceClient.setUserOrganization(
 						UserOrganizationUpdateRequest.create({ organizationId: newOrgId }),
 					)
-					await getUserOrganizations() // Refresh to get new active org
-					await getUserCredits() // Refresh credits for new org
+
+					// Fetch fresh data for the new organization
+					await fetchCreditBalance()
+
+					// Refresh organizations to get the updated active state from server
+					await getUserOrganizations()
 				} catch (error) {
 					console.error("Failed to update organization:", error)
+					// Reset pending ref on error
+					pendingOrganizationRef.current = null
 				} finally {
-					setIsSwitchingOrg(false) // Re-enable dropdown
+					setIsSwitchingProfile(false)
 				}
 			}
 		},
-		[activeOrganization],
+		[activeOrganization?.organizationId, fetchCreditBalance, getUserOrganizations, userOrganizations],
 	)
 
+	// Handle organization changes and initial load
+	useEffect(() => {
+		const loadData = async () => {
+			await getUserOrganizations()
+			await fetchCreditBalance()
+		}
+		loadData()
+	}, [activeOrganization?.organizationId])
+
+	// Periodic refresh
+	useEffect(() => {
+		const refreshData = async () => {
+			try {
+				if (clineUser?.uid) {
+					await Promise.all([getUserOrganizations(), fetchCreditBalance()])
+				}
+			} catch (error) {
+				console.error("Error during periodic refresh:", error)
+			}
+		}
+
+		const intervalId = setInterval(refreshData, 30000)
+		return () => clearInterval(intervalId)
+	}, [clineUser?.uid, getUserOrganizations, fetchCreditBalance])
+
+	// Determine the current dropdown value, considering pending changes
+	const dropdownValue =
+		pendingOrganizationRef.current !== null ? pendingOrganizationRef.current : activeOrganization?.organizationId || ""
+
 	return (
 		<div className="h-full flex flex-col">
 			{user ? (
@@ -276,10 +332,9 @@ export const ClineAccountView = () => {
 								<div className="flex gap-2 items-center mt-1">
 									{userOrganizations && (
 										<VSCodeDropdown
-											key={activeOrganization?.organizationId || "personal"}
-											currentValue={activeOrganization?.organizationId || ""}
+											currentValue={dropdownValue}
 											onChange={handleOrganizationChange}
-											disabled={isSwitchingOrg || isLoading}
+											disabled={isSwitchingProfile}
 											className="w-full">
 											<VSCodeOption value="">Personal</VSCodeOption>
 											{userOrganizations.map((org: UserOrganization) => (
@@ -305,55 +360,42 @@ export const ClineAccountView = () => {
 								Dashboard
 							</VSCodeButtonLink>
 						</div>
-						<VSCodeButton appearance="secondary" onClick={handleLogout} className="w-full min-[225px]:w-1/2">
+						<VSCodeButton appearance="secondary" onClick={() => handleSignOut()} className="w-full min-[225px]:w-1/2">
 							Log out
 						</VSCodeButton>
 					</div>
 
-					{/* Credit balance is not available for organization account */}
-					{activeOrganization === null && <VSCodeDivider className="w-full my-6" />}
+					<VSCodeDivider className="w-full my-6" />
 
-					{activeOrganization === null && (
-						<div className="w-full flex flex-col items-center">
-							<div className="text-sm text-[var(--vscode-descriptionForeground)] mb-3 font-azeret-mono font-light">
-								CURRENT BALANCE
-							</div>
+					<div
+						className="w-full flex flex-col items-center"
+						title={`Last updated: ${new Date(lastFetchTime).toLocaleTimeString()}`}>
+						<div className="text-sm text-[var(--vscode-descriptionForeground)] mb-3 font-azeret-mono font-light">
+							CURRENT BALANCE
+						</div>
 
-							<div className="text-4xl font-bold text-[var(--vscode-foreground)] mb-6 flex items-center gap-2">
-								{isLoading ? (
-									<div className="text-[var(--vscode-descriptionForeground)]">Loading...</div>
-								) : (
-									<>
-										{balance === null ? (
-											<span>----</span>
-										) : (
-											<>
-												<StyledCreditDisplay balance={balance} />
-											</>
-										)}
-										<VSCodeButton appearance="icon" className="mt-1" onClick={handleManualRefresh}>
-											<span className="codicon codicon-refresh"></span>
-										</VSCodeButton>
-									</>
-								)}
-							</div>
+						<div className="text-4xl font-bold text-[var(--vscode-foreground)] mb-6 flex items-center gap-2">
+							{balance === null ? <span>----</span> : <StyledCreditDisplay balance={balance} />}
+							<VSCodeButton appearance="icon" className="mt-1" onClick={handleManualRefresh}>
+								<span className="codicon codicon-refresh"></span>
+							</VSCodeButton>
+						</div>
 
-							<div className="w-full">
-								<VSCodeButtonLink href={clineUris.credits.href} className="w-full">
-									Add Credits
-								</VSCodeButtonLink>
-							</div>
+						<div className="w-full">
+							<VSCodeButtonLink href={clineUris.credits.href} className="w-full">
+								Add Credits
+							</VSCodeButtonLink>
 						</div>
-					)}
+					</div>
 
 					<VSCodeDivider className="mt-6 mb-3 w-full" />
 
 					<div className="flex-grow flex flex-col min-h-0 pb-[0px]">
 						<CreditsHistoryTable
-							isLoading={isLoading}
+							isLoading={isSwitchingProfile}
 							usageData={usageData}
 							paymentsData={paymentsData}
-							showPayments={!activeOrganization}
+							showPayments={!activeOrganization?.active}
 						/>
 					</div>
 				</div>
@@ -361,12 +403,12 @@ export const ClineAccountView = () => {
 				<div className="flex flex-col items-center pr-3">
 					<ClineLogoWhite className="size-16 mb-4" />
 
-					<p style={{}}>
+					<p>
 						Sign up for an account to get access to the latest models, billing dashboard to view usage and credits,
 						and more upcoming features.
 					</p>
 
-					<VSCodeButton onClick={handleLogin} className="w-full mb-4">
+					<VSCodeButton onClick={() => handleSignIn()} className="w-full mb-4">
 						Sign up with Cline
 					</VSCodeButton>