| 1 | package main |
| 2 | |
| 3 | import ( |
| 4 | "bytes" |
| 5 | "encoding/json" |
| 6 | "net/http" |
| 7 | "net/http/httptest" |
| 8 | "testing" |
| 9 | ) |
| 10 | |
| 11 | func postUsage(t *testing.T, base, token string, rows []botUsageRow) (int, map[string]any) { |
| 12 | t.Helper() |
| 13 | return postUsageReport(t, base, token, "aaaaaaaaaaaaaaaa", rows) |
| 14 | } |
| 15 | |
| 16 | func postUsageReport(t *testing.T, base, token, reportID string, rows []botUsageRow) (int, map[string]any) { |
| 17 | t.Helper() |
| 18 | raw, _ := json.Marshal(map[string]any{"token": token, "report_id": reportID, "rows": rows}) |
| 19 | resp, err := http.Post(base+"/api/bot-usage", "application/json", bytes.NewReader(raw)) |
| 20 | if err != nil { |
| 21 | t.Fatal(err) |
| 22 | } |
| 23 | defer resp.Body.Close() |
| 24 | var body map[string]any |
| 25 | json.NewDecoder(resp.Body).Decode(&body) |
| 26 | return resp.StatusCode, body |
| 27 | } |
| 28 | |
| 29 | // The relay end to end: a bot's report is validated, forwarded to the cloud |
| 30 | // meter with the instance's bearer token, and the cloud's budget state comes |
| 31 | // back to the sidecar untouched. |
| 32 | func TestBotUsageRelay(t *testing.T) { |
| 33 | ts, s := newTestServer(t) |
| 34 | agent := provisionNamed(t, ts.URL, "Eleven", true) |
| 35 | |
| 36 | var gotAuth string |
| 37 | var gotRows []botUsageRow |
| 38 | cloud := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 39 | gotAuth = r.Header.Get("Authorization") |
| 40 | var body struct { |
| 41 | Rows []botUsageRow `json:"rows"` |
| 42 | } |
| 43 | json.NewDecoder(r.Body).Decode(&body) |
| 44 | gotRows = body.Rows |
| 45 | writeJSON(w, http.StatusOK, map[string]any{"ok": true, "budget": map[string]any{"state": "low"}}) |
| 46 | })) |
| 47 | defer cloud.Close() |
| 48 | s.meterURL, s.meterToken = cloud.URL, "meter-tok" |
| 49 | |
| 50 | rows := []botUsageRow{ |
| 51 | {Day: "2026-07-14", Model: "claude-sonnet-5", Dimension: "llm_input_tokens", Quantity: 150}, |
| 52 | {Day: "2026-07-14", Dimension: "charter_runs", Quantity: 1}, |
| 53 | } |
| 54 | code, body := postUsage(t, ts.URL, agent.Token, rows) |
| 55 | if code != http.StatusOK { |
| 56 | t.Fatalf("relay: %d %v", code, body) |
| 57 | } |
| 58 | if gotAuth != "Bearer meter-tok" { |
| 59 | t.Fatalf("cloud auth = %q", gotAuth) |
| 60 | } |
| 61 | if len(gotRows) != 2 || gotRows[0].Quantity != 150 { |
| 62 | t.Fatalf("cloud rows = %+v", gotRows) |
| 63 | } |
| 64 | if budget, _ := body["budget"].(map[string]any); budget["state"] != "low" { |
| 65 | t.Fatalf("budget not passed through: %v", body) |
| 66 | } |
| 67 | } |
| 68 | |
| 69 | // Unconfigured (self-hosted / BYO key): accept, drop, and say "unmetered". |
| 70 | func TestBotUsageUnmetered(t *testing.T) { |
| 71 | ts, _ := newTestServer(t) |
| 72 | agent := provisionNamed(t, ts.URL, "Eleven", true) |
| 73 | code, body := postUsage(t, ts.URL, agent.Token, []botUsageRow{ |
| 74 | {Day: "2026-07-14", Dimension: "llm_requests", Quantity: 1}}) |
| 75 | if code != http.StatusOK { |
| 76 | t.Fatalf("unmetered: %d %v", code, body) |
| 77 | } |
| 78 | if budget, _ := body["budget"].(map[string]any); budget["state"] != "unmetered" { |
| 79 | t.Fatalf("want unmetered, got %v", body) |
| 80 | } |
| 81 | } |
| 82 | |
| 83 | // The gate: bots only, and malformed rows never leave the instance. |
| 84 | func TestBotUsageGateAndValidation(t *testing.T) { |
| 85 | ts, s := newTestServer(t) |
| 86 | alice := provisionNamed(t, ts.URL, "Alice", false) |
| 87 | agent := provisionNamed(t, ts.URL, "Eleven", true) |
| 88 | s.meterURL, s.meterToken = "http://cloud.invalid", "tok" // must never be dialed below |
| 89 | |
| 90 | good := []botUsageRow{{Day: "2026-07-14", Dimension: "llm_requests", Quantity: 1}} |
| 91 | if code, _ := postUsage(t, ts.URL, alice.Token, good); code != http.StatusForbidden { |
| 92 | t.Fatalf("member token: want 403, got %d", code) |
| 93 | } |
| 94 | if code, _ := postUsage(t, ts.URL, "nonsense", good); code != http.StatusForbidden { |
| 95 | t.Fatalf("stranger: want 403, got %d", code) |
| 96 | } |
| 97 | for name, rows := range map[string][]botUsageRow{ |
| 98 | "empty": {}, |
| 99 | "bad day": {{Day: "yesterday", Dimension: "llm_requests", Quantity: 1}}, |
| 100 | "bad dimension": {{Day: "2026-07-14", Dimension: "plaintext_words", Quantity: 1}}, |
| 101 | "bad quantity": {{Day: "2026-07-14", Dimension: "llm_requests", Quantity: 0}}, |
| 102 | } { |
| 103 | if code, _ := postUsage(t, ts.URL, agent.Token, rows); code != http.StatusBadRequest { |
| 104 | t.Fatalf("%s: want 400, got %d", name, code) |
| 105 | } |
| 106 | } |
| 107 | // A report without a well-formed dedupe id is refused too. |
| 108 | if code, _ := postUsageReport(t, ts.URL, agent.Token, "not hex!", good); code != http.StatusBadRequest { |
| 109 | t.Fatal("bad report id: want 400") |
| 110 | } |
| 111 | } |
| 112 | |
| 113 | // A cloud blip is a 502 — the sidecar keeps its counters and retries. |
| 114 | func TestBotUsageCloudDown(t *testing.T) { |
| 115 | ts, s := newTestServer(t) |
| 116 | agent := provisionNamed(t, ts.URL, "Eleven", true) |
| 117 | dead := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) |
| 118 | dead.Close() // connection refused from here on |
| 119 | s.meterURL, s.meterToken = dead.URL, "tok" |
| 120 | |
| 121 | code, _ := postUsage(t, ts.URL, agent.Token, []botUsageRow{ |
| 122 | {Day: "2026-07-14", Dimension: "llm_requests", Quantity: 1}}) |
| 123 | if code != http.StatusBadGateway { |
| 124 | t.Fatalf("cloud down: want 502, got %d", code) |
| 125 | } |
| 126 | } |
| 127 | |
amadan