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