import (
"bufio"
+ "cmp"
"context"
"encoding/json"
"errors"
// can grow in a probe run before the run stops.
probeQuotaLimit = 15.0
+ // probeCallLimit is the most model calls of a probe run.
+ probeCallLimit = 660
+
// probeLineBytes is the most bytes of one line of a probe log.
probeLineBytes = 1 << 20
)
// probe runs "onechat probe LOG COUNT". It appends the first COUNT messages of
// LOG, one JSON object {kind, text} per line, to a new chat in probe/ in the
-// data directory, and builds its tree with the real compactor. After each
-// message it waits for the builds, as a turn does. It stops at the first
-// failed call, or when the used percent of a quota window grows by more than
-// probeQuotaLimit points. At the end it reports the calls, the cached share,
-// the quota, the retries, and the line sizes.
+// data directory, and builds its tree with the real compactor. Before each
+// user message it waits for the builds, as a turn does. It stops at the first
+// failed call, after probeCallLimit calls, or when the used percent of a quota
+// window grows by more than probeQuotaLimit points. At the end it reports the
+// calls, the time per call, the calls at once, the cached share, the quota,
+// the retries, and the line sizes.
func probe(ctx context.Context, args []string, out io.Writer) error {
if len(args) != 2 {
return errors.New("usage: onechat probe LOG COUNT")
usagePath := filepath.Join(root, "usage.jsonl")
client := newClient(address, auth, "onechat-"+probeName, usagePath)
+ client.LimitCalls(probeCallLimit)
memory, err := openMemory(ctx, root, time.Now, newCompactor(client, probeName, config.Compactor))
if err != nil {
return err
return errors.Join(err, memory.Close())
}
- appended, stop := runProbe(ctx, memory, entries, usagePath)
+ run := probeRun{
+ memory: memory,
+ client: client,
+ usagePath: usagePath,
+ }
+
+ appended, stop := run.Run(ctx, entries)
sizes, err := modelLines(memory)
err = errors.Join(err, memory.Close())
if err != nil {
return writeProbeReport(out, appended, stop, records, sizes)
}
-// runProbe appends the entries and waits for the builds after each one. It
-// returns the number of messages appended and why the run stopped early, or
-// "" when it did not.
-func runProbe(ctx context.Context, memory *Memory, entries []probeEntry, usagePath string) (int, string) {
+// probeRun appends the messages of a probe run to its chat and stops the run
+// by the rules of probeStop.
+type probeRun struct {
+ memory *Memory
+ client *Client
+ usagePath string
+}
+
+// Run appends the entries. It waits for the builds only before a user message,
+// as a turn does, so up to maxBuilds builds run at once in between, and it
+// waits for all builds at the end. It returns the number of messages appended
+// and why the run stopped early, or "" when it did not. After a stop it waits
+// for the calls in flight, so no sent call is cancelled.
+func (r *probeRun) Run(ctx context.Context, entries []probeEntry) (int, string) {
+ appended, reason := r.append(ctx, entries)
+ if reason == "" {
+ return appended, ""
+ }
+
+ err := r.memory.Wait(ctx)
+ if err != nil {
+ reason += "; " + err.Error()
+ }
+
+ return appended, reason
+}
+
+// append appends the entries and stops at the first stop rule.
+func (r *probeRun) append(ctx context.Context, entries []probeEntry) (int, string) {
for n, entry := range entries {
- _, err := memory.Append(entry.Kind, entry.Text)
- if err == nil {
- err = memory.Wait(ctx)
+ if entry.Kind == "user" {
+ reason := r.wait(ctx)
+ if reason != "" {
+ return n, reason
+ }
}
+ _, err := r.memory.Append(entry.Kind, entry.Text)
if err != nil {
return n, err.Error()
}
- records, err := readUsageFile(usagePath)
- if err != nil {
- return n + 1, err.Error()
+ reason := r.check()
+ if reason != "" {
+ return n + 1, reason
}
+ }
- reason := probeStop(records)
+ return len(entries), r.wait(ctx)
+}
+
+// wait returns when no build runs, or when a stop rule ends the run. It checks
+// the rules each second, so a run stops soon also during a long wait.
+func (r *probeRun) wait(ctx context.Context) string {
+ for {
+ reason := r.check()
if reason != "" {
- return n + 1, reason
+ return reason
+ }
+
+ waitCtx, cancel := context.WithTimeout(ctx, time.Second)
+ err := r.memory.Wait(waitCtx)
+ cancel()
+ if err == nil {
+ return r.check()
+ }
+
+ if ctx.Err() != nil {
+ return ctx.Err().Error()
}
}
+}
- return len(entries), ""
+// check returns why the run must stop, or "" when it can go on.
+func (r *probeRun) check() string {
+ if r.client.LimitReached() {
+ return fmt.Sprintf("call limit of %d reached", r.client.limit)
+ }
+
+ records, err := readUsageFile(r.usagePath)
+ if err != nil {
+ return err.Error()
+ }
+
+ return probeStop(records)
}
// probeEntry is one message of a probe log.
lines = append(lines,
fmt.Sprintf("messages: %d", appended),
fmt.Sprintf("calls: %d (%d failed), %.2f per message", len(records), failed, ratio(len(records), appended)),
- fmt.Sprintf("model lines: %d, retries: %d", len(sizes), len(records)-failed-len(sizes)),
+ fmt.Sprintf("model lines: %d, retries: %d, %.2f calls per model line", len(sizes), len(records)-failed-len(sizes), ratio(len(records), len(sizes))),
fmt.Sprintf("tokens: %d input, %d cached (%.1f%%), %d output", input, cached, 100*ratio(cached, input), output),
)
lines = append(lines, sizeText(sizes))
}
+ if len(records) > 0 {
+ lines = append(lines, timeText(records))
+ }
+
if len(records) > 0 {
lines = append(lines,
"quota at the first call: "+quotaText(records[0]),
return fmt.Sprintf("line bytes: min %d, median %d, mean %.0f, max %d, %d over %d", sorted[0], sorted[len(sorted)/2], ratio(total, len(sorted)), sorted[len(sorted)-1], over, nodeBytes)
}
+// timeText describes the time per call and how many calls ran at once. A call
+// runs from its date for its seconds.
+func timeText(records []usageRecord) string {
+ type event struct {
+ at time.Time
+ delta int
+ }
+
+ var events []event
+ seconds := make([]float64, 0, len(records))
+ total := 0.0
+ for _, record := range records {
+ start, err := time.Parse(time.RFC3339Nano, record.Date)
+ if err != nil {
+ continue
+ }
+
+ length := time.Duration(record.Seconds * float64(time.Second))
+ events = append(events, event{
+ at: start,
+ delta: 1,
+ }, event{
+ at: start.Add(length),
+ delta: -1,
+ })
+ seconds = append(seconds, record.Seconds)
+ total += record.Seconds
+ }
+
+ if len(events) == 0 {
+ return "time: no dated calls"
+ }
+
+ slices.SortFunc(events, func(a, b event) int {
+ return cmp.Or(a.at.Compare(b.at), cmp.Compare(a.delta, b.delta))
+ })
+
+ running, most := 0, 0
+ for _, e := range events {
+ running += e.delta
+ most = max(most, running)
+ }
+
+ wall := events[len(events)-1].at.Sub(events[0].at).Seconds()
+ slices.Sort(seconds)
+
+ return fmt.Sprintf("time: %.0f s in all, %.1f s per call (median %.1f, max %.1f), calls at once: at most %d, %.1f on average", wall, total/float64(len(seconds)), seconds[len(seconds)/2], seconds[len(seconds)-1], most, total/max(wall, 1e-9))
+}
+
// quotaText returns the quota headers of record, sorted by name.
func quotaText(record usageRecord) string {
names := slices.Sorted(maps.Keys(record.Quota))
"path/filepath"
"strconv"
"strings"
+ "sync"
"testing"
+ "time"
"github.com/stretchr/testify/require"
)
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)
-
- appended, stop := runProbe(t.Context(), memory, probeLog(probeMessages), usagePath)
- sizes, err := modelLines(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.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.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=10")
}
+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 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, 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
// system, with the mix of kinds and sizes of syntheticMessage. The same count
// always gives the same log.
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"}
"path/filepath"
"strings"
"sync"
+ "sync/atomic"
"time"
)
errorBytes = 4096
)
+// errCallLimit tells that a client sent its limit of calls.
+var errCallLimit = errors.New("call limit reached")
+
// authorizer adds the credentials of a model call to the headers of its
// request.
type authorizer interface {
cacheKey string
usage *usageLog
http *http.Client
+ limit int64
+ calls atomic.Int64
}
// Request is one model call. Input holds the input items in order, such as
}
}
-// Create sends the call and reads its event stream to the end.
+// LimitCalls makes the client refuse each call after the first limit calls.
+func (c *Client) LimitCalls(limit int64) {
+ c.limit = limit
+}
+
+// LimitReached tells if the client refuses further calls.
+func (c *Client) LimitReached() bool {
+ return c.limit > 0 && c.calls.Load() >= c.limit
+}
+
+// Create sends the call and reads its event stream to the end. Over the call
+// limit it sends nothing and returns errCallLimit.
func (c *Client) Create(ctx context.Context, request Request) (Response, error) {
+ if c.calls.Add(1) > c.limit && c.limit > 0 {
+ return Response{}, errCallLimit
+ }
+
body, err := json.Marshal(requestBody{
Model: request.Model,
Instructions: request.Instructions,
require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
}
+func TestClientRefusesCallsOverTheLimit(t *testing.T) {
+ server := newFakeResponses(t)
+ client := newTestClient(t, server)
+ client.LimitCalls(2)
+
+ for range 2 {
+ require.False(t, client.LimitReached())
+ _, err := client.Create(t.Context(), testRequest())
+ require.NoError(t, err)
+ }
+
+ require.True(t, client.LimitReached())
+ _, err := client.Create(t.Context(), testRequest())
+ require.ErrorIs(t, err, errCallLimit)
+ require.Equal(t, 2, server.Calls())
+ require.Len(t, readUsage(t, client.usage.path), 2)
+}
+
func TestClientFailsWhenTheResponseFails(t *testing.T) {
testCases := []struct {
name string
projects / onechat / commitdiff