| 1 | package bakeoff |
| 2 | |
| 3 | import ( |
| 4 | "bufio" |
| 5 | "encoding/json" |
| 6 | "errors" |
| 7 | "io" |
| 8 | ) |
| 9 | |
| 10 | // Tokens is what one agent session cost. |
| 11 | type Tokens struct { |
| 12 | Input, Output, CacheRead, CacheWrite int64 |
| 13 | Turns int |
| 14 | Seconds float64 |
| 15 | CostUSD float64 |
| 16 | } |
| 17 | |
| 18 | // Add sums two sessions. |
| 19 | func (t Tokens) Add(o Tokens) Tokens { |
| 20 | return Tokens{t.Input + o.Input, t.Output + o.Output, t.CacheRead + o.CacheRead, t.CacheWrite + o.CacheWrite, |
| 21 | t.Turns + o.Turns, t.Seconds + o.Seconds, t.CostUSD + o.CostUSD} |
| 22 | } |
| 23 | |
| 24 | func scanLines(r io.Reader, f func([]byte)) error { |
| 25 | sc := bufio.NewScanner(r) |
| 26 | sc.Buffer(make([]byte, 1<<20), 1<<26) |
| 27 | for sc.Scan() { |
| 28 | f(sc.Bytes()) |
| 29 | } |
| 30 | return sc.Err() |
| 31 | } |
| 32 | |
| 33 | // ParseClaudeLog reads `claude -p --output-format stream-json` output. A |
| 34 | // session can emit several result events (it resumes itself when |
| 35 | // background work finishes): their per-model usage and cost are |
| 36 | // cumulative for the whole session, so the last one counts, while turns |
| 37 | // and duration are per result, so they are summed. Per-model usage, which |
| 38 | // includes any subagents, is preferred to the top-level usage. |
| 39 | func ParseClaudeLog(r io.Reader) (Tokens, error) { |
| 40 | type usage struct { |
| 41 | Input int64 `json:"input_tokens"` |
| 42 | Output int64 `json:"output_tokens"` |
| 43 | CacheRead int64 `json:"cache_read_input_tokens"` |
| 44 | CacheWrite int64 `json:"cache_creation_input_tokens"` |
| 45 | } |
| 46 | type modelUsage struct { |
| 47 | Input int64 `json:"inputTokens"` |
| 48 | Output int64 `json:"outputTokens"` |
| 49 | CacheRead int64 `json:"cacheReadInputTokens"` |
| 50 | CacheWrite int64 `json:"cacheCreationInputTokens"` |
| 51 | } |
| 52 | var last *Tokens |
| 53 | turns, seconds := 0, 0.0 |
| 54 | err := scanLines(r, func(b []byte) { |
| 55 | var e struct { |
| 56 | Type string `json:"type"` |
| 57 | Usage usage `json:"usage"` |
| 58 | ModelUsage map[string]modelUsage `json:"modelUsage"` |
| 59 | Turns int `json:"num_turns"` |
| 60 | DurationMs float64 `json:"duration_ms"` |
| 61 | Cost float64 `json:"total_cost_usd"` |
| 62 | } |
| 63 | if json.Unmarshal(b, &e) != nil || e.Type != "result" { |
| 64 | return |
| 65 | } |
| 66 | t := Tokens{Input: e.Usage.Input, Output: e.Usage.Output, CacheRead: e.Usage.CacheRead, CacheWrite: e.Usage.CacheWrite, |
| 67 | Turns: e.Turns, Seconds: e.DurationMs / 1000, CostUSD: e.Cost} |
| 68 | if len(e.ModelUsage) > 0 { |
| 69 | t.Input, t.Output, t.CacheRead, t.CacheWrite = 0, 0, 0, 0 |
| 70 | for _, m := range e.ModelUsage { |
| 71 | t.Input += m.Input |
| 72 | t.Output += m.Output |
| 73 | t.CacheRead += m.CacheRead |
| 74 | t.CacheWrite += m.CacheWrite |
| 75 | } |
| 76 | } |
| 77 | turns += e.Turns |
| 78 | seconds += e.DurationMs / 1000 |
| 79 | last = &t |
| 80 | }) |
| 81 | if err != nil { |
| 82 | return Tokens{}, err |
| 83 | } |
| 84 | if last == nil { |
| 85 | return Tokens{}, errors.New("no result event in claude log") |
| 86 | } |
| 87 | last.Turns, last.Seconds = turns, seconds |
| 88 | return *last, nil |
| 89 | } |
| 90 | |
| 91 | // ParseCodexLog sums the usage on every turn.completed event of |
| 92 | // `codex exec --json` output. Codex's input count includes cached input. |
| 93 | func ParseCodexLog(r io.Reader) (Tokens, error) { |
| 94 | var t Tokens |
| 95 | err := scanLines(r, func(b []byte) { |
| 96 | var e struct { |
| 97 | Type string `json:"type"` |
| 98 | Usage struct { |
| 99 | Input int64 `json:"input_tokens"` |
| 100 | Cached int64 `json:"cached_input_tokens"` |
| 101 | CacheWrite int64 `json:"cache_write_input_tokens"` |
| 102 | Output int64 `json:"output_tokens"` |
| 103 | } `json:"usage"` |
| 104 | } |
| 105 | if json.Unmarshal(b, &e) != nil || e.Type != "turn.completed" { |
| 106 | return |
| 107 | } |
| 108 | t.Input += e.Usage.Input |
| 109 | t.CacheRead += e.Usage.Cached |
| 110 | t.CacheWrite += e.Usage.CacheWrite |
| 111 | t.Output += e.Usage.Output |
| 112 | t.Turns++ |
| 113 | }) |
| 114 | if err == nil && t.Turns == 0 { |
| 115 | err = errors.New("no turn.completed event in codex log") |
| 116 | } |
| 117 | return t, err |
| 118 | } |
| 119 | |
| 120 | // Since turns a resumed session's cumulative totals into that phase's own: |
| 121 | // claude reports usage and cost for the whole session on resume, but turns |
| 122 | // and duration for the resumed run only. |
| 123 | func (t Tokens) Since(prev Tokens) Tokens { |
| 124 | return Tokens{t.Input - prev.Input, t.Output - prev.Output, t.CacheRead - prev.CacheRead, t.CacheWrite - prev.CacheWrite, |
| 125 | t.Turns, t.Seconds, t.CostUSD - prev.CostUSD} |
| 126 | } |
| 127 | |
amadan