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