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