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