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

   1 package main
   2
   3 import (
   4         "context"
   5         "encoding/json"
   6         "errors"
   7         "fmt"
   8         "os"
   9         "path/filepath"
  10         "slices"
  11         "strings"
  12         "sync"
  13 )
  14
  15 // openDaemonMessages connects the root and its children. It restores queued
  16 // messages before the loops start. A queue entry stays until its recipient
  17 // finishes the turn that takes it, so a stop can interrupt delivery or a turn.
  18 func openDaemonMessages(loops []*TurnLoop) (*DaemonMessages, error) {
  19         m := &DaemonMessages{
  20                 path:    filepath.Join(loops[0].agent.dir, "messages.json"),
  21                 ordered: append([]*TurnLoop(nil), loops...),
  22                 loops:   make(map[string]*TurnLoop),
  23         }
  24         for _, loop := range loops {
  25                 m.loops[loop.agent.config.Name] = loop
  26                 loop.messages = m
  27                 loop.tools.messages = m
  28                 if loop.agent.config.Parent == "" {
  29                         loop.logger = loop.logger.With("name", loop.agent.config.Name)
  30                 }
  31         }
  32
  33         for _, child := range loops[1:] {
  34                 err := loops[0].runs.subagents.Reserve(child.agent.config.Name)
  35                 if err != nil {
  36                         return nil, fmt.Errorf("reserve child name: %w", err)
  37                 }
  38         }
  39
  40         data, err := os.ReadFile(m.path)
  41         if err != nil && !errors.Is(err, os.ErrNotExist) {
  42                 return nil, fmt.Errorf("read agent messages: %w", err)
  43         }
  44
  45         if err == nil {
  46                 err = json.Unmarshal(data, &m.state)
  47                 if err != nil {
  48                         return nil, fmt.Errorf("decode agent messages: %w", err)
  49                 }
  50         }
  51
  52         if m.state.Next < 0 {
  53                 return nil, errors.New("invalid agent message sequence")
  54         }
  55
  56         previous := 0
  57         for _, message := range m.state.Pending {
  58                 source := m.loops[message.Source]
  59                 target := m.loops[message.Target]
  60                 if source == nil || target == nil || message.ID <= previous || message.ID > m.state.Next || (message.LogID != nil && *message.LogID < 0) {
  61                         return nil, errors.New("invalid queued agent message")
  62                 }
  63
  64                 if (message.Kind != "user" || target.agent.config.Parent != message.Source) && (message.Kind != "work" || source.agent.config.Parent != message.Target) {
  65                         return nil, errors.New("queued agents are not parent and child")
  66                 }
  67
  68                 previous = message.ID
  69         }
  70
  71         for i := range m.state.Pending {
  72                 err = m.deliver(&m.state.Pending[i])
  73                 if err != nil {
  74                         return nil, fmt.Errorf("restore agent message: %w", err)
  75                 }
  76         }
  77
  78         return m, nil
  79 }
  80
  81 // DaemonMessages owns the durable queue between the root and its children. The
  82 // mutex serializes queue writes and delivery. Backups take it before the locks
  83 // of the loops, memory, and subagents. A stored log ID makes delivery
  84 // idempotent, including a stop between the log append and the wake.
  85 type DaemonMessages struct {
  86         mutex   sync.Mutex
  87         path    string
  88         loops   map[string]*TurnLoop
  89         ordered []*TurnLoop
  90         state   DaemonMessageFile
  91         failure error
  92 }
  93
  94 // deliver records one queued message and wakes its recipient. The caller holds
  95 // the mutex, or owns the queue before the loops start.
  96 func (m *DaemonMessages) deliver(message *DaemonMessage) error {
  97         loop := m.loops[message.Target]
  98         logged, err := loop.memory.AppendOnce(message.Kind, message.Text, message.LogID, func(id int) error {
  99                 message.LogID = &id
 100                 return m.save()
 101         })
 102         if err != nil {
 103                 return fmt.Errorf("record agent message: %w", err)
 104         }
 105
 106         status := ""
 107         if message.Kind == "work" {
 108                 status = subagentFinished
 109         }
 110
 111         loop.notify(Event{
 112                 Kind:   message.Kind,
 113                 ID:     logged.ID,
 114                 Text:   message.Text,
 115                 Source: message.Source,
 116                 Status: status,
 117         })
 118         loop.Wake(logged)
 119         loop.logger.Info("message received", "from", message.Source, "id", message.ID)
 120
 121         return nil
 122 }
 123
 124 // save replaces the queue file under the mutex, or before the loops start.
 125 func (m *DaemonMessages) save() error {
 126         data, err := json.Marshal(m.state)
 127         if err != nil {
 128                 return fmt.Errorf("encode agent messages: %w", err)
 129         }
 130
 131         err = writeFileAtomic(m.path, append(data, '\n'))
 132         if err != nil {
 133                 return fmt.Errorf("save agent messages: %w", err)
 134         }
 135
 136         return nil
 137 }
 138
 139 // AddChild holds the queue lock while open makes a child's files and loop. The
 140 // loop joins backups before its first task can be delivered.
 141 func (m *DaemonMessages) AddChild(parent, name string, open func() (*TurnLoop, error)) (*TurnLoop, error) {
 142         m.mutex.Lock()
 143         defer m.mutex.Unlock()
 144
 145         if m.loops[name] != nil {
 146                 return nil, fmt.Errorf("name %q is already used", name)
 147         }
 148
 149         root := m.loops[parent]
 150         if root == nil || root.agent.config.Parent != "" {
 151                 return nil, errors.New("only the root can create children")
 152         }
 153
 154         if m.failure != nil {
 155                 return nil, fmt.Errorf("agent message queue stopped: %w", m.failure)
 156         }
 157
 158         loop, err := open()
 159         if err != nil {
 160                 return nil, fmt.Errorf("open child: %w", err)
 161         }
 162
 163         loop.messages = m
 164         loop.tools.messages = m
 165         m.loops[name] = loop
 166         m.ordered = append(m.ordered, loop)
 167
 168         return loop, nil
 169 }
 170
 171 // Child returns only a direct child of parent. Loops live until the daemon
 172 // stops, so the returned loop stays open for the caller's turn.
 173 func (m *DaemonMessages) Child(parent, name string) *TurnLoop {
 174         m.mutex.Lock()
 175         defer m.mutex.Unlock()
 176
 177         loop := m.loops[name]
 178         if loop == nil || loop.agent.config.Parent != parent {
 179                 return nil
 180         }
 181
 182         return loop
 183 }
 184
 185 // ChildState lists the names and responsibilities of a parent's children in
 186 // name order. It returns no heading when the parent has no children.
 187 func (m *DaemonMessages) ChildState(parent string) string {
 188         m.mutex.Lock()
 189         defer m.mutex.Unlock()
 190
 191         var lines []string
 192
 193         for name, loop := range m.loops {
 194                 if loop.agent.config.Parent == parent {
 195                         lines = append(lines, name+": "+loop.agent.config.Responsibility)
 196                 }
 197         }
 198
 199         if len(lines) == 0 {
 200                 return ""
 201         }
 202
 203         slices.Sort(lines)
 204
 205         return "\n\nChildren:\n" + strings.Join(lines, "\n")
 206 }
 207
 208 // SnapshotAgents locks the queue and all registered loops, including children
 209 // made during this daemon run, before snapshot copies their files.
 210 func (m *DaemonMessages) SnapshotAgents(ctx context.Context, snapshot func() error) error {
 211         return m.Snapshot(func() error {
 212                 return daemonSnapshot(ctx, m.ordered, snapshot)
 213         })
 214 }
 215
 216 // Tell queues a user message from parent to its own child, then wakes the
 217 // child. It is the daemon entry point for parent-to-child messages.
 218 func (m *DaemonMessages) Tell(parent, child, text string) error {
 219         m.mutex.Lock()
 220         defer m.mutex.Unlock()
 221
 222         source := m.loops[parent]
 223         target := m.loops[child]
 224         if source == nil || target == nil || target.agent.config.Parent != parent {
 225                 return errors.New("the recipient is not a child of the sender")
 226         }
 227
 228         if m.failure != nil {
 229                 return fmt.Errorf("agent message queue stopped: %w", m.failure)
 230         }
 231
 232         m.add(parent, child, "user", text)
 233
 234         err := m.save()
 235         if err == nil {
 236                 err = m.deliver(&m.state.Pending[len(m.state.Pending)-1])
 237         }
 238
 239         if err != nil {
 240                 return fmt.Errorf("send agent message: %w", m.failed(err))
 241         }
 242
 243         return nil
 244 }
 245
 246 // add appends a queue entry under the mutex. Its ID is never reused.
 247 func (m *DaemonMessages) add(source, target, kind, text string) {
 248         m.state.Next++
 249         m.state.Pending = append(m.state.Pending, DaemonMessage{
 250                 ID:     m.state.Next,
 251                 Source: source,
 252                 Target: target,
 253                 Kind:   kind,
 254                 Text:   text,
 255         })
 256 }
 257
 258 // failed stops all loops after a queue failure. Further sends must not change
 259 // the queue, because the outcome of the failed write is not known.
 260 func (m *DaemonMessages) failed(err error) error {
 261         m.failure = err
 262         for _, loop := range m.loops {
 263                 loop.fail(err)
 264         }
 265
 266         return err
 267 }
 268
 269 // Snapshot holds the queue lock while snapshot locks and copies agent files.
 270 // Its callback must take the memory locks after this lock, as delivery does.
 271 func (m *DaemonMessages) Snapshot(snapshot func() error) error {
 272         m.mutex.Lock()
 273         defer m.mutex.Unlock()
 274
 275         err := snapshot()
 276         if err != nil {
 277                 return fmt.Errorf("copy agent message snapshot: %w", err)
 278         }
 279
 280         return nil
 281 }
 282
 283 // Finish acknowledges the queued messages taken by a turn and queues its final
 284 // report in the same file replacement. A child with no final text sends
 285 // nothing. Call after a final reply, a failure, or a user interrupt, but not
 286 // when a daemon stop interrupts a model request.
 287 func (m *DaemonMessages) Finish(loop *TurnLoop, active []Message, reply string) error {
 288         m.mutex.Lock()
 289         defer m.mutex.Unlock()
 290
 291         if m.failure != nil {
 292                 return fmt.Errorf("agent message queue stopped: %w", m.failure)
 293         }
 294
 295         ids := make(map[int]bool, len(active))
 296         for _, message := range active {
 297                 ids[message.ID] = true
 298         }
 299
 300         name := loop.agent.config.Name
 301         before := len(m.state.Pending)
 302         m.state.Pending = slices.DeleteFunc(m.state.Pending, func(message DaemonMessage) bool {
 303                 return message.Target == name && message.LogID != nil && ids[*message.LogID]
 304         })
 305
 306         parent := loop.agent.config.Parent
 307         report := parent != "" && strings.TrimSpace(reply) != ""
 308         if report {
 309                 m.add(name, parent, "work", "["+name+"] "+reply)
 310         }
 311
 312         if !report && len(m.state.Pending) == before {
 313                 return nil
 314         }
 315
 316         err := m.save()
 317         if err == nil && report {
 318                 err = m.deliver(&m.state.Pending[len(m.state.Pending)-1])
 319         }
 320
 321         if err != nil {
 322                 return fmt.Errorf("finish agent messages: %w", m.failed(err))
 323         }
 324
 325         return nil
 326 }
 327
 328 // DaemonMessageFile holds the sequence and the unacknowledged deliveries.
 329 type DaemonMessageFile struct {
 330         Next    int             `json:"next"`
 331         Pending []DaemonMessage `json:"pending"`
 332 }
 333
 334 // DaemonMessage is one queued delivery. LogID is set before the log append. A
 335 // nil LogID means that delivery has not started.
 336 type DaemonMessage struct {
 337         ID     int    `json:"id"`
 338         Source string `json:"source"`
 339         Target string `json:"target"`
 340         Kind   string `json:"kind"`
 341         Text   string `json:"text"`
 342         LogID  *int   `json:"log_id,omitempty"`
 343 }