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