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

   1 package main
   2
   3 import (
   4         "bufio"
   5         "bytes"
   6         "context"
   7         "encoding/json"
   8         "errors"
   9         "fmt"
  10         "io"
  11         "net"
  12         "os"
  13         "path/filepath"
  14         "sync"
  15         "time"
  16 )
  17
  18 const (
  19         // daemonAgentDirName is the directory of the root agent in the data
  20         // directory.
  21         daemonAgentDirName = "agent"
  22
  23         // daemonClientEvents is the number of events that can wait for a slow
  24         // client before the daemon drops it.
  25         daemonClientEvents = 1024
  26 )
  27
  28 // DaemonHub holds the connected clients and sends each event to all of them. It
  29 // runs the WhatsApp commands of the clients on whatsapp. snapshot locks all
  30 // agents and subagents during a backup copy.
  31 type DaemonHub struct {
  32         whatsapp *WhatsApp
  33         snapshot func(context.Context, func() error) error
  34
  35         mutex   sync.Mutex
  36         clients map[*DaemonClient]bool
  37         workers sync.WaitGroup
  38 }
  39
  40 // DaemonClient is one connected client. Its events wait in a channel for the
  41 // writer, so a slow client does not stop the agent.
  42 type DaemonClient struct {
  43         conn   net.Conn
  44         events chan Event
  45         once   sync.Once
  46 }
  47
  48 // Accept serves each client of listener until the listener closes. A new client
  49 // gets a ready event with the name of the agent, and a busy event when a turn
  50 // runs.
  51 func (h *DaemonHub) Accept(ctx context.Context, listener net.Listener, l *TurnLoop) {
  52         for {
  53                 conn, err := listener.Accept()
  54                 if err != nil {
  55                         return
  56                 }
  57
  58                 client := &DaemonClient{
  59                         conn:   conn,
  60                         events: make(chan Event, daemonClientEvents),
  61                 }
  62
  63                 h.mutex.Lock()
  64
  65                 h.clients[client] = true
  66                 h.mutex.Unlock()
  67
  68                 h.workers.Add(2)
  69
  70                 go h.write(client)
  71                 go h.read(ctx, client, l)
  72
  73                 h.Send(client, Event{
  74                         Kind: "ready",
  75                         Text: l.agent.config.Name,
  76                 })
  77
  78                 if l.Busy() {
  79                         h.Send(client, Event{
  80                                 Kind: "busy",
  81                         })
  82                 }
  83         }
  84 }
  85
  86 // write sends the events of client as JSON lines until its channel closes.
  87 func (h *DaemonHub) write(client *DaemonClient) {
  88         defer h.workers.Done()
  89
  90         encoder := json.NewEncoder(client.conn)
  91         for event := range client.events {
  92                 err := encoder.Encode(event)
  93                 if err != nil {
  94                         h.drop(client)
  95                 }
  96         }
  97 }
  98
  99 // read runs the requests of client until it closes its connection.
 100 func (h *DaemonHub) read(ctx context.Context, client *DaemonClient, l *TurnLoop) {
 101         defer h.workers.Done()
 102         defer h.drop(client)
 103
 104         scanner := bufio.NewScanner(client.conn)
 105         scanner.Buffer(make([]byte, 64*1024), socketRequestBytes)
 106
 107         for ctx.Err() == nil && scanner.Scan() {
 108                 var request SocketRequest
 109
 110                 err := json.Unmarshal(scanner.Bytes(), &request)
 111                 if err == nil {
 112                         switch request.Type {
 113                         case "notify":
 114                                 var receipt Event
 115
 116                                 decoder := json.NewDecoder(bytes.NewReader(scanner.Bytes()))
 117                                 decoder.DisallowUnknownFields()
 118
 119                                 err = decoder.Decode(&request)
 120                                 if err == nil {
 121                                         receipt, err = l.messages.Notify(request.Notify)
 122                                 }
 123
 124                                 if err == nil {
 125                                         h.Send(client, receipt)
 126                                 }
 127                         case "backup":
 128                                 var size int64
 129
 130                                 size, err = backupData(ctx, filepath.Dir(l.agent.dir), request.Text, h.whatsapp, h.snapshot)
 131                                 if err == nil {
 132                                         h.Send(client, Event{
 133                                                 Kind: "backup",
 134                                                 Text: fmt.Sprintf("%s (%d bytes)", request.Text, size),
 135                                         })
 136                                 }
 137                         default:
 138                                 err = daemonRunRequest(l, h.whatsapp, request, func(text string) {
 139                                         h.Send(client, Event{
 140                                                 Kind: "notice",
 141                                                 Text: text,
 142                                         })
 143                                 })
 144                         }
 145                 }
 146
 147                 if err != nil {
 148                         h.Send(client, Event{
 149                                 Kind: "error",
 150                                 Text: err.Error(),
 151                         })
 152                 }
 153         }
 154 }
 155
 156 // Send gives event to one client, or drops the client when its events do not
 157 // fit.
 158 func (h *DaemonHub) Send(client *DaemonClient, event Event) {
 159         h.mutex.Lock()
 160         defer h.mutex.Unlock()
 161
 162         h.send(client, event)
 163 }
 164
 165 // drop removes client and closes its connection.
 166 func (h *DaemonHub) drop(client *DaemonClient) {
 167         h.mutex.Lock()
 168         defer h.mutex.Unlock()
 169
 170         h.remove(client)
 171 }
 172
 173 // send gives event to client. The caller holds the mutex.
 174 func (h *DaemonHub) send(client *DaemonClient, event Event) {
 175         if !h.clients[client] {
 176                 return
 177         }
 178
 179         select {
 180         case client.events <- event:
 181         default:
 182                 h.remove(client)
 183         }
 184 }
 185
 186 // remove removes client, closes its connection, and ends its writer. The caller
 187 // holds the mutex.
 188 func (h *DaemonHub) remove(client *DaemonClient) {
 189         delete(h.clients, client)
 190         client.once.Do(func() {
 191                 close(client.events)
 192
 193                 // The client is removed already, so a close failure cannot
 194                 // restore it.
 195                 client.conn.Close() // nolint: errcheck
 196         })
 197 }
 198
 199 // Broadcast gives event to every client.
 200 func (h *DaemonHub) Broadcast(event Event) {
 201         h.mutex.Lock()
 202         defer h.mutex.Unlock()
 203
 204         for client := range h.clients {
 205                 h.send(client, event)
 206         }
 207 }
 208
 209 // CloseAll drops every client and waits for their readers and writers.
 210 func (h *DaemonHub) CloseAll() {
 211         h.mutex.Lock()
 212
 213         for client := range h.clients {
 214                 h.remove(client)
 215         }
 216
 217         h.mutex.Unlock()
 218
 219         h.workers.Wait()
 220 }
 221
 222 // daemonRunRequest runs one request of a client. The answers of a WhatsApp
 223 // command go to answer.
 224 func daemonRunRequest(l *TurnLoop, whatsapp *WhatsApp, request SocketRequest, answer func(text string)) error {
 225         switch request.Type {
 226         case "send":
 227                 if request.Text == "" {
 228                         return errors.New("empty message")
 229                 }
 230
 231                 _, err := l.Submit(terminalSource, request.Text)
 232                 if err != nil {
 233                         return fmt.Errorf("submit terminal message: %w", err)
 234                 }
 235
 236                 return nil
 237         case "image":
 238                 if !filepath.IsAbs(request.Path) {
 239                         return errors.New("use /image PATH [caption]")
 240                 }
 241
 242                 data, err := imageReadFile(request.Path)
 243                 if err != nil {
 244                         return fmt.Errorf("read terminal image: %w", err)
 245                 }
 246
 247                 name, err := l.agent.Images().Save(data)
 248                 if err != nil {
 249                         return fmt.Errorf("store terminal image: %w", err)
 250                 }
 251
 252                 _, err = l.Submit(terminalSource, imageText([]string{name}, request.Text))
 253                 if err != nil {
 254                         return fmt.Errorf("submit terminal image: %w", err)
 255                 }
 256
 257                 return nil
 258         case "interrupt":
 259                 if !l.Interrupt() {
 260                         return errors.New("no turn runs")
 261                 }
 262
 263                 return nil
 264         case "whatsapp":
 265                 whatsapp.Command(request.Text, answer)
 266
 267                 return nil
 268         case "stop":
 269                 found, err := l.StopSubagent(request.Text)
 270                 if err == nil && !found {
 271                         err = fmt.Errorf("no subagent %s runs", request.Text)
 272                 }
 273
 274                 if err != nil {
 275                         return fmt.Errorf("stop subagent: %w", err)
 276                 }
 277
 278                 return nil
 279         default:
 280                 return fmt.Errorf("unknown request %q", request.Type)
 281         }
 282 }
 283
 284 // daemon runs "onechat daemon": it owns the root, its children, and their
 285 // files, and serves the clients on the Unix socket in the data directory until
 286 // ctx ends. When telegram.json exists, it also runs the Telegram bot. It also
 287 // keeps a linked WhatsApp device connected. It holds the lock of the token
 288 // file, so it is the only process that refreshes the tokens and opens the
 289 // WhatsApp session. Its log lines go to out, through the loggers of
 290 // daemonLogger.
 291 func daemon(ctx context.Context, out io.Writer) (err error) {
 292         dir, err := dataDir()
 293         if err != nil {
 294                 return fmt.Errorf("find data directory: %w", err)
 295         }
 296
 297         // Check the socket path before the daemon writes a file, so that a data
 298         // directory with a long path stays empty.
 299         socket := filepath.Join(dir, socketFileName)
 300         err = checkSocketPath(socket)
 301         if err != nil {
 302                 return fmt.Errorf("check daemon socket path: %w", err)
 303         }
 304
 305         config, err := loadConfig(dir)
 306         if err != nil {
 307                 return fmt.Errorf("load daemon config: %w", err)
 308         }
 309
 310         botSettings, hasBot, err := loadTelegramConfig(dir)
 311         if err != nil {
 312                 return fmt.Errorf("load Telegram config: %w", err)
 313         }
 314
 315         tokens, err := lockTokens(dir)
 316         if err != nil {
 317                 return fmt.Errorf("lock login tokens: %w", err)
 318         }
 319
 320         defer func() {
 321                 err = errors.Join(err, tokens.Close())
 322         }()
 323
 324         whatsappSetDeviceProps()
 325
 326         logger := daemonLogger(out, "onechat: ")
 327         whatsapp, err := openWhatsApp(ctx, dir, time.Now, daemonLogger(out, "whatsapp: "), whatsappNewClient)
 328         if err != nil {
 329                 return fmt.Errorf("open WhatsApp: %w", err)
 330         }
 331
 332         defer func() {
 333                 err = errors.Join(err, whatsapp.Close())
 334         }()
 335
 336         address, auth, err := modelEndpoint(config, dir)
 337         if err != nil {
 338                 return fmt.Errorf("configure model endpoint: %w", err)
 339         }
 340
 341         agents, err := openDaemonAgents(ctx, dir, address, auth, config.Compactor, time.Now, logger)
 342         if err != nil {
 343                 return fmt.Errorf("open daemon agents: %w", err)
 344         }
 345
 346         defer func() {
 347                 err = errors.Join(err, agents.Close())
 348         }()
 349
 350         whatsapp.inbox.voice = &VoiceTranscriber{config: config.Voice}
 351
 352         root := agents[0]
 353         agent := root.agent
 354         schedules := root.schedules
 355
 356         listener, err := daemonListenSocket(socket)
 357         if err != nil {
 358                 return fmt.Errorf("listen for daemon clients: %w", err)
 359         }
 360
 361         linked := "not linked"
 362         if whatsapp.isLinked() {
 363                 linked = "linked"
 364         }
 365
 366         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())
 367
 368         var telegram *TelegramBot
 369
 370         if hasBot {
 371                 telegram = newTelegramBot(newTelegramAPI(telegramAPIURL, botSettings.Token), botSettings.Owner, filepath.Join(agent.dir, telegramOffsetFileName), daemonLogger(out, "telegram: "))
 372
 373                 telegram.voice = &VoiceTranscriber{config: config.Voice}
 374
 375                 logger.Info("Telegram bot started")
 376         }
 377
 378         err = daemonServe(ctx, listener, telegram, whatsapp, schedules, func(notify func(Event)) (*TurnLoop, error) {
 379                 return root.Loop(whatsapp, notify)
 380         }, agents[1:]...)
 381         if err != nil {
 382                 return fmt.Errorf("serve daemon: %w", err)
 383         }
 384
 385         return nil
 386 }
 387
 388 // daemonListenSocket listens on the Unix socket at path, with mode 0600. The
 389 // caller owns the chat, so a socket file at path is left from a daemon that
 390 // ended, and daemonListenSocket removes it.
 391 func daemonListenSocket(path string) (net.Listener, error) {
 392         err := checkSocketPath(path)
 393         if err != nil {
 394                 return nil, fmt.Errorf("check daemon socket path: %w", err)
 395         }
 396
 397         err = os.Remove(path)
 398         if err != nil && !errors.Is(err, os.ErrNotExist) {
 399                 return nil, fmt.Errorf("remove old socket: %w", err)
 400         }
 401
 402         listener, err := net.Listen("unix", path)
 403         if err != nil {
 404                 return nil, fmt.Errorf("listen on socket: %w", err)
 405         }
 406
 407         err = os.Chmod(path, 0o600)
 408         if err != nil {
 409                 return nil, errors.Join(fmt.Errorf("set socket mode: %w", err), listener.Close())
 410         }
 411
 412         return listener, nil
 413 }
 414
 415 // daemonServe runs the loop that start makes and serves the clients of listener
 416 // until ctx ends or the loop fails. It closes listener. It runs whatsapp, which
 417 // logs its batches and notes in the loop and tells all clients its notices, and
 418 // a failure of whatsapp stops the agent. It runs schedules, which log their
 419 // runs in the loop, and a failure of schedules stops the agent too. When bot is
 420 // not nil, it runs the bot too, and a failure of the bot stops the agent. Each
 421 // child runs its own loop and schedules with no client events. A failed child
 422 // stops the daemon, and backups lock all loops before they copy data.
 423 func daemonServe(ctx context.Context, listener net.Listener, bot *TelegramBot, whatsapp *WhatsApp, schedules *Schedules, start func(notify func(Event)) (*TurnLoop, error), children ...*DaemonAgent) error {
 424         ctx, cancel := context.WithCancel(ctx)
 425         defer cancel()
 426
 427         hub := &DaemonHub{
 428                 whatsapp: whatsapp,
 429                 clients:  make(map[*DaemonClient]bool),
 430         }
 431
 432         notify := hub.Broadcast
 433         if bot != nil {
 434                 notify = func(event Event) {
 435                         hub.Broadcast(event)
 436                         bot.Notify(event)
 437                 }
 438         }
 439
 440         l, err := start(notify)
 441         if err != nil {
 442                 return errors.Join(err, listener.Close())
 443         }
 444
 445         if bot != nil {
 446                 l.tools.imageShow.bot = bot
 447         }
 448
 449         loops := []*TurnLoop{l}
 450         for _, child := range children {
 451                 loop, err := child.Loop(nil, func(Event) {
 452                 })
 453                 if err != nil {
 454                         return errors.Join(fmt.Errorf("start child %q: %w", child.agent.config.Name, err), listener.Close())
 455                 }
 456
 457                 loops = append(loops, loop)
 458         }
 459
 460         messages, err := openDaemonMessages(loops)
 461         if err != nil {
 462                 return errors.Join(fmt.Errorf("open agent messages: %w", err), listener.Close())
 463         }
 464
 465         hub.snapshot = func(ctx context.Context, snapshot func() error) error {
 466                 return messages.SnapshotAgents(ctx, snapshot)
 467         }
 468
 469         runtime := &DaemonChildren{
 470                 ctx:      ctx,
 471                 cancel:   cancel,
 472                 root:     l,
 473                 messages: messages,
 474         }
 475
 476         l.tools.children = runtime
 477
 478         var workers sync.WaitGroup
 479
 480         childErrors := make([]error, len(children))
 481         for i, child := range children {
 482                 workers.Add(1)
 483
 484                 go func() {
 485                         defer workers.Done()
 486
 487                         childErrors[i] = daemonRunChild(ctx, loops[i+1], child.schedules)
 488                         if childErrors[i] != nil {
 489                                 child.logger.Error("child stopped", "err", childErrors[i])
 490                         }
 491
 492                         cancel()
 493                 }()
 494         }
 495
 496         workers.Add(2)
 497
 498         go func() {
 499                 defer workers.Done()
 500
 501                 hub.Accept(ctx, listener, l)
 502         }()
 503
 504         var whatsappErr error
 505
 506         go func() {
 507                 defer workers.Done()
 508
 509                 whatsappErr = whatsapp.Run(ctx, l, notify)
 510
 511                 cancel()
 512         }()
 513
 514         var schedulesErr error
 515
 516         workers.Add(1)
 517
 518         go func() {
 519                 defer workers.Done()
 520
 521                 schedulesErr = schedules.Run(ctx, l)
 522
 523                 cancel()
 524         }()
 525
 526         var botErr error
 527
 528         if bot != nil {
 529                 workers.Add(1)
 530
 531                 go func() {
 532                         defer workers.Done()
 533
 534                         botErr = bot.Run(ctx, l, whatsapp)
 535
 536                         cancel()
 537                 }()
 538         }
 539
 540         err = l.Run(ctx)
 541         if errors.Is(err, context.Canceled) {
 542                 err = nil
 543         }
 544
 545         cancel()
 546
 547         err = errors.Join(err, listener.Close())
 548
 549         hub.CloseAll()
 550         workers.Wait()
 551
 552         return errors.Join(err, whatsappErr, schedulesErr, botErr, errors.Join(childErrors...), runtime.Close())
 553 }
 554
 555 // daemonRunChild runs a child's turn loop and its schedules until either stops.
 556 // Its events stay in its own log.
 557 func daemonRunChild(ctx context.Context, loop *TurnLoop, schedules *Schedules) error {
 558         ctx, cancel := context.WithCancel(ctx)
 559         defer cancel()
 560
 561         done := make(chan error, 1)
 562         go func() {
 563                 done <- schedules.Run(ctx, loop)
 564
 565                 cancel()
 566         }()
 567
 568         err := loop.Run(ctx)
 569         if errors.Is(err, context.Canceled) {
 570                 err = nil
 571         }
 572
 573         cancel()
 574
 575         return errors.Join(err, <-done)
 576 }