refactor(cli): simplify usage query by removing unnecessary useEffect

- Remove useEffect from useUsageQuery - now returns TanStack Query directly
- Update UsageBanner to read from query result instead of chat store
- Remove usageData and setUsageData from chat store (no longer needed)
- Update all tests to reflect new architecture
- Mark fetchAndUpdateUsage as deprecated

Data now flows directly from TanStack Query to components without
unnecessary store synchronization, improving code clarity and reactivity.

Author: brandonkachen

cli/src/__tests__/integration/usage-refresh-on-completion.test.ts

@@ -1,4 +1,4 @@
-import { describe, test, expect, beforeEach, afterEach, mock } from 'bun:test'
+import { describe, test, expect, beforeEach, afterEach, mock, spyOn } from 'bun:test'
 import { QueryClient } from '@tanstack/react-query'
 
 import { useChatStore } from '../../state/chat-store'
@@ -20,11 +20,10 @@ import * as authModule from '../../utils/auth'
  */
 describe('Usage Refresh on SDK Completion', () => {
   const originalFetch = globalThis.fetch
-  const originalGetAuthToken = authModule.getAuthToken
   const originalEnv = process.env.NEXT_PUBLIC_CODEBUFF_APP_URL
 
   let queryClient: QueryClient
-  let getAuthTokenMock: ReturnType<typeof mock>
+  let getAuthTokenSpy: ReturnType<typeof spyOn>
 
   beforeEach(() => {
     process.env.NEXT_PUBLIC_CODEBUFF_APP_URL = 'https://test.codebuff.local'
@@ -40,8 +39,7 @@ describe('Usage Refresh on SDK Completion', () => {
     })
 
     // Mock auth token
-    getAuthTokenMock = mock(() => 'test-token')
-    authModule.getAuthToken = getAuthTokenMock as any
+    getAuthTokenSpy = spyOn(authModule, 'getAuthToken').mockReturnValue('test-token')
 
     // Mock successful API response
     globalThis.fetch = mock(async () =>
@@ -59,7 +57,7 @@ describe('Usage Refresh on SDK Completion', () => {
 
   afterEach(() => {
     globalThis.fetch = originalFetch
-    authModule.getAuthToken = originalGetAuthToken
+    getAuthTokenSpy.mockRestore()
     process.env.NEXT_PUBLIC_CODEBUFF_APP_URL = originalEnv
     mock.restore()
   })
@@ -160,7 +158,7 @@ describe('Usage Refresh on SDK Completion', () => {
 
   describe('unauthenticated scenarios', () => {
     test('should not fetch when no auth token', () => {
-      getAuthTokenMock.mockReturnValue(undefined)
+      getAuthTokenSpy.mockReturnValue(undefined)
       useChatStore.getState().setIsUsageVisible(true)
 
       const fetchMock = mock(globalThis.fetch)

cli/src/components/usage-banner.tsx

@@ -15,11 +15,20 @@ export const UsageBanner = () => {
   const { terminalWidth } = useTerminalDimensions()
   const theme = useTheme()
   const isUsageVisible = useChatStore((state) => state.isUsageVisible)
-  const usageData = useChatStore((state) => state.usageData)
+  const sessionCreditsUsed = useChatStore((state) => state.sessionCreditsUsed)
   const setIsUsageVisible = useChatStore((state) => state.setIsUsageVisible)
 
   // Fetch usage data when banner is visible
-  useUsageQuery({ enabled: isUsageVisible })
+  const { data: apiData } = useUsageQuery({ enabled: isUsageVisible })
+
+  // Transform API data to usage data format
+  const usageData = apiData
+    ? {
+        sessionUsage: sessionCreditsUsed,
+        remainingBalance: apiData.remainingBalance,
+        nextQuotaReset: apiData.next_quota_reset,
+      }
+    : null
 
   // Auto-hide banner after 60 seconds
   useEffect(() => {

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

@@ -1,13 +1,23 @@
-import { describe, test, expect, beforeEach, afterEach, mock } from 'bun:test'
-import { renderHook, waitFor } from '@testing-library/react'
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
+import { renderHook, waitFor } from '@testing-library/react'
+import {
+  describe,
+  test,
+  expect,
+  beforeEach,
+  afterEach,
+  mock,
+  spyOn,
+} from 'bun:test'
 import React from 'react'
 
-import { fetchUsageData, useUsageQuery, useRefreshUsage } from '../use-usage-query'
 import { useChatStore } from '../../state/chat-store'
 import * as authModule from '../../utils/auth'
-
-import type { UsageData } from '../../state/chat-store'
+import {
+  fetchUsageData,
+  useUsageQuery,
+  useRefreshUsage,
+} from '../use-usage-query'
 
 describe('fetchUsageData', () => {
   const originalFetch = globalThis.fetch
@@ -25,18 +35,19 @@ describe('fetchUsageData', () => {
 
   test('should fetch usage data successfully', async () => {
     const mockResponse = {
-      type: 'usage-response',
+      type: 'usage-response' as const,
       usage: 100,
       remainingBalance: 500,
       balanceBreakdown: { free: 300, paid: 200 },
       next_quota_reset: '2024-02-01T00:00:00.000Z',
     }
 
-    globalThis.fetch = mock(async () =>
-      new Response(JSON.stringify(mockResponse), {
-        status: 200,
-        headers: { 'Content-Type': 'application/json' },
-      }),
+    globalThis.fetch = mock(
+      async () =>
+        new Response(JSON.stringify(mockResponse), {
+          status: 200,
+          headers: { 'Content-Type': 'application/json' },
+        }),
     )
 
     const result = await fetchUsageData({ authToken: 'test-token' })
@@ -63,7 +74,7 @@ describe('fetchUsageData', () => {
 
 describe('useUsageQuery', () => {
   let queryClient: QueryClient
-  const originalGetAuthToken = authModule.getAuthToken
+  let getAuthTokenSpy: ReturnType<typeof spyOn>
   const originalEnv = process.env.NEXT_PUBLIC_CODEBUFF_APP_URL
 
   function createWrapper() {
@@ -86,13 +97,15 @@ describe('useUsageQuery', () => {
   })
 
   afterEach(() => {
-    authModule.getAuthToken = originalGetAuthToken
+    getAuthTokenSpy?.mockRestore()
     process.env.NEXT_PUBLIC_CODEBUFF_APP_URL = originalEnv
     mock.restore()
   })
 
-  test('should fetch and update store when enabled', async () => {
-    authModule.getAuthToken = mock(() => 'test-token')
+  test('should fetch data when enabled', async () => {
+    getAuthTokenSpy = spyOn(authModule, 'getAuthToken').mockReturnValue(
+      'test-token',
+    )
 
     const mockResponse = {
       type: 'usage-response' as const,
@@ -101,11 +114,12 @@ describe('useUsageQuery', () => {
       next_quota_reset: '2024-02-01T00:00:00.000Z',
     }
 
-    globalThis.fetch = mock(async () =>
-      new Response(JSON.stringify(mockResponse), {
-        status: 200,
-        headers: { 'Content-Type': 'application/json' },
-      }),
+    globalThis.fetch = mock(
+      async () =>
+        new Response(JSON.stringify(mockResponse), {
+          status: 200,
+          headers: { 'Content-Type': 'application/json' },
+        }),
     )
 
     const { result } = renderHook(() => useUsageQuery(), {
@@ -118,7 +132,9 @@ describe('useUsageQuery', () => {
   })
 
   test('should not fetch when disabled', async () => {
-    authModule.getAuthToken = mock(() => 'test-token')
+    getAuthTokenSpy = spyOn(authModule, 'getAuthToken').mockReturnValue(
+      'test-token',
+    )
     const fetchMock = mock(async () => new Response('{}'))
     globalThis.fetch = fetchMock
 
@@ -133,7 +149,9 @@ describe('useUsageQuery', () => {
   })
 
   test('should not fetch when no auth token', async () => {
-    authModule.getAuthToken = mock(() => undefined)
+    getAuthTokenSpy = spyOn(authModule, 'getAuthToken').mockReturnValue(
+      undefined,
+    )
     const fetchMock = mock(async () => new Response('{}'))
     globalThis.fetch = fetchMock
 

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

@@ -1,7 +1,5 @@
-import React from 'react'
 import { useQuery, useQueryClient } from '@tanstack/react-query'
 
-import { useChatStore } from '../state/chat-store'
 import { getAuthToken } from '../utils/auth'
 import { logger as defaultLogger } from '../utils/logger'
 
@@ -70,17 +68,14 @@ export interface UseUsageQueryDeps {
 }
 
 /**
- * Hook to fetch and manage usage data
- * Automatically updates the chat store when new data arrives
+ * 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 setUsageData = useChatStore((state) => state.setUsageData)
-  const sessionCreditsUsed = useChatStore((state) => state.sessionCreditsUsed)
-
   const authToken = getAuthToken()
 
-  const query = useQuery({
+  return useQuery({
     queryKey: usageQueryKeys.current(),
     queryFn: () => fetchUsageData({ authToken: authToken!, logger }),
     enabled: enabled && !!authToken,
@@ -91,19 +86,6 @@ export function useUsageQuery(deps: UseUsageQueryDeps = {}) {
     refetchOnWindowFocus: false, // CLI doesn't have window focus
     refetchOnReconnect: false, // Don't auto-refetch on reconnect
   })
-
-  // Update store when data changes (replaces deprecated onSuccess)
-  React.useEffect(() => {
-    if (query.data) {
-      setUsageData({
-        sessionUsage: sessionCreditsUsed,
-        remainingBalance: query.data.remainingBalance,
-        nextQuotaReset: query.data.next_quota_reset,
-      })
-    }
-  }, [query.data, sessionCreditsUsed, setUsageData])
-
-  return query
 }
 
 /**

cli/src/state/chat-store.ts

@@ -31,7 +31,6 @@ export type ChatStoreState = {
   lastMessageMode: AgentMode | null
   sessionCreditsUsed: number
   runState: RunState | null
-  usageData: UsageData | null
   isUsageVisible: boolean
 }
 
@@ -61,19 +60,12 @@ type ChatStoreActions = {
   setLastMessageMode: (mode: AgentMode | null) => void
   addSessionCredits: (credits: number) => void
   setRunState: (runState: RunState | null) => void
-  setUsageData: (data: UsageData | null) => void
   setIsUsageVisible: (visible: boolean) => void
   reset: () => void
 }
 
 type ChatStore = ChatStoreState & ChatStoreActions
 
-export interface UsageData {
-  sessionUsage: number
-  remainingBalance: number | null
-  nextQuotaReset: string | null
-}
-
 const initialState: ChatStoreState = {
   messages: [],
   streamingAgents: new Set<string>(),
@@ -91,7 +83,6 @@ const initialState: ChatStoreState = {
   lastMessageMode: null,
   sessionCreditsUsed: 0,
   runState: null,
-  usageData: null,
   isUsageVisible: false,
 }
 
@@ -196,11 +187,6 @@ export const useChatStore = create<ChatStore>()(
         state.runState = runState ? castDraft(runState) : null
       }),
 
-    setUsageData: (data) =>
-      set((state) => {
-        state.usageData = data
-      }),
-
     setIsUsageVisible: (visible) =>
       set((state) => {
         state.isUsageVisible = visible
@@ -226,7 +212,6 @@ export const useChatStore = create<ChatStore>()(
         state.runState = initialState.runState
           ? castDraft(initialState.runState)
           : null
-        state.usageData = initialState.usageData
         state.isUsageVisible = initialState.isUsageVisible
       }),
   })),