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

   1 package main
   2
   3 import (
   4         "context"
   5         "encoding/json"
   6         "errors"
   7         "fmt"
   8         "log/slog"
   9         "maps"
  10         "os"
  11         "path/filepath"
  12         "slices"
  13         "strings"
  14         "sync"
  15         "time"
  16 )
  17
  18 const (
  19         // subagentNote is the first message of the log of a subagent.
  20         subagentNote = "subagent"
  21
  22         // subagentNoteEffort follows subagentNote when the subagent has its own
  23         // effort.
  24         subagentNoteEffort = " at effort "
  25
  26         // subagentWorkFileName is the file of the names of the running
  27         // subagents in the directory of an agent.
  28         subagentWorkFileName = "work.json"
  29 )
  30
  31 // The statuses of the end of a run, in the work event of its report.
  32 const (
  33         subagentFinished = "finished"
  34         subagentStopped  = "stopped"
  35         subagentFailed   = "failed"
  36 )
  37
  38 // errSubagentClosed tells that the runs of the subagents ended with the loop.
  39 var errSubagentClosed = errors.New("the daemon stops, so no subagent starts")
  40
  41 // SubagentRuns runs the subagents of one agent in the background. Each run is a
  42 // goroutine with its own context, which only Stop and Close end. A run is one
  43 // or more request chains. The first request of each chain holds the view of the
  44 // agent at the start of the run, then the state of the turn, then the whole log
  45 // of the subagent. A message of Tell that comes during a chain joins its next
  46 // request after the tool results. One that comes after the final answer starts
  47 // one more chain, which reads the log again. The report of the run goes to
  48 // report: the final answer of each chain, or the error of a failed call. The
  49 // logger gets one line at the end of a run: done, failed with the cause, or
  50 // stopped. Stop ends one run, and its report says that it stopped. At the stop
  51 // of the daemon the runs end with no report.
  52 //
  53 // work.json in the directory of the agent holds the names of the running
  54 // subagents. Each start and each end of a run writes it whole, but Close does
  55 // not, so after a stop of the daemon it names the runs that the stop cut, and
  56 // Resume goes on with them. The logs alone cannot tell this, because a stop
  57 // can cut a run just after its final reply.
  58 //
  59 // Spawn and Tell split the view with the layout of the agent, which has no
  60 // lock. They are tools, so they run in the goroutine of the turns, as each
  61 // split of a turn does.
  62 type SubagentRuns struct {
  63         agent     *Agent
  64         memory    *Memory
  65         subagents *Subagents
  66         layout    *ViewLayout
  67         client    ChainClient
  68         state     func() string
  69         report    func(name, text, status string)
  70         note      func(text string) error
  71         fail      func(err error)
  72         logger    *slog.Logger
  73         notify    func(Event)
  74
  75         mutex   sync.Mutex
  76         running map[string]*SubagentRun
  77         closed  bool
  78         workers sync.WaitGroup
  79 }
  80
  81 // Spawn starts one subagent for each task and returns their names at once. It
  82 // waits until the view has no pending summary, and gives the view to each
  83 // subagent. A nil effort is the effort of the agent.
  84 func (r *SubagentRuns) Spawn(ctx context.Context, tasks []string, effort *string) (string, error) {
  85         if len(tasks) == 0 {
  86                 return "", errors.New("no tasks")
  87         }
  88
  89         r.mutex.Lock()
  90
  91         closed := r.closed
  92         r.mutex.Unlock()
  93
  94         if closed {
  95                 return "", errSubagentClosed
  96         }
  97
  98         note := subagentNote
  99         runEffort := r.agent.config.Effort
 100         if effort != nil && !slices.Contains(promptEfforts, *effort) {
 101                 return "", fmt.Errorf("no effort %q: one of %s", *effort, strings.Join(promptEfforts, ", "))
 102         }
 103
 104         if effort != nil && *effort != runEffort {
 105                 note += subagentNoteEffort + *effort
 106                 runEffort = *effort
 107         }
 108
 109         err := r.memory.Settle(ctx)
 110         if err != nil {
 111                 return "", fmt.Errorf("wait for the summaries: %w", err)
 112         }
 113
 114         view, err := r.layout.Split(r.memory.ViewLines(nil))
 115         if err != nil {
 116                 return "", err
 117         }
 118
 119         names := make([]string, 0, len(tasks))
 120         for _, task := range tasks {
 121                 name, err := r.create(ctx, note, task, view, runEffort)
 122                 if err != nil && len(names) > 0 {
 123                         return "", fmt.Errorf("started %s, then: %w", strings.Join(names, ", "), err)
 124                 }
 125
 126                 if err != nil {
 127                         return "", err
 128                 }
 129
 130                 names = append(names, name)
 131         }
 132
 133         return "Started " + strings.Join(names, ", ") + ". Each one's report will reach you as a message when it finishes.", nil
 134 }
 135
 136 // Tell sends message to the subagent name. A running subagent reads it between
 137 // its tool calls, and its report answers it. A finished subagent goes on from
 138 // its log with a new run, with the view of now and the effort of its note.
 139 func (r *SubagentRuns) Tell(ctx context.Context, name, message string) (string, error) {
 140         r.mutex.Lock()
 141
 142         run, found := r.running[name]
 143         if found {
 144                 defer r.mutex.Unlock()
 145
 146                 told, err := run.subagent.Append("user", message)
 147                 if err != nil {
 148                         return "", err
 149                 }
 150
 151                 run.pending = append(run.pending, told)
 152
 153                 return "Sent. " + name + " reads it between its tool calls, and its report answers it.", nil
 154         }
 155
 156         closed := r.closed
 157         r.mutex.Unlock()
 158
 159         if closed {
 160                 return "", errSubagentClosed
 161         }
 162
 163         subagent, err := r.subagents.Get(ctx, name)
 164         if err != nil {
 165                 return "", err
 166         }
 167
 168         effort, err := r.effort(subagent)
 169         if err != nil {
 170                 return "", err
 171         }
 172
 173         err = r.memory.Settle(ctx)
 174         if err != nil {
 175                 return "", fmt.Errorf("wait for the summaries: %w", err)
 176         }
 177
 178         view, err := r.layout.Split(r.memory.ViewLines(nil))
 179         if err != nil {
 180                 return "", err
 181         }
 182
 183         _, err = subagent.Append("user", message)
 184         if err == nil {
 185                 err = subagent.Flush(ctx)
 186         }
 187
 188         if err != nil {
 189                 return "", err
 190         }
 191
 192         err = r.start(ctx, &SubagentRun{
 193                 subagent: subagent,
 194                 view:     view,
 195                 effort:   effort,
 196         })
 197         if err != nil {
 198                 return "", err
 199         }
 200
 201         return "Sent. " + name + "'s reply will reach you as a message.", nil
 202 }
 203
 204 // Stop ends the run of the subagent name and tells if it ran. The chat gets the
 205 // note "Stopped by the user: NAME." under the mutex, so the note comes before
 206 // the report of the run.
 207 func (r *SubagentRuns) Stop(name string) (bool, error) {
 208         r.mutex.Lock()
 209         defer r.mutex.Unlock()
 210
 211         run, found := r.running[name]
 212         if !found {
 213                 return false, nil
 214         }
 215
 216         err := r.note("Stopped by the user: " + name + ".")
 217         if err != nil {
 218                 return true, err
 219         }
 220
 221         run.stopped = true
 222         run.cancel()
 223
 224         return true, nil
 225 }
 226
 227 // Resume goes on with the runs that a stop of the daemon cut, as work.json
 228 // names them. The chat gets the note "The daemon restarted, cutting NAMES;
 229 // nothing was lost: go on.", and each cut subagent gets the same text as a
 230 // message of Tell. A subagent that cannot go on gets a failed report.
 231 func (r *SubagentRuns) Resume(ctx context.Context) error {
 232         names, err := r.load()
 233         if err != nil || len(names) == 0 {
 234                 return err
 235         }
 236
 237         text := "The daemon restarted, cutting " + strings.Join(names, ", ") + "; nothing was lost: go on."
 238         err = r.note(text)
 239         if err != nil {
 240                 return err
 241         }
 242
 243         for _, name := range names {
 244                 _, err = r.Tell(ctx, name, text)
 245                 if err != nil {
 246                         r.logger.Error("subagent lost at restart", "name", name, "err", err)
 247                         r.report(name, subagentReport(name, []string{"failed: " + err.Error()}), subagentFailed)
 248                 }
 249         }
 250
 251         // A subagent that cannot go on does not run, so work.json no longer
 252         // names it, and the next start does not try it again.
 253         r.mutex.Lock()
 254         defer r.mutex.Unlock()
 255
 256         return r.save()
 257 }
 258
 259 // Snapshot locks the run state, memory, and subagent logs while snapshot reads
 260 // the agent files. The lock order agrees with Stop, which logs under the run
 261 // lock.
 262 func (r *SubagentRuns) Snapshot(ctx context.Context, snapshot func() error) error {
 263         r.mutex.Lock()
 264         defer r.mutex.Unlock()
 265
 266         return r.memory.Snapshot(ctx, func() error {
 267                 return r.subagents.Snapshot(ctx, snapshot)
 268         })
 269 }
 270
 271 // Close ends every run and waits for them. No subagent starts after Close.
 272 func (r *SubagentRuns) Close() {
 273         r.mutex.Lock()
 274
 275         r.closed = true
 276         for _, run := range r.running {
 277                 run.cancel()
 278         }
 279
 280         r.mutex.Unlock()
 281
 282         r.workers.Wait()
 283 }
 284
 285 // create makes a new subagent with its note, its task, and its work directory,
 286 // and starts its run.
 287 func (r *SubagentRuns) create(ctx context.Context, note, task string, view []string, effort string) (string, error) {
 288         subagent, err := r.subagents.Create(ctx)
 289         if err != nil {
 290                 return "", err
 291         }
 292
 293         for _, message := range [][2]string{{"note", note}, {"user", task}} {
 294                 _, err = subagent.Append(message[0], message[1])
 295                 if err != nil {
 296                         return "", err
 297                 }
 298         }
 299
 300         err = subagent.Flush(ctx)
 301         if err != nil {
 302                 return "", err
 303         }
 304
 305         err = os.MkdirAll(subagent.WorkDir(), 0o700)
 306         if err != nil {
 307                 return "", fmt.Errorf("create work directory: %w", err)
 308         }
 309
 310         return subagent.name, r.start(ctx, &SubagentRun{
 311                 subagent: subagent,
 312                 view:     view,
 313                 effort:   effort,
 314         })
 315 }
 316
 317 // effort returns the effort of subagent from the note at the start of its log:
 318 // the effort of the agent, or the effort that the note names.
 319 func (r *SubagentRuns) effort(subagent *Subagent) (string, error) {
 320         note, err := subagent.First()
 321         if err != nil {
 322                 return "", err
 323         }
 324
 325         effort, own := strings.CutPrefix(note.Text, subagentNote+subagentNoteEffort)
 326         switch {
 327         case note.Kind != "note":
 328         case own:
 329                 return effort, nil
 330         case note.Text == subagentNote:
 331                 return r.agent.config.Effort, nil
 332         }
 333
 334         return "", fmt.Errorf("the log of %s does not start with a subagent note", subagent.name)
 335 }
 336
 337 // start runs run in the background and logs its start. Its context keeps the
 338 // values of ctx, but only Stop and Close end it.
 339 func (r *SubagentRuns) start(ctx context.Context, run *SubagentRun) error {
 340         ctx, run.cancel = context.WithCancel(context.WithoutCancel(ctx))
 341
 342         r.mutex.Lock()
 343
 344         _, found := r.running[run.subagent.name]
 345         if r.closed || found {
 346                 r.mutex.Unlock()
 347                 run.cancel()
 348
 349                 if found {
 350                         return fmt.Errorf("%s runs already", run.subagent.name)
 351                 }
 352
 353                 return errSubagentClosed
 354         }
 355
 356         r.running[run.subagent.name] = run
 357         err := r.save()
 358         if err != nil {
 359                 delete(r.running, run.subagent.name)
 360                 r.mutex.Unlock()
 361                 run.cancel()
 362
 363                 return err
 364         }
 365
 366         r.workers.Add(1)
 367         r.mutex.Unlock()
 368
 369         run.started = r.subagents.clock()
 370         r.logger.Info("subagent started", "name", run.subagent.name)
 371         r.notify(Event{
 372                 Kind: "subagent",
 373                 Text: run.subagent.name,
 374         })
 375
 376         go r.run(ctx, run)
 377
 378         return nil
 379 }
 380
 381 // load reads the names of work.json. A missing file names none.
 382 func (r *SubagentRuns) load() ([]string, error) {
 383         data, err := os.ReadFile(r.workPath())
 384         if errors.Is(err, os.ErrNotExist) {
 385                 return nil, nil
 386         }
 387
 388         if err != nil {
 389                 return nil, fmt.Errorf("read running subagents: %w", err)
 390         }
 391
 392         var file SubagentWorkFile
 393
 394         err = json.Unmarshal(data, &file)
 395         if err != nil {
 396                 return nil, fmt.Errorf("decode running subagents: %w", err)
 397         }
 398
 399         return file.Running, nil
 400 }
 401
 402 // save writes the names of the running subagents to work.json. The caller
 403 // holds the mutex.
 404 func (r *SubagentRuns) save() error {
 405         file := SubagentWorkFile{
 406                 Running: append([]string{}, slices.Sorted(maps.Keys(r.running))...),
 407         }
 408
 409         data, err := json.Marshal(file)
 410         if err != nil {
 411                 return fmt.Errorf("encode running subagents: %w", err)
 412         }
 413
 414         err = writeFileAtomic(r.workPath(), append(data, '\n'))
 415         if err != nil {
 416                 return fmt.Errorf("save running subagents: %w", err)
 417         }
 418
 419         return nil
 420 }
 421
 422 func (r *SubagentRuns) workPath() string {
 423         return filepath.Join(r.agent.dir, subagentWorkFileName)
 424 }
 425
 426 // run runs the chains of run and reports and logs its end with the length of
 427 // the run, unless Close ended it. A chain follows while messages wait after the final answer of the chain before
 428 // it, and no chain failed. A chain that Stop ends gives the reply "stopped".
 429 func (r *SubagentRuns) run(ctx context.Context, run *SubagentRun) {
 430         defer r.workers.Done()
 431         defer run.cancel()
 432
 433         var (
 434                 replies []string
 435                 failure error
 436         )
 437
 438         stopped := false
 439         status := subagentFinished
 440         for more := true; more; {
 441                 reply, err := r.chain(ctx, run)
 442
 443                 r.mutex.Lock()
 444
 445                 stopped = run.stopped
 446                 switch {
 447                 case err != nil && stopped:
 448                         reply = "stopped"
 449                 case err != nil:
 450                         reply = "failed: " + err.Error()
 451                         status = subagentFailed
 452                         failure = err
 453                 }
 454
 455                 more = err == nil && ctx.Err() == nil && len(run.pending) > 0
 456                 run.pending = nil
 457
 458                 var saveErr error
 459
 460                 if !more {
 461                         delete(r.running, run.subagent.name)
 462
 463                         if !r.closed {
 464                                 saveErr = r.save()
 465                         }
 466                 }
 467
 468                 r.mutex.Unlock()
 469
 470                 if saveErr != nil {
 471                         r.fail(saveErr)
 472                 }
 473
 474                 replies = append(replies, reply)
 475         }
 476
 477         if ctx.Err() != nil && !stopped {
 478                 return
 479         }
 480
 481         if stopped {
 482                 status = subagentStopped
 483         }
 484
 485         length := r.subagents.clock().Sub(run.started)
 486         switch status {
 487         case subagentStopped:
 488                 r.logger.Info("subagent stopped", "name", run.subagent.name, "after", length)
 489         case subagentFailed:
 490                 r.logger.Error("subagent failed", "name", run.subagent.name, "after", length, "err", failure)
 491         default:
 492                 r.logger.Info("subagent done", "name", run.subagent.name, "after", length)
 493         }
 494
 495         r.report(run.subagent.name, subagentReport(run.subagent.name, replies), status)
 496 }
 497
 498 // chain runs one request chain of run and returns its final answer.
 499 func (r *SubagentRuns) chain(ctx context.Context, run *SubagentRun) (string, error) {
 500         instructions, err := r.agent.SubagentPrompt()
 501         if err != nil {
 502                 return "", err
 503         }
 504
 505         log, err := run.subagent.Show()
 506         if err != nil {
 507                 return "", err
 508         }
 509
 510         input := make([]any, 0, len(run.view)+2)
 511         for _, text := range run.view {
 512                 input = append(input, modelUserMessage(text))
 513         }
 514
 515         input = append(input, modelUserMessage(r.state()), modelUserMessage(log))
 516         chain := &Chain{
 517                 client: r.client,
 518                 tools: &Toolbox{
 519                         log:       run.subagent,
 520                         chat:      r.memory,
 521                         subagents: r.subagents,
 522                         workdir:   run.subagent.WorkDir(),
 523                 },
 524                 request: ModelRequest{
 525                         Model:        r.agent.config.Model,
 526                         Effort:       run.effort,
 527                         Instructions: instructions,
 528                         Tools:        promptSubagentTools(),
 529                         Agent:        run.subagent.name,
 530                 },
 531                 take: func() []Message {
 532                         return r.take(run)
 533                 },
 534         }
 535
 536         return chain.Run(ctx, input)
 537 }
 538
 539 // take returns the messages of Tell that wait for a request of run, and
 540 // removes them.
 541 func (r *SubagentRuns) take(run *SubagentRun) []Message {
 542         r.mutex.Lock()
 543         defer r.mutex.Unlock()
 544
 545         messages := run.pending
 546         run.pending = nil
 547
 548         return messages
 549 }
 550
 551 // SubagentRun is one run of a subagent: the view that its requests get, its
 552 // effort, the end of its context, the messages of Tell that wait for a request,
 553 // whether Stop ended it, and the time it started. The mutex of SubagentRuns protects the messages
 554 // and the stop.
 555 type SubagentRun struct {
 556         subagent *Subagent
 557         view     []string
 558         effort   string
 559         cancel   context.CancelFunc
 560         pending  []Message
 561         stopped  bool
 562         started  time.Time
 563 }
 564
 565 // SubagentWorkFile is the content of work.json.
 566 type SubagentWorkFile struct {
 567         Running []string `json:"running"`
 568 }
 569
 570 // subagentReport returns the report of the subagent name with its replies: the
 571 // name in brackets, a new line, the replies, and the zoom that gives its whole
 572 // chat.
 573 func subagentReport(name string, replies []string) string {
 574         return "[" + name + "]\n" + strings.Join(replies, "\n\n") + "\n\nFull chat: zoom(\"" + name + "\")"
 575 }