l-ai-agent-kaya/internal/tasksws/client.go

// Package tasksws is a reconnecting WebSocket client for l-ai-tasks board events.
package tasksws

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"math/rand"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"github.com/gorilla/websocket"
)

// Event mirrors l-ai-tasks realtime.Event JSON.
type Event struct {
	Type     string    `json:"type"`
	Seq      int64     `json:"seq"`
	At       time.Time `json:"at"`
	Project  *Ref      `json:"project,omitempty"`
	Task     *TaskRef  `json:"task,omitempty"`
	Comment  *Comment  `json:"comment,omitempty"`
	Mentions []Mention `json:"mentions,omitempty"`
}

type Ref struct {
	ID   string `json:"id"`
	Slug string `json:"slug"`
	Name string `json:"name"`
}

type TaskRef struct {
	ID         string `json:"id"`
	Title      string `json:"title"`
	Status     string `json:"status"`
	AssigneeID string `json:"assignee_id,omitempty"`
}

type Comment struct {
	ID         string `json:"id"`
	Body       string `json:"body"`
	AuthorType string `json:"author_type"`
	AuthorID   string `json:"author_id"`
}

type Mention struct {
	Type string `json:"type"`
	ID   string `json:"id"`
}

// Handler is called for each board event (after ready).
type Handler func(ctx context.Context, ev Event) error

// Client reconnects forever until ctx cancel.
type Client struct {
	URL    string
	Token  string
	Agent  string // logical id for filtering (nora, coordinator, kaya)
	Log    *slog.Logger
	OnEvent Handler
}

// Run blocks until ctx is done.
func (c *Client) Run(ctx context.Context) {
	log := c.Log
	if log == nil {
		log = slog.Default()
	}
	var since int64
	attempt := 0
	for {
		if ctx.Err() != nil {
			return
		}
		err := c.session(ctx, &since)
		if ctx.Err() != nil {
			return
		}
		attempt++
		delay := backoff(attempt)
		log.Warn("tasks ws disconnected; reconnecting", "error", err, "attempt", attempt, "delay", delay, "since", since)
		select {
		case <-ctx.Done():
			return
		case <-time.After(delay):
		}
	}
}

func (c *Client) session(ctx context.Context, since *int64) error {
	u, err := url.Parse(c.URL)
	if err != nil {
		return err
	}
	q := u.Query()
	q.Set("since", strconv.FormatInt(*since, 10))
	u.RawQuery = q.Encode()

	hdr := http.Header{}
	if c.Token != "" {
		hdr.Set("Authorization", "Bearer "+c.Token)
	}
	dialer := websocket.Dialer{HandshakeTimeout: 15 * time.Second}
	conn, _, err := dialer.DialContext(ctx, u.String(), hdr)
	if err != nil {
		return err
	}
	defer conn.Close()

	_ = conn.SetReadDeadline(time.Now().Add(90 * time.Second))
	conn.SetPongHandler(func(string) error {
		_ = conn.SetReadDeadline(time.Now().Add(90 * time.Second))
		return nil
	})

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}
		_, data, err := conn.ReadMessage()
		if err != nil {
			return err
		}
		_ = conn.SetReadDeadline(time.Now().Add(90 * time.Second))

		var probe struct {
			Type string `json:"type"`
			Seq  int64  `json:"seq"`
		}
		if err := json.Unmarshal(data, &probe); err != nil {
			continue
		}
		if probe.Type == "ready" {
			continue
		}
		var ev Event
		if err := json.Unmarshal(data, &ev); err != nil {
			continue
		}
		if ev.Seq > *since {
			*since = ev.Seq
		}
		if c.OnEvent != nil {
			if err := c.OnEvent(ctx, ev); err != nil && c.Log != nil {
				c.Log.Warn("tasks ws handler", "error", err, "type", ev.Type)
			}
		}
	}
}

func backoff(attempt int) time.Duration {
	if attempt < 1 {
		attempt = 1
	}
	if attempt > 6 {
		attempt = 6
	}
	base := 200 * time.Millisecond
	d := base * time.Duration(1<<uint(attempt-1))
	if d > 5*time.Second {
		d = 5 * time.Second
	}
	jitter := time.Duration(rand.Int63n(int64(d/2) + 1))
	return d/2 + jitter
}

// MentionsSelf reports whether ev.mentions includes agentID.
func MentionsSelf(ev Event, agentID string) bool {
	for _, m := range ev.Mentions {
		if m.Type == "agent" && m.ID == agentID {
			return true
		}
	}
	return false
}

// ErrAuth is returned conceptually for logging; dial errors wrap status.
func FormatDial(url, agent string) string {
	return fmt.Sprintf("tasksws agent=%s url=%s", agent, url)
}