--- /dev/null
+package main
+
+import (
+ "bufio"
+ "context"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "io"
+ "maps"
+ "os"
+ "path/filepath"
+ "slices"
+ "strconv"
+ "strings"
+ "time"
+)
+
+const (
+ // probeName is the agent name of a probe run. The cache key comes from
+ // it.
+ probeName = "Probe"
+
+ // probeQuotaLimit is how many points the used percent of a quota window
+ // can grow in a probe run before the run stops.
+ probeQuotaLimit = 15.0
+
+ // 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.
+func probe(ctx context.Context, args []string, out io.Writer) error {
+ if len(args) != 2 {
+ return errors.New("usage: onechat probe LOG COUNT")
+ }
+
+ count, err := strconv.Atoi(args[1])
+ if err != nil || count <= 0 {
+ return fmt.Errorf("COUNT %q is not a positive number", args[1])
+ }
+
+ entries, err := readProbeLog(args[0], count)
+ if err != nil {
+ return err
+ }
+
+ dir, err := dataDir()
+ if err != nil {
+ return err
+ }
+
+ config, err := loadConfig(dir)
+ if err != nil {
+ return err
+ }
+
+ address, auth, err := modelEndpoint(config, dir)
+ if err != nil {
+ return err
+ }
+
+ root := filepath.Join(dir, "probe", time.Now().Format("20060102-150405"))
+ err = os.MkdirAll(root, 0o700)
+ if err != nil {
+ return fmt.Errorf("create probe directory: %w", err)
+ }
+
+ usagePath := filepath.Join(root, "usage.jsonl")
+ client := newClient(address, auth, "onechat-"+probeName, usagePath)
+ memory, err := openMemory(ctx, root, time.Now, newCompactor(client, probeName, config.Compactor))
+ if err != nil {
+ return err
+ }
+
+ _, err = fmt.Fprintf(out, "probe: %d messages into %s with %s at %s effort\n", len(entries), root, config.Compactor.Model, config.Compactor.Effort)
+ if err != nil {
+ return errors.Join(err, memory.Close())
+ }
+
+ appended, stop := runProbe(ctx, memory, entries, usagePath)
+ sizes, err := modelLines(memory)
+ err = errors.Join(err, memory.Close())
+ if err != nil {
+ return err
+ }
+
+ records, err := readUsageFile(usagePath)
+ if err != nil {
+ return err
+ }
+
+ 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) {
+ for n, entry := range entries {
+ _, err := memory.Append(entry.Kind, entry.Text)
+ if err == nil {
+ err = memory.Wait(ctx)
+ }
+
+ if err != nil {
+ return n, err.Error()
+ }
+
+ records, err := readUsageFile(usagePath)
+ if err != nil {
+ return n + 1, err.Error()
+ }
+
+ reason := probeStop(records)
+ if reason != "" {
+ return n + 1, reason
+ }
+ }
+
+ return len(entries), ""
+}
+
+// probeEntry is one message of a probe log.
+type probeEntry struct {
+ Kind string `json:"kind"`
+ Text string `json:"text"`
+}
+
+// readProbeLog reads the first count messages of the probe log at path.
+func readProbeLog(path string, count int) (entries []probeEntry, err error) {
+ file, err := os.Open(path)
+ if err != nil {
+ return nil, fmt.Errorf("open probe log: %w", err)
+ }
+
+ defer func() {
+ err = errors.Join(err, file.Close())
+ }()
+
+ scanner := bufio.NewScanner(file)
+ scanner.Buffer(make([]byte, 64*1024), probeLineBytes)
+ for len(entries) < count && scanner.Scan() {
+ var entry probeEntry
+ err = json.Unmarshal(scanner.Bytes(), &entry)
+ if err != nil {
+ return nil, fmt.Errorf("probe log line %d: %w", len(entries)+1, err)
+ }
+
+ entries = append(entries, entry)
+ }
+
+ err = scanner.Err()
+ if err != nil {
+ return nil, fmt.Errorf("read probe log: %w", err)
+ }
+
+ if len(entries) < count {
+ return nil, fmt.Errorf("probe log has %d messages, fewer than %d", len(entries), count)
+ }
+
+ return entries, nil
+}
+
+// probeStop returns why a probe run must stop: a failed call, or a quota window
+// whose used percent grew by more than probeQuotaLimit points since the first
+// call. It returns "" when the run can go on.
+func probeStop(records []usageRecord) string {
+ for _, record := range records {
+ if record.Error != "" {
+ return "call failed: " + record.Error
+ }
+ }
+
+ if len(records) == 0 {
+ return ""
+ }
+
+ first := usedPercents(records[0])
+ for name, used := range usedPercents(records[len(records)-1]) {
+ start, found := first[name]
+ if found && used-start > probeQuotaLimit {
+ return fmt.Sprintf("%s grew from %g to %g", name, start, used)
+ }
+ }
+
+ return ""
+}
+
+// modelLines returns the sizes of the nodes that the model wrote, in tree
+// order. A node that fits its source with no model call is not one of them.
+func modelLines(memory *Memory) ([]int, error) {
+ memory.mutex.Lock()
+ defer memory.mutex.Unlock()
+
+ keys, err := memory.tree.Keys()
+ if err != nil {
+ return nil, err
+ }
+
+ var sizes []int
+ for _, key := range keys {
+ node, _, err := memory.tree.Node(key.level, key.index)
+ if err != nil {
+ return nil, err
+ }
+
+ model, err := builtByModel(memory, key)
+ if err != nil {
+ return nil, err
+ }
+
+ if model {
+ sizes = append(sizes, node.Size)
+ }
+ }
+
+ return sizes, nil
+}
+
+// builtByModel tells if the node key needed a model call: a leaf whose source
+// is over nodeBytes, or a merge whose children do not fit in one node.
+func builtByModel(memory *Memory, key nodeKey) (bool, error) {
+ if key.level == 0 {
+ message, err := memory.log.Message(key.index)
+ if err != nil {
+ return false, err
+ }
+
+ _, free := freeLeaf(message)
+
+ return !free, nil
+ }
+
+ left, right, err := memory.children(key)
+ if err != nil {
+ return false, err
+ }
+
+ _, free := freeMerge(left, right)
+
+ return !free, nil
+}
+
+// readUsageFile reads the records of usage.jsonl at path.
+func readUsageFile(path string) (records []usageRecord, err error) {
+ file, err := os.Open(path)
+ if errors.Is(err, os.ErrNotExist) {
+ return nil, nil
+ }
+
+ if err != nil {
+ return nil, fmt.Errorf("open usage file: %w", err)
+ }
+
+ defer func() {
+ err = errors.Join(err, file.Close())
+ }()
+
+ scanner := bufio.NewScanner(file)
+ scanner.Buffer(make([]byte, 64*1024), probeLineBytes)
+ for scanner.Scan() {
+ var record usageRecord
+ err = json.Unmarshal(scanner.Bytes(), &record)
+ if err != nil {
+ return nil, fmt.Errorf("decode usage: %w", err)
+ }
+
+ records = append(records, record)
+ }
+
+ err = scanner.Err()
+ if err != nil {
+ return nil, fmt.Errorf("read usage file: %w", err)
+ }
+
+ return records, nil
+}
+
+// usedPercents returns the used percent of each quota window of record, by
+// header name.
+func usedPercents(record usageRecord) map[string]float64 {
+ used := make(map[string]float64)
+ for name, value := range record.Quota {
+ if !strings.HasSuffix(name, "-used-percent") {
+ continue
+ }
+
+ percent, err := strconv.ParseFloat(value, 64)
+ if err == nil {
+ used[name] = percent
+ }
+ }
+
+ return used
+}
+
+// writeProbeReport writes the result of a probe run.
+func writeProbeReport(out io.Writer, appended int, stop string, records []usageRecord, sizes []int) error {
+ var input, cached, output, failed int
+ for _, record := range records {
+ input += record.InputTokens
+ cached += record.CachedTokens
+ output += record.OutputTokens
+ if record.Error != "" {
+ failed++
+ }
+ }
+
+ var lines []string
+ if stop != "" {
+ lines = append(lines, "stopped early: "+stop)
+ }
+
+ 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("tokens: %d input, %d cached (%.1f%%), %d output", input, cached, 100*ratio(cached, input), output),
+ )
+
+ if len(sizes) > 0 {
+ lines = append(lines, sizeText(sizes))
+ }
+
+ if len(records) > 0 {
+ lines = append(lines,
+ "quota at the first call: "+quotaText(records[0]),
+ "quota at the last call: "+quotaText(records[len(records)-1]),
+ )
+ }
+
+ _, err := io.WriteString(out, strings.Join(lines, "\n")+"\n")
+
+ return err
+}
+
+// sizeText describes the sizes of the model lines.
+func sizeText(sizes []int) string {
+ sorted := slices.Sorted(slices.Values(sizes))
+ total, over := 0, 0
+ for _, size := range sorted {
+ total += size
+ if size > nodeBytes {
+ over++
+ }
+ }
+
+ 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)
+}
+
+// quotaText returns the quota headers of record, sorted by name.
+func quotaText(record usageRecord) string {
+ names := slices.Sorted(maps.Keys(record.Quota))
+ parts := make([]string, len(names))
+ for n, name := range names {
+ parts[n] = strings.TrimPrefix(name, "x-codex-") + "=" + record.Quota[name]
+ }
+
+ return strings.Join(parts, " ")
+}
+
+// ratio returns a / b, or 0 when b is 0.
+func ratio(a, b int) float64 {
+ if b == 0 {
+ return 0
+ }
+
+ return float64(a) / float64(b)
+}
--- /dev/null
+package main
+
+import (
+ "encoding/json"
+ "fmt"
+ "math/rand/v2"
+ "net/http"
+ "os"
+ "path/filepath"
+ "strconv"
+ "strings"
+ "testing"
+
+ "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
+// message {kind, text} per line, for "onechat probe".
+func TestProbeLogHasTheSyntheticMix(t *testing.T) {
+ messages := probeLog(probeMessages)
+ require.Len(t, messages, probeMessages)
+
+ counts := make(map[string]int)
+ long := 0
+ for _, message := range messages {
+ counts[message.Kind]++
+ require.True(t, kinds[message.Kind], message.Kind)
+ require.NotEmpty(t, message.Text)
+ for _, char := range message.Text {
+ require.True(t, char == '\n' || (char >= ' ' && char <= '~'), "character %q", char)
+ }
+
+ if len(message.Text) >= 30000 {
+ long++
+ }
+
+ if message.Kind == "work" {
+ require.True(t, strings.HasPrefix(message.Text, "["), message.Text)
+ }
+ }
+
+ require.Greater(t, counts["user"], 20)
+ require.Greater(t, counts["tool"], 300)
+ require.Greater(t, counts["echo"], 300)
+ require.Greater(t, long, 2)
+ require.Equal(t, messages, probeLog(probeMessages))
+
+ path := os.Getenv("ONECHAT_PROBE_LOG")
+ if path == "" {
+ return
+ }
+
+ var data strings.Builder
+ for _, message := range messages {
+ line, err := json.Marshal(message)
+ require.NoError(t, err)
+
+ data.Write(line)
+ data.WriteByte('\n')
+ }
+
+ 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))
+ }
+
+ 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)
+ require.NoError(t, err)
+ require.NoError(t, 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)
+ 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(), "quota at the first call: primary-used-percent=10")
+}
+
+// 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)),
+ })
+}
+
+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))]
+}
+
+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"}
+)
projects / onechat / commitdiff