[onechat] / daemon_message.go

package main

import (
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "os"
        "path/filepath"
        "slices"
        "strings"
        "sync"
)

// openDaemonMessages connects the root and its children. It restores queued
// messages before the loops start. A queue entry stays until its recipient
// finishes the turn that takes it, so a stop can interrupt delivery or a turn.
func openDaemonMessages(loops []*TurnLoop) (*DaemonMessages, error) {
        m := &DaemonMessages{
                path:    filepath.Join(loops[0].agent.dir, "messages.json"),
                ordered: append([]*TurnLoop(nil), loops...),
                loops:   make(map[string]*TurnLoop),
        }
        for _, loop := range loops {
                m.loops[loop.agent.config.Name] = loop
                loop.messages = m
                loop.tools.messages = m
                if loop.agent.config.Parent == "" {
                        loop.logger = loop.logger.With("name", loop.agent.config.Name)
                }
        }

        for _, child := range loops[1:] {
                err := loops[0].runs.subagents.Reserve(child.agent.config.Name)
                if err != nil {
                        return nil, fmt.Errorf("reserve child name: %w", err)
                }
        }

        data, err := os.ReadFile(m.path)
        if err != nil && !errors.Is(err, os.ErrNotExist) {
                return nil, fmt.Errorf("read agent messages: %w", err)
        }

        if err == nil {
                err = json.Unmarshal(data, &m.state)
                if err != nil {
                        return nil, fmt.Errorf("decode agent messages: %w", err)
                }
        }

        if m.state.Next < 0 {
                return nil, errors.New("invalid agent message sequence")
        }

        previous := 0
        for _, message := range m.state.Pending {
                source := m.loops[message.Source]
                target := m.loops[message.Target]
                if source == nil || target == nil || message.ID <= previous || message.ID > m.state.Next || (message.LogID != nil && *message.LogID < 0) {
                        return nil, errors.New("invalid queued agent message")
                }

                if (message.Kind != "user" || target.agent.config.Parent != message.Source) && (message.Kind != "work" || source.agent.config.Parent != message.Target) {
                        return nil, errors.New("queued agents are not parent and child")
                }

                previous = message.ID
        }

        for i := range m.state.Pending {
                err = m.deliver(&m.state.Pending[i])
                if err != nil {
                        return nil, fmt.Errorf("restore agent message: %w", err)
                }
        }

        return m, nil
}

// DaemonMessages owns the durable queue between the root and its children. The
// mutex serializes queue writes and delivery. Backups take it before the locks
// of the loops, memory, and subagents. A stored log ID makes delivery
// idempotent, including a stop between the log append and the wake.
type DaemonMessages struct {
        mutex   sync.Mutex
        path    string
        loops   map[string]*TurnLoop
        ordered []*TurnLoop
        state   DaemonMessageFile
        failure error
}

// deliver records one queued message and wakes its recipient. The caller holds
// the mutex, or owns the queue before the loops start.
func (m *DaemonMessages) deliver(message *DaemonMessage) error {
        loop := m.loops[message.Target]
        logged, err := loop.memory.AppendOnce(message.Kind, message.Text, message.LogID, func(id int) error {
                message.LogID = &id
                return m.save()
        })
        if err != nil {
                return fmt.Errorf("record agent message: %w", err)
        }

        status := ""
        if message.Kind == "work" {
                status = subagentFinished
        }

        loop.notify(Event{
                Kind:   message.Kind,
                ID:     logged.ID,
                Text:   message.Text,
                Source: message.Source,
                Status: status,
        })
        loop.Wake(logged)
        loop.logger.Info("message received", "from", message.Source, "id", message.ID)

        return nil
}

// save replaces the queue file under the mutex, or before the loops start.
func (m *DaemonMessages) save() error {
        data, err := json.Marshal(m.state)
        if err != nil {
                return fmt.Errorf("encode agent messages: %w", err)
        }

        err = writeFileAtomic(m.path, append(data, '\n'))
        if err != nil {
                return fmt.Errorf("save agent messages: %w", err)
        }

        return nil
}

// AddChild holds the queue lock while open makes a child's files and loop. The
// loop joins backups before its first task can be delivered.
func (m *DaemonMessages) AddChild(parent, name string, open func() (*TurnLoop, error)) (*TurnLoop, error) {
        m.mutex.Lock()
        defer m.mutex.Unlock()

        if m.loops[name] != nil {
                return nil, fmt.Errorf("name %q is already used", name)
        }

        root := m.loops[parent]
        if root == nil || root.agent.config.Parent != "" {
                return nil, errors.New("only the root can create children")
        }

        if m.failure != nil {
                return nil, fmt.Errorf("agent message queue stopped: %w", m.failure)
        }

        loop, err := open()
        if err != nil {
                return nil, fmt.Errorf("open child: %w", err)
        }

        loop.messages = m
        loop.tools.messages = m
        m.loops[name] = loop
        m.ordered = append(m.ordered, loop)

        return loop, nil
}

// Child returns only a direct child of parent. Loops live until the daemon
// stops, so the returned loop stays open for the caller's turn.
func (m *DaemonMessages) Child(parent, name string) *TurnLoop {
        m.mutex.Lock()
        defer m.mutex.Unlock()

        loop := m.loops[name]
        if loop == nil || loop.agent.config.Parent != parent {
                return nil
        }

        return loop
}

// ChildState lists the names and responsibilities of a parent's children in
// name order. It returns no heading when the parent has no children.
func (m *DaemonMessages) ChildState(parent string) string {
        m.mutex.Lock()
        defer m.mutex.Unlock()

        var lines []string

        for name, loop := range m.loops {
                if loop.agent.config.Parent == parent {
                        lines = append(lines, name+": "+loop.agent.config.Responsibility)
                }
        }

        if len(lines) == 0 {
                return ""
        }

        slices.Sort(lines)

        return "\n\nChildren:\n" + strings.Join(lines, "\n")
}

// SnapshotAgents locks the queue and all registered loops, including children
// made during this daemon run, before snapshot copies their files.
func (m *DaemonMessages) SnapshotAgents(ctx context.Context, snapshot func() error) error {
        return m.Snapshot(func() error {
                return daemonSnapshot(ctx, m.ordered, snapshot)
        })
}

// Tell queues a user message from parent to its own child, then wakes the
// child. It is the daemon entry point for parent-to-child messages.
func (m *DaemonMessages) Tell(parent, child, text string) error {
        m.mutex.Lock()
        defer m.mutex.Unlock()

        source := m.loops[parent]
        target := m.loops[child]
        if source == nil || target == nil || target.agent.config.Parent != parent {
                return errors.New("the recipient is not a child of the sender")
        }

        if m.failure != nil {
                return fmt.Errorf("agent message queue stopped: %w", m.failure)
        }

        m.add(parent, child, "user", text)

        err := m.save()
        if err == nil {
                err = m.deliver(&m.state.Pending[len(m.state.Pending)-1])
        }

        if err != nil {
                return fmt.Errorf("send agent message: %w", m.failed(err))
        }

        return nil
}

// add appends a queue entry under the mutex. Its ID is never reused.
func (m *DaemonMessages) add(source, target, kind, text string) {
        m.state.Next++
        m.state.Pending = append(m.state.Pending, DaemonMessage{
                ID:     m.state.Next,
                Source: source,
                Target: target,
                Kind:   kind,
                Text:   text,
        })
}

// failed stops all loops after a queue failure. Further sends must not change
// the queue, because the outcome of the failed write is not known.
func (m *DaemonMessages) failed(err error) error {
        m.failure = err
        for _, loop := range m.loops {
                loop.fail(err)
        }

        return err
}

// Snapshot holds the queue lock while snapshot locks and copies agent files.
// Its callback must take the memory locks after this lock, as delivery does.
func (m *DaemonMessages) Snapshot(snapshot func() error) error {
        m.mutex.Lock()
        defer m.mutex.Unlock()

        err := snapshot()
        if err != nil {
                return fmt.Errorf("copy agent message snapshot: %w", err)
        }

        return nil
}

// Finish acknowledges the queued messages taken by a turn and queues its final
// report in the same file replacement. A child with no final text sends
// nothing. Call after a final reply, a failure, or a user interrupt, but not
// when a daemon stop interrupts a model request.
func (m *DaemonMessages) Finish(loop *TurnLoop, active []Message, reply string) error {
        m.mutex.Lock()
        defer m.mutex.Unlock()

        if m.failure != nil {
                return fmt.Errorf("agent message queue stopped: %w", m.failure)
        }

        ids := make(map[int]bool, len(active))
        for _, message := range active {
                ids[message.ID] = true
        }

        name := loop.agent.config.Name
        before := len(m.state.Pending)
        m.state.Pending = slices.DeleteFunc(m.state.Pending, func(message DaemonMessage) bool {
                return message.Target == name && message.LogID != nil && ids[*message.LogID]
        })

        parent := loop.agent.config.Parent
        report := parent != "" && strings.TrimSpace(reply) != ""
        if report {
                m.add(name, parent, "work", "["+name+"] "+reply)
        }

        if !report && len(m.state.Pending) == before {
                return nil
        }

        err := m.save()
        if err == nil && report {
                err = m.deliver(&m.state.Pending[len(m.state.Pending)-1])
        }

        if err != nil {
                return fmt.Errorf("finish agent messages: %w", m.failed(err))
        }

        return nil
}

// DaemonMessageFile holds the sequence and the unacknowledged deliveries.
type DaemonMessageFile struct {
        Next    int             `json:"next"`
        Pending []DaemonMessage `json:"pending"`
}

// DaemonMessage is one queued delivery. LogID is set before the log append. A
// nil LogID means that delivery has not started.
type DaemonMessage struct {
        ID     int    `json:"id"`
        Source string `json:"source"`
        Target string `json:"target"`
        Kind   string `json:"kind"`
        Text   string `json:"text"`
        LogID  *int   `json:"log_id,omitempty"`
}