git.giacomolli.uk - onechat/blob - rebuild.go

   1 package main
   2
   3 import (
   4         "context"
   5         "encoding/json"
   6         "errors"
   7         "fmt"
   8         "io"
   9         "io/fs"
  10         "log/slog"
  11         "os"
  12         "path/filepath"
  13         "slices"
  14         "strings"
  15         "time"
  16 )
  17
  18 // rebuildCallLimit is the most model calls of a rebuild. A rebuild summarizes
  19 // the whole merged log, which can need more calls than a probe sample, and the
  20 // quota stop of the build run still guards the account.
  21 const rebuildCallLimit = 1500
  22
  23 // rebuild merges added messages with the log of agent, in date order, and
  24 // builds a new memory beside it. The caller holds the token lock. The model
  25 // uses the original agent prompt and cache key. A failed or stopped build
  26 // removes the new directory and leaves the original directory unchanged.
  27 // Success copies the other agent files, adds this run's usage, and swaps the
  28 // directories, with the original kept as a timestamped backup. Message IDs are
  29 // new positions in the merged log, but kind, text, and date stay exact. A child
  30 // keeps its staging directory and backup outside children so that neither is
  31 // read as a configured child.
  32 func rebuild(ctx context.Context, agent *Agent, added []DatedMessage, client *ModelClient, summarizer Summarizer, out io.Writer) (backup string, err error) {
  33         err = rebuildIdle(agent.dir)
  34         if err != nil {
  35                 return "", fmt.Errorf("check rebuild work: %w", err)
  36         }
  37
  38         old, err := rebuildReadLog(ctx, agent.dir)
  39         if err != nil {
  40                 return "", fmt.Errorf("read original log: %w", err)
  41         }
  42
  43         entries, err := rebuildMerge(old, added)
  44         if err != nil {
  45                 return "", fmt.Errorf("merge rebuild messages: %w", err)
  46         }
  47
  48         parent := filepath.Dir(agent.dir)
  49         backupDir := agent.dir
  50         if agent.config.Parent != "" {
  51                 parent = filepath.Dir(parent)
  52                 backupDir = filepath.Join(parent, agent.config.Name)
  53         }
  54
  55         root, err := os.MkdirTemp(parent, ".agent-rebuild-")
  56         if err != nil {
  57                 return "", fmt.Errorf("create rebuild directory: %w", err)
  58         }
  59
  60         defer func() {
  61                 closeErr := os.RemoveAll(root)
  62                 if closeErr != nil {
  63                         err = errors.Join(err, fmt.Errorf("remove rebuild directory: %w", closeErr))
  64                 }
  65         }()
  66
  67         memory, err := openMemory(ctx, root, time.Now, summarizer, slog.New(slog.DiscardHandler))
  68         if err != nil {
  69                 return "", fmt.Errorf("open rebuilt memory: %w", err)
  70         }
  71
  72         client.LimitCalls(rebuildCallLimit)
  73
  74         run := BuildRun{
  75                 memory:    memory,
  76                 client:    client,
  77                 usagePath: client.UsagePath(),
  78         }
  79
  80         appended, stop := run.Run(ctx, entries)
  81         sizes, err := memory.buildModelLines()
  82         if stop == "" {
  83                 err = errors.Join(err, memory.Settle(ctx), memory.rebuildComplete(), memory.Flush(ctx))
  84         }
  85
  86         err = errors.Join(err, memory.Close())
  87         records, usageErr := buildReadUsage(client.UsagePath())
  88         err = errors.Join(err, usageErr, buildWriteReport(out, appended, stop, records, sizes))
  89         if stop != "" {
  90                 err = errors.Join(err, fmt.Errorf("rebuild stopped: %s", stop))
  91         }
  92
  93         if err != nil {
  94                 return "", fmt.Errorf("finish rebuilt memory: %w", err)
  95         }
  96
  97         err = rebuildCopy(agent.dir, root)
  98         if err != nil {
  99                 return "", fmt.Errorf("copy other agent files: %w", err)
 100         }
 101
 102         usage, err := os.ReadFile(client.UsagePath())
 103         if errors.Is(err, os.ErrNotExist) {
 104                 err = nil
 105         }
 106
 107         if err != nil {
 108                 return "", fmt.Errorf("read rebuild usage: %w", err)
 109         }
 110
 111         if len(usage) > 0 {
 112                 err = rebuildAddUsage(filepath.Join(root, "usage.jsonl"), usage)
 113                 if err != nil {
 114                         return "", fmt.Errorf("append rebuild usage: %w", err)
 115                 }
 116         }
 117
 118         if ctx.Err() != nil {
 119                 return "", ctx.Err()
 120         }
 121
 122         backup = backupDir + "-backup-" + time.Now().Format("20060102-150405.000000000")
 123         err = os.Rename(agent.dir, backup)
 124         if err != nil {
 125                 return "", fmt.Errorf("back up agent: %w", err)
 126         }
 127
 128         err = os.Rename(root, agent.dir)
 129         if err != nil {
 130                 restoreErr := os.Rename(backup, agent.dir)
 131                 if restoreErr != nil {
 132                         err = errors.Join(err, fmt.Errorf("restore original agent: %w", restoreErr))
 133                 }
 134
 135                 return "", fmt.Errorf("install rebuilt agent: %w", err)
 136         }
 137
 138         return backup, nil
 139 }
 140
 141 // rebuildIdle refuses an agent with unfinished subagents, including ones that a
 142 // daemon restart must resume. A missing work file lists no work.
 143 func rebuildIdle(root string) error {
 144         data, err := os.ReadFile(filepath.Join(root, subagentWorkFileName))
 145         if errors.Is(err, os.ErrNotExist) {
 146                 return nil
 147         }
 148
 149         if err != nil {
 150                 return fmt.Errorf("read running subagents: %w", err)
 151         }
 152
 153         var work SubagentWorkFile
 154
 155         err = json.Unmarshal(data, &work)
 156         if err != nil {
 157                 return fmt.Errorf("decode running subagents: %w", err)
 158         }
 159
 160         if len(work.Running) > 0 {
 161                 return errors.New("rebuild refused: subagents have unfinished work")
 162         }
 163
 164         return nil
 165 }
 166
 167 // rebuildReadLog reads and validates the day files without a write or an index
 168 // repair. Even an incomplete final line is an error, not a repair.
 169 func rebuildReadLog(ctx context.Context, root string) ([]DatedMessage, error) {
 170         paths, err := filepath.Glob(filepath.Join(root, "main", "*.jsonl"))
 171         if err != nil {
 172                 return nil, fmt.Errorf("list message files: %w", err)
 173         }
 174
 175         slices.Sort(paths)
 176
 177         var entries []DatedMessage
 178
 179         for _, path := range paths {
 180                 var day uint32
 181
 182                 day, err = journalDayNumber(strings.TrimSuffix(filepath.Base(path), ".jsonl"))
 183                 if err != nil {
 184                         return nil, fmt.Errorf("parse message day: %w", err)
 185                 }
 186
 187                 err = journalScanFile(ctx, path, day, 0, false, func(line []byte, _ Position) error {
 188                         var message Message
 189
 190                         err = json.Unmarshal(line, &message)
 191                         if err != nil {
 192                                 return fmt.Errorf("decode message: %w", err)
 193                         }
 194
 195                         err = logValidateMessage(message)
 196                         if err != nil {
 197                                 return fmt.Errorf("validate message: %w", err)
 198                         }
 199
 200                         if message.ID != len(entries) {
 201                                 return fmt.Errorf("expected message %d, found %d", len(entries), message.ID)
 202                         }
 203
 204                         entries = append(entries, DatedMessage{
 205                                 Kind: message.Kind,
 206                                 Text: message.Text,
 207                                 Date: message.Date,
 208                         })
 209
 210                         return nil
 211                 })
 212                 if err != nil {
 213                         return nil, fmt.Errorf("read message file: %w", err)
 214                 }
 215         }
 216
 217         return entries, nil
 218 }
 219
 220 // rebuildMerge preserves the input order for equal dates, with old messages
 221 // before added messages. Each date must be explicit and valid.
 222 func rebuildMerge(old, added []DatedMessage) ([]DatedMessage, error) {
 223         entries := append(slices.Clone(old), added...)
 224         dates := make(map[string]time.Time)
 225         for _, entry := range entries {
 226                 err := logValidateMessage(Message{
 227                         Kind: entry.Kind,
 228                         Text: entry.Text,
 229                         Size: len(entry.Text),
 230                         Date: entry.Date,
 231                 })
 232                 if err != nil {
 233                         return nil, fmt.Errorf("validate dated message: %w", err)
 234                 }
 235
 236                 var date time.Time
 237
 238                 date, err = time.Parse(time.RFC3339Nano, entry.Date)
 239                 if err != nil {
 240                         return nil, fmt.Errorf("parse merge date: %w", err)
 241                 }
 242
 243                 dates[entry.Date] = date
 244         }
 245
 246         slices.SortStableFunc(entries, func(a, b DatedMessage) int {
 247                 return dates[a.Date].Compare(dates[b.Date])
 248         })
 249
 250         return entries, nil
 251 }
 252
 253 // rebuildCopy copies the other agent directories and regular files. It excludes
 254 // the memory and its lock, and refuses all other file types.
 255 func rebuildCopy(source, target string) error {
 256         err := filepath.WalkDir(source, func(path string, entry fs.DirEntry, err error) error {
 257                 if err != nil {
 258                         return fmt.Errorf("read agent entry: %w", err)
 259                 }
 260
 261                 relative, err := filepath.Rel(source, path)
 262                 if err != nil {
 263                         return fmt.Errorf("find agent relative path: %w", err)
 264                 }
 265
 266                 if relative == "." {
 267                         return nil
 268                 }
 269
 270                 switch relative {
 271                 case "main", "tree", "index":
 272                         return filepath.SkipDir
 273                 case "lock", "view.json":
 274                         return nil
 275                 }
 276
 277                 destination := filepath.Join(target, relative)
 278                 switch {
 279                 case entry.IsDir():
 280                         err = os.Mkdir(destination, 0o700)
 281                         if err != nil {
 282                                 return fmt.Errorf("create agent directory: %w", err)
 283                         }
 284                 case entry.Type().IsRegular():
 285                         var data []byte
 286
 287                         data, err = os.ReadFile(path)
 288                         if err != nil {
 289                                 return fmt.Errorf("read agent file: %w", err)
 290                         }
 291
 292                         err = os.WriteFile(destination, data, 0o600)
 293                         if err != nil {
 294                                 return fmt.Errorf("copy agent file: %w", err)
 295                         }
 296                 default:
 297                         return fmt.Errorf("copy agent: unsupported file type for %s", relative)
 298                 }
 299
 300                 return nil
 301         })
 302         if err != nil {
 303                 return fmt.Errorf("copy agent files: %w", err)
 304         }
 305
 306         return nil
 307 }
 308
 309 // rebuildAddUsage keeps both the original usage and the calls of this run.
 310 func rebuildAddUsage(path string, data []byte) (err error) {
 311         file, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
 312         if err != nil {
 313                 return fmt.Errorf("open rebuilt usage: %w", err)
 314         }
 315
 316         defer func() {
 317                 closeErr := file.Close()
 318                 if closeErr != nil {
 319                         err = errors.Join(err, fmt.Errorf("close rebuilt usage: %w", closeErr))
 320                 }
 321         }()
 322
 323         _, err = file.Write(data)
 324         if err != nil {
 325                 return fmt.Errorf("append rebuilt usage: %w", err)
 326         }
 327
 328         return nil
 329 }