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