1// Package agentrunner runs an optional agent process outside the relay daemon.2package agentrunner34import (5 "context"6 "encoding/json"7 "errors"8 "fmt"9 "os"10 "path/filepath"11 "sync"12 "time"13)1415var (16 ErrDuplicate = errors.New("agent runner: duplicate event")17 ErrUnauthorizedRoom = errors.New("agent runner: unauthorized room")18 ErrQueueFull = errors.New("agent runner: room queue full")19 ErrStopped = errors.New("agent runner: stopped")20)2122type Mention struct {23 Room string `json:"room"`24 EventID string `json:"event_id"`25 Author string `json:"author"`26 Prompt string `json:"prompt"`27 Root string `json:"root"`28 Kind int `json:"kind"`29 CreatedAt int64 `json:"created_at"`30}31type Handler func(context.Context, []Mention) error32type Options struct {33 AllowedRooms []string34 QueueLimit int35 StatePath string36 Handle Handler37 OnError func(error)38}39type roomQueue struct {40 pending, active []Mention41 cancel context.CancelFunc42 wake chan struct{}43}44type journal struct {45 Pending map[string][]Mention `json:"pending"`46 Seen []string `json:"seen"`47}48type Runner struct {49 mu sync.Mutex50 rooms map[string]*roomQueue51 seen map[string]bool52 order []string53 opts Options54 ctx context.Context55 cancel context.CancelFunc56 wg sync.WaitGroup57 stopped bool58}5960func New(opts Options) (*Runner, error) {61 if opts.Handle == nil || len(opts.AllowedRooms) == 0 || len(opts.AllowedRooms) > 32 {62 return nil, errors.New("agent runner: a handler and 1 to 32 rooms are required")63 }64 if opts.QueueLimit <= 0 {65 opts.QueueLimit = 3266 }67 if opts.QueueLimit > 256 {68 return nil, errors.New("agent runner: queue limit cannot exceed 256")69 }70 r := &Runner{opts: opts, rooms: map[string]*roomQueue{}, seen: map[string]bool{}}71 for _, room := range opts.AllowedRooms {72 if room == "" {73 return nil, ErrUnauthorizedRoom74 }75 r.rooms[room] = &roomQueue{wake: make(chan struct{}, 1)}76 }77 if err := r.restore(); err != nil {78 return nil, err79 }80 return r, nil81}82func (r *Runner) Start(ctx context.Context) {83 r.mu.Lock()84 defer r.mu.Unlock()85 if r.ctx != nil || r.stopped {86 return87 }88 r.ctx, r.cancel = context.WithCancel(ctx)89 for room, q := range r.rooms {90 r.wg.Add(1)91 go r.work(room, q)92 }93}94func (r *Runner) Enqueue(m Mention) error {95 if m.EventID == "" || m.Author == "" || len(m.Prompt) > 64*1024 {96 return errors.New("agent runner: invalid mention")97 }98 r.mu.Lock()99 defer r.mu.Unlock()100 if r.stopped {101 return ErrStopped102 }103 q, ok := r.rooms[m.Room]104 if !ok {105 return ErrUnauthorizedRoom106 }107 if r.seen[m.EventID] {108 return ErrDuplicate109 }110 if len(q.pending) >= r.opts.QueueLimit {111 return ErrQueueFull112 }113 q.pending = append(q.pending, m)114 r.seen[m.EventID] = true115 r.order = append(r.order, m.EventID)116 if err := r.save(); err != nil {117 q.pending = q.pending[:len(q.pending)-1]118 delete(r.seen, m.EventID)119 r.order = r.order[:len(r.order)-1]120 return err121 }122 select {123 case q.wake <- struct{}{}:124 default:125 }126 return nil127}128func (r *Runner) CancelFrom(room, author string) bool {129 r.mu.Lock()130 defer r.mu.Unlock()131 q := r.rooms[room]132 if q == nil {133 return false134 }135 allowed := false136 for _, m := range append(append([]Mention{}, q.active...), q.pending...) {137 if author == "" || m.Author == author {138 allowed = true139 break140 }141 }142 if !allowed {143 return false144 }145 filtered := q.pending[:0]146 for _, m := range q.pending {147 if author == "" || m.Author == author {148 continue149 }150 filtered = append(filtered, m)151 }152 q.pending = filtered153 if q.cancel != nil && anyAuthor(q.active, author) {154 q.cancel()155 }156 if err := r.save(); err != nil {157 r.report(err)158 }159 return true160}161func (r *Runner) Cancel(room string) error { r.CancelFrom(room, ""); return nil }162func (r *Runner) Stop() {163 r.mu.Lock()164 r.stopped = true165 if r.cancel != nil {166 r.cancel()167 }168 r.mu.Unlock()169 r.wg.Wait()170}171func (r *Runner) work(room string, q *roomQueue) {172 defer r.wg.Done()173 for {174 select {175 case <-r.ctx.Done():176 return177 default:178 }179 r.mu.Lock()180 if len(q.pending) == 0 {181 r.mu.Unlock()182 select {183 case <-r.ctx.Done():184 return185 case <-q.wake:186 }187 continue188 }189 batch := takeBatch(q.pending)190 q.pending = q.pending[len(batch):]191 q.active = batch192 ctx, cancel := context.WithCancel(r.ctx)193 q.cancel = cancel194 if err := r.save(); err != nil {195 r.report(err)196 cancel()197 q.pending = append(batch, q.pending...)198 q.active = nil199 q.cancel = nil200 r.mu.Unlock()201 return202 }203 r.mu.Unlock()204 err := r.opts.Handle(ctx, batch)205 cancel()206 r.mu.Lock()207 if r.ctx.Err() != nil {208 q.pending = append(batch, q.pending...)209 }210 q.active = nil211 q.cancel = nil212 if saveErr := r.save(); saveErr != nil {213 r.report(saveErr)214 }215 r.mu.Unlock()216 if err != nil && !errors.Is(err, context.Canceled) {217 r.report(fmt.Errorf("agent room %s: %w", room, err))218 }219 }220}221222func takeBatch(pending []Mention) []Mention {223 if len(pending) == 0 {224 return nil225 }226 root := pending[0].Root227 n := 1228 for n < len(pending) && pending[n].Root == root {229 n++230 }231 return append([]Mention(nil), pending[:n]...)232}233234func anyAuthor(events []Mention, author string) bool {235 if author == "" {236 return len(events) > 0237 }238 for _, event := range events {239 if event.Author == author {240 return true241 }242 }243 return false244}245func (r *Runner) report(err error) {246 if r.opts.OnError != nil {247 r.opts.OnError(err)248 }249}250func (r *Runner) restore() error {251 if r.opts.StatePath == "" {252 return nil253 }254 raw, err := os.ReadFile(r.opts.StatePath)255 if errors.Is(err, os.ErrNotExist) {256 return nil257 }258 if err != nil {259 return fmt.Errorf("read agent state: %w", err)260 }261 if len(raw) > 32<<20 {262 return errors.New("agent state exceeds 32 MiB")263 }264 var state journal265 if err := json.Unmarshal(raw, &state); err != nil {266 return fmt.Errorf("decode agent state: %w", err)267 }268 for _, id := range state.Seen {269 if !r.seen[id] {270 r.seen[id] = true271 r.order = append(r.order, id)272 }273 }274 for room, queue := range state.Pending {275 q := r.rooms[room]276 if q == nil {277 return errors.New("agent state contains an unconfigured room")278 }279 if len(queue) > r.opts.QueueLimit*2 {280 return ErrQueueFull281 }282 q.pending = queue283 }284 return nil285}286func (r *Runner) save() error {287 if r.opts.StatePath == "" {288 return nil289 }290 // Keep a bounded replay history while retaining every queued event.291 active := map[string]bool{}292 pending := map[string][]Mention{}293 for room, q := range r.rooms {294 pending[room] = append(append([]Mention{}, q.active...), q.pending...)295 for _, m := range pending[room] {296 active[m.EventID] = true297 }298 }299 if len(r.order) > 10000 {300 keep := make(map[string]bool, len(active)+10000)301 for id := range active {302 keep[id] = true303 }304 remaining := 10000 - len(active)305 if remaining < 0 {306 remaining = 0307 }308 for i := len(r.order) - 1; i >= 0 && remaining > 0; i-- {309 if !keep[r.order[i]] {310 keep[r.order[i]] = true311 remaining--312 }313 }314 kept := make([]string, 0, len(keep))315 for _, id := range r.order {316 if keep[id] {317 kept = append(kept, id)318 }319 }320 for id := range r.seen {321 if !keep[id] {322 delete(r.seen, id)323 }324 }325 r.order = kept326 }327 raw, err := json.Marshal(journal{Pending: pending, Seen: r.order})328 if err != nil {329 return err330 }331 dir := filepath.Dir(r.opts.StatePath)332 if err := os.MkdirAll(dir, 0700); err != nil {333 return err334 }335 file, err := os.CreateTemp(dir, ".tiny-agent-*")336 if err != nil {337 return err338 }339 name := file.Name()340 defer os.Remove(name)341 if _, err = file.Write(raw); err != nil {342 file.Close()343 return err344 }345 if err = file.Sync(); err != nil {346 file.Close()347 return err348 }349 if err = file.Close(); err != nil {350 return err351 }352 return os.Rename(name, r.opts.StatePath)353}354355// Drain waits until queued work has finished, or the caller cancels.356func (r *Runner) Drain(ctx context.Context) error {357 ticker := time.NewTicker(10 * time.Millisecond)358 defer ticker.Stop()359 for {360 r.mu.Lock()361 empty := true362 for _, q := range r.rooms {363 empty = empty && len(q.pending) == 0 && len(q.active) == 0364 }365 r.mu.Unlock()366 if empty {367 return nil368 }369 select {370 case <-ctx.Done():371 return ctx.Err()372 case <-ticker.C:373 }374 }375}