Skip to content

Commit 55ff21a

Browse files
brandonkachenbrandonkachen
committed
feat(cli): add real-time usage refresh when banner is visible
- Extract usage fetching logic into reusable utility function - Auto-refresh usage data after SDK run completion (only if banner visible) - Add comprehensive unit and integration tests (27 tests total) - Use dependency injection pattern for better testability
1 parent c4239ec commit 55ff21a

File tree

5 files changed

+662
-58
lines changed

5 files changed

+662
-58
lines changed

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

Lines changed: 257 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,257 @@
1+
import { describe, test, expect, beforeEach, afterEach, mock } from 'bun:test'
2+
3+
import { useChatStore } from '../../state/chat-store'
4+
import { fetchAndUpdateUsage } from '../../utils/fetch-usage'
5+
6+
import type { Logger } from '@codebuff/common/types/contracts/logger'
7+
8+
/**
9+
* Integration test for usage refresh on SDK run completion
10+
*
11+
* This test verifies the complete lifecycle:
12+
* 1. User opens usage banner (isUsageVisible = true)
13+
* 2. SDK run completes successfully
14+
* 3. Usage data is refreshed automatically
15+
* 4. Banner shows updated credit balance
16+
*
17+
* Also tests:
18+
* - No refresh when banner is closed (isUsageVisible = false)
19+
* - Error handling during background refresh
20+
* - Multiple sequential runs with banner open
21+
*/
22+
describe('Usage Refresh on SDK Completion', () => {
23+
let loggerMock: Logger
24+
let fetchMock: ReturnType<typeof mock>
25+
26+
const createMockResponse = (data: any, status: number = 200) => {
27+
return new Response(JSON.stringify(data), {
28+
status,
29+
headers: { 'Content-Type': 'application/json' },
30+
})
31+
}
32+
33+
beforeEach(() => {
34+
useChatStore.getState().reset()
35+
36+
loggerMock = {
37+
info: mock(() => {}),
38+
error: mock(() => {}),
39+
warn: mock(() => {}),
40+
debug: mock(() => {}),
41+
}
42+
43+
fetchMock = mock(async () =>
44+
createMockResponse({
45+
type: 'usage-response',
46+
usage: 100,
47+
remainingBalance: 850,
48+
next_quota_reset: '2024-03-01T00:00:00.000Z',
49+
}),
50+
)
51+
})
52+
53+
afterEach(() => {
54+
mock.restore()
55+
})
56+
57+
describe('banner visible scenarios', () => {
58+
test('should refresh usage data when banner is visible and run completes', async () => {
59+
useChatStore.getState().setIsUsageVisible(true)
60+
61+
const isUsageVisible = useChatStore.getState().isUsageVisible
62+
if (isUsageVisible) {
63+
await fetchAndUpdateUsage({
64+
getAuthToken: () => 'test-token',
65+
logger: loggerMock,
66+
fetch: fetchMock as any,
67+
})
68+
}
69+
70+
expect(fetchMock).toHaveBeenCalledTimes(1)
71+
const usageData = useChatStore.getState().usageData
72+
expect(usageData?.remainingBalance).toBe(850)
73+
})
74+
75+
test('should not show banner after background refresh', async () => {
76+
useChatStore.getState().setIsUsageVisible(true)
77+
78+
await fetchAndUpdateUsage({
79+
showBanner: false,
80+
getAuthToken: () => 'test-token',
81+
logger: loggerMock,
82+
fetch: fetchMock as any,
83+
})
84+
85+
expect(useChatStore.getState().isUsageVisible).toBe(true)
86+
})
87+
88+
test('should refresh multiple times for sequential runs', async () => {
89+
useChatStore.getState().setIsUsageVisible(true)
90+
91+
for (let i = 0; i < 3; i++) {
92+
if (useChatStore.getState().isUsageVisible) {
93+
await fetchAndUpdateUsage({
94+
getAuthToken: () => 'test-token',
95+
logger: loggerMock,
96+
fetch: fetchMock as any,
97+
})
98+
}
99+
}
100+
101+
expect(fetchMock).toHaveBeenCalledTimes(3)
102+
})
103+
104+
test('should update usage data with fresh values from API', async () => {
105+
useChatStore.getState().setUsageData({
106+
sessionUsage: 100,
107+
remainingBalance: 1000,
108+
nextQuotaReset: '2024-02-01T00:00:00.000Z',
109+
})
110+
useChatStore.getState().setIsUsageVisible(true)
111+
112+
await fetchAndUpdateUsage({
113+
getAuthToken: () => 'test-token',
114+
logger: loggerMock,
115+
fetch: fetchMock as any,
116+
})
117+
118+
const updatedData = useChatStore.getState().usageData
119+
expect(updatedData).not.toBeNull()
120+
expect(updatedData?.remainingBalance).toBe(850)
121+
expect(updatedData?.nextQuotaReset).toBe('2024-03-01T00:00:00.000Z')
122+
})
123+
})
124+
125+
describe('banner not visible scenarios', () => {
126+
test('should NOT refresh when banner is not visible', async () => {
127+
useChatStore.getState().setIsUsageVisible(false)
128+
129+
const isUsageVisible = useChatStore.getState().isUsageVisible
130+
if (isUsageVisible) {
131+
await fetchAndUpdateUsage({
132+
getAuthToken: () => 'test-token',
133+
logger: loggerMock,
134+
fetch: fetchMock as any,
135+
})
136+
}
137+
138+
expect(fetchMock).not.toHaveBeenCalled()
139+
})
140+
141+
test('should not refresh if banner was closed before run completed', async () => {
142+
useChatStore.getState().setIsUsageVisible(true)
143+
useChatStore.getState().setIsUsageVisible(false)
144+
145+
const isUsageVisible = useChatStore.getState().isUsageVisible
146+
if (isUsageVisible) {
147+
await fetchAndUpdateUsage({
148+
getAuthToken: () => 'test-token',
149+
logger: loggerMock,
150+
fetch: fetchMock as any,
151+
})
152+
}
153+
154+
expect(fetchMock).not.toHaveBeenCalled()
155+
})
156+
})
157+
158+
describe('error handling during refresh', () => {
159+
test('should handle API errors gracefully without crashing', async () => {
160+
fetchMock.mockImplementation(async () =>
161+
new Response('Server Error', { status: 500 }),
162+
)
163+
164+
useChatStore.getState().setIsUsageVisible(true)
165+
166+
await expect(
167+
fetchAndUpdateUsage({
168+
getAuthToken: () => 'test-token',
169+
logger: loggerMock,
170+
fetch: fetchMock as any,
171+
}),
172+
).resolves.toBe(false)
173+
174+
expect(useChatStore.getState().isUsageVisible).toBe(true)
175+
})
176+
177+
test('should handle network errors during background refresh', async () => {
178+
fetchMock.mockImplementation(async () => {
179+
throw new Error('Network failure')
180+
})
181+
182+
useChatStore.getState().setIsUsageVisible(true)
183+
184+
const result = await fetchAndUpdateUsage({
185+
getAuthToken: () => 'test-token',
186+
logger: loggerMock,
187+
fetch: fetchMock as any,
188+
})
189+
190+
expect(result).toBe(false)
191+
expect(loggerMock.error).toHaveBeenCalled()
192+
})
193+
194+
test('should continue working after failed refresh', async () => {
195+
useChatStore.getState().setIsUsageVisible(true)
196+
197+
fetchMock.mockImplementationOnce(async () =>
198+
new Response('Error', { status: 500 }),
199+
)
200+
await fetchAndUpdateUsage({
201+
getAuthToken: () => 'test-token',
202+
logger: loggerMock,
203+
fetch: fetchMock as any,
204+
})
205+
206+
fetchMock.mockImplementationOnce(async () =>
207+
createMockResponse({
208+
type: 'usage-response',
209+
usage: 200,
210+
remainingBalance: 800,
211+
next_quota_reset: null,
212+
}),
213+
)
214+
215+
const result = await fetchAndUpdateUsage({
216+
getAuthToken: () => 'test-token',
217+
logger: loggerMock,
218+
fetch: fetchMock as any,
219+
})
220+
221+
expect(result).toBe(true)
222+
})
223+
})
224+
225+
describe('unauthenticated user scenarios', () => {
226+
test('should not refresh when user is not authenticated', async () => {
227+
useChatStore.getState().setIsUsageVisible(true)
228+
229+
const result = await fetchAndUpdateUsage({
230+
getAuthToken: () => undefined,
231+
logger: loggerMock,
232+
fetch: fetchMock as any,
233+
})
234+
235+
expect(result).toBe(false)
236+
})
237+
})
238+
239+
describe('session credits tracking', () => {
240+
test('should include current session credits in refreshed data', async () => {
241+
useChatStore.getState().addSessionCredits(75)
242+
useChatStore.getState().addSessionCredits(25)
243+
244+
expect(useChatStore.getState().sessionCreditsUsed).toBe(100)
245+
246+
useChatStore.getState().setIsUsageVisible(true)
247+
await fetchAndUpdateUsage({
248+
getAuthToken: () => 'test-token',
249+
logger: loggerMock,
250+
fetch: fetchMock as any,
251+
})
252+
253+
const usageData = useChatStore.getState().usageData
254+
expect(usageData?.sessionUsage).toBe(100)
255+
})
256+
})
257+
})

‎cli/src/commands/usage.ts‎

Lines changed: 6 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,13 @@
1-
import { WEBSITE_URL } from '@codebuff/sdk'
2-
3-
import { useChatStore } from '../state/chat-store'
1+
import { fetchAndUpdateUsage } from '../utils/fetch-usage'
42
import { getAuthToken } from '../utils/auth'
5-
import { logger } from '../utils/logger'
63
import { getSystemMessage } from '../utils/message-history'
74

85
import type { PostUserMessageFn } from '../types/contracts/send-message'
96

10-
interface UsageResponse {
11-
type: 'usage-response'
12-
usage: number
13-
remainingBalance: number | null
14-
balanceBreakdown?: Record<string, number>
15-
next_quota_reset: string | null
16-
}
17-
187
export async function handleUsageCommand(): Promise<{
198
postUserMessage: PostUserMessageFn
209
}> {
2110
const authToken = getAuthToken()
22-
const sessionCreditsUsed = useChatStore.getState().sessionCreditsUsed
2311

2412
if (!authToken) {
2513
const postUserMessage: PostUserMessageFn = (prev) => [
@@ -29,56 +17,16 @@ export async function handleUsageCommand(): Promise<{
2917
return { postUserMessage }
3018
}
3119

32-
try {
33-
const response = await fetch(`${WEBSITE_URL}/api/v1/usage`, {
34-
method: 'POST',
35-
headers: {
36-
'Content-Type': 'application/json',
37-
},
38-
body: JSON.stringify({
39-
fingerprintId: 'cli-usage',
40-
authToken,
41-
}),
42-
})
43-
44-
if (!response.ok) {
45-
const errorText = await response.text().catch(() => 'Unknown error')
46-
logger.error(
47-
{ status: response.status, errorText },
48-
'Usage request failed',
49-
)
50-
51-
const postUserMessage: PostUserMessageFn = (prev) => [
52-
...prev,
53-
getSystemMessage(`Failed to fetch usage data: ${errorText}`),
54-
]
55-
return { postUserMessage }
56-
}
57-
58-
const data = (await response.json()) as UsageResponse
59-
60-
useChatStore.getState().setUsageData({
61-
sessionUsage: sessionCreditsUsed,
62-
remainingBalance: data.remainingBalance,
63-
nextQuotaReset: data.next_quota_reset,
64-
})
65-
useChatStore.getState().setIsUsageVisible(true)
66-
67-
const postUserMessage: PostUserMessageFn = (prev) => prev
68-
return { postUserMessage }
69-
} catch (error) {
70-
logger.error(
71-
{
72-
error: error instanceof Error ? error.message : String(error),
73-
errorStack: error instanceof Error ? error.stack : undefined,
74-
},
75-
'Error checking usage',
76-
)
20+
const success = await fetchAndUpdateUsage({ showBanner: true })
7721

22+
if (!success) {
7823
const postUserMessage: PostUserMessageFn = (prev) => [
7924
...prev,
8025
getSystemMessage('Error checking usage. Please try again later.'),
8126
]
8227
return { postUserMessage }
8328
}
29+
30+
const postUserMessage: PostUserMessageFn = (prev) => prev
31+
return { postUserMessage }
8432
}

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import {
99
} from '../utils/constants'
1010
import { createValidationErrorBlocks } from '../utils/create-validation-error-blocks'
1111
import { getErrorObject } from '../utils/error'
12+
import { fetchAndUpdateUsage } from '../utils/fetch-usage'
1213
import { formatTimestamp } from '../utils/helpers'
1314
import { loadAgentDefinitions } from '../utils/load-agent-definitions'
1415
import { getLoadedAgentsData } from '../utils/local-agent-registry'
@@ -18,6 +19,7 @@ import {
1819
loadMostRecentChatState,
1920
saveChatState,
2021
} from '../utils/run-state-storage'
22+
import { useChatStore } from '../state/chat-store'
2123
import { setCurrentChatId } from '../project-files'
2224

2325
import type { ElapsedTimeTracker } from './use-elapsed-time'
@@ -1642,6 +1644,18 @@ export const useSendMessage = ({
16421644
return
16431645
}
16441646

1647+
// Refresh usage data if the banner is currently visible
1648+
const isUsageVisible = useChatStore.getState().isUsageVisible
1649+
if (isUsageVisible) {
1650+
// Don't await - let it run in background to avoid blocking UI updates
1651+
fetchAndUpdateUsage({ showBanner: false }).catch((error) => {
1652+
logger.error(
1653+
{ error: getErrorObject(error) },
1654+
'Failed to refresh usage data after run completion',
1655+
)
1656+
})
1657+
}
1658+
16451659
setStreamStatus('idle')
16461660
if (resumeQueue) {
16471661
resumeQueue()

0 commit comments

Comments
 (0)