clone

a self-hosted multitenant nostr relay with sites, file storage, git hosting, and inbox and outbox delivery

main
tinyrelay / internal / agentrunner / service.go
service.go 7.5K
1package agentrunner23import (4 "context"5 "errors"6 "fmt"7 "slices"8 "strconv"9 "strings"10 "sync"11 "time"1213 "github.com/FelineStateMachine/tinyrelay/internal/event"14 "github.com/google/uuid"15)1617type ServiceOptions struct {18 RelayURL, Secret, StatePath string19 Rooms []string20 Process ProcessOptions21 Timeout time.Duration22 Log func(error)23}24type service struct {25 options ServiceOptions26 relay RelayClient27 queue *Runner28 ctx context.Context29 mu sync.Mutex30 decisions map[string]chan event.Event31}3233func Serve(ctx context.Context, opts ServiceOptions) error {34 if opts.Timeout <= 0 {35 opts.Timeout = 30 * time.Minute36 }37 pubkey, err := event.PublicKey(opts.Secret)38 if err != nil {39 return err40 }41 s := &service{options: opts, ctx: ctx, decisions: map[string]chan event.Event{}, relay: RelayClient{URL: opts.RelayURL, Secret: opts.Secret, PubKey: pubkey, Rooms: opts.Rooms, OnError: opts.Log}}42 queue, err := New(Options{AllowedRooms: opts.Rooms, StatePath: opts.StatePath, Handle: s.handle, OnError: opts.Log})43 if err != nil {44 return err45 }46 s.queue = queue47 s.relay.OnEvent = s.receive48 // Start queues only after a connection is established, so journal recovery49 // can publish status before it invokes the agent process.50 s.relay.OnReady = func() { queue.Start(ctx) }51 defer queue.Stop()52 err = s.relay.Run(ctx)53 if errors.Is(err, context.Canceled) {54 return nil55 }56 return err57}58func (s *service) receive(e event.Event) {59 s.mu.Lock()60 for _, id := range event.TagValues(e, "e") {61 if wait := s.decisions[id]; wait != nil {62 select {63 case wait <- e:64 default:65 }66 }67 }68 s.mu.Unlock()69 m, ok := s.relay.Mention(e)70 if !ok {71 return72 }73 if isCancelCommand(m.Prompt) {74 s.queue.CancelFrom(m.Room, m.Author)75 return76 }77 if err := s.queue.Enqueue(m); err != nil && !errors.Is(err, ErrDuplicate) && s.options.Log != nil {78 s.options.Log(err)79 }80}81func (s *service) handle(parent context.Context, batch []Mention) error {82 ctx, cancel := context.WithTimeout(parent, s.options.Timeout)83 defer cancel()84 first := batch[0]85 root, err := s.relay.ResolveRoot(ctx, first.Room, first.Root)86 if err != nil {87 return err88 }89 first.Root = root90 // Keep the canonical root on the batch so permission requests emitted91 // while this task runs stay attached to the same top-level conversation.92 batch[0].Root = root93 // The runner records its work as a long task it assigns to itself, so94 // the room's task card follows progress, the result or the error.95 tags := [][]string{{"h", first.Room}, {"e", first.Root, "", "root"}, {"p", s.relay.PubKey}, {"subject", shortText(first.Prompt, 100)}, {"expiration", strconv.FormatInt(time.Now().Add(s.options.Timeout).Unix(), 10)}}96 var prompt strings.Builder97 prompt.WriteString("These are messages addressed to you in a Nostr chat. Reply to the users' request. Room: " + first.Room + "\n\n")98 for _, m := range batch {99 fmt.Fprintf(&prompt, "Author: %s\nMessage: %s\n%s\n\n", m.Author, m.EventID, m.Prompt)100 }101 request, err := s.relay.Publish(ctx, event.Event{Kind: event.KIND_JOB_REQUEST, CreatedAt: first.CreatedAt, Tags: tags, Content: shortText(first.Prompt, 1000)})102 if err != nil {103 return err104 }105 answer := func(ctx context.Context, kind int, text string) error {106 _, err := s.relay.Publish(ctx, event.Event{Kind: kind, Tags: [][]string{{"h", first.Room}, {"e", request.ID}, {"p", request.PubKey}}, Content: shortText(text, 4000)})107 return err108 }109 if err := answer(ctx, event.KIND_JOB_PROGRESS, "Working on your request"); err != nil {110 return err111 }112 process := s.options.Process113 var last time.Time114 process.Progress = func(info string) {115 if time.Since(last) < 2*time.Second {116 return117 }118 last = time.Now()119 if err := answer(ctx, event.KIND_JOB_PROGRESS, info); err != nil {120 cancel()121 }122 }123 process.Permission = func(ctx context.Context, p Permission) (string, error) { return s.permission(ctx, batch, p) }124 heartbeatCtx, stopHeartbeat := context.WithCancel(ctx)125 heartbeatDone := make(chan struct{})126 go func() { defer close(heartbeatDone); s.heartbeat(heartbeatCtx, first.Room) }()127 output, runErr := RunProcess(ctx, process, prompt.String())128 stopHeartbeat()129 <-heartbeatDone130 finishCtx, finish := context.WithTimeout(s.ctx, 10*time.Second)131 defer finish()132 finalStatus := "idle"133 if runErr != nil {134 finalStatus = "offline"135 _ = s.publishPresence(finishCtx, first.Room, finalStatus)136 message := "Agent stopped: " + runErr.Error()137 if errors.Is(runErr, context.Canceled) {138 message = "Canceled"139 }140 if publishErr := answer(finishCtx, event.KIND_JOB_ERROR, shortText(message, 1000)); publishErr != nil {141 return errors.Join(runErr, publishErr)142 }143 return runErr144 }145 _ = s.publishPresence(finishCtx, first.Room, finalStatus)146 if output == "" {147 output = "Task completed."148 }149 if err := answer(ctx, event.KIND_JOB_RESULT, output); err != nil {150 return err151 }152 replyTags := [][]string{{"h", first.Room}, {"e", first.Root, "", "root"}}153 for _, m := range batch {154 if !slices.ContainsFunc(replyTags, func(tag []string) bool { return len(tag) > 1 && tag[0] == "p" && tag[1] == m.Author }) {155 replyTags = append(replyTags, []string{"p", m.Author})156 }157 }158 _, err = s.relay.Publish(ctx, event.Event{Kind: 9, Tags: replyTags, Content: output})159 return err160}161162func isCancelCommand(prompt string) bool {163 fields := strings.Fields(prompt)164 return len(fields) > 0 && (fields[0] == "/cancel" || fields[0] == "/stop")165}166167func (s *service) publishPresence(ctx context.Context, room, status string) error {168 _, err := s.relay.Publish(ctx, event.Event{Kind: 20001, Tags: [][]string{{"h", room}}, Content: status})169 return err170}171func (s *service) heartbeat(ctx context.Context, room string) {172 timer := time.NewTicker(30 * time.Second)173 defer timer.Stop()174 for {175 _, err := s.relay.Publish(ctx, event.Event{Kind: 20001, Tags: [][]string{{"h", room}}, Content: "working"})176 if err != nil {177 return178 }179 select {180 case <-ctx.Done():181 return182 case <-timer.C:183 }184 }185}186func (s *service) permission(ctx context.Context, batch []Mention, p Permission) (string, error) {187 allow, deny := "", ""188 for _, option := range p.Options {189 if option.Kind == "allow_once" {190 allow = option.ID191 }192 if option.Kind == "reject_once" {193 deny = option.ID194 }195 }196 if allow == "" {197 return deny, nil198 }199 first := batch[0]200 expires := time.Now().Add(10 * time.Minute).Unix()201 if deadline, ok := ctx.Deadline(); ok && deadline.Unix() < expires {202 expires = deadline.Unix()203 }204 request := event.Event{Kind: 9, CreatedAt: time.Now().Unix(), Tags: [][]string{{"h", first.Room}, {"e", first.Root, "", "root"}, {"request", "approve"}, {"d", uuid.NewString()}, {"expiration", strconv.FormatInt(expires, 10)}, {"subject", shortText(p.Title, 150)}, {"p", first.Author}}, Content: "Allow this tool action once?\n\n" + shortText(p.Title, 4000) + "\n\n" + p.Details}205 if err := event.Sign(&request, s.options.Secret); err != nil {206 return "", err207 }208 wait := make(chan event.Event, 1)209 s.mu.Lock()210 s.decisions[request.ID] = wait211 s.mu.Unlock()212 defer func() { s.mu.Lock(); delete(s.decisions, request.ID); s.mu.Unlock() }()213 if _, err := s.relay.Publish(ctx, request); err != nil {214 return "", err215 }216 for {217 select {218 case <-ctx.Done():219 return "", ctx.Err()220 case answer := <-wait:221 if answer.PubKey != first.Author || answer.Kind != 7 {222 continue223 }224 switch strings.TrimSpace(answer.Content) {225 case "+", "":226 return allow, nil227 case "-":228 return deny, nil229 }230 }231 }232}233func shortText(value string, size int) string {234 runes := []rune(value)235 if len(runes) > size {236 return string(runes[:size]) + "…"237 }238 return value239}