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

   1 package main
   2
   3 import (
   4         "encoding/json"
   5         "errors"
   6         "fmt"
   7         "os"
   8         "path/filepath"
   9         "slices"
  10         "strings"
  11         "sync"
  12 )
  13
  14 // openDaemonMessages connects the root and its children. It restores queued
  15 // messages before the loops start. A queue entry stays until its recipient
  16 // finishes the turn that takes it, so a stop can interrupt delivery or a turn.
  17 func openDaemonMessages(loops []*TurnLoop) (*DaemonMessages, error) {
  18         m := &DaemonMessages{
  19                 path:  filepath.Join(loops[0].agent.dir, "messages.json"),
  20                 loops: make(map[string]*TurnLoop),
  21         }
  22         for _, loop := range loops {
  23                 m.loops[loop.agent.config.Name] = loop
  24                 loop.messages = m
  25                 if loop.agent.config.Parent == "" {
  26                         loop.logger = loop.logger.With("name", loop.agent.config.Name)
  27                 }
  28         }
  29
  30         data, err := os.ReadFile(m.path)
  31         if err != nil && !errors.Is(err, os.ErrNotExist) {
  32                 return nil, fmt.Errorf("read agent messages: %w", err)
  33         }
  34
  35         if err == nil {
  36                 err = json.Unmarshal(data, &m.state)
  37                 if err != nil {
  38                         return nil, fmt.Errorf("decode agent messages: %w", err)
  39                 }
  40         }
  41
  42         if m.state.Next < 0 {
  43                 return nil, errors.New("invalid agent message sequence")
  44         }
  45
  46         previous := 0
  47         for _, message := range m.state.Pending {
  48                 source := m.loops[message.Source]
  49                 target := m.loops[message.Target]
  50                 if source == nil || target == nil || message.ID <= previous || message.ID > m.state.Next || (message.LogID != nil && *message.LogID < 0) {
  51                         return nil, errors.New("invalid queued agent message")
  52                 }
  53
  54                 if (message.Kind != "user" || target.agent.config.Parent != message.Source) && (message.Kind != "work" || source.agent.config.Parent != message.Target) {
  55                         return nil, errors.New("queued agents are not parent and child")
  56                 }
  57
  58                 previous = message.ID
  59         }
  60
  61         for i := range m.state.Pending {
  62                 err = m.deliver(&m.state.Pending[i])
  63                 if err != nil {
  64                         return nil, fmt.Errorf("restore agent message: %w", err)
  65                 }
  66         }
  67
  68         return m, nil
  69 }
  70
  71 // DaemonMessages owns the durable queue between the root and its children. The
  72 // mutex serializes queue writes and delivery. Backups take it before the locks
  73 // of the loops, memory, and subagents. A stored log ID makes delivery
  74 // idempotent, including a stop between the log append and the wake.
  75 type DaemonMessages struct {
  76         mutex   sync.Mutex
  77         path    string
  78         loops   map[string]*TurnLoop
  79         state   DaemonMessageFile
  80         failure error
  81 }
  82
  83 // deliver records one queued message and wakes its recipient. The caller holds
  84 // the mutex, or owns the queue before the loops start.
  85 func (m *DaemonMessages) deliver(message *DaemonMessage) error {
  86         loop := m.loops[message.Target]
  87         logged, err := loop.memory.AppendOnce(message.Kind, message.Text, message.LogID, func(id int) error {
  88                 message.LogID = &id
  89                 return m.save()
  90         })
  91         if err != nil {
  92                 return fmt.Errorf("record agent message: %w", err)
  93         }
  94
  95         status := ""
  96         if message.Kind == "work" {
  97                 status = subagentFinished
  98         }
  99
 100         loop.notify(Event{
 101                 Kind:   message.Kind,
 102                 ID:     logged.ID,
 103                 Text:   message.Text,
 104                 Source: message.Source,
 105                 Status: status,
 106         })
 107         loop.Wake(logged)
 108         loop.logger.Info("message received", "from", message.Source, "id", message.ID)
 109
 110         return nil
 111 }
 112
 113 // save replaces the queue file under the mutex, or before the loops start.
 114 func (m *DaemonMessages) save() error {
 115         data, err := json.Marshal(m.state)
 116         if err != nil {
 117                 return fmt.Errorf("encode agent messages: %w", err)
 118         }
 119
 120         err = writeFileAtomic(m.path, append(data, '\n'))
 121         if err != nil {
 122                 return fmt.Errorf("save agent messages: %w", err)
 123         }
 124
 125         return nil
 126 }
 127
 128 // Tell queues a user message from parent to its own child, then wakes the
 129 // child. It is the daemon entry point for parent-to-child messages.
 130 func (m *DaemonMessages) Tell(parent, child, text string) error {
 131         m.mutex.Lock()
 132         defer m.mutex.Unlock()
 133
 134         source := m.loops[parent]
 135         target := m.loops[child]
 136         if source == nil || target == nil || target.agent.config.Parent != parent {
 137                 return errors.New("the recipient is not a child of the sender")
 138         }
 139
 140         if m.failure != nil {
 141                 return fmt.Errorf("agent message queue stopped: %w", m.failure)
 142         }
 143
 144         m.add(parent, child, "user", text)
 145
 146         err := m.save()
 147         if err == nil {
 148                 err = m.deliver(&m.state.Pending[len(m.state.Pending)-1])
 149         }
 150
 151         if err != nil {
 152                 return fmt.Errorf("send agent message: %w", m.failed(err))
 153         }
 154
 155         return nil
 156 }
 157
 158 // add appends a queue entry under the mutex. Its ID is never reused.
 159 func (m *DaemonMessages) add(source, target, kind, text string) {
 160         m.state.Next++
 161         m.state.Pending = append(m.state.Pending, DaemonMessage{
 162                 ID:     m.state.Next,
 163                 Source: source,
 164                 Target: target,
 165                 Kind:   kind,
 166                 Text:   text,
 167         })
 168 }
 169
 170 // failed stops all loops after a queue failure. Further sends must not change
 171 // the queue, because the outcome of the failed write is not known.
 172 func (m *DaemonMessages) failed(err error) error {
 173         m.failure = err
 174         for _, loop := range m.loops {
 175                 loop.fail(err)
 176         }
 177
 178         return err
 179 }
 180
 181 // Snapshot holds the queue lock while snapshot locks and copies agent files.
 182 // Its callback must take the memory locks after this lock, as delivery does.
 183 func (m *DaemonMessages) Snapshot(snapshot func() error) error {
 184         m.mutex.Lock()
 185         defer m.mutex.Unlock()
 186
 187         err := snapshot()
 188         if err != nil {
 189                 return fmt.Errorf("copy agent message snapshot: %w", err)
 190         }
 191
 192         return nil
 193 }
 194
 195 // Finish acknowledges the queued messages taken by a turn and queues its final
 196 // report in the same file replacement. A child with no final text sends
 197 // nothing. Call after a final reply, a failure, or a user interrupt, but not
 198 // when a daemon stop interrupts a model request.
 199 func (m *DaemonMessages) Finish(loop *TurnLoop, active []Message, reply string) error {
 200         m.mutex.Lock()
 201         defer m.mutex.Unlock()
 202
 203         if m.failure != nil {
 204                 return fmt.Errorf("agent message queue stopped: %w", m.failure)
 205         }
 206
 207         ids := make(map[int]bool, len(active))
 208         for _, message := range active {
 209                 ids[message.ID] = true
 210         }
 211
 212         name := loop.agent.config.Name
 213         before := len(m.state.Pending)
 214         m.state.Pending = slices.DeleteFunc(m.state.Pending, func(message DaemonMessage) bool {
 215                 return message.Target == name && message.LogID != nil && ids[*message.LogID]
 216         })
 217
 218         parent := loop.agent.config.Parent
 219         report := parent != "" && strings.TrimSpace(reply) != ""
 220         if report {
 221                 m.add(name, parent, "work", "["+name+"] "+reply)
 222         }
 223
 224         if !report && len(m.state.Pending) == before {
 225                 return nil
 226         }
 227
 228         err := m.save()
 229         if err == nil && report {
 230                 err = m.deliver(&m.state.Pending[len(m.state.Pending)-1])
 231         }
 232
 233         if err != nil {
 234                 return fmt.Errorf("finish agent messages: %w", m.failed(err))
 235         }
 236
 237         return nil
 238 }
 239
 240 // DaemonMessageFile holds the sequence and the unacknowledged deliveries.
 241 type DaemonMessageFile struct {
 242         Next    int             `json:"next"`
 243         Pending []DaemonMessage `json:"pending"`
 244 }
 245
 246 // DaemonMessage is one queued delivery. LogID is set before the log append. A
 247 // nil LogID means that delivery has not started.
 248 type DaemonMessage struct {
 249         ID     int    `json:"id"`
 250         Source string `json:"source"`
 251         Target string `json:"target"`
 252         Kind   string `json:"kind"`
 253         Text   string `json:"text"`
 254         LogID  *int   `json:"log_id,omitempty"`
 255 }