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