1  import type { BetaUsage as Usage } from '@anthropic-ai/sdk/resources/beta/messages/messages.mjs'
  2  import chalk from 'chalk'
  3  import {
  4    addToTotalCostState,
  5    addToTotalLinesChanged,
  6    getCostCounter,
  7    getModelUsage,
  8    getSdkBetas,
  9    getSessionId,
 10    getTokenCounter,
 11    getTotalAPIDuration,
 12    getTotalAPIDurationWithoutRetries,
 13    getTotalCacheCreationInputTokens,
 14    getTotalCacheReadInputTokens,
 15    getTotalCostUSD,
 16    getTotalDuration,
 17    getTotalInputTokens,
 18    getTotalLinesAdded,
 19    getTotalLinesRemoved,
 20    getTotalOutputTokens,
 21    getTotalToolDuration,
 22    getTotalWebSearchRequests,
 23    getUsageForModel,
 24    hasUnknownModelCost,
 25    resetCostState,
 26    resetStateForTests,
 27    setCostStateForRestore,
 28    setHasUnknownModelCost,
 29  } from './bootstrap/state.js'
 30  import type { ModelUsage } from './entrypoints/agentSdkTypes.js'
 31  import {
 32    type AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
 33    logEvent,
 34  } from './services/analytics/index.js'
 35  import { getAdvisorUsage } from './utils/advisor.js'
 36  import {
 37    getCurrentProjectConfig,
 38    saveCurrentProjectConfig,
 39  } from './utils/config.js'
 40  import {
 41    getContextWindowForModel,
 42    getModelMaxOutputTokens,
 43  } from './utils/context.js'
 44  import { isFastModeEnabled } from './utils/fastMode.js'
 45  import { formatDuration, formatNumber } from './utils/format.js'
 46  import type { FpsMetrics } from './utils/fpsTracker.js'
 47  import { getCanonicalName } from './utils/model/model.js'
 48  import { calculateUSDCost } from './utils/modelCost.js'
 49  export {
 50    getTotalCostUSD as getTotalCost,
 51    getTotalDuration,
 52    getTotalAPIDuration,
 53    getTotalAPIDurationWithoutRetries,
 54    addToTotalLinesChanged,
 55    getTotalLinesAdded,
 56    getTotalLinesRemoved,
 57    getTotalInputTokens,
 58    getTotalOutputTokens,
 59    getTotalCacheReadInputTokens,
 60    getTotalCacheCreationInputTokens,
 61    getTotalWebSearchRequests,
 62    formatCost,
 63    hasUnknownModelCost,
 64    resetStateForTests,
 65    resetCostState,
 66    setHasUnknownModelCost,
 67    getModelUsage,
 68    getUsageForModel,
 69  }
 70  
 71  type StoredCostState = {
 72    totalCostUSD: number
 73    totalAPIDuration: number
 74    totalAPIDurationWithoutRetries: number
 75    totalToolDuration: number
 76    totalLinesAdded: number
 77    totalLinesRemoved: number
 78    lastDuration: number | undefined
 79    modelUsage: { [modelName: string]: ModelUsage } | undefined
 80  }
 81  
 82  /**
 83   * Gets stored cost state from project config for a specific session.
 84   * Returns the cost data if the session ID matches, or undefined otherwise.
 85   * Use this to read costs BEFORE overwriting the config with saveCurrentSessionCosts().
 86   */
 87  export function getStoredSessionCosts(
 88    sessionId: string,
 89  ): StoredCostState | undefined {
 90    const projectConfig = getCurrentProjectConfig()
 91  
 92    // Only return costs if this is the same session that was last saved
 93    if (projectConfig.lastSessionId !== sessionId) {
 94      return undefined
 95    }
 96  
 97    // Build model usage with context windows
 98    let modelUsage: { [modelName: string]: ModelUsage } | undefined
 99    if (projectConfig.lastModelUsage) {
100      modelUsage = Object.fromEntries(
101        Object.entries(projectConfig.lastModelUsage).map(([model, usage]) => [
102          model,
103          {
104            ...usage,
105            contextWindow: getContextWindowForModel(model, getSdkBetas()),
106            maxOutputTokens: getModelMaxOutputTokens(model).default,
107          },
108        ]),
109      )
110    }
111  
112    return {
113      totalCostUSD: projectConfig.lastCost ?? 0,
114      totalAPIDuration: projectConfig.lastAPIDuration ?? 0,
115      totalAPIDurationWithoutRetries:
116        projectConfig.lastAPIDurationWithoutRetries ?? 0,
117      totalToolDuration: projectConfig.lastToolDuration ?? 0,
118      totalLinesAdded: projectConfig.lastLinesAdded ?? 0,
119      totalLinesRemoved: projectConfig.lastLinesRemoved ?? 0,
120      lastDuration: projectConfig.lastDuration,
121      modelUsage,
122    }
123  }
124  
125  /**
126   * Restores cost state from project config when resuming a session.
127   * Only restores if the session ID matches the last saved session.
128   * @returns true if cost state was restored, false otherwise
129   */
130  export function restoreCostStateForSession(sessionId: string): boolean {
131    const data = getStoredSessionCosts(sessionId)
132    if (!data) {
133      return false
134    }
135    setCostStateForRestore(data)
136    return true
137  }
138  
139  /**
140   * Saves the current session's costs to project config.
141   * Call this before switching sessions to avoid losing accumulated costs.
142   */
143  export function saveCurrentSessionCosts(fpsMetrics?: FpsMetrics): void {
144    saveCurrentProjectConfig(current => ({
145      ...current,
146      lastCost: getTotalCostUSD(),
147      lastAPIDuration: getTotalAPIDuration(),
148      lastAPIDurationWithoutRetries: getTotalAPIDurationWithoutRetries(),
149      lastToolDuration: getTotalToolDuration(),
150      lastDuration: getTotalDuration(),
151      lastLinesAdded: getTotalLinesAdded(),
152      lastLinesRemoved: getTotalLinesRemoved(),
153      lastTotalInputTokens: getTotalInputTokens(),
154      lastTotalOutputTokens: getTotalOutputTokens(),
155      lastTotalCacheCreationInputTokens: getTotalCacheCreationInputTokens(),
156      lastTotalCacheReadInputTokens: getTotalCacheReadInputTokens(),
157      lastTotalWebSearchRequests: getTotalWebSearchRequests(),
158      lastFpsAverage: fpsMetrics?.averageFps,
159      lastFpsLow1Pct: fpsMetrics?.low1PctFps,
160      lastModelUsage: Object.fromEntries(
161        Object.entries(getModelUsage()).map(([model, usage]) => [
162          model,
163          {
164            inputTokens: usage.inputTokens,
165            outputTokens: usage.outputTokens,
166            cacheReadInputTokens: usage.cacheReadInputTokens,
167            cacheCreationInputTokens: usage.cacheCreationInputTokens,
168            webSearchRequests: usage.webSearchRequests,
169            costUSD: usage.costUSD,
170          },
171        ]),
172      ),
173      lastSessionId: getSessionId(),
174    }))
175  }
176  
177  function formatCost(cost: number, maxDecimalPlaces: number = 4): string {
178    return `$${cost > 0.5 ? round(cost, 100).toFixed(2) : cost.toFixed(maxDecimalPlaces)}`
179  }
180  
181  function formatModelUsage(): string {
182    const modelUsageMap = getModelUsage()
183    if (Object.keys(modelUsageMap).length === 0) {
184      return 'Usage:                 0 input, 0 output, 0 cache read, 0 cache write'
185    }
186  
187    // Accumulate usage by short name
188    const usageByShortName: { [shortName: string]: ModelUsage } = {}
189    for (const [model, usage] of Object.entries(modelUsageMap)) {
190      const shortName = getCanonicalName(model)
191      if (!usageByShortName[shortName]) {
192        usageByShortName[shortName] = {
193          inputTokens: 0,
194          outputTokens: 0,
195          cacheReadInputTokens: 0,
196          cacheCreationInputTokens: 0,
197          webSearchRequests: 0,
198          costUSD: 0,
199          contextWindow: 0,
200          maxOutputTokens: 0,
201        }
202      }
203      const accumulated = usageByShortName[shortName]
204      accumulated.inputTokens += usage.inputTokens
205      accumulated.outputTokens += usage.outputTokens
206      accumulated.cacheReadInputTokens += usage.cacheReadInputTokens
207      accumulated.cacheCreationInputTokens += usage.cacheCreationInputTokens
208      accumulated.webSearchRequests += usage.webSearchRequests
209      accumulated.costUSD += usage.costUSD
210    }
211  
212    let result = 'Usage by model:'
213    for (const [shortName, usage] of Object.entries(usageByShortName)) {
214      const usageString =
215        `  ${formatNumber(usage.inputTokens)} input, ` +
216        `${formatNumber(usage.outputTokens)} output, ` +
217        `${formatNumber(usage.cacheReadInputTokens)} cache read, ` +
218        `${formatNumber(usage.cacheCreationInputTokens)} cache write` +
219        (usage.webSearchRequests > 0
220          ? `, ${formatNumber(usage.webSearchRequests)} web search`
221          : '') +
222        ` (${formatCost(usage.costUSD)})`
223      result += `\n` + `${shortName}:`.padStart(21) + usageString
224    }
225    return result
226  }
227  
228  export function formatTotalCost(): string {
229    const costDisplay =
230      formatCost(getTotalCostUSD()) +
231      (hasUnknownModelCost()
232        ? ' (costs may be inaccurate due to usage of unknown models)'
233        : '')
234  
235    const modelUsageDisplay = formatModelUsage()
236  
237    return chalk.dim(
238      `Total cost:            ${costDisplay}\n` +
239        `Total duration (API):  ${formatDuration(getTotalAPIDuration())}
240  Total duration (wall): ${formatDuration(getTotalDuration())}
241  Total code changes:    ${getTotalLinesAdded()} ${getTotalLinesAdded() === 1 ? 'line' : 'lines'} added, ${getTotalLinesRemoved()} ${getTotalLinesRemoved() === 1 ? 'line' : 'lines'} removed
242  ${modelUsageDisplay}`,
243    )
244  }
245  
246  function round(number: number, precision: number): number {
247    return Math.round(number * precision) / precision
248  }
249  
250  function addToTotalModelUsage(
251    cost: number,
252    usage: Usage,
253    model: string,
254  ): ModelUsage {
255    const modelUsage = getUsageForModel(model) ?? {
256      inputTokens: 0,
257      outputTokens: 0,
258      cacheReadInputTokens: 0,
259      cacheCreationInputTokens: 0,
260      webSearchRequests: 0,
261      costUSD: 0,
262      contextWindow: 0,
263      maxOutputTokens: 0,
264    }
265  
266    modelUsage.inputTokens += usage.input_tokens
267    modelUsage.outputTokens += usage.output_tokens
268    modelUsage.cacheReadInputTokens += usage.cache_read_input_tokens ?? 0
269    modelUsage.cacheCreationInputTokens += usage.cache_creation_input_tokens ?? 0
270    modelUsage.webSearchRequests +=
271      usage.server_tool_use?.web_search_requests ?? 0
272    modelUsage.costUSD += cost
273    modelUsage.contextWindow = getContextWindowForModel(model, getSdkBetas())
274    modelUsage.maxOutputTokens = getModelMaxOutputTokens(model).default
275    return modelUsage
276  }
277  
278  export function addToTotalSessionCost(
279    cost: number,
280    usage: Usage,
281    model: string,
282  ): number {
283    const modelUsage = addToTotalModelUsage(cost, usage, model)
284    addToTotalCostState(cost, modelUsage, model)
285  
286    const attrs =
287      isFastModeEnabled() && usage.speed === 'fast'
288        ? { model, speed: 'fast' }
289        : { model }
290  
291    getCostCounter()?.add(cost, attrs)
292    getTokenCounter()?.add(usage.input_tokens, { ...attrs, type: 'input' })
293    getTokenCounter()?.add(usage.output_tokens, { ...attrs, type: 'output' })
294    getTokenCounter()?.add(usage.cache_read_input_tokens ?? 0, {
295      ...attrs,
296      type: 'cacheRead',
297    })
298    getTokenCounter()?.add(usage.cache_creation_input_tokens ?? 0, {
299      ...attrs,
300      type: 'cacheCreation',
301    })
302  
303    let totalCost = cost
304    for (const advisorUsage of getAdvisorUsage(usage)) {
305      const advisorCost = calculateUSDCost(advisorUsage.model, advisorUsage)
306      logEvent('tengu_advisor_tool_token_usage', {
307        advisor_model:
308          advisorUsage.model as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS,
309        input_tokens: advisorUsage.input_tokens,
310        output_tokens: advisorUsage.output_tokens,
311        cache_read_input_tokens: advisorUsage.cache_read_input_tokens ?? 0,
312        cache_creation_input_tokens:
313          advisorUsage.cache_creation_input_tokens ?? 0,
314        cost_usd_micros: Math.round(advisorCost * 1_000_000),
315      })
316      totalCost += addToTotalSessionCost(
317        advisorCost,
318        advisorUsage,
319        advisorUsage.model,
320      )
321    }
322    return totalCost
323  }