feat: show helpful usage stats when the user is below certain thresholds

Author: brandonkachen

cli/src/chat.tsx

@@ -39,6 +39,7 @@ import { useSuggestionMenuHandlers } from './hooks/use-suggestion-menu-handlers'
 import { useTerminalDimensions } from './hooks/use-terminal-dimensions'
 import { useTheme } from './hooks/use-theme'
 import { useTimeout } from './hooks/use-timeout'
+import { useUsageMonitor } from './hooks/use-usage-monitor'
 
 import { useChatStore } from './state/chat-store'
 import { getInputModeConfig } from './utils/input-modes'
@@ -118,6 +119,9 @@ export const Chat = ({
   // Subscribe to ask_user bridge to trigger form display
   useAskUserBridge()
 
+  // Monitor usage data and auto-show banner when thresholds are crossed
+  useUsageMonitor()
+
   const {
     inputValue,
     cursorPosition,

cli/src/components/chat-input-bar.tsx

@@ -20,10 +20,16 @@ import type { InputMode } from '../utils/input-modes'
 
 type Theme = ReturnType<typeof useTheme>
 
-const InputModeBanner = ({ inputMode }: { inputMode: InputMode }) => {
+const InputModeBanner = ({
+  inputMode,
+  usageBannerShowTime,
+}: {
+  inputMode: InputMode
+  usageBannerShowTime: number
+}) => {
   switch (inputMode) {
     case 'usage':
-      return <UsageBanner />
+      return <UsageBanner showTime={usageBannerShowTime} />
     case 'referral':
       return <ReferralBanner />
     default:
@@ -104,6 +110,16 @@ export const ChatInputBar = ({
   const inputMode = useChatStore((state) => state.inputMode)
   const setInputMode = useChatStore((state) => state.setInputMode)
 
+  const [usageBannerShowTime, setUsageBannerShowTime] = React.useState(
+    () => Date.now(),
+  )
+
+  React.useEffect(() => {
+    if (inputMode === 'usage') {
+      setUsageBannerShowTime(Date.now())
+    }
+  }, [inputMode])
+
   const modeConfig = getInputModeConfig(inputMode)
   const askUserState = useChatStore((state) => state.askUserState)
   const updateAskUserAnswer = useChatStore((state) => state.updateAskUserAnswer)
@@ -336,7 +352,10 @@ export const ChatInputBar = ({
           </box>
         </box>
       </box>
-      <InputModeBanner inputMode={inputMode} />
+      <InputModeBanner
+        inputMode={inputMode}
+        usageBannerShowTime={usageBannerShowTime}
+      />
     </>
   )
 }

cli/src/components/usage-banner.tsx

@@ -1,32 +1,41 @@
-import type { UserState } from '@codebuff/common/old-constants'
-import { useQuery } from '@tanstack/react-query'
-import React, { useEffect, useRef, useState } from 'react'
+import { useQuery, useQueryClient } from '@tanstack/react-query'
+import React, { useEffect } from 'react'
 
 import { BannerWrapper } from './banner-wrapper'
 import { useTheme } from '../hooks/use-theme'
 import { usageQueryKeys, useUsageQuery } from '../hooks/use-usage-query'
 import { useChatStore } from '../state/chat-store'
-import { getAuthToken } from '../utils/auth'
 import {
   getBannerColorLevel,
   generateUsageBannerText,
   generateLoadingBannerText,
-  shouldAutoShowBanner,
 } from '../utils/usage-banner-state'
 
 const MANUAL_SHOW_TIMEOUT = 60 * 1000 // 1 minute
-const AUTO_SHOW_TIMEOUT = 5 * 60 * 1000 // 5 minutes
+const USAGE_POLL_INTERVAL = 30 * 1000 // 30 seconds
 
-export const UsageBanner = () => {
+export const UsageBanner = ({ showTime }: { showTime: number }) => {
   const theme = useTheme()
+  const queryClient = useQueryClient()
   const sessionCreditsUsed = useChatStore((state) => state.sessionCreditsUsed)
-  const isChainInProgress = useChatStore((state) => state.isChainInProgress)
   const setInputMode = useChatStore((state) => state.setInputMode)
 
-  const [isAutoShown, setIsAutoShown] = useState(false)
-  const lastWarnedStateRef = useRef<UserState | null>(null)
+  const {
+    data: apiData,
+    isLoading,
+    isFetching,
+  } = useUsageQuery({
+    enabled: true,
+  })
 
-  const { data: apiData, isLoading, isFetching } = useUsageQuery({ enabled: true })
+  // Manual polling using setInterval - TanStack Query's refetchInterval doesn't work
+  // reliably in terminal environments even with focusManager configuration
+  useEffect(() => {
+    const interval = setInterval(() => {
+      queryClient.invalidateQueries({ queryKey: usageQueryKeys.current() })
+    }, USAGE_POLL_INTERVAL)
+    return () => clearInterval(interval)
+  }, [queryClient])
 
   const { data: cachedUsageData } = useQuery<{
     type: 'usage-response'
@@ -39,34 +48,13 @@ export const UsageBanner = () => {
     enabled: false,
   })
 
-  // Credit warning monitoring logic
-  useEffect(() => {
-    const authToken = getAuthToken()
-    const decision = shouldAutoShowBanner(
-      isChainInProgress,
-      !!authToken,
-      cachedUsageData?.remainingBalance ?? null,
-      lastWarnedStateRef.current,
-    )
-
-    if (decision.newWarningState !== lastWarnedStateRef.current) {
-      lastWarnedStateRef.current = decision.newWarningState
-    }
-
-    if (decision.shouldShow) {
-      setIsAutoShown(true)
-    }
-  }, [isChainInProgress, cachedUsageData])
-
-  // Auto-hide effect
+  // Auto-hide after timeout
   useEffect(() => {
-    const timeout = isAutoShown ? AUTO_SHOW_TIMEOUT : MANUAL_SHOW_TIMEOUT
     const timer = setTimeout(() => {
       setInputMode('default')
-      setIsAutoShown(false)
-    }, timeout)
+    }, MANUAL_SHOW_TIMEOUT)
     return () => clearTimeout(timer)
-  }, [isAutoShown, setInputMode])
+  }, [showTime, setInputMode])
 
   const activeData = apiData || cachedUsageData
   const isLoadingData = isLoading || isFetching

cli/src/hooks/use-send-message.ts

@@ -2,6 +2,8 @@ import {
   MAX_RETRIES_PER_MESSAGE,
   RETRY_BACKOFF_BASE_DELAY_MS,
   RETRY_BACKOFF_MAX_DELAY_MS,
+  isPaymentRequiredError,
+  ErrorCodes,
 } from '@codebuff/sdk'
 import { useQueryClient } from '@tanstack/react-query'
 import { has, isEqual } from 'lodash'
@@ -1769,15 +1771,53 @@ export const useSendMessage = ({
         })
 
         if (!runState.output || runState.output.type === 'error') {
-          logger.warn(
-            {
-              errorMessage:
-                runState.output?.type === 'error'
-                  ? runState.output.message
-                  : 'No output from agent run',
-            },
-            'Agent run failed',
-          )
+          const errorOutput = runState.output?.type === 'error' ? runState.output : null
+          const errorMessage = errorOutput?.message ?? 'No output from agent run'
+
+          logger.warn({ errorMessage, errorCode: errorOutput?.errorCode }, 'Agent run failed')
+
+          // Check if this is an out-of-credits error using the error code
+          const isOutOfCredits = errorOutput?.errorCode === ErrorCodes.PAYMENT_REQUIRED
+
+          if (isOutOfCredits) {
+            const appUrl = process.env.NEXT_PUBLIC_CODEBUFF_APP_URL || 'https://codebuff.com'
+            const paymentErrorMessage =
+              errorOutput?.message ??
+              `Out of credits. Please add credits at ${appUrl}/usage`
+            applyMessageUpdate((prev) =>
+              prev.map((msg) => {
+                if (msg.id !== aiMessageId) return msg
+                return {
+                  ...msg,
+                  content: paymentErrorMessage,
+                  blocks: undefined, // Clear blocks so content renders
+                  isComplete: true,
+                }
+              }),
+            )
+            // Show the usage banner so user can see their balance and renewal date
+            useChatStore.getState().setInputMode('usage')
+            // Refresh usage data to show current state
+            queryClient.invalidateQueries({ queryKey: usageQueryKeys.current() })
+          } else {
+            // Generic error - display the error message directly from SDK
+            applyMessageUpdate((prev) =>
+              prev.map((msg) => {
+                if (msg.id !== aiMessageId) return msg
+                return {
+                  ...msg,
+                  content: `**Error:** ${errorMessage}`,
+                  blocks: undefined, // Clear blocks so content renders
+                  isComplete: true,
+                }
+              }),
+            )
+          }
+
+          setStreamStatus('idle')
+          setCanProcessQueue(true)
+          updateChainInProgress(false)
+          timerController.stop('error')
           return
         }
 
@@ -1835,6 +1875,35 @@ export const useSendMessage = ({
 
         const errorMessage =
           error instanceof Error ? error.message : 'Unknown error occurred'
+
+        // Handle payment required (out of credits) specially
+        if (isPaymentRequiredError(error)) {
+          const appUrl = process.env.NEXT_PUBLIC_CODEBUFF_APP_URL || 'https://codebuff.com'
+          const paymentErrorMessage =
+            error instanceof Error && error.message
+              ? error.message
+              : `Out of credits. Please add credits at ${appUrl}/usage`
+
+          applyMessageUpdate((prev) =>
+            prev.map((msg) => {
+              if (msg.id !== aiMessageId) {
+                return msg
+              }
+              return {
+                ...msg,
+                content: paymentErrorMessage,
+                blocks: undefined, // Clear blocks so content renders
+                isComplete: true,
+              }
+            }),
+          )
+          // Show the usage banner so user can see their balance and renewal date
+          useChatStore.getState().setInputMode('usage')
+          // Refresh usage data to show current state
+          queryClient.invalidateQueries({ queryKey: usageQueryKeys.current() })
+          return
+        }
+
         applyMessageUpdate((prev) =>
           prev.map((msg) => {
             if (msg.id !== aiMessageId) {

cli/src/hooks/use-usage-monitor.ts

@@ -0,0 +1,50 @@
+import { useEffect, useRef } from 'react'
+
+import { getAuthToken } from '../utils/auth'
+import { useChatStore } from '../state/chat-store'
+import { useUsageQuery } from './use-usage-query'
+import { shouldAutoShowBanner } from '../utils/usage-banner-state'
+
+/**
+ * Hook that monitors usage data and auto-shows the usage banner
+ * when credit thresholds are crossed.
+ * 
+ * This should be placed in a component that's always mounted (like Chat)
+ * so monitoring happens continuously, not just when the banner is visible.
+ */
+export function useUsageMonitor() {
+  const isChainInProgress = useChatStore((state) => state.isChainInProgress)
+  const sessionCreditsUsed = useChatStore((state) => state.sessionCreditsUsed)
+  const setInputMode = useChatStore((state) => state.setInputMode)
+  const lastWarnedThresholdRef = useRef<number | null>(null)
+
+  // Query usage data - this will refetch when invalidated after message completion
+  const { data: usageData } = useUsageQuery({ enabled: true })
+
+  useEffect(() => {
+    // Only show after user has sent at least one message (to avoid overwhelming on app start)
+    if (sessionCreditsUsed === 0) {
+      return
+    }
+
+    const authToken = getAuthToken()
+    const remainingBalance = usageData?.remainingBalance ?? null
+
+    const decision = shouldAutoShowBanner(
+      isChainInProgress,
+      !!authToken,
+      remainingBalance,
+      lastWarnedThresholdRef.current,
+    )
+
+    // Update the last warned threshold
+    if (decision.newWarningThreshold !== lastWarnedThresholdRef.current) {
+      lastWarnedThresholdRef.current = decision.newWarningThreshold
+    }
+
+    // Show the usage banner if we should
+    if (decision.shouldShow) {
+      setInputMode('usage')
+    }
+  }, [isChainInProgress, usageData, sessionCreditsUsed, setInputMode])
+}

cli/src/hooks/use-usage-query.ts

@@ -65,26 +65,35 @@ export async function fetchUsageData({
 export interface UseUsageQueryDeps {
   logger?: Logger
   enabled?: boolean
+  refetchInterval?: number | false
+  refetchIntervalInBackground?: boolean
 }
 
 /**
  * Hook to fetch usage data from the API
  * Returns TanStack Query result directly - no store synchronization needed
  */
 export function useUsageQuery(deps: UseUsageQueryDeps = {}) {
-  const { logger = defaultLogger, enabled = true } = deps
+  const { 
+    logger = defaultLogger, 
+    enabled = true, 
+    refetchInterval = false,
+    refetchIntervalInBackground = false,
+  } = deps
   const authToken = getAuthToken()
 
   return useQuery({
     queryKey: usageQueryKeys.current(),
     queryFn: () => fetchUsageData({ authToken: authToken!, logger }),
     enabled: enabled && !!authToken,
-    staleTime: 30 * 1000, // 30 seconds - usage data changes as user makes requests
+    staleTime: 0, // Always consider data stale for immediate refetching
     gcTime: 5 * 60 * 1000, // 5 minutes
     retry: false, // Don't retry failed usage queries
-    refetchOnMount: false, // Don't auto-refetch on mount
+    refetchOnMount: 'always', // Always refetch on mount to get fresh data when banner opens
     refetchOnWindowFocus: false, // CLI doesn't have window focus
     refetchOnReconnect: false, // Don't auto-refetch on reconnect
+    refetchInterval, // Poll at specified interval (when banner is visible)
+    refetchIntervalInBackground, // Required for terminal environments without browser visibility API
   })
 }
 

cli/src/index.tsx

@@ -8,7 +8,7 @@ import { createRoot } from '@opentui/react'
 import { API_KEY_ENV_VAR } from '@codebuff/common/old-constants'
 import { getProjectFileTree } from '@codebuff/common/project-file-tree'
 import { validateAgents } from '@codebuff/sdk'
-import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
+import { QueryClient, QueryClientProvider, focusManager } from '@tanstack/react-query'
 import { Command } from 'commander'
 import React from 'react'
 
@@ -50,6 +50,15 @@ function loadPackageVersion(): string {
   return 'dev'
 }
 
+// Configure TanStack Query's focusManager for terminal environments
+// This is required because there's no browser visibility API in terminal apps
+// Without this, refetchInterval won't work because TanStack Query thinks the app is "unfocused"
+focusManager.setEventListener(() => {
+  // No-op: no event listeners in CLI environment (no window focus/visibility events)
+  return () => {}
+})
+focusManager.setFocused(true)
+
 function createQueryClient(): QueryClient {
   return new QueryClient({
     defaultOptions: {