thermograph/backend/daemon/internal/gateway/gateway.go

// Package gateway holds Discord's **gateway** bot — it reacts to @mentions
// (and DMs) in real time. The other two Discord paths need no persistent
// connection: slash commands arrive as signed HTTP POSTs and the daily post
// goes out over an incoming webhook, both handled on the Python side. This one
// is different — to answer an *ordinary* message that mentions the bot, we
// have to hold a live websocket to Discord's gateway and listen for
// MESSAGE_CREATE events.
//
// Single-connection invariant. Discord permits exactly one gateway connection
// per bot token (per shard). The Python predecessor enforced that with a
// leader election inside web/app.py; here it holds by deployment shape — the
// daemon runs as exactly one replica, and this package owns its only
// connection.
//
// Least privilege / no privileged intent. We subscribe to GUILDS +
// GUILD_MESSAGES + DIRECT_MESSAGES and act only on messages that @mention the
// bot or DM it. Discord delivers message *content* for exactly those cases
// even without the privileged MESSAGE_CONTENT intent, so the bot needs no
// privileged-intent portal review — a real design constraint, not trivia. To
// trigger on arbitrary keywords in channels, enable MESSAGE_CONTENT in the
// Developer Portal and extend triage; the rest of the machinery is unchanged.
//
// The gateway owns NO climate/grading logic. A mention with a city query is
// answered by calling back into the Python app (Grader / apiclient), which
// shares one grade builder with the /grade slash command so the two never
// drift.
package gateway

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"math/rand"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/coder/websocket"

	"thermograph/daemon/internal/config"
)

// Run is the package entrypoint: build a Bot from the resolved config and
// drive it for ctx's lifetime. The caller has already checked
// cfg.DiscordEnabled — an unconfigured deploy never gets this far, so an
// unconfigured deploy pays nothing.
func Run(ctx context.Context, cfg *config.Config, api Grader) error {
	return New(cfg.DiscordToken, api).Run(ctx)
}

// v10 JSON gateway. On a resume we reconnect to the session's
// resume_gateway_url (from the READY payload) instead, with the same query
// string appended.
const (
	gatewayURL = "wss://gateway.discord.gg/?v=10&encoding=json"
	gatewayQS  = "/?v=10&encoding=json"
)

// Gateway intents (bitfield). GUILDS gives guild/channel context;
// GUILD_MESSAGES and DIRECT_MESSAGES deliver message-create events in servers
// and DMs. We deliberately do NOT request MESSAGE_CONTENT (privileged) —
// Discord still includes content for messages that mention us or DM us, which
// is all we act on.
const (
	intentGuilds         = 1 << 0
	intentGuildMessages  = 1 << 9
	intentDirectMessages = 1 << 12
	intents              = intentGuilds | intentGuildMessages | intentDirectMessages
)

// Gateway opcodes (discord.com/developers/docs/topics/gateway-events#opcodes).
const (
	opDispatch       = 0
	opHeartbeat      = 1
	opIdentify       = 2
	opResume         = 6
	opReconnect      = 7
	opInvalidSession = 9
	opHello          = 10
	opHeartbeatACK   = 11
)

// isFatalClose reports whether a close code is unrecoverable — reconnecting
// just loops, so we stop: 4004 auth failed (bad token), 4010 invalid shard,
// 4011 sharding required, 4012 invalid API version, 4013 invalid intents,
// 4014 disallowed (privileged) intents.
func isFatalClose(code websocket.StatusCode) bool {
	switch code {
	case 4004, 4010, 4011, 4012, 4013, 4014:
		return true
	}
	return false
}

// Reconnect backoff bounds. A server-requested reconnect (op 7) resets to the
// start so a planned reconnect comes back promptly.
const (
	backoffStart = 1 * time.Second
	backoffMax   = 60 * time.Second
)

// nextBackoff doubles the reconnect delay, capped at backoffMax.
func nextBackoff(d time.Duration) time.Duration {
	d *= 2
	if d > backoffMax {
		d = backoffMax
	}
	return d
}

// maxConcurrentReplies bounds in-flight MESSAGE_CREATE handling. The Python
// version pushed each message onto a thread (asyncio.to_thread) so a slow
// grading lookup couldn't stall heartbeats, but its queue was unbounded; the
// semaphore here keeps the same never-block-the-read-loop property while
// making sure a flood of mentions can't spawn unbounded goroutines against
// the backend.
const maxConcurrentReplies = 4

// dialTimeout mirrors the Python's open_timeout=30.
const dialTimeout = 30 * time.Second

// maxFrameSize mirrors the Python's max_size=2**20 — a grade embed is tiny,
// so anything near a megabyte is not a frame we want to buffer.
const maxFrameSize = 1 << 20

const (
	modeIdentify = "identify"
	modeResume   = "resume"
)

// Grader is the one slice of the internal API the gateway needs: a city query
// in, gateway-ready Discord message JSON out. Satisfied by *apiclient.Client.
type Grader interface {
	Grade(ctx context.Context, query string) (json.RawMessage, error)
}

// Bot owns the single gateway connection and everything hanging off it.
type Bot struct {
	token string
	api   Grader
	sem   chan struct{} // bounds concurrent MESSAGE_CREATE handlers
	rest  string        // Discord REST base; swapped for a test server in tests
	http  *http.Client  // REST replies (not the gateway socket)
}

// New returns a Bot that authenticates with token and answers city queries
// through api.
func New(token string, api Grader) *Bot {
	return &Bot{
		token: token,
		api:   api,
		sem:   make(chan struct{}, maxConcurrentReplies),
		rest:  discordAPIBase,
		http:  &http.Client{Timeout: 30 * time.Second},
	}
}

// payload is the gateway envelope: opcode, event data, sequence number,
// dispatch type.
type payload struct {
	Op int             `json:"op"`
	D  json.RawMessage `json:"d"`
	S  *int64          `json:"s"`
	T  string          `json:"t"`
}

// session is the per-connection state that must survive a resume: the last
// sequence number (what RESUME replays from), the resume token/URL, and our
// own user id (for mention detection). seq is mutex-guarded because the read
// loop advances it while the heartbeat goroutine reads it; the other fields
// are only ever touched from the read/run loop.
type session struct {
	mu        sync.Mutex
	seq       int64
	haveSeq   bool
	sessionID string
	resumeURL string
	userID    string
}

func (s *session) setSeq(v int64) {
	s.mu.Lock()
	s.seq, s.haveSeq = v, true
	s.mu.Unlock()
}

func (s *session) lastSeq() (int64, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.seq, s.haveSeq
}

// wsConn is a thin JSON send/receive wrapper. coder/websocket allows
// concurrent calls to everything except Read, and only the read loop reads,
// so no extra locking is needed for the heartbeat goroutine's writes.
type wsConn struct {
	c *websocket.Conn
}

func (w *wsConn) send(ctx context.Context, v any) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}
	return w.c.Write(ctx, websocket.MessageText, data)
}

func (w *wsConn) read(ctx context.Context) ([]byte, error) {
	_, data, err := w.c.Read(ctx)
	return data, err
}

func (w *wsConn) close(code websocket.StatusCode, reason string) {
	// Best-effort: the connection may already be gone, and a failed close on
	// a dead socket changes nothing about what happens next.
	_ = w.c.Close(code, reason)
}

// jitter is the 0..1s randomness added to every reconnect sleep so a fleet of
// clients dropped by the same incident doesn't reconnect in lockstep.
func jitter() time.Duration {
	return time.Duration(rand.Float64() * float64(time.Second))
}

// sleepCtx sleeps for d, returning false if ctx was cancelled first.
func sleepCtx(ctx context.Context, d time.Duration) bool {
	select {
	case <-ctx.Done():
		return false
	case <-time.After(d):
		return true
	}
}

// Run maintains the gateway connection for the daemon's lifetime: connect,
// IDENTIFY (or RESUME), dispatch events, and reconnect with capped
// exponential backoff. It returns nil on context cancellation and an error
// only on a fatal close (bad token / disallowed intents) — reconnecting after
// those loops forever, so the caller must not restart us.
func (b *Bot) Run(ctx context.Context) error {
	sess := &session{}
	mode := modeIdentify
	backoff := backoffStart
	for {
		useResume := mode == modeResume && sess.resumeURL != ""
		url := gatewayURL
		if useResume {
			url = sess.resumeURL
		}
		dialCtx, cancel := context.WithTimeout(ctx, dialTimeout)
		conn, _, err := websocket.Dial(dialCtx, url, nil)
		cancel()
		if err != nil {
			if ctx.Err() != nil {
				log.Print("gateway: shutting down")
				return nil
			}
			// A failed dial carries no gateway close code, so it is never
			// fatal — back off and retry.
			mode = modeIdentify
			if sess.sessionID != "" {
				mode = modeResume
			}
			log.Printf("gateway: connect failed (%v); reconnecting in %s", err, backoff)
			if !sleepCtx(ctx, backoff+jitter()) {
				return nil
			}
			backoff = nextBackoff(backoff)
			continue
		}
		conn.SetReadLimit(maxFrameSize)
		ws := &wsConn{c: conn}

		next, err := b.connection(ctx, ws, sess, useResume)
		if ctx.Err() != nil {
			// Shutdown: close clean (1000) — we are gone for good, so the
			// session should not be held open for a RESUME that never comes.
			ws.close(websocket.StatusNormalClosure, "shutting down")
			log.Print("gateway: shutting down")
			return nil
		}
		if err != nil {
			if code := websocket.CloseStatus(err); isFatalClose(code) {
				log.Printf("gateway: fatal gateway close %d; stopping", code)
				return fmt.Errorf("gateway: fatal close code %d", code)
			}
			mode = modeIdentify
			if sess.sessionID != "" {
				mode = modeResume
			}
			log.Printf("gateway: connection dropped (%v); reconnecting in %s", err, backoff)
			if !sleepCtx(ctx, backoff+jitter()) {
				return nil
			}
			backoff = nextBackoff(backoff)
			continue
		}

		// Deliberate reconnect (op 7 / op 9): close it ourselves. Non-1000
		// when we intend to RESUME — Discord invalidates the session on a
		// 1000/1001 close, which would force a fresh IDENTIFY.
		if next == modeResume {
			ws.close(websocket.StatusCode(4000), "resuming")
		} else {
			ws.close(websocket.StatusNormalClosure, "reidentifying")
		}
		mode = next
		backoff = backoffStart // planned reconnect: come back promptly
	}
}

// parseHello validates the op-10 HELLO that opens every connection and returns
// its heartbeat interval.
//
// Every failure here is an ERROR rather than a clean "reidentify". That
// distinction is load-bearing: a clean return is Run's PLANNED-reconnect path,
// which resets the backoff to 1s and redials with no sleep at all. That is
// correct for a server-requested reconnect (op 7), but a gateway persistently
// answering with a malformed HELLO would turn it into a tight reconnect loop
// against Discord -- precisely what the capped backoff exists to prevent. The
// error path applies backoff and still picks the next mode from surviving
// session state.
func parseHello(raw []byte) (time.Duration, error) {
	var hello payload
	if err := json.Unmarshal(raw, &hello); err != nil {
		return 0, fmt.Errorf("gateway: unparseable HELLO: %w", err)
	}
	if hello.Op != opHello {
		return 0, fmt.Errorf("gateway: expected HELLO (op %d), got op %d", opHello, hello.Op)
	}
	var hd struct {
		HeartbeatInterval float64 `json:"heartbeat_interval"`
	}
	if err := json.Unmarshal(hello.D, &hd); err != nil {
		return 0, fmt.Errorf("gateway: unparseable HELLO payload: %w", err)
	}
	// A non-positive interval would busy-loop the heartbeat goroutine.
	if hd.HeartbeatInterval <= 0 {
		return 0, fmt.Errorf("gateway: invalid heartbeat_interval %v", hd.HeartbeatInterval)
	}
	return time.Duration(hd.HeartbeatInterval * float64(time.Millisecond)), nil
}

// connection drives one live connection until it ends. A clean return carries
// the mode for the next connect ("resume" to reconnect and RESUME, "identify"
// to start a fresh session); an error return means the socket died and the
// caller picks the next mode from the surviving session state.
func (b *Bot) connection(ctx context.Context, ws *wsConn, sess *session, resume bool) (string, error) {
	raw, err := ws.read(ctx)
	if err != nil {
		return "", err
	}
	interval, err := parseHello(raw)
	if err != nil {
		return modeIdentify, err
	}

	// The heartbeat lives exactly as long as this connection.
	hbCtx, stopHB := context.WithCancel(ctx)
	defer stopHB()
	acked := &atomic.Bool{}
	acked.Store(true)
	beatNow := make(chan struct{}, 1)
	go b.heartbeat(hbCtx, ws, interval, sess, acked, beatNow)

	seq, haveSeq := sess.lastSeq()
	if resume && sess.sessionID != "" && haveSeq {
		err = ws.send(ctx, map[string]any{"op": opResume, "d": map[string]any{
			"token":      b.token,
			"session_id": sess.sessionID,
			"seq":        seq,
		}})
	} else {
		err = ws.send(ctx, map[string]any{"op": opIdentify, "d": map[string]any{
			"token":      b.token,
			"intents":    intents,
			"properties": map[string]string{"os": "linux", "browser": "thermograph", "device": "thermograph"},
		}})
	}
	if err != nil {
		return "", err
	}

	for {
		data, err := ws.read(ctx)
		if err != nil {
			return "", err
		}
		var p payload
		if err := json.Unmarshal(data, &p); err != nil {
			log.Printf("gateway: dropping unparseable frame: %v", err)
			continue
		}
		if p.S != nil {
			// Every payload that carries a sequence number advances it — it
			// is what RESUME replays from.
			sess.setSeq(*p.S)
		}
		switch p.Op {
		case opDispatch:
			b.dispatch(ctx, p, sess)
		case opHeartbeat:
			// Server asked for an immediate beat; route it through the
			// heartbeat goroutine, which owns beat sending.
			select {
			case beatNow <- struct{}{}:
			default:
			}
		case opHeartbeatACK:
			acked.Store(true)
		case opReconnect:
			return modeResume, nil // planned reconnect — resume
		case opInvalidSession:
			// d == true => the session can still be resumed; d == false =>
			// it is dead and the next connect must IDENTIFY fresh.
			var resumable bool
			_ = json.Unmarshal(p.D, &resumable)
			if !resumable {
				sess.sessionID = ""
				sess.resumeURL = ""
			}
			// Discord wants a short randomized wait before re-authing.
			if !sleepCtx(ctx, time.Second+jitter()) {
				return "", ctx.Err()
			}
			if sess.sessionID != "" {
				return modeResume, nil
			}
			return modeIdentify, nil
		}
	}
}

// heartbeat sends op-1 heartbeats every interval, the first one jittered by
// interval*rand per Discord's guidance. If the previous beat was never ACKed
// the link is a zombie — a half-open TCP connection where our writes
// "succeed" but nothing comes back — so close with a non-1000 code (1000
// would invalidate the session) and stop; the next connect can then RESUME.
func (b *Bot) heartbeat(ctx context.Context, ws *wsConn, interval time.Duration, sess *session, acked *atomic.Bool, beatNow <-chan struct{}) {
	timer := time.NewTimer(time.Duration(rand.Float64() * float64(interval)))
	defer timer.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-beatNow:
			// Server-requested beat (op 1): send immediately, outside the
			// zombie accounting — it is not the reply to one of ours.
			if err := b.sendHeartbeat(ctx, ws, sess); err != nil {
				return // socket gone; the read loop owns reconnecting
			}
		case <-timer.C:
			if !acked.Load() {
				ws.close(websocket.StatusCode(4000), "heartbeat ack missed")
				return
			}
			acked.Store(false)
			if err := b.sendHeartbeat(ctx, ws, sess); err != nil {
				return
			}
			timer.Reset(interval)
		}
	}
}

func (b *Bot) sendHeartbeat(ctx context.Context, ws *wsConn, sess *session) error {
	var d any // null until the first sequence number arrives
	if seq, ok := sess.lastSeq(); ok {
		d = seq
	}
	return ws.send(ctx, map[string]any{"op": opHeartbeat, "d": d})
}

// dispatch handles an op-0 DISPATCH event we care about.
func (b *Bot) dispatch(ctx context.Context, p payload, sess *session) {
	switch p.T {
	case "READY":
		var d struct {
			SessionID        string `json:"session_id"`
			ResumeGatewayURL string `json:"resume_gateway_url"`
			User             struct {
				ID snowflake `json:"id"`
			} `json:"user"`
		}
		if err := json.Unmarshal(p.D, &d); err != nil {
			log.Printf("gateway: unparseable READY: %v", err)
			return
		}
		sess.sessionID = d.SessionID
		sess.resumeURL = ""
		if r := strings.TrimRight(d.ResumeGatewayURL, "/"); r != "" {
			sess.resumeURL = r + gatewayQS
		}
		sess.userID = string(d.User.ID)
		log.Printf("gateway: ready (user %s)", sess.userID)
	case "MESSAGE_CREATE":
		// Until READY tells us our own user id we can't detect mentions (or
		// guard against answering ourselves), so stay silent.
		if sess.userID == "" {
			return
		}
		var m gwMessage
		if err := json.Unmarshal(p.D, &m); err != nil {
			log.Printf("gateway: unparseable MESSAGE_CREATE: %v", err)
			return
		}
		// Grading is a synchronous parquet/DB round trip on the Python side,
		// so it must never run on the read loop — a slow lookup would stall
		// heartbeats (the Python used asyncio.to_thread for the same
		// reason). Each message gets its own goroutine, with the semaphore
		// bounding in-flight replies so a flood of mentions can't spawn
		// unbounded goroutines against the backend — an improvement over
		// to_thread's unbounded queue. Over the bound we drop and log rather
		// than block the read loop.
		botID := sess.userID
		select {
		case b.sem <- struct{}{}:
			go func() {
				defer func() { <-b.sem }()
				b.handleMessage(ctx, &m, botID)
			}()
		default:
			log.Printf("gateway: dropping message %s: too many replies in flight", m.ID)
		}
	}
}