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

   1 package main
   2
   3 import (
   4         "context"
   5         "errors"
   6         "fmt"
   7         "io/fs"
   8         "os"
   9         "path/filepath"
  10         "strconv"
  11         "strings"
  12         "sync"
  13         "time"
  14 )
  15
  16 const (
  17         // subagentDirName is the directory of the subagents in the directory of
  18         // an agent.
  19         subagentDirName = "subagents"
  20
  21         // subagentNameDraws is the most names that Create draws before it
  22         // fails.
  23         subagentNameDraws = 1000
  24 )
  25
  26 // newSubagents returns the subagents in dir, one directory for each name. The
  27 // clock gives the dates of their messages, and draw gives the names of new
  28 // subagents.
  29 func newSubagents(dir string, clock func() time.Time, draw func() string) *Subagents {
  30         return &Subagents{
  31                 dir:      dir,
  32                 clock:    clock,
  33                 draw:     draw,
  34                 open:     make(map[string]*Subagent),
  35                 reserved: make(map[string]bool),
  36         }
  37 }
  38
  39 // Subagents holds the subagents of one agent. Each subagent has a directory
  40 // named for it, which holds its log. A log takes an exclusive lock, so
  41 // Subagents opens the log of a name once and keeps it open until Close. Its
  42 // mutex protects the open logs.
  43 type Subagents struct {
  44         dir   string
  45         clock func() time.Time
  46         draw  func() string
  47
  48         mutex    sync.Mutex
  49         open     map[string]*Subagent
  50         reserved map[string]bool
  51 }
  52
  53 // Create makes the directory and the log of a new subagent. It draws a name
  54 // again when the directory of the name exists.
  55 func (s *Subagents) Create(ctx context.Context) (*Subagent, error) {
  56         err := os.MkdirAll(s.dir, 0o700)
  57         if err != nil {
  58                 return nil, fmt.Errorf("create subagent directory: %w", err)
  59         }
  60
  61         for range subagentNameDraws {
  62                 name := s.draw()
  63                 s.mutex.Lock()
  64
  65                 reserved := s.reserved[name]
  66                 s.mutex.Unlock()
  67
  68                 if reserved {
  69                         continue
  70                 }
  71
  72                 err = os.Mkdir(filepath.Join(s.dir, name), 0o700)
  73                 if errors.Is(err, fs.ErrExist) {
  74                         continue
  75                 }
  76
  77                 if err != nil {
  78                         return nil, fmt.Errorf("create subagent directory: %w", err)
  79                 }
  80
  81                 return s.Get(ctx, name)
  82         }
  83
  84         return nil, fmt.Errorf("no free subagent name after %d draws", subagentNameDraws)
  85 }
  86
  87 // Reserve keeps a child name out of future subagent draws. An existing subagent
  88 // directory or a reserved name is refused.
  89 func (s *Subagents) Reserve(name string) error {
  90         s.mutex.Lock()
  91         defer s.mutex.Unlock()
  92
  93         if s.reserved[name] {
  94                 return fmt.Errorf("name %q is already used", name)
  95         }
  96
  97         _, err := os.Lstat(filepath.Join(s.dir, name))
  98         if err == nil {
  99                 return fmt.Errorf("name %q is already used by a subagent", name)
 100         }
 101
 102         if !errors.Is(err, os.ErrNotExist) {
 103                 return fmt.Errorf("check subagent name: %w", err)
 104         }
 105
 106         s.reserved[name] = true
 107
 108         return nil
 109 }
 110
 111 // Release frees a reservation after a child fails to publish. It does not
 112 // remove a subagent directory or change an open log.
 113 func (s *Subagents) Release(name string) {
 114         s.mutex.Lock()
 115         defer s.mutex.Unlock()
 116
 117         delete(s.reserved, name)
 118 }
 119
 120 // Get returns the subagent name with its log open. A name that is not valid or
 121 // that has no directory is an unknown agent.
 122 func (s *Subagents) Get(ctx context.Context, name string) (*Subagent, error) {
 123         s.mutex.Lock()
 124         defer s.mutex.Unlock()
 125
 126         subagent, found := s.open[name]
 127         if found {
 128                 return subagent, nil
 129         }
 130
 131         unknown := fmt.Errorf("no agent %s", name)
 132         if !subagentNameValid(name) {
 133                 return nil, unknown
 134         }
 135
 136         root := filepath.Join(s.dir, name)
 137         info, err := os.Stat(root)
 138         if errors.Is(err, fs.ErrNotExist) {
 139                 return nil, unknown
 140         }
 141
 142         if err != nil {
 143                 return nil, fmt.Errorf("read subagent directory: %w", err)
 144         }
 145
 146         if !info.IsDir() {
 147                 return nil, unknown
 148         }
 149
 150         log, err := openLog(ctx, root, s.clock)
 151         if err != nil {
 152                 return nil, fmt.Errorf("open log of %s: %w", name, err)
 153         }
 154
 155         syncer, err := openSyncer(root)
 156         if err != nil {
 157                 return nil, errors.Join(fmt.Errorf("open log of %s: %w", name, err), log.Close())
 158         }
 159
 160         subagent = &Subagent{
 161                 name:   name,
 162                 dir:    root,
 163                 log:    log,
 164                 syncer: syncer,
 165         }
 166         s.open[name] = subagent
 167
 168         return subagent, nil
 169 }
 170
 171 // Snapshot flushes and locks the open logs while snapshot reads the subagent
 172 // files. snapshot must not call Subagents or an open Subagent.
 173 func (s *Subagents) Snapshot(ctx context.Context, snapshot func() error) error {
 174         s.mutex.Lock()
 175         defer s.mutex.Unlock()
 176
 177         for _, subagent := range s.open {
 178                 subagent.mutex.Lock()
 179                 defer subagent.mutex.Unlock()
 180
 181                 err := subagent.Flush(ctx)
 182                 if err != nil {
 183                         return fmt.Errorf("flush subagent snapshot: %w", err)
 184                 }
 185         }
 186
 187         err := snapshot()
 188         if err != nil {
 189                 return fmt.Errorf("copy subagent snapshot: %w", err)
 190         }
 191
 192         return nil
 193 }
 194
 195 // Close makes the writes of each open log durable and closes it.
 196 func (s *Subagents) Close() error {
 197         s.mutex.Lock()
 198         defer s.mutex.Unlock()
 199
 200         var err error
 201
 202         for name, subagent := range s.open {
 203                 err = errors.Join(err, subagent.close())
 204
 205                 delete(s.open, name)
 206         }
 207
 208         return err
 209 }
 210
 211 // Subagent is one subagent, its directory, and its log. Its mutex protects the
 212 // log, because the run of the subagent writes it while the agent reads it with
 213 // zoom.
 214 type Subagent struct {
 215         name   string
 216         dir    string
 217         mutex  sync.Mutex
 218         log    *Log
 219         syncer *Syncer
 220 }
 221
 222 // Append logs a message of the subagent.
 223 func (a *Subagent) Append(kind, text string) (Message, error) {
 224         a.mutex.Lock()
 225         defer a.mutex.Unlock()
 226
 227         message, err := a.log.Append(kind, text)
 228         if err != nil {
 229                 return Message{}, err
 230         }
 231
 232         a.syncer.Written()
 233
 234         return message, nil
 235 }
 236
 237 // Flush returns when every write before the call is durable.
 238 func (a *Subagent) Flush(ctx context.Context) error {
 239         return a.syncer.Flush(ctx)
 240 }
 241
 242 // First returns the first message of the log.
 243 func (a *Subagent) First() (Message, error) {
 244         a.mutex.Lock()
 245         defer a.mutex.Unlock()
 246
 247         return a.log.Message(0)
 248 }
 249
 250 // Show returns the whole log, each message as "i|kind: text", with a blank line
 251 // between messages.
 252 func (a *Subagent) Show() (string, error) {
 253         a.mutex.Lock()
 254         defer a.mutex.Unlock()
 255
 256         lines := make([]string, 0, a.log.Len())
 257         for id := range a.log.Len() {
 258                 message, err := a.log.Message(id)
 259                 if err != nil {
 260                         return "", err
 261                 }
 262
 263                 lines = append(lines, strconv.Itoa(id)+"|"+leafSource(message))
 264         }
 265
 266         return strings.Join(lines, "\n\n"), nil
 267 }
 268
 269 // WorkDir returns the directory where the tools of the subagent run: work in
 270 // its directory.
 271 func (a *Subagent) WorkDir() string {
 272         return filepath.Join(a.dir, "work")
 273 }
 274
 275 // close makes the writes durable and closes the log.
 276 func (a *Subagent) close() error {
 277         a.mutex.Lock()
 278         defer a.mutex.Unlock()
 279
 280         return errors.Join(a.syncer.Close(), a.log.Close())
 281 }