| 1 | package bot |
| 2 | |
| 3 | import ( |
| 4 | "encoding/json" |
| 5 | "net/http" |
| 6 | "net/http/httptest" |
| 7 | "testing" |
| 8 | ) |
| 9 | |
| 10 | func rowsByKey(rows []UsageRow) map[string]int64 { |
| 11 | out := map[string]int64{} |
| 12 | for _, r := range rows { |
| 13 | out[r.Day+"|"+r.Model+"|"+r.Dimension] = r.Quantity |
| 14 | } |
| 15 | return out |
| 16 | } |
| 17 | |
| 18 | // The accumulator's contract: calls merge into (day, model, dimension) |
| 19 | // cells, drain empties, and restore folds a failed flush back in losslessly. |
| 20 | func TestUsageMeterAccumulateDrainRestore(t *testing.T) { |
| 21 | m := newUsageMeter() |
| 22 | m.addLLM("claude-sonnet-5", anthUsage{InputTokens: 100, OutputTokens: 20, CacheReadInputTokens: 300}) |
| 23 | m.addLLM("claude-sonnet-5", anthUsage{InputTokens: 50, OutputTokens: 5, CacheCreationInputTokens: 7}) |
| 24 | m.addLLM("claude-haiku-4-5", anthUsage{InputTokens: 9}) |
| 25 | m.inc("charter_runs", 1) |
| 26 | |
| 27 | day := meterDay() |
| 28 | rows := m.drain() |
| 29 | got := rowsByKey(rows) |
| 30 | for key, want := range map[string]int64{ |
| 31 | day + "|claude-sonnet-5|llm_requests": 2, |
| 32 | day + "|claude-sonnet-5|llm_input_tokens": 150, |
| 33 | day + "|claude-sonnet-5|llm_output_tokens": 25, |
| 34 | day + "|claude-sonnet-5|llm_cache_read_tokens": 300, |
| 35 | day + "|claude-sonnet-5|llm_cache_write_tokens": 7, |
| 36 | day + "|claude-haiku-4-5|llm_requests": 1, |
| 37 | day + "|claude-haiku-4-5|llm_input_tokens": 9, |
| 38 | day + "||charter_runs": 1, |
| 39 | } { |
| 40 | if got[key] != want { |
| 41 | t.Errorf("%s = %d, want %d", key, got[key], want) |
| 42 | } |
| 43 | } |
| 44 | if m.drain() != nil { |
| 45 | t.Fatal("second drain should be empty") |
| 46 | } |
| 47 | |
| 48 | // A failed flush restores; new counts merge on top. |
| 49 | m.restore(rows) |
| 50 | m.addLLM("claude-sonnet-5", anthUsage{InputTokens: 1}) |
| 51 | again := rowsByKey(m.drain()) |
| 52 | if again[day+"|claude-sonnet-5|llm_input_tokens"] != 151 { |
| 53 | t.Fatalf("restore+add did not merge: %v", again) |
| 54 | } |
| 55 | } |
| 56 | |
| 57 | // A zero-usage block (a failed parse, an empty response) records only the |
| 58 | // request count — never zero-quantity rows. |
| 59 | func TestUsageMeterSkipsZeroes(t *testing.T) { |
| 60 | m := newUsageMeter() |
| 61 | m.addLLM("m", anthUsage{}) |
| 62 | rows := m.drain() |
| 63 | if len(rows) != 1 || rows[0].Dimension != "llm_requests" || rows[0].Quantity != 1 { |
| 64 | t.Fatalf("want just llm_requests=1, got %+v", rows) |
| 65 | } |
| 66 | } |
| 67 | |
| 68 | // The non-streamed path: the response's usage block reaches both the returned |
| 69 | // response and the record hook (the meter's tap). |
| 70 | func TestAnthropicUsageParsed(t *testing.T) { |
| 71 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 72 | json.NewEncoder(w).Encode(map[string]any{ |
| 73 | "content": []map[string]any{{"type": "text", "text": "hi"}}, |
| 74 | "stop_reason": "end_turn", |
| 75 | "usage": map[string]any{ |
| 76 | "input_tokens": 120, "output_tokens": 34, |
| 77 | "cache_creation_input_tokens": 8, "cache_read_input_tokens": 4000, |
| 78 | }, |
| 79 | }) |
| 80 | })) |
| 81 | defer srv.Close() |
| 82 | |
| 83 | var gotModel string |
| 84 | var got anthUsage |
| 85 | c := &anthropicClient{key: "k", base: srv.URL, http: srv.Client(), |
| 86 | record: func(model string, u anthUsage) { gotModel, got = model, u }} |
| 87 | resp, err := c.create(anthRequest{Model: "claude-sonnet-5", MaxTokens: 10}) |
| 88 | if err != nil { |
| 89 | t.Fatalf("create: %v", err) |
| 90 | } |
| 91 | want := anthUsage{InputTokens: 120, OutputTokens: 34, CacheCreationInputTokens: 8, CacheReadInputTokens: 4000} |
| 92 | if resp.Usage != want { |
| 93 | t.Fatalf("response usage = %+v, want %+v", resp.Usage, want) |
| 94 | } |
| 95 | if gotModel != "claude-sonnet-5" || got != want { |
| 96 | t.Fatalf("record hook got (%q, %+v)", gotModel, got) |
| 97 | } |
| 98 | } |
| 99 | |
| 100 | // The streamed path: input + cache counters ride message_start, the final |
| 101 | // (cumulative) output count rides message_delta. |
| 102 | func TestAnthropicUsageParsedStream(t *testing.T) { |
| 103 | stream := "event: message_start\n" + |
| 104 | `data: {"type":"message_start","message":{"usage":{"input_tokens":200,"output_tokens":1,"cache_read_input_tokens":900}}}` + "\n\n" + |
| 105 | "event: content_block_start\n" + |
| 106 | `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}` + "\n\n" + |
| 107 | "event: content_block_delta\n" + |
| 108 | `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ok"}}` + "\n\n" + |
| 109 | "event: message_delta\n" + |
| 110 | `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":57}}` + "\n\n" + |
| 111 | "event: message_stop\n" + |
| 112 | `data: {"type":"message_stop"}` + "\n\n" |
| 113 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 114 | w.Header().Set("Content-Type", "text/event-stream") |
| 115 | w.Write([]byte(stream)) |
| 116 | })) |
| 117 | defer srv.Close() |
| 118 | |
| 119 | var got anthUsage |
| 120 | c := &anthropicClient{key: "k", base: srv.URL, http: srv.Client(), |
| 121 | record: func(_ string, u anthUsage) { got = u }} |
| 122 | resp, err := c.createStream(anthRequest{Model: "m", MaxTokens: 10}) |
| 123 | if err != nil { |
| 124 | t.Fatalf("createStream: %v", err) |
| 125 | } |
| 126 | want := anthUsage{InputTokens: 200, OutputTokens: 57, CacheReadInputTokens: 900} |
| 127 | if resp.Usage != want { |
| 128 | t.Fatalf("stream usage = %+v, want %+v", resp.Usage, want) |
| 129 | } |
| 130 | if got != want { |
| 131 | t.Fatalf("record hook got %+v", got) |
| 132 | } |
| 133 | } |
| 134 | |
| 135 | // The outbox contract (the ack-loss fix): a failed flush keeps the report |
| 136 | // VERBATIM — same id, same rows — so the ledger can dedupe the replay, and |
| 137 | // usage counted during the outage rides the NEXT report, never the pending one. |
| 138 | func TestFlushUsageOutboxRetriesVerbatim(t *testing.T) { |
| 139 | type report struct { |
| 140 | ReportID string `json:"report_id"` |
| 141 | Rows []UsageRow `json:"rows"` |
| 142 | } |
| 143 | var got []report |
| 144 | fail := true |
| 145 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 146 | var rep report |
| 147 | json.NewDecoder(r.Body).Decode(&rep) |
| 148 | got = append(got, rep) |
| 149 | if fail { |
| 150 | http.Error(w, "cloud blip", http.StatusBadGateway) |
| 151 | return |
| 152 | } |
| 153 | json.NewEncoder(w).Encode(map[string]any{"ok": true, "budget": map[string]any{"state": "ok"}}) |
| 154 | })) |
| 155 | defer srv.Close() |
| 156 | |
| 157 | r := &Runtime{Base: srv.URL, Session: &Session{Token: "tok", Name: "Eleven"}, Logf: t.Logf} |
| 158 | r.client = srv.Client() |
| 159 | r.meter = newUsageMeter() |
| 160 | |
| 161 | r.meter.inc("charter_runs", 1) |
| 162 | r.flushUsage() // fails; report stays pending |
| 163 | r.meter.inc("charter_runs", 5) |
| 164 | r.flushUsage() // still failing; the SAME report is resent, new counts stay in the meter |
| 165 | fail = false |
| 166 | r.flushUsage() // succeeds: pending report lands |
| 167 | r.flushUsage() // then the accumulated counts land as a NEW report |
| 168 | |
| 169 | if len(got) != 4 { |
| 170 | t.Fatalf("attempts = %d, want 4", len(got)) |
| 171 | } |
| 172 | if got[0].ReportID == "" || got[0].ReportID != got[1].ReportID || got[1].ReportID != got[2].ReportID { |
| 173 | t.Fatalf("retries must reuse the report id: %q %q %q", got[0].ReportID, got[1].ReportID, got[2].ReportID) |
| 174 | } |
| 175 | if got[3].ReportID == got[2].ReportID { |
| 176 | t.Fatal("post-success flush must mint a fresh report id") |
| 177 | } |
| 178 | if len(got[2].Rows) != 1 || got[2].Rows[0].Quantity != 1 { |
| 179 | t.Fatalf("pending report mutated across retries: %+v", got[2].Rows) |
| 180 | } |
| 181 | if len(got[3].Rows) != 1 || got[3].Rows[0].Quantity != 5 { |
| 182 | t.Fatalf("outage-time counts must ride the next report: %+v", got[3].Rows) |
| 183 | } |
| 184 | if r.pendingUsage != nil { |
| 185 | t.Fatal("nothing should be pending after success") |
| 186 | } |
| 187 | if r.budgetState.Load() != "ok" { |
| 188 | t.Fatalf("budget state not stored: %v", r.budgetState.Load()) |
| 189 | } |
| 190 | } |
| 191 | |
| 192 | // Phase 3's sidecar half: an "exhausted" budget state gates the model paths |
| 193 | // (one honest notice, then silence) and recovery re-arms the notice. |
| 194 | func TestBudgetExhaustionGate(t *testing.T) { |
| 195 | r := &Runtime{Session: &Session{Name: "Eleven"}} |
| 196 | if r.budgetExhausted() { |
| 197 | t.Fatal("never-flushed must read as within budget") |
| 198 | } |
| 199 | r.budgetState.Store("exhausted") |
| 200 | if !r.budgetExhausted() { |
| 201 | t.Fatal("exhausted state not read") |
| 202 | } |
| 203 | if !r.noteBudgetExhausted() { |
| 204 | t.Fatal("first notice must fire") |
| 205 | } |
| 206 | if r.noteBudgetExhausted() { |
| 207 | t.Fatal("the notice is once-per-exhaustion") |
| 208 | } |
| 209 | // A task run refuses without burning tokens (ai is set, budget is not). |
| 210 | r.ai = &anthropicClient{} |
| 211 | if _, _, err := r.execTaskOnce("p", &TaskSpec{Prompt: "x"}, false); err == nil { |
| 212 | t.Fatal("exhausted budget must refuse task runs") |
| 213 | } |
| 214 | // Recovery (simulated flush storing a non-exhausted state) re-arms. |
| 215 | r.budgetState.Store("ok") |
| 216 | r.budgetNotified.Store(false) |
| 217 | if r.budgetExhausted() { |
| 218 | t.Fatal("ok state must clear the gate") |
| 219 | } |
| 220 | if !r.noteBudgetExhausted() { |
| 221 | t.Fatal("notice must re-arm after recovery") |
| 222 | } |
| 223 | } |
| 224 | |
amadan