git.giacomolli.uk - onechat/blob - probe_test.go

   1 package main
   2
   3 import (
   4         "encoding/json"
   5         "fmt"
   6         "math/rand/v2"
   7         "net/http"
   8         "os"
   9         "path/filepath"
  10         "strconv"
  11         "strings"
  12         "testing"
  13
  14         "github.com/stretchr/testify/require"
  15 )
  16
  17 // probeMessages is the length of the generated log for the live probe.
  18 const probeMessages = 1000
  19
  20 // TestProbeLogHasTheSyntheticMix checks that the generated log has readable
  21 // ASCII text with the mix of kinds and sizes of the synthetic log. When
  22 // ONECHAT_PROBE_LOG names a file, the test also writes the log there, one
  23 // message {kind, text} per line, for "onechat probe".
  24 func TestProbeLogHasTheSyntheticMix(t *testing.T) {
  25         messages := probeLog(probeMessages)
  26         require.Len(t, messages, probeMessages)
  27
  28         counts := make(map[string]int)
  29         long := 0
  30         for _, message := range messages {
  31                 counts[message.Kind]++
  32                 require.True(t, kinds[message.Kind], message.Kind)
  33                 require.NotEmpty(t, message.Text)
  34                 for _, char := range message.Text {
  35                         require.True(t, char == '\n' || (char >= ' ' && char <= '~'), "character %q", char)
  36                 }
  37
  38                 if len(message.Text) >= 30000 {
  39                         long++
  40                 }
  41
  42                 if message.Kind == "work" {
  43                         require.True(t, strings.HasPrefix(message.Text, "["), message.Text)
  44                 }
  45         }
  46
  47         require.Greater(t, counts["user"], 20)
  48         require.Greater(t, counts["tool"], 300)
  49         require.Greater(t, counts["echo"], 300)
  50         require.Greater(t, long, 2)
  51         require.Equal(t, messages, probeLog(probeMessages))
  52
  53         path := os.Getenv("ONECHAT_PROBE_LOG")
  54         if path == "" {
  55                 return
  56         }
  57
  58         var data strings.Builder
  59         for _, message := range messages {
  60                 line, err := json.Marshal(message)
  61                 require.NoError(t, err)
  62
  63                 data.Write(line)
  64                 data.WriteByte('\n')
  65         }
  66
  67         require.NoError(t, os.WriteFile(path, []byte(data.String()), 0o600))
  68 }
  69
  70 func TestProbeStopRules(t *testing.T) {
  71         quota := func(primary, secondary string) usageRecord {
  72                 return usageRecord{
  73                         Quota: map[string]string{
  74                                 "x-codex-primary-used-percent":   primary,
  75                                 "x-codex-secondary-used-percent": secondary,
  76                                 "x-codex-plan-type":              "plus",
  77                         },
  78                 }
  79         }
  80
  81         testCases := []struct {
  82                 name    string
  83                 records []usageRecord
  84                 want    string
  85         }{
  86                 {
  87                         name:    "no calls",
  88                         records: nil,
  89                         want:    "",
  90                 },
  91                 {
  92                         name:    "growth within the limit",
  93                         records: []usageRecord{quota("10", "40"), quota("25", "41")},
  94                         want:    "",
  95                 },
  96                 {
  97                         name:    "growth over the limit",
  98                         records: []usageRecord{quota("10", "40"), quota("12", "55.5")},
  99                         want:    "x-codex-secondary-used-percent grew from 40 to 55.5",
 100                 },
 101                 {
 102                         name: "failed call",
 103                         records: []usageRecord{quota("10", "40"), {
 104                                 Status: 400,
 105                                 Error:  "model call: HTTP 400",
 106                         }},
 107                         want: "call failed: model call: HTTP 400",
 108                 },
 109         }
 110
 111         for _, tc := range testCases {
 112                 t.Run(tc.name, func(t *testing.T) {
 113                         require.Equal(t, tc.want, probeStop(tc.records))
 114                 })
 115         }
 116 }
 117
 118 func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) {
 119         server := newFakeResponses(t)
 120         server.reply = func(w http.ResponseWriter, call fakeCall) {
 121                 w.Header().Set("X-Codex-Primary-Used-Percent", strconv.Itoa(10+call.number))
 122                 writeEvents(w, completedEvent([]any{messageItem(strings.Repeat("s", 400))}, 1000, 800, 50))
 123         }
 124
 125         root := t.TempDir()
 126         usagePath := filepath.Join(root, "usage.jsonl")
 127         client := newClient(server.URL, apiKey(testKey), "onechat-"+probeName, usagePath)
 128         memory, err := openMemory(t.Context(), root, newTestClock().Now, newCompactor(client, probeName, defaultConfig().Compactor))
 129         require.NoError(t, err)
 130
 131         appended, stop := runProbe(t.Context(), memory, probeLog(probeMessages), usagePath)
 132         sizes, err := modelLines(memory)
 133         require.NoError(t, err)
 134         require.NoError(t, memory.Close())
 135
 136         require.Contains(t, stop, "x-codex-primary-used-percent grew from 10 to")
 137         require.Less(t, appended, probeMessages)
 138         require.NotEmpty(t, sizes)
 139
 140         records, err := readUsageFile(usagePath)
 141         require.NoError(t, err)
 142         require.Greater(t, len(records), int(probeQuotaLimit))
 143
 144         var report strings.Builder
 145         require.NoError(t, writeProbeReport(&report, appended, stop, records, sizes))
 146         require.Contains(t, report.String(), "stopped early: ")
 147         require.Contains(t, report.String(), "cached (80.0%)")
 148         require.Contains(t, report.String(), "line bytes: min 400, median 400, mean 400, max 400, 0 over 512")
 149         require.Contains(t, report.String(), "quota at the first call: primary-used-percent=10")
 150 }
 151
 152 // probeLog returns count messages of made-up work on a small bakery order
 153 // system, with the mix of kinds and sizes of syntheticMessage. The same count
 154 // always gives the same log.
 155 func probeLog(count int) []probeEntry {
 156         g := probeWriter{
 157                 random: rand.New(rand.NewPCG(7, 11)),
 158         }
 159
 160         messages := make([]probeEntry, count)
 161         for id := range messages {
 162                 kind, text := g.message()
 163                 messages[id] = probeEntry{
 164                         Kind: kind,
 165                         Text: text,
 166                 }
 167         }
 168
 169         return messages
 170 }
 171
 172 // probeWriter writes readable made-up messages.
 173 type probeWriter struct {
 174         random *rand.Rand
 175 }
 176
 177 func (g *probeWriter) message() (string, string) {
 178         roll := g.random.IntN(100)
 179         switch {
 180         case roll < 5:
 181                 return "user", g.fill(g.userLine, 10+g.random.IntN(200))
 182         case roll < 10:
 183                 return "reply", g.fill(g.replyLine, 100+g.random.IntN(1500))
 184         case roll < 12:
 185                 return "note", g.fill(g.noteLine, 50+g.random.IntN(600))
 186         case roll < 13:
 187                 return "work", "[" + g.pick(people) + "] report: " + g.fill(g.replyLine, 300+g.random.IntN(3000))
 188         case roll < 56:
 189                 return "tool", g.toolCall(60 + g.random.IntN(400))
 190         case roll < 57:
 191                 return "echo", g.fill(g.outputLine, 30000+g.random.IntN(30000))
 192         default:
 193                 return "echo", g.fill(g.outputLine, 20+g.random.IntN(4000))
 194         }
 195 }
 196
 197 // fill joins lines of write until the text has about size bytes.
 198 func (g *probeWriter) fill(write func() string, size int) string {
 199         var text strings.Builder
 200         for text.Len() < size {
 201                 if text.Len() > 0 {
 202                         text.WriteByte(' ')
 203                 }
 204
 205                 text.WriteString(write())
 206         }
 207
 208         return text.String()
 209 }
 210
 211 func (g *probeWriter) userLine() string {
 212         return g.pick([]string{
 213                 fmt.Sprintf("Please check why the %s %s.", g.pick(things), g.pick(problems)),
 214                 fmt.Sprintf("Move the %s to %s before %s.", g.pick(things), g.pick(places), g.pick(days)),
 215                 fmt.Sprintf("Thanks, that works. Next, %s.", g.pick(tasks)),
 216                 fmt.Sprintf("No, keep the old %s. %s asked for it.", g.pick(things), g.pick(people)),
 217                 fmt.Sprintf("How many %s orders came in on %s?", g.pick(breads), g.pick(days)),
 218         })
 219 }
 220
 221 func (g *probeWriter) replyLine() string {
 222         return g.pick([]string{
 223                 fmt.Sprintf("I looked at %s. The %s %s because %s.", g.pick(files), g.pick(things), g.pick(problems), g.pick(causes)),
 224                 fmt.Sprintf("I changed %s, and the tests pass now.", g.pick(changes)),
 225                 fmt.Sprintf("There were %d %s orders on %s, %d of them late.", 10+g.random.IntN(90), g.pick(breads), g.pick(days), g.random.IntN(10)),
 226                 fmt.Sprintf("Open question: should %s get %s?", g.pick(people), g.pick(tasks)),
 227         })
 228 }
 229
 230 func (g *probeWriter) noteLine() string {
 231         return g.pick([]string{
 232                 fmt.Sprintf("%s prefers %s.", g.pick(people), g.pick(preferences)),
 233                 fmt.Sprintf("The bakery opens at %d and closes at %d on %s.", 5+g.random.IntN(3), 17+g.random.IntN(4), g.pick(days)),
 234         })
 235 }
 236
 237 func (g *probeWriter) outputLine() string {
 238         return g.pick([]string{
 239                 fmt.Sprintf("%02d:%02d order %d for %d %s, packed by %s.", 6+g.random.IntN(12), g.random.IntN(60), 1000+g.random.IntN(9000), 1+g.random.IntN(24), g.pick(breads), g.pick(people)),
 240                 fmt.Sprintf("%s:%d: %s %s.", g.pick(files), 1+g.random.IntN(400), g.pick(things), g.pick(problems)),
 241                 fmt.Sprintf("ok   bakery/%s %d.%03ds", g.pick(modules), g.random.IntN(3), g.random.IntN(1000)),
 242         })
 243 }
 244
 245 // toolCall returns a tool name and its JSON input of about size bytes.
 246 func (g *probeWriter) toolCall(size int) string {
 247         name := g.pick([]string{"shell", "read", "write", "edit"})
 248         var input map[string]string
 249         switch name {
 250         case "shell":
 251                 input = map[string]string{"command": "grep -n '" + g.pick(things) + "' " + g.pick(files)}
 252         case "read":
 253                 input = map[string]string{"path": g.pick(files)}
 254         default:
 255                 input = map[string]string{
 256                         "path": g.pick(files),
 257                         "text": g.fill(g.replyLine, max(0, size-80)),
 258                 }
 259         }
 260
 261         data, err := json.Marshal(input)
 262         if err != nil {
 263                 panic(err)
 264         }
 265
 266         return name + " " + string(data)
 267 }
 268
 269 func (g *probeWriter) pick(words []string) string {
 270         return words[g.random.IntN(len(words))]
 271 }
 272
 273 var (
 274         people      = []string{"Mara", "Joel", "Ines", "Tomas", "Priya", "Lena"}
 275         breads      = []string{"rye", "sourdough", "baguette", "brioche", "spelt", "seed loaf"}
 276         days        = []string{"Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday"}
 277         things      = []string{"order form", "price list", "delivery route", "stock count", "invoice export", "oven schedule"}
 278         problems    = []string{"shows the wrong total", "skips the last row", "is slow at noon", "loses the customer note", "rounds the tax down"}
 279         places      = []string{"the front page", "the weekly report", "the shared folder", "the admin menu"}
 280         tasks       = []string{"add a field for allergies", "send the invoices on Friday", "split the route in two", "print labels for the shelf"}
 281         causes      = []string{"the date has no time zone", "the cache keeps old prices", "the loop stops one row early", "the form trims the note"}
 282         changes     = []string{"the total to use cents", "the route to sort by street", "the export to add the tax line", "the form to keep the note"}
 283         preferences = []string{"short answers", "a call before a change to prices", "invoices as PDF", "reports in the morning"}
 284         files       = []string{"orders/form.go", "prices/list.go", "routes/plan.go", "stock/count.go", "export/invoice.go", "ovens/schedule.go"}
 285         modules     = []string{"orders", "prices", "routes", "stock", "export", "ovens"}
 286 )