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