| 1 | package agentrunner |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "errors" |
| 7 | "fmt" |
| 8 | "net/url" |
| 9 | "slices" |
| 10 | "strings" |
| 11 | "sync" |
| 12 | "time" |
| 13 | |
| 14 | "github.com/FelineStateMachine/tinyrelay/internal/event" |
| 15 | "github.com/coder/websocket" |
| 16 | ) |
| 17 | |
| 18 | type RelayClient struct { |
| 19 | URL, Secret, PubKey string |
| 20 | Rooms []string |
| 21 | OnEvent func(event.Event) |
| 22 | OnReady func() |
| 23 | OnError func(error) |
| 24 | mu sync.Mutex |
| 25 | conn *websocket.Conn |
| 26 | connection context.Context |
| 27 | pending map[string]chan error |
| 28 | queries map[string]chan queryResult |
| 29 | } |
| 30 | |
| 31 | type queryResult struct { |
| 32 | event event.Event |
| 33 | err error |
| 34 | } |
| 35 | |
| 36 | func (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 err |
| 44 | } |
| 45 | delay := time.Second |
| 46 | for ctx.Err() == nil { |
| 47 | err = c.connect(ctx) |
| 48 | if ctx.Err() != nil { |
| 49 | break |
| 50 | } |
| 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 | } |
| 65 | func (c *RelayClient) connect(ctx context.Context) error { |
| 66 | conn, _, err := websocket.Dial(ctx, c.URL, nil) |
| 67 | if err != nil { |
| 68 | return err |
| 69 | } |
| 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 = conn |
| 76 | c.connection = scope |
| 77 | 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 = nil |
| 84 | } |
| 85 | c.mu.Unlock() |
| 86 | }() |
| 87 | if err := c.subscribe(scope, conn); err != nil { |
| 88 | return err |
| 89 | } |
| 90 | authID := "" |
| 91 | for { |
| 92 | _, raw, err := conn.Read(scope) |
| 93 | if err != nil { |
| 94 | return err |
| 95 | } |
| 96 | var frame []json.RawMessage |
| 97 | if json.Unmarshal(raw, &frame) != nil || len(frame) < 2 { |
| 98 | continue |
| 99 | } |
| 100 | var kind string |
| 101 | json.Unmarshal(frame[0], &kind) |
| 102 | switch kind { |
| 103 | case "AUTH" : |
| 104 | var challenge string |
| 105 | if json.Unmarshal(frame[1], &challenge) != nil { |
| 106 | continue |
| 107 | } |
| 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 err |
| 111 | } |
| 112 | authID = auth.ID |
| 113 | if err := writeRelay(scope, conn, []any{"AUTH" , auth}); err != nil { |
| 114 | return err |
| 115 | } |
| 116 | case "OK" : |
| 117 | if len(frame) < 3 { |
| 118 | continue |
| 119 | } |
| 120 | var id string |
| 121 | var ok bool |
| 122 | var reason string |
| 123 | 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 err |
| 134 | } |
| 135 | if c.OnReady != nil { |
| 136 | c.OnReady() |
| 137 | } |
| 138 | continue |
| 139 | } |
| 140 | var result error |
| 141 | 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 string |
| 155 | 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 | continue |
| 166 | } |
| 167 | } |
| 168 | if subscription == "tiny-agent" && c.OnReady != nil { |
| 169 | c.OnReady() |
| 170 | } |
| 171 | case "EVENT" : |
| 172 | if len(frame) != 3 { |
| 173 | continue |
| 174 | } |
| 175 | e, err := event.Parse(frame[2]) |
| 176 | if err != nil { |
| 177 | continue |
| 178 | } |
| 179 | c.mu.Lock() |
| 180 | var subscription string |
| 181 | _ = 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 | continue |
| 190 | } |
| 191 | if subscription == "tiny-agent" && c.OnEvent != nil { |
| 192 | c.OnEvent(e) |
| 193 | } |
| 194 | case "CLOSED" : |
| 195 | var subscription string |
| 196 | _ = 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 | continue |
| 210 | } |
| 211 | if subscription != "tiny-agent" { |
| 212 | continue |
| 213 | } |
| 214 | if len(frame) > 2 { |
| 215 | var reason string |
| 216 | 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 | } |
| 224 | |
| 225 | func (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] = true |
| 231 | e, err := c.queryEvent(ctx, room, id) |
| 232 | if err != nil { |
| 233 | return "" , err |
| 234 | } |
| 235 | parent := event.RoomReplyRoot(e) |
| 236 | if parent == "" { |
| 237 | return e.ID, nil |
| 238 | } |
| 239 | id = parent |
| 240 | } |
| 241 | return "" , errors.New("thread nesting limit exceeded" ) |
| 242 | } |
| 243 | |
| 244 | func (c *RelayClient) queryEvent(ctx context.Context, room, id string) (event.Event, error) { |
| 245 | c.mu.Lock() |
| 246 | conn := c.conn |
| 247 | scope := c.connection |
| 248 | 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] = wait |
| 255 | c.mu.Unlock() |
| 256 | defer func() { |
| 257 | c.mu.Lock() |
| 258 | delete(c.queries, sub) |
| 259 | active := c.conn == conn && c.connection == scope |
| 260 | 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{}, err |
| 269 | } |
| 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.err |
| 277 | } |
| 278 | e := result.event |
| 279 | 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, nil |
| 289 | case <-ctx.Done(): |
| 290 | return event.Event{}, ctx.Err() |
| 291 | case <-scope.Done(): |
| 292 | return event.Event{}, errors.New("agent relay disconnected" ) |
| 293 | } |
| 294 | } |
| 295 | func (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 | } |
| 299 | func writeRelay(ctx context.Context, conn *websocket.Conn, value any) error { |
| 300 | raw, err := json.Marshal(value) |
| 301 | if err != nil { |
| 302 | return err |
| 303 | } |
| 304 | return conn.Write(ctx, websocket.MessageText, raw) |
| 305 | } |
| 306 | func (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, err |
| 312 | } |
| 313 | c.mu.Lock() |
| 314 | conn, scope := c.conn, c.connection |
| 315 | 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] = wait |
| 321 | 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, err |
| 327 | } |
| 328 | select { |
| 329 | case err := <-wait: |
| 330 | return e, err |
| 331 | case <-ctx.Done(): |
| 332 | return e, ctx.Err() |
| 333 | case <-scope.Done(): |
| 334 | return e, errors.New("agent relay disconnected before acknowledgement" ) |
| 335 | } |
| 336 | } |
| 337 | func (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{}, false |
| 340 | } |
| 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{}, false |
| 343 | } |
| 344 | root := event.RoomReplyRoot(e) |
| 345 | if root == "" { |
| 346 | root = e.ID |
| 347 | } |
| 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}, true |
| 349 | } |