Track session usage and print it in /usage

Author: jahooma

cli/src/chat.tsx

@@ -119,6 +119,7 @@ export const Chat = ({
     setHasReceivedPlanResponse,
     lastMessageMode,
     setLastMessageMode,
+    addSessionCredits,
     resetChatStore,
   } = useChatStore(
     useShallow((store) => ({
@@ -151,6 +152,7 @@ export const Chat = ({
       setHasReceivedPlanResponse: store.setHasReceivedPlanResponse,
       lastMessageMode: store.lastMessageMode,
       setLastMessageMode: store.setLastMessageMode,
+      addSessionCredits: store.addSessionCredits,
       resetChatStore: store.reset,
     })),
   )
@@ -403,6 +405,7 @@ export const Chat = ({
     setHasReceivedPlanResponse,
     lastMessageMode,
     setLastMessageMode,
+    addSessionCredits,
   })
 
   sendMessageRef.current = sendMessage

cli/src/commands/usage.ts

@@ -1,5 +1,6 @@
 import { BACKEND_URL } from '@codebuff/sdk'
 
+import { useChatStore } from '../state/chat-store'
 import { getAuthToken } from '../utils/auth'
 import { logger } from '../utils/logger'
 import type { PostUserMessageFn } from '../types/contracts/send-message'
@@ -16,6 +17,7 @@ export async function handleUsageCommand(): Promise<{
   postUserMessage: PostUserMessageFn
 }> {
   const authToken = getAuthToken()
+  const sessionCreditsUsed = useChatStore.getState().sessionCreditsUsed
 
   if (!authToken) {
     const postUserMessage: PostUserMessageFn = (prev) => [
@@ -64,7 +66,7 @@ export async function handleUsageCommand(): Promise<{
     const data = (await response.json()) as UsageResponse
 
     // Format the usage message similar to npm-app
-    let usageMessage = ''
+    let usageMessage = `Session usage: ${sessionCreditsUsed.toLocaleString()} credits used`
 
     if (data.remainingBalance !== null) {
       const remainingColor =
@@ -74,7 +76,7 @@ export async function handleUsageCommand(): Promise<{
             ? 'yellow'
             : 'green'
 
-      usageMessage = `Credits Remaining: ${data.remainingBalance.toLocaleString()}`
+      usageMessage += `\n\nCredits Remaining: ${data.remainingBalance.toLocaleString()}`
 
       // Add next quota reset info if available
       if (data.next_quota_reset) {
@@ -89,7 +91,7 @@ export async function handleUsageCommand(): Promise<{
         usageMessage += `\n\nFree credits will renew on ${dateDisplay}.`
       }
     } else {
-      usageMessage = 'Usage information not available.'
+      usageMessage += '\n\nTotal balance information not available.'
     }
 
     const postUserMessage: PostUserMessageFn = (prev) => [

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

@@ -83,7 +83,7 @@ const scrubPlanTagsInBlocks = (blocks: ContentBlock[]): ContentBlock[] => {
 /**
  * Auto-collapse thinking blocks to reduce UI clutter.
  * Tracks which thinking blocks have been collapsed to avoid duplicate collapses.
- * 
+ *
  * @param messageId - ID of the message containing the thinking block
  * @param agentId - Optional agent ID for nested agent thinking blocks
  * @param autoCollapsedThinkingIdsRef - Ref tracking which thinking IDs have been auto-collapsed
@@ -228,6 +228,7 @@ interface UseSendMessageOptions {
   setHasReceivedPlanResponse: (value: boolean) => void
   lastMessageMode: AgentMode | null
   setLastMessageMode: (mode: AgentMode | null) => void
+  addSessionCredits: (credits: number) => void
 }
 
 export const useSendMessage = ({
@@ -258,6 +259,7 @@ export const useSendMessage = ({
   setHasReceivedPlanResponse,
   lastMessageMode,
   setLastMessageMode,
+  addSessionCredits,
 }: UseSendMessageOptions): {
   sendMessage: SendMessageFn
   clearMessages: () => void
@@ -966,12 +968,12 @@ export const useSendMessage = ({
               return
             }
 
-            if (event.type === 'finish' && event.totalCost !== undefined) {
-              actualCredits = event.totalCost
-            }
-
-            if ('totalCost' in event && event.totalCost !== undefined) {
+            if (
+              event.type === 'finish' &&
+              typeof event.totalCost === 'number'
+            ) {
               actualCredits = event.totalCost
+              addSessionCredits(event.totalCost)
             }
 
             if (event.type === 'subagent_start') {
@@ -1474,11 +1476,15 @@ export const useSendMessage = ({
                               content,
                             }
                             // Determine status based on whether there's an error
-                            const hasError = has(result.value, 'errorMessage') && result.value.errorMessage
+                            const hasError =
+                              has(result.value, 'errorMessage') &&
+                              result.value.errorMessage
                             return {
                               ...block,
                               blocks: [resultTextBlock],
-                              status: hasError ? ('failed' as const) : ('complete' as const),
+                              status: hasError
+                                ? ('failed' as const)
+                                : ('complete' as const),
                             }
                           }
                         }
@@ -1658,6 +1664,7 @@ export const useSendMessage = ({
       setHasReceivedPlanResponse,
       lastMessageMode,
       setLastMessageMode,
+      addSessionCredits,
     ],
   )
 

cli/src/state/chat-store.ts

@@ -29,6 +29,7 @@ export type ChatStoreState = {
   agentMode: AgentMode
   hasReceivedPlanResponse: boolean
   lastMessageMode: AgentMode | null
+  sessionCreditsUsed: number
 }
 
 type ChatStoreActions = {
@@ -58,6 +59,7 @@ type ChatStoreActions = {
   toggleAgentMode: () => void
   setHasReceivedPlanResponse: (value: boolean) => void
   setLastMessageMode: (mode: AgentMode | null) => void
+  addSessionCredits: (credits: number) => void
   reset: () => void
 }
 
@@ -81,6 +83,7 @@ const initialState: ChatStoreState = {
   agentMode: 'DEFAULT',
   hasReceivedPlanResponse: false,
   lastMessageMode: null,
+  sessionCreditsUsed: 0,
 }
 
 export const useChatStore = create<ChatStore>()(
@@ -180,6 +183,11 @@ export const useChatStore = create<ChatStore>()(
         state.lastMessageMode = mode
       }),
 
+    addSessionCredits: (credits) =>
+      set((state) => {
+        state.sessionCreditsUsed += credits
+      }),
+
     reset: () =>
       set((state) => {
         state.messages = initialState.messages.slice()
@@ -197,6 +205,7 @@ export const useChatStore = create<ChatStore>()(
         state.agentMode = initialState.agentMode
         state.hasReceivedPlanResponse = initialState.hasReceivedPlanResponse
         state.lastMessageMode = initialState.lastMessageMode
+        state.sessionCreditsUsed = initialState.sessionCreditsUsed
       }),
   })),
 )