git.giacomolli.uk - onechat/blob - build_run_test.go

   1 package main
   2
   3 import (
   4         "net/http"
   5         "path/filepath"
   6         "strconv"
   7         "strings"
   8         "sync"
   9         "testing"
  10         "time"
  11
  12         "github.com/stretchr/testify/require"
  13 )
  14
  15 func TestBuildRunStopsOnFailedCallsOrQuotaGrowth(t *testing.T) {
  16         quota := func(primary, secondary string) ModelUsageRecord {
  17                 return ModelUsageRecord{
  18                         Quota: map[string]string{
  19                                 "x-codex-primary-used-percent":   primary,
  20                                 "x-codex-secondary-used-percent": secondary,
  21                                 "x-codex-plan-type":              "plus",
  22                         },
  23                 }
  24         }
  25
  26         testCases := []struct {
  27                 name    string
  28                 records []ModelUsageRecord
  29                 want    string
  30         }{
  31                 {
  32                         name:    "no calls",
  33                         records: nil,
  34                         want:    "",
  35                 },
  36                 {
  37                         name:    "growth within the limit",
  38                         records: []ModelUsageRecord{quota("10", "40"), quota("25", "41")},
  39                         want:    "",
  40                 },
  41                 {
  42                         name:    "growth over the limit",
  43                         records: []ModelUsageRecord{quota("10", "40"), quota("12", "55.5")},
  44                         want:    "x-codex-secondary-used-percent grew from 40 to 55.5",
  45                 },
  46                 {
  47                         name: "failed call",
  48                         records: []ModelUsageRecord{quota("10", "40"), {
  49                                 Status: 400,
  50                                 Error:  "model call: HTTP 400",
  51                         }},
  52                         want: "call failed: model call: HTTP 400",
  53                 },
  54         }
  55
  56         for _, tc := range testCases {
  57                 t.Run(tc.name, func(t *testing.T) {
  58                         require.Equal(t, tc.want, buildStop(tc.records))
  59                 })
  60         }
  61 }
  62
  63 func TestBuildRunStopsWhenTheQuotaGrows(t *testing.T) {
  64         server := newModelFake(t)
  65         server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
  66                 w.Header().Set("X-Codex-Primary-Used-Percent", strconv.Itoa(10+call.number))
  67                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem(strings.Repeat("s", 400))}, 1000, 800, 50))
  68         }
  69
  70         run := newTestBuildRun(t, server)
  71         appended, stop := run.Run(t.Context(), probeLog(probeMessages))
  72         sizes, err := run.memory.buildModelLines()
  73         require.NoError(t, err)
  74         require.NoError(t, run.memory.Close())
  75
  76         require.Regexp(t, `^x-codex-primary-used-percent grew from \d+ to \d+$`, stop)
  77         require.Less(t, appended, probeMessages)
  78         require.NotEmpty(t, sizes)
  79
  80         records, err := buildReadUsage(run.usagePath)
  81         require.NoError(t, err)
  82         require.Greater(t, len(records), int(buildQuotaLimit))
  83
  84         var report strings.Builder
  85
  86         require.NoError(t, buildWriteReport(&report, appended, stop, records, sizes))
  87         require.Contains(t, report.String(), "stopped early: ")
  88         require.Contains(t, report.String(), "cached (80.0%)")
  89         require.Contains(t, report.String(), "line bytes: min 400, median 400, mean 400, max 400, 0 over 512")
  90         require.Contains(t, report.String(), "calls at once: at most ")
  91         require.Contains(t, report.String(), "quota at the first call: primary-used-percent=")
  92 }
  93
  94 func TestBuildRunStopsAtTheCallLimit(t *testing.T) {
  95         server := newModelFake(t)
  96         run := newTestBuildRun(t, server)
  97         run.client.LimitCalls(5)
  98
  99         entries := make([]DatedMessage, 40)
 100         for n := range entries {
 101                 entries[n] = DatedMessage{
 102                         Kind: "echo",
 103                         Text: strings.Repeat("long output ", 100),
 104                 }
 105         }
 106
 107         _, stop := run.Run(t.Context(), entries)
 108         require.NoError(t, run.memory.Close())
 109         require.Equal(t, "call limit of 5 reached", stop)
 110         require.Equal(t, 5, server.Calls())
 111 }
 112
 113 func TestBuildRunWaitsOnlyBeforeUserMessages(t *testing.T) {
 114         gate := make(chan struct{})
 115
 116         var mutex sync.Mutex
 117
 118         running, most := 0, 0
 119         server := newModelFake(t)
 120         server.reply = func(w http.ResponseWriter, _ ModelFakeCall) {
 121                 mutex.Lock()
 122
 123                 running++
 124                 most = max(most, running)
 125                 mutex.Unlock()
 126
 127                 <-gate
 128
 129                 mutex.Lock()
 130
 131                 running--
 132                 mutex.Unlock()
 133
 134                 modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("a line")}, 100, 0, 5))
 135         }
 136
 137         run := newTestBuildRun(t, server)
 138         long := strings.Repeat("long output ", 100)
 139         entries := []DatedMessage{
 140                 {
 141                         Kind: "echo",
 142                         Text: long,
 143                 },
 144                 {
 145                         Kind: "tool",
 146                         Text: long,
 147                 },
 148                 {
 149                         Kind: "user",
 150                         Text: "next",
 151                 },
 152                 {
 153                         Kind: "reply",
 154                         Text: long,
 155                 },
 156         }
 157
 158         done := make(chan string)
 159         go func() {
 160                 _, stop := run.Run(t.Context(), entries)
 161                 done <- stop
 162         }()
 163
 164         require.Eventually(t, func() bool {
 165                 return server.Calls() == 2
 166         }, 10*time.Second, time.Millisecond)
 167
 168         appended := run.memory.Len()
 169         require.Equal(t, 2, appended)
 170
 171         close(gate)
 172         require.Empty(t, <-done)
 173         require.NoError(t, run.memory.Close())
 174
 175         mutex.Lock()
 176         defer mutex.Unlock()
 177
 178         require.Equal(t, 2, most)
 179         require.Equal(t, 3, server.Calls())
 180 }
 181
 182 func TestReadUsageFileSortsByCallStart(t *testing.T) {
 183         path := filepath.Join(t.TempDir(), "usage.jsonl")
 184         usage := ModelUsageLog{
 185                 path: path,
 186         }
 187
 188         for _, date := range []string{"2026-01-01T12:00:02Z", "2026-01-01T12:00:00.5Z", "2026-01-01T12:00:01.25Z"} {
 189                 require.NoError(t, usage.Write(ModelUsageRecord{
 190                         Date: date,
 191                 }))
 192         }
 193
 194         records, err := buildReadUsage(path)
 195         require.NoError(t, err)
 196         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})
 197 }
 198
 199 func TestTimeTextCountsTheCallsAtOnce(t *testing.T) {
 200         start := time.Date(2026, 1, 1, 12, 0, 0, 0, time.UTC)
 201         record := func(offset, seconds float64) ModelUsageRecord {
 202                 return ModelUsageRecord{
 203                         Date:    start.Add(time.Duration(offset * float64(time.Second))).Format(time.RFC3339Nano),
 204                         Seconds: seconds,
 205                 }
 206         }
 207
 208         text := buildTimeText([]ModelUsageRecord{record(0, 10), record(2, 4), record(3, 2), record(10, 10)})
 209         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)
 210 }
 211
 212 func newTestBuildRun(t *testing.T, server *ModelFake) *BuildRun {
 213         t.Helper()
 214
 215         root := t.TempDir()
 216         usagePath := filepath.Join(root, "usage.jsonl")
 217         client := newModelClient(server.URL, ModelAPIKey(modelTestKey), "onechat-"+probeName, usagePath)
 218         memory, err := openMemory(t.Context(), root, newTestClock().Now, newCompactor(client, probePrompt, defaultConfig().Compactor, promptTools(), &ImageStore{dir: t.TempDir()}), discardLog())
 219         require.NoError(t, err)
 220
 221         return &BuildRun{
 222                 memory:    memory,
 223                 client:    client,
 224                 usagePath: usagePath,
 225         }
 226 }