func (a *Agent) Prompt() (string, error) {
instructions, err := a.instructions()
if err != nil {
- return "", err
+ return "", fmt.Errorf("read prompt instructions: %w", err)
}
path, err := filepath.Abs(filepath.Join(a.dir, agentInstructionsFileName))
--- /dev/null
+package main
+
+import (
+ "bufio"
+ "cmp"
+ "context"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "io"
+ "maps"
+ "os"
+ "slices"
+ "strconv"
+ "strings"
+ "time"
+)
+
+// buildLineBytes is the most bytes of one line in a build input or usage log.
+const buildLineBytes = 1 << 20
+
+// buildQuotaLimit is the largest increase in used quota points of a run.
+const buildQuotaLimit = 15.0
+
+// buildCallLimit is the most model calls of a run.
+const buildCallLimit = 660
+
+// BuildRun appends messages, builds their tree, and stops the run by the rules
+// of buildStop.
+type BuildRun struct {
+ memory *Memory
+ client *ModelClient
+ usagePath string
+}
+
+// Run appends the entries. It waits for the builds only before a user message,
+// as a turn does, so up to buildMax 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 *BuildRun) Run(ctx context.Context, entries []DatedMessage) (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 *BuildRun) append(ctx context.Context, entries []DatedMessage) (int, string) {
+ for n, entry := range entries {
+ if entry.Kind == "user" {
+ reason := r.wait(ctx)
+ if reason != "" {
+ return n, reason
+ }
+ }
+
+ if ctx.Err() != nil {
+ return n, ctx.Err().Error()
+ }
+
+ _, err := r.memory.AppendDated(entry.Kind, entry.Text, entry.Date)
+ 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 *BuildRun) wait(ctx context.Context) string {
+ for {
+ reason := r.check()
+ if reason != "" {
+ return reason
+ }
+
+ err := r.waitBuilds(ctx)
+ 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 *BuildRun) check() string {
+ if r.client.LimitReached() {
+ return fmt.Sprintf("call limit of %d reached", r.client.CallLimit())
+ }
+
+ records, err := buildReadUsage(r.usagePath)
+ if err != nil {
+ return err.Error()
+ }
+
+ return buildStop(records)
+}
+
+// waitBuilds waits up to one second without leaving a child context open.
+func (r *BuildRun) waitBuilds(ctx context.Context) error {
+ ctx, cancel := context.WithTimeout(ctx, time.Second)
+ defer cancel()
+
+ return r.memory.Wait(ctx)
+}
+
+// DatedMessage is one message with an optional RFC 3339 date.
+type DatedMessage struct {
+ Kind string `json:"kind"`
+ Text string `json:"text"`
+ Date string `json:"date,omitempty"`
+}
+
+// buildReadUsage 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 buildReadUsage(path string) (records []ModelUsageRecord, 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() {
+ closeErr := file.Close()
+ if closeErr != nil {
+ err = errors.Join(err, fmt.Errorf("close build usage: %w", closeErr))
+ }
+ }()
+
+ scanner := bufio.NewScanner(file)
+ scanner.Buffer(make([]byte, 64*1024), buildLineBytes)
+ for scanner.Scan() {
+ var record ModelUsageRecord
+ 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 ModelUsageRecord) int {
+ return buildCallStart(a).Compare(buildCallStart(b))
+ })
+
+ return records, nil
+}
+
+// buildWriteReport writes the result of a build run.
+func buildWriteReport(out io.Writer, appended int, stop string, records []ModelUsageRecord, 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, buildRatio(len(records), appended)),
+ fmt.Sprintf("model lines: %d, retries: %d, %.2f calls per model line", len(sizes), len(records)-failed-len(sizes), buildRatio(len(records), len(sizes))),
+ fmt.Sprintf("tokens: %d input, %d cached (%.1f%%), %d output", input, cached, 100*buildRatio(cached, input), output),
+ )
+
+ if len(sizes) > 0 {
+ lines = append(lines, buildSizeText(sizes))
+ }
+
+ if len(records) > 0 {
+ lines = append(lines, buildTimeText(records))
+ }
+
+ if len(records) > 0 {
+ lines = append(lines,
+ "quota at the first call: "+buildQuotaText(records[0]),
+ "quota at the last call: "+buildQuotaText(records[len(records)-1]),
+ )
+ }
+
+ _, err := io.WriteString(out, strings.Join(lines, "\n")+"\n")
+ if err != nil {
+ return fmt.Errorf("write build report: %w", err)
+ }
+
+ return nil
+}
+
+// buildStop returns why a build run must stop: a failed call, or a quota window
+// whose used percent grew by more than buildQuotaLimit points since the first
+// call. It returns "" when the run can go on.
+func buildStop(records []ModelUsageRecord) string {
+ for _, record := range records {
+ if record.Error != "" {
+ return "call failed: " + record.Error
+ }
+ }
+
+ if len(records) == 0 {
+ return ""
+ }
+
+ first := buildUsedPercents(records[0])
+ for name, used := range buildUsedPercents(records[len(records)-1]) {
+ start, found := first[name]
+ if found && used-start > buildQuotaLimit {
+ return fmt.Sprintf("%s grew from %g to %g", name, start, used)
+ }
+ }
+
+ return ""
+}
+
+// buildCallStart returns the start of the call of record, or the zero time when
+// its date is not valid.
+func buildCallStart(record ModelUsageRecord) time.Time {
+ start, err := time.Parse(time.RFC3339Nano, record.Date)
+ if err != nil {
+ return time.Time{}
+ }
+
+ return start
+}
+
+// buildRatio returns a / b, or 0 when b is 0.
+func buildRatio(a, b int) float64 {
+ if b == 0 {
+ return 0
+ }
+
+ return float64(a) / float64(b)
+}
+
+// buildSizeText describes the sizes of the model lines.
+func buildSizeText(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], buildRatio(total, len(sorted)), sorted[len(sorted)-1], over, nodeBytes)
+}
+
+// buildTimeText describes the time per call and how many calls ran at once. A
+// call runs from its date for its seconds.
+func buildTimeText(records []ModelUsageRecord) 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))
+}
+
+// buildQuotaText returns the quota headers of record, sorted by name.
+func buildQuotaText(record ModelUsageRecord) 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, " ")
+}
+
+// buildUsedPercents returns the used percent of each quota window of record, by
+// header name.
+func buildUsedPercents(record ModelUsageRecord) 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
+}
--- /dev/null
+package main
+
+import (
+ "net/http"
+ "path/filepath"
+ "strconv"
+ "strings"
+ "sync"
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/require"
+)
+
+func TestBuildRunStopsOnFailedCallsOrQuotaGrowth(t *testing.T) {
+ quota := func(primary, secondary string) ModelUsageRecord {
+ return ModelUsageRecord{
+ 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 []ModelUsageRecord
+ want string
+ }{
+ {
+ name: "no calls",
+ records: nil,
+ want: "",
+ },
+ {
+ name: "growth within the limit",
+ records: []ModelUsageRecord{quota("10", "40"), quota("25", "41")},
+ want: "",
+ },
+ {
+ name: "growth over the limit",
+ records: []ModelUsageRecord{quota("10", "40"), quota("12", "55.5")},
+ want: "x-codex-secondary-used-percent grew from 40 to 55.5",
+ },
+ {
+ name: "failed call",
+ records: []ModelUsageRecord{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, buildStop(tc.records))
+ })
+ }
+}
+
+func TestBuildRunStopsWhenTheQuotaGrows(t *testing.T) {
+ server := newModelFake(t)
+ server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
+ w.Header().Set("X-Codex-Primary-Used-Percent", strconv.Itoa(10+call.number))
+ modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem(strings.Repeat("s", 400))}, 1000, 800, 50))
+ }
+
+ run := newTestBuildRun(t, server)
+ appended, stop := run.Run(t.Context(), probeLog(probeMessages))
+ sizes, err := run.memory.buildModelLines()
+ require.NoError(t, err)
+ require.NoError(t, run.memory.Close())
+
+ require.Regexp(t, `^x-codex-primary-used-percent grew from \d+ to \d+$`, stop)
+ require.Less(t, appended, probeMessages)
+ require.NotEmpty(t, sizes)
+
+ records, err := buildReadUsage(run.usagePath)
+ require.NoError(t, err)
+ require.Greater(t, len(records), int(buildQuotaLimit))
+
+ var report strings.Builder
+ require.NoError(t, buildWriteReport(&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=")
+}
+
+func TestBuildRunStopsAtTheCallLimit(t *testing.T) {
+ server := newModelFake(t)
+ run := newTestBuildRun(t, server)
+ run.client.LimitCalls(5)
+
+ entries := make([]DatedMessage, 40)
+ for n := range entries {
+ entries[n] = DatedMessage{
+ Kind: "echo",
+ Text: strings.Repeat("long output ", 100),
+ }
+ }
+
+ _, 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())
+}
+
+func TestBuildRunWaitsOnlyBeforeUserMessages(t *testing.T) {
+ gate := make(chan struct{})
+ var mutex sync.Mutex
+ running, most := 0, 0
+ server := newModelFake(t)
+ server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
+ mutex.Lock()
+ running++
+ most = max(most, running)
+ mutex.Unlock()
+
+ <-gate
+
+ mutex.Lock()
+ running--
+ mutex.Unlock()
+
+ modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("a line")}, 100, 0, 5))
+ }
+
+ run := newTestBuildRun(t, server)
+ long := strings.Repeat("long output ", 100)
+ entries := []DatedMessage{
+ {
+ Kind: "echo",
+ Text: long,
+ },
+ {
+ Kind: "tool",
+ Text: long,
+ },
+ {
+ Kind: "user",
+ Text: "next",
+ },
+ {
+ Kind: "reply",
+ Text: long,
+ },
+ }
+
+ done := make(chan string)
+ go func() {
+ _, stop := run.Run(t.Context(), entries)
+ done <- stop
+ }()
+
+ require.Eventually(t, func() bool {
+ return server.Calls() == 2
+ }, 10*time.Second, time.Millisecond)
+
+ appended := run.memory.Len()
+ 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 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 := buildReadUsage(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)
+ record := func(offset, seconds float64) ModelUsageRecord {
+ return ModelUsageRecord{
+ Date: start.Add(time.Duration(offset * float64(time.Second))).Format(time.RFC3339Nano),
+ Seconds: seconds,
+ }
+ }
+
+ text := buildTimeText([]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)
+}
+
+func newTestBuildRun(t *testing.T, server *ModelFake) *BuildRun {
+ t.Helper()
+
+ root := t.TempDir()
+ usagePath := filepath.Join(root, "usage.jsonl")
+ client := newModelClient(server.URL, ModelAPIKey(modelTestKey), "onechat-"+probeName, usagePath)
+ memory, err := openMemory(t.Context(), root, newTestClock().Now, newCompactor(client, probePrompt, defaultConfig().Compactor), discardLog())
+ require.NoError(t, err)
+
+ return &BuildRun{
+ memory: memory,
+ client: client,
+ usagePath: usagePath,
+ }
+}
}
// Log is the append-only message log of one chat. It writes one JSON line for
-// each message to main/YYYY-MM-DD.jsonl, named for the local day of the write.
-// It never edits or deletes a line. Each line goes to the operating system
-// before Append returns, so a crash of the process loses no message. An
+// each message to main/YYYY-MM-DD.jsonl, named for the local day of the
+// message. It never edits or deletes a line. Each line goes to the operating
+// system before Append returns, so a crash of the process loses no message. An
// exclusive lock on the lock file makes sure that only one process writes the
// log. After a failed write, the end of the file is not known, so the log
// refuses all writes until it is opened again. The caller serializes all calls.
// Append writes a new message to the log and returns it with its ID, size, and
// date.
func (l *Log) Append(kind, text string) (Message, error) {
+ return l.AppendDated(kind, text, "")
+}
+
+// AppendDated writes a message with date in RFC 3339 format. An empty date uses
+// the clock. The day file uses the local day of the given date.
+func (l *Log) AppendDated(kind, text, date string) (Message, error) {
if l.lock == nil {
- return Message{}, os.ErrClosed
+ return Message{}, fmt.Errorf("append dated message: %w", os.ErrClosed)
}
if l.failure != nil {
return Message{}, fmt.Errorf("log needs a reopen after a failed write: %w", l.failure)
}
+ var err error
now := l.clock()
+ if date == "" {
+ date = now.Format(time.RFC3339Nano)
+ } else {
+ now, err = time.Parse(time.RFC3339Nano, date)
+ if err != nil {
+ return Message{}, fmt.Errorf("parse message date: %w", err)
+ }
+
+ now = now.In(time.Local)
+ }
+
message := Message{
ID: l.index.count,
Kind: kind,
Text: text,
Size: len(text),
- Date: now.Format(time.RFC3339Nano),
+ Date: date,
}
- err := logValidateMessage(message)
+ err = logValidateMessage(message)
if err != nil {
- return Message{}, err
+ return Message{}, fmt.Errorf("validate dated message: %w", err)
}
line, err := json.Marshal(message)
return l.index.Append(record)
}
-// Message is one entry of the chat log. ID is the permanent position of the
-// message in the chat, from 0. Size is the UTF-8 byte length of Text. Date is
-// the time of the write in RFC 3339 format.
+// Message is one entry of the chat log. ID is its position in the log, from 0.
+// A rebuild assigns positions in the merged log. Size is the UTF-8 byte length
+// of Text. Date is the message time in RFC 3339 format, supplied at append or
+// taken from the clock.
type Message struct {
ID int `json:"i"`
Kind string `json:"kind"`
// Append logs a new message, adds it to the tree and the views, and starts the
// builds that are ready, also the failed ones.
func (m *Memory) Append(kind, text string) (Message, error) {
+ return m.AppendDated(kind, text, "")
+}
+
+// AppendDated appends a message with its RFC 3339 date and builds its nodes. An
+// empty date uses the clock, as Append does.
+func (m *Memory) AppendDated(kind, text, date string) (Message, error) {
m.mutex.Lock()
defer m.mutex.Unlock()
- message, err := m.log.Append(kind, text)
+ message, err := m.log.AppendDated(kind, text, date)
if err != nil {
- return Message{}, err
+ return Message{}, fmt.Errorf("append dated log message: %w", err)
}
defer m.syncer.Written()
err = m.addMessage(message)
if err != nil {
- return Message{}, err
+ return Message{}, fmt.Errorf("add dated message to memory: %w", err)
}
err = m.save()
if err != nil {
- return Message{}, err
+ return Message{}, fmt.Errorf("save dated message view: %w", err)
}
m.queue.Retry()
return m.context(level, index)
}
+// rebuildComplete checks the whole tree, not only the visible view. A compactor
+// can fail without an HTTP error, for example with an empty line. Such a failed
+// merge can be absent from the view and must prevent a swap.
+func (m *Memory) rebuildComplete() error {
+ m.mutex.Lock()
+ defer m.mutex.Unlock()
+
+ if len(m.queue.leaves) != 0 || len(m.queue.queued) != 0 {
+ return errors.New("rebuild failed: summary tree has unfinished nodes")
+ }
+
+ return nil
+}
+
+// buildModelLines 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 (m *Memory) buildModelLines() ([]int, error) {
+ m.mutex.Lock()
+ defer m.mutex.Unlock()
+
+ keys, err := m.tree.Keys()
+ if err != nil {
+ return nil, fmt.Errorf("list built nodes: %w", err)
+ }
+
+ var sizes []int
+ for _, key := range keys {
+ var node Node
+
+ node, _, err = m.tree.Node(key.level, key.index)
+ if err != nil {
+ return nil, fmt.Errorf("read built node: %w", err)
+ }
+
+ var model bool
+
+ model, err = m.buildBuiltByModel(key)
+ if err != nil {
+ return nil, fmt.Errorf("check model source: %w", err)
+ }
+
+ if model {
+ sizes = append(sizes, node.Size)
+ }
+ }
+
+ return sizes, nil
+}
+
+// buildBuiltByModel 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 (m *Memory) buildBuiltByModel(key NodeKey) (bool, error) {
+ if key.level == 0 {
+ message, err := m.log.Message(key.index)
+ if err != nil {
+ return false, fmt.Errorf("read leaf source: %w", err)
+ }
+
+ _, free := freeLeaf(message)
+
+ return !free, nil
+ }
+
+ left, right, err := m.children(key)
+ if err != nil {
+ return false, fmt.Errorf("read merge sources: %w", err)
+ }
+
+ _, free := freeMerge(left, right)
+
+ return !free, nil
+}
+
// Task is the build of the node (Level, Index). Context is the compaction view
// for the build. Input holds the source "kind: text" of a leaf, or the two
// child lines of a merge with their newlines as spaces.
return c.limit > 0 && c.calls.Load() >= c.limit
}
+// CallLimit returns the configured call limit, or zero for no limit.
+func (c *ModelClient) CallLimit() int64 {
+ return c.limit
+}
+
+// UsagePath returns the file that records the calls of this client.
+func (c *ModelClient) UsagePath() string {
+ return c.usage.path
+}
+
// Create sends the call and reads its event stream to the end. Over the call
// limit it sends nothing and returns errModelCallLimit.
func (c *ModelClient) Create(ctx context.Context, request ModelRequest) (ModelResponse, error) {
import (
"bufio"
- "cmp"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log"
- "maps"
"os"
"path/filepath"
- "slices"
"strconv"
- "strings"
"time"
)
// 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
+// the first failed call, after buildCallLimit 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
}
count, err := strconv.Atoi(args[1])
- if err != nil || count <= 0 {
+ 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 || builds < 1 || builds > buildMax {
+ 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 err
+ return fmt.Errorf("read probe input: %w", err)
}
dir, err := dataDir()
if err != nil {
- return err
+ return fmt.Errorf("find probe data directory: %w", err)
}
config, err := loadConfig(dir)
if err != nil {
- return err
+ return fmt.Errorf("load probe configuration: %w", err)
}
tokens, err := lockTokens(dir)
if err != nil {
- return err
+ return fmt.Errorf("lock probe tokens: %w", err)
}
defer func() {
- err = errors.Join(err, tokens.Close())
+ 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 err
+ return fmt.Errorf("select probe endpoint: %w", err)
}
root := filepath.Join(dir, "probe", time.Now().Format("20060102-150405"))
usagePath := filepath.Join(root, "usage.jsonl")
client := newModelClient(address, auth, "onechat-"+probeName, usagePath)
- client.LimitCalls(probeCallLimit)
+ client.LimitCalls(buildCallLimit)
summarizer := ProbeSlotSummarizer{
inner: newCompactor(client, probePrompt, config.Compactor),
slots: make(chan struct{}, builds),
// failed builds are not needed.
memory, err := openMemory(ctx, root, time.Now, summarizer, log.New(io.Discard, "", 0))
if err != nil {
- return err
+ 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 errors.Join(err, memory.Close())
+ return fmt.Errorf("write probe header: %w", errors.Join(err, memory.Close()))
}
- run := ProbeRun{
+ run := BuildRun{
memory: memory,
client: client,
usagePath: usagePath,
}
appended, stop := run.Run(ctx, entries)
- sizes, err := probeModelLines(memory)
+ sizes, err := memory.buildModelLines()
err = errors.Join(err, memory.Close())
if err != nil {
- return err
+ return fmt.Errorf("finish probe memory: %w", err)
}
- records, err := probeReadUsage(usagePath)
+ records, err := buildReadUsage(usagePath)
if err != nil {
- return err
+ return fmt.Errorf("read probe usage: %w", err)
}
- return probeWriteReport(out, appended, stop, records, sizes)
+ 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
return s.inner.Summarize(ctx, task)
}
-// 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 *ModelClient
- usagePath string
-}
-
-// Run appends the entries. It waits for the builds only before a user message,
-// as a turn does, so up to buildMax 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 := probeReadUsage(r.usagePath)
- if err != nil {
- return err.Error()
- }
-
- return probeStop(records)
-}
-
-// ProbeEntry is one message of a probe log.
-type ProbeEntry struct {
- Kind string `json:"kind"`
- Text string `json:"text"`
-}
-
-// probeReadLog reads the first count messages of the probe log at path.
-func probeReadLog(path string, count int) (entries []ProbeEntry, err error) {
+// 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() {
- err = errors.Join(err, file.Close())
+ 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), probeLineBytes)
+ scanner.Buffer(make([]byte, 64*1024), buildLineBytes)
for len(entries) < count && scanner.Scan() {
- var entry ProbeEntry
+ 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)
}
func probePrompt() (string, error) {
return prompt(probeName, "/PROBE/instructions.md", ""), nil
}
-
-// probeModelLines 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 probeModelLines(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 := probeBuiltByModel(memory, key)
- if err != nil {
- return nil, err
- }
-
- if model {
- sizes = append(sizes, node.Size)
- }
- }
-
- return sizes, nil
-}
-
-// probeReadUsage 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 probeReadUsage(path string) (records []ModelUsageRecord, 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 ModelUsageRecord
- 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 ModelUsageRecord) int {
- return probeCallStart(a).Compare(probeCallStart(b))
- })
-
- return records, nil
-}
-
-// probeWriteReport writes the result of a probe run.
-func probeWriteReport(out io.Writer, appended int, stop string, records []ModelUsageRecord, 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, probeRatio(len(records), appended)),
- fmt.Sprintf("model lines: %d, retries: %d, %.2f calls per model line", len(sizes), len(records)-failed-len(sizes), probeRatio(len(records), len(sizes))),
- fmt.Sprintf("tokens: %d input, %d cached (%.1f%%), %d output", input, cached, 100*probeRatio(cached, input), output),
- )
-
- if len(sizes) > 0 {
- lines = append(lines, probeSizeText(sizes))
- }
-
- if len(records) > 0 {
- lines = append(lines, probeTimeText(records))
- }
-
- if len(records) > 0 {
- lines = append(lines,
- "quota at the first call: "+probeQuotaText(records[0]),
- "quota at the last call: "+probeQuotaText(records[len(records)-1]),
- )
- }
-
- _, err := io.WriteString(out, strings.Join(lines, "\n")+"\n")
-
- return err
-}
-
-// 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 []ModelUsageRecord) string {
- for _, record := range records {
- if record.Error != "" {
- return "call failed: " + record.Error
- }
- }
-
- if len(records) == 0 {
- return ""
- }
-
- first := probeUsedPercents(records[0])
- for name, used := range probeUsedPercents(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 ""
-}
-
-// probeBuiltByModel 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 probeBuiltByModel(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
-}
-
-// probeCallStart returns the start of the call of record, or the zero time when
-// its date is not valid.
-func probeCallStart(record ModelUsageRecord) time.Time {
- start, err := time.Parse(time.RFC3339Nano, record.Date)
- if err != nil {
- return time.Time{}
- }
-
- return start
-}
-
-// probeRatio returns a / b, or 0 when b is 0.
-func probeRatio(a, b int) float64 {
- if b == 0 {
- return 0
- }
-
- return float64(a) / float64(b)
-}
-
-// probeSizeText describes the sizes of the model lines.
-func probeSizeText(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], probeRatio(total, len(sorted)), sorted[len(sorted)-1], over, nodeBytes)
-}
-
-// probeTimeText describes the time per call and how many calls ran at once. A
-// call runs from its date for its seconds.
-func probeTimeText(records []ModelUsageRecord) 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))
-}
-
-// probeQuotaText returns the quota headers of record, sorted by name.
-func probeQuotaText(record ModelUsageRecord) 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, " ")
-}
-
-// probeUsedPercents returns the used percent of each quota window of record, by
-// header name.
-func probeUsedPercents(record ModelUsageRecord) 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
-}
--- /dev/null
+package main
+
+import (
+ "encoding/json"
+ "fmt"
+ "math/rand/v2"
+ "strings"
+)
+
+// probeMessages is the length of the generated log for the live probe.
+const probeMessages = 1000
+
+// 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(probePeople) + "] 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(probeThings), g.pick(probeProblems)),
+ fmt.Sprintf("Move the %s to %s before %s.", g.pick(probeThings), g.pick(probePlaces), g.pick(probeDays)),
+ fmt.Sprintf("Thanks, that works. Next, %s.", g.pick(probeTasks)),
+ fmt.Sprintf("No, keep the old %s. %s asked for it.", g.pick(probeThings), g.pick(probePeople)),
+ fmt.Sprintf("How many %s orders came in on %s?", g.pick(probeBreads), g.pick(probeDays)),
+ })
+}
+
+func (g *ProbeWriter) replyLine() string {
+ return g.pick([]string{
+ fmt.Sprintf("I looked at %s. The %s %s because %s.", g.pick(probeFiles), g.pick(probeThings), g.pick(probeProblems), g.pick(probeCauses)),
+ fmt.Sprintf("I changed %s, and the tests pass now.", g.pick(probeChanges)),
+ fmt.Sprintf("There were %d %s orders on %s, %d of them late.", 10+g.random.IntN(90), g.pick(probeBreads), g.pick(probeDays), g.random.IntN(10)),
+ fmt.Sprintf("Open question: should %s get %s?", g.pick(probePeople), g.pick(probeTasks)),
+ })
+}
+
+func (g *ProbeWriter) noteLine() string {
+ return g.pick([]string{
+ fmt.Sprintf("%s prefers %s.", g.pick(probePeople), g.pick(probePreferences)),
+ 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(probeDays)),
+ })
+}
+
+func (g *ProbeWriter) pick(words []string) string {
+ return words[g.random.IntN(len(words))]
+}
+
+// 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(probeThings) + "' " + g.pick(probeFiles)}
+ case "read":
+ input = map[string]string{"path": g.pick(probeFiles)}
+ default:
+ input = map[string]string{
+ "path": g.pick(probeFiles),
+ "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) 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(probeBreads), g.pick(probePeople)),
+ fmt.Sprintf("%s:%d: %s %s.", g.pick(probeFiles), 1+g.random.IntN(400), g.pick(probeThings), g.pick(probeProblems)),
+ fmt.Sprintf("ok bakery/%s %d.%03ds", g.pick(probeModules), g.random.IntN(3), g.random.IntN(1000)),
+ })
+}
+
+// probeLog returns count messages of made-up work on a small bakery order
+// system, with the mix of kinds and sizes of replayMessage. The same count
+// always gives the same log.
+func probeLog(count int) []DatedMessage {
+ g := ProbeWriter{
+ random: rand.New(rand.NewPCG(7, 11)),
+ }
+
+ messages := make([]DatedMessage, count)
+ for id := range messages {
+ kind, text := g.message()
+ messages[id] = DatedMessage{
+ Kind: kind,
+ Text: text,
+ }
+ }
+
+ return messages
+}
+
+var (
+ probePeople = []string{"Mara", "Joel", "Ines", "Tomas", "Priya", "Lena"}
+ probeBreads = []string{"rye", "sourdough", "baguette", "brioche", "spelt", "seed loaf"}
+ probeDays = []string{"Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday"}
+ probeThings = []string{"order form", "price list", "delivery route", "stock count", "invoice export", "oven schedule"}
+ probeProblems = []string{"shows the wrong total", "skips the last row", "is slow at noon", "loses the customer note", "rounds the tax down"}
+ probePlaces = []string{"the front page", "the weekly report", "the shared folder", "the admin menu"}
+ probeTasks = []string{"add a field for allergies", "send the invoices on Friday", "split the route in two", "print labels for the shelf"}
+ probeCauses = []string{"the date has no time zone", "the cache keeps old prices", "the loop stops one row early", "the form trims the note"}
+ probeChanges = []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"}
+ probePreferences = []string{"short answers", "a call before a change to prices", "invoices as PDF", "reports in the morning"}
+ probeFiles = []string{"orders/form.go", "prices/list.go", "routes/plan.go", "stock/count.go", "export/invoice.go", "ovens/schedule.go"}
+ probeModules = []string{"orders", "prices", "routes", "stock", "export", "ovens"}
+)
import (
"encoding/json"
- "fmt"
- "math/rand/v2"
- "net/http"
"os"
"path/filepath"
- "strconv"
"strings"
"sync"
"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
require.NoError(t, os.WriteFile(path, []byte(data.String()), 0o600))
}
-func TestProbeStopRules(t *testing.T) {
- quota := func(primary, secondary string) ModelUsageRecord {
- return ModelUsageRecord{
- 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 []ModelUsageRecord
- want string
- }{
- {
- name: "no calls",
- records: nil,
- want: "",
- },
- {
- name: "growth within the limit",
- records: []ModelUsageRecord{quota("10", "40"), quota("25", "41")},
- want: "",
- },
- {
- name: "growth over the limit",
- records: []ModelUsageRecord{quota("10", "40"), quota("12", "55.5")},
- want: "x-codex-secondary-used-percent grew from 40 to 55.5",
- },
- {
- name: "failed call",
- records: []ModelUsageRecord{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 := newModelFake(t)
- server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
- w.Header().Set("X-Codex-Primary-Used-Percent", strconv.Itoa(10+call.number))
- modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem(strings.Repeat("s", 400))}, 1000, 800, 50))
- }
+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))
- run := newTestProbeRun(t, server)
- appended, stop := run.Run(t.Context(), probeLog(probeMessages))
- sizes, err := probeModelLines(run.memory)
+ entries, err := probeReadLog(path, 1)
require.NoError(t, err)
- require.NoError(t, run.memory.Close())
-
- require.Regexp(t, `^x-codex-primary-used-percent grew from \d+ to \d+$`, stop)
- require.Less(t, appended, probeMessages)
- require.NotEmpty(t, sizes)
-
- records, err := probeReadUsage(run.usagePath)
- require.NoError(t, err)
- require.Greater(t, len(records), int(probeQuotaLimit))
-
- var report strings.Builder
- require.NoError(t, probeWriteReport(&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=")
-}
-
-func TestProbeRunStopsAtTheCallLimit(t *testing.T) {
- server := newModelFake(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),
- }
- }
-
- _, 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())
-}
-
-func TestProbeRunWaitsOnlyBeforeUserMessages(t *testing.T) {
- gate := make(chan struct{})
- var mutex sync.Mutex
- running, most := 0, 0
- server := newModelFake(t)
- server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
- mutex.Lock()
- running++
- most = max(most, running)
- mutex.Unlock()
-
- <-gate
-
- mutex.Lock()
- running--
- mutex.Unlock()
-
- modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("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,
- },
- }
-
- done := make(chan string)
- go func() {
- _, stop := run.Run(t.Context(), entries)
- done <- stop
- }()
-
- require.Eventually(t, func() bool {
- return server.Calls() == 2
- }, 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())
+ require.Len(t, entries, 1)
+ require.Empty(t, entries[0].Date)
}
func TestSlotSummarizerLimitsTheBuildsAtOnce(t *testing.T) {
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 := probeReadUsage(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)
- record := func(offset, seconds float64) ModelUsageRecord {
- return ModelUsageRecord{
- Date: start.Add(time.Duration(offset * float64(time.Second))).Format(time.RFC3339Nano),
- Seconds: seconds,
- }
- }
-
- text := probeTimeText([]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)
-}
-
-// 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(probePeople) + "] 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(probeThings), g.pick(probeProblems)),
- fmt.Sprintf("Move the %s to %s before %s.", g.pick(probeThings), g.pick(probePlaces), g.pick(probeDays)),
- fmt.Sprintf("Thanks, that works. Next, %s.", g.pick(probeTasks)),
- fmt.Sprintf("No, keep the old %s. %s asked for it.", g.pick(probeThings), g.pick(probePeople)),
- fmt.Sprintf("How many %s orders came in on %s?", g.pick(probeBreads), g.pick(probeDays)),
- })
-}
-
-func (g *ProbeWriter) replyLine() string {
- return g.pick([]string{
- fmt.Sprintf("I looked at %s. The %s %s because %s.", g.pick(probeFiles), g.pick(probeThings), g.pick(probeProblems), g.pick(probeCauses)),
- fmt.Sprintf("I changed %s, and the tests pass now.", g.pick(probeChanges)),
- fmt.Sprintf("There were %d %s orders on %s, %d of them late.", 10+g.random.IntN(90), g.pick(probeBreads), g.pick(probeDays), g.random.IntN(10)),
- fmt.Sprintf("Open question: should %s get %s?", g.pick(probePeople), g.pick(probeTasks)),
- })
-}
-
-func (g *ProbeWriter) noteLine() string {
- return g.pick([]string{
- fmt.Sprintf("%s prefers %s.", g.pick(probePeople), g.pick(probePreferences)),
- 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(probeDays)),
- })
-}
-
-func (g *ProbeWriter) pick(words []string) string {
- return words[g.random.IntN(len(words))]
-}
-
-// 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(probeThings) + "' " + g.pick(probeFiles)}
- case "read":
- input = map[string]string{"path": g.pick(probeFiles)}
- default:
- input = map[string]string{
- "path": g.pick(probeFiles),
- "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) 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(probeBreads), g.pick(probePeople)),
- fmt.Sprintf("%s:%d: %s %s.", g.pick(probeFiles), 1+g.random.IntN(400), g.pick(probeThings), g.pick(probeProblems)),
- fmt.Sprintf("ok bakery/%s %d.%03ds", g.pick(probeModules), g.random.IntN(3), g.random.IntN(1000)),
- })
-}
-
-// probeLog returns count messages of made-up work on a small bakery order
-// system, with the mix of kinds and sizes of replayMessage. 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
-}
-
-func newTestProbeRun(t *testing.T, server *ModelFake) *ProbeRun {
- t.Helper()
-
- root := t.TempDir()
- usagePath := filepath.Join(root, "usage.jsonl")
- client := newModelClient(server.URL, ModelAPIKey(modelTestKey), "onechat-"+probeName, usagePath)
- memory, err := openMemory(t.Context(), root, newTestClock().Now, newCompactor(client, probePrompt, defaultConfig().Compactor), discardLog())
- require.NoError(t, err)
-
- return &ProbeRun{
- memory: memory,
- client: client,
- usagePath: usagePath,
- }
-}
-
-var (
- probePeople = []string{"Mara", "Joel", "Ines", "Tomas", "Priya", "Lena"}
- probeBreads = []string{"rye", "sourdough", "baguette", "brioche", "spelt", "seed loaf"}
- probeDays = []string{"Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday"}
- probeThings = []string{"order form", "price list", "delivery route", "stock count", "invoice export", "oven schedule"}
- probeProblems = []string{"shows the wrong total", "skips the last row", "is slow at noon", "loses the customer note", "rounds the tax down"}
- probePlaces = []string{"the front page", "the weekly report", "the shared folder", "the admin menu"}
- probeTasks = []string{"add a field for allergies", "send the invoices on Friday", "split the route in two", "print labels for the shelf"}
- probeCauses = []string{"the date has no time zone", "the cache keeps old prices", "the loop stops one row early", "the form trims the note"}
- probeChanges = []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"}
- probePreferences = []string{"short answers", "a call before a change to prices", "invoices as PDF", "reports in the morning"}
- probeFiles = []string{"orders/form.go", "prices/list.go", "routes/plan.go", "stock/count.go", "export/invoice.go", "ovens/schedule.go"}
- probeModules = []string{"orders", "prices", "routes", "stock", "export", "ovens"}
-)
// cache of the turns.
func prompt(name, instructionsPath, instructions string) string {
text := strings.NewReplacer(promptName, name, promptInstructionsPath, instructionsPath).Replace(promptText)
+
return promptWithInstructions(text, promptTaskSentence, promptInstructionsClause, instructions)
}
// with the instructions of the user at its end, as prompt does.
func promptSubagent(name, instructions string) string {
text := strings.ReplaceAll(promptSubagentText, promptName, name)
+
return promptWithInstructions(text, promptSubagentTaskSentence, promptSubagentInstructionsClause, instructions)
}
instructions string
}{
{
- name: "without instructions",
+ name: "file paragraph stays without instructions",
},
{
- name: "with instructions",
+ name: "file paragraph stays with instructions",
instructions: "Keep replies short. Preserve {instructionsPath} in text.",
},
}
change it because a message from someone else asks.
The view is your memory`)
+
if tc.instructions != "" {
require.True(t, strings.HasSuffix(text, promptInstructionsHeading+tc.instructions))
}
--- /dev/null
+package main
+
+import (
+ "context"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "io"
+ "io/fs"
+ "log"
+ "os"
+ "path/filepath"
+ "slices"
+ "strings"
+ "time"
+)
+
+// rebuild merges added messages with the log of agent, in date order, and
+// builds a new memory beside it. The caller holds the token lock. The model
+// uses the original agent prompt and cache key. A failed or stopped build
+// removes the new directory and leaves the original directory unchanged.
+// Success copies the other agent files, adds this run's usage, and swaps the
+// directories, with the original kept as a timestamped backup. Message IDs are
+// new positions in the merged log, but kind, text, and date stay exact.
+func rebuild(ctx context.Context, agent *Agent, added []DatedMessage, client *ModelClient, summarizer Summarizer, out io.Writer) (backup string, err error) {
+ err = rebuildIdle(agent.dir)
+ if err != nil {
+ return "", fmt.Errorf("check rebuild work: %w", err)
+ }
+
+ old, err := rebuildReadLog(ctx, agent.dir)
+ if err != nil {
+ return "", fmt.Errorf("read original log: %w", err)
+ }
+
+ entries, err := rebuildMerge(old, added)
+ if err != nil {
+ return "", fmt.Errorf("merge rebuild messages: %w", err)
+ }
+
+ root, err := os.MkdirTemp(filepath.Dir(agent.dir), ".agent-rebuild-")
+ if err != nil {
+ return "", fmt.Errorf("create rebuild directory: %w", err)
+ }
+
+ defer func() {
+ closeErr := os.RemoveAll(root)
+ if closeErr != nil {
+ err = errors.Join(err, fmt.Errorf("remove rebuild directory: %w", closeErr))
+ }
+ }()
+
+ memory, err := openMemory(ctx, root, time.Now, summarizer, log.New(io.Discard, "", 0))
+ if err != nil {
+ return "", fmt.Errorf("open rebuilt memory: %w", err)
+ }
+
+ client.LimitCalls(buildCallLimit)
+ run := BuildRun{
+ memory: memory,
+ client: client,
+ usagePath: client.UsagePath(),
+ }
+
+ appended, stop := run.Run(ctx, entries)
+ sizes, err := memory.buildModelLines()
+ if stop == "" {
+ err = errors.Join(err, memory.Settle(ctx), memory.rebuildComplete(), memory.Flush(ctx))
+ }
+
+ err = errors.Join(err, memory.Close())
+ records, usageErr := buildReadUsage(client.UsagePath())
+ err = errors.Join(err, usageErr, buildWriteReport(out, appended, stop, records, sizes))
+ if stop != "" {
+ err = errors.Join(err, fmt.Errorf("rebuild stopped: %s", stop))
+ }
+
+ if err != nil {
+ return "", fmt.Errorf("finish rebuilt memory: %w", err)
+ }
+
+ err = rebuildCopy(agent.dir, root)
+ if err != nil {
+ return "", fmt.Errorf("copy other agent files: %w", err)
+ }
+
+ usage, err := os.ReadFile(client.UsagePath())
+ if errors.Is(err, os.ErrNotExist) {
+ err = nil
+ }
+
+ if err != nil {
+ return "", fmt.Errorf("read rebuild usage: %w", err)
+ }
+
+ if len(usage) > 0 {
+ err = rebuildAddUsage(filepath.Join(root, "usage.jsonl"), usage)
+ if err != nil {
+ return "", fmt.Errorf("append rebuild usage: %w", err)
+ }
+ }
+
+ if ctx.Err() != nil {
+ return "", ctx.Err()
+ }
+
+ backup = agent.dir + "-backup-" + time.Now().Format("20060102-150405.000000000")
+ err = os.Rename(agent.dir, backup)
+ if err != nil {
+ return "", fmt.Errorf("back up agent: %w", err)
+ }
+
+ err = os.Rename(root, agent.dir)
+ if err != nil {
+ restoreErr := os.Rename(backup, agent.dir)
+ if restoreErr != nil {
+ err = errors.Join(err, fmt.Errorf("restore original agent: %w", restoreErr))
+ }
+
+ return "", fmt.Errorf("install rebuilt agent: %w", err)
+ }
+
+ return backup, nil
+}
+
+// rebuildIdle refuses an agent with unfinished subagents, including ones that a
+// daemon restart must resume. A missing work file lists no work.
+func rebuildIdle(root string) error {
+ data, err := os.ReadFile(filepath.Join(root, subagentWorkFileName))
+ if errors.Is(err, os.ErrNotExist) {
+ return nil
+ }
+
+ if err != nil {
+ return fmt.Errorf("read running subagents: %w", err)
+ }
+
+ var work SubagentWorkFile
+ err = json.Unmarshal(data, &work)
+ if err != nil {
+ return fmt.Errorf("decode running subagents: %w", err)
+ }
+
+ if len(work.Running) > 0 {
+ return errors.New("rebuild refused: subagents have unfinished work")
+ }
+
+ return nil
+}
+
+// rebuildReadLog reads and validates the day files without a write or an index
+// repair. Even an incomplete final line is an error, not a repair.
+func rebuildReadLog(ctx context.Context, root string) ([]DatedMessage, error) {
+ paths, err := filepath.Glob(filepath.Join(root, "main", "*.jsonl"))
+ if err != nil {
+ return nil, fmt.Errorf("list message files: %w", err)
+ }
+
+ slices.Sort(paths)
+ var entries []DatedMessage
+ for _, path := range paths {
+ var day uint32
+
+ day, err = journalDayNumber(strings.TrimSuffix(filepath.Base(path), ".jsonl"))
+ if err != nil {
+ return nil, fmt.Errorf("parse message day: %w", err)
+ }
+
+ err = journalScanFile(ctx, path, day, 0, false, func(line []byte, _ Position) error {
+ var message Message
+ err = json.Unmarshal(line, &message)
+ if err != nil {
+ return fmt.Errorf("decode message: %w", err)
+ }
+
+ err = logValidateMessage(message)
+ if err != nil {
+ return fmt.Errorf("validate message: %w", err)
+ }
+
+ if message.ID != len(entries) {
+ return fmt.Errorf("expected message %d, found %d", len(entries), message.ID)
+ }
+
+ entries = append(entries, DatedMessage{
+ Kind: message.Kind,
+ Text: message.Text,
+ Date: message.Date,
+ })
+
+ return nil
+ })
+ if err != nil {
+ return nil, fmt.Errorf("read message file: %w", err)
+ }
+ }
+
+ return entries, nil
+}
+
+// rebuildMerge preserves the input order for equal dates, with old messages
+// before added messages. Each date must be explicit and valid.
+func rebuildMerge(old, added []DatedMessage) ([]DatedMessage, error) {
+ entries := append(slices.Clone(old), added...)
+ dates := make(map[string]time.Time)
+ for _, entry := range entries {
+ err := logValidateMessage(Message{
+ Kind: entry.Kind,
+ Text: entry.Text,
+ Size: len(entry.Text),
+ Date: entry.Date,
+ })
+ if err != nil {
+ return nil, fmt.Errorf("validate dated message: %w", err)
+ }
+
+ var date time.Time
+ date, err = time.Parse(time.RFC3339Nano, entry.Date)
+ if err != nil {
+ return nil, fmt.Errorf("parse merge date: %w", err)
+ }
+
+ dates[entry.Date] = date
+ }
+
+ slices.SortStableFunc(entries, func(a, b DatedMessage) int {
+ return dates[a.Date].Compare(dates[b.Date])
+ })
+
+ return entries, nil
+}
+
+// rebuildCopy copies the other agent directories and regular files. It excludes
+// the memory and its lock, and refuses all other file types.
+func rebuildCopy(source, target string) error {
+ err := filepath.WalkDir(source, func(path string, entry fs.DirEntry, err error) error {
+ if err != nil {
+ return fmt.Errorf("read agent entry: %w", err)
+ }
+
+ relative, err := filepath.Rel(source, path)
+ if err != nil {
+ return fmt.Errorf("find agent relative path: %w", err)
+ }
+
+ if relative == "." {
+ return nil
+ }
+
+ switch relative {
+ case "main", "tree", "index":
+ return filepath.SkipDir
+ case "lock", "view.json":
+ return nil
+ }
+
+ destination := filepath.Join(target, relative)
+ switch {
+ case entry.IsDir():
+ err = os.Mkdir(destination, 0o700)
+ if err != nil {
+ return fmt.Errorf("create agent directory: %w", err)
+ }
+ case entry.Type().IsRegular():
+ var data []byte
+ data, err = os.ReadFile(path)
+ if err != nil {
+ return fmt.Errorf("read agent file: %w", err)
+ }
+
+ err = os.WriteFile(destination, data, 0o600)
+ if err != nil {
+ return fmt.Errorf("copy agent file: %w", err)
+ }
+ default:
+ return fmt.Errorf("copy agent: unsupported file type for %s", relative)
+ }
+
+ return nil
+ })
+ if err != nil {
+ return fmt.Errorf("copy agent files: %w", err)
+ }
+
+ return nil
+}
+
+// rebuildAddUsage keeps both the original usage and the calls of this run.
+func rebuildAddUsage(path string, data []byte) (err error) {
+ file, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
+ if err != nil {
+ return fmt.Errorf("open rebuilt usage: %w", err)
+ }
+
+ defer func() {
+ closeErr := file.Close()
+ if closeErr != nil {
+ err = errors.Join(err, fmt.Errorf("close rebuilt usage: %w", closeErr))
+ }
+ }()
+
+ _, err = file.Write(data)
+ if err != nil {
+ return fmt.Errorf("append rebuilt usage: %w", err)
+ }
+
+ return nil
+}
--- /dev/null
+package main
+
+import (
+ "fmt"
+ "io/fs"
+ "os"
+ "path/filepath"
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/require"
+)
+
+func newRebuildTestAgent(t *testing.T) *Agent {
+ t.Helper()
+
+ agent, err := openAgent(filepath.Join(t.TempDir(), "agent"))
+ require.NoError(t, err)
+
+ return agent
+}
+
+func rebuildTestLog(t *testing.T, root string, entries []DatedMessage) {
+ t.Helper()
+
+ memory, err := openMemory(t.Context(), root, time.Now, newFakeSummarizer(), discardLog())
+ require.NoError(t, err)
+
+ for _, entry := range entries {
+ _, err = memory.AppendDated(entry.Kind, entry.Text, entry.Date)
+ require.NoError(t, err)
+ }
+
+ require.NoError(t, memory.Wait(t.Context()))
+ require.NoError(t, memory.Close())
+}
+
+func rebuildTestSnapshot(t *testing.T, root string) map[string]string {
+ t.Helper()
+
+ files := make(map[string]string)
+ err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, err error) error {
+ if err != nil {
+ return fmt.Errorf("read snapshot entry: %w", err)
+ }
+
+ if entry.IsDir() {
+ return nil
+ }
+
+ relative, err := filepath.Rel(root, path)
+ if err != nil {
+ return fmt.Errorf("find snapshot relative path: %w", err)
+ }
+
+ data, err := os.ReadFile(path)
+ if err != nil {
+ return fmt.Errorf("read snapshot file: %w", err)
+ }
+
+ info, err := entry.Info()
+ if err != nil {
+ return fmt.Errorf("read snapshot mode: %w", err)
+ }
+
+ files[relative] = info.Mode().String() + string(data)
+
+ return nil
+ })
+ require.NoError(t, err)
+
+ return files
+}
--- /dev/null
+package main
+
+import (
+ "context"
+ "encoding/json"
+ "net/http"
+ "os"
+ "path/filepath"
+ "strings"
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/require"
+)
+
+func TestRebuildMergesDatesAndKeepsAgentFiles(t *testing.T) {
+ agent := newRebuildTestAgent(t)
+ old := []DatedMessage{
+ {
+ Kind: "note",
+ Text: "middle",
+ Date: "2025-03-02T12:00:00-03:00",
+ },
+ {
+ Kind: "user",
+ Text: "last",
+ Date: "2025-03-03T12:00:00Z",
+ },
+ }
+
+ rebuildTestLog(t, agent.dir, old)
+ require.NoError(t, os.MkdirAll(filepath.Join(agent.dir, "subagents", "BlueBird"), 0o700))
+
+ files := map[string]string{
+ "instructions.md": "Keep exact answers.",
+ "schedules.json": "[]\n",
+ "telegram-offset.json": "123\n",
+ "work.json": "{\"running\":[]}\n",
+ "usage.jsonl": "{\"model\":\"old\"}\n",
+ "subagents/BlueBird/report": "Finished.",
+ }
+
+ for name, text := range files {
+ require.NoError(t, os.WriteFile(filepath.Join(agent.dir, name), []byte(text), 0o600))
+ }
+
+ before := rebuildTestSnapshot(t, agent.dir)
+ server := newModelFake(t)
+ client := newModelClient(server.URL, ModelAPIKey(modelTestKey), agent.CacheKey(), filepath.Join(t.TempDir(), "usage.jsonl"))
+ var report strings.Builder
+ added := []DatedMessage{
+ {
+ Kind: "work",
+ Text: strings.Repeat("Earlier message. ", 80),
+ Date: "2025-03-01T10:00:00-03:00",
+ },
+ {
+ Kind: "note",
+ Text: "equal date comes after old",
+ Date: old[0].Date,
+ },
+ }
+
+ backup, err := rebuild(t.Context(), agent, added, client, newCompactor(client, agent.Prompt, defaultConfig().Compactor), &report)
+ require.NoError(t, err)
+ require.Equal(t, before, rebuildTestSnapshot(t, backup))
+
+ entries, err := rebuildReadLog(t.Context(), agent.dir)
+ require.NoError(t, err)
+ require.Equal(t, []DatedMessage{added[0], old[0], added[1], old[1]}, entries)
+
+ for name, text := range files {
+ var data []byte
+
+ data, err = os.ReadFile(filepath.Join(agent.dir, name))
+ require.NoError(t, err)
+
+ if name == "usage.jsonl" {
+ require.True(t, strings.HasPrefix(string(data), text))
+
+ continue
+ }
+
+ require.Equal(t, text, string(data))
+ }
+
+ require.Positive(t, server.Calls())
+
+ call := server.Call(0)
+ require.Equal(t, agent.CacheKey(), call.body["prompt_cache_key"])
+
+ prompt, err := agent.Prompt()
+ require.NoError(t, err)
+ require.Contains(t, string(call.raw), strings.Trim(string(rebuildTestJSON(t, prompt)), "\""))
+ require.Contains(t, report.String(), "cached (")
+
+ memory, err := openMemory(t.Context(), agent.dir, time.Now, newFakeSummarizer(), discardLog())
+ require.NoError(t, err)
+ require.NoError(t, memory.Settle(t.Context()))
+ require.NoError(t, memory.Close())
+}
+
+func TestRebuildFailureAndCancellationKeepOriginal(t *testing.T) {
+ testCases := []struct {
+ name string
+ cancel bool
+ }{
+ {
+ name: "failed call keeps the original",
+ },
+ {
+ name: "cancelled run keeps the original",
+ cancel: true,
+ },
+ }
+
+ for _, tc := range testCases {
+ t.Run(tc.name, func(t *testing.T) {
+ agent := newRebuildTestAgent(t)
+ rebuildTestLog(t, agent.dir, []DatedMessage{{
+ Kind: "user",
+ Text: "Keep this.",
+ Date: "2026-01-01T00:00:00Z",
+ }})
+
+ before := rebuildTestSnapshot(t, agent.dir)
+ server := newModelFake(t)
+ server.reply = func(w http.ResponseWriter, _ ModelFakeCall) {
+ w.WriteHeader(http.StatusBadRequest)
+ }
+
+ client := newModelClient(server.URL, ModelAPIKey(modelTestKey), agent.CacheKey(), filepath.Join(t.TempDir(), "usage.jsonl"))
+ ctx, cancel := context.WithCancel(t.Context())
+ defer cancel()
+
+ if tc.cancel {
+ cancel()
+ }
+
+ backup, err := rebuild(ctx, agent, []DatedMessage{{
+ Kind: "work",
+ Text: strings.Repeat("long ", 200),
+ Date: "2025-01-01T00:00:00Z",
+ }}, client, newCompactor(client, agent.Prompt, defaultConfig().Compactor), &strings.Builder{})
+ require.Error(t, err)
+ require.Empty(t, backup)
+ require.Equal(t, before, rebuildTestSnapshot(t, agent.dir))
+
+ pending, err := filepath.Glob(filepath.Join(filepath.Dir(agent.dir), ".agent-rebuild-*"))
+ require.NoError(t, err)
+ require.Empty(t, pending)
+ })
+ }
+}
+
+func TestRebuildKeepsOriginalWhenAnInvisibleMergeFails(t *testing.T) {
+ agent := newRebuildTestAgent(t)
+ old := []DatedMessage{
+ {
+ Kind: "user",
+ Text: strings.Repeat("a", 300),
+ Date: "2025-01-01T00:00:00Z",
+ },
+ {
+ Kind: "reply",
+ Text: strings.Repeat("b", 300),
+ Date: "2025-01-01T00:01:00Z",
+ },
+ }
+
+ rebuildTestLog(t, agent.dir, old)
+ before := rebuildTestSnapshot(t, agent.dir)
+ server := newModelFake(t)
+ server.reply = func(w http.ResponseWriter, _ ModelFakeCall) {
+ modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("")}, 100, 0, 5))
+ }
+
+ client := newModelClient(server.URL, ModelAPIKey(modelTestKey), agent.CacheKey(), filepath.Join(t.TempDir(), "usage.jsonl"))
+ _, err := rebuild(t.Context(), agent, nil, client, newCompactor(client, agent.Prompt, defaultConfig().Compactor), &strings.Builder{})
+ require.ErrorContains(t, err, "summary tree has unfinished nodes")
+ require.Equal(t, before, rebuildTestSnapshot(t, agent.dir))
+}
+
+func TestRebuildRefusesUnfinishedSubagents(t *testing.T) {
+ agent := newRebuildTestAgent(t)
+ require.NoError(t, os.WriteFile(filepath.Join(agent.dir, "work.json"), []byte(`{"running":["BlueBird"]}`), 0o600))
+
+ require.ErrorContains(t, rebuildIdle(agent.dir), "unfinished work")
+}
+
+func TestRebuildReadDoesNotRepairIncompleteLog(t *testing.T) {
+ agent := newRebuildTestAgent(t)
+ require.NoError(t, os.Mkdir(filepath.Join(agent.dir, "main"), 0o700))
+
+ path := filepath.Join(agent.dir, "main", "2026-01-01.jsonl")
+ require.NoError(t, os.WriteFile(path, []byte("incomplete"), 0o600))
+
+ _, err := rebuildReadLog(t.Context(), agent.dir)
+ require.ErrorContains(t, err, "incomplete line")
+
+ data, err := os.ReadFile(path)
+ require.NoError(t, err)
+ require.Equal(t, "incomplete", string(data))
+}
+
+func TestRebuildReadStopsWhenTheContextEnds(t *testing.T) {
+ agent := newRebuildTestAgent(t)
+ rebuildTestLog(t, agent.dir, []DatedMessage{{
+ Kind: "note",
+ Text: "x",
+ Date: "2026-01-01T00:00:00Z",
+ }})
+
+ ctx, cancel := context.WithCancel(t.Context())
+ defer cancel()
+
+ cancel()
+
+ _, err := rebuildReadLog(ctx, agent.dir)
+ require.ErrorIs(t, err, context.Canceled)
+}
+
+func rebuildTestJSON(t *testing.T, value string) []byte {
+ t.Helper()
+
+ data, err := json.Marshal(value)
+ require.NoError(t, err)
+
+ return data
+}
projects / onechat / commitdiff