From: Anderson A. Giacomolli <agiacomolli@gmail.com>
Date: Fri, 6 Nov 2026 18:56:54 +0000 (-0300)
Subject: Run the probe builds in parallel to a call limit
X-Git-Url: http://git.giacomolli.uk/onechat/commitdiff_plain/63278e5c97fbad815c37c02f17d562367f4460a7?ds=inline

Run the probe builds in parallel to a call limit

The probe waits for the builds only before a user message, as a turn
does, so up to eight builds run at once while the other messages come
in. During a wait it checks the stop rules each second, so a stop does
not wait for a long backlog. After a stop it waits for the calls in
flight, so no call that was sent is cancelled.

The client takes a call limit and refuses each call over it before it
sends anything, so the probe stops at 660 calls also with eight builds
at once. The report adds the calls per model line, the time per call,
and how many calls ran at once.

Issue: b14d537
---

diff --git a/probe.go b/probe.go
index acd4f6b..a1af30f 100644
--- a/probe.go
+++ b/probe.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"bufio"
+	"cmp"
 	"context"
 	"encoding/json"
 	"errors"
@@ -25,17 +26,21 @@ const (
 	// 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")
@@ -74,6 +79,7 @@ func probe(ctx context.Context, args []string, out io.Writer) error {
 
 	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
@@ -84,7 +90,13 @@ func probe(ctx context.Context, args []string, out io.Writer) error {
 		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 {
@@ -99,32 +111,91 @@ func probe(ctx context.Context, args []string, out io.Writer) error {
 	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.
@@ -321,7 +392,7 @@ func writeProbeReport(out io.Writer, appended int, stop string, records []usageR
 	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),
 	)
 
@@ -329,6 +400,10 @@ func writeProbeReport(out io.Writer, appended int, stop string, records []usageR
 		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]),
@@ -355,6 +430,55 @@ func sizeText(sizes []int) string {
 	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))
diff --git a/probe_test.go b/probe_test.go
index c7436dc..c1d08df 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"
 )
@@ -122,22 +124,17 @@ func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) {
 		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))
 
@@ -146,9 +143,109 @@ 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(), "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.
@@ -270,6 +367,22 @@ func (g *probeWriter) pick(words []string) string {
 	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"}
diff --git a/responses.go b/responses.go
index fddec73..3055e70 100644
--- a/responses.go
+++ b/responses.go
@@ -13,6 +13,7 @@ import (
 	"path/filepath"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
@@ -36,6 +37,9 @@ const (
 	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 {
@@ -63,6 +67,8 @@ type Client struct {
 	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
@@ -191,8 +197,23 @@ func newClient(url string, auth authorizer, cacheKey, usagePath string) *Client
 	}
 }
 
-// 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,
diff --git a/responses_test.go b/responses_test.go
index 6017ad7..9a1355a 100644
--- a/responses_test.go
+++ b/responses_test.go
@@ -197,6 +197,24 @@ func TestClientRecordsEachCall(t *testing.T) {
 	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