X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/ec38dadaaf127d8bd2ffb72df450e8d3f5ac88c0..016d46008fc9726d2617e4e3ea098113f34c7636:/probe_test.go diff --git a/probe_test.go b/probe_test.go index 265d49d..298e057 100644 --- a/probe_test.go +++ b/probe_test.go @@ -2,12 +2,8 @@ package main import ( "encoding/json" - "fmt" - "math/rand/v2" - "net/http" "os" "path/filepath" - "strconv" "strings" "sync" "testing" @@ -16,9 +12,6 @@ import ( "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 @@ -31,8 +24,9 @@ func TestProbeLogHasTheSyntheticMix(t *testing.T) { long := 0 for _, message := range messages { counts[message.Kind]++ - require.True(t, kinds[message.Kind], message.Kind) + require.True(t, logKinds[message.Kind], message.Kind) require.NotEmpty(t, message.Text) + for _, char := range message.Text { require.True(t, char == '\n' || (char >= ' ' && char <= '~'), "character %q", char) } @@ -58,6 +52,7 @@ func TestProbeLogHasTheSyntheticMix(t *testing.T) { } var data strings.Builder + for _, message := range messages { line, err := json.Marshal(message) require.NoError(t, err) @@ -69,168 +64,14 @@ func TestProbeLogHasTheSyntheticMix(t *testing.T) { require.NoError(t, os.WriteFile(path, []byte(data.String()), 0o600)) } -func TestProbeStopRules(t *testing.T) { - quota := func(primary, secondary string) ModelUsageRecord { - return ModelUsageRecord{ - 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 []ModelUsageRecord - want string - }{ - { - name: "no calls", - records: nil, - want: "", - }, - { - name: "growth within the limit", - records: []ModelUsageRecord{quota("10", "40"), quota("25", "41")}, - want: "", - }, - { - name: "growth over the limit", - records: []ModelUsageRecord{quota("10", "40"), quota("12", "55.5")}, - want: "x-codex-secondary-used-percent grew from 40 to 55.5", - }, - { - name: "failed call", - records: []ModelUsageRecord{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 := newModelFake(t) - server.reply = func(w http.ResponseWriter, call ModelFakeCall) { - w.Header().Set("X-Codex-Primary-Used-Percent", strconv.Itoa(10+call.number)) - modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem(strings.Repeat("s", 400))}, 1000, 800, 50)) - } - - run := newTestProbeRun(t, server) - appended, stop := run.Run(t.Context(), probeLog(probeMessages)) - sizes, err := probeModelLines(run.memory) - require.NoError(t, err) - require.NoError(t, run.memory.Close()) +func TestProbeReadIgnoresSourceDates(t *testing.T) { + path := filepath.Join(t.TempDir(), "probe.jsonl") + require.NoError(t, os.WriteFile(path, []byte(`{"kind":"user","text":"A probe message","date":"2025-01-01T00:00:00Z"}`+"\n"), 0o600)) - require.Regexp(t, `^x-codex-primary-used-percent grew from \d+ to \d+$`, stop) - require.Less(t, appended, probeMessages) - require.NotEmpty(t, sizes) - - records, err := probeReadUsage(run.usagePath) + entries, err := probeReadLog(path, 1) require.NoError(t, err) - require.Greater(t, len(records), int(probeQuotaLimit)) - - var report strings.Builder - require.NoError(t, probeWriteReport(&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(), "calls at once: at most ") - require.Contains(t, report.String(), "quota at the first call: primary-used-percent=") -} - -func TestProbeRunStopsAtTheCallLimit(t *testing.T) { - server := newModelFake(t) - run := newTestProbeRun(t, server) - run.client.LimitCalls(5) - - entries := make([]ProbeEntry, 40) - for n := range entries { - entries[n] = ProbeEntry{ - Kind: "echo", - Text: strings.Repeat("long output ", 100), - } - } - - _, stop := run.Run(t.Context(), entries) - require.NoError(t, run.memory.Close()) - require.Equal(t, "call limit of 5 reached", stop) - require.Equal(t, 5, server.Calls()) -} - -func TestProbeRunWaitsOnlyBeforeUserMessages(t *testing.T) { - gate := make(chan struct{}) - var mutex sync.Mutex - running, most := 0, 0 - server := newModelFake(t) - server.reply = func(w http.ResponseWriter, call ModelFakeCall) { - mutex.Lock() - running++ - most = max(most, running) - mutex.Unlock() - - <-gate - - mutex.Lock() - running-- - mutex.Unlock() - - modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("a line")}, 100, 0, 5)) - } - - run := newTestProbeRun(t, server) - long := strings.Repeat("long output ", 100) - entries := []ProbeEntry{ - { - Kind: "echo", - Text: long, - }, - { - Kind: "tool", - Text: long, - }, - { - Kind: "user", - Text: "next", - }, - { - Kind: "reply", - Text: long, - }, - } - - done := make(chan string) - go func() { - _, stop := run.Run(t.Context(), entries) - done <- stop - }() - - require.Eventually(t, func() bool { - return server.Calls() == 2 - }, 10*time.Second, time.Millisecond) - - run.memory.mutex.Lock() - appended := run.memory.log.Len() - run.memory.mutex.Unlock() - require.Equal(t, 2, appended) - - close(gate) - require.Empty(t, <-done) - require.NoError(t, run.memory.Close()) - - mutex.Lock() - defer mutex.Unlock() - - require.Equal(t, 2, most) - require.Equal(t, 3, server.Calls()) + require.Len(t, entries, 1) + require.Empty(t, entries[0].Date) } func TestSlotSummarizerLimitsTheBuildsAtOnce(t *testing.T) { @@ -242,6 +83,7 @@ func TestSlotSummarizerLimitsTheBuildsAtOnce(t *testing.T) { } var group sync.WaitGroup + for n := range 6 { group.Go(func() { _, err := summarizer.Summarize(t.Context(), Task{ @@ -261,185 +103,3 @@ func TestSlotSummarizerLimitsTheBuildsAtOnce(t *testing.T) { group.Wait() require.Equal(t, 2, inner.peak) } - -func TestReadUsageFileSortsByCallStart(t *testing.T) { - path := filepath.Join(t.TempDir(), "usage.jsonl") - usage := ModelUsageLog{ - path: path, - } - - for _, date := range []string{"2026-01-01T12:00:02Z", "2026-01-01T12:00:00.5Z", "2026-01-01T12:00:01.25Z"} { - require.NoError(t, usage.Write(ModelUsageRecord{ - Date: date, - })) - } - - records, err := probeReadUsage(path) - require.NoError(t, err) - 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}) -} - -func TestTimeTextCountsTheCallsAtOnce(t *testing.T) { - start := time.Date(2026, 1, 1, 12, 0, 0, 0, time.UTC) - record := func(offset, seconds float64) ModelUsageRecord { - return ModelUsageRecord{ - Date: start.Add(time.Duration(offset * float64(time.Second))).Format(time.RFC3339Nano), - Seconds: seconds, - } - } - - text := probeTimeText([]ModelUsageRecord{record(0, 10), record(2, 4), record(3, 2), record(10, 10)}) - 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) -} - -// 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(probePeople) + "] 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(probeThings), g.pick(probeProblems)), - fmt.Sprintf("Move the %s to %s before %s.", g.pick(probeThings), g.pick(probePlaces), g.pick(probeDays)), - fmt.Sprintf("Thanks, that works. Next, %s.", g.pick(probeTasks)), - fmt.Sprintf("No, keep the old %s. %s asked for it.", g.pick(probeThings), g.pick(probePeople)), - fmt.Sprintf("How many %s orders came in on %s?", g.pick(probeBreads), g.pick(probeDays)), - }) -} - -func (g *ProbeWriter) replyLine() string { - return g.pick([]string{ - fmt.Sprintf("I looked at %s. The %s %s because %s.", g.pick(probeFiles), g.pick(probeThings), g.pick(probeProblems), g.pick(probeCauses)), - fmt.Sprintf("I changed %s, and the tests pass now.", g.pick(probeChanges)), - fmt.Sprintf("There were %d %s orders on %s, %d of them late.", 10+g.random.IntN(90), g.pick(probeBreads), g.pick(probeDays), g.random.IntN(10)), - fmt.Sprintf("Open question: should %s get %s?", g.pick(probePeople), g.pick(probeTasks)), - }) -} - -func (g *ProbeWriter) noteLine() string { - return g.pick([]string{ - fmt.Sprintf("%s prefers %s.", g.pick(probePeople), g.pick(probePreferences)), - 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(probeDays)), - }) -} - -func (g *ProbeWriter) pick(words []string) string { - return words[g.random.IntN(len(words))] -} - -// 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(probeThings) + "' " + g.pick(probeFiles)} - case "read": - input = map[string]string{"path": g.pick(probeFiles)} - default: - input = map[string]string{ - "path": g.pick(probeFiles), - "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) 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(probeBreads), g.pick(probePeople)), - fmt.Sprintf("%s:%d: %s %s.", g.pick(probeFiles), 1+g.random.IntN(400), g.pick(probeThings), g.pick(probeProblems)), - fmt.Sprintf("ok bakery/%s %d.%03ds", g.pick(probeModules), g.random.IntN(3), g.random.IntN(1000)), - }) -} - -// 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 -} - -func newTestProbeRun(t *testing.T, server *ModelFake) *ProbeRun { - t.Helper() - - root := t.TempDir() - usagePath := filepath.Join(root, "usage.jsonl") - client := newModelClient(server.URL, ModelAPIKey(modelTestKey), "onechat-"+probeName, usagePath) - memory, err := openMemory(t.Context(), root, newTestClock().Now, newCompactor(client, probePrompt, defaultConfig().Compactor)) - require.NoError(t, err) - - return &ProbeRun{ - memory: memory, - client: client, - usagePath: usagePath, - } -} - -var ( - probePeople = []string{"Mara", "Joel", "Ines", "Tomas", "Priya", "Lena"} - probeBreads = []string{"rye", "sourdough", "baguette", "brioche", "spelt", "seed loaf"} - probeDays = []string{"Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday"} - probeThings = []string{"order form", "price list", "delivery route", "stock count", "invoice export", "oven schedule"} - probeProblems = []string{"shows the wrong total", "skips the last row", "is slow at noon", "loses the customer note", "rounds the tax down"} - probePlaces = []string{"the front page", "the weekly report", "the shared folder", "the admin menu"} - probeTasks = []string{"add a field for allergies", "send the invoices on Friday", "split the route in two", "print labels for the shelf"} - probeCauses = []string{"the date has no time zone", "the cache keeps old prices", "the loop stops one row early", "the form trims the note"} - probeChanges = []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"} - probePreferences = []string{"short answers", "a call before a change to prices", "invoices as PDF", "reports in the morning"} - probeFiles = []string{"orders/form.go", "prices/list.go", "routes/plan.go", "stock/count.go", "export/invoice.go", "ovens/schedule.go"} - probeModules = []string{"orders", "prices", "routes", "stock", "export", "ovens"} -)