Write private data archives with safe extraction

[onechat] / probe.go

package main

import (
        "bufio"
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "log/slog"
        "os"
        "path/filepath"
        "strconv"
        "time"
)

const (
        // probeName is the agent name of a probe run. The cache key comes from
        // it.
        probeName = "Probe"

        // probeCallLimit is the most model calls of a probe run.
        probeCallLimit = 660
)

// 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 buildQuotaLimit 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 buildMax,
// limits the builds that call the model at once in this run only. The core
// still queues up to buildMax builds. The run holds the lock of the token file,
// so it does not run while the daemon runs.
func probe(ctx context.Context, args []string, out io.Writer) (err 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 {
                return fmt.Errorf("parse probe count: %w", err)
        }

        if count <= 0 {
                return fmt.Errorf("COUNT %q is not a positive number", args[1])
        }

        builds := buildMax
        if len(args) == 3 {
                builds, err = strconv.Atoi(args[2])
                if err != nil {
                        return fmt.Errorf("parse probe builds: %w", err)
                }

                if builds < 1 || builds > buildMax {
                        return fmt.Errorf("BUILDS %q is not a number from 1 to %d", args[2], buildMax)
                }
        }

        entries, err := probeReadLog(args[0], count)
        if err != nil {
                return fmt.Errorf("read probe input: %w", err)
        }

        dir, err := dataDir()
        if err != nil {
                return fmt.Errorf("find probe data directory: %w", err)
        }

        config, err := loadConfig(dir)
        if err != nil {
                return fmt.Errorf("load probe configuration: %w", err)
        }

        tokens, err := lockTokens(dir)
        if err != nil {
                return fmt.Errorf("lock probe tokens: %w", err)
        }

        defer func() {
                closeErr := tokens.Close()
                if closeErr != nil {
                        err = errors.Join(err, fmt.Errorf("close probe token lock: %w", closeErr))
                }
        }()

        address, auth, err := modelEndpoint(config, dir)
        if err != nil {
                return fmt.Errorf("select probe endpoint: %w", 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 := newModelClient(address, auth, "onechat-"+probeName, usagePath)
        client.LimitCalls(probeCallLimit)

        summarizer := ProbeSlotSummarizer{
                inner: newCompactor(client, probePrompt, config.Compactor),
                slots: make(chan struct{}, builds),
        }

        // The report of the probe counts the failed calls, so the lines of the
        // failed builds are not needed.
        memory, err := openMemory(ctx, root, time.Now, summarizer, slog.New(slog.DiscardHandler))
        if err != nil {
                return fmt.Errorf("open probe memory: %w", 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 fmt.Errorf("write probe header: %w", errors.Join(err, memory.Close()))
        }

        run := BuildRun{
                memory:    memory,
                client:    client,
                usagePath: usagePath,
        }

        appended, stop := run.Run(ctx, entries)
        sizes, err := memory.buildModelLines()
        err = errors.Join(err, memory.Close())
        if err != nil {
                return fmt.Errorf("finish probe memory: %w", err)
        }

        records, err := buildReadUsage(usagePath)
        if err != nil {
                return fmt.Errorf("read probe usage: %w", err)
        }

        err = buildWriteReport(out, appended, stop, records, sizes)
        if err != nil {
                return fmt.Errorf("write probe report: %w", err)
        }

        return nil
}

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

// Summarize waits for a slot and builds the node with inner.
func (s ProbeSlotSummarizer) 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)
}

// probeReadLog reads the first count messages of the probe log at path. Source
// dates are ignored because a probe uses its own clock.
func probeReadLog(path string, count int) (entries []DatedMessage, err error) {
        file, err := os.Open(path)
        if err != nil {
                return nil, fmt.Errorf("open probe log: %w", err)
        }

        defer func() {
                closeErr := file.Close()
                if closeErr != nil {
                        err = errors.Join(err, fmt.Errorf("close probe log: %w", closeErr))
                }
        }()

        scanner := bufio.NewScanner(file)
        scanner.Buffer(make([]byte, 64*1024), buildLineBytes)

        for len(entries) < count && scanner.Scan() {
                var entry DatedMessage

                err = json.Unmarshal(scanner.Bytes(), &entry)
                if err != nil {
                        return nil, fmt.Errorf("probe log line %d: %w", len(entries)+1, err)
                }

                entry.Date = ""
                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
}

// probePrompt returns the system prompt of a probe run: the prompt of the agent
// probeName, without instructions and with a fixed instructions file path.
func probePrompt() (string, error) {
        return prompt(probeName, "/PROBE/instructions.md", ""), nil
}