"path/filepath"
"strconv"
"strings"
+ "sync"
"testing"
+ "time"
"github.com/stretchr/testify/require"
)
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)
}
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,
testCases := []struct {
name string
- records []usageRecord
+ records []ModelUsageRecord
want string
}{
{
},
{
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",
}},
}
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 := probeModelLines(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 := probeReadUsage(run.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.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(), "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=")
}
-// 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)),
+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),
+ }
}
- messages := make([]probeEntry, count)
- for id := range messages {
- kind, text := g.message()
- messages[id] = probeEntry{
- Kind: kind,
- Text: text,
+ _, 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 := ProbeSlotSummarizer{
+ 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 := 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,
}
}
- return messages
+ 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 {
+// ProbeWriter writes readable made-up messages.
+type ProbeWriter struct {
random *rand.Rand
}
-func (g *probeWriter) message() (string, string) {
+func (g *ProbeWriter) message() (string, string) {
roll := g.random.IntN(100)
switch {
case roll < 5:
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))
+ 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:
}
// fill joins lines of write until the text has about size bytes.
-func (g *probeWriter) fill(write func() string, size int) string {
+func (g *ProbeWriter) fill(write func() string, size int) string {
var text strings.Builder
for text.Len() < size {
if text.Len() > 0 {
return text.String()
}
-func (g *probeWriter) userLine() 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)),
+ 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 {
+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)),
+ 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 {
+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)),
+ 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) 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)),
- })
+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 {
+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)}
+ input = map[string]string{"command": "grep -n '" + g.pick(probeThings) + "' " + g.pick(probeFiles)}
case "read":
- input = map[string]string{"path": g.pick(files)}
+ input = map[string]string{"path": g.pick(probeFiles)}
default:
input = map[string]string{
- "path": g.pick(files),
+ "path": g.pick(probeFiles),
"text": g.fill(g.replyLine, max(0, size-80)),
}
}
return name + " " + string(data)
}
-func (g *probeWriter) pick(words []string) string {
- return words[g.random.IntN(len(words))]
+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 replayMessage. 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 (
- 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"}
+ 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"}
)
projects / onechat / blobdiff