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