[onechat] / probe.go

package main

import (
        "bufio"
        "cmp"
        "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

        // 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 [BUILDS]". 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.
// 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. BUILDS, from 1 to maxBuilds,
// limits the builds that call the model at once in this run only. The core
// still queues up to maxBuilds builds.
func probe(ctx context.Context, args []string, out io.Writer) error {
        if len(args) != 2 && len(args) != 3 {
                return errors.New("usage: onechat probe LOG COUNT [BUILDS]")
        }

        count, err := strconv.Atoi(args[1])
        if err != nil || count <= 0 {
                return fmt.Errorf("COUNT %q is not a positive number", args[1])
        }

        builds := maxBuilds
        if len(args) == 3 {
                builds, err = strconv.Atoi(args[2])
                if err != nil || builds < 1 || builds > maxBuilds {
                        return fmt.Errorf("BUILDS %q is not a number from 1 to %d", args[2], maxBuilds)
                }
        }

        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)
        client.LimitCalls(probeCallLimit)
        summarizer := slotSummarizer{
                inner: newCompactor(client, probeName, config.Compactor),
                slots: make(chan struct{}, builds),
        }

        memory, err := openMemory(ctx, root, time.Now, summarizer)
        if err != nil {
                return err
        }

        _, err = fmt.Fprintf(out, "probe: %d messages into %s with %s at %s effort, %d builds at once\n", len(entries), root, config.Compactor.Model, config.Compactor.Effort, builds)
        if err != nil {
                return errors.Join(err, memory.Close())
        }

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

        records, err := readUsageFile(usagePath)
        if err != nil {
                return err
        }

        return writeProbeReport(out, appended, stop, records, sizes)
}

// 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 {
                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()
                }

                reason := r.check()
                if reason != "" {
                        return n + 1, reason
                }
        }

        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 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()
                }
        }
}

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

// slotSummarizer lets at most cap(slots) builds call inner at once. A build
// past the limit waits for a slot.
type slotSummarizer struct {
        inner Summarizer
        slots chan struct{}
}

// Summarize waits for a slot and builds the node with inner.
func (s slotSummarizer) Summarize(ctx context.Context, task Task) (string, error) {
        select {
        case s.slots <- struct{}{}:
        case <-ctx.Done():
                return "", ctx.Err()
        }

        defer func() {
                <-s.slots
        }()

        return s.inner.Summarize(ctx, task)
}

// 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, in the order of the
// start of their calls. The file has them in the order of the end of their
// calls, which differs when calls run at once.
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)
        }

        slices.SortStableFunc(records, func(a, b usageRecord) int {
                return callStart(a).Compare(callStart(b))
        })

        return records, nil
}

// callStart returns the start of the call of record, or the zero time when its
// date is not valid.
func callStart(record usageRecord) time.Time {
        start, err := time.Parse(time.RFC3339Nano, record.Date)
        if err != nil {
                return time.Time{}
        }

        return start
}

// 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, %.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),
        )

        if len(sizes) > 0 {
                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]),
                        "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)
}

// 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 calls per minute, %.1f s per call (median %.1f, max %.1f), calls at once: at most %d, %.1f on average", wall, 60*float64(len(seconds))/max(wall, 1e-9), 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))
        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)
}