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