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