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         "sync"
  13         "testing"
  14         "time"
  15
  16         "github.com/stretchr/testify/require"
  17 )
  18
  19 // probeMessages is the length of the generated log for the live probe.
  20 const probeMessages = 1000
  21
  22 // TestProbeLogHasTheSyntheticMix checks that the generated log has readable
  23 // ASCII text with the mix of kinds and sizes of the synthetic log. When
  24 // ONECHAT_PROBE_LOG names a file, the test also writes the log there, one
  25 // message {kind, text} per line, for "onechat probe".
  26 func TestProbeLogHasTheSyntheticMix(t *testing.T) {
  27         messages := probeLog(probeMessages)
  28         require.Len(t, messages, probeMessages)
  29
  30         counts := make(map[string]int)
  31         long := 0
  32         for _, message := range messages {
  33                 counts[message.Kind]++
  34                 require.True(t, kinds[message.Kind], message.Kind)
  35                 require.NotEmpty(t, message.Text)
  36                 for _, char := range message.Text {
  37                         require.True(t, char == '\n' || (char >= ' ' && char <= '~'), "character %q", char)
  38                 }
  39
  40                 if len(message.Text) >= 30000 {
  41                         long++
  42                 }
  43
  44                 if message.Kind == "work" {
  45                         require.True(t, strings.HasPrefix(message.Text, "["), message.Text)
  46                 }
  47         }
  48
  49         require.Greater(t, counts["user"], 20)
  50         require.Greater(t, counts["tool"], 300)
  51         require.Greater(t, counts["echo"], 300)
  52         require.Greater(t, long, 2)
  53         require.Equal(t, messages, probeLog(probeMessages))
  54
  55         path := os.Getenv("ONECHAT_PROBE_LOG")
  56         if path == "" {
  57                 return
  58         }
  59
  60         var data strings.Builder
  61         for _, message := range messages {
  62                 line, err := json.Marshal(message)
  63                 require.NoError(t, err)
  64
  65                 data.Write(line)
  66                 data.WriteByte('\n')
  67         }
  68
  69         require.NoError(t, os.WriteFile(path, []byte(data.String()), 0o600))
  70 }
  71
  72 func TestProbeStopRules(t *testing.T) {
  73         quota := func(primary, secondary string) ModelUsageRecord {
  74                 return ModelUsageRecord{
  75                         Quota: map[string]string{
  76                                 "x-codex-primary-used-percent":   primary,
  77                                 "x-codex-secondary-used-percent": secondary,
  78                                 "x-codex-plan-type":              "plus",
  79                         },
  80                 }
  81         }
  82
  83         testCases := []struct {
  84                 name    string
  85                 records []ModelUsageRecord
  86                 want    string
  87         }{
  88                 {
  89                         name:    "no calls",
  90                         records: nil,
  91                         want:    "",
  92                 },
  93                 {
  94                         name:    "growth within the limit",
  95                         records: []ModelUsageRecord{quota("10", "40"), quota("25", "41")},
  96                         want:    "",
  97                 },
  98                 {
  99                         name:    "growth over the limit",
 100                         records: []ModelUsageRecord{quota("10", "40"), quota("12", "55.5")},
 101                         want:    "x-codex-secondary-used-percent grew from 40 to 55.5",
 102                 },
 103                 {
 104                         name: "failed call",
 105                         records: []ModelUsageRecord{quota("10", "40"), {
 106                                 Status: 400,
 107                                 Error:  "model call: HTTP 400",
 108                         }},
 109                         want: "call failed: model call: HTTP 400",
 110                 },
 111         }
 112
 113         for _, tc := range testCases {
 114                 t.Run(tc.name, func(t *testing.T) {
 115                         require.Equal(t, tc.want, probeStop(tc.records))
 116                 })
 117         }
 118 }
 119
 120 func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) {
 121         server := newModelFake(t)
 122         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
 123                 w.Header().Set("X-Codex-Primary-Used-Percent", strconv.Itoa(10+call.number))
 124                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem(strings.Repeat("s", 400))}, 1000, 800, 50))
 125         }
 126
 127         run := newTestProbeRun(t, server)
 128         appended, stop := run.Run(t.Context(), probeLog(probeMessages))
 129         sizes, err := modelLines(run.memory)
 130         require.NoError(t, err)
 131         require.NoError(t, run.memory.Close())
 132
 133         require.Regexp(t, `^x-codex-primary-used-percent grew from \d+ to \d+$`, stop)
 134         require.Less(t, appended, probeMessages)
 135         require.NotEmpty(t, sizes)
 136
 137         records, err := readUsageFile(run.usagePath)
 138         require.NoError(t, err)
 139         require.Greater(t, len(records), int(probeQuotaLimit))
 140
 141         var report strings.Builder
 142         require.NoError(t, writeProbeReport(&report, appended, stop, records, sizes))
 143         require.Contains(t, report.String(), "stopped early: ")
 144         require.Contains(t, report.String(), "cached (80.0%)")
 145         require.Contains(t, report.String(), "line bytes: min 400, median 400, mean 400, max 400, 0 over 512")
 146         require.Contains(t, report.String(), "calls at once: at most ")
 147         require.Contains(t, report.String(), "quota at the first call: primary-used-percent=")
 148 }
 149
 150 func TestProbeRunStopsAtTheCallLimit(t *testing.T) {
 151         server := newModelFake(t)
 152         run := newTestProbeRun(t, server)
 153         run.client.LimitCalls(5)
 154
 155         entries := make([]probeEntry, 40)
 156         for n := range entries {
 157                 entries[n] = probeEntry{
 158                         Kind: "echo",
 159                         Text: strings.Repeat("long output ", 100),
 160                 }
 161         }
 162
 163         _, stop := run.Run(t.Context(), entries)
 164         require.NoError(t, run.memory.Close())
 165         require.Equal(t, "call limit of 5 reached", stop)
 166         require.Equal(t, 5, server.Calls())
 167 }
 168
 169 func TestProbeRunWaitsOnlyBeforeUserMessages(t *testing.T) {
 170         gate := make(chan struct{})
 171         var mutex sync.Mutex
 172         running, most := 0, 0
 173         server := newModelFake(t)
 174         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
 175                 mutex.Lock()
 176                 running++
 177                 most = max(most, running)
 178                 mutex.Unlock()
 179
 180                 <-gate
 181
 182                 mutex.Lock()
 183                 running--
 184                 mutex.Unlock()
 185
 186                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("a line")}, 100, 0, 5))
 187         }
 188
 189         run := newTestProbeRun(t, server)
 190         long := strings.Repeat("long output ", 100)
 191         entries := []probeEntry{
 192                 {
 193                         Kind: "echo",
 194                         Text: long,
 195                 },
 196                 {
 197                         Kind: "tool",
 198                         Text: long,
 199                 },
 200                 {
 201                         Kind: "user",
 202                         Text: "next",
 203                 },
 204                 {
 205                         Kind: "reply",
 206                         Text: long,
 207                 },
 208         }
 209
 210         done := make(chan string)
 211         go func() {
 212                 _, stop := run.Run(t.Context(), entries)
 213                 done <- stop
 214         }()
 215
 216         require.Eventually(t, func() bool {
 217                 return server.Calls() == 2
 218         }, 10*time.Second, time.Millisecond)
 219
 220         run.memory.mutex.Lock()
 221         appended := run.memory.log.Len()
 222         run.memory.mutex.Unlock()
 223         require.Equal(t, 2, appended)
 224
 225         close(gate)
 226         require.Empty(t, <-done)
 227         require.NoError(t, run.memory.Close())
 228
 229         mutex.Lock()
 230         defer mutex.Unlock()
 231
 232         require.Equal(t, 2, most)
 233         require.Equal(t, 3, server.Calls())
 234 }
 235
 236 func TestSlotSummarizerLimitsTheBuildsAtOnce(t *testing.T) {
 237         inner := newFakeSummarizer()
 238         inner.gate = make(chan struct{})
 239         summarizer := slotSummarizer{
 240                 inner: inner,
 241                 slots: make(chan struct{}, 2),
 242         }
 243
 244         var group sync.WaitGroup
 245         for n := range 6 {
 246                 group.Go(func() {
 247                         _, err := summarizer.Summarize(t.Context(), Task{
 248                                 Level: 0,
 249                                 Index: n,
 250                                 Input: []string{"user: " + strings.Repeat("x", 600)},
 251                         })
 252                         require.NoError(t, err)
 253                 })
 254         }
 255
 256         require.Eventually(t, func() bool {
 257                 return inner.Active() == 2
 258         }, 10*time.Second, time.Millisecond)
 259
 260         close(inner.gate)
 261         group.Wait()
 262         require.Equal(t, 2, inner.peak)
 263 }
 264
 265 func TestReadUsageFileSortsByCallStart(t *testing.T) {
 266         path := filepath.Join(t.TempDir(), "usage.jsonl")
 267         usage := ModelUsageLog{
 268                 path: path,
 269         }
 270
 271         for _, date := range []string{"2026-01-01T12:00:02Z", "2026-01-01T12:00:00.5Z", "2026-01-01T12:00:01.25Z"} {
 272                 require.NoError(t, usage.Write(ModelUsageRecord{
 273                         Date: date,
 274                 }))
 275         }
 276
 277         records, err := readUsageFile(path)
 278         require.NoError(t, err)
 279         require.Equal(t, []string{"2026-01-01T12:00:00.5Z", "2026-01-01T12:00:01.25Z", "2026-01-01T12:00:02Z"}, []string{records[0].Date, records[1].Date, records[2].Date})
 280 }
 281
 282 func TestTimeTextCountsTheCallsAtOnce(t *testing.T) {
 283         start := time.Date(2026, 1, 1, 12, 0, 0, 0, time.UTC)
 284         record := func(offset, seconds float64) ModelUsageRecord {
 285                 return ModelUsageRecord{
 286                         Date:    start.Add(time.Duration(offset * float64(time.Second))).Format(time.RFC3339Nano),
 287                         Seconds: seconds,
 288                 }
 289         }
 290
 291         text := timeText([]ModelUsageRecord{record(0, 10), record(2, 4), record(3, 2), record(10, 10)})
 292         require.Equal(t, "time: 20 s in all, 12.0 calls per minute, 6.5 s per call (median 10.0, max 10.0), calls at once: at most 3, 1.3 on average", text)
 293 }
 294
 295 // probeLog returns count messages of made-up work on a small bakery order
 296 // system, with the mix of kinds and sizes of syntheticMessage. The same count
 297 // always gives the same log.
 298 func probeLog(count int) []probeEntry {
 299         g := probeWriter{
 300                 random: rand.New(rand.NewPCG(7, 11)),
 301         }
 302
 303         messages := make([]probeEntry, count)
 304         for id := range messages {
 305                 kind, text := g.message()
 306                 messages[id] = probeEntry{
 307                         Kind: kind,
 308                         Text: text,
 309                 }
 310         }
 311
 312         return messages
 313 }
 314
 315 // probeWriter writes readable made-up messages.
 316 type probeWriter struct {
 317         random *rand.Rand
 318 }
 319
 320 func (g *probeWriter) message() (string, string) {
 321         roll := g.random.IntN(100)
 322         switch {
 323         case roll < 5:
 324                 return "user", g.fill(g.userLine, 10+g.random.IntN(200))
 325         case roll < 10:
 326                 return "reply", g.fill(g.replyLine, 100+g.random.IntN(1500))
 327         case roll < 12:
 328                 return "note", g.fill(g.noteLine, 50+g.random.IntN(600))
 329         case roll < 13:
 330                 return "work", "[" + g.pick(people) + "] report: " + g.fill(g.replyLine, 300+g.random.IntN(3000))
 331         case roll < 56:
 332                 return "tool", g.toolCall(60 + g.random.IntN(400))
 333         case roll < 57:
 334                 return "echo", g.fill(g.outputLine, 30000+g.random.IntN(30000))
 335         default:
 336                 return "echo", g.fill(g.outputLine, 20+g.random.IntN(4000))
 337         }
 338 }
 339
 340 // fill joins lines of write until the text has about size bytes.
 341 func (g *probeWriter) fill(write func() string, size int) string {
 342         var text strings.Builder
 343         for text.Len() < size {
 344                 if text.Len() > 0 {
 345                         text.WriteByte(' ')
 346                 }
 347
 348                 text.WriteString(write())
 349         }
 350
 351         return text.String()
 352 }
 353
 354 func (g *probeWriter) userLine() string {
 355         return g.pick([]string{
 356                 fmt.Sprintf("Please check why the %s %s.", g.pick(things), g.pick(problems)),
 357                 fmt.Sprintf("Move the %s to %s before %s.", g.pick(things), g.pick(places), g.pick(days)),
 358                 fmt.Sprintf("Thanks, that works. Next, %s.", g.pick(tasks)),
 359                 fmt.Sprintf("No, keep the old %s. %s asked for it.", g.pick(things), g.pick(people)),
 360                 fmt.Sprintf("How many %s orders came in on %s?", g.pick(breads), g.pick(days)),
 361         })
 362 }
 363
 364 func (g *probeWriter) replyLine() string {
 365         return g.pick([]string{
 366                 fmt.Sprintf("I looked at %s. The %s %s because %s.", g.pick(files), g.pick(things), g.pick(problems), g.pick(causes)),
 367                 fmt.Sprintf("I changed %s, and the tests pass now.", g.pick(changes)),
 368                 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)),
 369                 fmt.Sprintf("Open question: should %s get %s?", g.pick(people), g.pick(tasks)),
 370         })
 371 }
 372
 373 func (g *probeWriter) noteLine() string {
 374         return g.pick([]string{
 375                 fmt.Sprintf("%s prefers %s.", g.pick(people), g.pick(preferences)),
 376                 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)),
 377         })
 378 }
 379
 380 func (g *probeWriter) outputLine() string {
 381         return g.pick([]string{
 382                 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)),
 383                 fmt.Sprintf("%s:%d: %s %s.", g.pick(files), 1+g.random.IntN(400), g.pick(things), g.pick(problems)),
 384                 fmt.Sprintf("ok   bakery/%s %d.%03ds", g.pick(modules), g.random.IntN(3), g.random.IntN(1000)),
 385         })
 386 }
 387
 388 // toolCall returns a tool name and its JSON input of about size bytes.
 389 func (g *probeWriter) toolCall(size int) string {
 390         name := g.pick([]string{"shell", "read", "write", "edit"})
 391         var input map[string]string
 392         switch name {
 393         case "shell":
 394                 input = map[string]string{"command": "grep -n '" + g.pick(things) + "' " + g.pick(files)}
 395         case "read":
 396                 input = map[string]string{"path": g.pick(files)}
 397         default:
 398                 input = map[string]string{
 399                         "path": g.pick(files),
 400                         "text": g.fill(g.replyLine, max(0, size-80)),
 401                 }
 402         }
 403
 404         data, err := json.Marshal(input)
 405         if err != nil {
 406                 panic(err)
 407         }
 408
 409         return name + " " + string(data)
 410 }
 411
 412 func (g *probeWriter) pick(words []string) string {
 413         return words[g.random.IntN(len(words))]
 414 }
 415
 416 func newTestProbeRun(t *testing.T, server *ModelFake) *probeRun {
 417         t.Helper()
 418
 419         root := t.TempDir()
 420         usagePath := filepath.Join(root, "usage.jsonl")
 421         client := newModelClient(server.URL, ModelAPIKey(modelTestKey), "onechat-"+probeName, usagePath)
 422         memory, err := openMemory(t.Context(), root, newTestClock().Now, newCompactor(client, probePrompt, defaultConfig().Compactor))
 423         require.NoError(t, err)
 424
 425         return &probeRun{
 426                 memory:    memory,
 427                 client:    client,
 428                 usagePath: usagePath,
 429         }
 430 }
 431
 432 var (
 433         people      = []string{"Mara", "Joel", "Ines", "Tomas", "Priya", "Lena"}
 434         breads      = []string{"rye", "sourdough", "baguette", "brioche", "spelt", "seed loaf"}
 435         days        = []string{"Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday"}
 436         things      = []string{"order form", "price list", "delivery route", "stock count", "invoice export", "oven schedule"}
 437         problems    = []string{"shows the wrong total", "skips the last row", "is slow at noon", "loses the customer note", "rounds the tax down"}
 438         places      = []string{"the front page", "the weekly report", "the shared folder", "the admin menu"}
 439         tasks       = []string{"add a field for allergies", "send the invoices on Friday", "split the route in two", "print labels for the shelf"}
 440         causes      = []string{"the date has no time zone", "the cache keeps old prices", "the loop stops one row early", "the form trims the note"}
 441         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"}
 442         preferences = []string{"short answers", "a call before a change to prices", "invoices as PDF", "reports in the morning"}
 443         files       = []string{"orders/form.go", "prices/list.go", "routes/plan.go", "stock/count.go", "export/invoice.go", "ovens/schedule.go"}
 444         modules     = []string{"orders", "prices", "routes", "stock", "export", "ovens"}
 445 )