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

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