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