1package agentrunner23import (4 "context"5 "encoding/json"6 "errors"7 "fmt"8 "net/url"9 "slices"10 "strings"11 "sync"12 "time"1314 "github.com/FelineStateMachine/tinyrelay/internal/event"15 "github.com/coder/websocket"16)1718type RelayClient struct {19 URL, Secret, PubKey string20 Rooms []string21 OnEvent func(event.Event)22 OnReady func()23 OnError func(error)24 mu sync.Mutex25 conn *websocket.Conn26 connection context.Context27 pending map[string]chan error28 queries map[string]chan queryResult29}3031type queryResult struct {32 event event.Event33 err error34}3536func (c *RelayClient) Run(ctx context.Context) error {37 parsed, err := url.Parse(c.URL)38 if err != nil || parsed.Host == "" || (parsed.Scheme != "ws" && parsed.Scheme != "wss") || parsed.User != nil || parsed.Fragment != "" {39 return errors.New("agent relay must be a ws or wss URL")40 }41 c.PubKey, err = event.PublicKey(c.Secret)42 if err != nil {43 return err44 }45 delay := time.Second46 for ctx.Err() == nil {47 err = c.connect(ctx)48 if ctx.Err() != nil {49 break50 }51 if c.OnError != nil {52 c.OnError(err)53 }54 timer := time.NewTimer(delay)55 select {56 case <-ctx.Done():57 timer.Stop()58 return ctx.Err()59 case <-timer.C:60 }61 delay = min(delay*2, 30*time.Second)62 }63 return ctx.Err()64}65func (c *RelayClient) connect(ctx context.Context) error {66 conn, _, err := websocket.Dial(ctx, c.URL, nil)67 if err != nil {68 return err69 }70 conn.SetReadLimit(1 << 20)71 scope, cancel := context.WithCancel(ctx)72 defer cancel()73 defer conn.CloseNow()74 c.mu.Lock()75 c.conn = conn76 c.connection = scope77 c.pending = map[string]chan error{}78 c.queries = make(map[string]chan queryResult)79 c.mu.Unlock()80 defer func() {81 c.mu.Lock()82 if c.conn == conn {83 c.conn = nil84 }85 c.mu.Unlock()86 }()87 if err := c.subscribe(scope, conn); err != nil {88 return err89 }90 authID := ""91 for {92 _, raw, err := conn.Read(scope)93 if err != nil {94 return err95 }96 var frame []json.RawMessage97 if json.Unmarshal(raw, &frame) != nil || len(frame) < 2 {98 continue99 }100 var kind string101 json.Unmarshal(frame[0], &kind)102 switch kind {103 case "AUTH":104 var challenge string105 if json.Unmarshal(frame[1], &challenge) != nil {106 continue107 }108 auth := event.Event{Kind: 22242, CreatedAt: time.Now().Unix(), Tags: [][]string{{"relay", c.URL}, {"challenge", challenge}}}109 if err := event.Sign(&auth, c.Secret); err != nil {110 return err111 }112 authID = auth.ID113 if err := writeRelay(scope, conn, []any{"AUTH", auth}); err != nil {114 return err115 }116 case "OK":117 if len(frame) < 3 {118 continue119 }120 var id string121 var ok bool122 var reason string123 json.Unmarshal(frame[1], &id)124 json.Unmarshal(frame[2], &ok)125 if len(frame) > 3 {126 json.Unmarshal(frame[3], &reason)127 }128 if id == authID {129 if !ok {130 return fmt.Errorf("agent authentication refused: %s", reason)131 }132 if err := c.subscribe(scope, conn); err != nil {133 return err134 }135 if c.OnReady != nil {136 c.OnReady()137 }138 continue139 }140 var result error141 if !ok && !strings.HasPrefix(reason, "duplicate:") {142 result = fmt.Errorf("relay refused event: %s", reason)143 }144 c.mu.Lock()145 wait := c.pending[id]146 c.mu.Unlock()147 if wait != nil {148 select {149 case wait <- result:150 default:151 }152 }153 case "EOSE":154 var subscription string155 if len(frame) > 1 {156 _ = json.Unmarshal(frame[1], &subscription)157 c.mu.Lock()158 query := c.queries[subscription]159 c.mu.Unlock()160 if query != nil {161 select {162 case query <- queryResult{err: errors.New("thread parent not found")}:163 default:164 }165 continue166 }167 }168 if subscription == "tiny-agent" && c.OnReady != nil {169 c.OnReady()170 }171 case "EVENT":172 if len(frame) != 3 {173 continue174 }175 e, err := event.Parse(frame[2])176 if err != nil {177 continue178 }179 c.mu.Lock()180 var subscription string181 _ = json.Unmarshal(frame[1], &subscription)182 query := c.queries[subscription]183 c.mu.Unlock()184 if query != nil {185 select {186 case query <- queryResult{event: e}:187 default:188 }189 continue190 }191 if subscription == "tiny-agent" && c.OnEvent != nil {192 c.OnEvent(e)193 }194 case "CLOSED":195 var subscription string196 _ = json.Unmarshal(frame[1], &subscription)197 c.mu.Lock()198 query := c.queries[subscription]199 c.mu.Unlock()200 if query != nil {201 reason := "agent query closed"202 if len(frame) > 2 {203 _ = json.Unmarshal(frame[2], &reason)204 }205 select {206 case query <- queryResult{err: errors.New(reason)}:207 default:208 }209 continue210 }211 if subscription != "tiny-agent" {212 continue213 }214 if len(frame) > 2 {215 var reason string216 json.Unmarshal(frame[2], &reason)217 if !strings.HasPrefix(reason, "auth-required:") {218 return fmt.Errorf("agent subscription closed: %s", reason)219 }220 }221 }222 }223}224225func (c *RelayClient) ResolveRoot(ctx context.Context, room, id string) (string, error) {226 ctx, cancel := context.WithTimeout(ctx, 15*time.Second)227 defer cancel()228 seen := map[string]bool{}229 for len(seen) < 64 && !seen[id] {230 seen[id] = true231 e, err := c.queryEvent(ctx, room, id)232 if err != nil {233 return "", err234 }235 parent := event.RoomReplyRoot(e)236 if parent == "" {237 return e.ID, nil238 }239 id = parent240 }241 return "", errors.New("thread nesting limit exceeded")242}243244func (c *RelayClient) queryEvent(ctx context.Context, room, id string) (event.Event, error) {245 c.mu.Lock()246 conn := c.conn247 scope := c.connection248 if conn == nil {249 c.mu.Unlock()250 return event.Event{}, errors.New("agent relay is disconnected")251 }252 sub := fmt.Sprintf("root-%d", time.Now().UnixNano())253 wait := make(chan queryResult, 1)254 c.queries[sub] = wait255 c.mu.Unlock()256 defer func() {257 c.mu.Lock()258 delete(c.queries, sub)259 active := c.conn == conn && c.connection == scope260 c.mu.Unlock()261 if active {262 closeCtx, cancel := context.WithTimeout(context.Background(), time.Second)263 defer cancel()264 _ = writeRelay(closeCtx, conn, []any{"CLOSE", sub})265 }266 }()267 if err := writeRelay(ctx, conn, []any{"REQ", sub, map[string]any{"ids": []string{id}, "#h": []string{room}, "limit": 1}}); err != nil {268 return event.Event{}, err269 }270 select {271 case result, ok := <-wait:272 if !ok {273 return event.Event{}, errors.New("thread parent not found")274 }275 if result.err != nil {276 return event.Event{}, result.err277 }278 e := result.event279 if e.ID != id {280 return event.Event{}, errors.New("thread parent ID mismatch")281 }282 if event.Tag(e, "h") != room {283 return event.Event{}, errors.New("thread parent is in another room")284 }285 if e.Kind != 9 && e.Kind != 11 && e.Kind != 12 && e.Kind != 40002 {286 return event.Event{}, errors.New("thread parent has an invalid kind")287 }288 return e, nil289 case <-ctx.Done():290 return event.Event{}, ctx.Err()291 case <-scope.Done():292 return event.Event{}, errors.New("agent relay disconnected")293 }294}295func (c *RelayClient) subscribe(ctx context.Context, conn *websocket.Conn) error {296 since := time.Now().Add(-5 * time.Minute).Unix()297 return writeRelay(ctx, conn, []any{"REQ", "tiny-agent", map[string]any{"kinds": []int{9, 11, 12}, "#h": c.Rooms, "#p": []string{c.PubKey}, "since": since}, map[string]any{"kinds": []int{7, 1111}, "#p": []string{c.PubKey}, "since": since}})298}299func writeRelay(ctx context.Context, conn *websocket.Conn, value any) error {300 raw, err := json.Marshal(value)301 if err != nil {302 return err303 }304 return conn.Write(ctx, websocket.MessageText, raw)305}306func (c *RelayClient) Publish(ctx context.Context, e event.Event) (event.Event, error) {307 if e.CreatedAt == 0 {308 e.CreatedAt = time.Now().Unix()309 }310 if err := event.Sign(&e, c.Secret); err != nil {311 return e, err312 }313 c.mu.Lock()314 conn, scope := c.conn, c.connection315 if conn == nil {316 c.mu.Unlock()317 return e, errors.New("agent relay is disconnected")318 }319 wait := make(chan error, 1)320 c.pending[e.ID] = wait321 c.mu.Unlock()322 defer func() { c.mu.Lock(); delete(c.pending, e.ID); c.mu.Unlock() }()323 ctx, cancel := context.WithTimeout(ctx, 15*time.Second)324 defer cancel()325 if err := writeRelay(ctx, conn, []any{"EVENT", e}); err != nil {326 return e, err327 }328 select {329 case err := <-wait:330 return e, err331 case <-ctx.Done():332 return e, ctx.Err()333 case <-scope.Done():334 return e, errors.New("agent relay disconnected before acknowledgement")335 }336}337func (c *RelayClient) Mention(e event.Event) (Mention, bool) {338 if e.PubKey == c.PubKey || !slices.Contains([]int{9, 11, 12}, e.Kind) || !slices.Contains(event.TagValues(e, "p"), c.PubKey) || !slices.Contains(c.Rooms, event.Tag(e, "h")) {339 return Mention{}, false340 }341 if e.CreatedAt > time.Now().Add(time.Minute).Unix() || e.CreatedAt < time.Now().Add(-24*time.Hour).Unix() || len(e.Content) > 64000 {342 return Mention{}, false343 }344 root := event.RoomReplyRoot(e)345 if root == "" {346 root = e.ID347 }348 return Mention{Room: event.Tag(e, "h"), EventID: e.ID, Author: e.PubKey, Prompt: e.Content, Root: root, Kind: e.Kind, CreatedAt: e.CreatedAt}, true349}