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