/** * Balance hook — uses the WebSocket gateway to receive instant updates when * a tx involving the current address is committed, with HTTP polling as a * graceful fallback for old nodes that don't expose /api/ws. * * Flow: * 1. On mount: immediate HTTP fetch so the UI has a non-zero balance ASAP * 2. Subscribe to `addr:` on the WS hub * 3. On every `tx` event, re-fetch balance (cheap — one Badger read server-side) * 4. If WS disconnects for >15s, fall back to 10-second polling until it reconnects */ import { useEffect, useCallback, useRef } from 'react'; import { getBalance } from '@/lib/api'; import { getWSClient } from '@/lib/ws'; import { useStore } from '@/lib/store'; const FALLBACK_POLL_INTERVAL = 10_000; // HTTP poll when WS is down const WS_GRACE_BEFORE_POLLING = 15_000; // don't start polling immediately on disconnect export function useBalance() { const keyFile = useStore(s => s.keyFile); const setBalance = useStore(s => s.setBalance); const refresh = useCallback(async () => { if (!keyFile) return; try { const bal = await getBalance(keyFile.pub_key); setBalance(bal); } catch { // transient — next call will retry } }, [keyFile, setBalance]); // --- fallback polling management --- const pollTimerRef = useRef | null>(null); const disconnectSinceRef = useRef(null); const disconnectTORef = useRef | null>(null); const startPolling = useCallback(() => { if (pollTimerRef.current) return; console.log('[useBalance] WS down for grace period — starting HTTP poll'); refresh(); pollTimerRef.current = setInterval(refresh, FALLBACK_POLL_INTERVAL); }, [refresh]); const stopPolling = useCallback(() => { if (pollTimerRef.current) { clearInterval(pollTimerRef.current); pollTimerRef.current = null; } if (disconnectTORef.current) { clearTimeout(disconnectTORef.current); disconnectTORef.current = null; } disconnectSinceRef.current = null; }, []); useEffect(() => { if (!keyFile) return; const ws = getWSClient(); // Immediate HTTP fetch so the UI is not empty while the WS hello arrives. refresh(); // Refresh balance whenever a tx for our address is committed. const offTx = ws.subscribe('addr:' + keyFile.pub_key, (frame) => { if (frame.event === 'tx') { refresh(); } }); // Manage fallback polling based on WS connection state. const offConn = ws.onConnectionChange((ok) => { if (ok) { stopPolling(); refresh(); // catch up anything we missed while disconnected } else if (disconnectTORef.current === null) { disconnectSinceRef.current = Date.now(); disconnectTORef.current = setTimeout(startPolling, WS_GRACE_BEFORE_POLLING); } }); ws.connect(); return () => { offTx(); offConn(); stopPolling(); }; }, [keyFile, refresh, startPolling, stopPolling]); return { refresh }; }