Implement per-subagent cost attribution and surface credits in interactive UI

Moves credit display to subagent selector and detailed trace view, wired to backend cost events and client store; excludes .agents changes.

Generated with Codebuff 🤖
Co-Authored-By: Codebuff <noreply@codebuff.com>

Author: brandonkachen

backend/src/llm-apis/message-cost-tracker.ts

@@ -420,6 +420,7 @@ async function sendCostResponseToClient(
   clientSessionId: string,
   userInputId: string,
   creditsUsed: number,
+  agentId?: string,
 ): Promise<void> {
   try {
     const clientEntry = Array.from(SWITCHBOARD.clients.entries()).find(
@@ -434,6 +435,7 @@ async function sendCostResponseToClient(
           type: 'message-cost-response',
           promptId: userInputId,
           credits: creditsUsed,
+          agentId,
         })
       } else {
         logger.warn(
@@ -543,6 +545,7 @@ export const saveMessage = async (value: {
   usesUserApiKey?: boolean
   chargeUser?: boolean
   costOverrideDollars?: number
+  agentId?: string
 }): Promise<number> =>
   withLoggerContext(
     {
@@ -612,6 +615,7 @@ export const saveMessage = async (value: {
         value.clientSessionId,
         value.userInputId,
         creditsUsed,
+        value.agentId,
       )
 
       const savedMessageResult = await insertMessageRecord({

backend/src/llm-apis/vercel-ai-sdk/ai-sdk.ts

@@ -72,6 +72,7 @@ export const promptAiSdkStream = async function* (
     chargeUser?: boolean
     thinkingBudget?: number
     userInputId: string
+    agentId?: string
     maxRetries?: number
     onCostCalculated?: (credits: number) => Promise<void>
   } & Omit<Parameters<typeof streamText>[0], 'model'>,
@@ -210,6 +211,7 @@ export const promptAiSdkStream = async function* (
     latencyMs: Date.now() - startTime,
     chargeUser: options.chargeUser ?? true,
     costOverrideDollars,
+    agentId: options.agentId,
   })
 
   // Call the cost callback if provided
@@ -229,6 +231,7 @@ export const promptAiSdk = async function (
     model: Model
     userId: string | undefined
     chargeUser?: boolean
+    agentId?: string
     onCostCalculated?: (credits: number) => Promise<void>
   } & Omit<Parameters<typeof generateText>[0], 'model'>,
 ): Promise<string> {
@@ -276,6 +279,7 @@ export const promptAiSdk = async function (
     finishedAt: new Date(),
     latencyMs: Date.now() - startTime,
     chargeUser: options.chargeUser ?? true,
+    agentId: options.agentId,
   })
 
   // Call the cost callback if provided
@@ -300,6 +304,7 @@ export const promptAiSdkStructured = async function <T>(options: {
   temperature?: number
   timeout?: number
   chargeUser?: boolean
+  agentId?: string
   onCostCalculated?: (credits: number) => Promise<void>
 }): Promise<T> {
   if (
@@ -350,6 +355,7 @@ export const promptAiSdkStructured = async function <T>(options: {
     finishedAt: new Date(),
     latencyMs: Date.now() - startTime,
     chargeUser: options.chargeUser ?? true,
+    agentId: options.agentId,
   })
 
   // Call the cost callback if provided

backend/src/prompt-agent-stream.ts

@@ -13,6 +13,7 @@ export const getAgentStreamFromTemplate = (params: {
   userInputId: string
   userId: string | undefined
   onCostCalculated?: (credits: number) => Promise<void>
+  agentId?: string
 
   template: AgentTemplate
 }) => {
@@ -22,6 +23,7 @@ export const getAgentStreamFromTemplate = (params: {
     userInputId,
     userId,
     onCostCalculated,
+    agentId,
     template,
   } = params
 
@@ -42,6 +44,7 @@ export const getAgentStreamFromTemplate = (params: {
       userId,
       maxOutputTokens: 32_000,
       onCostCalculated,
+      agentId,
     }
 
     // Add Gemini-specific options if needed

backend/src/run-agent-step.ts

@@ -10,7 +10,6 @@ import { renderToolResults } from '@codebuff/common/tools/utils'
 import { buildArray } from '@codebuff/common/util/array'
 import { generateCompactId } from '@codebuff/common/util/string'
 
-
 import { asyncAgentManager } from './async-agent-manager'
 import { getFileReadingUpdates } from './get-file-reading-updates'
 import { checkLiveUserInput } from './live-user-inputs'
@@ -253,19 +252,11 @@ export const runAgentStep = async (
     fingerprintId,
     userInputId,
     userId,
+    agentId: agentState.agentId,
     template: agentTemplate,
     onCostCalculated: async (credits: number) => {
       try {
         agentState.creditsUsed += credits
-        logger.debug(
-          {
-            agentId: agentState.agentId,
-            credits,
-            totalCredits: agentState.creditsUsed,
-          },
-          'Added LLM cost to agent state',
-        )
-        
         // Transactional cost attribution: ensure costs are actually deducted
         // This is already handled by the saveMessage function which calls updateUserCycleUsage
         // If that fails, the promise rejection will bubble up and halt agent execution
@@ -623,13 +614,5 @@ export const loopAgentSteps = async (
     throw error
   } finally {
     // Ensure costs are always captured, even on failure
-    logger.debug(
-      {
-        agentId: currentAgentState.agentId,
-        creditsUsed: currentAgentState.creditsUsed,
-        status: 'completed_or_failed',
-      },
-      'Agent execution completed with cost tracking',
-    )
   }
 }

common/src/actions.ts

@@ -102,11 +102,10 @@ export const MessageCostResponseSchema = z.object({
   type: z.literal('message-cost-response'),
   promptId: z.string(),
   credits: z.number(),
+  agentId: z.string().optional(),
 })
 export type MessageCostResponse = z.infer<typeof MessageCostResponseSchema>
 
-
-
 export const PromptResponseSchema = z.object({
   type: z.literal('prompt-response'),
   promptId: z.string(),

npm-app/src/cli-handlers/subagent-list.ts

@@ -1,7 +1,10 @@
 import { pluralize } from '@codebuff/common/util/string'
 import { green, yellow, cyan, magenta, bold, gray } from 'picocolors'
 
-import { getSubagentsChronological } from '../subagent-storage'
+import {
+  getSubagentsChronological,
+  type SubagentData,
+} from '../subagent-storage'
 import { enterSubagentBuffer } from './traces'
 import {
   ENTER_ALT_BUFFER,
@@ -19,14 +22,7 @@ let persistentSelectedIndex = -1 // -1 means not initialized
 let scrollOffset = 0
 let allContentLines: string[] = []
 let subagentLinePositions: number[] = []
-let subagentList: Array<{
-  agentId: string
-  agentType: string
-  prompt?: string
-  isActive: boolean
-  lastActivity: number
-  startTime: number
-}> = []
+let subagentList: SubagentData[] = []
 
 export function isInSubagentListMode(): boolean {
   return isInSubagentListBuffer
@@ -255,7 +251,11 @@ function buildAllContentLines() {
       const timeInfo = agent.isActive
         ? green(`[Active - ${startTime}]`)
         : gray(`[${startTime}]`)
-      const headerLine = `${agentInfo} ${timeInfo}`
+      const creditsDisplay =
+        agent.creditsUsed > 0
+          ? yellow(` (${pluralize(agent.creditsUsed, 'credit')})`)
+          : ''
+      const headerLine = `${agentInfo}${creditsDisplay} ${timeInfo}`
 
       const contentForBox = [headerLine, ...promptLines.map((p) => gray(p))]
 

npm-app/src/cli-handlers/traces.ts

@@ -19,6 +19,7 @@ import {
 } from '../utils/terminal'
 
 import type { SubagentData } from '../subagent-storage'
+import { logger } from '../utils/logger'
 
 // Add helpers to truncate to first line and reduce sections
 function firstLine(text: string): string {
@@ -52,6 +53,9 @@ export function isInSubagentBufferMode(): boolean {
  * Display a formatted list of traces with enhanced styling
  */
 export function displaySubagentList(agents: SubagentData[]) {
+  // Added: log list render
+  logger.debug({ count: agents.length }, 'Rendering subagent list!')
+
   console.log(bold(cyan('🤖 Available Traces')))
   console.log(gray(`Found ${pluralize(agents.length, 'trace')}`))
   console.log()
@@ -63,6 +67,7 @@ export function displaySubagentList(agents: SubagentData[]) {
       // Truncate prompt preview to first line
       const promptFirst = agent.prompt ? firstLine(agent.prompt) : '(no prompt)'
       const promptPreview = gray(promptFirst)
+
       console.log(
         `  ${status} ${bold(agent.agentId)} ${gray(`(${agent.agentType})`)}`,
       )
@@ -71,7 +76,6 @@ export function displaySubagentList(agents: SubagentData[]) {
     })
   }
 }
-
 export function enterSubagentBuffer(
   rl: any,
   agentId: string,
@@ -168,6 +172,16 @@ function updateSubagentContent() {
   if (agentData.prompt) {
     const promptLine = bold(gray(`Prompt: ${firstLine(agentData.prompt)}`))
     wrappedLines.push(...wrapLine(promptLine, terminalWidth))
+  }
+
+  // Add credits used if any
+  if (agentData.creditsUsed > 0) {
+    const creditsLine = yellow(`${pluralize(agentData.creditsUsed, 'credit')}`)
+    wrappedLines.push(...wrapLine(creditsLine, terminalWidth))
+  }
+
+  // Add a separator line if we added prompt or credits
+  if (agentData.prompt || agentData.creditsUsed > 0) {
     wrappedLines.push('')
   }
 

npm-app/src/client.ts

@@ -79,6 +79,7 @@ import {
   getAllSubagentIds,
   markSubagentInactive,
   storeSubagentChunk,
+  addCreditsByAgentId,
 } from './subagent-storage'
 import { handleToolCall } from './tool-handlers'
 import { identifyUser, trackEvent } from './utils/analytics'
@@ -872,6 +873,11 @@ export class Client {
         this.creditsByPromptId[response.promptId] = []
       }
       this.creditsByPromptId[response.promptId].push(response.credits)
+
+      // Attribute credits directly to the agentId (backend now always provides it)
+      if (response.agentId) {
+        addCreditsByAgentId(response.agentId, response.credits)
+      }
     })
 
     this.webSocket.subscribe('usage-response', (action) => {
@@ -908,10 +914,9 @@ export class Client {
     this.webSocket.subscribe('request-reconnect', () => {
       this.reconnectWhenNextIdle()
     })
-
     // Handle subagent streaming messages
     this.webSocket.subscribe('subagent-response-chunk', (action) => {
-      const { userInputId, agentId, agentType, chunk, prompt } = action
+      const { agentId, agentType, chunk, prompt } = action
 
       // Store the chunk locally
       storeSubagentChunk({ agentId, agentType, chunk, prompt })

npm-app/src/subagent-storage.ts

@@ -4,6 +4,7 @@ import * as path from 'path'
 import { endToolTag } from '@codebuff/common/tools/constants'
 
 import { getProjectRoot } from './project-files'
+import { logger } from './utils/logger'
 
 interface SubagentMessage {
   timestamp: number
@@ -19,6 +20,8 @@ export interface SubagentData {
   isActive: boolean
   lastActivity: number
   startTime: number
+  // Add: accumulated credits used by this subagent
+  creditsUsed: number
 }
 
 // Global storage for all subagent data
@@ -47,6 +50,11 @@ export function setTraceEnabled(enabled: boolean) {
       }
     } catch (error) {
       console.warn('Warning: Could not clear traces directory:', error)
+      // Restore: existing warning log for failed directory clearing
+      logger.warn(
+        { error: error instanceof Error ? error.message : String(error) },
+        'Failed clearing traces directory',
+      )
     }
   }
 }
@@ -115,6 +123,8 @@ export function storeSubagentChunk({
       isActive: true,
       lastActivity: now,
       startTime: now,
+      // Initialize new field
+      creditsUsed: 0,
     })
   }
 
@@ -137,6 +147,16 @@ export function storeSubagentChunk({
   writeToTraceLog(agentId, agentType, chunk)
 }
 
+/**
+ * Add credits directly by agentId
+ */
+export function addCreditsByAgentId(agentId: string, credits: number) {
+  const data = subagentStorage.get(agentId)
+  if (data) {
+    data.creditsUsed += credits
+  }
+}
+
 /**
  * Get all messages for a specific subagent
  */