+package main
+
+import (
+ "encoding/json"
+ "fmt"
+ "math/rand/v2"
+ "net/http"
+ "os"
+ "path/filepath"
+ "strconv"
+ "strings"
+ "testing"
+
+ "github.com/stretchr/testify/require"
+)
+
+// probeMessages is the length of the generated log for the live probe.
+const probeMessages = 1000
+
+// TestProbeLogHasTheSyntheticMix checks that the generated log has readable
+// ASCII text with the mix of kinds and sizes of the synthetic log. When
+// ONECHAT_PROBE_LOG names a file, the test also writes the log there, one
+// message {kind, text} per line, for "onechat probe".
+func TestProbeLogHasTheSyntheticMix(t *testing.T) {
+ messages := probeLog(probeMessages)
+ require.Len(t, messages, probeMessages)
+
+ counts := make(map[string]int)
+ long := 0
+ for _, message := range messages {
+ counts[message.Kind]++
+ require.True(t, kinds[message.Kind], message.Kind)
+ require.NotEmpty(t, message.Text)
+ for _, char := range message.Text {
+ require.True(t, char == '\n' || (char >= ' ' && char <= '~'), "character %q", char)
+ }
+
+ if len(message.Text) >= 30000 {
+ long++
+ }
+
+ if message.Kind == "work" {
+ require.True(t, strings.HasPrefix(message.Text, "["), message.Text)
+ }
+ }
+
+ require.Greater(t, counts["user"], 20)
+ require.Greater(t, counts["tool"], 300)
+ require.Greater(t, counts["echo"], 300)
+ require.Greater(t, long, 2)
+ require.Equal(t, messages, probeLog(probeMessages))
+
+ path := os.Getenv("ONECHAT_PROBE_LOG")
+ if path == "" {
+ return
+ }
+
+ var data strings.Builder
+ for _, message := range messages {
+ line, err := json.Marshal(message)
+ require.NoError(t, err)
+
+ data.Write(line)
+ data.WriteByte('\n')
+ }
+
+ require.NoError(t, os.WriteFile(path, []byte(data.String()), 0o600))
+}
+
+func TestProbeStopRules(t *testing.T) {
+ quota := func(primary, secondary string) usageRecord {
+ return usageRecord{
+ Quota: map[string]string{
+ "x-codex-primary-used-percent": primary,
+ "x-codex-secondary-used-percent": secondary,
+ "x-codex-plan-type": "plus",
+ },
+ }
+ }
+
+ testCases := []struct {
+ name string
+ records []usageRecord
+ want string
+ }{
+ {
+ name: "no calls",
+ records: nil,
+ want: "",
+ },
+ {
+ name: "growth within the limit",
+ records: []usageRecord{quota("10", "40"), quota("25", "41")},
+ want: "",
+ },
+ {
+ name: "growth over the limit",
+ records: []usageRecord{quota("10", "40"), quota("12", "55.5")},
+ want: "x-codex-secondary-used-percent grew from 40 to 55.5",
+ },
+ {
+ name: "failed call",
+ records: []usageRecord{quota("10", "40"), {
+ Status: 400,
+ Error: "model call: HTTP 400",
+ }},
+ want: "call failed: model call: HTTP 400",
+ },
+ }
+
+ for _, tc := range testCases {
+ t.Run(tc.name, func(t *testing.T) {
+ require.Equal(t, tc.want, probeStop(tc.records))
+ })
+ }
+}
+
+func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) {
+ server := newFakeResponses(t)
+ server.reply = func(w http.ResponseWriter, call fakeCall) {
+ w.Header().Set("X-Codex-Primary-Used-Percent", strconv.Itoa(10+call.number))
+ writeEvents(w, completedEvent([]any{messageItem(strings.Repeat("s", 400))}, 1000, 800, 50))
+ }
+
+ root := t.TempDir()
+ usagePath := filepath.Join(root, "usage.jsonl")
+ client := newClient(server.URL, apiKey(testKey), "onechat-"+probeName, usagePath)
+ memory, err := openMemory(t.Context(), root, newTestClock().Now, newCompactor(client, probeName, defaultConfig().Compactor))
+ require.NoError(t, err)
+
+ appended, stop := runProbe(t.Context(), memory, probeLog(probeMessages), usagePath)
+ sizes, err := modelLines(memory)
+ require.NoError(t, err)
+ require.NoError(t, memory.Close())
+
+ require.Contains(t, stop, "x-codex-primary-used-percent grew from 10 to")
+ require.Less(t, appended, probeMessages)
+ require.NotEmpty(t, sizes)
+
+ records, err := readUsageFile(usagePath)
+ require.NoError(t, err)
+ require.Greater(t, len(records), int(probeQuotaLimit))
+
+ var report strings.Builder
+ require.NoError(t, writeProbeReport(&report, appended, stop, records, sizes))
+ require.Contains(t, report.String(), "stopped early: ")
+ require.Contains(t, report.String(), "cached (80.0%)")
+ require.Contains(t, report.String(), "line bytes: min 400, median 400, mean 400, max 400, 0 over 512")
+ require.Contains(t, report.String(), "quota at the first call: primary-used-percent=10")
+}
+
+// probeLog returns count messages of made-up work on a small bakery order
+// system, with the mix of kinds and sizes of syntheticMessage. The same count
+// always gives the same log.
+func probeLog(count int) []probeEntry {
+ g := probeWriter{
+ random: rand.New(rand.NewPCG(7, 11)),
+ }
+
+ messages := make([]probeEntry, count)
+ for id := range messages {
+ kind, text := g.message()
+ messages[id] = probeEntry{
+ Kind: kind,
+ Text: text,
+ }
+ }
+
+ return messages
+}
+
+// probeWriter writes readable made-up messages.
+type probeWriter struct {
+ random *rand.Rand
+}
+
+func (g *probeWriter) message() (string, string) {
+ roll := g.random.IntN(100)
+ switch {
+ case roll < 5:
+ return "user", g.fill(g.userLine, 10+g.random.IntN(200))
+ case roll < 10:
+ return "reply", g.fill(g.replyLine, 100+g.random.IntN(1500))
+ case roll < 12:
+ return "note", g.fill(g.noteLine, 50+g.random.IntN(600))
+ case roll < 13:
+ return "work", "[" + g.pick(people) + "] report: " + g.fill(g.replyLine, 300+g.random.IntN(3000))
+ case roll < 56:
+ return "tool", g.toolCall(60 + g.random.IntN(400))
+ case roll < 57:
+ return "echo", g.fill(g.outputLine, 30000+g.random.IntN(30000))
+ default:
+ return "echo", g.fill(g.outputLine, 20+g.random.IntN(4000))
+ }
+}
+
+// fill joins lines of write until the text has about size bytes.
+func (g *probeWriter) fill(write func() string, size int) string {
+ var text strings.Builder
+ for text.Len() < size {
+ if text.Len() > 0 {
+ text.WriteByte(' ')
+ }
+
+ text.WriteString(write())
+ }
+
+ return text.String()
+}
+
+func (g *probeWriter) userLine() string {
+ return g.pick([]string{
+ fmt.Sprintf("Please check why the %s %s.", g.pick(things), g.pick(problems)),
+ fmt.Sprintf("Move the %s to %s before %s.", g.pick(things), g.pick(places), g.pick(days)),
+ fmt.Sprintf("Thanks, that works. Next, %s.", g.pick(tasks)),
+ fmt.Sprintf("No, keep the old %s. %s asked for it.", g.pick(things), g.pick(people)),
+ fmt.Sprintf("How many %s orders came in on %s?", g.pick(breads), g.pick(days)),
+ })
+}
+
+func (g *probeWriter) replyLine() string {
+ return g.pick([]string{
+ fmt.Sprintf("I looked at %s. The %s %s because %s.", g.pick(files), g.pick(things), g.pick(problems), g.pick(causes)),
+ fmt.Sprintf("I changed %s, and the tests pass now.", g.pick(changes)),
+ 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)),
+ fmt.Sprintf("Open question: should %s get %s?", g.pick(people), g.pick(tasks)),
+ })
+}
+
+func (g *probeWriter) noteLine() string {
+ return g.pick([]string{
+ fmt.Sprintf("%s prefers %s.", g.pick(people), g.pick(preferences)),
+ 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)),
+ })
+}
+
+func (g *probeWriter) outputLine() string {
+ return g.pick([]string{
+ 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)),
+ fmt.Sprintf("%s:%d: %s %s.", g.pick(files), 1+g.random.IntN(400), g.pick(things), g.pick(problems)),
+ fmt.Sprintf("ok bakery/%s %d.%03ds", g.pick(modules), g.random.IntN(3), g.random.IntN(1000)),
+ })
+}
+
+// toolCall returns a tool name and its JSON input of about size bytes.
+func (g *probeWriter) toolCall(size int) string {
+ name := g.pick([]string{"shell", "read", "write", "edit"})
+ var input map[string]string
+ switch name {
+ case "shell":
+ input = map[string]string{"command": "grep -n '" + g.pick(things) + "' " + g.pick(files)}
+ case "read":
+ input = map[string]string{"path": g.pick(files)}
+ default:
+ input = map[string]string{
+ "path": g.pick(files),
+ "text": g.fill(g.replyLine, max(0, size-80)),
+ }
+ }
+
+ data, err := json.Marshal(input)
+ if err != nil {
+ panic(err)
+ }
+
+ return name + " " + string(data)
+}
+
+func (g *probeWriter) pick(words []string) string {
+ return words[g.random.IntN(len(words))]
+}
+
+var (
+ people = []string{"Mara", "Joel", "Ines", "Tomas", "Priya", "Lena"}
+ breads = []string{"rye", "sourdough", "baguette", "brioche", "spelt", "seed loaf"}
+ days = []string{"Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday"}
+ things = []string{"order form", "price list", "delivery route", "stock count", "invoice export", "oven schedule"}
+ problems = []string{"shows the wrong total", "skips the last row", "is slow at noon", "loses the customer note", "rounds the tax down"}
+ places = []string{"the front page", "the weekly report", "the shared folder", "the admin menu"}
+ tasks = []string{"add a field for allergies", "send the invoices on Friday", "split the route in two", "print labels for the shelf"}
+ causes = []string{"the date has no time zone", "the cache keeps old prices", "the loop stops one row early", "the form trims the note"}
+ 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"}
+ preferences = []string{"short answers", "a call before a change to prices", "invoices as PDF", "reports in the morning"}
+ files = []string{"orders/form.go", "prices/list.go", "routes/plan.go", "stock/count.go", "export/invoice.go", "ovens/schedule.go"}
+ modules = []string{"orders", "prices", "routes", "stock", "export", "ovens"}
+)