X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/522d226230a2eae0a8a33ae0fbf32f96c6886b26..503486b8364f07ef44c06bcc1be566f62600b0eb:/probe_test.go diff --git a/probe_test.go b/probe_test.go index c7436dc..afc6e2f 100644 --- a/probe_test.go +++ b/probe_test.go @@ -9,7 +9,9 @@ import ( "path/filepath" "strconv" "strings" + "sync" "testing" + "time" "github.com/stretchr/testify/require" ) @@ -68,8 +70,8 @@ func TestProbeLogHasTheSyntheticMix(t *testing.T) { } func TestProbeStopRules(t *testing.T) { - quota := func(primary, secondary string) usageRecord { - return usageRecord{ + quota := func(primary, secondary string) ModelUsageRecord { + return ModelUsageRecord{ Quota: map[string]string{ "x-codex-primary-used-percent": primary, "x-codex-secondary-used-percent": secondary, @@ -80,7 +82,7 @@ func TestProbeStopRules(t *testing.T) { testCases := []struct { name string - records []usageRecord + records []ModelUsageRecord want string }{ { @@ -90,17 +92,17 @@ func TestProbeStopRules(t *testing.T) { }, { name: "growth within the limit", - records: []usageRecord{quota("10", "40"), quota("25", "41")}, + records: []ModelUsageRecord{quota("10", "40"), quota("25", "41")}, want: "", }, { name: "growth over the limit", - records: []usageRecord{quota("10", "40"), quota("12", "55.5")}, + 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: []usageRecord{quota("10", "40"), { + records: []ModelUsageRecord{quota("10", "40"), { Status: 400, Error: "model call: HTTP 400", }}, @@ -116,28 +118,23 @@ func TestProbeStopRules(t *testing.T) { } func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) { - server := newFakeResponses(t) - server.reply = func(w http.ResponseWriter, call fakeCall) { + server := newModelFake(t) + server.reply = func(w http.ResponseWriter, call ModelFakeCall) { 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)) + modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem(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) + run := newTestProbeRun(t, server) + appended, stop := run.Run(t.Context(), probeLog(probeMessages)) + sizes, err := modelLines(run.memory) require.NoError(t, err) - require.NoError(t, memory.Close()) + require.NoError(t, run.memory.Close()) - require.Contains(t, stop, "x-codex-primary-used-percent grew from 10 to") + 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 := readUsageFile(usagePath) + records, err := readUsageFile(run.usagePath) require.NoError(t, err) require.Greater(t, len(records), int(probeQuotaLimit)) @@ -146,7 +143,153 @@ func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) { 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") + 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()) +} + +func TestSlotSummarizerLimitsTheBuildsAtOnce(t *testing.T) { + inner := newFakeSummarizer() + inner.gate = make(chan struct{}) + summarizer := slotSummarizer{ + inner: inner, + slots: make(chan struct{}, 2), + } + + var group sync.WaitGroup + for n := range 6 { + group.Go(func() { + _, err := summarizer.Summarize(t.Context(), Task{ + Level: 0, + Index: n, + Input: []string{"user: " + strings.Repeat("x", 600)}, + }) + require.NoError(t, err) + }) + } + + require.Eventually(t, func() bool { + return inner.Active() == 2 + }, 10*time.Second, time.Millisecond) + + close(inner.gate) + 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 := readUsageFile(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 := timeText([]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) } // probeLog returns count messages of made-up work on a small bakery order @@ -270,6 +413,22 @@ func (g *probeWriter) pick(words []string) string { return words[g.random.IntN(len(words))] } +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 ( people = []string{"Mara", "Joel", "Ines", "Tomas", "Priya", "Lena"} breads = []string{"rye", "sourdough", "baguette", "brioche", "spelt", "seed loaf"}