feat: use sse to get balances

limitofzero · limitofzero · commit 129ae2700221 · 2026-01-05T01:53:08.000+04:00
diff --git a/apps/cowswap-frontend/src/modules/balancesAndAllowances/updaters/CommonPriorityBalancesAndAllowancesUpdater.tsx b/apps/cowswap-frontend/src/modules/balancesAndAllowances/updaters/CommonPriorityBalancesAndAllowancesUpdater.tsx
@@ -91,6 +91,7 @@ export function CommonPriorityBalancesAndAllowancesUpdater(): ReactNode {
         chainId={sourceChainId}
         isBffSwitchedOn={isBffSwitchedOn}
         isBffEnabled={isBffEnabled}
+        isSseEnabled={true}
         excludedTokens={priorityTokenAddresses}
         invalidateCacheTrigger={invalidateCacheTrigger}
       />
diff --git a/libs/balances-and-allowances/src/hooks/useSseBalances.ts b/libs/balances-and-allowances/src/hooks/useSseBalances.ts
@@ -0,0 +1,171 @@
+import { useCallback, useEffect, useRef, useState } from 'react'
+
+import { BALANCES_SSE_URL } from '@cowprotocol/common-const'
+import { SupportedChainId } from '@cowprotocol/cow-sdk'
+
+const RECONNECT_DELAY_MS = 3000
+const MAX_RECONNECT_ATTEMPTS = 5
+
+export interface SseBalancesState {
+  isConnected: boolean
+  isLoading: boolean
+  error: Error | null
+}
+
+export interface UseSseBalancesParams {
+  account: string | undefined
+  chainId: SupportedChainId
+  enabled: boolean
+  tokensListsUrls: string[]
+  customTokens?: string[]
+  onAllBalances: (balances: Record<string, string>) => void
+  onBalanceUpdate: (address: string, balance: string) => void
+  onError?: (error: Error) => void
+}
+
+const INITIAL_STATE: SseBalancesState = { isConnected: false, isLoading: false, error: null }
+
+async function createSession(
+  chainId: SupportedChainId,
+  account: string,
+  tokensListsUrls: string[],
+  customTokens?: string[],
+): Promise<void> {
+  const response = await fetch(`${BALANCES_SSE_URL}/${chainId}/sessions/${account}`, {
+    method: 'POST',
+    headers: { 'Content-Type': 'application/json' },
+    body: JSON.stringify({ tokensListsUrls, customTokens }),
+  })
+
+  if (!response.ok) {
+    throw new Error(`Session creation failed: ${response.status}`)
+  }
+}
+
+interface EventHandlers {
+  onAllBalances: (balances: Record<string, string>) => void
+  onBalanceUpdate: (address: string, balance: string) => void
+  onError?: (error: Error) => void
+  onOpen: () => void
+  onClose: () => void
+}
+
+function setupEventSource(url: string, handlers: EventHandlers): EventSource {
+  const { onAllBalances, onBalanceUpdate, onError, onOpen, onClose } = handlers
+  const es = new EventSource(url)
+
+  es.onopen = onOpen
+
+  es.addEventListener('all_balances', (e: MessageEvent): void => {
+    try {
+      const { balances } = JSON.parse(e.data)
+      if (balances && Object.keys(balances).length > 0) onAllBalances(balances)
+    } catch {
+      /* ignore */
+    }
+  })
+
+  es.addEventListener('balance_update', (e: MessageEvent): void => {
+    try {
+      const { address, balance } = JSON.parse(e.data)
+      if (address && balance != null) onBalanceUpdate(address.toLowerCase(), balance)
+    } catch {
+      /* ignore */
+    }
+  })
+
+  es.addEventListener('error', (e: MessageEvent): void => {
+    try {
+      const { message, code } = JSON.parse(e.data)
+      onError?.(new Error(`SSE Error ${code}: ${message}`))
+    } catch {
+      /* ignore */
+    }
+  })
+
+  es.onerror = (): void => {
+    if (es.readyState === EventSource.CLOSED) onClose()
+  }
+
+  return es
+}
+
+export function useSseBalances(params: UseSseBalancesParams): SseBalancesState {
+  const {
+    account,
+    chainId,
+    enabled,
+    tokensListsUrls,
+    customTokens = [],
+    onAllBalances,
+    onBalanceUpdate,
+    onError,
+  } = params
+
+  const [state, setState] = useState<SseBalancesState>(INITIAL_STATE)
+  const esRef = useRef<EventSource | null>(null)
+  const attemptsRef = useRef(0)
+  const timeoutRef = useRef<ReturnType<typeof setTimeout>>()
+
+  const cleanup = useCallback((): void => {
+    esRef.current?.close()
+    esRef.current = null
+    clearTimeout(timeoutRef.current)
+  }, [])
+
+  useEffect(() => {
+    if (!enabled || !account || tokensListsUrls.length === 0) {
+      cleanup()
+      setState(INITIAL_STATE)
+      return cleanup
+    }
+
+    let cancelled = false
+
+    const connect = async (): Promise<void> => {
+      if (cancelled) return
+      cleanup()
+      setState((s) => ({ ...s, isLoading: true, error: null }))
+
+      try {
+        await createSession(chainId, account, tokensListsUrls, customTokens)
+      } catch (e) {
+        const error = e instanceof Error ? e : new Error('Session failed')
+        setState({ ...INITIAL_STATE, error })
+        onError?.(error)
+        return
+      }
+
+      if (cancelled) return
+
+      esRef.current = setupEventSource(`${BALANCES_SSE_URL}/sse/${chainId}/balances/${account}`, {
+        onAllBalances,
+        onBalanceUpdate,
+        onError,
+        onOpen: (): void => {
+          attemptsRef.current = 0
+          setState({ isConnected: true, isLoading: false, error: null })
+        },
+        onClose: (): void => {
+          setState((s) => ({ ...s, isConnected: false }))
+          if (attemptsRef.current < MAX_RECONNECT_ATTEMPTS) {
+            attemptsRef.current++
+            timeoutRef.current = setTimeout(() => void connect(), RECONNECT_DELAY_MS * 2 ** (attemptsRef.current - 1))
+          } else {
+            const err = new Error('SSE: Max reconnect attempts')
+            setState((s) => ({ ...s, error: err }))
+            onError?.(err)
+          }
+        },
+      })
+    }
+
+    void connect()
+    return (): void => {
+      cancelled = true
+      cleanup()
+    }
+  }, [account, chainId, enabled, tokensListsUrls, customTokens, cleanup, onAllBalances, onBalanceUpdate, onError])
+
+  return state
+}
diff --git a/libs/balances-and-allowances/src/index.ts b/libs/balances-and-allowances/src/index.ts
@@ -1,6 +1,9 @@
-// Updater
+// Updaters
 export { BalancesAndAllowancesUpdater } from './updaters/BalancesAndAllowancesUpdater'
 export { PriorityTokensUpdater, PRIORITY_TOKENS_REFRESH_INTERVAL } from './updaters/PriorityTokensUpdater'
+export { BalancesBffUpdater } from './updaters/BalancesBffUpdater'
+export { BalancesRpcCallUpdater } from './updaters/BalancesRpcCallUpdater'
+export { BalancesSseUpdater } from './updaters/BalancesSseUpdater'
 
 // Hooks
 export { useTokensBalances } from './hooks/useTokensBalances'
@@ -14,15 +17,20 @@ export { useUpdateTokenBalance } from './hooks/useUpdateTokenBalance'
 export { useTokenAllowances } from './hooks/useTokenAllowances'
 export { useBalancesAndAllowances } from './hooks/useBalancesAndAllowances'
 export { useTradeSpenderAddress } from './hooks/useTradeSpenderAddress'
+export { useSseBalances } from './hooks/useSseBalances'
+
+// State hooks
 export { useIsBffFailed } from './state/isBffFailedAtom'
-export { BalancesBffUpdater } from './updaters/BalancesBffUpdater'
-export { BalancesRpcCallUpdater } from './updaters/BalancesRpcCallUpdater'
-export type { BalancesAndAllowances } from './types/balances-and-allowances'
-export * from './utils/isBffSupportedNetwork'
+export { useIsSseFailed } from './state/isSseFailedAtom'
 
 // Types
+export type { BalancesAndAllowances } from './types/balances-and-allowances'
 export type { BalancesState } from './state/balancesAtom'
 export type { AllowancesState } from './hooks/useTokenAllowances'
+export type { SseBalancesState, UseSseBalancesParams } from './hooks/useSseBalances'
+
+// Utils
+export * from './utils/isBffSupportedNetwork'
 
 // Consts
 export { DEFAULT_BALANCES_STATE } from './state/balancesAtom'
diff --git a/libs/balances-and-allowances/src/state/isSseFailedAtom.ts b/libs/balances-and-allowances/src/state/isSseFailedAtom.ts
@@ -0,0 +1,12 @@
+import { useSetAtom } from 'jotai'
+import { atom, useAtomValue } from 'jotai/index'
+
+export const isSseFailedAtom = atom(false)
+
+export function useIsSseFailed(): boolean {
+  return useAtomValue(isSseFailedAtom)
+}
+
+export function useSetIsSseFailed(): (value: boolean) => void {
+  return useSetAtom(isSseFailedAtom)
+}
diff --git a/libs/balances-and-allowances/src/updaters/BalancesAndAllowancesUpdater.tsx b/libs/balances-and-allowances/src/updaters/BalancesAndAllowancesUpdater.tsx
@@ -2,7 +2,7 @@ import { ReactNode, useEffect, useMemo } from 'react'
 
 import { LpToken, NATIVE_CURRENCIES } from '@cowprotocol/common-const'
 import type { SupportedChainId } from '@cowprotocol/cow-sdk'
-import { useAllActiveTokens } from '@cowprotocol/tokens'
+import { useAllActiveTokens, useListsEnabledState } from '@cowprotocol/tokens'
 
 import ms from 'ms.macro'
 import { SWRConfiguration } from 'swr'
@@ -11,11 +11,13 @@ import { BalancesBffUpdater } from './BalancesBffUpdater'
 import { BalancesCacheUpdater } from './BalancesCacheUpdater'
 import { BalancesResetUpdater } from './BalancesResetUpdater'
 import { BalancesRpcCallUpdater } from './BalancesRpcCallUpdater'
+import { BalancesSseUpdater } from './BalancesSseUpdater'
 
 import { BASIC_MULTICALL_SWR_CONFIG } from '../consts'
 import { useNativeTokenBalance } from '../hooks/useNativeTokenBalance'
 import { useSwrConfigWithPauseForNetwork } from '../hooks/useSwrConfigWithPauseForNetwork'
 import { useUpdateTokenBalance } from '../hooks/useUpdateTokenBalance'
+import { useIsSseFailed } from '../state/isSseFailedAtom'
 
 // A small gap between balances and allowances refresh intervals is needed to avoid high load to the node at the same time
 const RPC_BALANCES_SWR_CONFIG: SWRConfiguration = { ...BASIC_MULTICALL_SWR_CONFIG, refreshInterval: ms`31s` }
@@ -27,8 +29,12 @@ export interface BalancesAndAllowancesUpdaterProps {
   chainId: SupportedChainId
   invalidateCacheTrigger: number
   excludedTokens: Set<string>
+  /** @deprecated Use isSseEnabled instead */
   isBffSwitchedOn: boolean
+  /** @deprecated Use isSseEnabled instead */
   isBffEnabled?: boolean
+  /** Enable SSE-based real-time balance updates */
+  isSseEnabled?: boolean
 }
 
 export function BalancesAndAllowancesUpdater({
@@ -38,8 +44,10 @@ export function BalancesAndAllowancesUpdater({
   isBffSwitchedOn,
   excludedTokens,
   isBffEnabled,
+  isSseEnabled = false,
 }: BalancesAndAllowancesUpdaterProps): ReactNode {
   const updateTokenBalance = useUpdateTokenBalance()
+  const isSseFailed = useIsSseFailed()
 
   const allTokens = useAllActiveTokens()
   const { data: nativeTokenBalance } = useNativeTokenBalance(account, chainId)
@@ -55,31 +63,57 @@ export function BalancesAndAllowancesUpdater({
     }, [])
   }, [allTokens, chainId])
 
+  // Get enabled token list URLs for SSE
+  const listsEnabledState = useListsEnabledState()
+  const tokensListsUrls = useMemo(() => {
+    return Object.entries(listsEnabledState)
+      .filter(([, isEnabled]) => isEnabled === true)
+      .map(([url]) => url)
+  }, [listsEnabledState])
+
   const rpcBalancesSwrConfig = useSwrConfigWithPauseForNetwork(chainId, account, RPC_BALANCES_SWR_CONFIG)
+
+  // Determine which updater to use
+  const hasSseTokenLists = tokensListsUrls.length > 0
+  const useSse = isSseEnabled && !isSseFailed && hasSseTokenLists
+  const useBff = !isSseEnabled && isBffEnabled
+  const useRpcFallback = (!isBffSwitchedOn || !isBffEnabled) && !useSse
+
   // Add native token balance to the store as well
   useEffect(() => {
-    if (isBffSwitchedOn) return
+    if (isBffSwitchedOn || isSseEnabled) return
 
     const nativeToken = NATIVE_CURRENCIES[chainId]
 
     if (nativeToken && nativeTokenBalance) {
       updateTokenBalance(nativeToken.address, nativeTokenBalance)
     }
-  }, [isBffSwitchedOn, nativeTokenBalance, chainId, updateTokenBalance])
-
-  const enableRpcFallback = !isBffSwitchedOn || !isBffEnabled
+  }, [isBffSwitchedOn, isSseEnabled, nativeTokenBalance, chainId, updateTokenBalance])
 
   return (
     <>
-      {isBffEnabled && (
+      {/* SSE-based real-time updates (preferred) */}
+      {useSse && (
+        <BalancesSseUpdater
+          account={account}
+          chainId={chainId}
+          tokenAddresses={tokenAddresses}
+          tokensListsUrls={tokensListsUrls}
+        />
+      )}
+
+      {/* Legacy BFF polling (deprecated, for backward compatibility) */}
+      {useBff && (
         <BalancesBffUpdater
           account={account}
           chainId={chainId}
           invalidateCacheTrigger={invalidateCacheTrigger}
           tokenAddresses={tokenAddresses}
         />
       )}
-      {enableRpcFallback && (
+
+      {/* RPC fallback when SSE/BFF fails or is disabled */}
+      {(useRpcFallback || isSseFailed) && (
         <BalancesRpcCallUpdater
           account={account}
           chainId={chainId}
@@ -88,6 +122,7 @@ export function BalancesAndAllowancesUpdater({
           setLoadingState
         />
       )}
+
       <BalancesResetUpdater chainId={chainId} account={account} />
       <BalancesCacheUpdater chainId={chainId} account={account} excludedTokens={excludedTokens} />
     </>
diff --git a/libs/balances-and-allowances/src/updaters/BalancesSseUpdater.tsx b/libs/balances-and-allowances/src/updaters/BalancesSseUpdater.tsx
diff --git a/libs/common-const/src/bff.ts b/libs/common-const/src/bff.ts
