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 holds the lock of the
  31 // token file, so it is the only process that refreshes the tokens.
  32 func daemon(ctx context.Context, out io.Writer) (err error) {
  33         dir, err := dataDir()
  34         if err != nil {
  35                 return err
  36         }
  37
  38         // Check the socket path before the daemon writes a file, so that a data
  39         // directory with a long path stays empty.
  40         socket := filepath.Join(dir, socketFileName)
  41         err = checkSocketPath(socket)
  42         if err != nil {
  43                 return err
  44         }
  45
  46         config, err := loadConfig(dir)
  47         if err != nil {
  48                 return err
  49         }
  50
  51         botSettings, hasBot, err := loadTelegramConfig(dir)
  52         if err != nil {
  53                 return err
  54         }
  55
  56         tokens, err := lockTokens(dir)
  57         if err != nil {
  58                 return err
  59         }
  60
  61         defer func() {
  62                 err = errors.Join(err, tokens.Close())
  63         }()
  64
  65         address, auth, err := modelEndpoint(config, dir)
  66         if err != nil {
  67                 return err
  68         }
  69
  70         agent, err := openAgent(filepath.Join(dir, daemonAgentDirName))
  71         if err != nil {
  72                 return err
  73         }
  74
  75         client := newModelClient(address, auth, agent.CacheKey(), filepath.Join(agent.dir, "usage.jsonl"))
  76         memory, err := openMemory(ctx, agent.dir, time.Now, newCompactor(client, agent.Prompt, config.Compactor))
  77         if err != nil {
  78                 return err
  79         }
  80
  81         defer func() {
  82                 err = errors.Join(err, memory.Close())
  83         }()
  84
  85         listener, err := daemonListenSocket(socket)
  86         if err != nil {
  87                 return err
  88         }
  89
  90         _, 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)
  91         if err != nil {
  92                 return errors.Join(err, listener.Close())
  93         }
  94
  95         var telegram *TelegramBot
  96         if hasBot {
  97                 telegram = newTelegramBot(newTelegramAPI(telegramAPIURL, botSettings.Token), botSettings.Owner, filepath.Join(agent.dir, telegramOffsetFileName), log.New(out, "telegram: ", log.LstdFlags))
  98                 _, err = fmt.Fprintln(out, "onechat: the Telegram bot waits for messages of its owner")
  99                 if err != nil {
 100                         return errors.Join(err, listener.Close())
 101                 }
 102         }
 103
 104         return daemonServe(ctx, listener, telegram, func(notify func(Event)) (*TurnLoop, error) {
 105                 return newTurnLoop(agent, memory, client, time.Now, notify)
 106         })
 107 }
 108
 109 // DaemonHub holds the connected clients and sends each event to all of them.
 110 type DaemonHub struct {
 111         mutex   sync.Mutex
 112         clients map[*DaemonClient]bool
 113         workers sync.WaitGroup
 114 }
 115
 116 // Accept serves each client of listener until the listener closes. A new
 117 // client gets a ready event with the name of the agent, and a busy event when a
 118 // turn runs.
 119 func (h *DaemonHub) Accept(ctx context.Context, listener net.Listener, l *TurnLoop) {
 120         for {
 121                 conn, err := listener.Accept()
 122                 if err != nil {
 123                         return
 124                 }
 125
 126                 client := &DaemonClient{
 127                         conn:   conn,
 128                         events: make(chan Event, daemonClientEvents),
 129                 }
 130
 131                 h.mutex.Lock()
 132                 h.clients[client] = true
 133                 h.mutex.Unlock()
 134
 135                 h.workers.Add(2)
 136                 go h.write(client)
 137                 go h.read(ctx, client, l)
 138
 139                 h.Send(client, Event{
 140                         Kind: "ready",
 141                         Text: l.agent.config.Name,
 142                 })
 143
 144                 if l.Busy() {
 145                         h.Send(client, Event{
 146                                 Kind: "busy",
 147                         })
 148                 }
 149         }
 150 }
 151
 152 // write sends the events of client as JSON lines until its channel closes.
 153 func (h *DaemonHub) write(client *DaemonClient) {
 154         defer h.workers.Done()
 155
 156         encoder := json.NewEncoder(client.conn)
 157         for event := range client.events {
 158                 err := encoder.Encode(event)
 159                 if err != nil {
 160                         h.drop(client)
 161                 }
 162         }
 163 }
 164
 165 // read runs the requests of client until it closes its connection.
 166 func (h *DaemonHub) read(ctx context.Context, client *DaemonClient, l *TurnLoop) {
 167         defer h.workers.Done()
 168         defer h.drop(client)
 169
 170         scanner := bufio.NewScanner(client.conn)
 171         scanner.Buffer(make([]byte, 64*1024), socketRequestBytes)
 172         for ctx.Err() == nil && scanner.Scan() {
 173                 var request SocketRequest
 174                 err := json.Unmarshal(scanner.Bytes(), &request)
 175                 if err == nil {
 176                         err = daemonRunRequest(l, request)
 177                 }
 178
 179                 if err != nil {
 180                         h.Send(client, Event{
 181                                 Kind: "error",
 182                                 Text: err.Error(),
 183                         })
 184                 }
 185         }
 186 }
 187
 188 // Send gives event to one client, or drops the client when its events do not
 189 // fit.
 190 func (h *DaemonHub) Send(client *DaemonClient, event Event) {
 191         h.mutex.Lock()
 192         defer h.mutex.Unlock()
 193
 194         h.send(client, event)
 195 }
 196
 197 // drop removes client and closes its connection.
 198 func (h *DaemonHub) drop(client *DaemonClient) {
 199         h.mutex.Lock()
 200         defer h.mutex.Unlock()
 201
 202         h.remove(client)
 203 }
 204
 205 // send gives event to client. The caller holds the mutex.
 206 func (h *DaemonHub) send(client *DaemonClient, event Event) {
 207         if !h.clients[client] {
 208                 return
 209         }
 210
 211         select {
 212         case client.events <- event:
 213         default:
 214                 h.remove(client)
 215         }
 216 }
 217
 218 // remove removes client, closes its connection, and ends its writer. The
 219 // caller holds the mutex.
 220 func (h *DaemonHub) remove(client *DaemonClient) {
 221         delete(h.clients, client)
 222         client.once.Do(func() {
 223                 close(client.events)
 224                 _ = client.conn.Close()
 225         })
 226 }
 227
 228 // Broadcast gives event to every client.
 229 func (h *DaemonHub) Broadcast(event Event) {
 230         h.mutex.Lock()
 231         defer h.mutex.Unlock()
 232
 233         for client := range h.clients {
 234                 h.send(client, event)
 235         }
 236 }
 237
 238 // CloseAll drops every client and waits for their readers and writers.
 239 func (h *DaemonHub) CloseAll() {
 240         h.mutex.Lock()
 241         for client := range h.clients {
 242                 h.remove(client)
 243         }
 244         h.mutex.Unlock()
 245
 246         h.workers.Wait()
 247 }
 248
 249 // DaemonClient is one connected client. Its events wait in a channel for the
 250 // writer, so a slow client does not stop the agent.
 251 type DaemonClient struct {
 252         conn   net.Conn
 253         events chan Event
 254         once   sync.Once
 255 }
 256
 257 // daemonListenSocket listens on the Unix socket at path, with mode 0600. The
 258 // caller owns the chat, so a socket file at path is left from a daemon that
 259 // ended, and daemonListenSocket removes it.
 260 func daemonListenSocket(path string) (net.Listener, error) {
 261         err := checkSocketPath(path)
 262         if err != nil {
 263                 return nil, err
 264         }
 265
 266         err = os.Remove(path)
 267         if err != nil && !errors.Is(err, os.ErrNotExist) {
 268                 return nil, fmt.Errorf("remove old socket: %w", err)
 269         }
 270
 271         listener, err := net.Listen("unix", path)
 272         if err != nil {
 273                 return nil, fmt.Errorf("listen on socket: %w", err)
 274         }
 275
 276         err = os.Chmod(path, 0o600)
 277         if err != nil {
 278                 return nil, errors.Join(fmt.Errorf("set socket mode: %w", err), listener.Close())
 279         }
 280
 281         return listener, nil
 282 }
 283
 284 // daemonServe runs the loop that start makes and serves the clients of listener
 285 // until ctx ends or the loop fails. It closes listener. When bot is not nil, it
 286 // runs the bot too, and a failure of the bot stops the agent.
 287 func daemonServe(ctx context.Context, listener net.Listener, bot *TelegramBot, start func(notify func(Event)) (*TurnLoop, error)) error {
 288         ctx, cancel := context.WithCancel(ctx)
 289         defer cancel()
 290
 291         hub := &DaemonHub{
 292                 clients: make(map[*DaemonClient]bool),
 293         }
 294
 295         notify := hub.Broadcast
 296         if bot != nil {
 297                 notify = func(event Event) {
 298                         hub.Broadcast(event)
 299                         bot.Notify(event)
 300                 }
 301         }
 302
 303         l, err := start(notify)
 304         if err != nil {
 305                 return errors.Join(err, listener.Close())
 306         }
 307
 308         var workers sync.WaitGroup
 309         workers.Add(1)
 310         go func() {
 311                 defer workers.Done()
 312
 313                 hub.Accept(ctx, listener, l)
 314         }()
 315
 316         var botErr error
 317         if bot != nil {
 318                 workers.Add(1)
 319                 go func() {
 320                         defer workers.Done()
 321
 322                         botErr = bot.Run(ctx, l)
 323                         cancel()
 324                 }()
 325         }
 326
 327         err = l.Run(ctx)
 328         if errors.Is(err, context.Canceled) {
 329                 err = nil
 330         }
 331
 332         cancel()
 333         err = errors.Join(err, listener.Close())
 334         hub.CloseAll()
 335         workers.Wait()
 336
 337         return errors.Join(err, botErr)
 338 }
 339
 340 // daemonRunRequest runs one request of a client.
 341 func daemonRunRequest(l *TurnLoop, request SocketRequest) error {
 342         switch request.Type {
 343         case "send":
 344                 if request.Text == "" {
 345                         return errors.New("empty message")
 346                 }
 347
 348                 _, err := l.Submit(terminalSource, request.Text)
 349
 350                 return err
 351         case "interrupt":
 352                 if !l.Interrupt() {
 353                         return errors.New("no turn runs")
 354                 }
 355
 356                 return nil
 357         default:
 358                 return fmt.Errorf("unknown request %q", request.Type)
 359         }
 360 }