1  import chalk from 'chalk'
 2  import { useEffect } from 'react'
 3  import { formatDuration } from './utils/format.js'
 4  import {
 5    getCurrentProjectConfig,
 6    saveCurrentProjectConfig,
 7  } from './utils/config.js'
 8  import { SESSION_ID } from './utils/log.js'
 9  
10  // DO NOT ADD MORE STATE HERE OR BORIS WILL CURSE YOU
11  const STATE: {
12    totalCost: number
13    totalAPIDuration: number
14    startTime: number
15  } = {
16    totalCost: 0,
17    totalAPIDuration: 0,
18    startTime: Date.now(),
19  }
20  
21  export function addToTotalCost(cost: number, duration: number): void {
22    STATE.totalCost += cost
23    STATE.totalAPIDuration += duration
24  }
25  
26  export function getTotalCost(): number {
27    return STATE.totalCost
28  }
29  
30  export function getTotalDuration(): number {
31    return Date.now() - STATE.startTime
32  }
33  
34  export function getTotalAPIDuration(): number {
35    return STATE.totalAPIDuration
36  }
37  
38  function formatCost(cost: number): string {
39    return `$${cost > 0.5 ? round(cost, 100).toFixed(2) : cost.toFixed(4)}`
40  }
41  
42  export function formatTotalCost(): string {
43    return chalk.grey(
44      `Total cost: ${formatCost(STATE.totalCost)}
45  Total duration (API): ${formatDuration(STATE.totalAPIDuration)}
46  Total duration (wall): ${formatDuration(getTotalDuration())}`,
47    )
48  }
49  
50  export function useCostSummary(): void {
51    useEffect(() => {
52      const f = () => {
53        process.stdout.write('\n' + formatTotalCost() + '\n')
54  
55        // Save last cost and duration to project config
56        const projectConfig = getCurrentProjectConfig()
57        saveCurrentProjectConfig({
58          ...projectConfig,
59          lastCost: STATE.totalCost,
60          lastAPIDuration: STATE.totalAPIDuration,
61          lastDuration: getTotalDuration(),
62          lastSessionId: SESSION_ID,
63        })
64      }
65      process.on('exit', f)
66      return () => {
67        process.off('exit', f)
68      }
69    }, [])
70  }
71  
72  function round(number: number, precision: number): number {
73    return Math.round(number * precision) / precision
74  }
75  
76  // Only used in tests
77  export function resetStateForTests(): void {
78    if (process.env.NODE_ENV !== 'test') {
79      throw new Error('resetStateForTests can only be called in tests')
80    }
81    STATE.startTime = Date.now()
82    STATE.totalCost = 0
83    STATE.totalAPIDuration = 0
84  }