Back up the daemon and restore its data

[onechat] / probe_test.go

diff --git a/probe_test.go b/probe_test.go
index c1d08dff6c3992060c9a07d76dcc44004a266783..298e05733a0295d6309349e37fdf219281ed83ac 100644 (file)
--- a/probe_test.go
+++ b/probe_test.go
@@ -2,12 +2,8 @@ package main
 
 import (
        "encoding/json"
 
 import (
        "encoding/json"

-       "fmt"
-       "math/rand/v2"
-       "net/http"

        "os"
        "path/filepath"
        "os"
        "path/filepath"

-       "strconv"

        "strings"
        "sync"
        "testing"
        "strings"
        "sync"
        "testing"


@@ -16,9 +12,6 @@ import (
        "github.com/stretchr/testify/require"
 )
 
        "github.com/stretchr/testify/require"
 )
 

-// probeMessages is the length of the generated log for the live probe.
-const probeMessages = 1000
-

 // TestProbeLogHasTheSyntheticMix checks that the generated log has readable
 // ASCII text with the mix of kinds and sizes of the synthetic log. When
 // ONECHAT_PROBE_LOG names a file, the test also writes the log there, one
 // TestProbeLogHasTheSyntheticMix checks that the generated log has readable
 // ASCII text with the mix of kinds and sizes of the synthetic log. When
 // ONECHAT_PROBE_LOG names a file, the test also writes the log there, one


@@ -31,8 +24,9 @@ func TestProbeLogHasTheSyntheticMix(t *testing.T) {
        long := 0
        for _, message := range messages {
                counts[message.Kind]++
        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)
                require.NotEmpty(t, message.Text)

+

                for _, char := range message.Text {
                        require.True(t, char == '\n' || (char >= ' ' && char <= '~'), "character %q", char)
                }
                for _, char := range message.Text {
                        require.True(t, char == '\n' || (char >= ' ' && char <= '~'), "character %q", char)
                }


@@ -58,6 +52,7 @@ func TestProbeLogHasTheSyntheticMix(t *testing.T) {
        }
 
        var data strings.Builder
        }
 
        var data strings.Builder

+

        for _, message := range messages {
                line, err := json.Marshal(message)
                require.NoError(t, err)
        for _, message := range messages {
                line, err := json.Marshal(message)
                require.NoError(t, err)


@@ -69,331 +64,42 @@ func TestProbeLogHasTheSyntheticMix(t *testing.T) {
        require.NoError(t, os.WriteFile(path, []byte(data.String()), 0o600))
 }
 
        require.NoError(t, os.WriteFile(path, []byte(data.String()), 0o600))
 }
 

-func TestProbeStopRules(t *testing.T) {
-       quota := func(primary, secondary string) usageRecord {
-               return usageRecord{
-                       Quota: map[string]string{
-                               "x-codex-primary-used-percent":   primary,
-                               "x-codex-secondary-used-percent": secondary,
-                               "x-codex-plan-type":              "plus",
-                       },
-               }
-       }
-
-       testCases := []struct {
-               name    string
-               records []usageRecord
-               want    string
-       }{
-               {
-                       name:    "no calls",
-                       records: nil,
-                       want:    "",
-               },
-               {
-                       name:    "growth within the limit",
-                       records: []usageRecord{quota("10", "40"), quota("25", "41")},
-                       want:    "",
-               },
-               {
-                       name:    "growth over the limit",
-                       records: []usageRecord{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"), {
-                               Status: 400,
-                               Error:  "model call: HTTP 400",
-                       }},
-                       want: "call failed: model call: HTTP 400",
-               },
-       }
-
-       for _, tc := range testCases {
-               t.Run(tc.name, func(t *testing.T) {
-                       require.Equal(t, tc.want, probeStop(tc.records))
-               })
-       }
-}
-
-func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) {
-       server := newFakeResponses(t)
-       server.reply = func(w http.ResponseWriter, call fakeCall) {
-               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))
-       }
-
-       run := newTestProbeRun(t, server)
-       appended, stop := run.Run(t.Context(), probeLog(probeMessages))
-       sizes, err := modelLines(run.memory)
-       require.NoError(t, err)
-       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)
+func TestProbeReadIgnoresSourceDates(t *testing.T) {
+       path := filepath.Join(t.TempDir(), "probe.jsonl")
+       require.NoError(t, os.WriteFile(path, []byte(`{"kind":"user","text":"A probe message","date":"2025-01-01T00:00:00Z"}`+"\n"), 0o600))

 
 

-       records, err := readUsageFile(run.usagePath)
+       entries, err := probeReadLog(path, 1)

        require.NoError(t, err)
        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.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")
+       require.Len(t, entries, 1)
+       require.Empty(t, entries[0].Date)

 }
 
 }
 

-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),
-               }
+func TestSlotSummarizerLimitsTheBuildsAtOnce(t *testing.T) {
+       inner := newFakeSummarizer()
+       inner.gate = make(chan struct{})
+       summarizer := ProbeSlotSummarizer{
+               inner: inner,
+               slots: make(chan struct{}, 2),

        }
 
        }
 

-       _, 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())
-}
+       var group sync.WaitGroup

 
 

-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,
-               },
+       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)
+               })

        }
 
        }
 

-       done := make(chan string)
-       go func() {
-               _, stop := run.Run(t.Context(), entries)
-               done <- stop
-       }()
-

        require.Eventually(t, func() bool {
        require.Eventually(t, func() bool {

-               return server.Calls() == 2
+               return inner.Active() == 2

        }, 10*time.Second, time.Millisecond)
 
        }, 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.
-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
-}
-
-// probeWriter writes readable made-up messages.
-type probeWriter struct {
-       random *rand.Rand
-}
-
-func (g *probeWriter) message() (string, string) {
-       roll := g.random.IntN(100)
-       switch {
-       case roll < 5:
-               return "user", g.fill(g.userLine, 10+g.random.IntN(200))
-       case roll < 10:
-               return "reply", g.fill(g.replyLine, 100+g.random.IntN(1500))
-       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))
-       case roll < 56:
-               return "tool", g.toolCall(60 + g.random.IntN(400))
-       case roll < 57:
-               return "echo", g.fill(g.outputLine, 30000+g.random.IntN(30000))
-       default:
-               return "echo", g.fill(g.outputLine, 20+g.random.IntN(4000))
-       }
-}
-
-// fill joins lines of write until the text has about size bytes.
-func (g *probeWriter) fill(write func() string, size int) string {
-       var text strings.Builder
-       for text.Len() < size {
-               if text.Len() > 0 {
-                       text.WriteByte(' ')
-               }
-
-               text.WriteString(write())
-       }
-
-       return text.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)),
-       })
-}
-
-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)),
-       })
-}
-
-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)),
-       })
+       close(inner.gate)
+       group.Wait()
+       require.Equal(t, 2, inner.peak)

 }
 }

-
-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)),
-       })
-}
-
-// toolCall returns a tool name and its JSON input of about size bytes.
-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)}
-       case "read":
-               input = map[string]string{"path": g.pick(files)}
-       default:
-               input = map[string]string{
-                       "path": g.pick(files),
-                       "text": g.fill(g.replyLine, max(0, size-80)),
-               }
-       }
-
-       data, err := json.Marshal(input)
-       if err != nil {
-               panic(err)
-       }
-
-       return name + " " + string(data)
-}
-
-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"}
-       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"}
-)