synapbus/internal/api/agent_events.go

package api

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

	"github.com/synapbus/synapbus/internal/agents"
	"github.com/synapbus/synapbus/internal/messaging"
)

const (
	// maxAgentConnections is the maximum number of simultaneous
	// /api/agent-events connections per agent.
	maxAgentConnections = 5
	// maxAgentBacklog is the maximum number of missed events replayed on
	// reconnect; beyond this a single resync_required event is sent instead.
	maxAgentBacklog = 200
	// agentSubBuffer is the per-connection live event buffer. A connection
	// whose buffer is full is dropped rather than blocking the broadcaster.
	agentSubBuffer = 64
	// agentWriteTimeout bounds a single write to a slow client.
	agentWriteTimeout = 10 * time.Second
)

// AgentEventBacklog loads body-free metadata for messages visible to an agent,
// used to replay events after a reconnect with Last-Event-ID.
type AgentEventBacklog interface {
	ListEventMetaAfter(ctx context.Context, agentName string, afterID int64, limit int) ([]*messaging.EventMeta, error)
}

// AgentMessageEvent is the payload of a new_message event on the agent stream.
// It carries metadata only, never the message body.
type AgentMessageEvent struct {
	MessageID int64  `json:"message_id"`
	Channel   string `json:"channel,omitempty"`
	FromAgent string `json:"from_agent,omitempty"`
	ToAgent   string `json:"to_agent,omitempty"`
	Subject   string `json:"subject,omitempty"`
}

// agentSub is one connected /api/agent-events client.
type agentSub struct {
	ch chan AgentMessageEvent
}

// SetAgentEventBacklog configures the source used to replay missed events.
func (h *SSEHub) SetAgentEventBacklog(b AgentEventBacklog) {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.agentBacklog = b
}

// hasAgentSubs reports whether any /api/agent-events connection is open.
func (h *SSEHub) hasAgentSubs() bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.agentSubs) > 0
}

// BroadcastAgentMessage delivers a new_message event to every connection of the
// named agents. Connections that cannot keep up are dropped.
func (h *SSEHub) BroadcastAgentMessage(agentNames []string, ev AgentMessageEvent) {
	type slowSub struct {
		name string
		sub  *agentSub
	}
	var slow []slowSub

	h.mu.RLock()
	for _, name := range agentNames {
		for sub := range h.agentSubs[name] {
			select {
			case sub.ch <- ev:
			default:
				slow = append(slow, slowSub{name, sub})
			}
		}
	}
	h.mu.RUnlock()

	for _, s := range slow {
		h.logger.Warn("dropping slow agent SSE client", "agent", s.name)
		h.removeAgentSub(s.name, s.sub)
	}
}

// addAgentSub registers a connection for the agent. It returns nil when the
// agent already has maxAgentConnections open connections.
func (h *SSEHub) addAgentSub(agentName string) *agentSub {
	h.mu.Lock()
	defer h.mu.Unlock()

	set := h.agentSubs[agentName]
	if len(set) >= maxAgentConnections {
		return nil
	}
	if set == nil {
		set = make(map[*agentSub]struct{})
		h.agentSubs[agentName] = set
	}
	sub := &agentSub{ch: make(chan AgentMessageEvent, agentSubBuffer)}
	set[sub] = struct{}{}
	h.logger.Info("agent SSE client connected", "agent", agentName)
	return sub
}

func (h *SSEHub) removeAgentSub(agentName string, sub *agentSub) {
	h.mu.Lock()
	defer h.mu.Unlock()

	set := h.agentSubs[agentName]
	if _, ok := set[sub]; ok {
		delete(set, sub)
		close(sub.ch)
		h.logger.Info("agent SSE client disconnected", "agent", agentName)
	}
	if len(set) == 0 {
		delete(h.agentSubs, agentName)
	}
}

// closeAgentSubsLocked disconnects all agent subscriptions. Caller holds h.mu.
func (h *SSEHub) closeAgentSubsLocked() {
	for name, set := range h.agentSubs {
		for sub := range set {
			close(sub.ch)
		}
		delete(h.agentSubs, name)
	}
}

// HandleAgentEvents handles GET /api/agent-events. It must be mounted behind
// the agent API-key middleware; the authenticated agent is read from the
// request context. Only metadata for messages visible to that agent is sent.
func (h *SSEHub) HandleAgentEvents(w http.ResponseWriter, r *http.Request) {
	agent, ok := agents.AgentFromContext(r.Context())
	if !ok || agent == nil {
		writeAgentEventsError(w, http.StatusUnauthorized, "unauthorized", "An agent API key is required")
		return
	}

	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "Streaming not supported", http.StatusInternalServerError)
		return
	}

	var lastID int64
	hasLast := false
	if v := r.Header.Get("Last-Event-ID"); v != "" {
		id, err := strconv.ParseInt(v, 10, 64)
		if err != nil || id < 0 {
			writeAgentEventsError(w, http.StatusBadRequest, "bad_request", "Last-Event-ID must be a message id")
			return
		}
		lastID, hasLast = id, true
	}

	// Subscribe before replaying so no message falls between the replay query
	// and the live stream; duplicates are removed by id below.
	sub := h.addAgentSub(agent.Name)
	if sub == nil {
		writeAgentEventsError(w, http.StatusTooManyRequests, "too_many_connections",
			fmt.Sprintf("Agent %s already has the maximum of %d event stream connections", agent.Name, maxAgentConnections))
		return
	}
	defer h.removeAgentSub(agent.Name, sub)

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("X-Accel-Buffering", "no")

	rc := http.NewResponseController(w)
	write := func(id int64, eventType string, data any) bool {
		_ = rc.SetWriteDeadline(time.Now().Add(agentWriteTimeout))
		return writeSSEEvent(w, flusher, id, eventType, data)
	}

	if !write(0, "connected", map[string]any{
		"agent":     agent.Name,
		"timestamp": time.Now().Format(time.RFC3339),
	}) {
		return
	}

	// replayedUpTo is fixed once the replay ends: live events at or below it
	// are duplicates of the replay. It is never advanced by live events, because
	// concurrent senders may broadcast ids out of order (11 before 10).
	replayedUpTo := lastID
	if hasLast {
		h.mu.RLock()
		backlog := h.agentBacklog
		h.mu.RUnlock()

		if backlog != nil {
			metas, err := backlog.ListEventMetaAfter(r.Context(), agent.Name, lastID, maxAgentBacklog+1)
			if err != nil || len(metas) > maxAgentBacklog {
				if err != nil {
					h.logger.Warn("agent SSE backlog failed", "agent", agent.Name, "error", err)
				}
				if !write(0, "resync_required", map[string]any{
					"after_id": lastID,
					"limit":    maxAgentBacklog,
				}) {
					return
				}
			} else {
				for _, m := range metas {
					ev := AgentMessageEvent{MessageID: m.MessageID, Subject: m.Subject}
					if m.Channel != "" {
						ev.Channel = m.Channel
					} else {
						ev.FromAgent, ev.ToAgent = m.FromAgent, m.ToAgent
					}
					if !write(m.MessageID, "new_message", ev) {
						return
					}
					replayedUpTo = m.MessageID
				}
			}
		}
	}

	interval := h.heartbeat
	if interval <= 0 {
		interval = 30 * time.Second
	}
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	ctx := r.Context()
	for {
		select {
		case <-ctx.Done():
			return
		case ev, ok := <-sub.ch:
			if !ok {
				return
			}
			if ev.MessageID <= replayedUpTo {
				continue
			}
			if !write(ev.MessageID, "new_message", ev) {
				return
			}
		case <-ticker.C:
			if !write(0, "heartbeat", map[string]any{
				"timestamp": time.Now().Format(time.RFC3339),
			}) {
				return
			}
		}
	}
}

// writeAgentEventsError writes a JSON error body encoded with encoding/json.
func writeAgentEventsError(w http.ResponseWriter, status int, code, message string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(map[string]string{"error": code, "message": message})
}

// writeSSEEvent writes one SSE frame. A positive id is emitted as the "id:"
// line. It reports whether the write succeeded.
func writeSSEEvent(w http.ResponseWriter, flusher http.Flusher, id int64, eventType string, data any) bool {
	jsonData, err := json.Marshal(data)
	if err != nil {
		return false
	}
	if id > 0 {
		if _, err := fmt.Fprintf(w, "id: %d\n", id); err != nil {
			return false
		}
	}
	if _, err := fmt.Fprintf(w, "event: %s\ndata: %s\n\n", eventType, jsonData); err != nil {
		return false
	}
	flusher.Flush()
	return true
}