]> git.giacomolli.uk - onechat/blobdiff - probe_test.go
Make the chat writes durable by group commit
[onechat] / probe_test.go
index c7436dcea94b507fdabd59ec23abf31837ee7492..0651281368c3e194ab15711986384c226fbc484e 100644 (file)
@@ -9,7 +9,9 @@ import (
        "path/filepath"
        "strconv"
        "strings"
        "path/filepath"
        "strconv"
        "strings"
+       "sync"
        "testing"
        "testing"
+       "time"
 
        "github.com/stretchr/testify/require"
 )
 
        "github.com/stretchr/testify/require"
 )
@@ -122,22 +124,17 @@ func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) {
                writeEvents(w, completedEvent([]any{messageItem(strings.Repeat("s", 400))}, 1000, 800, 50))
        }
 
                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))
+       run := newTestProbeRun(t, server)
+       appended, stop := run.Run(t.Context(), probeLog(probeMessages))
+       sizes, err := modelLines(run.memory)
        require.NoError(t, err)
        require.NoError(t, err)
+       require.NoError(t, run.memory.Close())
 
 
-       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.Regexp(t, `^x-codex-primary-used-percent grew from \d+ to \d+$`, stop)
        require.Less(t, appended, probeMessages)
        require.NotEmpty(t, sizes)
 
        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))
 
        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(), "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 := newFakeResponses(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 := newFakeResponses(t)
+       server.reply = func(w http.ResponseWriter, call fakeCall) {
+               mutex.Lock()
+               running++
+               most = max(most, running)
+               mutex.Unlock()
+
+               <-gate
+
+               mutex.Lock()
+               running--
+               mutex.Unlock()
+
+               writeEvents(w, completedEvent([]any{messageItem("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 := usageLog{
+               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(usageRecord{
+                       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) usageRecord {
+               return usageRecord{
+                       Date:    start.Add(time.Duration(offset * float64(time.Second))).Format(time.RFC3339Nano),
+                       Seconds: seconds,
+               }
+       }
+
+       text := timeText([]usageRecord{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
 }
 
 // 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))]
 }
 
        return words[g.random.IntN(len(words))]
 }
 
+func newTestProbeRun(t *testing.T, server *fakeResponses) *probeRun {
+       t.Helper()
+
+       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)
+
+       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"}
 var (
        people      = []string{"Mara", "Joel", "Ines", "Tomas", "Priya", "Lena"}
        breads      = []string{"rye", "sourdough", "baguette", "brioche", "spelt", "seed loaf"}