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.
  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         whatsapp, err := openWhatsApp(ctx, dir, time.Now, log.New(out, "whatsapp: ", log.LstdFlags), whatsappNewClient)
  68         if err != nil {
  69                 return err
  70         }
  71
  72         defer func() {
  73                 err = errors.Join(err, whatsapp.Close())
  74         }()
  75
  76         address, auth, err := modelEndpoint(config, dir)
  77         if err != nil {
  78                 return err
  79         }
  80
  81         agent, err := openAgent(filepath.Join(dir, daemonAgentDirName))
  82         if err != nil {
  83                 return err
  84         }
  85
  86         client := newModelClient(address, auth, agent.CacheKey(), filepath.Join(agent.dir, "usage.jsonl"))
  87         memory, err := openMemory(ctx, agent.dir, time.Now, newCompactor(client, agent.Prompt, config.Compactor))
  88         if err != nil {
  89                 return err
  90         }
  91
  92         defer func() {
  93                 err = errors.Join(err, memory.Close())
  94         }()
  95
  96         schedules, err := openSchedules(filepath.Join(agent.dir, scheduleFileName), time.Now)
  97         if err != nil {
  98                 return err
  99         }
 100
 101         subagents := newSubagents(filepath.Join(agent.dir, subagentDirName), time.Now, subagentNameRandom)
 102         defer func() {
 103                 err = errors.Join(err, subagents.Close())
 104         }()
 105
 106         listener, err := daemonListenSocket(socket)
 107         if err != nil {
 108                 return err
 109         }
 110
 111         _, err = fmt.Fprintf(out, "onechat: %s on %s, %s at %s effort, in %s\n", agent.config.Name, listener.Addr(), agent.config.Model, agent.config.Effort, agent.config.Workdir)
 112         if err != nil {
 113                 return errors.Join(err, listener.Close())
 114         }
 115
 116         var telegram *TelegramBot
 117         if hasBot {
 118                 telegram = newTelegramBot(newTelegramAPI(telegramAPIURL, botSettings.Token), botSettings.Owner, filepath.Join(agent.dir, telegramOffsetFileName), log.New(out, "telegram: ", log.LstdFlags))
 119                 _, err = fmt.Fprintln(out, "onechat: the Telegram bot waits for messages of its owner")
 120                 if err != nil {
 121                         return errors.Join(err, listener.Close())
 122                 }
 123         }
 124
 125         return daemonServe(ctx, listener, telegram, whatsapp, schedules, func(notify func(Event)) (*TurnLoop, error) {
 126                 return newTurnLoop(agent, memory, subagents, whatsapp, schedules, client, time.Now, notify)
 127         })
 128 }
 129
 130 // DaemonHub holds the connected clients and sends each event to all of them. It
 131 // runs the WhatsApp commands of the clients on whatsapp.
 132 type DaemonHub struct {
 133         whatsapp *WhatsApp
 134
 135         mutex   sync.Mutex
 136         clients map[*DaemonClient]bool
 137         workers sync.WaitGroup
 138 }
 139
 140 // Accept serves each client of listener until the listener closes. A new
 141 // client gets a ready event with the name of the agent, and a busy event when a
 142 // turn runs.
 143 func (h *DaemonHub) Accept(ctx context.Context, listener net.Listener, l *TurnLoop) {
 144         for {
 145                 conn, err := listener.Accept()
 146                 if err != nil {
 147                         return
 148                 }
 149
 150                 client := &DaemonClient{
 151                         conn:   conn,
 152                         events: make(chan Event, daemonClientEvents),
 153                 }
 154
 155                 h.mutex.Lock()
 156                 h.clients[client] = true
 157                 h.mutex.Unlock()
 158
 159                 h.workers.Add(2)
 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         for ctx.Err() == nil && scanner.Scan() {
 197                 var request SocketRequest
 198                 err := json.Unmarshal(scanner.Bytes(), &request)
 199                 if err == nil {
 200                         err = daemonRunRequest(l, h.whatsapp, request, func(text string) {
 201                                 h.Send(client, Event{
 202                                         Kind: "notice",
 203                                         Text: text,
 204                                 })
 205                         })
 206                 }
 207
 208                 if err != nil {
 209                         h.Send(client, Event{
 210                                 Kind: "error",
 211                                 Text: err.Error(),
 212                         })
 213                 }
 214         }
 215 }
 216
 217 // Send gives event to one client, or drops the client when its events do not
 218 // fit.
 219 func (h *DaemonHub) Send(client *DaemonClient, event Event) {
 220         h.mutex.Lock()
 221         defer h.mutex.Unlock()
 222
 223         h.send(client, event)
 224 }
 225
 226 // drop removes client and closes its connection.
 227 func (h *DaemonHub) drop(client *DaemonClient) {
 228         h.mutex.Lock()
 229         defer h.mutex.Unlock()
 230
 231         h.remove(client)
 232 }
 233
 234 // send gives event to client. The caller holds the mutex.
 235 func (h *DaemonHub) send(client *DaemonClient, event Event) {
 236         if !h.clients[client] {
 237                 return
 238         }
 239
 240         select {
 241         case client.events <- event:
 242         default:
 243                 h.remove(client)
 244         }
 245 }
 246
 247 // remove removes client, closes its connection, and ends its writer. The
 248 // caller holds the mutex.
 249 func (h *DaemonHub) remove(client *DaemonClient) {
 250         delete(h.clients, client)
 251         client.once.Do(func() {
 252                 close(client.events)
 253                 _ = client.conn.Close()
 254         })
 255 }
 256
 257 // Broadcast gives event to every client.
 258 func (h *DaemonHub) Broadcast(event Event) {
 259         h.mutex.Lock()
 260         defer h.mutex.Unlock()
 261
 262         for client := range h.clients {
 263                 h.send(client, event)
 264         }
 265 }
 266
 267 // CloseAll drops every client and waits for their readers and writers.
 268 func (h *DaemonHub) CloseAll() {
 269         h.mutex.Lock()
 270         for client := range h.clients {
 271                 h.remove(client)
 272         }
 273         h.mutex.Unlock()
 274
 275         h.workers.Wait()
 276 }
 277
 278 // DaemonClient is one connected client. Its events wait in a channel for the
 279 // writer, so a slow client does not stop the agent.
 280 type DaemonClient struct {
 281         conn   net.Conn
 282         events chan Event
 283         once   sync.Once
 284 }
 285
 286 // daemonListenSocket listens on the Unix socket at path, with mode 0600. The
 287 // caller owns the chat, so a socket file at path is left from a daemon that
 288 // ended, and daemonListenSocket removes it.
 289 func daemonListenSocket(path string) (net.Listener, error) {
 290         err := checkSocketPath(path)
 291         if err != nil {
 292                 return nil, err
 293         }
 294
 295         err = os.Remove(path)
 296         if err != nil && !errors.Is(err, os.ErrNotExist) {
 297                 return nil, fmt.Errorf("remove old socket: %w", err)
 298         }
 299
 300         listener, err := net.Listen("unix", path)
 301         if err != nil {
 302                 return nil, fmt.Errorf("listen on socket: %w", err)
 303         }
 304
 305         err = os.Chmod(path, 0o600)
 306         if err != nil {
 307                 return nil, errors.Join(fmt.Errorf("set socket mode: %w", err), listener.Close())
 308         }
 309
 310         return listener, nil
 311 }
 312
 313 // daemonServe runs the loop that start makes and serves the clients of listener
 314 // until ctx ends or the loop fails. It closes listener. It runs whatsapp, which
 315 // logs its batches and notes in the loop and tells all clients its notices, and
 316 // a failure of whatsapp stops the agent. It runs schedules, which log their
 317 // runs in the loop, and a failure of schedules stops the agent too. When bot is
 318 // not nil, it runs the bot too, and a failure of the bot stops the agent.
 319 func daemonServe(ctx context.Context, listener net.Listener, bot *TelegramBot, whatsapp *WhatsApp, schedules *Schedules, start func(notify func(Event)) (*TurnLoop, error)) error {
 320         ctx, cancel := context.WithCancel(ctx)
 321         defer cancel()
 322
 323         hub := &DaemonHub{
 324                 whatsapp: whatsapp,
 325                 clients:  make(map[*DaemonClient]bool),
 326         }
 327
 328         notify := hub.Broadcast
 329         if bot != nil {
 330                 notify = func(event Event) {
 331                         hub.Broadcast(event)
 332                         bot.Notify(event)
 333                 }
 334         }
 335
 336         l, err := start(notify)
 337         if err != nil {
 338                 return errors.Join(err, listener.Close())
 339         }
 340
 341         var workers sync.WaitGroup
 342         workers.Add(2)
 343         go func() {
 344                 defer workers.Done()
 345
 346                 hub.Accept(ctx, listener, l)
 347         }()
 348
 349         var whatsappErr error
 350         go func() {
 351                 defer workers.Done()
 352
 353                 whatsappErr = whatsapp.Run(ctx, l, notify)
 354                 cancel()
 355         }()
 356
 357         var schedulesErr error
 358         workers.Add(1)
 359         go func() {
 360                 defer workers.Done()
 361
 362                 schedulesErr = schedules.Run(ctx, l)
 363                 cancel()
 364         }()
 365
 366         var botErr error
 367         if bot != nil {
 368                 workers.Add(1)
 369                 go func() {
 370                         defer workers.Done()
 371
 372                         botErr = bot.Run(ctx, l, whatsapp)
 373                         cancel()
 374                 }()
 375         }
 376
 377         err = l.Run(ctx)
 378         if errors.Is(err, context.Canceled) {
 379                 err = nil
 380         }
 381
 382         cancel()
 383         err = errors.Join(err, listener.Close())
 384         hub.CloseAll()
 385         workers.Wait()
 386
 387         return errors.Join(err, whatsappErr, schedulesErr, botErr)
 388 }
 389
 390 // daemonRunRequest runs one request of a client. The answers of a WhatsApp
 391 // command go to answer.
 392 func daemonRunRequest(l *TurnLoop, whatsapp *WhatsApp, request SocketRequest, answer func(text string)) error {
 393         switch request.Type {
 394         case "send":
 395                 if request.Text == "" {
 396                         return errors.New("empty message")
 397                 }
 398
 399                 _, err := l.Submit(terminalSource, request.Text)
 400
 401                 return err
 402         case "interrupt":
 403                 if !l.Interrupt() {
 404                         return errors.New("no turn runs")
 405                 }
 406
 407                 return nil
 408         case "whatsapp":
 409                 whatsapp.Command(request.Text, answer)
 410
 411                 return nil
 412         case "stop":
 413                 found, err := l.StopSubagent(request.Text)
 414                 if err == nil && !found {
 415                         err = fmt.Errorf("no subagent %s runs", request.Text)
 416                 }
 417
 418                 return err
 419         default:
 420                 return fmt.Errorf("unknown request %q", request.Type)
 421         }
 422 }
 423
 424 // durationText returns d for a log line: rounded to the second, or to the
 425 // millisecond when it is shorter than a second.
 426 func durationText(d time.Duration) string {
 427         if d < time.Second {
 428                 return d.Round(time.Millisecond).String()
 429         }
 430
 431         return d.Round(time.Second).String()
 432 }