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 // Reserve keeps a child name out of future subagent draws. An existing subagent
  99 // directory or a reserved name is refused.
 100 func (s *Subagents) Reserve(name string) error {
 101         s.mutex.Lock()
 102         defer s.mutex.Unlock()
 103
 104         if s.reserved[name] {
 105                 return fmt.Errorf("name %q is already used", name)
 106         }
 107
 108         _, err := os.Lstat(filepath.Join(s.dir, name))
 109         if err == nil {
 110                 return fmt.Errorf("name %q is already used by a subagent", name)
 111         }
 112
 113         if !errors.Is(err, os.ErrNotExist) {
 114                 return fmt.Errorf("check subagent name: %w", err)
 115         }
 116
 117         s.reserved[name] = true
 118
 119         return nil
 120 }
 121
 122 // Release frees a reservation after a child fails to publish. It does not
 123 // remove a subagent directory or change an open log.
 124 func (s *Subagents) Release(name string) {
 125         s.mutex.Lock()
 126         defer s.mutex.Unlock()
 127
 128         delete(s.reserved, name)
 129 }
 130
 131 // Get returns the subagent name with its log open. A name that is not valid or
 132 // that has no directory is an unknown agent.
 133 func (s *Subagents) Get(ctx context.Context, name string) (*Subagent, error) {
 134         s.mutex.Lock()
 135         defer s.mutex.Unlock()
 136
 137         subagent, found := s.open[name]
 138         if found {
 139                 return subagent, nil
 140         }
 141
 142         unknown := fmt.Errorf("no agent %s", name)
 143         if !subagentNameValid(name) {
 144                 return nil, unknown
 145         }
 146
 147         root := filepath.Join(s.dir, name)
 148         info, err := os.Stat(root)
 149         if errors.Is(err, fs.ErrNotExist) {
 150                 return nil, unknown
 151         }
 152
 153         if err != nil {
 154                 return nil, fmt.Errorf("read subagent directory: %w", err)
 155         }
 156
 157         if !info.IsDir() {
 158                 return nil, unknown
 159         }
 160
 161         log, err := openLog(ctx, root, s.clock)
 162         if err != nil {
 163                 return nil, fmt.Errorf("open log of %s: %w", name, err)
 164         }
 165
 166         syncer, err := openSyncer(root)
 167         if err != nil {
 168                 return nil, errors.Join(fmt.Errorf("open log of %s: %w", name, err), log.Close())
 169         }
 170
 171         subagent = &Subagent{
 172                 name:   name,
 173                 dir:    root,
 174                 log:    log,
 175                 syncer: syncer,
 176         }
 177
 178         s.open[name] = subagent
 179
 180         return subagent, nil
 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 // Append logs a message of the subagent.
 224 func (a *Subagent) Append(kind, text string) (Message, error) {
 225         a.mutex.Lock()
 226         defer a.mutex.Unlock()
 227
 228         message, err := a.log.Append(kind, text)
 229         if err != nil {
 230                 return Message{}, err
 231         }
 232
 233         a.syncer.Written()
 234
 235         return message, nil
 236 }
 237
 238 // Flush returns when every write before the call is durable.
 239 func (a *Subagent) Flush(ctx context.Context) error {
 240         return a.syncer.Flush(ctx)
 241 }
 242
 243 // First returns the first message of the log.
 244 func (a *Subagent) First() (Message, error) {
 245         a.mutex.Lock()
 246         defer a.mutex.Unlock()
 247
 248         return a.log.Message(0)
 249 }
 250
 251 // Show returns the whole log, each message as "i|kind: text", with a blank line
 252 // between messages.
 253 func (a *Subagent) Show() (string, error) {
 254         a.mutex.Lock()
 255         defer a.mutex.Unlock()
 256
 257         lines := make([]string, 0, a.log.Len())
 258         for id := range a.log.Len() {
 259                 message, err := a.log.Message(id)
 260                 if err != nil {
 261                         return "", err
 262                 }
 263
 264                 lines = append(lines, strconv.Itoa(id)+"|"+leafSource(message))
 265         }
 266
 267         return strings.Join(lines, "\n\n"), nil
 268 }
 269
 270 // WorkDir returns the directory where the tools of the subagent run: work in
 271 // its directory.
 272 func (a *Subagent) WorkDir() string {
 273         return filepath.Join(a.dir, "work")
 274 }
 275
 276 // close makes the writes durable and closes the log.
 277 func (a *Subagent) close() error {
 278         a.mutex.Lock()
 279         defer a.mutex.Unlock()
 280
 281         return errors.Join(a.syncer.Close(), a.log.Close())
 282 }