byre/internal/server/sse.go

package server

import (
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"g.fc.run/ai/byre/internal/api"
)

const backfillPage = 500

func (s *Server) StreamEvents(w http.ResponseWriter, r *http.Request, params api.StreamEventsParams) {
	last := int64(0)
	if params.After != nil {
		last = *params.After
	}
	if params.LastEventID != nil {
		last = *params.LastEventID
	}
	viewer := actor(r)
	lead := params.Lead != nil && *params.Lead
	taskID := ""
	if params.TaskId != nil {
		taskID = *params.TaskId
		_, err := s.store.GetTask(r.Context(), taskID)
		if err != nil {
			writeStoreError(w, err)
			return
		}
	}
	backfill := func(after int64) ([]api.Event, error) {
		switch {
		case lead:
			return s.store.LeadEvents(r.Context(), viewer, after, backfillPage)
		case taskID != "":
			return s.store.TaskEvents(r.Context(), taskID, after, backfillPage)
		}
		return s.store.EventsAfter(r.Context(), viewer, after, backfillPage)
	}
	wanted := func(ev api.Event) bool {
		owned := ev.Owner != nil && *ev.Owner == viewer
		switch {
		case lead:
			return owned && ev.AgentSessionId != nil
		case taskID != "":
			return ev.TaskId != nil && *ev.TaskId == taskID
		}
		return ev.Owner == nil || owned
	}

	rc := http.NewResponseController(w)
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.WriteHeader(http.StatusOK)

	live := s.hub.subscribe()
	defer s.clientsChanged()
	defer s.hub.unsubscribe(live)
	s.clientsChanged()

	for {
		events, err := backfill(last)
		if err != nil {
			return
		}
		for _, ev := range events {
			if !writeEvent(w, ev) {
				return
			}
			last = ev.Seq
		}
		if len(events) < backfillPage {
			break
		}
	}

	_, err := fmt.Fprint(w, "event: ready\ndata: ready\n\n")
	if err != nil {
		return
	}
	if rc.Flush() != nil {
		return
	}

	ping := time.NewTicker(15 * time.Second)
	defer ping.Stop()

	for {
		select {
		case ev, ok := <-live:
			if !ok {
				return
			}
			if ev.Seq != 0 && ev.Seq <= last {
				continue
			}
			if !wanted(ev) {
				continue
			}
			if !writeEvent(w, ev) {
				return
			}
			if ev.Seq != 0 {
				last = ev.Seq
			}

		case <-ping.C:
			_, err := fmt.Fprint(w, "event: ping\ndata: ping\n\n")
			if err != nil {
				return
			}

		case <-r.Context().Done():
			return
		}

		if rc.Flush() != nil {
			return
		}
	}
}

func writeEvent(w http.ResponseWriter, ev api.Event) bool {
	data, err := json.Marshal(ev)
	if err != nil {
		return false
	}
	if ev.Seq == 0 {
		_, err = fmt.Fprintf(w, "data: %s\n\n", data)
		return err == nil
	}
	_, err = fmt.Fprintf(w, "id: %d\ndata: %s\n\n", ev.Seq, data)
	return err == nil
}