[onechat] / daemon.go

package main

import (
        "bufio"
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "net"
        "os"
        "path/filepath"
        "sync"
        "time"
)

const (
        // daemonAgentDirName is the directory of the root agent in the data
        // directory.
        daemonAgentDirName = "agent"

        // daemonClientEvents is the number of events that can wait for a slow
        // client before the daemon drops it.
        daemonClientEvents = 1024
)

// DaemonHub holds the connected clients and sends each event to all of them. It
// runs the WhatsApp commands of the clients on whatsapp. snapshot locks all
// agents and subagents during a backup copy.
type DaemonHub struct {
        whatsapp *WhatsApp
        snapshot func(context.Context, func() error) error

        mutex   sync.Mutex
        clients map[*DaemonClient]bool
        workers sync.WaitGroup
}

// DaemonClient is one connected client. Its events wait in a channel for the
// writer, so a slow client does not stop the agent.
type DaemonClient struct {
        conn   net.Conn
        events chan Event
        once   sync.Once
}

// Accept serves each client of listener until the listener closes. A new client
// gets a ready event with the name of the agent, and a busy event when a turn
// runs.
func (h *DaemonHub) Accept(ctx context.Context, listener net.Listener, l *TurnLoop) {
        for {
                conn, err := listener.Accept()
                if err != nil {
                        return
                }

                client := &DaemonClient{
                        conn:   conn,
                        events: make(chan Event, daemonClientEvents),
                }

                h.mutex.Lock()

                h.clients[client] = true
                h.mutex.Unlock()

                h.workers.Add(2)

                go h.write(client)
                go h.read(ctx, client, l)

                h.Send(client, Event{
                        Kind: "ready",
                        Text: l.agent.config.Name,
                })

                if l.Busy() {
                        h.Send(client, Event{
                                Kind: "busy",
                        })
                }
        }
}

// write sends the events of client as JSON lines until its channel closes.
func (h *DaemonHub) write(client *DaemonClient) {
        defer h.workers.Done()

        encoder := json.NewEncoder(client.conn)
        for event := range client.events {
                err := encoder.Encode(event)
                if err != nil {
                        h.drop(client)
                }
        }
}

// read runs the requests of client until it closes its connection.
func (h *DaemonHub) read(ctx context.Context, client *DaemonClient, l *TurnLoop) {
        defer h.workers.Done()
        defer h.drop(client)

        scanner := bufio.NewScanner(client.conn)
        scanner.Buffer(make([]byte, 64*1024), socketRequestBytes)

        for ctx.Err() == nil && scanner.Scan() {
                var request SocketRequest

                err := json.Unmarshal(scanner.Bytes(), &request)
                if err == nil {
                        if request.Type == "backup" {
                                var size int64

                                size, err = backupData(ctx, filepath.Dir(l.agent.dir), request.Text, h.whatsapp, h.snapshot)
                                if err == nil {
                                        h.Send(client, Event{
                                                Kind: "backup",
                                                Text: fmt.Sprintf("%s (%d bytes)", request.Text, size),
                                        })
                                }
                        } else {
                                err = daemonRunRequest(l, h.whatsapp, request, func(text string) {
                                        h.Send(client, Event{
                                                Kind: "notice",
                                                Text: text,
                                        })
                                })
                        }
                }

                if err != nil {
                        h.Send(client, Event{
                                Kind: "error",
                                Text: err.Error(),
                        })
                }
        }
}

// Send gives event to one client, or drops the client when its events do not
// fit.
func (h *DaemonHub) Send(client *DaemonClient, event Event) {
        h.mutex.Lock()
        defer h.mutex.Unlock()

        h.send(client, event)
}

// drop removes client and closes its connection.
func (h *DaemonHub) drop(client *DaemonClient) {
        h.mutex.Lock()
        defer h.mutex.Unlock()

        h.remove(client)
}

// send gives event to client. The caller holds the mutex.
func (h *DaemonHub) send(client *DaemonClient, event Event) {
        if !h.clients[client] {
                return
        }

        select {
        case client.events <- event:
        default:
                h.remove(client)
        }
}

// remove removes client, closes its connection, and ends its writer. The caller
// holds the mutex.
func (h *DaemonHub) remove(client *DaemonClient) {
        delete(h.clients, client)
        client.once.Do(func() {
                close(client.events)

                // The client is removed already, so a close failure cannot
                // restore it.
                client.conn.Close() // nolint: errcheck
        })
}

// Broadcast gives event to every client.
func (h *DaemonHub) Broadcast(event Event) {
        h.mutex.Lock()
        defer h.mutex.Unlock()

        for client := range h.clients {
                h.send(client, event)
        }
}

// CloseAll drops every client and waits for their readers and writers.
func (h *DaemonHub) CloseAll() {
        h.mutex.Lock()

        for client := range h.clients {
                h.remove(client)
        }

        h.mutex.Unlock()

        h.workers.Wait()
}

// daemonRunRequest runs one request of a client. The answers of a WhatsApp
// command go to answer.
func daemonRunRequest(l *TurnLoop, whatsapp *WhatsApp, request SocketRequest, answer func(text string)) error {
        switch request.Type {
        case "send":
                if request.Text == "" {
                        return errors.New("empty message")
                }

                _, err := l.Submit(terminalSource, request.Text)

                return err
        case "image":
                if !filepath.IsAbs(request.Path) {
                        return errors.New("use /image PATH [caption]")
                }

                data, err := imageReadFile(request.Path)
                if err != nil {
                        return fmt.Errorf("read terminal image: %w", err)
                }

                name, err := l.agent.Images().Save(data)
                if err != nil {
                        return fmt.Errorf("store terminal image: %w", err)
                }

                _, err = l.Submit(terminalSource, imageText([]string{name}, request.Text))
                if err != nil {
                        return fmt.Errorf("submit terminal image: %w", err)
                }

                return nil
        case "interrupt":
                if !l.Interrupt() {
                        return errors.New("no turn runs")
                }

                return nil
        case "whatsapp":
                whatsapp.Command(request.Text, answer)

                return nil
        case "stop":
                found, err := l.StopSubagent(request.Text)
                if err == nil && !found {
                        err = fmt.Errorf("no subagent %s runs", request.Text)
                }

                return err
        default:
                return fmt.Errorf("unknown request %q", request.Type)
        }
}

// daemon runs "onechat daemon": it owns the root, its children, and their
// files, and serves the clients on the Unix socket in the data directory until
// ctx ends. When telegram.json exists, it also runs the Telegram bot. It also
// keeps a linked WhatsApp device connected. It holds the lock of the token
// file, so it is the only process that refreshes the tokens and opens the
// WhatsApp session. Its log lines go to out, through the loggers of
// daemonLogger.
func daemon(ctx context.Context, out io.Writer) (err error) {
        dir, err := dataDir()
        if err != nil {
                return err
        }

        // Check the socket path before the daemon writes a file, so that a data
        // directory with a long path stays empty.
        socket := filepath.Join(dir, socketFileName)
        err = checkSocketPath(socket)
        if err != nil {
                return err
        }

        config, err := loadConfig(dir)
        if err != nil {
                return err
        }

        botSettings, hasBot, err := loadTelegramConfig(dir)
        if err != nil {
                return err
        }

        tokens, err := lockTokens(dir)
        if err != nil {
                return err
        }

        defer func() {
                err = errors.Join(err, tokens.Close())
        }()

        whatsappSetDeviceProps()

        logger := daemonLogger(out, "onechat: ")
        whatsapp, err := openWhatsApp(ctx, dir, time.Now, daemonLogger(out, "whatsapp: "), whatsappNewClient)
        if err != nil {
                return err
        }

        defer func() {
                err = errors.Join(err, whatsapp.Close())
        }()

        address, auth, err := modelEndpoint(config, dir)
        if err != nil {
                return err
        }

        agents, err := openDaemonAgents(ctx, dir, address, auth, config.Compactor, time.Now, logger)
        if err != nil {
                return err
        }

        defer func() {
                err = errors.Join(err, agents.Close())
        }()

        whatsapp.inbox.voice = &VoiceTranscriber{config: config.Voice}

        root := agents[0]
        agent := root.agent
        schedules := root.schedules

        listener, err := daemonListenSocket(socket)
        if err != nil {
                return err
        }

        linked := "not linked"
        if whatsapp.isLinked() {
                linked = "linked"
        }

        logger.Info("ready", "version", version, "socket", listener.Addr(), "model", agent.config.Model, "effort", agent.config.Effort, "workdir", agent.config.Workdir, "whatsapp", linked, "schedules", schedules.Len(), "children", agents.ChildNames())

        var telegram *TelegramBot

        if hasBot {
                telegram = newTelegramBot(newTelegramAPI(telegramAPIURL, botSettings.Token), botSettings.Owner, filepath.Join(agent.dir, telegramOffsetFileName), daemonLogger(out, "telegram: "))

                telegram.voice = &VoiceTranscriber{config: config.Voice}

                logger.Info("Telegram bot started")
        }

        return daemonServe(ctx, listener, telegram, whatsapp, schedules, func(notify func(Event)) (*TurnLoop, error) {
                return root.Loop(whatsapp, notify)
        }, agents[1:]...)
}

// daemonListenSocket listens on the Unix socket at path, with mode 0600. The
// caller owns the chat, so a socket file at path is left from a daemon that
// ended, and daemonListenSocket removes it.
func daemonListenSocket(path string) (net.Listener, error) {
        err := checkSocketPath(path)
        if err != nil {
                return nil, err
        }

        err = os.Remove(path)
        if err != nil && !errors.Is(err, os.ErrNotExist) {
                return nil, fmt.Errorf("remove old socket: %w", err)
        }

        listener, err := net.Listen("unix", path)
        if err != nil {
                return nil, fmt.Errorf("listen on socket: %w", err)
        }

        err = os.Chmod(path, 0o600)
        if err != nil {
                return nil, errors.Join(fmt.Errorf("set socket mode: %w", err), listener.Close())
        }

        return listener, nil
}

// daemonServe runs the loop that start makes and serves the clients of listener
// until ctx ends or the loop fails. It closes listener. It runs whatsapp, which
// logs its batches and notes in the loop and tells all clients its notices, and
// a failure of whatsapp stops the agent. It runs schedules, which log their
// runs in the loop, and a failure of schedules stops the agent too. When bot is
// not nil, it runs the bot too, and a failure of the bot stops the agent. Each
// child runs its own loop and schedules with no client events. A failed child
// stops the daemon, and backups lock all loops before they copy data.
func daemonServe(ctx context.Context, listener net.Listener, bot *TelegramBot, whatsapp *WhatsApp, schedules *Schedules, start func(notify func(Event)) (*TurnLoop, error), children ...*DaemonAgent) error {
        ctx, cancel := context.WithCancel(ctx)
        defer cancel()

        hub := &DaemonHub{
                whatsapp: whatsapp,
                clients:  make(map[*DaemonClient]bool),
        }

        notify := hub.Broadcast
        if bot != nil {
                notify = func(event Event) {
                        hub.Broadcast(event)
                        bot.Notify(event)
                }
        }

        l, err := start(notify)
        if err != nil {
                return errors.Join(err, listener.Close())
        }

        if bot != nil {
                l.tools.imageShow.bot = bot
        }

        loops := []*TurnLoop{l}
        for _, child := range children {
                loop, err := child.Loop(nil, func(Event) {
                })
                if err != nil {
                        return errors.Join(fmt.Errorf("start child %q: %w", child.agent.config.Name, err), listener.Close())
                }

                loops = append(loops, loop)
        }

        messages, err := openDaemonMessages(loops)
        if err != nil {
                return errors.Join(fmt.Errorf("open agent messages: %w", err), listener.Close())
        }

        hub.snapshot = func(ctx context.Context, snapshot func() error) error {
                return messages.SnapshotAgents(ctx, snapshot)
        }

        runtime := &DaemonChildren{
                ctx:      ctx,
                cancel:   cancel,
                root:     l,
                messages: messages,
        }

        l.tools.children = runtime

        var workers sync.WaitGroup

        childErrors := make([]error, len(children))
        for i, child := range children {
                workers.Add(1)

                go func() {
                        defer workers.Done()

                        childErrors[i] = daemonRunChild(ctx, loops[i+1], child.schedules)
                        if childErrors[i] != nil {
                                child.logger.Error("child stopped", "err", childErrors[i])
                        }

                        cancel()
                }()
        }

        workers.Add(2)

        go func() {
                defer workers.Done()

                hub.Accept(ctx, listener, l)
        }()

        var whatsappErr error

        go func() {
                defer workers.Done()

                whatsappErr = whatsapp.Run(ctx, l, notify)

                cancel()
        }()

        var schedulesErr error

        workers.Add(1)

        go func() {
                defer workers.Done()

                schedulesErr = schedules.Run(ctx, l)

                cancel()
        }()

        var botErr error

        if bot != nil {
                workers.Add(1)

                go func() {
                        defer workers.Done()

                        botErr = bot.Run(ctx, l, whatsapp)

                        cancel()
                }()
        }

        err = l.Run(ctx)
        if errors.Is(err, context.Canceled) {
                err = nil
        }

        cancel()

        err = errors.Join(err, listener.Close())

        hub.CloseAll()
        workers.Wait()

        return errors.Join(err, whatsappErr, schedulesErr, botErr, errors.Join(childErrors...), runtime.Close())
}

// daemonRunChild runs a child's turn loop and its schedules until either stops.
// Its events stay in its own log.
func daemonRunChild(ctx context.Context, loop *TurnLoop, schedules *Schedules) error {
        ctx, cancel := context.WithCancel(ctx)
        defer cancel()

        done := make(chan error, 1)
        go func() {
                done <- schedules.Run(ctx, loop)

                cancel()
        }()

        err := loop.Run(ctx)
        if errors.Is(err, context.Canceled) {
                err = nil
        }

        cancel()

        return errors.Join(err, <-done)
}