Name and order the model client code

[onechat] / probe_test.go

diff --git a/probe_test.go b/probe_test.go
index c1d08dff6c3992060c9a07d76dcc44004a266783..afc6e2f8abb00c7648b09308e2e01585054d60f0 100644 (file)
--- a/probe_test.go
+++ b/probe_test.go
@@ -70,8 +70,8 @@ func TestProbeLogHasTheSyntheticMix(t *testing.T) {
 }
 
 func TestProbeStopRules(t *testing.T) {
 }
 
 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,
                        Quota: map[string]string{
                                "x-codex-primary-used-percent":   primary,
                                "x-codex-secondary-used-percent": secondary,


@@ -82,7 +82,7 @@ func TestProbeStopRules(t *testing.T) {
 
        testCases := []struct {
                name    string
 
        testCases := []struct {
                name    string

-               records []usageRecord
+               records []ModelUsageRecord

                want    string
        }{
                {
                want    string
        }{
                {


@@ -92,17 +92,17 @@ func TestProbeStopRules(t *testing.T) {
                },
                {
                        name:    "growth within the limit",
                },
                {
                        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",
                        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",
                        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",
                        }},
                                Status: 400,
                                Error:  "model call: HTTP 400",
                        }},


@@ -118,10 +118,10 @@ func TestProbeStopRules(t *testing.T) {
 }
 
 func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) {
 }
 
 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))
                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))

        }
 
        run := newTestProbeRun(t, server)
        }
 
        run := newTestProbeRun(t, server)


@@ -130,7 +130,7 @@ func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) {
        require.NoError(t, err)
        require.NoError(t, run.memory.Close())
 
        require.NoError(t, err)
        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)
 
        require.Less(t, appended, probeMessages)
        require.NotEmpty(t, sizes)
 


@@ -144,11 +144,11 @@ func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) {
        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(), "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")
+       require.Contains(t, report.String(), "quota at the first call: primary-used-percent=")

 }
 
 func TestProbeRunStopsAtTheCallLimit(t *testing.T) {
 }
 
 func TestProbeRunStopsAtTheCallLimit(t *testing.T) {

-       server := newFakeResponses(t)
+       server := newModelFake(t)

        run := newTestProbeRun(t, server)
        run.client.LimitCalls(5)
 
        run := newTestProbeRun(t, server)
        run.client.LimitCalls(5)
 


@@ -170,8 +170,8 @@ func TestProbeRunWaitsOnlyBeforeUserMessages(t *testing.T) {
        gate := make(chan struct{})
        var mutex sync.Mutex
        running, most := 0, 0
        gate := make(chan struct{})
        var mutex sync.Mutex
        running, most := 0, 0

-       server := newFakeResponses(t)
-       server.reply = func(w http.ResponseWriter, call fakeCall) {
+       server := newModelFake(t)
+       server.reply = func(w http.ResponseWriter, call ModelFakeCall) {

                mutex.Lock()
                running++
                most = max(most, running)
                mutex.Lock()
                running++
                most = max(most, running)


@@ -183,7 +183,7 @@ func TestProbeRunWaitsOnlyBeforeUserMessages(t *testing.T) {
                running--
                mutex.Unlock()
 
                running--
                mutex.Unlock()
 

-               writeEvents(w, completedEvent([]any{messageItem("a line")}, 100, 0, 5))
+               modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("a line")}, 100, 0, 5))

        }
 
        run := newTestProbeRun(t, server)
        }
 
        run := newTestProbeRun(t, server)


@@ -233,17 +233,63 @@ func TestProbeRunWaitsOnlyBeforeUserMessages(t *testing.T) {
        require.Equal(t, 3, server.Calls())
 }
 
        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 := 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 := 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)
 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{
+       record := func(offset, seconds float64) ModelUsageRecord {
+               return ModelUsageRecord{

                        Date:    start.Add(time.Duration(offset * float64(time.Second))).Format(time.RFC3339Nano),
                        Seconds: seconds,
                }
        }
 
                        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)
+       text := timeText([]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)

 }
 
 // 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


@@ -367,13 +413,13 @@ 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 {
+func newTestProbeRun(t *testing.T, server *ModelFake) *probeRun {

        t.Helper()
 
        root := t.TempDir()
        usagePath := filepath.Join(root, "usage.jsonl")
        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))
+       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{
        require.NoError(t, err)
 
        return &probeRun{