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