core/server/events/pg_broadcast.go

package events

// pg_broadcast.go — Postgres LISTEN/NOTIFY adapter for multi-pod fan-out.
//
// Problem: the in-process Bus is per-pod. With N backend replicas, an event
// published on pod-0 (e.g. a chat message completion) is never seen by clients
// connected to pod-1.
//
// Solution: for every event that should reach clients (DirToClient | DirBoth),
// publish via pg_notify so every pod's listener goroutine receives it and
// re-publishes into its local Bus → Hub → WebSocket clients.
//
// Key invariant: re-published (remote) events use publishLocal, which skips the
// broadcastHook, preventing infinite re-broadcast loops.
//
// No-op when running SQLite (single process, in-process Bus is sufficient).

import (
	"encoding/json"
	"fmt"
	"log"
	"time"

	"github.com/lib/pq"

	"switchboard-core/database"
)

const pgNotifyChannel = "switchboard_events"

// maxNotifyBytes is the practical PG NOTIFY payload limit (8192 byte hard limit,
// leave headroom for encoding overhead).
const maxNotifyBytes = 7900

// StartPGBroadcast wires the Bus to Postgres LISTEN/NOTIFY.
// Call once after bus and database are both initialized.
// No-op for SQLite deployments.
func StartPGBroadcast(bus *Bus) {
	if !database.IsPostgres() {
		log.Println("[pg_broadcast] SQLite mode — skipping cross-pod broadcast")
		return
	}

	// Start the listener goroutine (auto-reconnects on failure).
	go broadcastListenLoop(bus)

	// Hook into the Bus: after every local Publish, fan out to other pods.
	bus.SetBroadcastHook(func(e Event) {
		// Only events that clients care about need to cross pod boundaries.
		dir := RouteFor(e.Label)
		if dir == DirLocal {
			return
		}

		// tool.result.* arrives from a client WS on a specific pod — that pod
		// handles WaitFor() locally. No need (or benefit) to broadcast.
		if dir == DirFromClient {
			return
		}

		data, err := json.Marshal(e)
		if err != nil {
			return
		}
		if len(data) > maxNotifyBytes {
			// Event payload too large for NOTIFY — this should never happen
			// in practice (we notify IDs/metadata, not message bodies).
			log.Printf("[pg_broadcast] event %q payload %d bytes exceeds limit — skipping broadcast", e.Label, len(data))
			return
		}

		if _, err := database.DB.Exec("SELECT pg_notify($1, $2)", pgNotifyChannel, string(data)); err != nil {
			log.Printf("[pg_broadcast] pg_notify failed: %v", err)
		}
	})

	log.Println("[pg_broadcast] cross-pod broadcast enabled via Postgres LISTEN/NOTIFY")
}

// broadcastListenLoop runs forever, reconnecting on error.
func broadcastListenLoop(bus *Bus) {
	dsn := database.DSN()
	for {
		if err := runBroadcastListener(bus, dsn); err != nil {
			log.Printf("[pg_broadcast] listener error: %v — reconnecting in 5s", err)
		}
		time.Sleep(5 * time.Second)
	}
}

// runBroadcastListener opens a dedicated pq listener connection and feeds
// received notifications into the local Bus via publishLocal.
// Returns an error when the connection is lost (caller retries).
func runBroadcastListener(bus *Bus, dsn string) error {
	reportErr := func(ev pq.ListenerEventType, err error) {
		if err != nil {
			log.Printf("[pg_broadcast] pq listener event=%d err=%v", ev, err)
		}
	}

	l := pq.NewListener(dsn, 5*time.Second, time.Minute, reportErr)
	if err := l.Listen(pgNotifyChannel); err != nil {
		return fmt.Errorf("LISTEN %s: %w", pgNotifyChannel, err)
	}
	defer l.Close()

	log.Printf("[pg_broadcast] listener ready on channel %q", pgNotifyChannel)

	for {
		select {
		case n, ok := <-l.Notify:
			if !ok {
				return fmt.Errorf("notify channel closed")
			}
			if n == nil {
				// Keepalive ping from pq — ignore.
				continue
			}

			var e Event
			if err := json.Unmarshal([]byte(n.Extra), &e); err != nil {
				log.Printf("[pg_broadcast] malformed notification: %v", err)
				continue
			}

			// publishLocal bypasses the broadcastHook — no re-broadcast loop.
			bus.publishLocal(e)
		}
	}
}