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", "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                         err = daemonRunRequest(l, h.whatsapp, request, func(text string) {
 208                                 h.Send(client, Event{
 209                                         Kind: "notice",
 210                                         Text: text,
 211                                 })
 212                         })
 213                 }
 214
 215                 if err != nil {
 216                         h.Send(client, Event{
 217                                 Kind: "error",
 218                                 Text: err.Error(),
 219                         })
 220                 }
 221         }
 222 }
 223
 224 // Send gives event to one client, or drops the client when its events do not
 225 // fit.
 226 func (h *DaemonHub) Send(client *DaemonClient, event Event) {
 227         h.mutex.Lock()
 228         defer h.mutex.Unlock()
 229
 230         h.send(client, event)
 231 }
 232
 233 // drop removes client and closes its connection.
 234 func (h *DaemonHub) drop(client *DaemonClient) {
 235         h.mutex.Lock()
 236         defer h.mutex.Unlock()
 237
 238         h.remove(client)
 239 }
 240
 241 // send gives event to client. The caller holds the mutex.
 242 func (h *DaemonHub) send(client *DaemonClient, event Event) {
 243         if !h.clients[client] {
 244                 return
 245         }
 246
 247         select {
 248         case client.events <- event:
 249         default:
 250                 h.remove(client)
 251         }
 252 }
 253
 254 // remove removes client, closes its connection, and ends its writer. The
 255 // caller holds the mutex.
 256 func (h *DaemonHub) remove(client *DaemonClient) {
 257         delete(h.clients, client)
 258         client.once.Do(func() {
 259                 close(client.events)
 260
 261                 _ = client.conn.Close()
 262         })
 263 }
 264
 265 // Broadcast gives event to every client.
 266 func (h *DaemonHub) Broadcast(event Event) {
 267         h.mutex.Lock()
 268         defer h.mutex.Unlock()
 269
 270         for client := range h.clients {
 271                 h.send(client, event)
 272         }
 273 }
 274
 275 // CloseAll drops every client and waits for their readers and writers.
 276 func (h *DaemonHub) CloseAll() {
 277         h.mutex.Lock()
 278
 279         for client := range h.clients {
 280                 h.remove(client)
 281         }
 282
 283         h.mutex.Unlock()
 284
 285         h.workers.Wait()
 286 }
 287
 288 // DaemonClient is one connected client. Its events wait in a channel for the
 289 // writer, so a slow client does not stop the agent.
 290 type DaemonClient struct {
 291         conn   net.Conn
 292         events chan Event
 293         once   sync.Once
 294 }
 295
 296 // daemonListenSocket listens on the Unix socket at path, with mode 0600. The
 297 // caller owns the chat, so a socket file at path is left from a daemon that
 298 // ended, and daemonListenSocket removes it.
 299 func daemonListenSocket(path string) (net.Listener, error) {
 300         err := checkSocketPath(path)
 301         if err != nil {
 302                 return nil, err
 303         }
 304
 305         err = os.Remove(path)
 306         if err != nil && !errors.Is(err, os.ErrNotExist) {
 307                 return nil, fmt.Errorf("remove old socket: %w", err)
 308         }
 309
 310         listener, err := net.Listen("unix", path)
 311         if err != nil {
 312                 return nil, fmt.Errorf("listen on socket: %w", err)
 313         }
 314
 315         err = os.Chmod(path, 0o600)
 316         if err != nil {
 317                 return nil, errors.Join(fmt.Errorf("set socket mode: %w", err), listener.Close())
 318         }
 319
 320         return listener, nil
 321 }
 322
 323 // daemonServe runs the loop that start makes and serves the clients of listener
 324 // until ctx ends or the loop fails. It closes listener. It runs whatsapp, which
 325 // logs its batches and notes in the loop and tells all clients its notices, and
 326 // a failure of whatsapp stops the agent. It runs schedules, which log their
 327 // runs in the loop, and a failure of schedules stops the agent too. When bot is
 328 // not nil, it runs the bot too, and a failure of the bot stops the agent.
 329 func daemonServe(ctx context.Context, listener net.Listener, bot *TelegramBot, whatsapp *WhatsApp, schedules *Schedules, start func(notify func(Event)) (*TurnLoop, error)) error {
 330         ctx, cancel := context.WithCancel(ctx)
 331         defer cancel()
 332
 333         hub := &DaemonHub{
 334                 whatsapp: whatsapp,
 335                 clients:  make(map[*DaemonClient]bool),
 336         }
 337
 338         notify := hub.Broadcast
 339         if bot != nil {
 340                 notify = func(event Event) {
 341                         hub.Broadcast(event)
 342                         bot.Notify(event)
 343                 }
 344         }
 345
 346         l, err := start(notify)
 347         if err != nil {
 348                 return errors.Join(err, listener.Close())
 349         }
 350
 351         var workers sync.WaitGroup
 352
 353         workers.Add(2)
 354
 355         go func() {
 356                 defer workers.Done()
 357
 358                 hub.Accept(ctx, listener, l)
 359         }()
 360
 361         var whatsappErr error
 362
 363         go func() {
 364                 defer workers.Done()
 365
 366                 whatsappErr = whatsapp.Run(ctx, l, notify)
 367
 368                 cancel()
 369         }()
 370
 371         var schedulesErr error
 372
 373         workers.Add(1)
 374
 375         go func() {
 376                 defer workers.Done()
 377
 378                 schedulesErr = schedules.Run(ctx, l)
 379
 380                 cancel()
 381         }()
 382
 383         var botErr error
 384
 385         if bot != nil {
 386                 workers.Add(1)
 387
 388                 go func() {
 389                         defer workers.Done()
 390
 391                         botErr = bot.Run(ctx, l, whatsapp)
 392
 393                         cancel()
 394                 }()
 395         }
 396
 397         err = l.Run(ctx)
 398         if errors.Is(err, context.Canceled) {
 399                 err = nil
 400         }
 401
 402         cancel()
 403
 404         err = errors.Join(err, listener.Close())
 405
 406         hub.CloseAll()
 407         workers.Wait()
 408
 409         return errors.Join(err, whatsappErr, schedulesErr, botErr)
 410 }
 411
 412 // daemonRunRequest runs one request of a client. The answers of a WhatsApp
 413 // command go to answer.
 414 func daemonRunRequest(l *TurnLoop, whatsapp *WhatsApp, request SocketRequest, answer func(text string)) error {
 415         switch request.Type {
 416         case "send":
 417                 if request.Text == "" {
 418                         return errors.New("empty message")
 419                 }
 420
 421                 _, err := l.Submit(terminalSource, request.Text)
 422
 423                 return err
 424         case "interrupt":
 425                 if !l.Interrupt() {
 426                         return errors.New("no turn runs")
 427                 }
 428
 429                 return nil
 430         case "whatsapp":
 431                 whatsapp.Command(request.Text, answer)
 432
 433                 return nil
 434         case "stop":
 435                 found, err := l.StopSubagent(request.Text)
 436                 if err == nil && !found {
 437                         err = fmt.Errorf("no subagent %s runs", request.Text)
 438                 }
 439
 440                 return err
 441         default:
 442                 return fmt.Errorf("unknown request %q", request.Type)
 443         }
 444 }