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